Add color blend node.
This commit is contained in:
+3
-2
@@ -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']
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
opencv-python
|
||||
|
||||
Reference in New Issue
Block a user