v1.1.3: 修复 PurgeVRAM_UTK 和 CheckMask_UTK 节点问题，完成模块化重构
- 修复 PurgeVRAM_UTK 节点类型不匹配问题，创建 AnyType 类实现 - 修复 CheckMask_UTK 节点 NoneType 错误，添加空值检查 - 完成模块化重构：将 common_utils.py 功能拆分到不同目录 - 创建 image_converters.py、color_utils.py、logging_utils.py、any_type.py - 删除 common_utils.py，避免依赖冲突 - 更新所有相关文件的导入语句 - 提高代码可维护性和模块化程度
This commit is contained in:
Cyber Dick Lang
+42
-1
@@ -54,4 +54,45 @@
|
||||
- 不要自动拉取外部仓库或依赖。
|
||||
- 不要更改用户未授权的文件。
|
||||
- 不要随意更改项目结构。
|
||||
- 不要在 requirements.txt/pyproject.toml 中写死所有依赖的精确版本号。
|
||||
- 不要在 requirements.txt/pyproject.toml 中写死所有依赖的精确版本号。
|
||||
|
||||
# ComfyUI-UniversalToolkit 开发规范（AI_CODING_RULES）
|
||||
|
||||
## 2024-07-13 重要改进与规则
|
||||
|
||||
### 1. 节点注册与导入（ComfyUI v3官方规范）
|
||||
- 每个节点文件只导出自己的 `NODE_CLASS_MAPPINGS` 和 `NODE_DISPLAY_NAME_MAPPINGS`。
|
||||
- `__init__.py` 必须静态导入所有节点注册字典，禁止动态try/except导入和动态合并。
|
||||
- 只导出 `NODE_CLASS_MAPPINGS`、`NODE_DISPLAY_NAME_MAPPINGS`，对齐官方插件加载机制。
|
||||
- 节点注册顺序清晰，所有节点都必须被静态合并进主注册字典。
|
||||
|
||||
### 2. 节点分组与命名
|
||||
- 每个节点类必须有唯一且规范的 `CATEGORY` 属性，分组如 `UniversalToolkit/Image`、`UniversalToolkit/Mask`、`UniversalToolkit/Audio`、`UniversalToolkit/Tools`。
|
||||
- 节点类名、注册名、显示名必须唯一，全部带 `_UTK` 后缀，禁止与原生节点或其它插件重名。
|
||||
- 节点显示名统一加 `(UTK)` 后缀，保证界面分组风格一致。
|
||||
|
||||
### 3. pyproject.toml 规范
|
||||
- 必须包含 `[project]` 和 `[tool.comfy]` 两大段，字段严格对齐官方文档：
|
||||
- `name`、`description`、`version`、`license`、`dependencies`、`Repository`
|
||||
- `[tool.comfy]` 下 `PublisherId`、`DisplayName`、`Icon` 必填
|
||||
- `PublisherId` 必须与Registry注册一致，`DisplayName`为插件在ComfyUI-Manager/Registry中的显示名
|
||||
|
||||
### 4. 目录结构与分层
|
||||
- 按功能分为 `nodes/image`、`nodes/mask`、`nodes/audio`、`nodes/tools` 四大目录，每个节点独立py文件
|
||||
- 禁止使用绝对导入和跨目录utils模块，所有依赖应在本插件目录下
|
||||
|
||||
### 5. 节点输入输出与兼容性
|
||||
- 所有节点输入输出shape、类型、参数名、返回名必须严格遵循ComfyUI官方节点开发规范
|
||||
- 禁止随意更改节点实现、参数、shape，所有节点必须与ComfyUI原生节点和其它插件100%兼容
|
||||
- 任何涉及shape、类型、参数的修正，必须优先保证与ComfyUI主程序和主流插件生态兼容
|
||||
|
||||
### 6. 其它重要约定
|
||||
- 禁止在`__init__.py`中做复杂逻辑或动态注册，推荐静态声明所有节点映射
|
||||
- 禁止用旧版的`register_node`、`register_nodes`等动态注册API
|
||||
- 所有节点分组、命名、注册、导入、依赖、pyproject.toml等必须随时对齐ComfyUI官方最新规范
|
||||
|
||||
---
|
||||
|
||||
**本规范为ComfyUI-UniversalToolkit插件开发的最高准则，所有贡献者和维护者必须严格遵守。**
|
||||
|
||||
（如有新规范或官方更新，须第一时间同步修订本文件）
|
||||
+82
-79
@@ -8,13 +8,29 @@ A comprehensive toolkit for ComfyUI that provides various utility nodes for imag
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
__version__ = "1.1.2"
|
||||
__version__ = "1.1.3"
|
||||
__author__ = "CyberDickLang"
|
||||
__email__ = "286878701@qq.com"
|
||||
__url__ = "https://github.com/whmc76"
|
||||
|
||||
# 更新日志
|
||||
CHANGELOG = {
|
||||
"1.1.3": [
|
||||
"修复 PurgeVRAM_UTK 节点类型不匹配问题：",
|
||||
"- 创建 AnyType 类实现，参考 ComfyUI-LayerStyle 项目",
|
||||
"- 解决 'received_type(IMAGE) mismatch input_type(*)' 错误",
|
||||
"- 支持接受任何类型输入并正确返回",
|
||||
"修复 CheckMask_UTK 节点 NoneType 错误：",
|
||||
"- 添加空值检查，防止 mask 为 None 时出错",
|
||||
"- 改进 tensor2pil 转换失败的处理",
|
||||
"- 增强错误处理和日志输出",
|
||||
"完成模块化重构：",
|
||||
"- 将 common_utils.py 功能拆分到不同目录",
|
||||
"- 创建 image_converters.py、color_utils.py、logging_utils.py、any_type.py",
|
||||
"- 删除 common_utils.py，避免依赖冲突",
|
||||
"- 更新所有相关文件的导入语句",
|
||||
"- 提高代码可维护性和模块化程度"
|
||||
],
|
||||
"1.1.2": [
|
||||
"修复 DepthMapBlur_UTK 节点 kernel size 类型和 OpenCV 奇数断言问题，保证所有模糊核为正奇数，完全兼容 ComfyUI 规范。",
|
||||
"修正 EmptyUnitGenerator_UTK 输出 shape，所有节点输入输出严格遵循 ComfyUI 官方规范。",
|
||||
@@ -133,29 +149,22 @@ CHANGELOG = {
|
||||
# 导入节点模块
|
||||
try:
|
||||
# 工具类节点
|
||||
from .nodes.tools.show_nodes import NODE_CLASS_MAPPINGS as SHOW_NODES_MAPPINGS
|
||||
from .nodes.tools.show_nodes import NODE_DISPLAY_NAME_MAPPINGS as SHOW_NODES_DISPLAY_MAPPINGS
|
||||
from .nodes.tools.show_nodes import NODE_CLASS_MAPPINGS as SHOW_NODES_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as SHOW_NODES_DISPLAY_MAPPINGS
|
||||
|
||||
# 音频节点
|
||||
from .nodes.audio.audio_crop_process import NODE_CLASS_MAPPINGS as AUDIO_CROP_MAPPINGS
|
||||
from .nodes.audio.audio_crop_process import NODE_DISPLAY_NAME_MAPPINGS as AUDIO_CROP_DISPLAY_MAPPINGS
|
||||
from .nodes.audio.audio_crop_process import NODE_CLASS_MAPPINGS as AUDIO_CROP_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as AUDIO_CROP_DISPLAY_MAPPINGS
|
||||
|
||||
# 掩码节点
|
||||
from .nodes.mask.mask_operations import NODE_CLASS_MAPPINGS as MASK_OPERATIONS_MAPPINGS
|
||||
from .nodes.mask.mask_operations import NODE_DISPLAY_NAME_MAPPINGS as MASK_OPERATIONS_DISPLAY_MAPPINGS
|
||||
from .nodes.mask.mask_operations import NODE_CLASS_MAPPINGS as MASK_OPERATIONS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as MASK_OPERATIONS_DISPLAY_MAPPINGS
|
||||
|
||||
# 图像节点
|
||||
from .nodes.image.image_concatenate_multi import NODE_CLASS_MAPPINGS as CONCATENATE_MULTI_MAPPINGS
|
||||
from .nodes.image.image_concatenate_multi import NODE_DISPLAY_NAME_MAPPINGS as CONCATENATE_MULTI_DISPLAY_MAPPINGS
|
||||
from .nodes.image.image_concatenate_multi import NODE_CLASS_MAPPINGS as CONCATENATE_MULTI_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as CONCATENATE_MULTI_DISPLAY_MAPPINGS
|
||||
|
||||
from .nodes.image.image_pad_for_outpaint_masked import NODE_CLASS_MAPPINGS as PAD_OUTPAINT_MAPPINGS
|
||||
from .nodes.image.image_pad_for_outpaint_masked import NODE_DISPLAY_NAME_MAPPINGS as PAD_OUTPAINT_DISPLAY_MAPPINGS
|
||||
from .nodes.image.image_pad_for_outpaint_masked import NODE_CLASS_MAPPINGS as PAD_OUTPAINT_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as PAD_OUTPAINT_DISPLAY_MAPPINGS
|
||||
|
||||
from .nodes.image.image_and_mask_preview import NODE_CLASS_MAPPINGS as AND_MASK_PREVIEW_MAPPINGS
|
||||
from .nodes.image.image_and_mask_preview import NODE_DISPLAY_NAME_MAPPINGS as AND_MASK_PREVIEW_DISPLAY_MAPPINGS
|
||||
from .nodes.image.image_and_mask_preview import NODE_CLASS_MAPPINGS as AND_MASK_PREVIEW_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as AND_MASK_PREVIEW_DISPLAY_MAPPINGS
|
||||
|
||||
from .nodes.image.imitation_hue_node import NODE_CLASS_MAPPINGS as IMITATION_HUE_MAPPINGS
|
||||
from .nodes.image.imitation_hue_node import NODE_DISPLAY_NAME_MAPPINGS as IMITATION_HUE_DISPLAY_MAPPINGS
|
||||
from .nodes.image.imitation_hue_node import NODE_CLASS_MAPPINGS as IMITATION_HUE_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as IMITATION_HUE_DISPLAY_MAPPINGS
|
||||
|
||||
except ImportError as e:
|
||||
print(f"导入错误: {e}")
|
||||
@@ -177,116 +186,101 @@ except ImportError as e:
|
||||
|
||||
# 尝试导入其他可能有依赖的节点
|
||||
try:
|
||||
from .nodes.tools.fill_masked_area import NODE_CLASS_MAPPINGS as FILL_MASKED_MAPPINGS
|
||||
from .nodes.tools.fill_masked_area import NODE_DISPLAY_NAME_MAPPINGS as FILL_MASKED_DISPLAY_MAPPINGS
|
||||
from .nodes.tools.fill_masked_area import NODE_CLASS_MAPPINGS as FILL_MASKED_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as FILL_MASKED_DISPLAY_MAPPINGS
|
||||
except ImportError:
|
||||
FILL_MASKED_MAPPINGS = {}
|
||||
FILL_MASKED_DISPLAY_MAPPINGS = {}
|
||||
|
||||
try:
|
||||
from .nodes.audio.load_audio import NODE_CLASS_MAPPINGS as LOAD_AUDIO_MAPPINGS
|
||||
from .nodes.audio.load_audio import NODE_DISPLAY_NAME_MAPPINGS as LOAD_AUDIO_DISPLAY_MAPPINGS
|
||||
from .nodes.audio.load_audio import NODE_CLASS_MAPPINGS as LOAD_AUDIO_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as LOAD_AUDIO_DISPLAY_MAPPINGS
|
||||
except ImportError:
|
||||
LOAD_AUDIO_MAPPINGS = {}
|
||||
LOAD_AUDIO_DISPLAY_MAPPINGS = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.empty_unit_generator import NODE_CLASS_MAPPINGS as IMAGE_GENERATOR_MAPPINGS
|
||||
from .nodes.image.empty_unit_generator import NODE_DISPLAY_NAME_MAPPINGS as IMAGE_GENERATOR_DISPLAY_MAPPINGS
|
||||
from .nodes.image.empty_unit_generator import NODE_CLASS_MAPPINGS as EMPTY_UNIT_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as EMPTY_UNIT_DISPLAY
|
||||
except ImportError:
|
||||
IMAGE_GENERATOR_MAPPINGS = {}
|
||||
IMAGE_GENERATOR_DISPLAY_MAPPINGS = {}
|
||||
EMPTY_UNIT_MAPPINGS = {}
|
||||
EMPTY_UNIT_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.image_ratio_detector import NODE_CLASS_MAPPINGS as IMAGE_DETECTOR_MAPPINGS
|
||||
from .nodes.image.image_ratio_detector import NODE_DISPLAY_NAME_MAPPINGS as IMAGE_DETECTOR_DISPLAY_MAPPINGS
|
||||
from .nodes.image.image_ratio_detector import NODE_CLASS_MAPPINGS as RATIO_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as RATIO_DISPLAY
|
||||
except ImportError:
|
||||
IMAGE_DETECTOR_MAPPINGS = {}
|
||||
IMAGE_DETECTOR_DISPLAY_MAPPINGS = {}
|
||||
RATIO_MAPPINGS = {}
|
||||
RATIO_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.depth_map_blur import NODE_CLASS_MAPPINGS as DEPTH_BLUR_MAPPINGS
|
||||
from .nodes.image.depth_map_blur import NODE_DISPLAY_NAME_MAPPINGS as DEPTH_BLUR_DISPLAY_MAPPINGS
|
||||
from .nodes.image.depth_map_blur import NODE_CLASS_MAPPINGS as DEPTH_BLUR_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as DEPTH_BLUR_DISPLAY
|
||||
except ImportError:
|
||||
DEPTH_BLUR_MAPPINGS = {}
|
||||
DEPTH_BLUR_DISPLAY_MAPPINGS = {}
|
||||
DEPTH_BLUR_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.image_concatenate import NODE_CLASS_MAPPINGS as CONCATENATE_MAPPINGS
|
||||
from .nodes.image.image_concatenate import NODE_DISPLAY_NAME_MAPPINGS as CONCATENATE_DISPLAY_MAPPINGS
|
||||
from .nodes.image.image_concatenate import NODE_CLASS_MAPPINGS as CONCAT_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as CONCAT_DISPLAY
|
||||
except ImportError:
|
||||
CONCATENATE_MAPPINGS = {}
|
||||
CONCATENATE_DISPLAY_MAPPINGS = {}
|
||||
CONCAT_MAPPINGS = {}
|
||||
CONCAT_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.image_scale_by_aspect_ratio import NODE_CLASS_MAPPINGS as SCALE_ASPECT_MAPPINGS
|
||||
from .nodes.image.image_scale_by_aspect_ratio import NODE_DISPLAY_NAME_MAPPINGS as SCALE_ASPECT_DISPLAY_MAPPINGS
|
||||
from .nodes.image.image_scale_by_aspect_ratio import NODE_CLASS_MAPPINGS as SCALE_ASPECT_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as SCALE_ASPECT_DISPLAY
|
||||
except ImportError:
|
||||
SCALE_ASPECT_MAPPINGS = {}
|
||||
SCALE_ASPECT_DISPLAY_MAPPINGS = {}
|
||||
SCALE_ASPECT_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.image_mask_scale_as import NODE_CLASS_MAPPINGS as MASK_SCALE_MAPPINGS
|
||||
from .nodes.image.image_mask_scale_as import NODE_DISPLAY_NAME_MAPPINGS as MASK_SCALE_DISPLAY_MAPPINGS
|
||||
from .nodes.image.image_mask_scale_as import NODE_CLASS_MAPPINGS as MASK_SCALE_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as MASK_SCALE_DISPLAY
|
||||
except ImportError:
|
||||
MASK_SCALE_MAPPINGS = {}
|
||||
MASK_SCALE_DISPLAY_MAPPINGS = {}
|
||||
MASK_SCALE_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.image_scale_restore import NODE_CLASS_MAPPINGS as SCALE_RESTORE_MAPPINGS
|
||||
from .nodes.image.image_scale_restore import NODE_DISPLAY_NAME_MAPPINGS as SCALE_RESTORE_DISPLAY_MAPPINGS
|
||||
from .nodes.image.image_scale_restore import NODE_CLASS_MAPPINGS as SCALE_RESTORE_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as SCALE_RESTORE_DISPLAY
|
||||
except ImportError:
|
||||
SCALE_RESTORE_MAPPINGS = {}
|
||||
SCALE_RESTORE_DISPLAY_MAPPINGS = {}
|
||||
SCALE_RESTORE_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.image_remove_alpha import NODE_CLASS_MAPPINGS as REMOVE_ALPHA_MAPPINGS
|
||||
from .nodes.image.image_remove_alpha import NODE_DISPLAY_NAME_MAPPINGS as REMOVE_ALPHA_DISPLAY_MAPPINGS
|
||||
from .nodes.image.image_remove_alpha import NODE_CLASS_MAPPINGS as REMOVE_ALPHA_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as REMOVE_ALPHA_DISPLAY
|
||||
except ImportError:
|
||||
REMOVE_ALPHA_MAPPINGS = {}
|
||||
REMOVE_ALPHA_DISPLAY_MAPPINGS = {}
|
||||
REMOVE_ALPHA_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.image_combine_alpha import NODE_CLASS_MAPPINGS as COMBINE_ALPHA_MAPPINGS
|
||||
from .nodes.image.image_combine_alpha import NODE_DISPLAY_NAME_MAPPINGS as COMBINE_ALPHA_DISPLAY_MAPPINGS
|
||||
from .nodes.image.image_combine_alpha import NODE_CLASS_MAPPINGS as COMBINE_ALPHA_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as COMBINE_ALPHA_DISPLAY
|
||||
except ImportError:
|
||||
COMBINE_ALPHA_MAPPINGS = {}
|
||||
COMBINE_ALPHA_DISPLAY_MAPPINGS = {}
|
||||
COMBINE_ALPHA_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.check_mask import NODE_CLASS_MAPPINGS as CHECK_MASK_MAPPINGS
|
||||
from .nodes.image.check_mask import NODE_DISPLAY_NAME_MAPPINGS as CHECK_MASK_DISPLAY_MAPPINGS
|
||||
from .nodes.image.check_mask import NODE_CLASS_MAPPINGS as CHECK_MASK_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as CHECK_MASK_DISPLAY
|
||||
except ImportError:
|
||||
CHECK_MASK_MAPPINGS = {}
|
||||
CHECK_MASK_DISPLAY_MAPPINGS = {}
|
||||
CHECK_MASK_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.purge_vram import NODE_CLASS_MAPPINGS as PURGE_VRAM_MAPPINGS
|
||||
from .nodes.image.purge_vram import NODE_DISPLAY_NAME_MAPPINGS as PURGE_VRAM_DISPLAY_MAPPINGS
|
||||
from .nodes.image.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_MAPPINGS = {}
|
||||
PURGE_VRAM_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.crop_by_mask import NODE_CLASS_MAPPINGS as CROP_MASK_MAPPINGS
|
||||
from .nodes.image.crop_by_mask import NODE_DISPLAY_NAME_MAPPINGS as CROP_MASK_DISPLAY_MAPPINGS
|
||||
from .nodes.image.crop_by_mask import NODE_CLASS_MAPPINGS as CROP_MASK_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as CROP_MASK_DISPLAY
|
||||
except ImportError:
|
||||
CROP_MASK_MAPPINGS = {}
|
||||
CROP_MASK_DISPLAY_MAPPINGS = {}
|
||||
CROP_MASK_DISPLAY = {}
|
||||
|
||||
try:
|
||||
from .nodes.image.restore_crop_box import NODE_CLASS_MAPPINGS as RESTORE_CROP_MAPPINGS
|
||||
from .nodes.image.restore_crop_box import NODE_DISPLAY_NAME_MAPPINGS as RESTORE_CROP_DISPLAY_MAPPINGS
|
||||
from .nodes.image.restore_crop_box import NODE_CLASS_MAPPINGS as RESTORE_CROP_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as RESTORE_CROP_DISPLAY
|
||||
except ImportError:
|
||||
RESTORE_CROP_MAPPINGS = {}
|
||||
RESTORE_CROP_DISPLAY_MAPPINGS = {}
|
||||
RESTORE_CROP_DISPLAY = {}
|
||||
|
||||
# 合并所有节点映射
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_CLASS_MAPPINGS.update(IMAGE_GENERATOR_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(IMAGE_DETECTOR_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(EMPTY_UNIT_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(RATIO_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(DEPTH_BLUR_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(CONCATENATE_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(CONCAT_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(CONCATENATE_MULTI_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(PAD_OUTPAINT_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(AND_MASK_PREVIEW_MAPPINGS)
|
||||
@@ -302,29 +296,29 @@ NODE_CLASS_MAPPINGS.update(CROP_MASK_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(RESTORE_CROP_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(SHOW_NODES_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(FILL_MASKED_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(MASK_OPERATIONS_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(LOAD_AUDIO_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(AUDIO_CROP_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(MASK_OPERATIONS_MAPPINGS)
|
||||
|
||||
# 合并显示名称映射
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(IMAGE_GENERATOR_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(IMAGE_DETECTOR_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(DEPTH_BLUR_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(CONCATENATE_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(EMPTY_UNIT_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(RATIO_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(DEPTH_BLUR_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(CONCAT_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(CONCATENATE_MULTI_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(PAD_OUTPAINT_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(AND_MASK_PREVIEW_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(IMITATION_HUE_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(SCALE_ASPECT_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(MASK_SCALE_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(SCALE_RESTORE_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(REMOVE_ALPHA_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(COMBINE_ALPHA_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(CHECK_MASK_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(PURGE_VRAM_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(CROP_MASK_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(RESTORE_CROP_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(SCALE_ASPECT_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(MASK_SCALE_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(SCALE_RESTORE_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(REMOVE_ALPHA_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(COMBINE_ALPHA_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(CHECK_MASK_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(PURGE_VRAM_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(CROP_MASK_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(RESTORE_CROP_DISPLAY)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(SHOW_NODES_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(FILL_MASKED_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(LOAD_AUDIO_DISPLAY_MAPPINGS)
|
||||
@@ -371,4 +365,13 @@ __all__ = [
|
||||
"__email__",
|
||||
"__url__",
|
||||
"CHANGELOG",
|
||||
]
|
||||
]
|
||||
|
||||
# 调试：输出所有注册节点的分组属性
|
||||
if __name__ == "__main__":
|
||||
with open("node_category_report.txt", "w", encoding="utf-8") as f:
|
||||
f.write("=== UniversalToolkit 节点分组属性清单 ===\n")
|
||||
for k, v in NODE_CLASS_MAPPINGS.items():
|
||||
cat = getattr(v, 'CATEGORY', '无CATEGORY')
|
||||
f.write(f"{k}: CATEGORY = {cat}\n")
|
||||
print("节点分组清单已导出到 node_category_report.txt")
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
=== UniversalToolkit 节点分组属性清单 ===
|
||||
@@ -13,7 +13,7 @@ import torch
|
||||
FLOAT_MAX = 99999999999999999.0
|
||||
|
||||
class AudioCropProcessUTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Audio"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
|
||||
@@ -18,7 +18,7 @@ import librosa
|
||||
FLOAT_MAX = 99999999999999999.0
|
||||
|
||||
class LoadAudioPlusFromPath_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Audio"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
|
||||
@@ -2,26 +2,31 @@
|
||||
Check Mask Node
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
Checks if a mask is valid based on white area percentage.
|
||||
Check if a mask is valid and provide information about it.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
from ...common_utils import log, tensor2pil, pil2tensor
|
||||
from PIL import Image
|
||||
import cv2
|
||||
|
||||
from ..tools.logging_utils import log
|
||||
from .image_converters import tensor2pil, pil2tensor
|
||||
|
||||
def mask_white_area(mask, white_point):
|
||||
"""Calculate the percentage of white area in mask"""
|
||||
if mask is None:
|
||||
return 0.0
|
||||
mask_array = np.array(mask)
|
||||
white_pixels = np.sum(mask_array > white_point)
|
||||
total_pixels = mask_array.size
|
||||
return white_pixels / total_pixels if total_pixels > 0 else 0
|
||||
|
||||
class CheckMask_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -41,14 +46,24 @@ class CheckMask_UTK:
|
||||
|
||||
def check_mask(self, mask, white_point, area_percent,):
|
||||
|
||||
if mask is None:
|
||||
log("CheckMask_UTK: mask is None", message_type="warning")
|
||||
return (False,)
|
||||
|
||||
if mask.dim() == 2:
|
||||
mask = torch.unsqueeze(mask, 0)
|
||||
mask = tensor2pil(mask[0])
|
||||
if mask.width * mask.height > 262144:
|
||||
|
||||
mask_pil = tensor2pil(mask[0])
|
||||
if mask_pil is None:
|
||||
log("CheckMask_UTK: Failed to convert mask to PIL", message_type="warning")
|
||||
return (False,)
|
||||
|
||||
if mask_pil.width * mask_pil.height > 262144:
|
||||
target_width = 512
|
||||
target_height = int(target_width * mask.height / mask.width)
|
||||
mask = mask.resize((target_width, target_height), Image.LANCZOS)
|
||||
ret = mask_white_area(mask, white_point) * 100 > area_percent
|
||||
target_height = int(target_width * mask_pil.height / mask_pil.width)
|
||||
mask_pil = mask_pil.resize((target_width, target_height), Image.LANCZOS)
|
||||
|
||||
ret = mask_white_area(mask_pil, white_point) * 100 > area_percent
|
||||
log(f"CheckMask_UTK:{ret}", message_type="finish")
|
||||
return (ret,)
|
||||
|
||||
|
||||
@@ -1,133 +1,29 @@
|
||||
"""
|
||||
ComfyUI Universal Toolkit - Common Utilities
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
Color Utilities for UniversalToolkit
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Common utility functions and constants used across all nodes.
|
||||
Color processing utilities for UniversalToolkit.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import re
|
||||
import math
|
||||
import random
|
||||
import os
|
||||
import json
|
||||
import cv2
|
||||
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
|
||||
|
||||
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 tensor2pil(image):
|
||||
"""将torch张量转换为PIL图像"""
|
||||
if image.dim() == 4:
|
||||
image = image.squeeze(0)
|
||||
if image.dim() == 3:
|
||||
if image.shape[0] == 1: # 灰度图
|
||||
image = image.squeeze(0)
|
||||
image = (image * 255).clamp(0, 255).to(torch.uint8)
|
||||
return Image.fromarray(image.cpu().numpy(), mode='L')
|
||||
else: # RGB图
|
||||
image = image.permute(1, 2, 0)
|
||||
image = (image * 255).clamp(0, 255).to(torch.uint8)
|
||||
return Image.fromarray(image.cpu().numpy(), mode='RGB')
|
||||
return None
|
||||
|
||||
def pil2tensor(image):
|
||||
"""将PIL图像转换为torch张量"""
|
||||
if image.mode == 'L':
|
||||
image = image.convert('RGB')
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = torch.from_numpy(image)
|
||||
if image.dim() == 3:
|
||||
image = image.permute(2, 0, 1)
|
||||
return image
|
||||
|
||||
def image2mask(image):
|
||||
"""将PIL图像转换为掩码张量"""
|
||||
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 tensor2cv2(image: torch.Tensor) -> np.array:
|
||||
"""将torch张量转换为OpenCV格式"""
|
||||
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 pil2cv2(pil_img):
|
||||
"""将PIL图像转换为OpenCV格式"""
|
||||
np_img_array = np.asarray(pil_img)
|
||||
return cv2.cvtColor(np_img_array, cv2.COLOR_RGB2BGR)
|
||||
|
||||
def cv22pil(cv2_img):
|
||||
"""将OpenCV图像转换为PIL图像"""
|
||||
cv2_img = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)
|
||||
return Image.fromarray(cv2_img)
|
||||
|
||||
def tensor2np(tensor):
|
||||
"""将torch张量转换为numpy数组"""
|
||||
if len(tensor.shape) == 3: # Single image
|
||||
return np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8)
|
||||
else: # Batch of images
|
||||
return [np.clip(255.0 * t.cpu().numpy(), 0, 255).astype(np.uint8) for t in tensor]
|
||||
|
||||
# Color transfer functions for ImitationHueNode
|
||||
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:
|
||||
@@ -139,6 +35,7 @@ def adjust_brightness(image, factor, mask=None):
|
||||
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:
|
||||
@@ -150,6 +47,7 @@ def adjust_saturation(image, factor, mask=None):
|
||||
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:
|
||||
@@ -161,6 +59,7 @@ def adjust_contrast(image, factor, mask=None):
|
||||
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)
|
||||
@@ -206,6 +105,7 @@ def adjust_tone(source, target, tone_strength=0.7, mask=None):
|
||||
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)
|
||||
|
||||
@@ -54,7 +54,7 @@ def draw_rect(image, x, y, width, height, line_color="#FF0000", line_width=2):
|
||||
return image
|
||||
|
||||
class CropByMask_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
@@ -15,7 +15,7 @@ import numpy as np
|
||||
import folder_paths
|
||||
|
||||
class DepthMapBlur_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
Empty Unit Generator Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Generates empty images, masks, and latents with various preset ratios and configurations.
|
||||
Generate empty units for testing and development.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
@@ -12,10 +12,12 @@ import torch
|
||||
import numpy as np
|
||||
import re
|
||||
from PIL import Image
|
||||
from ..common_utils import log
|
||||
import cv2
|
||||
|
||||
from ..tools.logging_utils import log
|
||||
|
||||
class EmptyUnitGenerator_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -108,5 +110,5 @@ NODE_CLASS_MAPPINGS = {
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"EmptyUnitGenerator_UTK": "Empty Unit Generator",
|
||||
"EmptyUnitGenerator_UTK": "Empty Unit Generator (UTK)",
|
||||
}
|
||||
@@ -9,9 +9,12 @@ Preview an image or a mask, when both inputs are used composites the mask on top
|
||||
"""
|
||||
|
||||
import random
|
||||
from nodes import SaveImage, ImageCompositeMasked
|
||||
from nodes import SaveImage
|
||||
import folder_paths
|
||||
|
||||
# 导入本地的 ImageCompositeMasked 实现
|
||||
from .image_composite_masked import ImageCompositeMasked
|
||||
|
||||
class ImageAndMaskPreview_UTK(SaveImage):
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_temp_directory()
|
||||
@@ -37,7 +40,7 @@ class ImageAndMaskPreview_UTK(SaveImage):
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("composite",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
DESCRIPTION = """
|
||||
Preview an image or a mask, when both inputs are used
|
||||
composites the mask on top of the image.
|
||||
|
||||
@@ -31,7 +31,7 @@ def image_channel_merge(channels, mode):
|
||||
return Image.merge(mode, channels)
|
||||
|
||||
class ImageCombineAlpha_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
"""
|
||||
Image Composite Masked Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Local implementation of ImageCompositeMasked for UniversalToolkit.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
class ImageCompositeMasked:
|
||||
"""
|
||||
Local implementation of ImageCompositeMasked for compositing images with masks.
|
||||
Based on KJNodes implementation.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def composite(self, image, mask_image, x, y, resize_mask, mask):
|
||||
"""
|
||||
Composite an image with a mask at specified position.
|
||||
|
||||
Args:
|
||||
image: Base image tensor (B, H, W, C)
|
||||
mask_image: Mask image tensor (B, H, W, C)
|
||||
x: X offset
|
||||
y: Y offset
|
||||
resize_mask: Whether to resize mask to match image
|
||||
mask: Alpha mask tensor (B, H, W)
|
||||
|
||||
Returns:
|
||||
Composited image tensor
|
||||
"""
|
||||
if image is None:
|
||||
return (mask_image,)
|
||||
|
||||
if mask_image is None:
|
||||
return (image,)
|
||||
|
||||
# Ensure tensors are on the same device
|
||||
device = image.device
|
||||
mask_image = mask_image.to(device)
|
||||
mask = mask.to(device) if mask is not None else None
|
||||
|
||||
# Get dimensions
|
||||
batch_size, image_height, image_width, channels = image.shape
|
||||
mask_batch_size, mask_height, mask_width, mask_channels = mask_image.shape
|
||||
|
||||
# Handle batch size mismatch
|
||||
if batch_size != mask_batch_size:
|
||||
if batch_size == 1:
|
||||
mask_image = mask_image[:1]
|
||||
elif mask_batch_size == 1:
|
||||
mask_image = mask_image.expand(batch_size, -1, -1, -1)
|
||||
else:
|
||||
raise ValueError("Batch sizes must match or one must be 1")
|
||||
|
||||
# Resize mask if needed
|
||||
if resize_mask and (mask_height != image_height or mask_width != image_width):
|
||||
mask_image = F.interpolate(
|
||||
mask_image.permute(0, 3, 1, 2), # (B, C, H, W)
|
||||
size=(image_height, image_width),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).permute(0, 2, 3, 1) # (B, H, W, C)
|
||||
|
||||
# Apply mask if provided
|
||||
if mask is not None:
|
||||
if mask.shape[1:] != (image_height, image_width):
|
||||
mask = F.interpolate(
|
||||
mask.unsqueeze(1), # (B, 1, H, W)
|
||||
size=(image_height, image_width),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).squeeze(1) # (B, H, W)
|
||||
|
||||
# Expand mask to match channels
|
||||
mask = mask.unsqueeze(-1).expand(-1, -1, -1, channels)
|
||||
mask_image = mask_image * mask
|
||||
|
||||
# Calculate crop region
|
||||
if x < 0:
|
||||
crop_x = -x
|
||||
x = 0
|
||||
else:
|
||||
crop_x = 0
|
||||
|
||||
if y < 0:
|
||||
crop_y = -y
|
||||
y = 0
|
||||
else:
|
||||
crop_y = 0
|
||||
|
||||
# Crop mask image if needed
|
||||
if crop_x > 0 or crop_y > 0:
|
||||
mask_image = mask_image[:, crop_y:, crop_x:, :]
|
||||
|
||||
# Calculate final dimensions
|
||||
mask_height, mask_width = mask_image.shape[1:3]
|
||||
|
||||
# Check bounds
|
||||
if x + mask_width > image_width:
|
||||
mask_width = image_width - x
|
||||
mask_image = mask_image[:, :, :mask_width, :]
|
||||
|
||||
if y + mask_height > image_height:
|
||||
mask_height = image_height - y
|
||||
mask_image = mask_image[:, :mask_height, :, :]
|
||||
|
||||
# Create output image
|
||||
result = image.clone()
|
||||
|
||||
# Composite mask image onto result
|
||||
if mask is not None:
|
||||
# Use alpha blending
|
||||
alpha = mask[:, y:y+mask_height, x:x+mask_width, :]
|
||||
result[:, y:y+mask_height, x:x+mask_width, :] = (
|
||||
result[:, y:y+mask_height, x:x+mask_width, :] * (1 - alpha) +
|
||||
mask_image * alpha
|
||||
)
|
||||
else:
|
||||
# Direct replacement
|
||||
result[:, y:y+mask_height, x:x+mask_width, :] = mask_image
|
||||
|
||||
return (result,)
|
||||
@@ -13,7 +13,7 @@ import torch.nn.functional as F
|
||||
from ..common_utils import log
|
||||
|
||||
class ImageConcatenate_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -201,5 +201,5 @@ NODE_CLASS_MAPPINGS = {
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageConcatenate_UTK": "Image Concatenate",
|
||||
"ImageConcatenate_UTK": "Image Concatenate (UTK)",
|
||||
}
|
||||
@@ -11,7 +11,7 @@ Concatenates multiple images in various directions and layouts.
|
||||
import torch
|
||||
|
||||
class ImageConcatenateMulti_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
@@ -0,0 +1,72 @@
|
||||
"""
|
||||
Image Converters for UniversalToolkit
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Image format conversion utilities for UniversalToolkit.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image
|
||||
|
||||
def tensor2pil(image):
|
||||
"""将torch张量转换为PIL图像"""
|
||||
if image.dim() == 4:
|
||||
image = image.squeeze(0)
|
||||
if image.dim() == 3:
|
||||
if image.shape[0] == 1: # 灰度图
|
||||
image = image.squeeze(0)
|
||||
image = (image * 255).clamp(0, 255).to(torch.uint8)
|
||||
return Image.fromarray(image.cpu().numpy(), mode='L')
|
||||
else: # RGB图
|
||||
image = image.permute(1, 2, 0)
|
||||
image = (image * 255).clamp(0, 255).to(torch.uint8)
|
||||
return Image.fromarray(image.cpu().numpy(), mode='RGB')
|
||||
return None
|
||||
|
||||
def pil2tensor(image):
|
||||
"""将PIL图像转换为torch张量"""
|
||||
if image.mode == 'L':
|
||||
image = image.convert('RGB')
|
||||
image = np.array(image).astype(np.float32) / 255.0
|
||||
image = torch.from_numpy(image)
|
||||
if image.dim() == 3:
|
||||
image = image.permute(2, 0, 1)
|
||||
return image
|
||||
|
||||
def image2mask(image):
|
||||
"""将PIL图像转换为掩码张量"""
|
||||
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 tensor2cv2(image: torch.Tensor) -> np.array:
|
||||
"""将torch张量转换为OpenCV格式"""
|
||||
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 pil2cv2(pil_img):
|
||||
"""将PIL图像转换为OpenCV格式"""
|
||||
np_img_array = np.asarray(pil_img)
|
||||
return cv2.cvtColor(np_img_array, cv2.COLOR_RGB2BGR)
|
||||
|
||||
def cv22pil(cv2_img):
|
||||
"""将OpenCV图像转换为PIL图像"""
|
||||
cv2_img = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)
|
||||
return Image.fromarray(cv2_img)
|
||||
|
||||
def tensor2np(tensor):
|
||||
"""将torch张量转换为numpy数组"""
|
||||
if len(tensor.shape) == 3: # Single image
|
||||
return np.clip(255.0 * tensor.cpu().numpy(), 0, 255).astype(np.uint8)
|
||||
else: # Batch of images
|
||||
return [np.clip(255.0 * t.cpu().numpy(), 0, 255).astype(np.uint8) for t in tensor]
|
||||
@@ -52,7 +52,7 @@ def fit_resize_image(image, target_width, target_height, fit_mode, resize_sample
|
||||
return image.resize((target_width, target_height), resize_sampler)
|
||||
|
||||
class ImageMaskScaleAs_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
@@ -14,7 +14,7 @@ import torch.nn.functional as F
|
||||
MAX_RESOLUTION = 8192
|
||||
|
||||
class ImagePadForOutpaintMasked_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
@@ -1,18 +1,20 @@
|
||||
"""
|
||||
Image Ratio Detector Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Detects and analyzes image aspect ratios and dimensions.
|
||||
Detect the aspect ratio of an image.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import math
|
||||
from ..common_utils import log
|
||||
|
||||
from ..tools.logging_utils import log
|
||||
|
||||
class ImageRatioDetector_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -73,5 +75,5 @@ NODE_CLASS_MAPPINGS = {
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageRatioDetector_UTK": "Image Ratio Detector",
|
||||
"ImageRatioDetector_UTK": "Image Ratio Detector (UTK)",
|
||||
}
|
||||
@@ -13,7 +13,7 @@ from PIL import Image
|
||||
from ...common_utils import log, tensor2pil, pil2tensor
|
||||
|
||||
class ImageRemoveAlpha_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
@@ -54,7 +54,7 @@ def fit_resize_image(image, target_width, target_height, fit_mode, resize_sample
|
||||
return image.resize((target_width, target_height), resize_sampler)
|
||||
|
||||
class ImageScaleByAspectRatio_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
@@ -13,7 +13,7 @@ from PIL import Image
|
||||
from ...common_utils import log, tensor2pil, pil2tensor, image2mask
|
||||
|
||||
class ImageScaleRestore_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
@@ -12,187 +12,8 @@ import torch
|
||||
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 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
|
||||
@@ -217,34 +38,49 @@ class ImitationHueNode_UTK:
|
||||
},
|
||||
}
|
||||
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "imitation_hue"
|
||||
DESCRIPTION = """
|
||||
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):
|
||||
for img in imitation_image:
|
||||
img_cv1 = tensor2cv2(img)
|
||||
|
||||
for img in target_image:
|
||||
img_cv2 = tensor2cv2(img)
|
||||
|
||||
img_cv3 = None
|
||||
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
|
||||
if mask is not None:
|
||||
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,)
|
||||
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,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
|
||||
@@ -2,15 +2,20 @@
|
||||
Purge VRAM Node
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
Purges GPU memory and optionally unloads models.
|
||||
Purge GPU memory to free up VRAM.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch.cuda
|
||||
import torch
|
||||
import gc
|
||||
from ...common_utils import log
|
||||
|
||||
from ..tools.logging_utils import log
|
||||
from ..tools.any_type import AnyType
|
||||
|
||||
# 创建 AnyType 实例
|
||||
any = AnyType("*")
|
||||
|
||||
def clear_memory():
|
||||
"""Clear GPU memory"""
|
||||
@@ -19,13 +24,13 @@ def clear_memory():
|
||||
gc.collect()
|
||||
|
||||
class PurgeVRAM_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"anything": ("*", {}),
|
||||
"anything": (any, {}),
|
||||
"purge_cache": ("BOOLEAN", {"default": True}),
|
||||
"purge_models": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
@@ -33,7 +38,7 @@ class PurgeVRAM_UTK:
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("*",)
|
||||
RETURN_TYPES = (any,)
|
||||
RETURN_NAMES = ("anything",)
|
||||
FUNCTION = "purge_vram"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
@@ -13,7 +13,7 @@ from PIL import Image
|
||||
from ...common_utils import log, tensor2pil, pil2tensor, image2mask
|
||||
|
||||
class RestoreCropBox_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Image"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
@@ -11,7 +11,7 @@ Performs logical operations on masks.
|
||||
import torch
|
||||
|
||||
class MaskAnd_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Mask"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"mask1": ("MASK",), "mask2": ("MASK",)}}
|
||||
@@ -25,7 +25,7 @@ class MaskAnd_UTK:
|
||||
return (mask1 * mask2,)
|
||||
|
||||
class MaskSub_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Mask"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"mask1": ("MASK",), "mask2": ("MASK",)}}
|
||||
@@ -39,7 +39,7 @@ class MaskSub_UTK:
|
||||
return (torch.clamp(mask1 - mask2, 0, 1),)
|
||||
|
||||
class MaskAdd_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Mask"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"mask1": ("MASK",), "mask2": ("MASK",)}}
|
||||
|
||||
@@ -0,0 +1,16 @@
|
||||
"""
|
||||
AnyType for UniversalToolkit
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
AnyType class for accepting any input type in nodes.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
class AnyType(str):
|
||||
"""A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
|
||||
def __eq__(self, __value: object) -> bool:
|
||||
return True
|
||||
def __ne__(self, __value: object) -> bool:
|
||||
return False
|
||||
@@ -32,7 +32,7 @@ def mask_blur(mask, feathering):
|
||||
return mask
|
||||
|
||||
class FillMaskedArea_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Tools"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
"""
|
||||
Logging Utilities for UniversalToolkit
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Logging utilities for UniversalToolkit.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
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}")
|
||||
@@ -11,7 +11,7 @@ Display and preview nodes for various data types.
|
||||
import torch
|
||||
|
||||
class Show_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Tools"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"input": ("STRING", "INT", "FLOAT", "LIST", "MASK", "IMAGE", "LATENT")}}
|
||||
@@ -39,7 +39,7 @@ class Show_UTK:
|
||||
return tuple(outs)
|
||||
|
||||
class ShowInt_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Tools"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"int_val": ("INT",)}}
|
||||
@@ -53,7 +53,7 @@ class ShowInt_UTK:
|
||||
return (int_val,)
|
||||
|
||||
class ShowFloat_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Tools"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"float_val": ("FLOAT",)}}
|
||||
@@ -67,7 +67,7 @@ class ShowFloat_UTK:
|
||||
return (float_val,)
|
||||
|
||||
class ShowList_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Tools"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"list_val": ("LIST",)}}
|
||||
@@ -81,7 +81,7 @@ class ShowList_UTK:
|
||||
return (list_val,)
|
||||
|
||||
class ShowText_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Tools"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"text": ("STRING",)}}
|
||||
@@ -95,7 +95,7 @@ class ShowText_UTK:
|
||||
return (text,)
|
||||
|
||||
class PreviewMask_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
CATEGORY = "UniversalToolkit/Tools"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"mask": ("MASK",)}}
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "universaltoolkit"
|
||||
description = "A comprehensive toolkit based on ComfyUI, providing image, mask, audio, and tools nodes, fully modular and v3 compatible."
|
||||
version = "1.1.2"
|
||||
version = "1.1.3"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["torch", "numpy", "Pillow", "opencv-python", "scipy", "tqdm"]
|
||||
requires-python = ">=3.8"
|
||||
|
||||
Reference in New Issue
Block a user