diff --git a/cfgstar.py b/cfgstar.py
index 39ead92..c90c2d2 100644
--- a/cfgstar.py
+++ b/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,)
 
diff --git a/pyproject.toml b/pyproject.toml
index 629ef65..cb232ea 100644
--- a/pyproject.toml
+++ b/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]