78 lines
2.5 KiB
Python
78 lines
2.5 KiB
Python
# Adapt from https://github.com/Stability-AI/stability-ComfyUI-nodes/blob/master/color_blend.py
|
|
|
|
# Color blend node by Yam Levi
|
|
# Property of Stability AI
|
|
import cv2
|
|
import torch
|
|
import numpy as np
|
|
|
|
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 = "Art Venture/Post Processing"
|
|
|
|
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, *_ = bw_layer.shape
|
|
tensor_output = torch.empty_like(bw_layer)
|
|
|
|
image1 = bw_layer.cpu()
|
|
image2 = color_layer.cpu()
|
|
if image1.shape != 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,)
|