From aa45c98ac7abff939b5cc985c0de0dd3d12d9212 Mon Sep 17 00:00:00 2001
From: MMaker <mmaker@mmaker.moe>
Date: Wed, 15 Nov 2023 20:32:42 -0500
Subject: [PATCH] ComfyUI: Add color blend

---
 __init__.py                     |  2 ++
 mmaker_color_enhance_comfyui.py | 31 ++++++++++++++++++++++++++++++-
 mmaker_color_enhance_core.py    | 10 ++++++++++
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/__init__.py b/__init__.py
index 5d52004..5fc72c5 100644
--- a/__init__.py
+++ b/__init__.py
@@ -2,8 +2,10 @@ from . import mmaker_color_enhance_comfyui
 
 NODE_CLASS_MAPPINGS = {
     "MMakerColorEnhance": mmaker_color_enhance_comfyui.ColorEnhanceComfyNode,
+    "MMakerColorBlend": mmaker_color_enhance_comfyui.ColorBlendComfyNode,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "MMakerColorEnhance": "Color Enhance",
+    "MMakerColorBlend": "Color Blend",
 }
diff --git a/mmaker_color_enhance_comfyui.py b/mmaker_color_enhance_comfyui.py
index 6954740..be8dcd4 100644
--- a/mmaker_color_enhance_comfyui.py
+++ b/mmaker_color_enhance_comfyui.py
@@ -1,7 +1,7 @@
 import torch
 import torchvision.transforms.functional as tf
 import torchvision.transforms.v2 as v2
-from .mmaker_color_enhance_core import color_enhance
+from .mmaker_color_enhance_core import color_enhance, color_blend
 
 class ColorEnhanceComfyNode:
     def __init__(self):
@@ -30,3 +30,32 @@ class ColorEnhanceComfyNode:
             images.append(edited_image)
 
         return (torch.stack(images).permute(0, 2, 3, 1),)
+
+class ColorBlendComfyNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "image_blend": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "apply_color_enhance"
+    CATEGORY = "postprocessing/Effects"
+
+    def apply_color_enhance(self, image: torch.Tensor, image_blend: torch.Tensor):
+        images = []
+        image_blend = v2.ToDtype(dtype=torch.uint8, scale=True)(image_blend).squeeze().detach().cpu().numpy()
+
+        for img in image:
+            edited_image = v2.ToDtype(dtype=torch.uint8, scale=True)(img).squeeze()
+            edited_image = color_blend(edited_image.detach().cpu().numpy(), image_blend)
+            edited_image = tf.to_tensor(edited_image)
+            images.append(edited_image)
+
+        return (torch.stack(images).permute(0, 2, 3, 1),)
diff --git a/mmaker_color_enhance_core.py b/mmaker_color_enhance_core.py
index e9188bb..258a323 100644
--- a/mmaker_color_enhance_core.py
+++ b/mmaker_color_enhance_core.py
@@ -10,5 +10,15 @@ def color_enhance(arr, strength: float = 1) -> Image.Image:
     lch[:, :, 1] *= 100/(lerp(100, lch[:, :, 1].max(), strength)) # Normalize chroma component
     return Image.fromarray(np.array(skimage.color.lab2rgb(lab=skimage.color.lch2lab(lch=lch)) * 255, dtype=np.uint8))
 
+def color_blend(arr1, arr2) -> Image.Image:
+    hsv1 = skimage.color.rgb2hsv(rgb=np.array(arr1, dtype=np.uint8))
+    hsv2 = skimage.color.rgb2hsv(rgb=np.array(arr2, dtype=np.uint8))
+
+    hsv1[..., 0] = hsv2[..., 0]
+    hsv1[..., 1] = hsv2[..., 1]
+
+    out = skimage.color.hsv2rgb(hsv1)
+    return Image.fromarray(np.array(out * 255, dtype=np.uint8))
+
 def lerp(a: float, b: float, t: float) -> float:
     return (1 - t) * a + t * b