diff --git a/comfy_extras/nodes_advanced_samplers.py b/comfy_extras/nodes_advanced_samplers.py index 820c250ef..fc8196170 100644 --- a/comfy_extras/nodes_advanced_samplers.py +++ b/comfy_extras/nodes_advanced_samplers.py @@ -82,7 +82,6 @@ def sample_euler_pp(model, x, sigmas, extra_args=None, callback=None, disable=No x = x + d * dt return x - class SamplerEulerCFGpp: @classmethod def INPUT_TYPES(s): @@ -102,11 +101,61 @@ class SamplerEulerCFGpp: sampler = comfy.samplers.ksampler("euler_cfg_pp") return (sampler, ) +@torch.no_grad() +def sample_heun_pp(model, x, sigmas, extra_args=None, callback=None, disable=None): + extra_args = {} if extra_args is None else extra_args + + temp = [0] + def post_cfg_function(args): + temp[0] = args["uncond_denoised"] + return args["denoised"] + + model_options = extra_args.get("model_options", {}).copy() + extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True) + + s_in = x.new_ones([x.shape[0]]) + for i in trange(len(sigmas) - 1, disable=disable): + sigma_hat = sigmas[i] + denoised = model(x, sigma_hat * s_in, **extra_args) + d = to_d(x - denoised + temp[0], sigmas[i], denoised) + if callback is not None: + callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised}) + dt = sigmas[i + 1] - sigma_hat + if sigmas[i + 1] == 0: + # Euler method + x = x + d * dt + else: + # Heun's method + x_2 = x + d * dt + denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args) + d_2 = to_d(x_2, sigmas[i + 1], denoised_2) + d_prime = (d + d_2) / 2 + x = x + d_prime * dt + return x + +class SamplerHeunCFGpp: + @classmethod + def INPUT_TYPES(s): + return {"required": + {} + } + RETURN_TYPES = ("SAMPLER",) + # CATEGORY = "sampling/custom_sampling/samplers" + CATEGORY = "_for_testing" + + FUNCTION = "get_sampler" + + def get_sampler(self): + sampler = comfy.samplers.KSAMPLER(sample_heun_pp) + return (sampler, ) + NODE_CLASS_MAPPINGS = { "SamplerLCMUpscale": SamplerLCMUpscale, "SamplerEulerCFGpp": SamplerEulerCFGpp, + "SamplerHeunCFGpp": SamplerHeunCFGpp, } NODE_DISPLAY_NAME_MAPPINGS = { "SamplerEulerCFGpp": "SamplerEulerCFG++", + "SamplerHeunCFGpp": "SamplerHeunCFG++", }