From 0f5768e038343a8eb96074c8c7f3911abf12461e Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 4 Apr 2024 23:38:57 -0400
Subject: [PATCH] Fix missing arguments in cfg_function.

---
 comfy/samplers.py                    | 4 ++--
 comfy_extras/nodes_custom_sampler.py | 7 +++++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/comfy/samplers.py b/comfy/samplers.py
index 08cbab75..415a35cc 100644
--- a/comfy/samplers.py
+++ b/comfy/samplers.py
@@ -231,7 +231,7 @@ def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options):
     logging.warning("WARNING: The comfy.samplers.calc_cond_uncond_batch function is deprecated please use the calc_cond_batch one instead.")
     return tuple(calc_cond_batch(model, [cond, uncond], x_in, timestep, model_options))
 
-def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}):
+def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}, cond=None, uncond=None):
     if "sampler_cfg_function" in model_options:
         args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
                 "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
@@ -256,7 +256,7 @@ def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_option
 
     conds = [cond, uncond_]
     out = calc_cond_batch(model, conds, x, timestep, model_options)
-    return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options)
+    return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options, cond=cond, uncond=uncond_)
 
 
 class KSamplerX0Inpaint:
diff --git a/comfy_extras/nodes_custom_sampler.py b/comfy_extras/nodes_custom_sampler.py
index fa113192..a5d791d9 100644
--- a/comfy_extras/nodes_custom_sampler.py
+++ b/comfy_extras/nodes_custom_sampler.py
@@ -436,8 +436,11 @@ class Guider_DualCFG(comfy.samplers.CFGGuider):
         self.inner_set_conds({"positive": positive, "middle": middle, "negative": negative})
 
     def predict_noise(self, x, timestep, model_options={}, seed=None):
-        out = comfy.samplers.calc_cond_batch(self.inner_model, [self.conds.get("negative", None), self.conds.get("middle", None), self.conds.get("positive", None)], x, timestep, model_options)
-        return comfy.samplers.cfg_function(self.inner_model, out[1], out[0], self.cfg2, x, timestep, model_options=model_options) + (out[2] - out[1]) * self.cfg1
+        negative_cond = self.conds.get("negative", None)
+        middle_cond = self.conds.get("middle", None)
+
+        out = comfy.samplers.calc_cond_batch(self.inner_model, [negative_cond, middle_cond, self.conds.get("positive", None)], x, timestep, model_options)
+        return comfy.samplers.cfg_function(self.inner_model, out[1], out[0], self.cfg2, x, timestep, model_options=model_options, cond=middle_cond, uncond=negative_cond) + (out[2] - out[1]) * self.cfg1
 
 class DualCFGGuider:
     @classmethod