From 08007223338697b90e2228f64d5fcda20955dfcd Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Wed, 17 May 2023 01:34:43 -0400
Subject: [PATCH] Create sampler_rescalecfg.py

---
 sampler_rescalecfg.py | 34 ++++++++++++++++++++++++++++++++++
 1 file changed, 34 insertions(+)
 create mode 100644 sampler_rescalecfg.py

diff --git a/sampler_rescalecfg.py b/sampler_rescalecfg.py
new file mode 100644
index 0000000..6e1289f
--- /dev/null
+++ b/sampler_rescalecfg.py
@@ -0,0 +1,34 @@
+import torch
+
+
+class RescaleClassifierFreeGuidance:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "multiplier": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}),
+                              }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
+
+    CATEGORY = "custom_node_experiments"
+
+    def patch(self, model, multiplier):
+        
+        def sampler_tonemap_reinhard(cond, uncond, cond_scale):
+            x_cfg = uncond + cond_scale * (cond - uncond)
+            ro_pos = torch.std(cond, dim=(1,2,3), keepdim=True)
+            ro_cfg = torch.std(x_cfg, dim=(1,2,3), keepdim=True)
+
+            x_rescaled = x_cfg * (ro_pos / ro_cfg)
+            x_final = multiplier * x_rescaled + (1.0 - multiplier) * x_cfg
+
+            return x_final
+
+        m = model.clone()
+        m.set_model_sampler_cfg_function(sampler_tonemap_reinhard)
+        return (m, )
+
+
+NODE_CLASS_MAPPINGS = {
+    "RescaleClassifierFreeGuidanceTest": RescaleClassifierFreeGuidance,
+}