diff --git a/AI_CODING_RULES.md b/AI_CODING_RULES.md new file mode 100644 index 0000000..f59d430 --- /dev/null +++ b/AI_CODING_RULES.md @@ -0,0 +1,56 @@ +# AI 协作编程规则（AI Coding Rules） + +## 总则 +- 所有代码必须遵循本项目的风格和结构。 +- 不允许自动拉取、克隆、下载外部项目或依赖，除非用户明确要求。 +- 只允许在本地已有文件和依赖范围内进行开发和修改。 +- 遵循用户指令，优先满足用户需求。 +- 保持代码风格、接口、参数、行为与用户项目一致。 + +## 参考项目代码的强制要求 +- 当用户指定参考某个项目（如GitHub仓库、第三方节点等）时，**必须严格复制该项目的相关代码实现**，保持原有逻辑、参数、shape处理、异常处理等完全一致。 +- **不得自行"优化"或"简化"实现，不得随意更改shape、类型、参数声明、默认值、边界处理等。** +- 如需适配本地环境，仅允许在不影响核心逻辑的前提下做最小必要的兼容性调整，并需明确告知用户。 +- 如遇到与主线ComfyUI不兼容的情况，优先保持参考项目的原始行为，并向用户说明。 + +## 代码风格 +- 遵循 PEP8 代码风格。 +- 类、函数、变量命名需简洁明了，使用英文。 +- 关键逻辑必须有中英文注释。 +- 遵循PEP8及项目原有风格。 +- 变量、函数、类命名与参考项目保持一致。 + +## 节点开发 +- 节点参数需有默认值、类型、范围说明。 +- 输入输出类型必须与 ComfyUI 规范一致。 +- 新增节点需在 `__init__.py` 注册，并补充到文档。 +- 节点参数、UI、输出类型、行为与参考项目完全一致。 +- 不得随意增删参数或更改默认值。 + +## 依赖管理 +- 避免严格的依赖版本限制。除非有兼容性或安全性要求，否则建议只指定主版本或不指定版本。 +- 依赖声明与参考项目一致，不随意更改依赖版本。 + +## 变更限制 +- 不允许删除或覆盖用户已有的自定义节点。 +- 不允许修改依赖包的源码。 +- 不允许自动生成或修改测试数据文件，除非用户要求。 +- 仅在用户明确要求时才可对参考项目代码做自定义扩展或优化。 +- 所有变更需在注释中注明原因。 + +## 其它 +- 如需引入第三方实现，必须先征得用户同意。 +- 任何自动化操作前，需先说明理由和影响。 + +## 示例 +**应该做：** +- 优化已有节点的参数校验和注释。 +- 按照原项目实现方式重构 fill mask 相关逻辑。 +- 新增节点时同步更新注册和文档。 +- 依赖项如无特殊需求，不要写死具体小版本号。 + +**不应该做：** +- 不要自动拉取外部仓库或依赖。 +- 不要更改用户未授权的文件。 +- 不要随意更改项目结构。 +- 不要在 requirements.txt/pyproject.toml 中写死所有依赖的精确版本号。 \ No newline at end of file diff --git a/__init__.py b/__init__.py index 789bfd4..8ce42f4 100644 --- a/__init__.py +++ b/__init__.py @@ -8,13 +8,42 @@ A comprehensive toolkit for ComfyUI that provides various utility nodes for imag :license: MIT, see LICENSE for more details. """ -__version__ = "1.0.3" +__version__ = "1.0.7" __author__ = "CyberDickLang" __email__ = "286878701@qq.com" __url__ = "https://github.com/whmc76" # 更新日志 CHANGELOG = { + "1.0.7": [ + "改进 ImagePadForOutpaintMasked (UTK) 节点:", + "- 新增数据模式（data_mode）参数，支持 'pixel' 和 'percent' 两种模式", + "- 在 'percent' 模式下，允许输入大于100的百分比", + "- 新增背景颜色（background_color）预设选项", + "改进 ImageAndMaskPreview (UTK) 节点:", + "- 将颜色输入从手动输入字符串改为预设颜色下拉菜单", + ], + "1.0.6": [ + "新增 ImageAndMaskPreview_UTK 节点，用于同时预览图像和掩码：", + "- 支持叠加模式（overlay）：在图像上叠加彩色掩码", + "- 支持并排模式（side_by_side）：图像和掩码并排显示", + "- 支持单独显示模式（mask_only/image_only）", + "- 支持多种掩码颜色和透明度调节" + ], + "1.0.5": [ + "新增 ImagePadForOutpaintMasked_UTK 节点，用于外绘时扩展图像尺寸：", + "- 支持上下左右四个方向的独立扩展", + "- 支持多种背景颜色（黑、白、灰、透明）", + "- 支持掩码边缘羽化效果", + "- 自动生成对应的掩码用于后续处理" + ], + "1.0.4": [ + "新增 FillMaskedArea_UTK 节点，支持三种填充模式：", + "- neutral: 使用灰色填充，适合添加全新内容", + "- telea: 基于 Telea 算法的边界填充", + "- navier-stokes: 基于流体动力学的边界填充", + "添加 opencv-python 依赖支持" + ], "1.0.3": [ "新增 MaskAnd_UTK、MaskSub_UTK、MaskAdd_UTK 三个mask像素级运算节点 (UTK)", "修正节点注册与显示名风格统一，完善导入路径" @@ -45,8 +74,8 @@ CHANGELOG = { ] } -from .nodes.image_nodes_utk import EmptyUnitGenerator_UTK, ImageRatioDetector_UTK, DepthMapBlur_UTK, ImageConcatenate_UTK, ImageConcatenateMulti_UTK -from .nodes.tool_nodes_utk import ShowInt_UTK, ShowFloat_UTK, ShowList_UTK, ShowText_UTK, PreviewMask_UTK +from .nodes.image_nodes_utk import EmptyUnitGenerator_UTK, ImageRatioDetector_UTK, DepthMapBlur_UTK, ImageConcatenate_UTK, ImageConcatenateMulti_UTK, ImagePadForOutpaintMasked_UTK, ImageAndMaskPreview_UTK +from .nodes.tool_nodes_utk import ShowInt_UTK, ShowFloat_UTK, ShowList_UTK, ShowText_UTK, PreviewMask_UTK, FillMaskedArea_UTK from .nodes.audio_nodes_utk import LoadAudioPlusFromPath_UTK, AudioCropProcessUTK from .nodes.mask_nodes_utk import MaskAnd_UTK, MaskSub_UTK, MaskAdd_UTK @@ -58,9 +87,12 @@ NODE_CLASS_MAPPINGS = { "ShowList_UTK": ShowList_UTK, "ShowText_UTK": ShowText_UTK, "PreviewMask_UTK": PreviewMask_UTK, + "FillMaskedArea_UTK": FillMaskedArea_UTK, + "ImageAndMaskPreview_UTK": ImageAndMaskPreview_UTK, "DepthMapBlur_UTK": DepthMapBlur_UTK, "ImageConcatenate_UTK": ImageConcatenate_UTK, "ImageConcatenateMulti_UTK": ImageConcatenateMulti_UTK, + "ImagePadForOutpaintMasked_UTK": ImagePadForOutpaintMasked_UTK, "LoadAudioPlusFromPath_UTK": LoadAudioPlusFromPath_UTK, "AudioCropProcessUTK": AudioCropProcessUTK, "MaskAnd_UTK": MaskAnd_UTK, @@ -76,9 +108,12 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ShowList_UTK": "Show List (UTK)", "ShowText_UTK": "Show Text (UTK)", "PreviewMask_UTK": "Preview Mask (UTK)", + "FillMaskedArea_UTK": "Fill Masked Area (UTK)", + "ImageAndMaskPreview_UTK": "Image And Mask Preview (UTK)", "DepthMapBlur_UTK": "Depth Map Blur", "ImageConcatenate_UTK": "Image Concatenate", "ImageConcatenateMulti_UTK": "Image Concatenate Multi", + "ImagePadForOutpaintMasked_UTK": "Image Pad For Outpaint Masked (UTK)", "LoadAudioPlusFromPath_UTK": "Load Audio Plus From Path (UTK)", "AudioCropProcessUTK": "Audio Crop Process (UTK)", "MaskAnd_UTK": "Mask And (UTK)", @@ -93,6 +128,9 @@ NODE_CATEGORIES = { "DepthMapBlur_UTK", "ImageConcatenate_UTK", "ImageConcatenateMulti_UTK", + "ImagePadForOutpaintMasked_UTK", + "FillMaskedArea_UTK", + "ImageAndMaskPreview_UTK", "LoadAudioPlusFromPath_UTK", "AudioCropProcessUTK", "MaskAnd_UTK", diff --git a/nodes/image_nodes_utk.py b/nodes/image_nodes_utk.py index df51fa9..cdc4e54 100644 --- a/nodes/image_nodes_utk.py +++ b/nodes/image_nodes_utk.py @@ -1,9 +1,41 @@ import torch +import torch.nn.functional as F import numpy as np import re import math +import random +import os +import json from comfy.utils import ProgressBar, common_upscale from PIL import Image +from PIL.PngImagePlugin import PngInfo + +# Import ComfyUI modules with fallbacks +MAX_RESOLUTION = 8192 +SaveImage = None +ImageCompositeMasked = None +args = None +folder_paths = None + +try: + from nodes import MAX_RESOLUTION, SaveImage +except ImportError: + pass + +try: + from comfy_extras.nodes_mask import ImageCompositeMasked +except ImportError: + pass + +try: + from comfy.cli_args import args +except ImportError: + pass + +try: + import folder_paths +except ImportError: + pass class EmptyUnitGenerator_UTK: CATEGORY = "UniversalToolkit" @@ -578,4 +610,148 @@ class ImageConcatenateMulti_UTK: if direction != 'up': y_offset += h - return (output,) \ No newline at end of file + return (output,) + +class ImagePadForOutpaintMasked_UTK: + CATEGORY = "UniversalToolkit" + + @classmethod + def INPUT_TYPES(cls): + color_options = ["gray", "white", "black", "red", "green", "blue", "yellow", "cyan", "magenta"] + return { + "required": { + "image": ("IMAGE",), + "data_mode": (["pixel", "percent"], {"default": "pixel"}), + "left": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), + "top": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), + "right": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), + "bottom": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}), + "feathering": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), + "background_color": (color_options, {"default": "gray"}), + }, + "optional": { + "mask": ("MASK",), + } + } + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "expand_image" + + def expand_image(self, image, data_mode, left, top, right, bottom, feathering, background_color, mask=None): + B, H, W, C = image.size() + # 处理 pad 参数 + if data_mode == "percent": + left = int(W * left / 100) + right = int(W * right / 100) + top = int(H * top / 100) + bottom = int(H * bottom / 100) + # 背景色映射 + color_map = { + "gray": [0.5, 0.5, 0.5], + "white": [1.0, 1.0, 1.0], + "black": [0.0, 0.0, 0.0], + "red": [1.0, 0.0, 0.0], + "green": [0.0, 1.0, 0.0], + "blue": [0.0, 0.0, 1.0], + "yellow": [1.0, 1.0, 0.0], + "cyan": [0.0, 1.0, 1.0], + "magenta": [1.0, 0.0, 1.0], + } + bg_rgb = color_map.get(background_color, [0.5, 0.5, 0.5]) + # 新图像 + new_image = torch.ones((B, H + top + bottom, W + left + right, C), dtype=torch.float32) + for i in range(C): + new_image[:, :, :, i] = bg_rgb[i] + new_image[:, top:top + H, left:left + W, :] = image + # 掩码逻辑与原实现一致 + if mask is not None: + if torch.allclose(mask, torch.zeros_like(mask)): + print("Warning: The incoming mask is fully black. Handling it as None.") + mask = None + if mask is None: + new_mask = torch.ones((B, H + top + bottom, W + left + right), dtype=torch.float32) + t = torch.zeros((B, H, W), dtype=torch.float32) + else: + mask = F.pad(mask, (left, right, top, bottom), mode='constant', value=0) + mask = 1 - mask + t = torch.zeros_like(mask) + if feathering > 0 and feathering * 2 < H and feathering * 2 < W: + for i in range(H): + for j in range(W): + dt = i if top != 0 else H + db = H - i if bottom != 0 else H + dl = j if left != 0 else W + dr = W - j if right != 0 else W + d = min(dt, db, dl, dr) + if d >= feathering: + continue + v = (feathering - d) / feathering + if mask is None: + t[:, i, j] = v * v + else: + t[:, top + i, left + j] = v * v + if mask is None: + new_mask[:, top:top + H, left:left + W] = t + return (new_image, new_mask,) + else: + return (new_image, mask,) + +class ImageAndMaskPreview_UTK(SaveImage): + def __init__(self): + self.output_dir = folder_paths.get_temp_directory() + self.type = "temp" + self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) + self.compress_level = 4 + + @classmethod + def INPUT_TYPES(s): + colors = ["red", "green", "blue", "yellow", "cyan", "magenta", "white", "black"] + return { + "required": { + "mask_opacity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "mask_color": (colors, {"default": "red"}), + "pass_through": ("BOOLEAN", {"default": False}), + }, + "optional": { + "image": ("IMAGE",), + "mask": ("MASK",), + }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("composite",) + FUNCTION = "execute" + CATEGORY = "UniversalToolkit" + DESCRIPTION = """ +Preview an image or a mask, when both inputs are used +composites the mask on top of the image. +with pass_through on the preview is disabled and the +composite is returned from the composite slot instead, +this allows for the preview to be passed for video combine +nodes for example. +""" + + def execute(self, mask_opacity, mask_color, pass_through, filename_prefix="ComfyUI", image=None, mask=None, prompt=None, extra_pnginfo=None): + if mask is not None and image is None: + preview = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + elif mask is None and image is not None: + preview = image + elif mask is not None and image is not None: + mask_adjusted = mask * mask_opacity + mask_image = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3).clone() + + color_map = { + "red": [255, 0, 0], "green": [0, 255, 0], "blue": [0, 0, 255], + "yellow": [255, 255, 0], "cyan": [0, 255, 255], "magenta": [255, 0, 255], + "white": [255, 255, 255], "black": [0, 0, 0] + } + color_list = color_map.get(mask_color, [255, 0, 0]) + + mask_image[:, :, :, 0] = color_list[0] / 255 # Red channel + mask_image[:, :, :, 1] = color_list[1] / 255 # Green channel + mask_image[:, :, :, 2] = color_list[2] / 255 # Blue channel + + preview, = ImageCompositeMasked.composite(self, image, mask_image, 0, 0, True, mask_adjusted) + if pass_through: + return (preview, ) + return(self.save_images(preview, filename_prefix, prompt, extra_pnginfo)) \ No newline at end of file diff --git a/nodes/tool_nodes_utk.py b/nodes/tool_nodes_utk.py index bd1bb28..f843b40 100644 --- a/nodes/tool_nodes_utk.py +++ b/nodes/tool_nodes_utk.py @@ -1,4 +1,7 @@ import torch +import numpy as np +import cv2 +from scipy.ndimage import binary_erosion, gaussian_filter class Show_UTK: CATEGORY = "UniversalToolkit" @@ -96,4 +99,107 @@ class PreviewMask_UTK: def show(self, mask=None): if mask is None: mask = torch.zeros([1, 1, 64, 64], dtype=torch.float32) - return (mask,) \ No newline at end of file + return (mask,) + +# mask二值化，阈值0.5 +def mask_floor(mask): + return (mask > 0.5).astype(np.float32) + +# 腐蚀操作，kernel为feathering +def mask_erosion(mask, feathering): + if feathering > 0: + structure = np.ones((feathering, feathering), dtype=np.uint8) + return binary_erosion(mask, structure=structure).astype(np.float32) + return mask + +# 高斯模糊，sigma=feathering/3 +def mask_blur(mask, feathering): + if feathering > 0: + sigma = feathering / 3.0 + return gaussian_filter(mask, sigma=sigma) + return mask + +class FillMaskedArea_UTK: + CATEGORY = "UniversalToolkit" + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "mask": ("MASK",), + "fill_mode": (["neutral", "telea", "navier-stokes"], {"default": "neutral"}), + "feathering": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1, "label": "Feathering (羽化/边缘过渡)"}), + } + } + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "fill_masked" + IS_PREVIEW = True + + def fill_masked(self, image, mask, fill_mode, feathering): + # 支持batch + if isinstance(image, torch.Tensor): + if image.dim() == 4: + batch_size = image.shape[0] + results = [] + for i in range(batch_size): + img_np = image[i].cpu().numpy() + mask_np = mask[i].cpu().numpy() if mask.dim() == 4 else mask.cpu().numpy() + result = self._fill_single_image(img_np, mask_np, fill_mode, feathering) + results.append(result) + return (torch.from_numpy(np.stack(results)).float(),) + else: + img_np = image.cpu().numpy() + mask_np = mask.cpu().numpy() + result = self._fill_single_image(img_np, mask_np, fill_mode, feathering) + return (torch.from_numpy(result).unsqueeze(0).float(),) + else: + result = self._fill_single_image(image, mask, fill_mode, feathering) + return (torch.from_numpy(result).unsqueeze(0).float(),) + + def _fill_single_image(self, image, mask, fill_mode, feathering): + # [C,H,W] -> [H,W,C] + if image.shape[0] <= 4: + image = np.transpose(image, (1, 2, 0)) + if mask.ndim == 3 and mask.shape[0] == 1: + mask = mask[0] + elif mask.ndim == 3 and mask.shape[2] == 1: + mask = mask[:, :, 0] + # 归一化到0-1 + if mask.max() > 1.0: + mask = mask / 255.0 + # 1. mask二值化 + mask_bin = mask_floor(mask) + # 2. 腐蚀+高斯羽化 + if feathering > 0: + mask_eroded = mask_erosion(mask_bin, feathering) + mask_feathered = mask_blur(mask_eroded, feathering) + else: + mask_feathered = mask_bin + alpha = np.clip(mask_feathered, 0, 1) + # 3. neutral模式 + if fill_mode == "neutral": + result = image.astype(np.float32) / 255.0 if image.dtype != np.float32 else image.copy() + gray = np.ones_like(result) * 0.5 + out = result * (1 - alpha[..., None]) + gray * alpha[..., None] + return np.clip(out, 0, 1) + # 4. inpaint模式 + if image.dtype != np.uint8: + img_uint8 = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.astype(np.uint8) + else: + img_uint8 = image.copy() + mask_uint8 = (alpha > 0.5).astype(np.uint8) + method = cv2.INPAINT_TELEA if fill_mode == "telea" else cv2.INPAINT_NS + if img_uint8.shape[2] == 3: + img_bgr = cv2.cvtColor(img_uint8, cv2.COLOR_RGB2BGR) + filled = cv2.inpaint(img_bgr, mask_uint8, 3, method) + filled = cv2.cvtColor(filled, cv2.COLOR_BGR2RGB) + else: + filled = cv2.inpaint(img_uint8, mask_uint8, 3, method) + filled = filled.astype(np.float32) / 255.0 + result = image.astype(np.float32) / 255.0 if image.dtype != np.float32 else image.copy() + out = result * (1 - alpha[..., None]) + filled * alpha[..., None] + return np.clip(out, 0, 1) + +class ImageAndMaskPreview_UTK: + pass # This class is now implemented in image_nodes_utk.py \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 6d9bb35..75d4df1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ Pillow numpy torch librosa -torchaudio \ No newline at end of file +torchaudio +opencv-python \ No newline at end of file