166 lines
8.6 KiB
Python
166 lines
8.6 KiB
Python
from __future__ import annotations
|
|
import torch
|
|
import comfy.samplers
|
|
import comfy.sample
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
import latent_preview
|
|
import node_helpers
|
|
import os
|
|
import folder_paths
|
|
import numpy as np
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
import json
|
|
|
|
# Add debug print for output directory
|
|
output_dir = folder_paths.get_output_directory()
|
|
print(f"Debug - Output directory path: {output_dir}")
|
|
print(f"Debug - Output directory exists: {os.path.exists(output_dir)}")
|
|
print(f"Debug - Output directory is writable: {os.access(output_dir, os.W_OK) if os.path.exists(output_dir) else False}")
|
|
|
|
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)
|
|
|
|
if disable_noise:
|
|
noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
|
|
else:
|
|
batch_inds = latent["batch_index"] if "batch_index" in latent else None
|
|
noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
|
|
|
|
noise_mask = None
|
|
if "noise_mask" in latent:
|
|
noise_mask = latent["noise_mask"]
|
|
|
|
# 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
|
|
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,
|
|
denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step,
|
|
force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
|
|
|
|
out = latent.copy()
|
|
out["samples"] = samples
|
|
return (out, callback.get_results()) # Return the results from the callback
|
|
|
|
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()
|
|
self.type = "output"
|
|
self.prefix_append = ""
|
|
self.compress_level = 4
|
|
self.results = [] # Store results for later retrieval
|
|
# 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 get_results(self):
|
|
return {"ui": {"images": self.results}}
|
|
|
|
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])
|
|
|
|
# Use the same approach as ComfyDeployOutputImage for saving files
|
|
filename_prefix = f"preview_step_{step:03d}"
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
|
filename_prefix,
|
|
self.output_dir,
|
|
image.size[0], # width
|
|
image.size[1] # height
|
|
)
|
|
|
|
# Create metadata
|
|
metadata = PngInfo()
|
|
metadata.add_text("step", str(step))
|
|
metadata.add_text("total_steps", str(total_steps))
|
|
|
|
# Save the image with metadata
|
|
file = f"{filename}_{counter:05}_.png"
|
|
file_path = os.path.join(full_output_folder, file)
|
|
|
|
# Save with PNG compression
|
|
image.save(file_path, pnginfo=metadata, compress_level=self.compress_level)
|
|
|
|
# Create result entry similar to ComfyDeployOutputImage
|
|
result = {
|
|
"filename": file,
|
|
"subfolder": subfolder,
|
|
"type": self.type,
|
|
"output_id": "preview_images"
|
|
}
|
|
|
|
# Add to results list
|
|
self.results.append(result)
|
|
|
|
print(f"Saved preview image to: {file_path}")
|
|
|
|
class KSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True, "tooltip": "The random seed used for creating the noise."}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}),
|
|
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The algorithm used when sampling, this can affect the quality, speed, and style of the generated output."}),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
|
|
"positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}),
|
|
"negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
|
|
"latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}),
|
|
"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 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 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, 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
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Fun KSampler": "Fun KSampler"
|
|
}
|