285 lines
14 KiB
Python
285 lines
14 KiB
Python
import torch
|
|
import numpy as np
|
|
import comfy.samplers
|
|
import comfy.sample
|
|
import comfy.utils
|
|
import traceback
|
|
import os
|
|
import folder_paths
|
|
|
|
class CFGExplorer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent": ("LATENT",),
|
|
"vae": ("VAE",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg_start": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.5}),
|
|
"cfg_end": ("FLOAT", {"default": 20.0, "min": 0.0, "max": 100.0, "step": 0.5}),
|
|
"cfg_steps": ("INT", {"default": 5, "min": 2, "max": 50}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler"}),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "normal"}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "LATENT",)
|
|
RETURN_NAMES = ("images", "latents")
|
|
FUNCTION = "generate"
|
|
CATEGORY = "VariationLab"
|
|
|
|
def generate(self, model, positive, negative, latent, vae, seed, steps, cfg_start, cfg_end, cfg_steps,
|
|
sampler_name, scheduler, denoise):
|
|
try:
|
|
# Generate CFG values
|
|
cfg_values = np.linspace(cfg_start, cfg_end, cfg_steps)
|
|
|
|
results = []
|
|
latent_results = []
|
|
|
|
# Generate images for each CFG value
|
|
noise = comfy.sample.prepare_noise(latent["samples"], seed)
|
|
noise_mask = None
|
|
|
|
for cfg in cfg_values:
|
|
# Sample the latent
|
|
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent["samples"],
|
|
denoise=denoise, noise_mask=noise_mask)
|
|
latent_results.append({"samples": samples.clone()})
|
|
|
|
# Decode the latent
|
|
image = vae.decode(samples)
|
|
results.append(image)
|
|
|
|
# Stack all images into a batch
|
|
batched_images = torch.cat(results, dim=0)
|
|
batched_latents = {"samples": torch.cat([x["samples"] for x in latent_results], dim=0)}
|
|
|
|
return (batched_images, batched_latents)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
# Return a small error image
|
|
error_image = torch.zeros(1, 64, 64, 3)
|
|
error_latent = {"samples": torch.zeros(1, 4, 8, 8)}
|
|
return (error_image, error_latent)
|
|
|
|
|
|
class StepExplorer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent": ("LATENT",),
|
|
"vae": ("VAE",),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps_start": ("INT", {"default": 5, "min": 1, "max": 500}),
|
|
"steps_end": ("INT", {"default": 50, "min": 1, "max": 500}),
|
|
"step_count": ("INT", {"default": 5, "min": 2, "max": 20}),
|
|
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.5}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler"}),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "normal"}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "LATENT",)
|
|
RETURN_NAMES = ("images", "latents")
|
|
FUNCTION = "generate"
|
|
CATEGORY = "VariationLab"
|
|
|
|
def generate(self, model, positive, negative, latent, vae, seed, steps_start, steps_end, step_count,
|
|
cfg, sampler_name, scheduler, denoise):
|
|
try:
|
|
# Generate step values (rounded to integers)
|
|
step_values = np.linspace(steps_start, steps_end, step_count, dtype=int)
|
|
|
|
results = []
|
|
latent_results = []
|
|
|
|
# Generate images for each step count
|
|
noise = comfy.sample.prepare_noise(latent["samples"], seed)
|
|
noise_mask = None
|
|
|
|
for steps in step_values:
|
|
# Sample the latent
|
|
samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent["samples"],
|
|
denoise=denoise, noise_mask=noise_mask)
|
|
latent_results.append({"samples": samples.clone()})
|
|
|
|
# Decode the latent
|
|
image = vae.decode(samples)
|
|
results.append(image)
|
|
|
|
# Stack all images into a batch
|
|
batched_images = torch.cat(results, dim=0)
|
|
batched_latents = {"samples": torch.cat([x["samples"] for x in latent_results], dim=0)}
|
|
|
|
return (batched_images, batched_latents)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
# Return a small error image
|
|
error_image = torch.zeros(1, 64, 64, 3)
|
|
error_latent = {"samples": torch.zeros(1, 4, 8, 8)}
|
|
return (error_image, error_latent)
|
|
|
|
|
|
class CheckpointExplorer:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
checkpoints = folder_paths.get_filename_list("checkpoints")
|
|
return {
|
|
"required": {
|
|
"latent": ("LATENT",),
|
|
"base_positive_prompt": ("STRING", {"multiline": True, "default": "a photograph of a person"}),
|
|
"base_negative_prompt": ("STRING", {"multiline": True, "default": "blurry, low quality"}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.5}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
},
|
|
"optional": {
|
|
"checkpoint1": (checkpoints,),
|
|
"checkpoint1_clip_skip": ("INT", {"default": -1, "min": -24, "max": -1, "step": 1}),
|
|
"checkpoint1_positive_suffix": ("STRING", {"multiline": True, "default": ""}),
|
|
"checkpoint1_negative_suffix": ("STRING", {"multiline": True, "default": ""}),
|
|
"checkpoint1_sampler": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler"}),
|
|
"checkpoint1_scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "normal"}),
|
|
"checkpoint1_steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"checkpoint1_cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.5}),
|
|
|
|
"checkpoint2": (checkpoints,),
|
|
"checkpoint2_clip_skip": ("INT", {"default": -1, "min": -24, "max": -1, "step": 1}),
|
|
"checkpoint2_positive_suffix": ("STRING", {"multiline": True, "default": ""}),
|
|
"checkpoint2_negative_suffix": ("STRING", {"multiline": True, "default": ""}),
|
|
"checkpoint2_sampler": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler"}),
|
|
"checkpoint2_scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "normal"}),
|
|
"checkpoint2_steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"checkpoint2_cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.5}),
|
|
|
|
"checkpoint3": (checkpoints,),
|
|
"checkpoint3_clip_skip": ("INT", {"default": -1, "min": -24, "max": -1, "step": 1}),
|
|
"checkpoint3_positive_suffix": ("STRING", {"multiline": True, "default": ""}),
|
|
"checkpoint3_negative_suffix": ("STRING", {"multiline": True, "default": ""}),
|
|
"checkpoint3_sampler": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler"}),
|
|
"checkpoint3_scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "normal"}),
|
|
"checkpoint3_steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"checkpoint3_cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.5}),
|
|
|
|
"checkpoint4": (checkpoints,),
|
|
"checkpoint4_clip_skip": ("INT", {"default": -1, "min": -24, "max": -1, "step": 1}),
|
|
"checkpoint4_positive_suffix": ("STRING", {"multiline": True, "default": ""}),
|
|
"checkpoint4_negative_suffix": ("STRING", {"multiline": True, "default": ""}),
|
|
"checkpoint4_sampler": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler"}),
|
|
"checkpoint4_scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "normal"}),
|
|
"checkpoint4_steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"checkpoint4_cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.5}),
|
|
|
|
"checkpoint5": (checkpoints,),
|
|
"checkpoint5_clip_skip": ("INT", {"default": -1, "min": -24, "max": -1, "step": 1}),
|
|
"checkpoint5_positive_suffix": ("STRING", {"multiline": True, "default": ""}),
|
|
"checkpoint5_negative_suffix": ("STRING", {"multiline": True, "default": ""}),
|
|
"checkpoint5_sampler": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler"}),
|
|
"checkpoint5_scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "normal"}),
|
|
"checkpoint5_steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
|
"checkpoint5_cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.5}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "VariationLab"
|
|
|
|
def generate(self, latent, base_positive_prompt, base_negative_prompt, seed, steps, cfg, denoise, **kwargs):
|
|
try:
|
|
results = []
|
|
checkpoint_info = []
|
|
|
|
# Collect checkpoint info
|
|
for i in range(1, 6):
|
|
checkpoint_key = f"checkpoint{i}"
|
|
if checkpoint_key in kwargs and kwargs[checkpoint_key] is not None:
|
|
checkpoint_name = kwargs[checkpoint_key]
|
|
clip_skip = kwargs.get(f"{checkpoint_key}_clip_skip", -1)
|
|
pos_suffix = kwargs.get(f"{checkpoint_key}_positive_suffix", "")
|
|
neg_suffix = kwargs.get(f"{checkpoint_key}_negative_suffix", "")
|
|
|
|
# Get checkpoint-specific sampler settings or use defaults
|
|
sampler = kwargs.get(f"{checkpoint_key}_sampler", "euler")
|
|
scheduler = kwargs.get(f"{checkpoint_key}_scheduler", "normal")
|
|
checkpoint_steps = kwargs.get(f"{checkpoint_key}_steps", steps)
|
|
checkpoint_cfg = kwargs.get(f"{checkpoint_key}_cfg", cfg)
|
|
|
|
checkpoint_info.append({
|
|
"name": checkpoint_name,
|
|
"clip_skip": clip_skip,
|
|
"positive_prompt": base_positive_prompt + (f", {pos_suffix}" if pos_suffix else ""),
|
|
"negative_prompt": base_negative_prompt + (f", {neg_suffix}" if neg_suffix else ""),
|
|
"sampler": sampler,
|
|
"scheduler": scheduler,
|
|
"steps": checkpoint_steps,
|
|
"cfg": checkpoint_cfg
|
|
})
|
|
|
|
# No checkpoints specified
|
|
if not checkpoint_info:
|
|
error_image = torch.zeros(1, 64, 64, 3)
|
|
return (error_image,)
|
|
|
|
# Generate image for each checkpoint
|
|
for info in checkpoint_info:
|
|
checkpoint_path = folder_paths.get_full_path("checkpoints", info["name"])
|
|
|
|
# Load model
|
|
model, clip, vae = comfy.sd.load_checkpoint_guess_config(checkpoint_path)
|
|
|
|
# Apply clip skip
|
|
if info["clip_skip"] != -1:
|
|
clip = comfy.sd.CLIP(clip.tokenizer, clip.transformer, clip.device, int(info["clip_skip"]))
|
|
|
|
# Encode positive and negative prompts
|
|
positive = comfy.sd.encode_prompt(clip, info["positive_prompt"])
|
|
negative = comfy.sd.encode_prompt(clip, info["negative_prompt"])
|
|
|
|
# Sample the latent
|
|
noise = comfy.sample.prepare_noise(latent["samples"], seed)
|
|
samples = comfy.sample.sample(model, noise, info["steps"], info["cfg"],
|
|
info["sampler"], info["scheduler"],
|
|
positive, negative, latent["samples"],
|
|
denoise=denoise, noise_mask=None)
|
|
|
|
# Decode the latent
|
|
image = vae.decode(samples)
|
|
results.append(image)
|
|
|
|
# Stack all images into a batch
|
|
batched_images = torch.cat(results, dim=0)
|
|
|
|
return (batched_images,)
|
|
except Exception as e:
|
|
traceback.print_exc()
|
|
# Return a small error image
|
|
error_image = torch.zeros(1, 64, 64, 3)
|
|
return (error_image,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CFGExplorer": CFGExplorer,
|
|
"StepExplorer": StepExplorer,
|
|
"CheckpointExplorer": CheckpointExplorer,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CFGExplorer": "VariationLab: CFG Explorer",
|
|
"StepExplorer": "VariationLab: Step Explorer",
|
|
"CheckpointExplorer": "VariationLab: Checkpoint Explorer",
|
|
}
|