173 lines
7.1 KiB
Python
173 lines
7.1 KiB
Python
import torch
|
|
|
|
class PSBlendNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image1": ("IMAGE",),
|
|
"image2": ("IMAGE",),
|
|
"blend_mode": (["Normal", "Dissolve", "Darken", "Multiply", "Color Burn", "Linear Burn", "Darker Color", "Lighten", "Screen", "Color Dodge", "Linear Dodge", "Lighter Color", "Overlay", "Soft Light", "Hard Light", "Vivid Light", "Linear Light", "Pin Light", "Hard Mix", "Difference", "Exclusion", "Subtract", "Divide", "Hue", "Saturation", "Color", "Luminosity"],),
|
|
"opacity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "blend_images"
|
|
|
|
CATEGORY = "image/blending"
|
|
|
|
def blend_images(self, image1, image2, blend_mode, opacity):
|
|
# GPU availability check (remains the same)
|
|
if torch.cuda.is_available():
|
|
print("PS Blend Node is using CUDA.")
|
|
device = torch.device("cuda")
|
|
elif torch.backends.mps.is_available():
|
|
print("PS Blend Node is using MPS.")
|
|
device = torch.device("mps")
|
|
else:
|
|
print("No GPU acceleration available, using CPU")
|
|
device = torch.device("cpu")
|
|
|
|
# Move tensors to the appropriate device
|
|
img1 = image1.to(device)
|
|
img2 = image2.to(device)
|
|
|
|
|
|
# Ensure images have the same shape
|
|
if img1.shape != img2.shape:
|
|
raise ValueError("Images must have the same dimensions")
|
|
|
|
# Helper functions (update to handle RGBA)
|
|
def lum(img):
|
|
return 0.3 * img[..., 0] + 0.59 * img[..., 1] + 0.11 * img[..., 2]
|
|
|
|
def sat(img):
|
|
min_val, _ = torch.min(img[..., :3], dim=-1)
|
|
max_val, _ = torch.max(img[..., :3], dim=-1)
|
|
return max_val - min_val
|
|
|
|
def set_lum(c, l):
|
|
d = l - lum(c)
|
|
c_rgb = c[..., :3] + d.unsqueeze(-1)
|
|
return torch.cat([clip_color(c_rgb), c[..., 3:4]], dim=-1)
|
|
|
|
def set_sat(c, s):
|
|
c_rgb = c[..., :3]
|
|
c_min, _ = torch.min(c_rgb, dim=-1, keepdim=True)
|
|
c_max, _ = torch.max(c_rgb, dim=-1, keepdim=True)
|
|
c_mid = c_rgb - c_min
|
|
c_mid *= s.unsqueeze(-1) / (c_max - c_min + 1e-6)
|
|
c_rgb_new = clip_color(c_mid + c_min)
|
|
return torch.cat([c_rgb_new, c[..., 3:4]], dim=-1)
|
|
|
|
def clip_color(c):
|
|
l = lum(c).unsqueeze(-1)
|
|
n = torch.min(c, dim=-1, keepdim=True)[0]
|
|
x = torch.max(c, dim=-1, keepdim=True)[0]
|
|
c_clipped = torch.clamp(c, 0, 1)
|
|
c_clipped = torch.where(n < 0, l + ((c - l) * l) / (l - n + 1e-6), c_clipped)
|
|
c_clipped = torch.where(x > 1, l + ((c - l) * (1 - l)) / (x - l + 1e-6), c_clipped)
|
|
return c_clipped
|
|
|
|
def hue_blend(img1, img2):
|
|
lum2 = lum(img2)
|
|
sat2 = sat(img2)
|
|
result = set_lum(set_sat(img1, sat2), lum2)
|
|
return result
|
|
|
|
# Blend mode calculations (now including alpha)
|
|
if blend_mode == "Normal":
|
|
result = img1
|
|
elif blend_mode == "Dissolve":
|
|
mask = torch.rand_like(img1) < opacity
|
|
result = torch.where(mask, img1, img2)
|
|
elif blend_mode == "Darken":
|
|
result = torch.min(img1, img2)
|
|
elif blend_mode == "Multiply":
|
|
result = img1 * img2
|
|
elif blend_mode == "Color Burn":
|
|
result = torch.where(
|
|
img1 == 1,
|
|
img2,
|
|
torch.where(
|
|
img1 > 0,
|
|
1 - torch.clamp((1 - img2) / img1, 0, 1),
|
|
torch.zeros_like(img2)
|
|
)
|
|
)
|
|
elif blend_mode == "Linear Burn":
|
|
result = img1 + img2 - 1
|
|
elif blend_mode == "Darker Color":
|
|
result = torch.where(lum(img1).unsqueeze(-1) < lum(img2).unsqueeze(-1), img1, img2)
|
|
elif blend_mode == "Lighten":
|
|
result = torch.max(img1, img2)
|
|
elif blend_mode == "Screen":
|
|
result = 1 - (1 - img1) * (1 - img2)
|
|
elif blend_mode == "Color Dodge":
|
|
result = torch.where(img2 < 1, torch.min(torch.ones_like(img1), img1 / (1 - img2)), torch.ones_like(img1))
|
|
elif blend_mode == "Linear Dodge":
|
|
result = img1 + img2
|
|
elif blend_mode == "Lighter Color":
|
|
result = torch.where(lum(img1).unsqueeze(-1) > lum(img2).unsqueeze(-1), img1, img2)
|
|
elif blend_mode == "Overlay":
|
|
result = torch.where(img2 < 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
|
|
elif blend_mode == "Soft Light":
|
|
d = torch.where(img2 <= 0.25,
|
|
((16 * img2 - 12) * img2 + 4) * img2,
|
|
torch.sqrt(img2))
|
|
result = torch.where(img1 <= 0.5,
|
|
img2 - (1 - 2 * img1) * img2 * (1 - img2),
|
|
img2 + (2 * img1 - 1) * (d - img2))
|
|
elif blend_mode == "Hard Light":
|
|
result = torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
|
|
elif blend_mode == "Vivid Light":
|
|
result = torch.where(img1 < 0.5,
|
|
1 - (1 - img2) / (2 * img1 + 1e-8),
|
|
img2 / (2 * (1 - img1) + 1e-8))
|
|
elif blend_mode == "Linear Light":
|
|
result = img2 + 2 * img1 - 1
|
|
elif blend_mode == "Pin Light":
|
|
result = torch.where(img1 < 0.5,
|
|
torch.min(img2, 2 * img1),
|
|
torch.max(img2, 2 * img1 - 1))
|
|
elif blend_mode == "Hard Mix":
|
|
sum_img = img1 + img2
|
|
# Apply the Hard Mix formula
|
|
result = torch.where(sum_img >= 1, torch.ones_like(img1), torch.zeros_like(img1))
|
|
elif blend_mode == "Difference":
|
|
result = torch.abs(img1 - img2)
|
|
elif blend_mode == "Exclusion":
|
|
result = img1 + img2 - 2 * img1 * img2
|
|
elif blend_mode == "Subtract":
|
|
result = img2 - img1
|
|
elif blend_mode == "Divide":
|
|
result = img2 / (img1 + 1e-8)
|
|
elif blend_mode == "Hue":
|
|
result = hue_blend(img1, img2)
|
|
elif blend_mode == "Saturation":
|
|
result = set_lum(set_sat(img2, sat(img1)), lum(img2))
|
|
elif blend_mode == "Color":
|
|
result = set_lum(img1, lum(img2))
|
|
elif blend_mode == "Luminosity":
|
|
result = set_lum(img2, lum(img1))
|
|
else:
|
|
print(f"Unsupported blend mode: {blend_mode}")
|
|
result = img1 # Default to using the top image if blend mode is not recognized
|
|
|
|
# Apply opacity
|
|
opacity_tensor = torch.tensor(opacity, device=device).view(1, 1, 1, 1)
|
|
result = img2 * (1 - opacity_tensor) + result * opacity_tensor
|
|
|
|
# Ensure result is clamped between 0 and 1
|
|
result = torch.clamp(result, 0, 1)
|
|
|
|
# Debug print
|
|
self.print_tensor_shape(result, "final result")
|
|
|
|
return (result,)
|
|
|
|
# Debug method remains the same
|
|
def print_tensor_shape(self, tensor, name):
|
|
print(f"{name} shape: {tensor.shape}")
|