import torch import torch.nn.functional as F import numpy as np from scipy import ndimage # --- 原有的文本节点 --- class MultipleTextInputs: @classmethod def INPUT_TYPES(s): return { "required": { "inputcount": ("INT", {"default": 3, "min": 2, "max": 1000, "step": 1}), "string_1": ("STRING", {"default": '', "forceInput": True}), "delimiter": ("STRING", {"default": '', "multiline": False}), }, "optional": { "string_2": ("STRING", {"default": '', "forceInput": True}), "string_3": ("STRING", {"default": '', "forceInput": True}), } } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("string",) FUNCTION = "combine" CATEGORY = "MyTemplate" def combine(self, inputcount, delimiter, **kwargs): result_list = [] if "string_1" in kwargs: result_list.append(kwargs["string_1"]) for c in range(1, inputcount): key = f"string_{c + 1}" new_string = kwargs.get(key, "") if new_string: result_list.append(new_string) final_string = delimiter.join(result_list) return (final_string,) # --- 原有的 RGB 遮罩节点 (保留) --- class MultipleMasksToImage: def __init__(self): self.colors = self._generate_color_palette() def _generate_color_palette(self): levels = [0, 255, 128, 192, 64] palette = [] for r in levels: for g in levels: for b in levels: if r == 0 and g == 0 and b == 0: continue palette.append((r, g, b)) def sort_key(color): r, g, b = color non_zero = sum(1 for c in [r, g, b] if c > 0) is_primary = 1 if non_zero == 1 and 255 in [r, g, b] else 0 score = 0 if is_primary: score += 1000 elif non_zero == 1: score += 500 elif non_zero == 2 and 255 in [r, g, b]: score += 300 brightness = r + g + b return (-score, -brightness) palette.sort(key=sort_key) return palette @classmethod def INPUT_TYPES(s): return { "required": { "inputcount": ("INT", { "default": 3, "min": 2, "max": 100, "step": 1, "tooltip": "Number of mask inputs. Click 'Update inputs' to refresh ports." }), "mask_1": ("MASK", { "tooltip": "The base mask. All other masks will be resized to match this mask's dimensions." }), "binarize": ("BOOLEAN", { "default": False, "label_on": "Enabled", "label_off": "Disabled", "tooltip": "If enabled, converts soft mask edges (gray) to hard edges (black/white)." }), "binarize_threshold": ("FLOAT", { "default": 0.2, "min": 0.01, "max": 0.99, "step": 0.01, "tooltip": "The pixel value threshold. Pixels brighter than this become 1, darker become 0." }), "remove_noise": ("BOOLEAN", { "default": False, "label_on": "Enabled", "label_off": "Disabled", "tooltip": "If enabled, removes small isolated pixel islands (dust/noise) from the mask." }), }, "optional": { "mask_2": ("MASK",), "mask_3": ("MASK",), } } RETURN_TYPES = ("IMAGE", "STRING") RETURN_NAMES = ("image", "color_info") FUNCTION = "combine_masks" CATEGORY = "MyTemplate" def combine_masks(self, inputcount, mask_1, binarize, binarize_threshold, remove_noise, **kwargs): # 1. 收集输入 (Collection) raw_inputs = [mask_1] for c in range(1, inputcount): key = f"mask_{c + 1}" m = kwargs.get(key, None) if m is not None: raw_inputs.append(m) # 2. 深度展平 (Deep Flattening & Batch Unpacking) # 将 List 和 Batch Tensor 全部拆解为单独的 [1, H, W] Tensor masks_to_process = [] for item in raw_inputs: # Handle List input if isinstance(item, list): sub_items = item else: sub_items = [item] for m in sub_items: # 统一维度为 [B, H, W] if m.dim() == 2: m = m.unsqueeze(0) # Handle Batch > 1 batch_count = m.shape[0] if batch_count == 1: masks_to_process.append(m) else: # 拆分 Batch 为多个独立 Mask for i in range(batch_count): masks_to_process.append(m[i:i+1]) if not masks_to_process: # 防御性编程，理论上必有 mask_1 raise ValueError("No masks provided.") # 检查颜色池是否足够 num_masks = len(masks_to_process) if num_masks > len(self.colors): raise ValueError(f"Too many masks! Maximum supported is {len(self.colors)}, but got {num_masks} (including batch items).") # 3. 确定基准尺寸 (使用展平后的第一个 Mask) ref_mask = masks_to_process[0] # ref_mask 已经是 [1, H, W] target_h, target_w = ref_mask.shape[1], ref_mask.shape[2] # 创建纯黑底图 [1, H, W, 3] (RGB) output_image = torch.zeros((1, target_h, target_w, 3), dtype=torch.float32, device=ref_mask.device) used_colors_rgb = [] used_colors_hex = [] # 4. 遍历处理并叠加 for i, mask in enumerate(masks_to_process): # 获取颜色 r_int, g_int, b_int = self.colors[i] color_float = (r_int / 255.0, g_int / 255.0, b_int / 255.0) color_tensor = torch.tensor(color_float, device=mask.device, dtype=torch.float32) used_colors_rgb.append(f"({r_int}, {g_int}, {b_int})") used_colors_hex.append(f"0x{r_int:02x}{g_int:02x}{b_int:02x}") # A. 缩放 # mask 是 [1, H, W]，interpolate 需要 [1, 1, H, W] mask_in = mask.unsqueeze(1) if mask.shape[1] != target_h or mask.shape[2] != target_w: mask_in = F.interpolate(mask_in, size=(target_h, target_w), mode="nearest") # 取回 [H, W] current_mask = mask_in[0, 0, :, :] # B. 二值化 if binarize: current_mask = (current_mask > binarize_threshold).float() # C. 去噪 if remove_noise: np_mask = current_mask.cpu().numpy() binary_mask = np_mask > 0.1 labeled, num_features = ndimage.label(binary_mask) component_sizes = ndimage.sum(binary_mask, labeled, range(num_features + 1)) too_small = component_sizes < 25 remove_pixel = too_small[labeled] np_mask[remove_pixel] = 0 current_mask = torch.from_numpy(np_mask).to(mask.device) # D. 叠加颜色 # output_image 是 [1, H, W, 3] # current_mask 是 [H, W] mask_3d = current_mask.unsqueeze(-1) # [H, W, 1] mask_indices = mask_3d > 0.01 # 扩展维度以匹配 output_image [1, H, W, 3] # mask_indices: [H, W, 1] -> [1, H, W, 1] -> broadcast to [1, H, W, 3] mask_indices_expanded = mask_indices.unsqueeze(0).expand(-1, -1, -1, 3) # color_tensor: [3] -> [1, 1, 1, 3] color_target = color_tensor.view(1, 1, 1, 3).expand(1, target_h, target_w, 3) # 使用 torch.where 进行叠加 output_image = torch.where(mask_indices_expanded, color_target, output_image) rgb_str = ", ".join(used_colors_rgb) hex_str = " , ".join(used_colors_hex) info_string = f"including colors :\nRGB : {rgb_str}\nHex color : {hex_str}" return (output_image, info_string) class MultipleImageStitch: @classmethod def INPUT_TYPES(s): return { "required": { "inputcount": ("INT", {"default": 3, "min": 2, "max": 50, "step": 1, "tooltip": "Number of inputs."}), "image_1": ("IMAGE", {"tooltip": "Base image for dimensions."}), "direction": (["right", "down"], {"default": "right", "tooltip": "Stitch direction."}), "size": ("INT", {"default": -1, "min": -1, "max": 8192, "step": 1, "tooltip": "Resize result. -1 to keep original."}), }, "optional": { "image_2": ("IMAGE",), "image_3": ("IMAGE",), } } RETURN_TYPES = ("IMAGE", "IMAGE") RETURN_NAMES = ("stitched_image", "image_list") OUTPUT_IS_LIST = (False, True) FUNCTION = "process_images" CATEGORY = "MyTemplate" def process_images(self, inputcount, image_1, direction, size, **kwargs): # 1. 初始收集 raw_inputs = [image_1] for c in range(1, inputcount): key = f"image_{c + 1}" val = kwargs.get(key, None) if val is not None: raw_inputs.append(val) # 2. 深度展平 (Handle List & Handle Batch) flat_images = [] for item in raw_inputs: # 2.1 先解开 Python List # (有些节点 output_is_list=True 会传过来 list) if isinstance(item, list): sub_items = item else: sub_items = [item] # 2.2 再解开 Tensor Batch for img in sub_items: # ComfyUI 图片通常是 [B, H, W, C] # 容错：如果是 [H, W, C] 先升维 if img.dim() == 3: img = img.unsqueeze(0) # 核心逻辑：检测 Batch 维度 batch_count = img.shape[0] if batch_count == 1: flat_images.append(img) else: # 如果 Batch > 1，拆分为多个 [1, H, W, C] for i in range(batch_count): # 使用切片 [i:i+1] 保持维度为 4D [1, H, W, C] flat_images.append(img[i:i+1]) if not flat_images: raise ValueError("No images provided.") # ============================================================ # [Alpha 通道兼容处理] # ============================================================ max_channels = 3 for img in flat_images: if img.shape[-1] > max_channels: max_channels = img.shape[-1] normalized_images = [] for img in flat_images: current_c = img.shape[-1] if current_c < max_channels: if current_c == 3 and max_channels == 4: alpha = torch.ones((img.shape[0], img.shape[1], img.shape[2], 1), device=img.device, dtype=img.dtype) img = torch.cat((img, alpha), dim=-1) elif current_c == 1: img = img.repeat(1, 1, 1, 3) if max_channels == 4: alpha = torch.ones((img.shape[0], img.shape[1], img.shape[2], 1), device=img.device, dtype=img.dtype) img = torch.cat((img, alpha), dim=-1) normalized_images.append(img) flat_images = normalized_images # ============================================================ # 3. 确定基准尺寸 ref_img = flat_images[0] ref_h, ref_w = ref_img.shape[1], ref_img.shape[2] processed_images = [] # 4. 遍历处理 (缩放) for img in flat_images: img_permuted = img.permute(0, 3, 1, 2) if direction == "right": if img.shape[1] != ref_h: ratio = ref_h / img.shape[1] new_w = int(img.shape[2] * ratio) img_stitch = F.interpolate(img_permuted, size=(ref_h, new_w), mode="bilinear", align_corners=False) else: img_stitch = img_permuted else: if img.shape[2] != ref_w: ratio = ref_w / img.shape[2] new_h = int(img.shape[1] * ratio) img_stitch = F.interpolate(img_permuted, size=(new_h, ref_w), mode="bilinear", align_corners=False) else: img_stitch = img_permuted processed_images.append(img_stitch.permute(0, 2, 3, 1)) # 5. 拼接 cat_dim = 2 if direction == "right" else 1 # 因为我们之前已经把 Batch 拆散了，这里 processed_images 全都是 Batch=1 # 所以直接 cat 肯定不会报错 (只要长宽处理对了) stitched = torch.cat(processed_images, dim=cat_dim) # 6. Global Resize if size > 0: final_h, final_w = stitched.shape[1], stitched.shape[2] stitched_permuted = stitched.permute(0, 3, 1, 2) if direction == "right": scale_factor = size / final_h new_w = int(final_w * scale_factor) stitched_permuted = F.interpolate(stitched_permuted, size=(size, new_w), mode="bilinear", align_corners=False) else: scale_factor = size / final_w new_h = int(final_h * scale_factor) stitched_permuted = F.interpolate(stitched_permuted, size=(new_h, size), mode="bilinear", align_corners=False) stitched = stitched_permuted.permute(0, 2, 3, 1) # 返回: (Stitched Image, List of [1, H, W, C] tensors) return (stitched, flat_images) class MultipleMaskStitch: @classmethod def INPUT_TYPES(s): return { "required": { "inputcount": ("INT", {"default": 3, "min": 2, "max": 50, "step": 1, "tooltip": "Number of inputs."}), "mask_1": ("MASK", {"tooltip": "Base mask for dimensions."}), "direction": (["right", "down"], {"default": "right"}), "size": ("INT", {"default": -1, "min": -1, "max": 8192, "step": 1}), }, "optional": { "mask_2": ("MASK",), "mask_3": ("MASK",), } } RETURN_TYPES = ("MASK", "MASK") RETURN_NAMES = ("stitched_mask", "mask_list") OUTPUT_IS_LIST = (False, True) FUNCTION = "process_masks" CATEGORY = "MyTemplate" def process_masks(self, inputcount, mask_1, direction, size, **kwargs): # 1. 初始收集 raw_inputs = [mask_1] for c in range(1, inputcount): key = f"mask_{c + 1}" val = kwargs.get(key, None) if val is not None: raw_inputs.append(val) # 2. 深度展平 (Handle List & Handle Batch) flat_masks = [] for item in raw_inputs: if isinstance(item, list): sub_items = item else: sub_items = [item] for m in sub_items: # Mask 可能是 [H, W] 或 [B, H, W] if m.dim() == 2: m = m.unsqueeze(0) # 统一变成 [1, H, W] batch_count = m.shape[0] if batch_count == 1: flat_masks.append(m) else: # 拆分 Batch -> 多个 [1, H, W] for i in range(batch_count): flat_masks.append(m[i:i+1]) if not flat_masks: raise ValueError("No masks provided.") # 3. 确定基准 ref_mask = flat_masks[0] ref_h, ref_w = ref_mask.shape[1], ref_mask.shape[2] processed_masks = [] # 4. 遍历处理 for m in flat_masks: # m 已经是 [1, H, W] m_in = m.unsqueeze(1) # [1, 1, H, W] for interpolate if direction == "right": if m.shape[1] != ref_h: ratio = ref_h / m.shape[1] new_w = int(m.shape[2] * ratio) m_stitch = F.interpolate(m_in, size=(ref_h, new_w), mode="nearest") else: m_stitch = m_in else: if m.shape[2] != ref_w: ratio = ref_w / m.shape[2] new_h = int(m.shape[1] * ratio) m_stitch = F.interpolate(m_in, size=(new_h, ref_w), mode="nearest") else: m_stitch = m_in processed_masks.append(m_stitch.squeeze(1)) # Back to [1, H, W] # 5. 拼接 cat_dim = 2 if direction == "right" else 1 stitched = torch.cat(processed_masks, dim=cat_dim) # 6. Global Resize if size > 0: stitched_in = stitched.unsqueeze(1) final_h, final_w = stitched.shape[1], stitched.shape[2] if direction == "right": scale_factor = size / final_h new_w = int(final_w * scale_factor) stitched_in = F.interpolate(stitched_in, size=(size, new_w), mode="nearest") else: scale_factor = size / final_w new_h = int(final_h * scale_factor) stitched_in = F.interpolate(stitched_in, size=(new_h, size), mode="nearest") stitched = stitched_in.squeeze(1) return (stitched, flat_masks) NODE_CLASS_MAPPINGS = { "MultipleTextInputs": MultipleTextInputs, "MultipleMasksToImage": MultipleMasksToImage, "MultipleImageStitch": MultipleImageStitch, "MultipleMaskStitch": MultipleMaskStitch } NODE_DISPLAY_NAME_MAPPINGS = { "MultipleTextInputs": "Multiple Text Inputs", "MultipleMasksToImage": "Multiple Masks to RGB Image", "MultipleImageStitch": "Multiple Image Stitch", "MultipleMaskStitch": "Multiple Mask Stitch" }