Release 0.3.1

Should fix #9

__init__.py

@@ -5,14 +5,14 @@ import torch
 cos = torch.nn.CosineSimilarity(dim=1)
 
 
-class Guider_AdaptiveGuidance(comfy.samplers.CFGGuider):
+# shared structure for adaptive guiders
+class AdaptiveGuider(object):
+    cfg_start_timestep = 1000.0
     threshold_timestep = 0
     uz_scale = 0.0
 
-    def set_cfg(self, cfg):
-        self.cfg = cfg
-
-    def set_threshold(self, threshold):
+    def set_threshold(self, threshold, start_at):
+        self.cfg_start_timestep = start_at
         self.threshold = threshold
 
     def set_uncond_zero_scale(self, scale):
@@ -25,27 +25,41 @@ class Guider_AdaptiveGuidance(comfy.samplers.CFGGuider):
         cond -= cond.mean()
         return x - (cond / cond.std() ** 0.5) * self.uz_scale
 
+    def check_similarity(self, ts, cond_pred, uncond_pred):
+        if not self.threshold >= 1.0:
+            sim = cos(cond_pred.reshape(1, -1), uncond_pred.reshape(1, -1)).item()
+            if sim >= self.threshold:
+                print(f"AdaptiveGuider: Cosine similarity {sim:.4f} exceeds threshold, setting CFG to 1.0")
+                self.threshold_timestep = ts
+
     def predict_noise(self, x, timestep, model_options={}, seed=None):
-        cond = self.conds.get("positive")
-        uncond = self.conds.get("negative")
         ts = timestep[0].item()
-        if self.threshold_timestep > ts or self.cfg == 1.0:
+        if ts > self.cfg_start_timestep or self.threshold_timestep > ts or self.cfg == 1.0:
             if self.uz_scale > 0.0:
                 model_options = model_options.copy()
                 model_options["sampler_cfg_function"] = self.zero_cond
+            cond = self.conds.get("positive")
+            uncond = self.conds.get("negative")
             return comfy.samplers.sampling_function(
                 self.inner_model, x, timestep, uncond, cond, 1.0, model_options=model_options, seed=seed
             )
         self.threshold_timestep = 0
-        uncond_pred, cond_pred = comfy.samplers.calc_cond_batch(
-            self.inner_model, [uncond, cond], x, timestep, model_options
-        )
-        if not self.threshold >= 1.0:
-            # Is this reshape correct? It at least gives a scalar value...
-            sim = cos(cond_pred.reshape(1, -1), uncond_pred.reshape(1, -1)).item()
-            if sim >= self.threshold:
-                print(f"\nAdaptiveGuidance: Cosine similarity {sim:.4f} exceeds threshold, setting CFG to 1.0")
-                self.threshold_timestep = ts
+        conds = self.calc_conds(x, timestep, model_options)
+        self.check_similarity(ts, conds[0], conds[1])
+        return self.calc_cfg(conds, x, timestep, model_options)
+
+
+class Guider_AdaptiveGuidance(AdaptiveGuider, comfy.samplers.CFGGuider):
+    def calc_conds(self, x, timestep, model_options):
+        cond = self.conds.get("positive")
+        uncond = self.conds.get("negative")
+        return comfy.samplers.calc_cond_batch(self.inner_model, [cond, uncond], x, timestep, model_options)
+
+    def calc_cfg(self, conds, x, timestep, model_options):
+        cond = self.conds.get("positive")
+        uncond = self.conds.get("negative")
+        cond_pred, uncond_pred = conds
+
         return comfy.samplers.cfg_function(
             self.inner_model,
             cond_pred,
@@ -59,82 +73,21 @@ class Guider_AdaptiveGuidance(comfy.samplers.CFGGuider):
         )
 
 
-class AdaptiveGuidanceGuider:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "model": ("MODEL",),
-                "positive": ("CONDITIONING",),
-                "negative": ("CONDITIONING",),
-                "threshold": ("FLOAT", {"default": 0.990, "min": 0.90, "max": 1.0, "step": 0.001, "round": 0.001}),
-                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
-            },
-            "optional": {"uncond_zero_scale": ("FLOAT", {"default": 0.0, "max": 2.0, "step": 0.01})},
-        }
-
-    RETURN_TYPES = ("GUIDER",)
-    FUNCTION = "get_guider"
-
-    CATEGORY = "sampling/custom_sampling/guiders"
-
-    def get_guider(self, model, positive, negative, threshold, cfg, uncond_zero_scale=0.0):
-        g = Guider_AdaptiveGuidance(model)
-        g.set_conds(positive, negative)
-        g.set_threshold(threshold)
-        g.set_uncond_zero_scale(uncond_zero_scale)
-        g.set_cfg(cfg)
-
-        return (g,)
-
-
-class Guider_PerpNegAG(comfy_extras.nodes_perpneg.Guider_PerpNeg):
-    threshold_timestep = 0
-    uz_scale = 0.0
-
-    def set_threshold(self, threshold):
-        self.threshold = threshold
-
-    def set_uncond_zero_scale(self, scale):
-        self.uz_scale = scale
-
-    def zero_cond(self, args):
-        cond = args["cond_denoised"]
-        x = args["input"]
-        x -= x.mean()
-        cond -= cond.mean()
-        return x - (cond / cond.std() ** 0.5) * self.uz_scale
-
-    def predict_noise(self, x, timestep, model_options={}, seed=None):
+class Guider_PerpNegAG(AdaptiveGuider, comfy_extras.nodes_perpneg.Guider_PerpNeg):
+    def calc_conds(self, x, timestep, model_options):
         cond = self.conds.get("positive")
         uncond = self.conds.get("negative")
