190 lines
7.1 KiB
Python
190 lines
7.1 KiB
Python
import torch
|
|
import comfy.samplers
|
|
import comfy.sample
|
|
import comfy.model_management
|
|
import comfy.utils
|
|
import latent_preview
|
|
|
|
|
|
def _common_ksampler(model, seed, steps, cfg, sampler_name, scheduler,
|
|
positive, negative, latent,
|
|
denoise=1.0, disable_noise=False,
|
|
start_step=None, last_step=None, force_full_denoise=False):
|
|
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.get("batch_index", None)
|
|
noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
|
|
|
|
noise_mask = latent.get("noise_mask", None)
|
|
|
|
callback = latent_preview.prepare_callback(model, steps)
|
|
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
|
|
|
|
|
|
class DualModelKSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_a": ("MODEL",),
|
|
"model_b": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent_image": ("LATENT",),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 200}),
|
|
"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,),
|
|
"switch_step": ("INT", {"default": 5, "min": 1, "max": 199,
|
|
"tooltip": "Phase1->Phase2 switch step"}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
},
|
|
"optional": {
|
|
"switch_end": ("INT", {"default": 0, "min": 0, "max": 199,
|
|
"tooltip": "0=2-phase, >0=3-phase sandwich"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "STRING",)
|
|
RETURN_NAMES = ("latent", "info",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling/custom"
|
|
|
|
def sample(self, model_a, model_b, positive, negative, latent_image,
|
|
seed, steps, cfg, sampler_name, scheduler, switch_step, denoise,
|
|
switch_end=0):
|
|
|
|
sandwich_mode = switch_end > switch_step
|
|
|
|
if sandwich_mode:
|
|
info = f"Sandwich: A(1-{switch_step}) -> B({switch_step+1}-{switch_end}) -> A({switch_end+1}-{steps})"
|
|
else:
|
|
info = f"2-Phase: A(1-{switch_step}) -> B({switch_step+1}-{steps})"
|
|
|
|
print(f"\n{'='*50}\nDualModelKSampler: {info}\n{'='*50}")
|
|
|
|
start_idx = max(0, int(round(steps * (1.0 - denoise))))
|
|
|
|
# Phase 1: model_a
|
|
print(f"\n Phase 1: model_a step {start_idx+1} ~ {switch_step}")
|
|
phase1_latent = _common_ksampler(
|
|
model_a, seed, steps, cfg, sampler_name, scheduler,
|
|
positive, negative, latent_image,
|
|
denoise=denoise, disable_noise=False,
|
|
start_step=start_idx, last_step=switch_step,
|
|
force_full_denoise=False,
|
|
)
|
|
|
|
# Phase 2: model_b
|
|
p2_end = switch_end if sandwich_mode else steps
|
|
print(f"\n Phase 2: model_b step {switch_step+1} ~ {p2_end}")
|
|
phase2_latent = _common_ksampler(
|
|
model_b, seed, steps, cfg, sampler_name, scheduler,
|
|
positive, negative, phase1_latent,
|
|
denoise=denoise, disable_noise=True,
|
|
start_step=switch_step, last_step=p2_end,
|
|
force_full_denoise=(not sandwich_mode),
|
|
)
|
|
|
|
# Phase 3: model_a (sandwich only)
|
|
if sandwich_mode:
|
|
print(f"\n Phase 3: model_a step {switch_end+1} ~ {steps}")
|
|
phase3_latent = _common_ksampler(
|
|
model_a, seed, steps, cfg, sampler_name, scheduler,
|
|
positive, negative, phase2_latent,
|
|
denoise=denoise, disable_noise=True,
|
|
start_step=switch_end, last_step=steps,
|
|
force_full_denoise=True,
|
|
)
|
|
final_latent = phase3_latent
|
|
else:
|
|
final_latent = phase2_latent
|
|
|
|
print(f"\n DualModelKSampler done!\n")
|
|
return (final_latent, info,)
|
|
|
|
|
|
class DualModelKSamplerBatch:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_a": ("MODEL",),
|
|
"model_b": ("MODEL",),
|
|
"positive": ("CONDITIONING",),
|
|
"negative": ("CONDITIONING",),
|
|
"latent_image": ("LATENT",),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff}),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 200}),
|
|
"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,),
|
|
"switch_steps_csv": ("STRING", {"default": "3,5,7,10,15",
|
|
"tooltip": "Comma-separated switch points"}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "STRING",)
|
|
RETURN_NAMES = ("latents", "info",)
|
|
FUNCTION = "sample_batch"
|
|
CATEGORY = "sampling/custom"
|
|
|
|
def sample_batch(self, model_a, model_b, positive, negative, latent_image,
|
|
seed, steps, cfg, sampler_name, scheduler, switch_steps_csv, denoise):
|
|
|
|
switch_points = [int(x.strip()) for x in switch_steps_csv.split(",")]
|
|
single_node = DualModelKSampler()
|
|
all_samples = []
|
|
all_info = []
|
|
|
|
for sp in switch_points:
|
|
result_latent, result_info = single_node.sample(
|
|
model_a, model_b, positive, negative, latent_image,
|
|
seed, steps, cfg, sampler_name, scheduler, sp, denoise,
|
|
switch_end=0,
|
|
)
|
|
all_samples.append(result_latent["samples"])
|
|
all_info.append(f"sw={sp}")
|
|
|
|
batched = torch.cat(all_samples, dim=0)
|
|
info_str = " | ".join(all_info)
|
|
return ({"samples": batched}, info_str,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DualModelKSampler": DualModelKSampler,
|
|
"DualModelKSamplerBatch": DualModelKSamplerBatch,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DualModelKSampler": "Dual Model KSampler (Step Switch)",
|
|
"DualModelKSamplerBatch": "Dual Model KSampler Batch",
|
|
}
|