diff --git a/core/gradient.py b/core/gradient.py new file mode 100644 index 0000000..a742275 --- /dev/null +++ b/core/gradient.py @@ -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) + + +