From 0a72baba13cc3e19cdd77c25508810a3a1084f93 Mon Sep 17 00:00:00 2001 From: GodOfNothing Date: Sun, 23 Mar 2025 15:35:26 +0300 Subject: [PATCH 1/6] Added lcm_scalewise --- comfy/k_diffusion/sampling.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index 5b8d8000d..c9f4c04d4 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -929,6 +929,30 @@ def sample_lcm(model, x, sigmas, extra_args=None, callback=None, disable=None, n return x +# x0 = + +@torch.no_grad() +def sample_lcm_scalewise(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None): + extra_args = {} if extra_args is None else extra_args + seed = extra_args.get("seed", None) + scales = extra_args.get("scales", None) + if scales: + assert len(scales) == len(sigmas) - 1, "Number of scales must be equal to number of sampling steps minus one." + noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler + s_in = x.new_ones([x.shape[0]]) + for i in trange(len(sigmas) - 1, disable=disable): + denoised = model(x, sigmas[i] * s_in, **extra_args) + if callback is not None: + callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) + + x = denoised + if sigmas[i + 1] > 0: + if scales: + # Interpolate to next scale + x = nn.functional.interpolate(x, size=scales[i + 1], mode='bicubic') + x = model.inner_model.inner_model.model_sampling.noise_scaling(sigmas[i + 1], noise_sampler(sigmas[i], sigmas[i + 1]), x) + return x + @torch.no_grad() def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.): From bdc8c2a8c7919b289504fd529647805e751253e8 Mon Sep 17 00:00:00 2001 From: GodOfNothing Date: Sun, 23 Mar 2025 15:54:49 +0300 Subject: [PATCH 2/6] Added lcm_scalewise to KSAMPLER_NAMES --- comfy/samplers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index 10728bd1f..142d938ba 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -708,7 +708,7 @@ class Sampler: KSAMPLER_NAMES = ["euler", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_cfg_pp", "heun", "heunpp2","dpm_2", "dpm_2_ancestral", "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_2s_ancestral_cfg_pp", "dpmpp_sde", "dpmpp_sde_gpu", - "dpmpp_2m", "dpmpp_2m_cfg_pp", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", + "dpmpp_2m", "dpmpp_2m_cfg_pp", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "lcm_scalewise", "ipndm", "ipndm_v", "deis", "res_multistep", "res_multistep_cfg_pp", "res_multistep_ancestral", "res_multistep_ancestral_cfg_pp", "gradient_estimation", "er_sde"] From ae0b0da8b8f7fde4184d480cce5c3b45e78543c1 Mon Sep 17 00:00:00 2001 From: GodOfNothing Date: Mon, 24 Mar 2025 14:26:39 +0300 Subject: [PATCH 3/6] Added SamplerLCMScalewise node --- comfy/k_diffusion/sampling.py | 24 ---------- comfy_extras/nodes_advanced_samplers.py | 60 +++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 24 deletions(-) diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py index c9f4c04d4..5b8d8000d 100644 --- a/comfy/k_diffusion/sampling.py +++ b/comfy/k_diffusion/sampling.py @@ -929,30 +929,6 @@ def sample_lcm(model, x, sigmas, extra_args=None, callback=None, disable=None, n return x -# x0 = - -@torch.no_grad() -def sample_lcm_scalewise(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None): - extra_args = {} if extra_args is None else extra_args - seed = extra_args.get("seed", None) - scales = extra_args.get("scales", None) - if scales: - assert len(scales) == len(sigmas) - 1, "Number of scales must be equal to number of sampling steps minus one." - noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler - s_in = x.new_ones([x.shape[0]]) - for i in trange(len(sigmas) - 1, disable=disable): - denoised = model(x, sigmas[i] * s_in, **extra_args) - if callback is not None: - callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) - - x = denoised - if sigmas[i + 1] > 0: - if scales: - # Interpolate to next scale - x = nn.functional.interpolate(x, size=scales[i + 1], mode='bicubic') - x = model.inner_model.inner_model.model_sampling.noise_scaling(sigmas[i + 1], noise_sampler(sigmas[i], sigmas[i + 1]), x) - return x - @torch.no_grad() def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.): diff --git a/comfy_extras/nodes_advanced_samplers.py b/comfy_extras/nodes_advanced_samplers.py index 5fbb096fb..9db4b23ca 100644 --- a/comfy_extras/nodes_advanced_samplers.py +++ b/comfy_extras/nodes_advanced_samplers.py @@ -1,5 +1,6 @@ import comfy.samplers import comfy.utils +from comfy.k_diffusion.sampling import default_noise_sampler import torch import numpy as np from tqdm.auto import trange @@ -54,6 +55,64 @@ class SamplerLCMUpscale: scale_steps = None sampler = comfy.samplers.KSAMPLER(sample_lcm_upscale, extra_options={"total_upscale": scale_ratio, "upscale_steps": scale_steps, "upscale_method": upscale_method}) return (sampler, ) + + +@torch.no_grad() +def sample_lcm_scalewise(model, x, sigmas, extra_args=None, callback=None, disable=None, total_upscale=2.0, upscale_method="bislerp", upscale_steps=None): + extra_args = {} if extra_args is None else extra_args + seed = extra_args.get("seed", None) + + if upscale_steps is None: + upscale_steps = max(len(sigmas) // 2 + 1, 2) + else: + upscale_steps += 1 + upscale_steps = min(upscale_steps, len(sigmas) + 1) + + upscales = np.linspace(1.0, total_upscale, upscale_steps)[1:] + + orig_shape = x.size() + s_in = x.new_ones([x.shape[0]]) + for i in trange(len(sigmas) - 1, disable=disable): + denoised = model(x, sigmas[i] * s_in, **extra_args) + if callback is not None: + callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised}) + + x = denoised + if i < len(upscales): + x = comfy.utils.common_upscale(x, round(orig_shape[-1] * upscales[i]), round(orig_shape[-2] * upscales[i]), upscale_method, "disabled") + + if sigmas[i + 1] > 0: + # Since the size of noise if changing, noise_sampler has to be redefined each time + noise_sampler = default_noise_sampler(x, seed=seed) + # Noise using the model's scheduler + x = model.inner_model.inner_model.model_sampling.noise_scaling(sigmas[i + 1], noise_sampler(sigmas[i], sigmas[i + 1]), x) + return x + + +class SamplerLCMScalewise: + upscale_methods = ["bicubic", "bilinear", "nearest-exact"] + + @classmethod + def INPUT_TYPES(s): + return { + "required": + { + "scale_ratio": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 4.0, "step": 0.25}), + "scale_steps": ("INT", {"default": -1, "min": -1, "max": 1000, "step": 1}), + "upscale_method": (s.upscale_methods,), + } + } + RETURN_TYPES = ("SAMPLER",) + CATEGORY = "sampling/custom_sampling/samplers" + + FUNCTION = "get_sampler" + + def get_sampler(self, scale_ratio, scale_steps, upscale_method): + if scale_steps < 0: + scale_steps = None + sampler = comfy.samplers.KSAMPLER(sample_lcm_scalewise, extra_options={"total_upscale": scale_ratio, "upscale_steps": scale_steps, "upscale_method": upscale_method}) + return (sampler, ) + from comfy.k_diffusion.sampling import to_d import comfy.model_patcher @@ -103,6 +162,7 @@ class SamplerEulerCFGpp: NODE_CLASS_MAPPINGS = { "SamplerLCMUpscale": SamplerLCMUpscale, + "SamplerLCMScalewise": SamplerLCMScalewise, "SamplerEulerCFGpp": SamplerEulerCFGpp, } From 9f9db7fc2977560388b35bcb671da3109687e888 Mon Sep 17 00:00:00 2001 From: GodOfNothing Date: Mon, 24 Mar 2025 16:17:39 +0300 Subject: [PATCH 4/6] Removed lcm_scalewise from KSAMPLER_NAMES --- comfy/samplers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy/samplers.py b/comfy/samplers.py index 142d938ba..10728bd1f 100644 --- a/comfy/samplers.py +++ b/comfy/samplers.py @@ -708,7 +708,7 @@ class Sampler: KSAMPLER_NAMES = ["euler", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_cfg_pp", "heun", "heunpp2","dpm_2", "dpm_2_ancestral", "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_2s_ancestral_cfg_pp", "dpmpp_sde", "dpmpp_sde_gpu", - "dpmpp_2m", "dpmpp_2m_cfg_pp", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "lcm_scalewise", + "dpmpp_2m", "dpmpp_2m_cfg_pp", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "ipndm", "ipndm_v", "deis", "res_multistep", "res_multistep_cfg_pp", "res_multistep_ancestral", "res_multistep_ancestral_cfg_pp", "gradient_estimation", "er_sde"] From c868cb205573b164c45c8f26d7a9e039d6dfadd3 Mon Sep 17 00:00:00 2001 From: GodOfNothing Date: Mon, 24 Mar 2025 18:09:10 +0300 Subject: [PATCH 5/6] Added option to specify list of scales as comma-separated string --- comfy_extras/nodes_advanced_samplers.py | 27 +++++++++++-------------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/comfy_extras/nodes_advanced_samplers.py b/comfy_extras/nodes_advanced_samplers.py index 9db4b23ca..75ed86f78 100644 --- a/comfy_extras/nodes_advanced_samplers.py +++ b/comfy_extras/nodes_advanced_samplers.py @@ -58,18 +58,14 @@ class SamplerLCMUpscale: @torch.no_grad() -def sample_lcm_scalewise(model, x, sigmas, extra_args=None, callback=None, disable=None, total_upscale=2.0, upscale_method="bislerp", upscale_steps=None): +def sample_lcm_scalewise(model, x, sigmas, extra_args=None, callback=None, disable=None, upscales=None, upscale_method="bicubic"): extra_args = {} if extra_args is None else extra_args seed = extra_args.get("seed", None) - if upscale_steps is None: - upscale_steps = max(len(sigmas) // 2 + 1, 2) - else: - upscale_steps += 1 - upscale_steps = min(upscale_steps, len(sigmas) + 1) - - upscales = np.linspace(1.0, total_upscale, upscale_steps)[1:] - + if upscales is not None: + # Resolution is increased on each step except the last one + assert len(upscales) == len(sigmas) - 2 + orig_shape = x.size() s_in = x.new_ones([x.shape[0]]) for i in trange(len(sigmas) - 1, disable=disable): @@ -97,8 +93,7 @@ class SamplerLCMScalewise: return { "required": { - "scale_ratio": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 4.0, "step": 0.25}), - "scale_steps": ("INT", {"default": -1, "min": -1, "max": 1000, "step": 1}), + "upscales": ("STRING", {"default": ""}), "upscale_method": (s.upscale_methods,), } } @@ -107,10 +102,12 @@ class SamplerLCMScalewise: FUNCTION = "get_sampler" - def get_sampler(self, scale_ratio, scale_steps, upscale_method): - if scale_steps < 0: - scale_steps = None - sampler = comfy.samplers.KSAMPLER(sample_lcm_scalewise, extra_options={"total_upscale": scale_ratio, "upscale_steps": scale_steps, "upscale_method": upscale_method}) + def get_sampler(self, upscales, upscale_method): + # Turn comma-separated list into string + upscales = [float(value) for value in upscales.split(',')] + if len(upscales) == 0: + upscales = None + sampler = comfy.samplers.KSAMPLER(sample_lcm_scalewise, extra_options={"upscales": upscales, "upscale_method": upscale_method}) return (sampler, ) From 618a7a3feaebf4652b28c0294666d715124884cb Mon Sep 17 00:00:00 2001 From: GodOfNothing Date: Mon, 24 Mar 2025 18:25:26 +0300 Subject: [PATCH 6/6] Added sanity check that upscales is non-decreasing sequence --- comfy_extras/nodes_advanced_samplers.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/comfy_extras/nodes_advanced_samplers.py b/comfy_extras/nodes_advanced_samplers.py index 75ed86f78..b4deec04c 100644 --- a/comfy_extras/nodes_advanced_samplers.py +++ b/comfy_extras/nodes_advanced_samplers.py @@ -102,9 +102,18 @@ class SamplerLCMScalewise: FUNCTION = "get_sampler" + def _validate_upscales(self, upscales): + if not upscales: + return + + for i in range(1, len(upscales)): + if upscales[i] < upscales[i-1]: + raise ValueError("`upscales` is expected to be non-decreasing sequence of numbers") + def get_sampler(self, upscales, upscale_method): # Turn comma-separated list into string upscales = [float(value) for value in upscales.split(',')] + self._validate_upscales(upscales) if len(upscales) == 0: upscales = None sampler = comfy.samplers.KSAMPLER(sample_lcm_scalewise, extra_options={"upscales": upscales, "upscale_method": upscale_method})