diff --git a/comfy_extras/nodes_mahiro.py b/comfy_extras/nodes_mahiro.py
index 7972401c0..1a84ff5b7 100644
--- a/comfy_extras/nodes_mahiro.py
+++ b/comfy_extras/nodes_mahiro.py
@@ -28,8 +28,8 @@ class Mahiro:
             sim = F.cosine_similarity(normu, normm).mean()
             simsc = 2 * (sim+1)
             wm = (simsc*cfg + (4-simsc)*leap) / 4
-            return wm
-        m.set_model_sampler_post_cfg_function(mahiro_normd)
+            return args['input'] - wm
+        m.set_model_sampler_cfg_function(mahiro_normd)
         return (m, )
 
 NODE_CLASS_MAPPINGS = {