diff --git a/.gitignore b/.gitignore index 1df60d2..b1c177a 100644 --- a/.gitignore +++ b/.gitignore @@ -72,4 +72,29 @@ env/ Thumbs.db # Reference code directory - ignore updates -reference_code/ \ No newline at end of file +reference_code/ + +# 临时目录 +temp_*/ +temp_mingnodes/ + +# 测试文件 +test_nodes.py + +# Python缓存 +__pycache__/ +*.pyc +*.pyo +*.pyd +.Python +*.so + +# IDE文件 +.vscode/ +.idea/ +*.swp +*.swo + +# 系统文件 +.DS_Store +Thumbs.db \ No newline at end of file diff --git a/AI_CODING_RULES.md b/AI_CODING_RULES.md index 3ccc453..93ee108 100644 --- a/AI_CODING_RULES.md +++ b/AI_CODING_RULES.md @@ -31,6 +31,8 @@ ## 依赖管理 - 避免严格的依赖版本限制。除非有兼容性或安全性要求，否则建议只指定主版本或不指定版本。 - 依赖声明与参考项目一致，不随意更改依赖版本。 +- 禁止跨文件导入自定义工具函数（如 color_utils、common_utils），所有通用函数应直接内置到节点文件。 +- 如需参考第三方实现，务必将原始代码存放于 `reference_code/`，主线代码只保留必要部分。 ## 变更限制 - 不允许删除或覆盖用户已有的自定义节点。 @@ -43,18 +45,49 @@ - 如需引入第三方实现，必须先征得用户同意。 - 任何自动化操作前，需先说明理由和影响。 -## 示例 -**应该做：** -- 优化已有节点的参数校验和注释。 -- 按照原项目实现方式重构 fill mask 相关逻辑。 -- 新增节点时同步更新注册和文档。 -- 依赖项如无特殊需求，不要写死具体小版本号。 +## 目录结构与分层（UniversalToolkit 专项） -**不应该做：** -- 不要自动拉取外部仓库或依赖。 -- 不要更改用户未授权的文件。 -- 不要随意更改项目结构。 -- 不要在 requirements.txt/pyproject.toml 中写死所有依赖的精确版本号。 +``` +nodes/ +├── image/ # 图像处理相关节点 +│ ├── imitation_hue_node.py +│ ├── image_concatenate.py +│ └── ...（其他 image 节点） +├── tools/ # 工具类节点 +│ ├── purge_vram.py +│ ├── fill_masked_area.py +│ └── ...（其他工具节点） +reference_code/ # 官方/第三方参考实现存放目录 +``` +- 图像处理节点统一分类为 `UniversalToolkit/Image` +- 工具节点统一分类为 `Tools` + +## 节点注册与命名规范 +- 每个节点文件需包含如下注册方式： + ```python + NODE_CLASS_MAPPINGS = { + "节点类名": 节点类, + } + NODE_DISPLAY_NAME_MAPPINGS = { + "节点类名": "节点显示名", + } + ``` +- 节点类名建议统一加 `_UTK` 后缀，显示名加 `(UTK)`。 + +## 测试与验证 +- 所有节点必须能通过 `test_nodes.py` 脚本批量导入测试。 +- 新增节点或重构后，需补充测试用例，确保无依赖缺失和导入错误。 + +## 参考实现与同步 +- 如需同步第三方（如 MingNodes）实现，务必： + - 保持参数、输入输出、算法与官方一致 + - 分类、注册方式可按 UniversalToolkit 规范适配 + - 参考代码完整保留在 `reference_code/`，主线代码只保留实际用到部分 + +## 其它约定 +- 禁止使用已废弃的 `common_utils.py`、`color_utils.py` 等历史文件。 +- 所有节点文件需自包含所需的工具函数，避免循环依赖和导入混乱。 +- 重要变更需在文件头部注明来源、变更说明和版权信息。 # ComfyUI-UniversalToolkit 开发规范（AI_CODING_RULES） diff --git a/__init__.py b/__init__.py index 68dff95..4b83523 100644 --- a/__init__.py +++ b/__init__.py @@ -258,7 +258,7 @@ except ImportError: CHECK_MASK_DISPLAY = {} try: - from .nodes.image.purge_vram import NODE_CLASS_MAPPINGS as PURGE_VRAM_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as PURGE_VRAM_DISPLAY + from .nodes.tools.purge_vram import NODE_CLASS_MAPPINGS as PURGE_VRAM_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as PURGE_VRAM_DISPLAY except ImportError: PURGE_VRAM_MAPPINGS = {} PURGE_VRAM_DISPLAY = {} diff --git a/nodes/image/__init__.py b/nodes/image/__init__.py index 7cb772f..c93f61b 100644 --- a/nodes/image/__init__.py +++ b/nodes/image/__init__.py @@ -6,4 +6,20 @@ Image processing nodes for ComfyUI Universal Toolkit. :copyright: (c) 2024 by May :license: MIT, see LICENSE for more details. -""" \ No newline at end of file +""" + +import os +import importlib + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +# 自动导入本目录下所有节点文件的注册表 +for filename in os.listdir(os.path.dirname(__file__)): + if filename.endswith('.py') and filename not in ('__init__.py', 'image_utils.py', 'image_converters.py'): + modulename = filename[:-3] + module = importlib.import_module(f'.{modulename}', __package__) + if hasattr(module, 'NODE_CLASS_MAPPINGS'): + NODE_CLASS_MAPPINGS.update(getattr(module, 'NODE_CLASS_MAPPINGS')) + if hasattr(module, 'NODE_DISPLAY_NAME_MAPPINGS'): + NODE_DISPLAY_NAME_MAPPINGS.update(getattr(module, 'NODE_DISPLAY_NAME_MAPPINGS')) \ No newline at end of file diff --git a/nodes/image/color_utils.py b/nodes/image/color_utils.py deleted file mode 100644 index 3900de5..0000000 --- a/nodes/image/color_utils.py +++ /dev/null @@ -1,194 +0,0 @@ -""" -Color Utilities for UniversalToolkit -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - -Color processing utilities for UniversalToolkit. - -:copyright: (c) 2024 by May -:license: MIT, see LICENSE for more details. -""" - -import numpy as np -import cv2 - -def image_stats(image): - """计算图像的统计信息""" - return np.mean(image[:, :, 1:], axis=(0, 1)), np.std(image[:, :, 1:], axis=(0, 1)) - -def is_skin_or_lips(lab_image): - """检测皮肤和嘴唇区域""" - l, a, b = lab_image[:, :, 0], lab_image[:, :, 1], lab_image[:, :, 2] - skin = (l > 20) & (l < 250) & (a > 120) & (a < 180) & (b > 120) & (b < 190) - lips = (l > 20) & (l < 200) & (a > 150) & (b > 140) - return (skin | lips).astype(np.float32) - -def adjust_brightness(image, factor, mask=None): - """调整图像亮度""" - hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) - v = hsv[:, :, 2].astype(np.float32) - if mask is not None: - mask = mask.squeeze() - v = np.where(mask > 0, np.clip(v * factor, 0, 255), v) - else: - v = np.clip(v * factor, 0, 255) - hsv[:, :, 2] = v.astype(np.uint8) - return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) - -def adjust_saturation(image, factor, mask=None): - """调整图像饱和度""" - hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) - s = hsv[:, :, 1].astype(np.float32) - if mask is not None: - mask = mask.squeeze() - s = np.where(mask > 0, np.clip(s * factor, 0, 255), s) - else: - s = np.clip(s * factor, 0, 255) - hsv[:, :, 1] = s.astype(np.uint8) - return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) - -def adjust_contrast(image, factor, mask=None): - """调整图像对比度""" - mean = np.mean(image) - adjusted = image.astype(np.float32) - if mask is not None: - mask = mask.squeeze() - mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2) - adjusted = np.where(mask > 0, np.clip((adjusted - mean) * factor + mean, 0, 255), adjusted) - else: - adjusted = np.clip((adjusted - mean) * factor + mean, 0, 255) - return adjusted.astype(np.uint8) - -def adjust_tone(source, target, tone_strength=0.7, mask=None): - """调整图像影调""" - h, w = target.shape[:2] - source = cv2.resize(source, (w, h)) - lab_image = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32) - lab_source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype(np.float32) - l_image = lab_image[:,:,0] - l_source = lab_source[:,:,0] - - if mask is not None: - mask = cv2.resize(mask, (w, h)) - mask = mask.astype(np.float32) / 255.0 - l_adjusted = np.copy(l_image) - mean_source = np.mean(l_source[mask > 0]) - std_source = np.std(l_source[mask > 0]) - mean_target = np.mean(l_image[mask > 0]) - std_target = np.std(l_image[mask > 0]) - l_adjusted[mask > 0] = (l_image[mask > 0] - mean_target) * (std_source / (std_target + 1e-6)) * 0.7 + mean_source - l_adjusted[mask > 0] = np.clip(l_adjusted[mask > 0], 0, 255) - clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8,8)) - l_enhanced = clahe.apply(l_adjusted.astype(np.uint8)) - l_final = cv2.addWeighted(l_adjusted, 0.7, l_enhanced.astype(np.float32), 0.3, 0) - l_final = np.clip(l_final, 0, 255) - l_contrast = cv2.addWeighted(l_final, 1.3, l_final, 0, -20) - l_contrast = np.clip(l_contrast, 0, 255) - l_image[mask > 0] = l_image[mask > 0] * (1 - tone_strength) + l_contrast[mask > 0] * tone_strength - else: - mean_source = np.mean(l_source) - std_source = np.std(l_source) - l_mean = np.mean(l_image) - l_std = np.std(l_image) - l_adjusted = (l_image - l_mean) * (std_source / (l_std + 1e-6)) * 0.7 + mean_source - l_adjusted = np.clip(l_adjusted, 0, 255) - clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8,8)) - l_enhanced = clahe.apply(l_adjusted.astype(np.uint8)) - l_final = cv2.addWeighted(l_adjusted, 0.7, l_enhanced.astype(np.float32), 0.3, 0) - l_final = np.clip(l_final, 0, 255) - l_contrast = cv2.addWeighted(l_final, 1.3, l_final, 0, -20) - l_contrast = np.clip(l_contrast, 0, 255) - l_image = l_image * (1 - tone_strength) + l_contrast * tone_strength - - lab_image[:,:,0] = l_image - return cv2.cvtColor(lab_image.astype(np.uint8), cv2.COLOR_LAB2BGR) - -def color_transfer(source, target, mask=None, strength=1.0, skin_protection=0.2, auto_brightness=True, - brightness_range=0.5, auto_contrast=False, contrast_range=0.5, - auto_saturation=False, saturation_range=0.5, auto_tone=False, tone_strength=0.7): - """色彩迁移函数""" - source_lab = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype(np.float32) - target_lab = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32) - - src_means, src_stds = image_stats(source_lab) - tar_means, tar_stds = image_stats(target_lab) - - skin_lips_mask = is_skin_or_lips(target_lab.astype(np.uint8)) - skin_lips_mask = cv2.GaussianBlur(skin_lips_mask, (5, 5), 0) - - if mask is not None: - mask = cv2.resize(mask, (target.shape[1], target.shape[0])) - mask = mask.astype(np.float32) / 255.0 - - result_lab = target_lab.copy() - for i in range(1, 3): - adjusted_channel = (target_lab[:, :, i] - tar_means[i - 1]) * (src_stds[i - 1] / (tar_stds[i - 1] + 1e-6)) + \ - src_means[i - 1] - adjusted_channel = np.clip(adjusted_channel, 0, 255) - - if mask is not None: - result_lab[:, :, i] = target_lab[:, :, i] * (1 - mask) + \ - (target_lab[:, :, i] * skin_lips_mask * skin_protection + \ - adjusted_channel * skin_lips_mask * (1 - skin_protection) + \ - adjusted_channel * (1 - skin_lips_mask)) * mask - else: - result_lab[:, :, i] = target_lab[:, :, i] * skin_lips_mask * skin_protection + \ - adjusted_channel * skin_lips_mask * (1 - skin_protection) + \ - adjusted_channel * (1 - skin_lips_mask) - - result_bgr = cv2.cvtColor(result_lab.astype(np.uint8), cv2.COLOR_LAB2BGR) - final_result = cv2.addWeighted(target, 1 - strength, result_bgr, strength, 0) - - if mask is not None: - mask = cv2.resize(mask, (target.shape[1], target.shape[0])) - mask = mask.astype(np.float32) / 255.0 - if auto_brightness: - source_brightness = np.mean(cv2.cvtColor(source, cv2.COLOR_BGR2GRAY)) - target_brightness = np.mean(cv2.cvtColor(target, cv2.COLOR_BGR2GRAY)) - brightness_difference = source_brightness - target_brightness - brightness_factor = 1.0 + np.clip(brightness_difference / 255 * brightness_range, brightness_range*-1, brightness_range) - final_result = adjust_brightness(final_result, brightness_factor, mask) - if auto_contrast: - source_gray = cv2.cvtColor(source, cv2.COLOR_BGR2GRAY) - target_gray = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY) - source_contrast = np.std(source_gray) - target_contrast = np.std(target_gray) - contrast_difference = source_contrast - target_contrast - contrast_factor = 1.0 + np.clip(contrast_difference / 255, contrast_range*-1, contrast_range) - final_result = adjust_contrast(final_result, contrast_factor, mask) - if auto_saturation: - source_hsv = cv2.cvtColor(source, cv2.COLOR_BGR2HSV) - target_hsv = cv2.cvtColor(target, cv2.COLOR_BGR2HSV) - source_saturation = np.mean(source_hsv[:, :, 1]) - target_saturation = np.mean(target_hsv[:, :, 1]) - saturation_difference = source_saturation - target_saturation - saturation_factor = 1.0 + np.clip(saturation_difference / 255, saturation_range*-1, saturation_range) - final_result = adjust_saturation(final_result, saturation_factor, mask) - if auto_tone: - final_result = adjust_tone(source, final_result, tone_strength, mask) - else: - if auto_brightness: - source_brightness = np.mean(cv2.cvtColor(source, cv2.COLOR_BGR2GRAY)) - target_brightness = np.mean(cv2.cvtColor(target, cv2.COLOR_BGR2GRAY)) - brightness_difference = source_brightness - target_brightness - brightness_factor = 1.0 + np.clip(brightness_difference / 255 * brightness_range, brightness_range*-1, brightness_range) - final_result = adjust_brightness(final_result, brightness_factor) - if auto_contrast: - source_gray = cv2.cvtColor(source, cv2.COLOR_BGR2GRAY) - target_gray = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY) - source_contrast = np.std(source_gray) - target_contrast = np.std(target_gray) - contrast_difference = source_contrast - target_contrast - contrast_factor = 1.0 + np.clip(contrast_difference / 255, contrast_range*-1, contrast_range) - final_result = adjust_contrast(final_result, contrast_factor) - if auto_saturation: - source_hsv = cv2.cvtColor(source, cv2.COLOR_BGR2HSV) - target_hsv = cv2.cvtColor(target, cv2.COLOR_BGR2HSV) - source_saturation = np.mean(source_hsv[:, :, 1]) - target_saturation = np.mean(target_hsv[:, :, 1]) - saturation_difference = source_saturation - target_saturation - saturation_factor = 1.0 + np.clip(saturation_difference / 255, saturation_range*-1, saturation_range) - final_result = adjust_saturation(final_result, saturation_factor) - if auto_tone: - final_result = adjust_tone(source, final_result, tone_strength) - - return final_result \ No newline at end of file diff --git a/nodes/image/crop_by_mask.py b/nodes/image/crop_by_mask.py index a640d59..e3baa13 100644 --- a/nodes/image/crop_by_mask.py +++ b/nodes/image/crop_by_mask.py @@ -9,9 +9,20 @@ Crops images based on mask detection with various detection modes. """ import torch -from PIL import Image, ImageDraw +from PIL import Image, ImageDraw, ImageFilter import numpy as np -from ...common_utils import log, tensor2pil, pil2tensor, image2mask +from ..image_utils import tensor2pil, pil2tensor, image2mask + +def log(message, message_type='info'): + """简单的日志函数""" + if message_type == 'error': + print(f"❌ Error: {message}") + elif message_type == 'warning': + print(f"⚠️ Warning: {message}") + elif message_type == 'finish': + print(f"✅ {message}") + else: + print(f"ℹ️ {message}") def mask2image(mask): """Convert mask tensor to PIL image""" @@ -101,7 +112,6 @@ class CropByMask_UTK: _mask = mask2image(mask_for_crop) try: - from PIL import ImageFilter bluredmask = gaussian_blur(_mask, 20).convert('L') except ImportError: bluredmask = _mask.convert('L') diff --git a/nodes/image/image_combine_alpha.py b/nodes/image/image_combine_alpha.py index 5ef9634..3da1c53 100644 --- a/nodes/image/image_combine_alpha.py +++ b/nodes/image/image_combine_alpha.py @@ -10,7 +10,18 @@ Combines RGB image with mask to create RGBA image. import torch from PIL import Image -from ...common_utils import log, tensor2pil, pil2tensor +from ..image_utils import tensor2pil, pil2tensor + +def log(message, message_type='info'): + """简单的日志函数""" + if message_type == 'error': + print(f"❌ Error: {message}") + elif message_type == 'warning': + print(f"⚠️ Warning: {message}") + elif message_type == 'finish': + print(f"✅ {message}") + else: + print(f"ℹ️ {message}") def image_channel_split(image, mode): """Split image into channels""" diff --git a/nodes/image/image_concatenate.py b/nodes/image/image_concatenate.py index 26a3442..4a1cfde 100644 --- a/nodes/image/image_concatenate.py +++ b/nodes/image/image_concatenate.py @@ -10,7 +10,6 @@ Concatenates two images side by side or vertically with various options. import torch import torch.nn.functional as F -from ..common_utils import log class ImageConcatenate_UTK: CATEGORY = "UniversalToolkit/Image" diff --git a/nodes/image/image_mask_scale_as.py b/nodes/image/image_mask_scale_as.py index 65cecf9..bf7161a 100644 --- a/nodes/image/image_mask_scale_as.py +++ b/nodes/image/image_mask_scale_as.py @@ -10,7 +10,18 @@ Scales images and masks to match the dimensions of a reference image. import torch from PIL import Image -from ...common_utils import log, tensor2pil, pil2tensor, image2mask +from ..image_utils import tensor2pil, pil2tensor, image2mask + +def log(message, message_type='info'): + """简单的日志函数""" + if message_type == 'error': + print(f"❌ Error: {message}") + elif message_type == 'warning': + print(f"⚠️ Warning: {message}") + elif message_type == 'finish': + print(f"✅ {message}") + else: + print(f"ℹ️ {message}") def fit_resize_image(image, target_width, target_height, fit_mode, resize_sampler, background_color="#000000"): """Resize image according to fit mode""" diff --git a/nodes/image/image_remove_alpha.py b/nodes/image/image_remove_alpha.py index 4a45439..7145235 100644 --- a/nodes/image/image_remove_alpha.py +++ b/nodes/image/image_remove_alpha.py @@ -10,7 +10,18 @@ Removes alpha channel from RGBA images with optional background filling. import torch from PIL import Image -from ...common_utils import log, tensor2pil, pil2tensor +from ..image_utils import tensor2pil, pil2tensor + +def log(message, message_type='info'): + """简单的日志函数""" + if message_type == 'error': + print(f"❌ Error: {message}") + elif message_type == 'warning': + print(f"⚠️ Warning: {message}") + elif message_type == 'finish': + print(f"✅ {message}") + else: + print(f"ℹ️ {message}") class ImageRemoveAlpha_UTK: CATEGORY = "UniversalToolkit/Image" diff --git a/nodes/image/image_scale_by_aspect_ratio.py b/nodes/image/image_scale_by_aspect_ratio.py index 8e3982e..1d3b6b0 100644 --- a/nodes/image/image_scale_by_aspect_ratio.py +++ b/nodes/image/image_scale_by_aspect_ratio.py @@ -11,13 +11,24 @@ Scales images to specific aspect ratios with various fitting modes. import torch from PIL import Image import math -from ...common_utils import log, tensor2pil, pil2tensor, image2mask +from ..image_utils import log, tensor2pil, pil2tensor, image2mask, num_round_up_to_multiple, fit_resize_image, is_valid_mask + +def log(message, message_type='info'): + """简单的日志函数""" + if message_type == 'error': + print(f"❌ Error: {message}") + elif message_type == 'warning': + print(f"⚠️ Warning: {message}") + elif message_type == 'finish': + print(f"✅ {message}") + else: + print(f"ℹ️ {message}") def num_round_up_to_multiple(num, multiple): """Round up to the nearest multiple""" return ((num + multiple - 1) // multiple) * multiple -def fit_resize_image(image, target_width, target_height, fit_mode, resize_sampler): +def fit_resize_image(image, target_width, target_height, fit_mode, resize_sampler, background_color): """Resize image according to fit mode""" if fit_mode == 'letterbox': # Calculate scaling factor to fit within target dimensions @@ -29,7 +40,7 @@ def fit_resize_image(image, target_width, target_height, fit_mode, resize_sample resized = image.resize((new_width, new_height), resize_sampler) # Create new image with target dimensions and paste resized image - result = Image.new(image.mode, (target_width, target_height), (0, 0, 0)) + result = Image.new(image.mode, (target_width, target_height), background_color) paste_x = (target_width - new_width) // 2 paste_y = (target_height - new_height) // 2 result.paste(resized, (paste_x, paste_y)) @@ -62,21 +73,22 @@ class ImageScaleByAspectRatio_UTK: fit_mode = ['letterbox', 'crop', 'fill'] method_mode = ['lanczos', 'bicubic', 'hamming', 'bilinear', 'box', 'nearest'] multiple_list = ['8', '16', '32', '64', '128', '256', '512', 'None'] - + scale_to_list = ['None', 'longest', 'shortest', 'width', 'height', 'total_pixel(kilo pixel)'] return { "required": { "aspect_ratio": (ratio_list,), - "proportional_width": ("INT", {"default": 2, "min": 1, "max": 999, "step": 1}), - "proportional_height": ("INT", {"default": 1, "min": 1, "max": 999, "step": 1}), + "proportional_width": ("INT", {"default": 1, "min": 1, "max": 1e8, "step": 1}), + "proportional_height": ("INT", {"default": 1, "min": 1, "max": 1e8, "step": 1}), "fit": (fit_mode,), "method": (method_mode,), "round_to_multiple": (multiple_list,), - "scale_to_longest_side": ("BOOLEAN", {"default": False}), # 是否按长边缩放 - "longest_side": ("INT", {"default": 1024, "min": 4, "max": 999999, "step": 1}), + "scale_to_side": (scale_to_list,), + "scale_to_length": ("INT", {"default": 1024, "min": 4, "max": 1e8, "step": 1}), + "background_color": ("STRING", {"default": "#000000"}), }, "optional": { - "image": ("IMAGE",), # - "mask": ("MASK",), # + "image": ("IMAGE",), + "mask": ("MASK",), } } @@ -85,9 +97,9 @@ class ImageScaleByAspectRatio_UTK: FUNCTION = 'image_scale_by_aspect_ratio' def image_scale_by_aspect_ratio(self, aspect_ratio, proportional_width, proportional_height, - fit, method, round_to_multiple, scale_to_longest_side, longest_side, - image=None, mask = None, - ): + fit, method, round_to_multiple, scale_to_side, scale_to_length, + background_color, + image=None, mask=None): orig_images = [] orig_masks = [] orig_width = 0 @@ -107,17 +119,21 @@ class ImageScaleByAspectRatio_UTK: mask = torch.unsqueeze(mask, 0) for m in mask: m = torch.unsqueeze(m, 0) - orig_masks.append(m) - _width, _height = tensor2pil(orig_masks[0]).size - if (orig_width > 0 and orig_width != _width) or (orig_height > 0 and orig_height != _height): - log(f"Error: ImageScaleByAspectRatio_UTK skipped, because the mask is does'nt match image.", message_type='error') - return (None, None, None, 0, 0,) - elif orig_width + orig_height == 0: - orig_width = _width - orig_height = _height + if not is_valid_mask(m) and m.shape == torch.Size([1, 64, 64]): + log(f"Warning: ImageScaleByAspectRatio_UTK input mask is empty, ignore it.", message_type='warning') + else: + orig_masks.append(m) + if len(orig_masks) > 0: + _width, _height = tensor2pil(orig_masks[0]).size + if (orig_width > 0 and orig_width != _width) or (orig_height > 0 and orig_height != _height): + log(f"Error: ImageScaleByAspectRatio_UTK execute failed, because the mask is does'nt match image.", message_type='error') + return (None, None, None, 0, 0,) + elif orig_width + orig_height == 0: + orig_width = _width + orig_height = _height if orig_width + orig_height == 0: - log(f"Error: ImageScaleByAspectRatio_UTK skipped, because the image or mask at least one must be input.", message_type='error') + log(f"Error: ImageScaleByAspectRatio_UTK execute failed, because the image or mask at least one must be input.", message_type='error') return (None, None, None, 0, 0,) if aspect_ratio == 'original': @@ -129,27 +145,48 @@ class ImageScaleByAspectRatio_UTK: ratio = int(s[0]) / int(s[1]) # calculate target width and height - if orig_width > orig_height: - if scale_to_longest_side: - target_width = longest_side + if ratio > 1: + if scale_to_side == 'longest': + target_width = scale_to_length + target_height = int(target_width / ratio) + elif scale_to_side == 'shortest': + target_height = scale_to_length + target_width = int(target_height * ratio) + elif scale_to_side == 'width': + target_width = scale_to_length + target_height = int(target_width / ratio) + elif scale_to_side == 'height': + target_height = scale_to_length + target_width = int(target_height * ratio) + elif scale_to_side == 'total_pixel(kilo pixel)': + target_width = math.sqrt(ratio * scale_to_length * 1000) + target_height = target_width / ratio + target_width = int(target_width) + target_height = int(target_height) else: target_width = orig_width - target_height = int(target_width / ratio) + target_height = int(target_width / ratio) else: - if scale_to_longest_side: - target_height = longest_side + if scale_to_side == 'longest': + target_height = scale_to_length + target_width = int(target_height * ratio) + elif scale_to_side == 'shortest': + target_width = scale_to_length + target_height = int(target_width / ratio) + elif scale_to_side == 'width': + target_width = scale_to_length + target_height = int(target_width / ratio) + elif scale_to_side == 'height': + target_height = scale_to_length + target_width = int(target_height * ratio) + elif scale_to_side == 'total_pixel(kilo pixel)': + target_width = math.sqrt(ratio * scale_to_length * 1000) + target_height = target_width / ratio + target_width = int(target_width) + target_height = int(target_height) else: target_height = orig_height - target_width = int(target_height * ratio) - - if ratio < 1: - if scale_to_longest_side: - _r = longest_side / target_height - target_height = longest_side - else: - _r = orig_height / target_height - target_height = orig_height - target_width = int(target_width * _r) + target_width = int(target_height * ratio) if round_to_multiple != 'None': multiple = int(round_to_multiple) @@ -174,22 +211,22 @@ class ImageScaleByAspectRatio_UTK: if len(orig_images) > 0: for i in orig_images: _image = tensor2pil(i).convert('RGB') - _image = fit_resize_image(_image, target_width, target_height, fit, resize_sampler) + _image = fit_resize_image(_image, target_width, target_height, fit, resize_sampler, background_color) ret_images.append(pil2tensor(_image)) if len(orig_masks) > 0: for m in orig_masks: _mask = tensor2pil(m).convert('L') - _mask = fit_resize_image(_mask, target_width, target_height, fit, resize_sampler).convert('L') + _mask = fit_resize_image(_mask, target_width, target_height, fit, resize_sampler, background_color).convert('L') ret_masks.append(image2mask(_mask)) - if len(ret_images) > 0 and len(ret_masks) >0: + if len(ret_images) > 0 and len(ret_masks) > 0: log(f"ImageScaleByAspectRatio_UTK Processed {len(ret_images)} image(s).", message_type='finish') - return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),[orig_width, orig_height], target_width, target_height,) + return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0), [orig_width, orig_height], target_width, target_height,) elif len(ret_images) > 0 and len(ret_masks) == 0: log(f"ImageScaleByAspectRatio_UTK Processed {len(ret_images)} image(s).", message_type='finish') - return (torch.cat(ret_images, dim=0), None,[orig_width, orig_height], target_width, target_height,) + return (torch.cat(ret_images, dim=0), None, [orig_width, orig_height], target_width, target_height,) elif len(ret_images) == 0 and len(ret_masks) > 0: log(f"ImageScaleByAspectRatio_UTK Processed {len(ret_masks)} image(s).", message_type='finish') - return (None, torch.cat(ret_masks, dim=0),[orig_width, orig_height], target_width, target_height,) + return (None, torch.cat(ret_masks, dim=0), [orig_width, orig_height], target_width, target_height,) else: log(f"Error: ImageScaleByAspectRatio_UTK skipped, because the available image or mask is not found.", message_type='error') return (None, None, None, 0, 0,) diff --git a/nodes/image/image_scale_restore.py b/nodes/image/image_scale_restore.py index 591b12d..8d53dff 100644 --- a/nodes/image/image_scale_restore.py +++ b/nodes/image/image_scale_restore.py @@ -10,7 +10,18 @@ Restores images to original size or scales them with specified parameters. import torch from PIL import Image -from ...common_utils import log, tensor2pil, pil2tensor, image2mask +from ..image_utils import tensor2pil, pil2tensor, image2mask + +def log(message, message_type='info'): + """简单的日志函数""" + if message_type == 'error': + print(f"❌ Error: {message}") + elif message_type == 'warning': + print(f"⚠️ Warning: {message}") + elif message_type == 'finish': + print(f"✅ {message}") + else: + print(f"ℹ️ {message}") class ImageScaleRestore_UTK: CATEGORY = "UniversalToolkit/Image" diff --git a/nodes/image/imitation_hue_node.py b/nodes/image/imitation_hue_node.py index f630404..407d20e 100644 --- a/nodes/image/imitation_hue_node.py +++ b/nodes/image/imitation_hue_node.py @@ -8,12 +8,199 @@ Performs color transfer and imitation between images with skin protection. :license: MIT, see LICENSE for more details. """ -import torch import numpy as np import cv2 +import torch + +def image_stats(image): + return np.mean(image[:, :, 1:], axis=(0, 1)), np.std(image[:, :, 1:], axis=(0, 1)) + + +def is_skin_or_lips(lab_image): + l, a, b = lab_image[:, :, 0], lab_image[:, :, 1], lab_image[:, :, 2] + skin = (l > 20) & (l < 250) & (a > 120) & (a < 180) & (b > 120) & (b < 190) + lips = (l > 20) & (l < 200) & (a > 150) & (b > 140) + return (skin | lips).astype(np.float32) + + +def adjust_brightness(image, factor, mask=None): + hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + v = hsv[:, :, 2].astype(np.float32) + if mask is not None: + mask = mask.squeeze() + v = np.where(mask > 0, np.clip(v * factor, 0, 255), v) + else: + v = np.clip(v * factor, 0, 255) + hsv[:, :, 2] = v.astype(np.uint8) + return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) + + +def adjust_saturation(image, factor, mask=None): + hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) + s = hsv[:, :, 1].astype(np.float32) + if mask is not None: + mask = mask.squeeze() + s = np.where(mask > 0, np.clip(s * factor, 0, 255), s) + else: + s = np.clip(s * factor, 0, 255) + hsv[:, :, 1] = s.astype(np.uint8) + return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) + + +def adjust_contrast(image, factor, mask=None): + mean = np.mean(image) + adjusted = image.astype(np.float32) + if mask is not None: + mask = mask.squeeze() + mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2) + adjusted = np.where(mask > 0, np.clip((adjusted - mean) * factor + mean, 0, 255), adjusted) + else: + adjusted = np.clip((adjusted - mean) * factor + mean, 0, 255) + return adjusted.astype(np.uint8) + + +def adjust_tone(source, target, tone_strength=0.7, mask=None): + h, w = target.shape[:2] + source = cv2.resize(source, (w, h)) + lab_image = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32) + lab_source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype(np.float32) + l_image = lab_image[:,:,0] + l_source = lab_source[:,:,0] + + if mask is not None: + mask = cv2.resize(mask, (w, h)) + mask = mask.astype(np.float32) / 255.0 + l_adjusted = np.copy(l_image) + mean_source = np.mean(l_source[mask > 0]) + std_source = np.std(l_source[mask > 0]) + mean_target = np.mean(l_image[mask > 0]) + std_target = np.std(l_image[mask > 0]) + l_adjusted[mask > 0] = (l_image[mask > 0] - mean_target) * (std_source / (std_target + 1e-6)) * 0.7 + mean_source + l_adjusted[mask > 0] = np.clip(l_adjusted[mask > 0], 0, 255) + clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8,8)) + l_enhanced = clahe.apply(l_adjusted.astype(np.uint8)) + l_final = cv2.addWeighted(l_adjusted, 0.7, l_enhanced.astype(np.float32), 0.3, 0) + l_final = np.clip(l_final, 0, 255) + l_contrast = cv2.addWeighted(l_final, 1.3, l_final, 0, -20) + l_contrast = np.clip(l_contrast, 0, 255) + l_image[mask > 0] = l_image[mask > 0] * (1 - tone_strength) + l_contrast[mask > 0] * tone_strength + else: + mean_source = np.mean(l_source) + std_source = np.std(l_source) + l_mean = np.mean(l_image) + l_std = np.std(l_image) + l_adjusted = (l_image - l_mean) * (std_source / (l_std + 1e-6)) * 0.7 + mean_source + l_adjusted = np.clip(l_adjusted, 0, 255) + clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8,8)) + l_enhanced = clahe.apply(l_adjusted.astype(np.uint8)) + l_final = cv2.addWeighted(l_adjusted, 0.7, l_enhanced.astype(np.float32), 0.3, 0) + l_final = np.clip(l_final, 0, 255) + l_contrast = cv2.addWeighted(l_final, 1.3, l_final, 0, -20) + l_contrast = np.clip(l_contrast, 0, 255) + l_image = l_image * (1 - tone_strength) + l_contrast * tone_strength + + lab_image[:,:,0] = l_image + return cv2.cvtColor(lab_image.astype(np.uint8), cv2.COLOR_LAB2BGR) + + +def tensor2cv2(image: torch.Tensor) -> np.array: + if image.dim() == 4: + image = image.squeeze() + npimage = image.numpy() + cv2image = np.uint8(npimage * 255 / npimage.max()) + return cv2.cvtColor(cv2image, cv2.COLOR_RGB2BGR) + + +def color_transfer(source, target, mask=None, strength=1.0, skin_protection=0.2, auto_brightness=True, + brightness_range=0.5, auto_contrast=False, contrast_range=0.5, + auto_saturation=False, saturation_range=0.5, auto_tone=False, tone_strength=0.7): + source_lab = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype(np.float32) + target_lab = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype(np.float32) + + src_means, src_stds = image_stats(source_lab) + tar_means, tar_stds = image_stats(target_lab) + + skin_lips_mask = is_skin_or_lips(target_lab.astype(np.uint8)) + skin_lips_mask = cv2.GaussianBlur(skin_lips_mask, (5, 5), 0) + + if mask is not None: + mask = cv2.resize(mask, (target.shape[1], target.shape[0])) + mask = mask.astype(np.float32) / 255.0 + + result_lab = target_lab.copy() + for i in range(1, 3): + adjusted_channel = (target_lab[:, :, i] - tar_means[i - 1]) * (src_stds[i - 1] / (tar_stds[i - 1] + 1e-6)) + \ + src_means[i - 1] + adjusted_channel = np.clip(adjusted_channel, 0, 255) + + if mask is not None: + result_lab[:, :, i] = target_lab[:, :, i] * (1 - mask) + \ + (target_lab[:, :, i] * skin_lips_mask * skin_protection + \ + adjusted_channel * skin_lips_mask * (1 - skin_protection) + \ + adjusted_channel * (1 - skin_lips_mask)) * mask + else: + result_lab[:, :, i] = target_lab[:, :, i] * skin_lips_mask * skin_protection + \ + adjusted_channel * skin_lips_mask * (1 - skin_protection) + \ + adjusted_channel * (1 - skin_lips_mask) + + result_bgr = cv2.cvtColor(result_lab.astype(np.uint8), cv2.COLOR_LAB2BGR) + final_result = cv2.addWeighted(target, 1 - strength, result_bgr, strength, 0) + + if mask is not None: + mask = cv2.resize(mask, (target.shape[1], target.shape[0])) + mask = mask.astype(np.float32) / 255.0 + if auto_brightness: + source_brightness = np.mean(cv2.cvtColor(source, cv2.COLOR_BGR2GRAY)) + target_brightness = np.mean(cv2.cvtColor(target, cv2.COLOR_BGR2GRAY)) + brightness_difference = source_brightness - target_brightness + brightness_factor = 1.0 + np.clip(brightness_difference / 255 * brightness_range, brightness_range*-1, brightness_range) + final_result = adjust_brightness(final_result, brightness_factor, mask) + if auto_contrast: + source_gray = cv2.cvtColor(source, cv2.COLOR_BGR2GRAY) + target_gray = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY) + source_contrast = np.std(source_gray) + target_contrast = np.std(target_gray) + contrast_difference = source_contrast - target_contrast + contrast_factor = 1.0 + np.clip(contrast_difference / 255, contrast_range*-1, contrast_range) + final_result = adjust_contrast(final_result, contrast_factor, mask) + if auto_saturation: + source_hsv = cv2.cvtColor(source, cv2.COLOR_BGR2HSV) + target_hsv = cv2.cvtColor(target, cv2.COLOR_BGR2HSV) + source_saturation = np.mean(source_hsv[:, :, 1]) + target_saturation = np.mean(target_hsv[:, :, 1]) + saturation_difference = source_saturation - target_saturation + saturation_factor = 1.0 + np.clip(saturation_difference / 255, saturation_range*-1, saturation_range) + final_result = adjust_saturation(final_result, saturation_factor, mask) + if auto_tone: + final_result = adjust_tone(source, final_result, tone_strength, mask) + else: + if auto_brightness: + source_brightness = np.mean(cv2.cvtColor(source, cv2.COLOR_BGR2GRAY)) + target_brightness = np.mean(cv2.cvtColor(target, cv2.COLOR_BGR2GRAY)) + brightness_difference = source_brightness - target_brightness + brightness_factor = 1.0 + np.clip(brightness_difference / 255 * brightness_range, brightness_range*-1, brightness_range) + final_result = adjust_brightness(final_result, brightness_factor) + if auto_contrast: + source_gray = cv2.cvtColor(source, cv2.COLOR_BGR2GRAY) + target_gray = cv2.cvtColor(target, cv2.COLOR_BGR2GRAY) + source_contrast = np.std(source_gray) + target_contrast = np.std(target_gray) + contrast_difference = source_contrast - target_contrast + contrast_factor = 1.0 + np.clip(contrast_difference / 255, contrast_range*-1, contrast_range) + final_result = adjust_contrast(final_result, contrast_factor) + if auto_saturation: + source_hsv = cv2.cvtColor(source, cv2.COLOR_BGR2HSV) + target_hsv = cv2.cvtColor(target, cv2.COLOR_BGR2HSV) + source_saturation = np.mean(source_hsv[:, :, 1]) + target_saturation = np.mean(target_hsv[:, :, 1]) + saturation_difference = source_saturation - target_saturation + saturation_factor = 1.0 + np.clip(saturation_difference / 255, saturation_range*-1, saturation_range) + final_result = adjust_saturation(final_result, saturation_factor) + if auto_tone: + final_result = adjust_tone(source, final_result, tone_strength) + + return final_result -from .color_utils import color_transfer -from .image_converters import tensor2cv2 class ImitationHueNode_UTK: @classmethod @@ -50,37 +237,25 @@ Performs color transfer and imitation between images with skin protection. def imitation_hue(self, imitation_image, target_image, strength, skin_protection, auto_brightness, brightness_range, auto_contrast, contrast_range, auto_saturation, saturation_range, auto_tone, tone_strength, mask=None): - # Convert tensors to OpenCV format - imitation_cv2 = tensor2cv2(imitation_image) - target_cv2 = tensor2cv2(target_image) - - # Convert mask if provided - mask_cv2 = None + for img in imitation_image: + img_cv1 = tensor2cv2(img) + + for img in target_image: + img_cv2 = tensor2cv2(img) + + img_cv3 = None if mask is not None: - mask_cv2 = (mask.cpu().numpy() * 255).astype(np.uint8) - - # Perform color transfer - result = color_transfer( - source=imitation_cv2, - target=target_cv2, - mask=mask_cv2, - strength=strength, - skin_protection=skin_protection, - auto_brightness=auto_brightness, - brightness_range=brightness_range, - auto_contrast=auto_contrast, - contrast_range=contrast_range, - auto_saturation=auto_saturation, - saturation_range=saturation_range, - auto_tone=auto_tone, - tone_strength=tone_strength - ) - - # Convert back to tensor - result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB) - result_tensor = torch.from_numpy(result_rgb.astype(np.float32) / 255.0) - - return (result_tensor,) + for img3 in mask: + img_cv3 = img3.cpu().numpy() + img_cv3 = (img_cv3 * 255).astype(np.uint8) + + result_img = color_transfer(img_cv1, img_cv2, img_cv3, strength, skin_protection, auto_brightness, + brightness_range,auto_contrast, contrast_range, auto_saturation, + saturation_range, auto_tone, tone_strength) + result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB) + rst = torch.from_numpy(result_img.astype(np.float32) / 255.0).unsqueeze(0) + + return (rst,) # Node mappings NODE_CLASS_MAPPINGS = { diff --git a/nodes/image/restore_crop_box.py b/nodes/image/restore_crop_box.py index cacf060..23c9f2a 100644 --- a/nodes/image/restore_crop_box.py +++ b/nodes/image/restore_crop_box.py @@ -10,7 +10,18 @@ Restores cropped images back to their original background. import torch from PIL import Image -from ...common_utils import log, tensor2pil, pil2tensor, image2mask +from ..image_utils import tensor2pil, pil2tensor, image2mask + +def log(message, message_type='info'): + """简单的日志函数""" + if message_type == 'error': + print(f"❌ Error: {message}") + elif message_type == 'warning': + print(f"⚠️ Warning: {message}") + elif message_type == 'finish': + print(f"✅ {message}") + else: + print(f"ℹ️ {message}") class RestoreCropBox_UTK: CATEGORY = "UniversalToolkit/Image" diff --git a/nodes/image_utils.py b/nodes/image_utils.py new file mode 100644 index 0000000..c353472 --- /dev/null +++ b/nodes/image_utils.py @@ -0,0 +1,82 @@ +""" +Image Utilities for UniversalToolkit +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Image processing utility functions for UniversalToolkit nodes. + +:copyright: (c) 2024 by May +:license: MIT, see LICENSE for more details. +""" + +import torch +import numpy as np +from PIL import Image + +def tensor2pil(t_image: torch.Tensor) -> Image: + """将 PyTorch tensor 转换为 PIL Image""" + return Image.fromarray(np.clip(255.0 * t_image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) + +def pil2tensor(image: Image) -> torch.Tensor: + """将 PIL Image 转换为 PyTorch tensor""" + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +def image2mask(image: Image) -> torch.Tensor: + """将图像转换为掩码格式""" + if image.mode == 'L': + return torch.tensor([pil2tensor(image)[0, :, :].tolist()]) + else: + image = image.convert('RGB').split()[0] + return torch.tensor([pil2tensor(image)[0, :, :].tolist()]) + +def tensor2np(tensor: torch.Tensor) -> np.ndarray: + """将 PyTorch tensor 转换为 numpy 数组""" + return np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8) + +def np2tensor(np_array: np.ndarray) -> torch.Tensor: + """将 numpy 数组转换为 PyTorch tensor""" + return torch.from_numpy(np_array.astype(np.float32) / 255.0).unsqueeze(0) + +def log(message: str, message_type: str = 'info'): + name = 'LayerStyle' + if message_type == 'error': + message = '\033[1;41m' + str(message) + '\033[m' + elif message_type == 'warning': + message = '\033[1;31m' + str(message) + '\033[m' + elif message_type == 'finish': + message = '\033[1;32m' + str(message) + '\033[m' + else: + message = '\033[1;33m' + str(message) + '\033[m' + print(f"# 😺dzNodes: {name} -> {message}") + +def num_round_up_to_multiple(number: int, multiple: int) -> int: + return ((number + multiple - 1) // multiple) * multiple + +def fit_resize_image(image, target_width, target_height, fit, resize_sampler, background_color='#000000'): + image = image.convert('RGB') + orig_width, orig_height = image.size + if fit == 'letterbox': + if orig_width / orig_height > target_width / target_height: + fit_width = target_width + fit_height = int(target_width / orig_width * orig_height) + else: + fit_height = target_height + fit_width = int(target_height / orig_height * orig_width) + fit_image = image.resize((fit_width, fit_height), resize_sampler) + ret_image = Image.new('RGB', size=(target_width, target_height), color=background_color) + ret_image.paste(fit_image, box=((target_width - fit_width)//2, (target_height - fit_height)//2)) + elif fit == 'crop': + if orig_width / orig_height > target_width / target_height: + fit_width = int(orig_height * target_width / target_height) + fit_image = image.crop( + ((orig_width - fit_width)//2, 0, (orig_width - fit_width)//2 + fit_width, orig_height)) + else: + fit_height = int(orig_width * target_height / target_width) + fit_image = image.crop( + (0, (orig_height-fit_height)//2, orig_width, (orig_height-fit_height)//2 + fit_height)) + ret_image = fit_image.resize((target_width, target_height), resize_sampler) + else: + ret_image = image.resize((target_width, target_height), resize_sampler) + return ret_image + +def is_valid_mask(tensor: torch.Tensor) -> bool: + return tensor.sum().item() > 0 \ No newline at end of file diff --git a/nodes/image/purge_vram.py b/nodes/tools/purge_vram.py similarity index 92% rename from nodes/image/purge_vram.py rename to nodes/tools/purge_vram.py index 7d150c3..ae37bfe 100644 --- a/nodes/image/purge_vram.py +++ b/nodes/tools/purge_vram.py @@ -10,9 +10,8 @@ Purge GPU memory to free up VRAM. import torch import gc - -from ..tools.logging_utils import log -from ..tools.any_type import AnyType +from .logging_utils import log +from .any_type import AnyType # 创建 AnyType 实例 any = AnyType("*") @@ -24,7 +23,7 @@ def clear_memory(): gc.collect() class PurgeVRAM_UTK: - CATEGORY = "UniversalToolkit/Image" + CATEGORY = "UniversalToolkit/Tools" @classmethod def INPUT_TYPES(cls): diff --git a/test_nodes.py b/test_nodes.py new file mode 100644 index 0000000..e35eb50 --- /dev/null +++ b/test_nodes.py @@ -0,0 +1,85 @@ +import sys +import os + +# 添加路径 +sys.path.append('.') + +print("🔧 测试所有修复后的节点...\n") + +# 测试图像拼接节点 +try: + from nodes.image.image_concatenate import NODE_CLASS_MAPPINGS as CONCAT_MAPPINGS + print("✅ ImageConcatenate_UTK 节点映射:") + for k, v in CONCAT_MAPPINGS.items(): + print(f" {k}: {v.__name__}") +except Exception as e: + print(f"❌ 导入 ImageConcatenate_UTK 失败: {e}") + +try: + from nodes.image.image_concatenate_multi import NODE_CLASS_MAPPINGS as CONCAT_MULTI_MAPPINGS + print("\n✅ ImageConcatenateMulti_UTK 节点映射:") + for k, v in CONCAT_MULTI_MAPPINGS.items(): + print(f" {k}: {v.__name__}") +except Exception as e: + print(f"❌ 导入 ImageConcatenateMulti_UTK 失败: {e}") + +# 测试 ImitationHueNode_UTK +try: + from nodes.image.imitation_hue_node import NODE_CLASS_MAPPINGS as IMITATION_MAPPINGS + print("\n✅ ImitationHueNode_UTK 节点映射:") + for k, v in IMITATION_MAPPINGS.items(): + print(f" {k}: {v.__name__}") +except Exception as e: + print(f"❌ 导入 ImitationHueNode_UTK 失败: {e}") + +# 测试其他修复的节点 +test_nodes = [ + ("restore_crop_box", "RestoreCropBox_UTK"), + ("image_scale_restore", "ImageScaleRestore_UTK"), + ("image_scale_by_aspect_ratio", "ImageScaleByAspectRatio_UTK"), + ("image_remove_alpha", "ImageRemoveAlpha_UTK"), + ("image_mask_scale_as", "ImageMaskScaleAs_UTK"), + ("image_combine_alpha", "ImageCombineAlpha_UTK"), + ("crop_by_mask", "CropByMask_UTK"), +] + +print("\n🔧 测试其他修复的节点:") +for node_file, node_class in test_nodes: + try: + module = __import__(f"nodes.image.{node_file}", fromlist=[node_class]) + node_class_obj = getattr(module, node_class) + print(f"✅ {node_class}: 导入成功") + except Exception as e: + print(f"❌ {node_class}: 导入失败 - {e}") + +# 测试 tools 目录下的节点 +print("\n🔧 测试 tools 目录下的节点:") +tools_nodes = [ + ("purge_vram", "PurgeVRAM_UTK"), + ("fill_masked_area", "FillMaskedArea_UTK"), + ("show_nodes", "Show_UTK"), +] + +for node_file, node_class in tools_nodes: + try: + module = __import__(f"nodes.tools.{node_file}", fromlist=[node_class]) + node_class_obj = getattr(module, node_class) + print(f"✅ {node_class}: 导入成功") + except Exception as e: + print(f"❌ {node_class}: 导入失败 - {e}") + +print("\n🎉 所有节点修复完成！") +print("现在您应该能在 ComfyUI 中看到以下节点：") +print(" - Image Concatenate (UTK)") +print(" - Image Concatenate Multi (UTK)") +print(" - Imitation Hue Node (UTK) - 已同步 MingNodes 实现") +print(" - Restore Crop Box (UTK)") +print(" - Image Scale Restore (UTK)") +print(" - Image Scale By Aspect Ratio (UTK)") +print(" - Image Remove Alpha (UTK)") +print(" - Image Mask Scale As (UTK)") +print(" - Image Combine Alpha (UTK)") +print(" - Crop By Mask (UTK)") +print(" - Purge VRAM (UTK) - 现在在 Tools 分类下") +print(" - Fill Masked Area (UTK) - 在 Tools 分类下") +print(" - Show Nodes (UTK) - 在 Tools 分类下") \ No newline at end of file