From 466aaa130287be719024589272be7746fef59801 Mon Sep 17 00:00:00 2001
From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com>
Date: Mon, 17 Apr 2023 12:20:29 -0400
Subject: [PATCH] Add a test node for simple sampler tonemapping.

---
 sampler_tonemap.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 sampler_tonemap.py

diff --git a/sampler_tonemap.py b/sampler_tonemap.py
new file mode 100644
index 0000000..d2f8cbd
--- /dev/null
+++ b/sampler_tonemap.py
@@ -0,0 +1,41 @@
+import torch
+
+
+class ModelSamplerTonemapNoiseTest:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "multiplier": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.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):
+            noise_pred = (cond - uncond)
+            noise_pred_vector_magnitude = (torch.linalg.vector_norm(noise_pred, dim=(1)) + 0.0000000001)[:,None]
+            noise_pred /= noise_pred_vector_magnitude
+
+            mean = torch.mean(noise_pred_vector_magnitude, dim=(1,2,3), keepdim=True)
+            std = torch.std(noise_pred_vector_magnitude, dim=(1,2,3), keepdim=True)
+
+            top = (std * 3 + mean) * multiplier
+
+            #reinhard
+            noise_pred_vector_magnitude *= (1.0 / top)
+            new_magnitude = noise_pred_vector_magnitude / (noise_pred_vector_magnitude + 1.0)
+            new_magnitude *= top
+
+            return uncond + noise_pred * new_magnitude * cond_scale
+
+        m = model.clone()
+        m.set_model_sampler_cfg_function(sampler_tonemap_reinhard)
+        return (m, )
+
+
+NODE_CLASS_MAPPINGS = {
+    "ModelSamplerTonemapNoiseTest": ModelSamplerTonemapNoiseTest,
+}