from comfy_extras.nodes_mask import ImageCompositeMasked import torch import torch.nn.functional as F class ImageCompositeWatermark(ImageCompositeMasked): @classmethod def INPUT_TYPES(s): return { "required": { "destination": ("IMAGE",), "watermark": ("IMAGE",), "position": (["top_right", "top_center", "top_left", "bottom_right", "bottom_center", "bottom_left"], {"default": "bottom_right"}), "resize_ratio": ("FLOAT", {"default": 1, "min": 0, "max": 10, "step": 0.05}), "margin": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), }, "optional": { "mask": ("MASK",), "enabled": ("BOOLEAN", {"default": True, "label_on": "enabled", "label_off": "disabled"}), "invert_mask": ("BOOLEAN", {"default": False, "label_on": "enabled", "label_off": "disabled"}), } } RETURN_TYPES = ("IMAGE",) FUNCTION = "composite_watermark" CATEGORY = "utils/image" def composite_watermark(self, destination, watermark, position, resize_ratio, margin, mask=None, enabled=True, invert_mask=False): if not enabled: return (destination,) dest_h, dest_w = destination.shape[1:3] water_h, water_w = watermark.shape[1:3] scale = 1 if water_h > dest_h or water_w > dest_w: # 计算需要的缩放比例 scale_h = dest_h / water_h scale_w = dest_w / water_w scale = min(scale_h, scale_w) if resize_ratio != 1 or scale != 1: watermark = torch.nn.functional.interpolate( watermark.movedim(-1, 1), scale_factor=resize_ratio * scale, mode="bicubic", antialias=True).movedim(1, -1).clamp(0.0, 1.0) if mask is not None: mask = torch.nn.functional.interpolate(mask.unsqueeze( 0), scale_factor=resize_ratio * scale, mode="bicubic", antialias=True).squeeze(0).clamp(0.0, 1.0) water_h, water_w = watermark.shape[1:3] # 计算y坐标 if position.startswith("top"): y = margin else: # bottom positions y = dest_h - water_h - margin x = 0 # 根据position计算x坐标 if position.endswith("left"): x = margin elif position.endswith("center"): x = (dest_w - water_w) // 2 elif position.endswith("right"): x = dest_w - water_w - margin if invert_mask and mask is not None: mask = 1.0 - mask return self.composite(destination, watermark, x, y, False, mask) class ImageTransition: @classmethod def INPUT_TYPES(s): return { "required": { "first_image": ("IMAGE",), "last_image": ("IMAGE",), "frames": ("INT", {"default": 24, "min": 2, "max": 120, "step": 1}), "transition_type": (["uniform", "smooth"], {"default": "uniform"}), "smooth_effect": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1}), } } RETURN_TYPES = ("IMAGE",) FUNCTION = "generate_transition" CATEGORY = "utils/image" def generate_transition(self, first_image, last_image, frames, transition_type="uniform", smooth_effect=1.0): # 生成插值权重 if transition_type == "uniform": weights = torch.linspace(0, 1, frames) else: # sigmoid x = torch.linspace(-20, 20, frames) weights = torch.sigmoid(x / smooth_effect) # 创建输出张量列表 output_frames = [] # 生成过渡帧 for w in weights: # 使用权重进行插值 transition_frame = first_image * (1 - w) + last_image * w output_frames.append(transition_frame) # 将所有帧拼接在一起 result = torch.cat(output_frames, dim=0) return (result,) class ImageMaskColorAverage: @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "mask": ("MASK",), } } RETURN_TYPES = ("INT", "STRING") RETURN_NAMES = ("COLOR_DEC", "COLOR_HEX") FUNCTION = "calculate_average_color" CATEGORY = "utils/image" def calculate_average_color(self, image, mask): # 确保mask是二维的 if len(mask.shape) > 2: mask = mask.squeeze() # 将mask扩展为与图像相同的通道数 expanded_mask = mask.unsqueeze(-1).expand(-1, -1, 3) # 计算mask区域的像素总数 pixel_count = torch.sum(mask) if pixel_count == 0: # 如果mask中没有选中区域，返回黑色 return (0, "#000000") # 计算mask区域的颜色总和 masked_image = image * expanded_mask.unsqueeze(0) color_sum = torch.sum(masked_image, dim=[0, 1, 2]) # 计算平均颜色 avg_color = color_sum / pixel_count # 转换为0-255范围的整数 r = int(avg_color[0].item() * 255) g = int(avg_color[1].item() * 255) b = int(avg_color[2].item() * 255) # 生成十六进制颜色代码 hex_color = f"#{r:02x}{g:02x}{b:02x}" # 计算十进制颜色值 (R*65536 + G*256 + B) dec_color = r * 65536 + g * 256 + b return (dec_color, hex_color) class ImagesConcanateToGrid: @classmethod def INPUT_TYPES(s): return {"required": { "image1": ("IMAGE",), "direction": ( ['right', 'down', ], { "default": 'right' }), "dimension_number": ("INT", {"default": 2, "min": 1, "max": 20, "step": 1}), } } RETURN_TYPES = ("IMAGE",) FUNCTION = "concanate" CATEGORY = "utils/image" def concanate(self, image1, direction='right', dimension_number=2): # 检查图像维度，如果不是4维直接返回 if len(image1.shape) != 4: return (image1,) batch_size = image1.shape[0] # 如果批次大小为1，直接返回 if batch_size == 1: return (image1,) # 将批次图像分离为单独的图像列表 images = [image1[i:i+1] for i in range(batch_size)] # 根据方向计算行数和列数 if direction == 'right': cols = dimension_number rows = (batch_size + cols - 1) // cols # 向上取整 else: # direction == 'down' rows = dimension_number cols = (batch_size + rows - 1) // rows # 向上取整 # 创建网格来存放图像 grid_rows = [] for row in range(rows): row_images = [] for col in range(cols): idx = row * cols + col if direction == 'right' else col * rows + row if idx < len(images): row_images.append(images[idx]) else: # 如果没有足够的图像，用黑色图像填充 black_image = torch.zeros_like(images[0]) row_images.append(black_image) # 水平拼接当前行的图像 if row_images: row_concat = torch.cat(row_images, dim=2) # 在宽度维度拼接 grid_rows.append(row_concat) # 垂直拼接所有行 if grid_rows: result = torch.cat(grid_rows, dim=1) # 在高度维度拼接 else: result = image1 return (result,) class NeedImageSizeAndCount: @classmethod def INPUT_TYPES(s): return {"required": {"image": ("IMAGE",)}} RETURN_TYPES = ("INT", "INT", "INT") RETURN_NAMES = ("width", "height", "count") FUNCTION = "get_image_size_and_count" CATEGORY = "utils/image" def get_image_size_and_count(self, image): return (image.shape[2], image.shape[1], image.shape[0]) NODE_CLASS_MAPPINGS = { #image "ImageCompositeWatermark": ImageCompositeWatermark, "ImageTransition": ImageTransition, "ImageMaskColorAverage": ImageMaskColorAverage, "ImagesConcanateToGrid": ImagesConcanateToGrid, "NeedImageSizeAndCount": NeedImageSizeAndCount, } NODE_DISPLAY_NAME_MAPPINGS = { # Image "ImageCompositeWatermark": "Image Composite Watermark", "ImageTransition": "Image Transition", "ImageMaskColorAverage": "Image Mask Color Average", "ImagesConcanateToGrid": "Images Concanate To Grid", "NeedImageSizeAndCount": "get Image Size And Count by utils", }