136 lines
4.9 KiB
Python
136 lines
4.9 KiB
Python
import torch
|
||
import numpy as np
|
||
from PIL import Image
|
||
from blend_modes import soft_light, lighten_only, dodge, addition, darken_only, multiply, hard_light, difference, subtract, grain_extract, grain_merge, divide, overlay, normal
|
||
import comfy.utils
|
||
|
||
class BlendImageNode:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"base_image": ("IMAGE", {"tooltip": "The base image to blend."}),
|
||
"overlay_image": ("IMAGE", {"tooltip": "The overlay image."}),
|
||
"blend_mode": ([
|
||
"soft_light", "lighten_only", "dodge", "addition", "darken_only",
|
||
"multiply", "hard_light", "difference", "subtract",
|
||
"grain_extract", "grain_merge", "divide", "overlay", "normal"
|
||
], {"default": "soft_light"}),
|
||
"opacity": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01
|
||
}),
|
||
"position": ([
|
||
"top_left", "top_center", "top_right",
|
||
"mid_left", "mid_center", "mid_right",
|
||
"bottom_left", "bottom_center", "bottom_right"
|
||
], {"default": "bottom_right", "tooltip": "Anchor point for placing the overlay."}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
FUNCTION = "blend_image"
|
||
CATEGORY = "image"
|
||
|
||
def pad_overlay_to_position(self, overlay, base_h, base_w, position):
|
||
"""Pad the overlay image to the full base image size, anchored according to position."""
|
||
overlay_h, overlay_w, _ = overlay.shape
|
||
pad_top, pad_bottom = 0, base_h - overlay_h
|
||
pad_left, pad_right = 0, base_w - overlay_w
|
||
|
||
# Vertical alignment
|
||
if "top" in position:
|
||
pad_top = 0
|
||
pad_bottom = base_h - overlay_h
|
||
elif "mid" in position:
|
||
pad_top = (base_h - overlay_h) // 2
|
||
pad_bottom = base_h - overlay_h - pad_top
|
||
elif "bottom" in position:
|
||
pad_top = base_h - overlay_h
|
||
pad_bottom = 0
|
||
|
||
# Horizontal alignment
|
||
if "left" in position:
|
||
pad_left = 0
|
||
pad_right = base_w - overlay_w
|
||
elif "center" in position:
|
||
pad_left = (base_w - overlay_w) // 2
|
||
pad_right = base_w - overlay_w - pad_left
|
||
elif "right" in position:
|
||
pad_left = base_w - overlay_w
|
||
pad_right = 0
|
||
|
||
return np.pad(
|
||
overlay,
|
||
((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)),
|
||
mode='constant',
|
||
constant_values=0
|
||
)
|
||
|
||
def to_rgba(self, img):
|
||
"""Ensure a (H, W, 3) or (H, W, 4) float32 image [0.0–1.0] is converted to (H, W, 4) RGBA [0.0–255.0]"""
|
||
h, w, c = img.shape
|
||
if c == 3:
|
||
alpha = np.ones((h, w, 1), dtype=np.float32)
|
||
img = np.concatenate([img, alpha], axis=2)
|
||
return (img * 255.0).astype(np.float32)
|
||
|
||
def from_rgba(self, img):
|
||
"""Convert (H, W, 4) float32 [0–255] back to RGB [0–1]"""
|
||
rgb = img[..., :3]
|
||
return np.clip(rgb / 255.0, 0.0, 1.0)
|
||
|
||
def blend_image(self, base_image, overlay_image, blend_mode, opacity, position):
|
||
# Get blending function
|
||
blend_functions = {
|
||
"soft_light": soft_light,
|
||
"lighten_only": lighten_only,
|
||
"dodge": dodge,
|
||
"addition": addition,
|
||
"darken_only": darken_only,
|
||
"multiply": multiply,
|
||
"hard_light": hard_light,
|
||
"difference": difference,
|
||
"subtract": subtract,
|
||
"grain_extract": grain_extract,
|
||
"grain_merge": grain_merge,
|
||
"divide": divide,
|
||
"overlay": overlay,
|
||
"normal": normal
|
||
}
|
||
blend_fn = blend_functions.get(blend_mode, soft_light)
|
||
|
||
# Prepare overlay image
|
||
overlay_np = overlay_image[0].cpu().numpy()
|
||
overlay_h, overlay_w, _ = overlay_np.shape
|
||
overlay_rgba = self.to_rgba(overlay_np)
|
||
|
||
image_count = base_image.shape[0]
|
||
pbar = comfy.utils.ProgressBar(image_count)
|
||
|
||
blended_images = []
|
||
for i in range(image_count):
|
||
base_np = base_image[i].cpu().numpy()
|
||
|
||
base_h, base_w, _ = base_np.shape
|
||
|
||
if overlay_h > base_h or overlay_w > base_w:
|
||
raise ValueError("Overlay image must be smaller than or equal to base image dimensions.")
|
||
|
||
base_rgba = self.to_rgba(base_np)
|
||
padded_overlay = self.pad_overlay_to_position(overlay_rgba, base_h, base_w, position)
|
||
|
||
blended = blend_fn(base_rgba, padded_overlay, opacity)
|
||
blended_rgb = self.from_rgba(blended)
|
||
blended_images.append(blended_rgb)
|
||
|
||
pbar.update(1)
|
||
|
||
print(f"Stacking {len(blended_images)} images...")
|
||
blended_tensor = torch.from_numpy(np.stack(blended_images)) # .unsqueeze(0)
|
||
|
||
return (blended_tensor,)
|
||
|