134 lines
6.2 KiB
Python
134 lines
6.2 KiB
Python
from __future__ import annotations
|
|
import torch
|
|
import comfy.samplers
|
|
import comfy.sample
|
|
import comfy.utils
|
|
import comfy.model_management
|
|
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
|
|
|
class KsamplerDF(ComfyNodeABC):
|
|
"""
|
|
Custom KSampler with Refiner functionality for ComfyUI.
|
|
Splits steps between main sampler and refiner sampler.
|
|
Example: 12 total steps with refiner at step 6 = 6 main + 6 refiner steps.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> InputTypeDict:
|
|
return {
|
|
"required": {
|
|
"model": (IO.MODEL, {"tooltip": "The model used for denoising the input latent."}),
|
|
"positive": (IO.CONDITIONING, {"tooltip": "The conditioning describing the attributes to include."}),
|
|
"negative": (IO.CONDITIONING, {"tooltip": "The conditioning describing the attributes to exclude."}),
|
|
"latent_image": (IO.LATENT, {"tooltip": "The latent image to denoise."}),
|
|
"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": "Total number of steps (main + refiner)."}),
|
|
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale."}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler", "tooltip": "The algorithm used for the main sampling phase."}),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "normal", "tooltip": "The scheduler for the main sampling phase."}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Main denoise amount."}),
|
|
"refiner_at_step": ("INT", {"default": 10, "min": 0, "max": 10000, "tooltip": "Step at which to switch to refiner sampler."}),
|
|
"sampler_refiner": (comfy.samplers.KSampler.SAMPLERS, {"default": "euler", "tooltip": "The algorithm used for the refiner phase."}),
|
|
"scheduler_refiner": (comfy.samplers.KSampler.SCHEDULERS, {"default": "normal", "tooltip": "The scheduler for the refiner phase."}),
|
|
"denoise_refiner": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Refiner denoise - controls how much the refiner interferes (0.0 = no effect, 1.0 = full effect)."}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (IO.LATENT,)
|
|
OUTPUT_TOOLTIPS = ("The denoised latent image.",)
|
|
FUNCTION = "sample"
|
|
CATEGORY = "sampling"
|
|
DESCRIPTION = "KSampler with Refiner - Splits steps between main sampler and refiner sampler with independent denoise control."
|
|
|
|
def sample(self, model, positive, negative, latent_image, seed, steps, cfg,
|
|
sampler_name, scheduler, denoise, refiner_at_step, sampler_refiner,
|
|
scheduler_refiner, denoise_refiner):
|
|
|
|
import nodes
|
|
|
|
# Clamp refiner_at_step to valid range
|
|
refiner_at_step = max(0, min(refiner_at_step, steps))
|
|
|
|
# Calculate steps for each phase
|
|
main_steps = refiner_at_step
|
|
refiner_steps = steps - refiner_at_step
|
|
|
|
current_latent = latent_image
|
|
|
|
# Phase 1: Main sampling (if main_steps > 0)
|
|
if main_steps > 0:
|
|
print(f"[KSampler DF] Main phase: {main_steps} steps with {sampler_name}/{scheduler}, denoise={denoise}")
|
|
result = nodes.common_ksampler(
|
|
model=model,
|
|
seed=seed,
|
|
steps=main_steps,
|
|
cfg=cfg,
|
|
sampler_name=sampler_name,
|
|
scheduler=scheduler,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=current_latent,
|
|
denoise=denoise,
|
|
disable_noise=False,
|
|
start_step=None,
|
|
last_step=None,
|
|
force_full_denoise=False # Don't fully denoise, refiner continues
|
|
)
|
|
current_latent = result[0]
|
|
|
|
# Phase 2: Refiner sampling (if refiner_steps > 0 and denoise_refiner > 0)
|
|
if refiner_steps > 0 and denoise_refiner > 0:
|
|
print(f"[KSampler DF] Refiner phase: {refiner_steps} steps with {sampler_refiner}/{scheduler_refiner}, denoise={denoise_refiner}")
|
|
result = nodes.common_ksampler(
|
|
model=model,
|
|
seed=seed,
|
|
steps=refiner_steps,
|
|
cfg=cfg,
|
|
sampler_name=sampler_refiner,
|
|
scheduler=scheduler_refiner,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=current_latent,
|
|
denoise=denoise_refiner,
|
|
disable_noise=False,
|
|
start_step=None,
|
|
last_step=None,
|
|
force_full_denoise=True # Final phase, full denoise
|
|
)
|
|
current_latent = result[0]
|
|
elif main_steps > 0:
|
|
# If no refiner, we need to do a final denoise on main result
|
|
# Run one more step to finalize
|
|
print(f"[KSampler DF] Finalizing without refiner...")
|
|
result = nodes.common_ksampler(
|
|
model=model,
|
|
seed=seed,
|
|
steps=1,
|
|
cfg=cfg,
|
|
sampler_name=sampler_name,
|
|
scheduler=scheduler,
|
|
positive=positive,
|
|
negative=negative,
|
|
latent=current_latent,
|
|
denoise=0.0,
|
|
disable_noise=True,
|
|
start_step=None,
|
|
last_step=None,
|
|
force_full_denoise=True
|
|
)
|
|
current_latent = result[0]
|
|
|
|
print(f"[KSampler DF] Complete! Total: {steps} steps ({main_steps} main + {refiner_steps} refiner)")
|
|
|
|
return (current_latent,)
|
|
|
|
|
|
# Register the node
|
|
NODE_CLASS_MAPPINGS = {
|
|
"KsamplerDF": KsamplerDF,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"KsamplerDF": "KSampler DF (with Refiner)",
|
|
}
|