nodes/DevDuckFace-KsamplerDF
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

v1.0.2

Browse Source 
v1.0.2
This commit is contained in:
DuckFace
2026-01-27 23:06:01 -03:00
committed by GitHub
parent e53a88a7d0
commit 4c80ed99c5
1 changed files with 113 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes.py
+113 -36
View File
@@ -5,16 +5,24 @@ import comfy.sample
import comfy.utils
import comfy.model_management
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
import folder_paths
from spandrel import ModelLoader, ImageModelDescriptor
class KsamplerDF(ComfyNodeABC):
"""
Custom KSampler with Refiner functionality for ComfyUI.
Custom KSampler with Refiner and Upscale 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.
Includes latent upscale between phases and optional model upscale at the end.
"""
upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]
@classmethod
def INPUT_TYPES(cls) -> InputTypeDict:
# Get list of available upscale models
upscale_models_list = ["none"] + folder_paths.get_filename_list("upscale_models")
return {
"required": {
"model": (IO.MODEL, {"tooltip": "The model used for denoising the input latent."}),
@@ -23,26 +31,56 @@ class KsamplerDF(ComfyNodeABC):
"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."}),
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale for main phase."}),
"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."}),
# Latent Upscale between main and refiner
"latent_upscale_enabled": (["enable", "disable"], {"default": "disable", "tooltip": "Enable latent upscale between main and refiner phases."}),
"latent_upscale_method": (cls.upscale_methods, {"default": "bislerp", "tooltip": "Upscale method for latent."}),
"latent_upscale_factor": ("FLOAT", {"default": 1.5, "min": 0.1, "max": 8.0, "step": 0.1, "tooltip": "Scale factor for latent upscale."}),
# Refiner settings
"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."}),
"cfg_refiner": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale for 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)."}),
# VAE for decoding latent to image
"vae": ("VAE", {"tooltip": "VAE for decoding latent to image output."}),
# Model Upscale at the end - select model directly
"upscale_model_name": (upscale_models_list, {"default": "none", "tooltip": "Select upscale model (e.g., RealESRGAN, ESRGAN). Select 'none' to disable."}),
},
}
RETURN_TYPES = (IO.LATENT,)
OUTPUT_TOOLTIPS = ("The denoised latent image.",)
RETURN_TYPES = (IO.LATENT, "IMAGE")
RETURN_NAMES = ("LATENT", "IMAGE")
OUTPUT_TOOLTIPS = ("The denoised latent image.", "The decoded/upscaled image.")
FUNCTION = "sample"
CATEGORY = "sampling"
DESCRIPTION = "KSampler with Refiner - Splits steps between main sampler and refiner sampler with independent denoise control."
DESCRIPTION = "KSampler with Refiner and Upscale - Splits steps between main and refiner, with latent upscale and optional model upscale."
def load_upscale_model(self, model_name):
"""Load an upscale model from the upscale_models folder."""
model_path = folder_paths.get_full_path_or_raise("upscale_models", model_name)
sd = comfy.utils.load_torch_file(model_path, safe_load=True)
if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
sd = comfy.utils.state_dict_prefix_replace(sd, {"module.": ""})
out = ModelLoader().load_from_state_dict(sd).eval()
if not isinstance(out, ImageModelDescriptor):
raise Exception("Upscale model must be a single-image model.")
return out
def sample(self, model, positive, negative, latent_image, seed, steps, cfg,
sampler_name, scheduler, denoise, refiner_at_step, sampler_refiner,
scheduler_refiner, denoise_refiner):
sampler_name, scheduler, denoise,
latent_upscale_enabled, latent_upscale_method, latent_upscale_factor,
refiner_at_step, sampler_refiner, scheduler_refiner, cfg_refiner, denoise_refiner,
vae, upscale_model_name):
import nodes
@@ -54,10 +92,11 @@ class KsamplerDF(ComfyNodeABC):
refiner_steps = steps - refiner_at_step
current_latent = latent_image
output_image = None
# 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}")
print(f"[KSampler DF] Main phase: {main_steps} steps with {sampler_name}/{scheduler}, CFG={cfg}, denoise={denoise}")
result = nodes.common_ksampler(
model=model,
seed=seed,
@@ -76,14 +115,24 @@ class KsamplerDF(ComfyNodeABC):
)
current_latent = result[0]
# Phase 2: Refiner sampling (if refiner_steps > 0 and denoise_refiner > 0)
# Phase 2: Latent Upscale (if enabled and before refiner)
if latent_upscale_enabled == "enable" and refiner_steps > 0:
print(f"[KSampler DF] Latent upscale: {latent_upscale_method} x{latent_upscale_factor}")
samples = current_latent["samples"]
width = round(samples.shape[-1] * latent_upscale_factor)
height = round(samples.shape[-2] * latent_upscale_factor)
upscaled = comfy.utils.common_upscale(samples, width, height, latent_upscale_method, "disabled")
current_latent = current_latent.copy()
current_latent["samples"] = upscaled
# Phase 3: 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}")
print(f"[KSampler DF] Refiner phase: {refiner_steps} steps with {sampler_refiner}/{scheduler_refiner}, CFG={cfg_refiner}, denoise={denoise_refiner}")
result = nodes.common_ksampler(
model=model,
seed=seed,
steps=refiner_steps,
cfg=cfg,
cfg=cfg_refiner,
sampler_name=sampler_refiner,
scheduler=scheduler_refiner,
positive=positive,
@@ -96,31 +145,59 @@ class KsamplerDF(ComfyNodeABC):
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)")
print(f"[KSampler DF] Sampling complete! Total: {steps} steps ({main_steps} main + {refiner_steps} refiner)")
# Decode latent to image using VAE
samples = current_latent["samples"]
output_image = vae.decode(samples)
print(f"[KSampler DF] Decoded latent to image")
# Phase 4: Model Upscale (if a model is selected)
if upscale_model_name != "none":
print(f"[KSampler DF] Loading upscale model: {upscale_model_name}")
upscale_model = self.load_upscale_model(upscale_model_name)
print(f"[KSampler DF] Applying model upscale...")
# Apply upscale model
device = comfy.model_management.get_torch_device()
memory_required = comfy.model_management.module_size(upscale_model.model)
memory_required += (512 * 512 * 3) * output_image.element_size() * max(upscale_model.scale, 1.0) * 384.0
memory_required += output_image.nelement() * output_image.element_size()
comfy.model_management.free_memory(memory_required, device)
upscale_model.to(device)
in_img = output_image.movedim(-1, -3).to(device)
tile = 512
overlap = 32
oom = True
try:
while oom:
try:
steps_upscale = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(
in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap
)
pbar = comfy.utils.ProgressBar(steps_upscale)
s = comfy.utils.tiled_scale(
in_img, lambda a: upscale_model(a),
tile_x=tile, tile_y=tile, overlap=overlap,
upscale_amount=upscale_model.scale, pbar=pbar
)
oom = False
except comfy.model_management.OOM_EXCEPTION as e:
tile //= 2
if tile < 128:
raise e
finally:
upscale_model.to("cpu")
output_image = torch.clamp(s.movedim(-3, -1), min=0, max=1.0)
print(f"[KSampler DF] Model upscale complete!")
return (current_latent,)
return (current_latent, output_image)
# Register the node
@@ -129,5 +206,5 @@ NODE_CLASS_MAPPINGS = {
}
NODE_DISPLAY_NAME_MAPPINGS = {
"KsamplerDF": "KSampler DF (with Refiner)",
"KsamplerDF": "KSampler DF (with Refiner & Upscale)",
}