From dc322e44fe1c6c8c89400a3a70fb044559faed62 Mon Sep 17 00:00:00 2001
From: Ken Simpson <ksimpson@mailchannels.com>
Date: Fri, 2 Feb 2024 18:35:13 -0800
Subject: [PATCH] Added a README

---
 README.md   | 23 +++++++++++++++++
 __init__.py |  3 +++
 nodes.py    | 74 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 100 insertions(+)
 create mode 100644 README.md
 create mode 100644 __init__.py
 create mode 100644 nodes.py

diff --git a/README.md b/README.md
new file mode 100644
index 0000000..4d4ee06
--- /dev/null
+++ b/README.md
@@ -0,0 +1,23 @@
+# ComfyUI Tensor Operations Nodes
+
+This repo contains nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) that implement some helpful operations on tensors, such as normalization.
+
+## Updates
+
+### February 2nd, 2024
+
+- Initial release. Added `Image Match Normalize` and `Latent Match Normalize` nodes.
+
+## Nodes
+
+### Image Match Normalize
+
+This node returns a normalized version of the target image using the mean and standard deviation of each color channel of the source image. If you want the color and brightness of your image to match the colour and brightness of another image,
+this is the node for you.
+
+### Latent Match Normalize
+
+This node returns a normalized version of the target latent using the mean and standard deviation of each channel of
+the source latent. Latents encode color information differently than images and you may find that normalizing an image
+by instead normalizing its latent representation results in a "better" result that is closer to the coloring of the
+source.
\ No newline at end of file
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..e74c2f2
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS']
\ No newline at end of file
diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..bf159cf
--- /dev/null
+++ b/nodes.py
@@ -0,0 +1,74 @@
+import logging
+import torch
+
+@torch.no_grad()
+def match_normalize(target_tensor, source_tensor, dimensions=4):
+    "Adjust target_tensor based on source_tensor's mean and stddev"   
+    if len(target_tensor.shape) != dimensions:
+        raise ValueError("source_latent must have four dimensions")
+    if len(source_tensor.shape) != dimensions:
+        raise ValueError("target_latent must have four dimensions")
+    
+    logging.warning(f"target_tensor.shape = {target_tensor.shape}")
+
+    # Calculate the mean and std of target tensor
+    tgt_mean = target_tensor.mean(dim=[2, 3], keepdim=True)
+    tgt_std = target_tensor.std(dim=[2, 3], keepdim=True)
+    
+    # Calculate the mean and std of source tensor
+    src_mean = source_tensor.mean(dim=[2, 3], keepdim=True)
+    src_std = source_tensor.std(dim=[2, 3], keepdim=True)
+    
+    # Normalize target tensor to have mean=0 and std=1, then rescale
+    normalized_tensor = (target_tensor.clone() - tgt_mean) / tgt_std * src_std + src_mean
+    
+    return normalized_tensor
+
+class LatentMatchNormalize:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"source_latent": ("LATENT", ),
+                     "target_latent": ("LATENT", )}}
+    RETURN_TYPES = ("LATENT",)
+    CATEGORY = "tensor_ops"
+
+    FUNCTION = "latent_match_normalize"
+
+    @torch.no_grad()
+    def latent_match_normalize(self, source_latent, target_latent):       
+        normalized_latent = match_normalize(target_latent["samples"], source_latent["samples"], dimensions=4)
+
+        return_latent = source_latent.copy()
+        return_latent["samples"] = normalized_latent
+        return (return_latent,)
+    
+class ImageMatchNormalize:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"source_image": ("IMAGE", ),
+                     "target_image": ("IMAGE", )}}
+    RETURN_TYPES = ("IMAGE",)
+    CATEGORY = "tensor_ops"
+
+    FUNCTION = "image_match_normalize"
+
+    @torch.no_grad()
+    def image_match_normalize(self, source_image, target_image):
+        # image shape is [B, H, W, C], but the normalize function needs [B, C, H, W]
+        logging.warning(f"source.shape = {source_image.shape}")
+        source = source_image.permute(0,3,1,2)
+        target = target_image.permute(0,3,1,2)
+        
+        normalized = match_normalize(target, source, dimensions=4)
+
+        logging.warning(f"normalized.shape = {normalized.shape}")
+        normalized_image = normalized.permute(0,2,3,1)
+        logging.warning(f"normalized_image.shape = {normalized_image.shape}")
+        return (normalized_image,)
+    
+NODE_CLASS_MAPPINGS = {
+    "Image Match Normalize": ImageMatchNormalize,
+    "Latent Match Normalize": LatentMatchNormalize
+}
\ No newline at end of file