250 lines
8.5 KiB
Python
250 lines
8.5 KiB
Python
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",
|
|
}
|