545 lines
20 KiB
Python
545 lines
20 KiB
Python
"""
|
||
WBLESS 渐变工具节点
|
||
|
||
这个模块实现了强大的渐变生成功能，支持多种渐变类型和自定义参数。
|
||
"""
|
||
|
||
import time
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image, ImageDraw
|
||
import math
|
||
from typing import Tuple, List, Dict, Any
|
||
|
||
from cozy_comfyui.node import CozyBaseNode
|
||
|
||
# 颜色映射，参考 Text Block 节点的实现
|
||
COLOR_MAPPING = {
|
||
"white": (255, 255, 255),
|
||
"black": (0, 0, 0),
|
||
"red": (255, 0, 0),
|
||
"green": (0, 255, 0),
|
||
"blue": (0, 0, 255),
|
||
"yellow": (255, 255, 0),
|
||
"cyan": (0, 255, 255),
|
||
"magenta": (255, 0, 255),
|
||
"orange": (255, 165, 0),
|
||
"purple": (128, 0, 128),
|
||
"pink": (255, 192, 203),
|
||
"brown": (160, 85, 15),
|
||
"gray": (128, 128, 128),
|
||
"lightgray": (211, 211, 211),
|
||
"darkgray": (102, 102, 102),
|
||
"olive": (128, 128, 0),
|
||
"lime": (0, 128, 0),
|
||
"teal": (0, 128, 128),
|
||
"navy": (0, 0, 128),
|
||
"maroon": (128, 0, 0),
|
||
"fuchsia": (255, 0, 128),
|
||
"aqua": (0, 255, 128),
|
||
"silver": (192, 192, 192),
|
||
"gold": (255, 215, 0),
|
||
"turquoise": (64, 224, 208),
|
||
"lavender": (230, 230, 250),
|
||
"violet": (238, 130, 238),
|
||
"coral": (255, 127, 80),
|
||
"indigo": (75, 0, 130),
|
||
"custom": (255, 255, 255) # 自定义颜色的占位符
|
||
}
|
||
|
||
|
||
class GradientNode(CozyBaseNode):
|
||
"""
|
||
渐变工具节点 - 生成各种类型的渐变图像
|
||
|
||
支持的渐变类型：
|
||
- 线性渐变 (Linear)
|
||
- 径向渐变 (Radial)
|
||
- 角度渐变 (Angular)
|
||
- 菱形渐变 (Diamond)
|
||
- 椭圆渐变 (Elliptical)
|
||
|
||
输出说明：
|
||
- image: 渐变与输入图像混合后的结果
|
||
- mask: 基于渐变透明度的遮罩，可用于后处理
|
||
|
||
反转选项：
|
||
- invert_alpha: 反转透明度渐变（交换起始和结束透明度）
|
||
- invert_colors: 反转渐变颜色（交换起始和结束颜色）
|
||
"""
|
||
|
||
NAME = "Gradient"
|
||
|
||
# 渐变类型选项
|
||
GRADIENT_TYPES = [
|
||
"linear", # 线性渐变
|
||
"radial", # 径向渐变
|
||
"angular", # 角度渐变
|
||
"diamond", # 菱形渐变
|
||
"elliptical" # 椭圆渐变
|
||
]
|
||
|
||
|
||
|
||
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
||
# 定义颜色选项，custom放在第一位
|
||
COLORS = ["custom"] + [color for color in COLOR_MAPPING.keys() if color != "custom"]
|
||
|
||
return {
|
||
"required": {
|
||
# 渐变类型
|
||
"gradient_type": (cls.GRADIENT_TYPES, {
|
||
"default": "linear"
|
||
}),
|
||
|
||
# 旋转角度（度）
|
||
"rotation_angle": ("FLOAT", {
|
||
"default": 0.0,
|
||
"min": -360.0,
|
||
"max": 360.0,
|
||
"step": 1.0,
|
||
"display": "number"
|
||
}),
|
||
|
||
# 渐变位置设置
|
||
"start_position": ("FLOAT", {
|
||
"default": 0.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"display": "number"
|
||
}),
|
||
"end_position": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"display": "number"
|
||
}),
|
||
"center_position": ("FLOAT", {
|
||
"default": 0.5,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"display": "number"
|
||
}),
|
||
|
||
# 起始颜色参数
|
||
"start_color": (COLORS, {"default": "custom"}),
|
||
"start_color_hex": ("STRING", {
|
||
"default": "#000000",
|
||
"multiline": False
|
||
}),
|
||
|
||
# 结束颜色参数
|
||
"end_color": (COLORS, {"default": "custom"}),
|
||
"end_color_hex": ("STRING", {
|
||
"default": "#FFFFFF",
|
||
"multiline": False
|
||
}),
|
||
|
||
# 透明度参数
|
||
"start_alpha": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"display": "number"
|
||
}),
|
||
"end_alpha": ("FLOAT", {
|
||
"default": 1.0,
|
||
"min": 0.0,
|
||
"max": 1.0,
|
||
"step": 0.01,
|
||
"display": "number"
|
||
}),
|
||
|
||
# 输入图像，用于获取尺寸
|
||
"image": ("IMAGE", {}),
|
||
|
||
# 反转选项
|
||
"invert_alpha": ("BOOLEAN", {
|
||
"default": False,
|
||
"label_on": "true",
|
||
"label_off": "false"
|
||
}),
|
||
"invert_colors": ("BOOLEAN", {
|
||
"default": False,
|
||
"label_on": "true",
|
||
"label_off": "false"
|
||
}),
|
||
},
|
||
"optional": {
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "MASK")
|
||
RETURN_NAMES = ("image", "mask")
|
||
FUNCTION = "generate_gradient"
|
||
CATEGORY = f"🌈WBLESS"
|
||
|
||
OUTPUT_NODE = False
|
||
|
||
# 输出说明:
|
||
# - image: 渐变与输入图像混合后的结果
|
||
# - mask: 基于渐变透明度的遮罩
|
||
# 反转选项:
|
||
# - invert_alpha: 反转透明度渐变（交换起始和结束透明度）
|
||
# - invert_colors: 反转渐变颜色（交换起始和结束颜色）
|
||
|
||
@classmethod
|
||
def IS_CHANGED(cls, *args, **kwargs):
|
||
"""强制禁用缓存，确保每次都重新生成渐变"""
|
||
return time.time()
|
||
|
||
def hex_to_rgb(self, hex_color) -> Tuple[int, int, int]:
|
||
"""将十六进制颜色转换为RGB"""
|
||
# 处理列表格式的输入
|
||
if isinstance(hex_color, list):
|
||
hex_color = hex_color[0] if hex_color else "#000000"
|
||
|
||
# 确保是字符串类型
|
||
if not isinstance(hex_color, str):
|
||
hex_color = str(hex_color)
|
||
|
||
hex_color = hex_color.lstrip('#')
|
||
if len(hex_color) != 6:
|
||
# 默认颜色
|
||
return (0, 0, 0)
|
||
try:
|
||
return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
|
||
except ValueError:
|
||
return (0, 0, 0)
|
||
|
||
def get_color_values(self, color_name, color_hex) -> Tuple[int, int, int]:
|
||
"""获取颜色的RGB值，支持预设颜色和自定义十六进制颜色"""
|
||
# 处理可能的列表类型输入
|
||
if isinstance(color_name, list):
|
||
color_name = color_name[0] if color_name else "black"
|
||
if isinstance(color_hex, list):
|
||
color_hex = color_hex[0] if color_hex else "#000000"
|
||
|
||
if color_name == "custom":
|
||
return self.hex_to_rgb(color_hex)
|
||
else:
|
||
return COLOR_MAPPING.get(color_name, (0, 0, 0)) # 默认为黑色
|
||
|
||
|
||
|
||
def interpolate_color(self, color1: Tuple[int, int, int], alpha1: float,
|
||
color2: Tuple[int, int, int], alpha2: float,
|
||
t: float) -> Tuple[int, int, int, int]:
|
||
"""在两种颜色之间进行RGBA线性插值"""
|
||
r = int(color1[0] * (1 - t) + color2[0] * t)
|
||
g = int(color1[1] * (1 - t) + color2[1] * t)
|
||
b = int(color1[2] * (1 - t) + color2[2] * t)
|
||
a = int((alpha1 * (1 - t) + alpha2 * t) * 255)
|
||
return (r, g, b, a)
|
||
|
||
|
||
|
||
def get_gradient_value(self, x: int, y: int, width: int, height: int,
|
||
gradient_type: str, rotation_angle: float,
|
||
start_position: float, end_position: float, center_position: float) -> float:
|
||
"""计算指定位置的渐变值（0-1）"""
|
||
|
||
if gradient_type == "linear":
|
||
# 线性渐变 - 类似PS的渐变工具
|
||
angle_rad = math.radians(rotation_angle)
|
||
cos_a = math.cos(angle_rad)
|
||
sin_a = math.sin(angle_rad)
|
||
|
||
# 图像中心点
|
||
img_cx, img_cy = width / 2, height / 2
|
||
|
||
# 将坐标转换到图像中心为原点的系统
|
||
dx, dy = x - img_cx, y - img_cy
|
||
|
||
# 投影到渐变方向
|
||
projection = dx * cos_a + dy * sin_a
|
||
|
||
# 根据角度确定主要方向的尺寸
|
||
# 0度=水平，90度=垂直
|
||
angle_normalized = abs(rotation_angle % 180)
|
||
if angle_normalized <= 45 or angle_normalized >= 135:
|
||
# 主要是水平方向
|
||
main_dimension = width
|
||
else:
|
||
# 主要是垂直方向
|
||
main_dimension = height
|
||
|
||
# 计算起始和结束的实际位置
|
||
half_dim = main_dimension / 2
|
||
start_pos = (start_position - 0.5) * main_dimension
|
||
end_pos = (end_position - 0.5) * main_dimension
|
||
|
||
# 计算在起始和结束位置之间的比例
|
||
if abs(end_pos - start_pos) < 1e-6:
|
||
# 避免除零
|
||
raw_t = 0.5
|
||
else:
|
||
raw_t = (projection - start_pos) / (end_pos - start_pos)
|
||
|
||
# 应用中心点偏移（类似PS中的中间拉杆）
|
||
if center_position != 0.5:
|
||
if raw_t <= center_position:
|
||
# 在中心点左侧，压缩到0-0.5
|
||
if center_position > 0:
|
||
t = (raw_t / center_position) * 0.5
|
||
else:
|
||
t = 0.0
|
||
else:
|
||
# 在中心点右侧，压缩到0.5-1.0
|
||
if center_position < 1.0:
|
||
t = 0.5 + ((raw_t - center_position) / (1.0 - center_position)) * 0.5
|
||
else:
|
||
t = 1.0
|
||
else:
|
||
t = raw_t
|
||
|
||
elif gradient_type == "radial":
|
||
# 径向渐变 - 从图像中心向外
|
||
img_cx, img_cy = width / 2, height / 2
|
||
max_radius = math.sqrt(width**2 + height**2) / 2
|
||
distance = math.sqrt((x - img_cx)**2 + (y - img_cy)**2)
|
||
|
||
# 应用起始和结束位置
|
||
start_radius = start_position * max_radius
|
||
end_radius = end_position * max_radius
|
||
|
||
if abs(end_radius - start_radius) < 1e-6:
|
||
raw_t = 0.5
|
||
else:
|
||
raw_t = (distance - start_radius) / (end_radius - start_radius)
|
||
|
||
# 应用中心点偏移
|
||
if center_position != 0.5:
|
||
if raw_t <= center_position:
|
||
if center_position > 0:
|
||
t = (raw_t / center_position) * 0.5
|
||
else:
|
||
t = 0.0
|
||
else:
|
||
if center_position < 1.0:
|
||
t = 0.5 + ((raw_t - center_position) / (1.0 - center_position)) * 0.5
|
||
else:
|
||
t = 1.0
|
||
else:
|
||
t = raw_t
|
||
|
||
elif gradient_type == "angular":
|
||
# 角度渐变 - 围绕图像中心
|
||
img_cx, img_cy = width / 2, height / 2
|
||
angle = math.atan2(y - img_cy, x - img_cx)
|
||
raw_t = (angle + math.pi) / (2 * math.pi)
|
||
|
||
# 角度渐变的起始和结束位置表示角度范围
|
||
start_angle = start_position * 2 * math.pi
|
||
end_angle = end_position * 2 * math.pi
|
||
current_angle = raw_t * 2 * math.pi
|
||
|
||
if abs(end_angle - start_angle) < 1e-6:
|
||
raw_t = 0.5
|
||
else:
|
||
raw_t = (current_angle - start_angle) / (end_angle - start_angle)
|
||
|
||
# 应用中心点偏移
|
||
if center_position != 0.5:
|
||
if raw_t <= center_position:
|
||
if center_position > 0:
|
||
t = (raw_t / center_position) * 0.5
|
||
else:
|
||
t = 0.0
|
||
else:
|
||
if center_position < 1.0:
|
||
t = 0.5 + ((raw_t - center_position) / (1.0 - center_position)) * 0.5
|
||
else:
|
||
t = 1.0
|
||
else:
|
||
t = raw_t
|
||
|
||
elif gradient_type == "diamond":
|
||
# 菱形渐变 - 从图像中心
|
||
img_cx, img_cy = width / 2, height / 2
|
||
dx = abs(x - img_cx) / (width / 2)
|
||
dy = abs(y - img_cy) / (height / 2)
|
||
raw_t = max(dx, dy)
|
||
|
||
# 应用起始和结束位置
|
||
if abs(end_position - start_position) < 1e-6:
|
||
raw_t = 0.5
|
||
else:
|
||
raw_t = (raw_t - start_position) / (end_position - start_position)
|
||
|
||
# 应用中心点偏移
|
||
if center_position != 0.5:
|
||
if raw_t <= center_position:
|
||
if center_position > 0:
|
||
t = (raw_t / center_position) * 0.5
|
||
else:
|
||
t = 0.0
|
||
else:
|
||
if center_position < 1.0:
|
||
t = 0.5 + ((raw_t - center_position) / (1.0 - center_position)) * 0.5
|
||
else:
|
||
t = 1.0
|
||
else:
|
||
t = raw_t
|
||
|
||
elif gradient_type == "elliptical":
|
||
# 椭圆渐变 - 从图像中心
|
||
img_cx, img_cy = width / 2, height / 2
|
||
dx = (x - img_cx) / (width / 2)
|
||
dy = (y - img_cy) / (height / 2)
|
||
raw_t = math.sqrt(dx**2 + dy**2)
|
||
|
||
# 应用起始和结束位置
|
||
if abs(end_position - start_position) < 1e-6:
|
||
raw_t = 0.5
|
||
else:
|
||
raw_t = (raw_t - start_position) / (end_position - start_position)
|
||
|
||
# 应用中心点偏移
|
||
if center_position != 0.5:
|
||
if raw_t <= center_position:
|
||
if center_position > 0:
|
||
t = (raw_t / center_position) * 0.5
|
||
else:
|
||
t = 0.0
|
||
else:
|
||
if center_position < 1.0:
|
||
t = 0.5 + ((raw_t - center_position) / (1.0 - center_position)) * 0.5
|
||
else:
|
||
t = 1.0
|
||
else:
|
||
t = raw_t
|
||
|
||
else:
|
||
t = 0.0
|
||
|
||
# 限制在[0, 1]范围内
|
||
return max(0.0, min(1.0, t))
|
||
|
||
def generate_gradient(self, gradient_type: str, rotation_angle: float,
|
||
start_position: float, end_position: float, center_position: float,
|
||
start_color: str, start_color_hex: str, end_color: str, end_color_hex: str,
|
||
start_alpha: float, end_alpha: float, image,
|
||
invert_alpha: bool, invert_colors: bool) -> Tuple[torch.Tensor, torch.Tensor]:
|
||
"""生成渐变图像，支持透明度和颜色反转"""
|
||
|
||
# 从输入图像获取尺寸
|
||
if isinstance(image, list):
|
||
# 如果是列表，取第一个元素
|
||
img_tensor = image[0]
|
||
else:
|
||
img_tensor = image
|
||
|
||
# 处理不同的张量形状
|
||
if len(img_tensor.shape) == 4:
|
||
# 形状: (batch, height, width, channels)
|
||
img_tensor = img_tensor[0] # 取第一个批次
|
||
height, width = img_tensor.shape[:2]
|
||
elif len(img_tensor.shape) == 3:
|
||
# 形状: (height, width, channels)
|
||
height, width = img_tensor.shape[:2]
|
||
else:
|
||
raise ValueError(f"Unsupported image tensor shape: {img_tensor.shape}")
|
||
|
||
# 处理可能是列表格式的参数
|
||
if isinstance(start_alpha, list):
|
||
start_alpha = start_alpha[0] if start_alpha else 1.0
|
||
if isinstance(end_alpha, list):
|
||
end_alpha = end_alpha[0] if end_alpha else 1.0
|
||
if isinstance(gradient_type, list):
|
||
gradient_type = gradient_type[0] if gradient_type else "linear"
|
||
if isinstance(rotation_angle, list):
|
||
rotation_angle = rotation_angle[0] if rotation_angle else 0.0
|
||
if isinstance(start_position, list):
|
||
start_position = start_position[0] if start_position else 0.0
|
||
if isinstance(end_position, list):
|
||
end_position = end_position[0] if end_position else 1.0
|
||
if isinstance(center_position, list):
|
||
center_position = center_position[0] if center_position else 0.5
|
||
if isinstance(invert_alpha, list):
|
||
invert_alpha = invert_alpha[0] if invert_alpha else False
|
||
if isinstance(invert_colors, list):
|
||
invert_colors = invert_colors[0] if invert_colors else False
|
||
|
||
# 解析颜色
|
||
start_rgb = self.get_color_values(start_color, start_color_hex)
|
||
end_rgb = self.get_color_values(end_color, end_color_hex)
|
||
|
||
# 应用颜色反转 - 交换起始和结束颜色
|
||
if invert_colors:
|
||
start_rgb, end_rgb = end_rgb, start_rgb
|
||
|
||
# 应用透明度反转 - 交换起始和结束透明度
|
||
if invert_alpha:
|
||
start_alpha, end_alpha = end_alpha, start_alpha
|
||
|
||
# 将输入图像转换为numpy数组
|
||
# 确保张量在CPU上并且形状正确
|
||
if img_tensor.device.type != 'cpu':
|
||
img_tensor = img_tensor.cpu()
|
||
|
||
# 确保张量是3维的 (height, width, channels)
|
||
if len(img_tensor.shape) == 4:
|
||
img_tensor = img_tensor[0] # 移除批次维度
|
||
|
||
input_array = (img_tensor.numpy() * 255).astype(np.uint8)
|
||
|
||
# 创建渐变遮罩数组 - 使用RGBA格式
|
||
gradient_array = np.zeros((height, width, 4), dtype=np.uint8)
|
||
|
||
# 生成渐变
|
||
for y in range(height):
|
||
for x in range(width):
|
||
# 获取渐变值
|
||
t = self.get_gradient_value(
|
||
x, y, width, height, gradient_type, rotation_angle,
|
||
start_position, end_position, center_position
|
||
)
|
||
|
||
# 计算颜色 - 使用完整的RGBA渐变
|
||
color = self.interpolate_color(start_rgb, start_alpha, end_rgb, end_alpha, t)
|
||
gradient_array[y, x] = color
|
||
|
||
# 将输入图像转换为RGBA格式（如果需要）
|
||
if input_array.shape[2] == 3:
|
||
# RGB转RGBA，添加完全不透明的alpha通道
|
||
alpha_channel = np.full((height, width, 1), 255, dtype=np.uint8)
|
||
input_array = np.concatenate([input_array, alpha_channel], axis=2)
|
||
|
||
# 执行alpha混合：将渐变叠加到输入图像上
|
||
# 使用标准的alpha混合公式：result = foreground * alpha + background * (1 - alpha)
|
||
gradient_alpha = gradient_array[:, :, 3:4] / 255.0 # 渐变的alpha通道
|
||
input_alpha = input_array[:, :, 3:4] / 255.0 # 输入图像的alpha通道
|
||
|
||
# 计算最终的颜色
|
||
blended_rgb = (gradient_array[:, :, :3] * gradient_alpha +
|
||
input_array[:, :, :3] * (1 - gradient_alpha))
|
||
|
||
# 计算最终的alpha通道
|
||
blended_alpha = gradient_alpha + input_alpha * (1 - gradient_alpha)
|
||
|
||
# 合并RGB和Alpha
|
||
final_array = np.concatenate([blended_rgb, blended_alpha * 255], axis=2).astype(np.uint8)
|
||
|
||
# 转换为torch张量
|
||
image_tensor = torch.from_numpy(final_array).float() / 255.0
|
||
image_tensor = image_tensor.unsqueeze(0) # 添加批次维度
|
||
|
||
# 生成遮罩 - 基于渐变的透明度
|
||
# 提取渐变的alpha通道作为遮罩
|
||
mask_array = gradient_array[:, :, 3] # 只取alpha通道
|
||
mask_tensor = torch.from_numpy(mask_array).float() / 255.0
|
||
mask_tensor = mask_tensor.unsqueeze(0) # 添加批次维度
|
||
|
||
return (image_tensor, mask_tensor)
|
||
|
||
|
||
|