From 3923884d2e7cc0b8ef7111a26afadd4722ee488a Mon Sep 17 00:00:00 2001
From: Jedrzej Kosinski <kosinkadink1@gmail.com>
Date: Fri, 9 Feb 2024 04:57:13 -0600
Subject: [PATCH] Fixed cfg being set to 1.0 on KSampler causing issues with
 custom_cfg

---
 animatediff/nodes.py           |  2 +-
 animatediff/sample_settings.py | 25 -------------------------
 animatediff/sampling.py        |  3 ++-
 3 files changed, 3 insertions(+), 27 deletions(-)

diff --git a/animatediff/nodes.py b/animatediff/nodes.py
index e16d3ee..ac8d630 100644
--- a/animatediff/nodes.py
+++ b/animatediff/nodes.py
@@ -134,5 +134,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     # Deprecated Nodes
     "AnimateDiffLoaderV1": "AnimateDiff Loader [DEPRECATED] 🎭🅐🅓",
     "ADE_AnimateDiffLoaderV1Advanced": "AnimateDiff Loader (Advanced) [DEPRECATED] 🎭🅐🅓",
-    "ADE_AnimateDiffCombine": "DO NOT USE, USE VideoCombine from ComfyUI-VideoHelperSuite instead! AnimateDiff Combine [DEPRECATED, DO NOT USE] 🎭🅐🅓",
+    "ADE_AnimateDiffCombine": "AnimateDiff Combine [DEPRECATED, Use Video Combine (VHS) Instead!] 🎭🅐🅓",
 }
diff --git a/animatediff/sample_settings.py b/animatediff/sample_settings.py
index 51d9790..e757e07 100644
--- a/animatediff/sample_settings.py
+++ b/animatediff/sample_settings.py
@@ -449,31 +449,6 @@ class FreeInitOptions(IterationOptions):
             raise ValueError(f"FreeInit init_type '{self.init_type}' is not recognized.")
 
 
-class CustomCFG:
-    def __init__(self, cfg_multival: Union[float, Tensor]):
-        self.cfg_multival = cfg_multival
-        self.masks = None
-        # TODO: add support for cfg keyframes
-    
-    def initialize_timesteps(self, model: BaseModel):
-        pass
-
-    def patch_model(self, model: ModelPatcher) -> ModelPatcher:
-        def evolved_custom_cfg(args):
-            cond: Tensor = args["cond"]
-            uncond: Tensor = args["uncond"]
-            # cond scale is based purely off of CustomCFG - cond_scale input in sampler is ignored!
-            cond_scale = self.cfg_multival
-            if isinstance(cond_scale, Tensor):
-                cond_scale = prepare_mask_batch(cond_scale.to(cond.dtype).to(cond.device), cond.shape)
-                cond_scale = extend_to_batch_size(cond_scale, cond.shape[0])
-            return uncond + (cond - uncond) * cond_scale
-
-        model = model.clone()
-        model.set_model_sampler_cfg_function(evolved_custom_cfg)
-        return model
-
-
 class CustomCFGKeyframe:
     def __init__(self, cfg_multival: Union[float, Tensor], start_percent=0.0, guarantee_steps=1):
         self.cfg_multival = cfg_multival
diff --git a/animatediff/sampling.py b/animatediff/sampling.py
index 0d64318..eee218b 100644
--- a/animatediff/sampling.py
+++ b/animatediff/sampling.py
@@ -367,7 +367,8 @@ def evolved_sampling_function(model, x, timestep, uncond, cond, cond_scale, mode
     if ADGS.sample_settings.custom_cfg is not None:
         ADGS.sample_settings.custom_cfg.prepare_current_keyframe(t=timestep)
 
-    if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
+    # never use cfg1 optimization if using custom_cfg (since can have timesteps and such)
+    if ADGS.sample_settings.custom_cfg is None and math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
         uncond_ = None
     else:
         uncond_ = uncond