Added Ksampler Fun which outputs preview steps to the output folder for comfy deploy

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