-        ts = timestep[0].item()
-        if self.threshold_timestep > ts or self.cfg == 1.0:
-            if self.uz_scale > 0.0:
-                model_options = model_options.copy()
-                model_options["sampler_cfg_function"] = self.zero_cond
-            return comfy.samplers.sampling_function(
-                self.inner_model, x, timestep, uncond, cond, 1.0, model_options=model_options, seed=seed
-            )
-        self.threshold_timestep = 0
-
-        # From comfy_extras.nodes_perpneg - Guider_PerpNeg
-        # No need for calculating perp-neg when skipping negative
         empty_cond = self.conds.get("empty_negative_prompt")
-        (cond_pred, uncond_pred, empty_cond_pred) = comfy.samplers.calc_cond_batch(
-            self.inner_model, [cond, uncond, empty_cond], x, timestep, model_options
-        )
+        return comfy.samplers.calc_cond_batch(self.inner_model, [cond, uncond, empty_cond], x, timestep, model_options)
+
+    def calc_cfg(self, conds, x, timestep, model_options):
+        cond = self.conds.get("positive")
+        uncond = self.conds.get("negative")
+        empty_cond = self.conds.get("empty_negative_prompt")
+        cond_pred, uncond_pred, empty_cond_pred = conds
         cfg_result = comfy_extras.nodes_perpneg.perp_neg(
             x, cond_pred, uncond_pred, empty_cond_pred, self.neg_scale, self.cfg
         )
-
-        if not self.threshold >= 1.0:
-            sim = cos(cond_pred.reshape(1, -1), uncond_pred.reshape(1, -1)).item()
-            if sim >= self.threshold:
-                print(f"\nPerpNegAG: Cosine similarity {sim:.4f} exceeds threshold, setting CFG to 1.0")
-                self.threshold_timestep = ts
-
-        # From comfy_extras.nodes_perpneg - Guider_PerpNeg
         for fn in model_options.get("sampler_post_cfg_function", []):
             args = {
                 "denoised": cfg_result,
@@ -155,6 +108,39 @@ class Guider_PerpNegAG(comfy_extras.nodes_perpneg.Guider_PerpNeg):
         return cfg_result
 
 
+class AdaptiveGuidanceGuider:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "threshold": ("FLOAT", {"default": 0.990, "min": 0.90, "max": 1.0, "step": 0.0001, "round": 0.0001}),
+                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
+            },
+            "optional": {
+                "uncond_zero_scale": ("FLOAT", {"default": 0.0, "max": 2.0, "step": 0.01}),
+                "cfg_start_pct": ("FLOAT", {"default": 0.0, "max": 1.0, "step": 0.01}),
+            },
+        }
+
+    RETURN_TYPES = ("GUIDER",)
+    FUNCTION = "get_guider"
+
+    CATEGORY = "sampling/custom_sampling/guiders"
+
+    def get_guider(self, model, positive, negative, threshold, cfg, uncond_zero_scale=0.0, cfg_start_pct=0.0):
+        cfg_start_timestep = model.get_model_object("model_sampling").percent_to_sigma(cfg_start_pct)
+        g = Guider_AdaptiveGuidance(model)
+        g.set_conds(positive, negative)
+        g.set_threshold(threshold, cfg_start_timestep)
+        g.set_uncond_zero_scale(uncond_zero_scale)
+        g.set_cfg(cfg)
+
+        return (g,)
+
+
 class PerpNegAGGuider:
     @classmethod
     def INPUT_TYPES(s):
@@ -164,11 +150,14 @@ class PerpNegAGGuider:
                 "positive": ("CONDITIONING",),
                 "negative": ("CONDITIONING",),
                 "empty_conditioning": ("CONDITIONING",),
-                "threshold": ("FLOAT", {"default": 0.990, "min": 0.90, "max": 1.0, "step": 0.001, "round": 0.001}),
+                "threshold": ("FLOAT", {"default": 0.990, "min": 0.90, "max": 1.0, "step": 0.0001, "round": 0.0001}),
                 "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
                 "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
             },
-            "optional": {"uncond_zero_scale": ("FLOAT", {"default": 0.0, "max": 2.0, "step": 0.01})},
+            "optional": {
+                "uncond_zero_scale": ("FLOAT", {"default": 0.0, "max": 2.0, "step": 0.01}),
+                "cfg_start_pct": ("FLOAT", {"default": 0.0, "max": 1.0, "step": 0.01}),
+            },
         }
 
     RETURN_TYPES = ("GUIDER",)
@@ -177,11 +166,21 @@ class PerpNegAGGuider:
     CATEGORY = "sampling/custom_sampling/guiders"
 
     def get_guider(
-        self, model, positive, negative, empty_conditioning, threshold, cfg, neg_scale, uncond_zero_scale=0.0
+        self,
+        model,
+        positive,
+        negative,
+        empty_conditioning,
+        threshold,
+        cfg,
+        neg_scale,
+        uncond_zero_scale=0.0,
+        cfg_start_pct=0.0,
     ):
+        cfg_start_timestep = model.get_model_object("model_sampling").percent_to_sigma(cfg_start_pct)
         g = Guider_PerpNegAG(model)
         g.set_conds(positive, negative, empty_conditioning)
-        g.set_threshold(threshold)
+        g.set_threshold(threshold, cfg_start_timestep)
         g.set_uncond_zero_scale(uncond_zero_scale)
         g.set_cfg(cfg, neg_scale)
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-adaptive-guidance"
 description = "An implementation of adaptive guidance for ComfyUI\nSee https://bcv-uniandes.github.io/adaptiveguidance-wp/"
-version = "0.2.1"
+version = "0.3.1"
 license = { text = "GNU General Public License v3.0" }
 
 [project.urls]