Add files via upload
This commit is contained in:
LaoMaoBoss
@@ -0,0 +1,503 @@
|
||||
"""
|
||||
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)
|
||||
"""
|
||||
|
||||
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", {}),
|
||||
},
|
||||
"optional": {
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
RETURN_NAMES = ("image", "mask")
|
||||
FUNCTION = "generate_gradient"
|
||||
CATEGORY = f"🌈WBLESS"
|
||||
|
||||
@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) -> Tuple[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}")
|
||||
|
||||
# 解析颜色
|
||||
start_rgb = self.get_color_values(start_color, start_color_hex)
|
||||
end_rgb = self.get_color_values(end_color, end_color_hex)
|
||||
|
||||
# 处理可能是列表格式的参数
|
||||
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
|
||||
|
||||
# 将输入图像转换为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)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user