import os import torch import json import folder_paths from datetime import datetime import comfy.sample import comfy.samplers import comfy.utils import comfy.model_management import latent_preview # ============================================================ # 1. StepwiseKSampler # comfy.sample.sample()을 그대로 사용하되, # callback에서 스텝별 latent를 저장합니다. # ============================================================ class StepwiseKSampler: @classmethod def INPUT_TYPES(cls): return { "required": { "model": ("MODEL",), "positive": ("CONDITIONING",), "negative": ("CONDITIONING",), "latent_image": ("LATENT",), "seed": ("INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff}), "steps": ("INT", {"default": 20, "min": 1, "max": 10000}), "cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1}), "sampler_name": (comfy.samplers.KSampler.SAMPLERS,), "scheduler": (comfy.samplers.KSampler.SCHEDULERS,), "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), "save_steps": ("STRING", {"default": "all", "multiline": False}), "save_prefix": ("STRING", {"default": "stepwise", "multiline": False}), } } RETURN_TYPES = ("LATENT", "STRING",) RETURN_NAMES = ("latent", "save_dir",) FUNCTION = "sample" CATEGORY = "sampling/stepwise" OUTPUT_NODE = True def sample(self, model, positive, negative, latent_image, seed, steps, cfg, sampler_name, scheduler, denoise, save_steps, save_prefix): # save_steps 파싱 if save_steps.strip().lower() == "all": steps_to_save = set(range(1, steps + 1)) else: steps_to_save = set() for s in save_steps.split(","): s = s.strip() if s.isdigit(): steps_to_save.add(int(s)) # 저장 디렉토리 생성 output_dir = folder_paths.get_output_directory() timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") save_dir = os.path.join(output_dir, f"{save_prefix}_{timestamp}") os.makedirs(save_dir, exist_ok=True) # latent 준비 — fix_empty_latent_channels 호출 (FLUX 16ch 대응) latent = latent_image.copy() latent_samples = latent["samples"] latent_samples = comfy.sample.fix_empty_latent_channels(model, latent_samples) latent["samples"] = latent_samples # 노이즈 생성 batch_inds = latent.get("batch_index", None) noise = comfy.sample.prepare_noise(latent_samples, seed, batch_inds) # noise_mask noise_mask = None if "noise_mask" in latent: noise_mask = latent["noise_mask"] # 콜백: latent_preview 기본 콜백 + 우리의 저장 로직 x0_output = {} preview_callback = latent_preview.prepare_callback(model, steps, x0_output) def step_callback(step, x0, x, total_steps): # 기본 프리뷰 콜백 호출 if preview_callback is not None: preview_callback(step, x0, x, total_steps) step_num = step + 1 # 0-indexed -> 1-indexed if step_num in steps_to_save: save_path = os.path.join(save_dir, f"step_{step_num:04d}.pt") save_data = { "samples": x.cpu().clone(), "step": step_num, "total_steps": total_steps, } if x0 is not None: save_data["x0_prediction"] = x0.cpu().clone() torch.save(save_data, save_path) # comfy.sample.sample() 호출 — KSampler와 동일한 경로 samples = comfy.sample.sample( model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_samples, denoise=denoise, disable_noise=False, start_step=None, last_step=None, force_full_denoise=True, noise_mask=noise_mask, sigmas=None, callback=step_callback, disable_pbar=False, seed=seed ) # 최종 latent 저장 torch.save({ "samples": samples.cpu().clone(), "step": steps, "total_steps": steps, }, os.path.join(save_dir, "step_final.pt")) # 메타데이터 저장 metadata = { "seed": seed, "steps": steps, "cfg": cfg, "sampler": sampler_name, "scheduler": scheduler, "denoise": denoise, "saved_steps": sorted(list(steps_to_save)), "timestamp": timestamp, } with open(os.path.join(save_dir, "metadata.json"), "w") as f: json.dump(metadata, f, indent=2) print(f"[StepwiseKSampler] Saved {len(steps_to_save)} steps to {save_dir}") out_latent = latent.copy() out_latent["samples"] = samples return (out_latent, save_dir,) # ============================================================ # 2. StepwiseLatentLoader # ============================================================ class StepwiseLatentLoader: @classmethod def INPUT_TYPES(cls): return { "required": { "save_dir": ("STRING", {"default": "", "multiline": False}), "step_number": ("INT", {"default": 1, "min": 1, "max": 10000}), "load_type": (["noisy_sample", "x0_prediction"],), } } RETURN_TYPES = ("LATENT",) RETURN_NAMES = ("latent",) FUNCTION = "load" CATEGORY = "sampling/stepwise" def load(self, save_dir, step_number, load_type): file_path = os.path.join(save_dir, f"step_{step_number:04d}.pt") if not os.path.exists(file_path): file_path = os.path.join(save_dir, "step_final.pt") if not os.path.exists(file_path): raise FileNotFoundError(f"Step {step_number} not found in {save_dir}") data = torch.load(file_path, map_location="cpu", weights_only=True) if load_type == "x0_prediction" and "x0_prediction" in data: samples = data["x0_prediction"] else: samples = data["samples"] return ({"samples": samples},) # ============================================================ # 3. StepwiseLatentPreview # ============================================================ class StepwiseLatentPreview: @classmethod def INPUT_TYPES(cls): return { "required": { "vae": ("VAE",), "save_dir": ("STRING", {"default": "", "multiline": False}), "preview_steps": ("STRING", {"default": "1,5,10,15,20", "multiline": False}), "use_x0": ("BOOLEAN", {"default": True}), } } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("grid_image",) FUNCTION = "preview" CATEGORY = "sampling/stepwise" def preview(self, vae, save_dir, preview_steps, use_x0): step_numbers = [] for s in preview_steps.split(","): s = s.strip() if s.isdigit(): step_numbers.append(int(s)) step_numbers.sort() decoded_images = [] for step_num in step_numbers: file_path = os.path.join(save_dir, f"step_{step_num:04d}.pt") if not os.path.exists(file_path): continue data = torch.load(file_path, map_location="cpu", weights_only=True) if use_x0 and "x0_prediction" in data: samples = data["x0_prediction"] else: samples = data["samples"] decoded = vae.decode(samples) decoded_images.append(decoded) if not decoded_images: raise ValueError(f"No valid steps found in {save_dir}") grid = torch.cat(decoded_images, dim=2) # W 방향으로 concat return (grid,) # ============================================================ # 등록 # ============================================================ NODE_CLASS_MAPPINGS = { "StepwiseKSampler": StepwiseKSampler, "StepwiseLatentLoader": StepwiseLatentLoader, "StepwiseLatentPreview": StepwiseLatentPreview, } NODE_DISPLAY_NAME_MAPPINGS = { "StepwiseKSampler": "Stepwise KSampler (Save Latents)", "StepwiseLatentLoader": "Stepwise Latent Loader", "StepwiseLatentPreview": "Stepwise Latent Preview Grid", }