488 lines
18 KiB
Python
488 lines
18 KiB
Python
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"
|
||
} |