From 0a1b7acefc3706a63e5ce573d92d50651725ec55 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Fri, 18 Aug 2023 14:35:27 -0400
Subject: [PATCH] Add color blend node.

---
 __init__.py      |  5 ++--
 color_blend.py   | 72 ++++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt |  1 +
 3 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 color_blend.py

diff --git a/__init__.py b/__init__.py
index c1dc126..dce71f3 100644
--- a/__init__.py
+++ b/__init__.py
@@ -3,6 +3,7 @@ import os
 
 node_list = [ #Add list of .py files containing nodes here
     "control_lora_create",
+    "color_blend",
 ]
 
 NODE_CLASS_MAPPINGS = {}
@@ -11,7 +12,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {}
 for module_name in node_list:
     imported_module = importlib.import_module(".{}".format(module_name), __name__)
 
-    NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **control_lora_create.NODE_CLASS_MAPPINGS}
-    NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **control_lora_create.NODE_DISPLAY_NAME_MAPPINGS}
+    NODE_CLASS_MAPPINGS = {**NODE_CLASS_MAPPINGS, **imported_module.NODE_CLASS_MAPPINGS}
+    NODE_DISPLAY_NAME_MAPPINGS = {**NODE_DISPLAY_NAME_MAPPINGS, **imported_module.NODE_DISPLAY_NAME_MAPPINGS}
 
 __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
diff --git a/color_blend.py b/color_blend.py
new file mode 100644
index 0000000..955271f
--- /dev/null
+++ b/color_blend.py
@@ -0,0 +1,72 @@
+# Color blend node by Yam Levi
+# Property of Stability AI
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+import comfy.utils
+
+def color_blend(bw_layer,color_layer):
+    # Convert the color layer to LAB color space
+    color_lab = cv2.cvtColor(color_layer, cv2.COLOR_BGR2Lab)
+    # Convert the black and white layer to grayscale
+    bw_layer_gray = cv2.cvtColor(bw_layer, cv2.COLOR_BGR2GRAY)
+    # Replace the luminosity (L) channel in the color image with the black and white luminosity
+    _, color_a, color_b = cv2.split(color_lab)
+    blended_lab = cv2.merge((bw_layer_gray, color_a, color_b))
+    # Convert the blended LAB image back to BGR color space
+    blended_result = cv2.cvtColor(blended_lab, cv2.COLOR_Lab2BGR)
+    return blended_result
+
+class ColorBlend:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "bw_layer": ("IMAGE",),
+                "color_layer": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "color_blending_mode"
+
+    CATEGORY = "stability/image/postprocessing"
+
+    def color_blending_mode(self, bw_layer, color_layer):
+        if bw_layer.shape[0] < color_layer.shape[0]:
+            bw_layer = bw_layer.repeat(color_layer.shape[0], 1, 1, 1)[:color_layer.shape[0]]
+        if bw_layer.shape[0] > color_layer.shape[0]:
+            color_layer = color_layer.repeat(bw_layer.shape[0], 1, 1, 1)[:bw_layer.shape[0]]
+
+        batch_size, height, width, _ = bw_layer.shape
+        tensor_output = torch.empty_like(bw_layer)
+
+        image1 = bw_layer.cpu()
+        image2 = color_layer.cpu()
+        if image1.shape != image2.shape:
+            #print(image1.shape)
+            #print(image2.shape)
+            image2 = image2.permute(0, 3, 1, 2)
+            image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center')
+            image2 = image2.permute(0, 2, 3, 1)
+        image1  = (image1 * 255).to(torch.uint8).numpy()
+        image2 = (image2 * 255).to(torch.uint8).numpy()
+
+        for i in range(batch_size):
+            blend = color_blend(image1[i],image2[i])
+            blend = np.stack([blend])
+            tensor_output[i:i+1] = (torch.from_numpy(blend.transpose(0, 3, 1, 2))/255.0).permute(0, 2, 3, 1)
+
+        return (tensor_output,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "ColorBlend": ColorBlend
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ColorBlend": "Color Blend"
+}
diff --git a/requirements.txt b/requirements.txt
index e69de29..0dd006b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -0,0 +1 @@
+opencv-python