From cc06d81002f542d67952c8d5e020314c44710a87 Mon Sep 17 00:00:00 2001
From: yoinked <pepe.dannyboy@gmail.com>
Date: Sun, 8 Dec 2024 22:25:17 -0800
Subject: [PATCH] Add MaHiRo (improved CFG)

long explanation of what it is is [here](https://huggingface.co/spaces/yoinked/blue-arxiv) (2024-1208.1)


note: if the node name has encoding issues (utf 8/whatever), id suggest to replace the face at the end with `(>w<)`
---
 comfy_extras/nodes_mahiro.py | 40 ++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)
 create mode 100644 comfy_extras/nodes_mahiro.py

diff --git a/comfy_extras/nodes_mahiro.py b/comfy_extras/nodes_mahiro.py
new file mode 100644
index 000000000..5c882acea
--- /dev/null
+++ b/comfy_extras/nodes_mahiro.py
@@ -0,0 +1,40 @@
+import torch
+import torch.nn.functional as F
+
+class Mahiro:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"model": ("MODEL",),
+                            }}
+    RETURN_TYPES = ("MODEL",)
+    RETURN_NAMES = ("patched_model",)
+    FUNCTION = "patch"
+    CATEGORY = "_for_testing"
+    def patch(self, model):
+        m = model.clone()
+        def mahiro_normd(args):
+            scale: float = args['cond_scale']
+            cond_p: torch.Tensor = args['cond_denoised']
+            uncond_p: torch.Tensor = args['uncond_denoised']
+            #naive leap
+            leap = cond_p * scale
+            #sim with uncond leap
+            u_leap = uncond_p * scale
+            cfg = uncond_p + (cond_p - uncond_p) * scale
+            merge = (leap + cfg) / 2
+            normu = torch.sqrt(u_leap.abs()) * u_leap.sign()
+            normm = torch.sqrt(merge.abs()) * merge.sign()
+            sim = F.cosine_similarity(normu, normm).mean()
+            simsc = 2 * (sim+1)
+            wm = (simsc*cfg + (4-simsc)*leap) / 4
+            return args['input'] - wm
+        m.set_model_sampler_cfg_function(mahiro_normd)
+        return (m, )
+
+NODE_CLASS_MAPPINGS = {
+    "Mahiro": Mahiro
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Mahiro": "Mahiro is so cute that she deserves a better guidance function!! (。・ω・。)",
+}