From 97c7a402d7c121b901dc0d0721e3ae163e774ea7 Mon Sep 17 00:00:00 2001
From: missionfloyd <missionfloyd@users.noreply.github.com>
Date: Sun, 9 Apr 2023 19:08:17 -0600
Subject: [PATCH] Merge channels node

---
 comfy_extras/nodes_post_processing.py | 52 +++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py
index b58267c12..9a8cdc846 100644
--- a/comfy_extras/nodes_post_processing.py
+++ b/comfy_extras/nodes_post_processing.py
@@ -395,6 +395,45 @@ class Split:
 
         return (result_r, result_g, result_b)
 
+class Merge:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "red": ("IMAGE",),
+                "green": ("IMAGE",),
+                "blue": ("IMAGE",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "merge"
+
+    CATEGORY = "image/postprocessing"
+
+    def merge(self, red: torch.Tensor, green: torch.Tensor, blue: torch.Tensor):
+        batch_size, height, width, _ = red.shape
+        result = torch.zeros_like(red)
+
+        for b in range(batch_size):
+            img_r = (red[b] * 255).to(torch.uint8).numpy()
+            img_g = (green[b] * 255).to(torch.uint8).numpy()
+            img_b = (blue[b] * 255).to(torch.uint8).numpy()
+            pil_image_r = Image.fromarray(img_r, mode='RGB').convert("L")
+            pil_image_g = Image.fromarray(img_g, mode='RGB').convert("L")
+            pil_image_b = Image.fromarray(img_b, mode='RGB').convert("L")
+
+            output_image = Image.merge("RGB", (pil_image_r, pil_image_g, pil_image_b))
+
+            output_array = torch.tensor(np.array(output_image.convert("RGB"))).float() / 255
+            result[b] = output_array
+
+        return (result,)
+
+
 
 NODE_CLASS_MAPPINGS = {
     "ImageBlend": Blend,
@@ -405,4 +444,17 @@ NODE_CLASS_MAPPINGS = {
     "ImageRotate": Rotate,
     "ImageGetChannel": GetChannel,
     "ImageSplit": Split,
+    "ImageMerge": Merge,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS  = {
+    "ImageBlend": "Blend Images",
+    "ImageBlur": "Blur Image",
+    "ImageQuantize": "Quantize Image",
+    "ImageSharpen": "Sharpen Image",
+    "ImageTranspose": "Transpose",
+    "ImageRotate": "Rotate",
+    "ImageGetChannel": "Extract Channel",
+    "ImageSplit": "Split Channels",
+    "ImageMerge": "Merge Channels",
 }