Implement as a post cfg function

cfgstar.py

@@ -20,17 +20,16 @@ class CFGStar:
             cond = args["cond_denoised"]
             uncond = args["uncond_denoised"]
             cond_scale = args["cond_scale"]
-            x = args["input"]
 
             dot_product = torch.sum(cond * uncond, dim=(-1, -2))
             squared_norm = torch.sum(uncond * uncond, dim=(-1, -2)) + 1e-8
             s = dot_product / squared_norm
             uncond_scaled = uncond * s[:, :, None, None]
 
-            return x - (uncond_scaled + cond_scale * (cond - uncond_scaled))
+            return uncond_scaled + cond_scale * (cond - uncond_scaled)
 
         m = model.clone()
-        m.set_model_sampler_cfg_function(custom_cfg_function)
+        m.set_model_sampler_post_cfg_function(custom_cfg_function)
 
         return (m,)
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_cfgstar"
 description = "A per channel implementation of the scaled CFG from this paper: https://arxiv.org/abs/2503.18886"
-version = "1.0.0"
+version = "1.1.0"
 license = { text = "GNU General Public License v3.0" }
 
 [project.urls]