diff --git a/__pycache__/ksampler_fun.cpython-312.pyc b/__pycache__/ksampler_fun.cpython-312.pyc index a6c7905..fd9340c 100644 Binary files a/__pycache__/ksampler_fun.cpython-312.pyc and b/__pycache__/ksampler_fun.cpython-312.pyc differ 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