From 90d346367781f51f13c5a9bfa46491c2afa1bb1f Mon Sep 17 00:00:00 2001
From: nofunstudio <jimmy@nofun.io>
Date: Wed, 9 Apr 2025 10:27:18 -0600
Subject: [PATCH] Added Ksampler Fun which outputs preview steps to the output
 folder for comfy deploy

---
 __pycache__/ksampler_fun.cpython-312.pyc | Bin 7276 -> 7832 bytes
 ksampler_fun.py                          |  87 +++++++++++++----------
 2 files changed, 49 insertions(+), 38 deletions(-)

diff --git a/__pycache__/ksampler_fun.cpython-312.pyc b/__pycache__/ksampler_fun.cpython-312.pyc
index a6c79057fb4eb4511755004a8ac950bfe792749c..fd9340c9882628fc1848d4c5eee8a7f7c86a6329 100644
GIT binary patch
delta 3459
zcmbtWU2Idy6`p(TYhVBW+p!a42S|bg`Qfi2OOhp!K=`Sf<+lY^?v3wFu;bWa?oGDI
zUT=bgDz>fEM57g<rGiS^78<lfi(0Wy-Aa}E(3f~a70cy;2P$=?stpB|7WJWLuI<1g
zq&{>kpP4!5%$YN1zI)F6b<aY#{SP*q89~F>|2guf(7e5dxijCCoiU60%gk9Piw+@?
zyNpD`B1^RpI_uJSM&R-5Jhy0)+$rX)M>L5h=)F^H%<?5wk1-{Skj>^4K}qGZavbK|
z;W1%EO1h}H(!9{?5MrMG8g{vgRGt@^An@`o;?zY*)abFRNOe{B1NH#3iE)u7y_d`~
zrhqGLt(tn-t1fbrsSmod+eE#}Pa@2zyr`#Uvy6!H9?kT>8SlF;l~Fy5us>~iH^{sq
z=vG#6nj>D7Rh`oisk~41iH6Y-L`*k=^8V~T#Zz91kZ4qWs$&|?u<mJq46@!))d|Z3
z)h!m8@`AiyWkn-3s`@7x>{b0~y4P#1s#9IXAf>jdqVlR=HK}M(U*_a%6!LdAT$A2V
z)p*%kQ9cMq-o4!ozUWlF%Bcp`sJc|MYEiAKO|`3DwMGr9%pHEwU9lHZ`E)(0H`e-U
z)sR{%GF>2m9u2c6P|QSruIn{oii|HA@-s^YuEcF^HkLRghm#jbueI0d7{t<KN_y{r
zkjV@Q$x(928g`ZR$=vu<tggh3<wPliedH5s*WQvLBPdc<Dd{E!3EROcc2MG^gu(?=
zw@SA0vc-g~NaJ$JERUwfiQm@ZbYN*rNM%#mk+O(f<a}+uSG5zVtRmq_A%p3iVG28J
zpf+{62}Q1h5xHdZlR2A#JmZelScfOFiBVY?8_!4>?*TD&O1?sVW&3Dm&42=R3+aM?
z%j`FbM*l+n7slFRI8xw4kE5Gz_THSi-%@Bgv=SXy=ED!e^&dCfGB5Qn?;Tv(_|8iB
z_%a_VTAXv+KkQttYb{t>i`EUR)`o($;ihy;_-X2awf&g^IU@^`%T1m4cYe;TI1bJZ
zJPy=deSiM_)j&rfKt8j-$38{mZ@P{;Tn}y~IUv<W@t}QB{~<|zTL^9&B`N?%NHJB_
z(q_2ZkPn133R+WW$Y~=H&d)KbK5ZZuc@N_TP0vgv3BUp^X;3{CrD+^rYqM4*P-Thi
zOYBaeDHPMiu@QvRN6!eDUe{aEcrIGGr|6iZT*!%0VNytCgrSVIEjkR)0j=neust~%
zRdUhciA*MXy0<?%oWs?1a?Faifl<7Rk`_uDfRG<?+h^=$TqP=7lx!oCQk_V+WUGk5
z+L)k}OcgbdaN`&d6XQyHS;`E*$-YdZ9;FR;z^KG00And7k(jO-tOB2;4UU<eAeAK`
zv#98=yQ2SrrMPj|75xJ~QuO#?$`^eh=*$m8jjN%(h0xxWP&dpSkAqFC!QF-6?j`B|
zso$LY?A&Vi=|cDEmEakW+KU#)4|_iBS+>+aw7Qnvi9cGuy{1FfjkptbdKD2CL~mGw
z#COPF3=wXQ!4Z|Xm>48bp6oRqV=vW{i^k@V2-%^UzRtsfydiJQ^I2I%kWRhiUyP4j
zH<~%4YD{}5df2p&CWr_yMBOxmBfWm{uw$huHRK^%N5Z`6s8I=5xo=8SU`5N(RmkLx
zkc%17@MFDb{E04aQljgtH)_CaAniKuo^)fC{jyh0)pO*{;BD-sq&iZ~q`-IB_;@)B
z=#LEcA0iG@L%**KBMnfSK^xP=Try^Ji4iP_ct4$O5Tzt!uS5ksO0Ej2M2Z$T+(OQn
z-inmioLth63(AF(KAlTtaSvUqmxW0wX2E^ru4&hdy@I$z1%aGqpu{9f{BSBGWrZ=R
zJfdb~3mDK4S0?MYOTuUJ72K5*Zfhb@me7sm7eD}k7RTJKRZnxl)BJ^{`JG4M`Wxr2
zom&mJ6v8bl;kGOB;+EKz_(CV-BS=H;z}4`4_#*&n;A+o&&$6eb*tYZT2e&_1Z989R
zJHPyHa#<Q)X&Whcx6Q?G_TS$Ah<x9C{F@uL-rBXKyVbnZw_Mvf$3LtM-{`p3aii;6
z*Dd{B$8E>`nw8r9b9~Y1yK0^{FYI~X-26O<-1T#jH5;;SeQwppFJv3I{9E0DO?~aU
z-?ejnon*Hq*3}P8EC9KUG^}e;(DI;0c7Q}%cp3eJHE{y-)bHUM%9{0LhXb2w0nk2&
z{r5g`v@$_XG{Da$Hj0kx`e9aVx^7D`P~*0cKRZ8W!sL6dI%bso!gbg=Mz_vVLV=7&
zKXn_Ky4RRzKvLpxZXzq<(R-Uban|<+4LnM|>pM6z-5H-48<KD|HyqV!8qIOEszrw+
zK!gNL0V^~EQa%7@J?FK?Y+9DXKI#em1aK22TPUG;$IX;%rKFvb7D}|r+Db=lKw@qV
z`=KxKl|m=u5fb({FdgK8ztM0KMDoPFod4_0zG7taQvLnj&!shlzQvqmpCHt9iiJt!
zH2b919<;CNoKEu^a+=IqsADWv=r)wVcpG^X*g12ik{6>I$mOU&4+w`6qhbn6NhLEy
zkEvXcqEM5^p)5m|&xjD637LlAWN-h8J}vPNNeUG6=*0;klTxNyF-<#_LxAG+FHgym
ztl1gIxyh6W8_>rFQEQZr|79ylLsV_FAZke2LHJ<PjB2~XEZWLKzy}+T>0Xa0qrDYk
zcH9SQ|MvyS)S2}UzEK;M;Z9FH?Kpa({6J|f-`_Ye<!<u(V9OD>nH)RCK10YkO|Qu`
z&1&~TU7n<_y+;YH=}<g}=r~D<2&805W&~L-7h|lF9iihhbURE-*ehS$2%S=L0{-$o
zAeT{*b3S2RocXB^8C+|0wAP5cTV{_HJ?*o{)>sr6WZ})^2+R&Vas+1wfCDd!*`v?&
z3|H|`GdXQNIFXH3j~DAPYEB=F#}B^ZxUXbRB!(xH2`nWN81AW(tL@3$*zi;(CRoOP
z(pmepZu~YVl&o2rw=YPUaS0^Ub4=g0I8MnR)zA$dK8E#lN?(2DyfQ*Qt=+xptLJ~g
zcLAs-GExto14a6<mnCAZYYlH?0?!AS2-ENc0eQZk*}zQ02Se%FesZL4gPxY?{{oqx
BPs#uQ

delta 2841
zcmbtWTTC0-89p<fvB%)cfG=RKc7S9j5VM3XAqmMY3t>qhge2QGi@J4O&v0>!P0tJ@
zps}D5E79!@dQ=6pTfs^zl@)2b+w>u?k^0c8QiV&@Sfi*)A1hTYEKyUX5AA=(p0Jzl
zL#rOy-?{w%`Okm;+nLoff9dr6*6l7q&_dt;``B+PRy>u=>dKLHs~~fW%!tUMH;~LP
zBH6yhQZ9r>%4CO9Ho=UP%MRHAy>FJ4oGbbICb7U^q94$P5VOAt+ioLWkeLuHWkvE=
zrk<Q;f~BG*5Z1oNG%zArU}|k*(nZ?bKA8QIiG<41Ei$JIGYIn_fuj=AjEu77Af<g?
z)ET{e3r?XeTL$Ys8Tg9aym_8a^A+d_uAsBJcMkM<ezIQ+SX8aT8p#ge^ILMan8JEC
zptG`l4osy7W*F?(18g!xRXtr?wip>?+DrhOIiDXCWRV`!1wEiUEy39e4UwO^YKYzD
z=UgaTX<Fv8E!b8D=i78n=XJa8&_%sOFV$VTTleTby+W_lnL8z0<wdPRx{$0VzZ0s0
z*(yDxSIJBlP@qSn>{TR{l3&}p9e<{Gz$`f~oH+ZeX>>XryP?J>r&0=zcBL|jcuGAR
z9mdMcgmSYdo=SZqp16_V3I~`-$6!Wg|BvJyp_SF_<h{}waF^euxh=M&hiY+1W~@8w
zC?;4$&0&p_<zTbG@MlA%NRrt*Hft|~W|71V8D*B9v8jT3>YEZ{2+FMqRf}bANF|0n
z71xxsM*hSH{I;36Vz3$2;HKi*xWOee6KV1{XIqmn8NZ>(6IhivgHx5%D5gPyX^dhT
z04lv8Ib$E57h<uA^n@0R%|-t!Ges@_Mk?k4zyiuwg_nB2JCHxzxzzi>)sU~Kf@%Mw
zy72ng+Sqz(Ep?~sVe??Fd2pwBXnQ!ZQzrxc(j%AW2LsZ#2UfXt@0$1H#y_}D?1MOz
z;3*cH9eUFdHKen&-iANb$tPt&NGBh8EC}R#rGsRz)Gz8bD;TqO-JbMP&yaSPpEC8+
zC7GQAu}Nx$7n(+mVzy+^8XfcpVzm&B-4+mKqj`2ffP7_S{yk2%|JaswXf@X6WF3jP
zd`yJ8pE*Pm$L+A(ofULDX@zxxWZfqmjv-7hO76Rx8hiyWQB4glk(rt`crAky<Ax)h
ziH+g7j02?EbE*zgZ86PZTuGtY+@Qv16v>He$)e{=c79?qKBkbICot~>YS?PHj`D)%
zArVsli{}gc&9Ga|0f6VdzGdxUMKlLLZ}f4b=_l16R6mTg<sxl6k@jUzzUA1mXLUB`
z5v3<~6s*4e_R8Boisi|Y_sX?Ht()4%J)6mU+;;t$<^D(Yk@cHvH`nzweY5xHmp;05
z@7hj%-*SK6AG&>E<-+Q%2mY6y`cSZGIkxLYV(U|vIew<nguKtb-qFk2e#P>=rR2CM
zy-^Q~_yoX52#4v|yt4O&oX1`uHNanW0W6|VS>D;X?-89+VN(|oxfE39EO(U0hk=Fs
zjO+d(UnAlsm`K$IKbeu06z(O*_^y^hj?whOXDFPd@G6DZD0BlD?qXV#5BX#Bf#M{l
zYRVL(p{C%Oc&gxkTGa6C6nZGo@GxwdX-%bmwjB8dec__6&*n1zI&AU4U%dro5uD}=
ztx7pxBl(4WDGqL_4R0_TuWVk-HJse3IknA&9))V(S2nsgrnjR#J9WJ~p>x|@Fz;|J
zee+wh+m%Oij-z>@>Y)(H36YJa&BzDR1EJ-qy^@#q5QPpT*00WNA31rieaHJM*>b$Y
z;xM@_;u2cJ%{TJ+U0+$0aUOo(HNim9K0-#yKV$~TzshSEnbi2si=)7a$0!U?knc|W
z98Ap%tF1sXc$}F|%fPYfA7TSvz=37*&tTtW$%QXdQN0wdP`FCrO$uM8@D&PQrSKMo
z13zAasl?-R&>2GU8By_d^81P=rjGoh;;@~TLG`V>jg=Rf#=Vd`<k_`(#gbj*b(WZc
zB(Zps%A^2eANeHIezm(8jA1iN)o>guVVb^jSj&X9aU~4io&txr{>PIL9?hg=c*=0!
zfl3qUvAOS%i&gJdMibAUC?He6n<jsmc&d-g-(j|Bs6zv@vqPYa8LH+L88c8MH?9Tp
z|6?pOB|jINW>kW@<bJi!u7G&8gBaDv&OSj%Ot8BMITH{guue6~(Br27z-66ue3HT`
z0HZXKimU3RqK#){JWVdwT$!gfz^5sMDbPH}*DW{?1C`!z5#@Ptj}`fneH*fu?b6Y1
zBl5qrFqrqf3=o`MxRm!+Ec8G2h8FsFS>(IUE(|>37`~{PDRj=&H=Pa-TrGS{NOqj*
z0O5T@LwyB>{YFVFHae|MV<i^DV0UDe4a1YjOpeY*p9jx5a=o^r$M6*1a_g-(TxlpL
zv2i6erC_yC;OSd~B?`3c8n)r{gP5aJn$!@4JGHMgeE#;&gf6H`(<$XFeiwM?mykL_
en(LzFCz*<;GfXX$VD}Kk-Pg!WT^$z&j(-C{L#%iJ

diff --git a/ksampler_fun.py b/ksampler_fun.py
index 701f759..6f2615d 100644
--- a/ksampler_fun.py
+++ b/ksampler_fun.py
@@ -11,7 +11,7 @@ import folder_paths
 import numpy as np
 from PIL import Image
 
-def fun_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, vae, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, preview_steps=5):
+def fun_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, vae, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, preview_steps=5, skip_steps=0):
     latent_image = latent["samples"]
     latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image)
 
@@ -25,41 +25,18 @@ def fun_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, neg
     if "noise_mask" in latent:
         noise_mask = latent["noise_mask"]
 
-    # Calculate which steps to capture previews
-    preview_step_interval = max(1, steps // preview_steps)
-    preview_steps_list = list(range(0, steps, preview_step_interval))
-    if steps - 1 not in preview_steps_list:
-        preview_steps_list.append(steps - 1)
+    # Calculate which steps to capture previews, starting after skip_steps
+    remaining_steps = steps - skip_steps
+    if remaining_steps <= 0:
+        preview_steps_list = []
+    else:
+        preview_step_interval = max(1, remaining_steps // preview_steps)
+        preview_steps_list = list(range(skip_steps, steps, preview_step_interval))
+        if steps - 1 not in preview_steps_list:
+            preview_steps_list.append(steps - 1)
 
     # Create a custom callback to capture intermediate images
-    class PreviewCallback:
-        def __init__(self, latent, preview_steps_list, vae):
-            self.latent = latent
-            self.preview_steps_list = preview_steps_list
-            self.vae = vae
-            self.output_dir = os.path.join(folder_paths.get_output_directory(), "ksampler_previews")
-            os.makedirs(self.output_dir, exist_ok=True)
-
-        def __call__(self, step, denoised, x, total_steps):
-            if step in self.preview_steps_list:
-                # Create a copy of the current state
-                current_latent = self.latent.copy()
-                current_latent["samples"] = x
-                
-                # Decode the latent to image
-                with torch.no_grad():
-                    image = self.vae.decode(x)
-                    image = image.cpu().numpy()
-                    image = (image * 255).astype(np.uint8)
-                    image = Image.fromarray(image[0])
-                    
-                    # Save the image
-                    filename = f"preview_step_{step:03d}.png"
-                    filepath = os.path.join(self.output_dir, filename)
-                    image.save(filepath)
-
-    # Set up the callback with VAE
-    callback = PreviewCallback(latent, preview_steps_list, vae)
+    callback = PreviewCallback(latent, preview_steps_list, vae, model)
 
     disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
     samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
@@ -70,6 +47,39 @@ def fun_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, neg
     out["samples"] = samples
     return (out, None)  # Return None for preview_latents since we're saving directly
 
+class PreviewCallback:
+    def __init__(self, latent, preview_steps_list, vae, model):
+        self.latent = latent
+        self.preview_steps_list = preview_steps_list
+        self.vae = vae
+        self.model = model
+        self.output_dir = folder_paths.get_output_directory()
+        # Initialize previewer using the model's latent_format
+        self.previewer = latent_preview.get_previewer("cpu", model.model.latent_format)
+        if self.previewer is None:
+            print("Warning: No preview method available, falling back to full VAE for previews")
+
+    def __call__(self, step, denoised, x, total_steps):
+        if step in self.preview_steps_list:
+            # Create a copy of the current state
+            current_latent = self.latent.copy()
+            current_latent["samples"] = x
+            
+            # Decode the latent to image using previewer if available, otherwise fall back to full VAE
+            with torch.no_grad():
+                if self.previewer is not None:
+                    preview_format, image, _ = self.previewer.decode_latent_to_preview_image("JPEG", x)
+                else:
+                    image = self.vae.decode(x)
+                    image = image.cpu().numpy()
+                    image = (image * 255).astype(np.uint8)
+                    image = Image.fromarray(image[0])
+                
+                # Save the image
+                filename = f"preview_step_{step:03d}.png"
+                filepath = os.path.join(self.output_dir, filename)
+                image.save(filepath)
+
 class KSampler:
     @classmethod
     def INPUT_TYPES(s):
@@ -87,19 +97,20 @@ class KSampler:
                 "vae": ("VAE", {"tooltip": "The VAE used for decoding the latent images."}),
                 "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising applied, lower values will maintain the structure of the initial image allowing for image to image sampling."}),
                 "preview_steps": ("INT", {"default": 5, "min": 1, "max": 20, "tooltip": "Number of intermediate previews to generate during sampling."}),
+                "skip_steps": ("INT", {"default": 0, "min": 0, "max": 10000, "tooltip": "Number of steps to skip before starting to save previews."}),
             }
         }
 
     RETURN_TYPES = ("LATENT", "LATENT")
     RETURN_NAMES = ("latent", "preview_latents")
-    OUTPUT_TOOLTIPS = ("The final denoised latent.", "Preview images are saved to the outputs/ksampler_previews folder.")
+    OUTPUT_TOOLTIPS = ("The final denoised latent.", "Preview images are saved directly to the outputs folder using TAESD for better quality.")
     FUNCTION = "sample"
 
     CATEGORY = "sampling"
-    DESCRIPTION = "Uses the provided model, positive and negative conditioning to denoise the latent image. Saves intermediate previews to the outputs folder."
+    DESCRIPTION = "Uses the provided model, positive and negative conditioning to denoise the latent image. Saves intermediate previews to the outputs folder using TAESD for better quality."
 
-    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, vae, denoise=1.0, preview_steps=5):
-        return fun_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, vae, denoise=denoise, preview_steps=preview_steps)
+    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, vae, denoise=1.0, preview_steps=5, skip_steps=0):
+        return fun_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, vae, denoise=denoise, preview_steps=preview_steps, skip_steps=skip_steps)
 
 NODE_CLASS_MAPPINGS = {
     "Fun KSampler": KSampler