From 73ca7800192f6bc2e2194966b316b7cc3028e09a Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Sun, 23 Jun 2024 13:21:18 -0400
Subject: [PATCH] Add SamplerEulerCFG++ node.

This node should match the DDIM implementation of CFG++ when "regular" is
selected.

"alternative" is a slightly different take on CFG++
---
 comfy/model_patcher.py                  | 10 +++-
 comfy_extras/nodes_advanced_samplers.py | 74 +++++++++++++++++++++++++
 2 files changed, 81 insertions(+), 3 deletions(-)

diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
index bf8787768a7..b949031e998 100644
--- a/comfy/model_patcher.py
+++ b/comfy/model_patcher.py
@@ -51,6 +51,12 @@ def set_model_options_patch_replace(model_options, patch, name, block_name, numb
     model_options["transformer_options"] = to
     return model_options
 
+def set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=False):
+    model_options["sampler_post_cfg_function"] = model_options.get("sampler_post_cfg_function", []) + [post_cfg_function]
+    if disable_cfg1_optimization:
+        model_options["disable_cfg1_optimization"] = True
+    return model_options
+
 class ModelPatcher:
     def __init__(self, model, load_device, offload_device, size=0, current_device=None, weight_inplace_update=False):
         self.size = size
@@ -122,9 +128,7 @@ def set_model_sampler_cfg_function(self, sampler_cfg_function, disable_cfg1_opti
             self.model_options["disable_cfg1_optimization"] = True
 
     def set_model_sampler_post_cfg_function(self, post_cfg_function, disable_cfg1_optimization=False):
-        self.model_options["sampler_post_cfg_function"] = self.model_options.get("sampler_post_cfg_function", []) + [post_cfg_function]
-        if disable_cfg1_optimization:
-            self.model_options["disable_cfg1_optimization"] = True
+        self.model_options = set_model_options_post_cfg_function(self.model_options, post_cfg_function, disable_cfg1_optimization)
 
     def set_model_unet_function_wrapper(self, unet_wrapper_function: UnetWrapperFunction):
         self.model_options["model_function_wrapper"] = unet_wrapper_function
diff --git a/comfy_extras/nodes_advanced_samplers.py b/comfy_extras/nodes_advanced_samplers.py
index d973def816b..7aa89c6447f 100644
--- a/comfy_extras/nodes_advanced_samplers.py
+++ b/comfy_extras/nodes_advanced_samplers.py
@@ -56,6 +56,80 @@ def get_sampler(self, scale_ratio, scale_steps, upscale_method):
         sampler = comfy.samplers.KSAMPLER(sample_lcm_upscale, 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
+
+@torch.no_grad()
+def sample_euler_cfgpp(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, sigma_hat, temp[0])
+        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
+        x = denoised + sigmas[i + 1] * d
+    return x
+
+@torch.no_grad()
+def sample_euler_cfgpp_alt(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], sigma_hat, 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
+        # Euler method
+        x = x + d * dt
+    return x
+
+class SamplerEulerCFGpp:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"version": (["regular", "alternative"],),}
+               }
+    RETURN_TYPES = ("SAMPLER",)
+    # CATEGORY = "sampling/custom_sampling/samplers"
+    CATEGORY = "_for_testing"
+
+    FUNCTION = "get_sampler"
+
+    def get_sampler(self, version):
+        if version == "regular":
+            sampler = comfy.samplers.KSAMPLER(sample_euler_cfgpp)
+        else:
+            sampler = comfy.samplers.KSAMPLER(sample_euler_cfgpp_alt)
+        return (sampler, )
+
 NODE_CLASS_MAPPINGS = {
     "SamplerLCMUpscale": SamplerLCMUpscale,
+    "SamplerEulerCFGpp": SamplerEulerCFGpp,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SamplerEulerCFGpp": "SamplerEulerCFG++",
 }