拆分代码结构
This commit is contained in:
Cyber Dick Lang
@@ -23,6 +23,7 @@
|
||||
## 节点开发
|
||||
- 节点参数需有默认值、类型、范围说明。
|
||||
- 输入输出类型必须与 ComfyUI 规范一致。
|
||||
- **所有节点的输入和输出shape、类型、参数名、返回名，必须严格遵循ComfyUI官方节点开发规范，保证与ComfyUI原生节点和其他插件节点100%兼容。不得自定义非标准shape或类型。**
|
||||
- 新增节点需在 `__init__.py` 注册，并补充到文档。
|
||||
- 节点参数、UI、输出类型、行为与参考项目完全一致。
|
||||
- 不得随意增删参数或更改默认值。
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
# ComfyUI Universal Toolkit 重构计划
|
||||
|
||||
## 概述
|
||||
|
||||
本项目将按照 ComfyUI-LayerStyle 项目的模块化组织方式，将大型节点文件拆分成多个独立的功能模块，提高代码的可维护性和模块化程度。
|
||||
|
||||
## 重构目标
|
||||
|
||||
1. **模块化组织**：将大型节点文件按功能拆分成独立模块
|
||||
2. **代码复用**：创建共用工具函数，避免重复代码
|
||||
3. **易于维护**：每个节点独立文件，便于修改和调试
|
||||
4. **清晰结构**：按功能分类组织代码结构
|
||||
|
||||
## 目录结构
|
||||
|
||||
```
|
||||
nodes/
|
||||
├── common_utils.py # 共用工具函数
|
||||
├── image/ # 图像处理节点
|
||||
│ ├── __init__.py
|
||||
│ ├── empty_unit_generator.py
|
||||
│ ├── image_ratio_detector.py
|
||||
│ ├── depth_map_blur.py
|
||||
│ ├── image_concatenate.py
|
||||
│ ├── image_concatenate_multi.py
|
||||
│ ├── image_pad_for_outpaint.py
|
||||
│ ├── image_and_mask_preview.py
|
||||
│ ├── imitation_hue_node.py
|
||||
│ ├── image_scale_by_aspect_ratio.py
|
||||
│ ├── image_mask_scale_as.py
|
||||
│ ├── image_scale_restore.py
|
||||
│ ├── image_remove_alpha.py
|
||||
│ ├── image_combine_alpha.py
|
||||
│ ├── check_mask.py
|
||||
│ ├── purge_vram.py
|
||||
│ ├── crop_by_mask.py
|
||||
│ └── restore_crop_box.py
|
||||
├── tools/ # 工具类节点
|
||||
│ ├── __init__.py
|
||||
│ ├── show_int.py
|
||||
│ ├── show_float.py
|
||||
│ ├── show_list.py
|
||||
│ ├── show_text.py
|
||||
│ ├── preview_mask.py
|
||||
│ └── fill_masked_area.py
|
||||
├── mask/ # 掩码处理节点
|
||||
│ ├── __init__.py
|
||||
│ ├── mask_and.py
|
||||
│ ├── mask_sub.py
|
||||
│ └── mask_add.py
|
||||
├── audio/ # 音频处理节点
|
||||
│ ├── __init__.py
|
||||
│ ├── load_audio_plus_from_path.py
|
||||
│ └── audio_crop_process.py
|
||||
├── image_nodes_utk.py # 原有文件（逐步迁移后删除）
|
||||
├── tool_nodes_utk.py # 原有文件（逐步迁移后删除）
|
||||
├── mask_nodes_utk.py # 原有文件（逐步迁移后删除）
|
||||
└── audio_nodes_utk.py # 原有文件（逐步迁移后删除）
|
||||
```
|
||||
|
||||
## 重构步骤
|
||||
|
||||
### 第一阶段：基础架构 ✅
|
||||
- [x] 创建目录结构
|
||||
- [x] 创建共用工具函数 `common_utils.py`
|
||||
- [x] 创建模块初始化文件
|
||||
|
||||
### 第二阶段：图像节点迁移 🔄
|
||||
- [x] EmptyUnitGenerator_UTK
|
||||
- [x] ImageRatioDetector_UTK
|
||||
- [x] DepthMapBlur_UTK
|
||||
- [x] ImageConcatenate_UTK
|
||||
- [ ] ImageConcatenateMulti_UTK
|
||||
- [ ] ImagePadForOutpaintMasked_UTK
|
||||
- [ ] ImageAndMaskPreview_UTK
|
||||
- [ ] ImitationHueNode_UTK
|
||||
- [ ] ImageScaleByAspectRatio_UTK
|
||||
- [ ] ImageMaskScaleAs_UTK
|
||||
- [ ] ImageScaleRestore_UTK
|
||||
- [ ] ImageRemoveAlpha_UTK
|
||||
- [ ] ImageCombineAlpha_UTK
|
||||
- [ ] CheckMask_UTK
|
||||
- [ ] PurgeVRAM_UTK
|
||||
- [ ] CropByMask_UTK
|
||||
- [ ] RestoreCropBox_UTK
|
||||
|
||||
### 第三阶段：工具节点迁移
|
||||
- [ ] ShowInt_UTK
|
||||
- [ ] ShowFloat_UTK
|
||||
- [ ] ShowList_UTK
|
||||
- [ ] ShowText_UTK
|
||||
- [ ] PreviewMask_UTK
|
||||
- [ ] FillMaskedArea_UTK
|
||||
|
||||
### 第四阶段：掩码节点迁移
|
||||
- [ ] MaskAnd_UTK
|
||||
- [ ] MaskSub_UTK
|
||||
- [ ] MaskAdd_UTK
|
||||
|
||||
### 第五阶段：音频节点迁移
|
||||
- [ ] LoadAudioPlusFromPath_UTK
|
||||
- [ ] AudioCropProcessUTK
|
||||
|
||||
### 第六阶段：清理和优化
|
||||
- [ ] 删除原有大型文件
|
||||
- [ ] 更新文档
|
||||
- [ ] 测试所有节点功能
|
||||
- [ ] 优化导入结构
|
||||
|
||||
## 节点文件模板
|
||||
|
||||
每个节点文件应遵循以下模板：
|
||||
|
||||
```python
|
||||
"""
|
||||
节点名称
|
||||
~~~~~~~~
|
||||
|
||||
节点功能描述
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
# 其他必要的导入
|
||||
from ..common_utils import log, tensor2pil, pil2tensor # 使用相对导入
|
||||
|
||||
class NodeName_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
# 输入参数定义
|
||||
},
|
||||
"optional": {
|
||||
# 可选参数定义
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("TYPE1", "TYPE2")
|
||||
RETURN_NAMES = ("name1", "name2")
|
||||
FUNCTION = "function_name"
|
||||
|
||||
def function_name(self, param1, param2, ...):
|
||||
# 节点实现逻辑
|
||||
return (output1, output2)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"NodeName_UTK": NodeName_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NodeName_UTK": "Node Display Name",
|
||||
}
|
||||
```
|
||||
|
||||
## 导入结构
|
||||
|
||||
主 `__init__.py` 文件应使用以下导入结构：
|
||||
|
||||
```python
|
||||
# 导入新的模块化节点
|
||||
from .nodes.image.node_file import NODE_CLASS_MAPPINGS as NODE_MAPPINGS
|
||||
from .nodes.image.node_file import NODE_DISPLAY_NAME_MAPPINGS as NODE_DISPLAY_MAPPINGS
|
||||
|
||||
# 合并所有映射
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_CLASS_MAPPINGS.update(NODE_MAPPINGS)
|
||||
# ... 其他映射
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(NODE_DISPLAY_MAPPINGS)
|
||||
# ... 其他显示名称映射
|
||||
```
|
||||
|
||||
## 优势
|
||||
|
||||
1. **模块化**：每个节点独立文件，便于维护
|
||||
2. **可扩展**：新增节点只需创建新文件
|
||||
3. **代码复用**：共用函数避免重复代码
|
||||
4. **清晰结构**：按功能分类，易于理解
|
||||
5. **冲突减少**：独立文件减少合并冲突
|
||||
|
||||
## 注意事项
|
||||
|
||||
1. **相对导入**：使用 `from ..common_utils import` 进行相对导入
|
||||
2. **节点映射**：每个文件都要包含 `NODE_CLASS_MAPPINGS` 和 `NODE_DISPLAY_NAME_MAPPINGS`
|
||||
3. **文档注释**：每个文件都要有清晰的文档字符串
|
||||
4. **功能测试**：迁移后要测试节点功能是否正常
|
||||
5. **渐进迁移**：逐步迁移，保持项目可用性
|
||||
|
||||
## 完成状态
|
||||
|
||||
- [x] 基础架构搭建
|
||||
- [x] 共用工具函数创建
|
||||
- [x] 部分图像节点迁移
|
||||
- [ ] 完整节点迁移
|
||||
- [ ] 测试和优化
|
||||
- [ ] 文档更新
|
||||
+259
-49
@@ -8,13 +8,59 @@ A comprehensive toolkit for ComfyUI that provides various utility nodes for imag
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
__version__ = "1.0.8"
|
||||
__version__ = "1.1.2"
|
||||
__author__ = "CyberDickLang"
|
||||
__email__ = "286878701@qq.com"
|
||||
__url__ = "https://github.com/whmc76"
|
||||
|
||||
# 更新日志
|
||||
CHANGELOG = {
|
||||
"1.1.2": [
|
||||
"修复 DepthMapBlur_UTK 节点 kernel size 类型和 OpenCV 奇数断言问题，保证所有模糊核为正奇数，完全兼容 ComfyUI 规范。",
|
||||
"修正 EmptyUnitGenerator_UTK 输出 shape，所有节点输入输出严格遵循 ComfyUI 官方规范。",
|
||||
"完善节点导入路径为相对导入，兼容 ComfyUI v3 插件机制。"
|
||||
],
|
||||
"1.1.1": [
|
||||
"项目重构：",
|
||||
"- 将大型节点文件按功能拆分成多个独立模块",
|
||||
"- 创建共用工具函数文件 common_utils.py",
|
||||
"- 按功能分类：image/、tools/、mask/、audio/",
|
||||
"- 提高代码可维护性和模块化程度",
|
||||
"- 参考 ComfyUI-LayerStyle 项目的模块化组织方式",
|
||||
"新增多个图像处理节点：",
|
||||
"- ImageScaleByAspectRatio_UTK: 按比例缩放图像",
|
||||
"- ImageMaskScaleAs_UTK: 按参考图像缩放",
|
||||
"- ImageScaleRestore_UTK: 恢复图像尺寸",
|
||||
"- ImageRemoveAlpha_UTK: 移除Alpha通道",
|
||||
"- ImageCombineAlpha_UTK: 合并Alpha通道",
|
||||
"- CheckMask_UTK: 检查掩码有效性",
|
||||
"- PurgeVRAM_UTK: 清理显存",
|
||||
"- CropByMask_UTK: 基于掩码裁剪",
|
||||
"- RestoreCropBox_UTK: 恢复裁剪框",
|
||||
],
|
||||
"1.1.0": [
|
||||
"新增 CropByMaskV3_UTK 节点（掩码裁剪节点）:",
|
||||
"- 基于 ComfyUI-LayerStyle 项目集成",
|
||||
"- 支持三种检测模式：mask_area、min_bounding_rect、max_inscribed_rect",
|
||||
"- 支持掩码反转功能",
|
||||
"- 支持四个方向的预留边距设置",
|
||||
"- 支持整除倍数调整",
|
||||
"- 提供裁剪框预览功能",
|
||||
"新增 RestoreCropBox_UTK 节点（恢复裁剪框节点）:",
|
||||
"- 基于 ComfyUI-LayerStyle 项目集成",
|
||||
"- 支持将裁剪后的图像恢复到原始背景",
|
||||
"- 支持掩码反转功能",
|
||||
"- 支持自定义裁剪框位置",
|
||||
"- 自动处理批量图像和掩码",
|
||||
],
|
||||
"1.0.9": [
|
||||
"新增 PurgeVRAM_UTK 节点（显存清理节点）:",
|
||||
"- 基于 ComfyUI-LayerStyle 项目集成",
|
||||
"- 支持清理GPU缓存和模型内存",
|
||||
"- 支持选择性清理（缓存/模型）",
|
||||
"- 添加内存清理辅助函数",
|
||||
"- 支持任意类型输入，保持数据流连续性",
|
||||
],
|
||||
"1.0.8": [
|
||||
"新增 ImitationHueNode_UTK 节点（追色节点）:",
|
||||
"- 基于 ComfyUI-MingNodes 项目集成",
|
||||
@@ -84,54 +130,206 @@ CHANGELOG = {
|
||||
]
|
||||
}
|
||||
|
||||
from .nodes.image_nodes_utk import EmptyUnitGenerator_UTK, ImageRatioDetector_UTK, DepthMapBlur_UTK, ImageConcatenate_UTK, ImageConcatenateMulti_UTK, ImagePadForOutpaintMasked_UTK, ImageAndMaskPreview_UTK, ImitationHueNode_UTK
|
||||
from .nodes.tool_nodes_utk import ShowInt_UTK, ShowFloat_UTK, ShowList_UTK, ShowText_UTK, PreviewMask_UTK, FillMaskedArea_UTK
|
||||
from .nodes.audio_nodes_utk import LoadAudioPlusFromPath_UTK, AudioCropProcessUTK
|
||||
from .nodes.mask_nodes_utk import MaskAnd_UTK, MaskSub_UTK, MaskAdd_UTK
|
||||
# 导入节点模块
|
||||
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.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.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.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_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_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.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
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"EmptyUnitGenerator_UTK": EmptyUnitGenerator_UTK,
|
||||
"ImageRatioDetector_UTK": ImageRatioDetector_UTK,
|
||||
"ShowInt_UTK": ShowInt_UTK,
|
||||
"ShowFloat_UTK": ShowFloat_UTK,
|
||||
"ShowList_UTK": ShowList_UTK,
|
||||
"ShowText_UTK": ShowText_UTK,
|
||||
"PreviewMask_UTK": PreviewMask_UTK,
|
||||
"FillMaskedArea_UTK": FillMaskedArea_UTK,
|
||||
"ImageAndMaskPreview_UTK": ImageAndMaskPreview_UTK,
|
||||
"DepthMapBlur_UTK": DepthMapBlur_UTK,
|
||||
"ImageConcatenate_UTK": ImageConcatenate_UTK,
|
||||
"ImageConcatenateMulti_UTK": ImageConcatenateMulti_UTK,
|
||||
"ImagePadForOutpaintMasked_UTK": ImagePadForOutpaintMasked_UTK,
|
||||
"LoadAudioPlusFromPath_UTK": LoadAudioPlusFromPath_UTK,
|
||||
"AudioCropProcessUTK": AudioCropProcessUTK,
|
||||
"MaskAnd_UTK": MaskAnd_UTK,
|
||||
"MaskSub_UTK": MaskSub_UTK,
|
||||
"MaskAdd_UTK": MaskAdd_UTK,
|
||||
"ImitationHueNode_UTK": ImitationHueNode_UTK,
|
||||
}
|
||||
except ImportError as e:
|
||||
print(f"导入错误: {e}")
|
||||
# 如果模块化导入失败，使用空字典
|
||||
SHOW_NODES_MAPPINGS = {}
|
||||
SHOW_NODES_DISPLAY_MAPPINGS = {}
|
||||
AUDIO_CROP_MAPPINGS = {}
|
||||
AUDIO_CROP_DISPLAY_MAPPINGS = {}
|
||||
MASK_OPERATIONS_MAPPINGS = {}
|
||||
MASK_OPERATIONS_DISPLAY_MAPPINGS = {}
|
||||
CONCATENATE_MULTI_MAPPINGS = {}
|
||||
CONCATENATE_MULTI_DISPLAY_MAPPINGS = {}
|
||||
PAD_OUTPAINT_MAPPINGS = {}
|
||||
PAD_OUTPAINT_DISPLAY_MAPPINGS = {}
|
||||
AND_MASK_PREVIEW_MAPPINGS = {}
|
||||
AND_MASK_PREVIEW_DISPLAY_MAPPINGS = {}
|
||||
IMITATION_HUE_MAPPINGS = {}
|
||||
IMITATION_HUE_DISPLAY_MAPPINGS = {}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"EmptyUnitGenerator_UTK": "Empty Unit Generator",
|
||||
"ImageRatioDetector_UTK": "Image Ratio Detector",
|
||||
"ShowInt_UTK": "Show Int (UTK)",
|
||||
"ShowFloat_UTK": "Show Float (UTK)",
|
||||
"ShowList_UTK": "Show List (UTK)",
|
||||
"ShowText_UTK": "Show Text (UTK)",
|
||||
"PreviewMask_UTK": "Preview Mask (UTK)",
|
||||
"FillMaskedArea_UTK": "Fill Masked Area (UTK)",
|
||||
"ImageAndMaskPreview_UTK": "Image And Mask Preview (UTK)",
|
||||
"DepthMapBlur_UTK": "Depth Map Blur",
|
||||
"ImageConcatenate_UTK": "Image Concatenate",
|
||||
"ImageConcatenateMulti_UTK": "Image Concatenate Multi",
|
||||
"ImagePadForOutpaintMasked_UTK": "Image Pad For Outpaint Masked (UTK)",
|
||||
"LoadAudioPlusFromPath_UTK": "Load Audio Plus From Path (UTK)",
|
||||
"AudioCropProcessUTK": "Audio Crop Process (UTK)",
|
||||
"MaskAnd_UTK": "Mask And (UTK)",
|
||||
"MaskSub_UTK": "Mask Sub (UTK)",
|
||||
"MaskAdd_UTK": "Mask Add (UTK)",
|
||||
"ImitationHueNode_UTK": "Imitation Hue Node (UTK)",
|
||||
}
|
||||
# 尝试导入其他可能有依赖的节点
|
||||
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
|
||||
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
|
||||
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
|
||||
except ImportError:
|
||||
IMAGE_GENERATOR_MAPPINGS = {}
|
||||
IMAGE_GENERATOR_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
IMAGE_DETECTOR_MAPPINGS = {}
|
||||
IMAGE_DETECTOR_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
DEPTH_BLUR_MAPPINGS = {}
|
||||
DEPTH_BLUR_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
CONCATENATE_MAPPINGS = {}
|
||||
CONCATENATE_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
SCALE_ASPECT_MAPPINGS = {}
|
||||
SCALE_ASPECT_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
MASK_SCALE_MAPPINGS = {}
|
||||
MASK_SCALE_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
SCALE_RESTORE_MAPPINGS = {}
|
||||
SCALE_RESTORE_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
REMOVE_ALPHA_MAPPINGS = {}
|
||||
REMOVE_ALPHA_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
COMBINE_ALPHA_MAPPINGS = {}
|
||||
COMBINE_ALPHA_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
CHECK_MASK_MAPPINGS = {}
|
||||
CHECK_MASK_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
PURGE_VRAM_MAPPINGS = {}
|
||||
PURGE_VRAM_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
CROP_MASK_MAPPINGS = {}
|
||||
CROP_MASK_DISPLAY_MAPPINGS = {}
|
||||
|
||||
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
|
||||
except ImportError:
|
||||
RESTORE_CROP_MAPPINGS = {}
|
||||
RESTORE_CROP_DISPLAY_MAPPINGS = {}
|
||||
|
||||
# 合并所有节点映射
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_CLASS_MAPPINGS.update(IMAGE_GENERATOR_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(IMAGE_DETECTOR_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(DEPTH_BLUR_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(CONCATENATE_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(CONCATENATE_MULTI_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(PAD_OUTPAINT_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(AND_MASK_PREVIEW_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(IMITATION_HUE_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(SCALE_ASPECT_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(MASK_SCALE_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(SCALE_RESTORE_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(REMOVE_ALPHA_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(COMBINE_ALPHA_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(CHECK_MASK_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(PURGE_VRAM_MAPPINGS)
|
||||
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(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(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(SHOW_NODES_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(FILL_MASKED_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(LOAD_AUDIO_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(AUDIO_CROP_DISPLAY_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(MASK_OPERATIONS_DISPLAY_MAPPINGS)
|
||||
|
||||
NODE_CATEGORIES = {
|
||||
"UniversalToolkit": [
|
||||
@@ -149,6 +347,18 @@ NODE_CATEGORIES = {
|
||||
"MaskSub_UTK",
|
||||
"MaskAdd_UTK",
|
||||
"ImitationHueNode_UTK",
|
||||
"ImageScaleByAspectRatio_UTK",
|
||||
"ImageMaskScaleAs_UTK",
|
||||
"ImageScaleRestore_UTK",
|
||||
"ImageRemoveAlpha_UTK",
|
||||
"ImageCombineAlpha_UTK",
|
||||
"CheckMask_UTK",
|
||||
"PurgeVRAM_UTK",
|
||||
"CropByMask_UTK",
|
||||
"RestoreCropBox_UTK",
|
||||
"Show_UTK",
|
||||
"ShowFloat_UTK",
|
||||
"PreviewMask_UTK",
|
||||
]
|
||||
}
|
||||
|
||||
@@ -160,5 +370,5 @@ __all__ = [
|
||||
"__author__",
|
||||
"__email__",
|
||||
"__url__",
|
||||
"CHANGELOG"
|
||||
]
|
||||
"CHANGELOG",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
"""
|
||||
Audio Crop Process Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Processes audio with cropping, resampling, and gain adjustments.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
FLOAT_MAX = 99999999999999999.0
|
||||
|
||||
class AudioCropProcessUTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("AUDIO",),
|
||||
"gain_db": ("FLOAT", {"default": 0, "min": -100, "max": 100}),
|
||||
"offset_seconds": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}),
|
||||
"duration_seconds": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}),
|
||||
"resample_to_hz": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}),
|
||||
"make_stereo": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("AUDIO", "INT", "INT", "FLOAT")
|
||||
RETURN_NAMES = ("audio", "sample_rate", "channels", "duration")
|
||||
FUNCTION = "execute"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, *args):
|
||||
return args
|
||||
|
||||
def execute(
|
||||
self,
|
||||
audio,
|
||||
gain_db: float,
|
||||
offset_seconds: float,
|
||||
duration_seconds: float,
|
||||
resample_to_hz: float,
|
||||
make_stereo: bool,
|
||||
):
|
||||
waveform = audio["waveform"] # [B, C, N]
|
||||
sample_rate = int(audio["sample_rate"])
|
||||
# 裁剪offset和duration
|
||||
start = int(offset_seconds * sample_rate)
|
||||
end = int(start + duration_seconds * sample_rate) if duration_seconds > 0 else waveform.shape[2]
|
||||
waveform = waveform[:, :, start:end]
|
||||
# 重采样
|
||||
if resample_to_hz > 0 and int(resample_to_hz) != sample_rate:
|
||||
import torchaudio
|
||||
waveform = torchaudio.functional.resample(waveform, sample_rate, int(resample_to_hz))
|
||||
sample_rate = int(resample_to_hz)
|
||||
# 增益
|
||||
if gain_db != 0.0:
|
||||
gain_scalar = 10 ** (gain_db / 20)
|
||||
waveform = waveform * gain_scalar
|
||||
# 强制立体声
|
||||
if make_stereo and waveform.shape[1] == 1:
|
||||
waveform = torch.cat([waveform, waveform], dim=1)
|
||||
elif make_stereo and waveform.shape[1] != 2:
|
||||
raise ValueError(f"Input audio has {waveform.shape[1]} channels, cannot convert to stereo (2 channels)")
|
||||
channels = int(waveform.shape[1])
|
||||
duration_val = float(waveform.shape[2] / sample_rate) if sample_rate > 0 else 0.0
|
||||
return ({"sample_rate": sample_rate, "waveform": waveform}, sample_rate, channels, duration_val)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AudioCropProcessUTK": AudioCropProcessUTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AudioCropProcessUTK": "Audio Crop Process (UTK)",
|
||||
}
|
||||
@@ -1,3 +1,13 @@
|
||||
"""
|
||||
Load Audio Node
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
Loads audio files from path with various processing options.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
@@ -80,57 +90,11 @@ class LoadAudioPlusFromPath_UTK:
|
||||
duration_val = float(mix.shape[2] / sample_rate) if sample_rate > 0 else 0.0
|
||||
return ({"sample_rate": sample_rate, "waveform": mix}, sample_rate, channels, duration_val)
|
||||
|
||||
class AudioCropProcessUTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("AUDIO",),
|
||||
"gain_db": ("FLOAT", {"default": 0, "min": -100, "max": 100}),
|
||||
"offset_seconds": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}),
|
||||
"duration_seconds": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}),
|
||||
"resample_to_hz": ("FLOAT", {"default": 0, "min": 0, "max": FLOAT_MAX}),
|
||||
"make_stereo": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("AUDIO", "INT", "INT", "FLOAT")
|
||||
RETURN_NAMES = ("audio", "sample_rate", "channels", "duration")
|
||||
FUNCTION = "execute"
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"LoadAudioPlusFromPath_UTK": LoadAudioPlusFromPath_UTK,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, *args):
|
||||
return args
|
||||
|
||||
def execute(
|
||||
self,
|
||||
audio,
|
||||
gain_db: float,
|
||||
offset_seconds: float,
|
||||
duration_seconds: float,
|
||||
resample_to_hz: float,
|
||||
make_stereo: bool,
|
||||
):
|
||||
waveform = audio["waveform"] # [B, C, N]
|
||||
sample_rate = int(audio["sample_rate"])
|
||||
# 裁剪offset和duration
|
||||
start = int(offset_seconds * sample_rate)
|
||||
end = int(start + duration_seconds * sample_rate) if duration_seconds > 0 else waveform.shape[2]
|
||||
waveform = waveform[:, :, start:end]
|
||||
# 重采样
|
||||
if resample_to_hz > 0 and int(resample_to_hz) != sample_rate:
|
||||
import torchaudio
|
||||
waveform = torchaudio.functional.resample(waveform, sample_rate, int(resample_to_hz))
|
||||
sample_rate = int(resample_to_hz)
|
||||
# 增益
|
||||
if gain_db != 0.0:
|
||||
gain_scalar = 10 ** (gain_db / 20)
|
||||
waveform = waveform * gain_scalar
|
||||
# 强制立体声
|
||||
if make_stereo and waveform.shape[1] == 1:
|
||||
waveform = torch.cat([waveform, waveform], dim=1)
|
||||
elif make_stereo and waveform.shape[1] != 2:
|
||||
raise ValueError(f"Input audio has {waveform.shape[1]} channels, cannot convert to stereo (2 channels)")
|
||||
channels = int(waveform.shape[1])
|
||||
duration_val = float(waveform.shape[2] / sample_rate) if sample_rate > 0 else 0.0
|
||||
return ({"sample_rate": sample_rate, "waveform": waveform}, sample_rate, channels, duration_val)
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"LoadAudioPlusFromPath_UTK": "Load Audio Plus From Path (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,294 @@
|
||||
"""
|
||||
ComfyUI Universal Toolkit - Common Utilities
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Common utility functions and constants used across all nodes.
|
||||
|
||||
: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:
|
||||
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
|
||||
@@ -0,0 +1,9 @@
|
||||
"""
|
||||
ComfyUI Universal Toolkit - Image Nodes
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Image processing nodes for ComfyUI Universal Toolkit.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
@@ -0,0 +1,62 @@
|
||||
"""
|
||||
Check Mask Node
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
Checks if a mask is valid based on white area percentage.
|
||||
|
||||
: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
|
||||
|
||||
def mask_white_area(mask, white_point):
|
||||
"""Calculate the percentage of white area in mask"""
|
||||
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"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK",), #
|
||||
"white_point": ("INT", {"default": 1, "min": 1, "max": 254, "step": 1}), # 用于判断mask是否有效的白点值，高于此值被计入有效
|
||||
"area_percent": ("INT", {"default": 1, "min": 1, "max": 99, "step": 1}), # 区域百分比，低于此则mask判定无效
|
||||
},
|
||||
"optional": { #
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("BOOLEAN",)
|
||||
RETURN_NAMES = ('bool',)
|
||||
FUNCTION = 'check_mask'
|
||||
|
||||
def check_mask(self, mask, white_point, area_percent,):
|
||||
|
||||
if mask.dim() == 2:
|
||||
mask = torch.unsqueeze(mask, 0)
|
||||
mask = tensor2pil(mask[0])
|
||||
if mask.width * mask.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
|
||||
log(f"CheckMask_UTK:{ret}", message_type="finish")
|
||||
return (ret,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CheckMask_UTK": CheckMask_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CheckMask_UTK": "Check Mask (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,149 @@
|
||||
"""
|
||||
Crop By Mask Node
|
||||
~~~~~~~~~~~~~~~~
|
||||
|
||||
Crops images based on mask detection with various detection modes.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from PIL import Image, ImageDraw
|
||||
import numpy as np
|
||||
from ...common_utils import log, tensor2pil, pil2tensor, image2mask
|
||||
|
||||
def mask2image(mask):
|
||||
"""Convert mask tensor to PIL image"""
|
||||
return tensor2pil(mask).convert('L')
|
||||
|
||||
def gaussian_blur(image, radius):
|
||||
"""Apply Gaussian blur to image"""
|
||||
return image.filter(ImageFilter.GaussianBlur(radius=radius))
|
||||
|
||||
def min_bounding_rect(mask):
|
||||
"""Find minimum bounding rectangle of mask"""
|
||||
mask_array = np.array(mask)
|
||||
coords = np.where(mask_array > 0)
|
||||
if len(coords[0]) == 0:
|
||||
return (0, 0, mask.width, mask.height)
|
||||
|
||||
y_min, y_max = coords[0].min(), coords[0].max()
|
||||
x_min, x_max = coords[1].min(), coords[1].max()
|
||||
|
||||
return (x_min, y_min, x_max - x_min, y_max - y_min)
|
||||
|
||||
def max_inscribed_rect(mask):
|
||||
"""Find maximum inscribed rectangle of mask"""
|
||||
# Simplified implementation - returns bounding rect
|
||||
return min_bounding_rect(mask)
|
||||
|
||||
def mask_area(mask):
|
||||
"""Find area of mask"""
|
||||
# Simplified implementation - returns bounding rect
|
||||
return min_bounding_rect(mask)
|
||||
|
||||
def num_round_up_to_multiple(num, multiple):
|
||||
"""Round up to the nearest multiple"""
|
||||
return ((num + multiple - 1) // multiple) * multiple
|
||||
|
||||
def draw_rect(image, x, y, width, height, line_color="#FF0000", line_width=2):
|
||||
"""Draw rectangle on image"""
|
||||
draw = ImageDraw.Draw(image)
|
||||
draw.rectangle([x, y, x + width, y + height], outline=line_color, width=line_width)
|
||||
return image
|
||||
|
||||
class CropByMask_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
detect_mode = ['min_bounding_rect', 'max_inscribed_rect', 'mask_area']
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", ), #
|
||||
"mask_for_crop": ("MASK",),
|
||||
"invert_mask": ("BOOLEAN", {"default": False}), # 反转mask#
|
||||
"detect": (detect_mode,),
|
||||
"top_reserve": ("INT", {"default": 20, "min": -9999, "max": 9999, "step": 1}),
|
||||
"bottom_reserve": ("INT", {"default": 20, "min": -9999, "max": 9999, "step": 1}),
|
||||
"left_reserve": ("INT", {"default": 20, "min": -9999, "max": 9999, "step": 1}),
|
||||
"right_reserve": ("INT", {"default": 20, "min": -9999, "max": 9999, "step": 1}),
|
||||
},
|
||||
"optional": {
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "BOX", "IMAGE",)
|
||||
RETURN_NAMES = ("croped_image", "croped_mask", "crop_box", "box_preview")
|
||||
FUNCTION = 'crop_by_mask'
|
||||
|
||||
def crop_by_mask(self, image, mask_for_crop, invert_mask, detect,
|
||||
top_reserve, bottom_reserve, left_reserve, right_reserve
|
||||
):
|
||||
|
||||
ret_images = []
|
||||
ret_masks = []
|
||||
l_images = []
|
||||
l_masks = []
|
||||
|
||||
for l in image:
|
||||
l_images.append(torch.unsqueeze(l, 0))
|
||||
if mask_for_crop.dim() == 2:
|
||||
mask_for_crop = torch.unsqueeze(mask_for_crop, 0)
|
||||
# 如果有多张mask输入，使用第一张
|
||||
if mask_for_crop.shape[0] > 1:
|
||||
log(f"Warning: Multiple mask inputs, using the first.", message_type='warning')
|
||||
mask_for_crop = torch.unsqueeze(mask_for_crop[0], 0)
|
||||
if invert_mask:
|
||||
mask_for_crop = 1 - mask_for_crop
|
||||
l_masks.append(tensor2pil(torch.unsqueeze(mask_for_crop, 0)).convert('L'))
|
||||
|
||||
_mask = mask2image(mask_for_crop)
|
||||
try:
|
||||
from PIL import ImageFilter
|
||||
bluredmask = gaussian_blur(_mask, 20).convert('L')
|
||||
except ImportError:
|
||||
bluredmask = _mask.convert('L')
|
||||
|
||||
x = 0
|
||||
y = 0
|
||||
width = 0
|
||||
height = 0
|
||||
if detect == "min_bounding_rect":
|
||||
(x, y, width, height) = min_bounding_rect(bluredmask)
|
||||
elif detect == "max_inscribed_rect":
|
||||
(x, y, width, height) = max_inscribed_rect(bluredmask)
|
||||
else:
|
||||
(x, y, width, height) = mask_area(_mask)
|
||||
|
||||
width = num_round_up_to_multiple(width, 8)
|
||||
height = num_round_up_to_multiple(height, 8)
|
||||
log(f"CropByMask_UTK: Box detected. x={x},y={y},width={width},height={height}")
|
||||
canvas_width, canvas_height = tensor2pil(torch.unsqueeze(image[0], 0)).convert('RGB').size
|
||||
x1 = x - left_reserve if x - left_reserve > 0 else 0
|
||||
y1 = y - top_reserve if y - top_reserve > 0 else 0
|
||||
x2 = x + width + right_reserve if x + width + right_reserve < canvas_width else canvas_width
|
||||
y2 = y + height + bottom_reserve if y + height + bottom_reserve < canvas_height else canvas_height
|
||||
preview_image = tensor2pil(mask_for_crop).convert('RGB')
|
||||
preview_image = draw_rect(preview_image, x, y, width, height, line_color="#F00000", line_width=(width+height)//100)
|
||||
preview_image = draw_rect(preview_image, x1, y1, x2 - x1, y2 - y1,
|
||||
line_color="#00F000", line_width=(width+height)//200)
|
||||
crop_box = (x1, y1, x2, y2)
|
||||
for i in range(len(l_images)):
|
||||
_canvas = tensor2pil(l_images[i]).convert('RGB')
|
||||
_mask = l_masks[0]
|
||||
ret_images.append(pil2tensor(_canvas.crop(crop_box)))
|
||||
ret_masks.append(image2mask(_mask.crop(crop_box)))
|
||||
|
||||
log(f"CropByMask_UTK Processed {len(ret_images)} image(s).", message_type='finish')
|
||||
return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0), list(crop_box), pil2tensor(preview_image),)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CropByMask_UTK": CropByMask_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CropByMask_UTK": "Crop By Mask (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
"""
|
||||
Depth Map Blur Node
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Applies depth-based blur effects to images using depth maps.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
import folder_paths
|
||||
|
||||
class DepthMapBlur_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"depth_map": ("IMAGE",),
|
||||
"blur_strength": ("FLOAT", {
|
||||
"default": 64.0,
|
||||
"min": 0.0,
|
||||
"max": 256.0,
|
||||
"step": 1.0
|
||||
}),
|
||||
"focal_depth": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}),
|
||||
"focus_spread": ("FLOAT", {
|
||||
"default": 1,
|
||||
"min": 1.0,
|
||||
"max": 8.0,
|
||||
"step": 0.1
|
||||
}),
|
||||
"steps": ("INT", {
|
||||
"default": 5,
|
||||
"min": 1,
|
||||
"max": 32,
|
||||
}),
|
||||
"focal_range": ("FLOAT", {
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.01
|
||||
}),
|
||||
"mask_blur": ("INT", {
|
||||
"default": 1,
|
||||
"min": 1,
|
||||
"max": 127,
|
||||
"step": 2
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE","MASK")
|
||||
RETURN_NAMES = ()
|
||||
FUNCTION = "depthblur_image"
|
||||
DESCRIPTION = """
|
||||
blur_strength: Represents the blur strength. This parameter controls the overall intensity of the blur effect; the higher the value, the more blurred the image becomes.
|
||||
|
||||
focal_depth: Represents the focal depth. This parameter is used to determine which depth level in the image should remain sharp, while other levels are blurred based on depth differences.
|
||||
|
||||
focus_spread: Represents the focus spread range. This parameter controls the size of the blur transition area near the focal depth; the larger the value, the wider the transition area, and the smoother the blur effect spreads around the focus.
|
||||
|
||||
steps: Represents the number of steps in the blur process. This parameter determines the calculation precision of the blur effect; the more steps, the finer the blur effect, but this also increases the computational load.
|
||||
|
||||
focal_range: Represents the focal range. This parameter is used to adjust the depth range within the focal depth that remains sharp; the larger the value, the wider the area around the focal depth that remains sharp.
|
||||
|
||||
mask_blur: Represents the mask blur strength for blurring the depth map. This parameter controls the intensity of the depth map's blur treatment, used for preprocessing the depth map before calculating the final blur effect, to achieve a more natural blur transition.
|
||||
"""
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
def depthblur_image(self, image: torch.Tensor, depth_map: torch.Tensor, blur_strength: float, focal_depth: float, focus_spread:float, steps: int, focal_range: float, mask_blur: int):
|
||||
batch_size, height, width, _ = image.shape
|
||||
image_result = torch.zeros_like(image)
|
||||
mask_result = torch.zeros((batch_size, height, width), dtype=torch.float32)
|
||||
|
||||
for b in range(batch_size):
|
||||
tensor_image = image[b].numpy()
|
||||
tensor_image_depth = depth_map[b].numpy()
|
||||
|
||||
# Apply blur
|
||||
blur_image,depth_mask = self.apply_depthblur(tensor_image, tensor_image_depth, blur_strength, focal_depth, focus_spread, steps, focal_range, mask_blur)
|
||||
|
||||
tensor_image = torch.from_numpy(blur_image).unsqueeze(0)
|
||||
tensor_mask = torch.from_numpy(depth_mask).unsqueeze(0)
|
||||
|
||||
image_result[b] = tensor_image
|
||||
mask_result[b] = tensor_mask
|
||||
|
||||
return (image_result,mask_result)
|
||||
|
||||
def apply_depthblur(self, image, depth_map, blur_strength, focal_depth, focus_spread, steps, focal_range, mask_blur):
|
||||
def make_odd(x):
|
||||
x = int(round(x))
|
||||
return x if x % 2 == 1 else x + 1
|
||||
|
||||
# Normalize the input image if needed
|
||||
needs_normalization = image.max() > 1
|
||||
if needs_normalization:
|
||||
image = image.astype(np.float32) / 255
|
||||
|
||||
# Normalize the depth map if needed
|
||||
depth_map = depth_map.astype(np.float32) / 255 if depth_map.max() > 1 else depth_map
|
||||
|
||||
# Resize depth map to match the image dimensions
|
||||
depth_map_resized = cv2.resize(depth_map, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_LINEAR)
|
||||
if len(depth_map_resized.shape) > 2:
|
||||
depth_map_resized = cv2.cvtColor(depth_map_resized, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# Adjust the depth map based on the focal plane
|
||||
depth_mask = np.abs(depth_map_resized - focal_depth)
|
||||
depth_mask = np.clip(depth_mask / np.max(depth_mask), 0, 1)
|
||||
|
||||
# Process the depth_mask
|
||||
depth_mask[depth_mask < focal_range] = 0
|
||||
depth_mask[depth_mask >= focal_range] = (depth_mask[depth_mask >= focal_range] - focal_range) / (1 - focal_range)
|
||||
|
||||
# Apply mask blur
|
||||
mask_blur = max(1, make_odd(mask_blur))
|
||||
depth_mask = cv2.GaussianBlur(depth_mask, (mask_blur, mask_blur), 0)
|
||||
|
||||
# Generate blurred versions of the image
|
||||
blur_ksize = max(1, make_odd(blur_strength))
|
||||
blurred_images = [cv2.GaussianBlur(image, (blur_ksize, blur_ksize), sigmaX=0) for _ in range(steps)]
|
||||
|
||||
# Use the adjusted depth map as a mask for applying blurred images
|
||||
# 这里简单实现：直接用最重的模糊图和原图按mask混合
|
||||
final_image = image * (1 - depth_mask[..., None]) + blurred_images[-1] * depth_mask[..., None]
|
||||
|
||||
# Convert back to original range if the image was normalized
|
||||
if needs_normalization:
|
||||
final_image = np.clip(final_image * 255, 0, 255).astype(np.uint8)
|
||||
|
||||
return final_image, depth_mask
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DepthMapBlur_UTK": DepthMapBlur_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DepthMapBlur_UTK": "Depth Map Blur (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
"""
|
||||
Empty Unit Generator Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Generates empty images, masks, and latents with various preset ratios and configurations.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
import re
|
||||
from PIL import Image
|
||||
from ..common_utils import log
|
||||
|
||||
class EmptyUnitGenerator_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
ratio_options = [
|
||||
"custom",
|
||||
"SD1.5 - 1:1 square 512x512",
|
||||
"SD1.5 - 2:3 portrait 512x768",
|
||||
"SD1.5 - 3:4 portrait 512x682",
|
||||
"SD1.5 - 3:2 landscape 768x512",
|
||||
"SD1.5 - 4:3 landscape 682x512",
|
||||
"SD1.5 - 16:9 cinema 910x512",
|
||||
"SD1.5 - 1.85:1 cinema 952x512",
|
||||
"SD1.5 - 2:1 cinema 1024x512",
|
||||
"SDXL - 1:1 square 1024x1024",
|
||||
"SDXL - 3:4 portrait 896x1152",
|
||||
"SDXL - 5:8 portrait 832x1216",
|
||||
"SDXL - 9:16 portrait 768x1344",
|
||||
"SDXL - 9:21 portrait 640x1536",
|
||||
"SDXL - 4:3 landscape 1152x896",
|
||||
"SDXL - 3:2 landscape 1216x832",
|
||||
"SDXL - 16:9 landscape 1344x768",
|
||||
"SDXL - 21:9 landscape 1536x640",
|
||||
]
|
||||
latent_type_options = ["standard", "sd3", "hunyuan", "ltx"]
|
||||
return {
|
||||
"required": {
|
||||
"width": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 8, "label": "Width (custom only)"}),
|
||||
"height": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 8, "label": "Height (custom only)"}),
|
||||
"ratio": (ratio_options, {"default": ratio_options[9], "label": "Resolution/Ratio"}),
|
||||
"scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 8.0, "step": 0.1, "label": "Scale (放大倍数)"}),
|
||||
"divisor": ("INT", {"default": 8, "min": 1, "max": 512, "step": 1, "label": "Divisor (整除裁切)"}),
|
||||
"image_color": (["white", "black", "gray", "red", "green", "blue"], {"default": "white"}),
|
||||
"batch": ("INT", {"default": 1, "min": 1, "max": 16, "label": "Batch 数量"}),
|
||||
"latent_type": (latent_type_options, {"default": "standard", "label": "Latent类型"}),
|
||||
},
|
||||
"optional": {},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "LATENT", "INT", "INT")
|
||||
RETURN_NAMES = ("image", "mask", "latent", "width", "height")
|
||||
FUNCTION = "generate"
|
||||
|
||||
def generate(self, width, height, ratio, scale, divisor, image_color, batch, latent_type):
|
||||
if ratio == "custom":
|
||||
w = width
|
||||
h = height
|
||||
else:
|
||||
m = re.search(r"(\d+)x(\d+)", ratio)
|
||||
if m:
|
||||
w, h = int(m.group(1)), int(m.group(2))
|
||||
else:
|
||||
w, h = 1024, 1024
|
||||
w = max(1, int(round(w * scale)))
|
||||
h = max(1, int(round(h * scale)))
|
||||
if divisor > 1:
|
||||
w = (w // divisor) * divisor
|
||||
h = (h // divisor) * divisor
|
||||
COLOR_OPTIONS = {
|
||||
"white": (255, 255, 255),
|
||||
"black": (0, 0, 0),
|
||||
"gray": (128, 128, 128),
|
||||
"red": (255, 0, 0),
|
||||
"green": (0, 255, 0),
|
||||
"blue": (0, 0, 255),
|
||||
}
|
||||
color_rgb = COLOR_OPTIONS[image_color]
|
||||
images = []
|
||||
for _ in range(batch):
|
||||
img = torch.from_numpy(np.array(Image.new("RGB", (w, h), color_rgb))).float() / 255.0
|
||||
img = img.permute(2, 0, 1)
|
||||
images.append(img)
|
||||
images = torch.stack(images, dim=0).permute(0, 2, 3, 1)
|
||||
mask_value = color_rgb[0] / 255.0
|
||||
masks = torch.ones([batch, 1, h, w], dtype=torch.float32) * mask_value
|
||||
latent_channels = {
|
||||
"standard": 4,
|
||||
"sd3": 8,
|
||||
"hunyuan": 8,
|
||||
"ltx": 16,
|
||||
}.get(latent_type, 4)
|
||||
latent = {
|
||||
"samples": torch.zeros([batch, latent_channels, h // 8, w // 8], dtype=torch.float32),
|
||||
"batch_index_list": None
|
||||
}
|
||||
return images, masks, latent, w, h
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"EmptyUnitGenerator_UTK": EmptyUnitGenerator_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"EmptyUnitGenerator_UTK": "Empty Unit Generator",
|
||||
}
|
||||
@@ -0,0 +1,82 @@
|
||||
"""
|
||||
Image And Mask Preview Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Preview an image or a mask, when both inputs are used composites the mask on top of the image.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import random
|
||||
from nodes import SaveImage, ImageCompositeMasked
|
||||
import folder_paths
|
||||
|
||||
class ImageAndMaskPreview_UTK(SaveImage):
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_temp_directory()
|
||||
self.type = "temp"
|
||||
self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
|
||||
self.compress_level = 4
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
colors = ["red", "green", "blue", "yellow", "cyan", "magenta", "white", "black"]
|
||||
return {
|
||||
"required": {
|
||||
"mask_opacity": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"mask_color": (colors, {"default": "red"}),
|
||||
"pass_through": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",),
|
||||
"mask": ("MASK",),
|
||||
},
|
||||
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
||||
}
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("composite",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "UniversalToolkit"
|
||||
DESCRIPTION = """
|
||||
Preview an image or a mask, when both inputs are used
|
||||
composites the mask on top of the image.
|
||||
with pass_through on the preview is disabled and the
|
||||
composite is returned from the composite slot instead,
|
||||
this allows for the preview to be passed for video combine
|
||||
nodes for example.
|
||||
"""
|
||||
|
||||
def execute(self, mask_opacity, mask_color, pass_through, filename_prefix="ComfyUI", image=None, mask=None, prompt=None, extra_pnginfo=None):
|
||||
if mask is not None and image is None:
|
||||
preview = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
|
||||
elif mask is None and image is not None:
|
||||
preview = image
|
||||
elif mask is not None and image is not None:
|
||||
mask_adjusted = mask * mask_opacity
|
||||
mask_image = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3).clone()
|
||||
|
||||
color_map = {
|
||||
"red": [255, 0, 0], "green": [0, 255, 0], "blue": [0, 0, 255],
|
||||
"yellow": [255, 255, 0], "cyan": [0, 255, 255], "magenta": [255, 0, 255],
|
||||
"white": [255, 255, 255], "black": [0, 0, 0]
|
||||
}
|
||||
color_list = color_map.get(mask_color, [255, 0, 0])
|
||||
|
||||
mask_image[:, :, :, 0] = color_list[0] / 255 # Red channel
|
||||
mask_image[:, :, :, 1] = color_list[1] / 255 # Green channel
|
||||
mask_image[:, :, :, 2] = color_list[2] / 255 # Blue channel
|
||||
|
||||
preview, = ImageCompositeMasked.composite(self, image, mask_image, 0, 0, True, mask_adjusted)
|
||||
if pass_through:
|
||||
return (preview, )
|
||||
return(self.save_images(preview, filename_prefix, prompt, extra_pnginfo))
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageAndMaskPreview_UTK": ImageAndMaskPreview_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageAndMaskPreview_UTK": "Image And Mask Preview (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,83 @@
|
||||
"""
|
||||
Image Combine Alpha Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Combines RGB image with mask to create RGBA image.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from ...common_utils import log, tensor2pil, pil2tensor
|
||||
|
||||
def image_channel_split(image, mode):
|
||||
"""Split image into channels"""
|
||||
if mode == 'RGB':
|
||||
return image.split()
|
||||
elif mode == 'RGBA':
|
||||
return image.split()
|
||||
else:
|
||||
return image.split()
|
||||
|
||||
def image_channel_merge(channels, mode):
|
||||
"""Merge channels into image"""
|
||||
if mode == 'RGBA':
|
||||
return Image.merge('RGBA', channels)
|
||||
elif mode == 'RGB':
|
||||
return Image.merge('RGB', channels)
|
||||
else:
|
||||
return Image.merge(mode, channels)
|
||||
|
||||
class ImageCombineAlpha_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"RGB_image": ("IMAGE", ), #
|
||||
"mask": ("MASK",), #
|
||||
},
|
||||
"optional": {
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("RGBA_image",)
|
||||
FUNCTION = 'image_combine_alpha'
|
||||
|
||||
def image_combine_alpha(self, RGB_image, mask):
|
||||
|
||||
ret_images = []
|
||||
input_images = []
|
||||
input_masks = []
|
||||
|
||||
for i in RGB_image:
|
||||
input_images.append(torch.unsqueeze(i, 0))
|
||||
if mask.dim() == 2:
|
||||
mask = torch.unsqueeze(mask, 0)
|
||||
for m in mask:
|
||||
input_masks.append(torch.unsqueeze(m, 0))
|
||||
|
||||
max_batch = max(len(input_images), len(input_masks))
|
||||
for i in range(max_batch):
|
||||
_image = input_images[i] if i < len(input_images) else input_images[-1]
|
||||
_mask = input_masks[i] if i < len(input_masks) else input_masks[-1]
|
||||
r, g, b, _ = image_channel_split(tensor2pil(_image).convert('RGB'), 'RGB')
|
||||
ret_image = image_channel_merge((r, g, b, tensor2pil(_mask).convert('L')), 'RGBA')
|
||||
|
||||
ret_images.append(pil2tensor(ret_image))
|
||||
|
||||
log(f"ImageCombineAlpha_UTK Processed {len(ret_images)} image(s).", message_type='finish')
|
||||
return (torch.cat(ret_images, dim=0),)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageCombineAlpha_UTK": ImageCombineAlpha_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageCombineAlpha_UTK": "Image Combine Alpha (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
"""
|
||||
Image Concatenate Node
|
||||
~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Concatenates two images side by side or vertically with various options.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from ..common_utils import log
|
||||
|
||||
class ImageConcatenate_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image1": ("IMAGE",),
|
||||
"image2": ("IMAGE",),
|
||||
"direction": (
|
||||
[ 'right',
|
||||
'down',
|
||||
'left',
|
||||
'up',
|
||||
'auto',
|
||||
],
|
||||
{
|
||||
"default": 'auto'
|
||||
}),
|
||||
"match_image_size": ("BOOLEAN", {"default": True}),
|
||||
"max_size": ("INT", {"default": 4096, "min": 64, "max": 8192, "step": 64}),
|
||||
"background_color": (["black", "white", "gray", "transparent"], {"default": "black"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "concatenate"
|
||||
|
||||
def concatenate(self, image1, image2, direction, match_image_size, max_size, background_color):
|
||||
# Check if the batch sizes are different
|
||||
batch_size1 = image1.shape[0]
|
||||
batch_size2 = image2.shape[0]
|
||||
|
||||
# 保证batch维度一致，补齐到最大batch
|
||||
if batch_size1 != batch_size2:
|
||||
max_batch_size = max(batch_size1, batch_size2)
|
||||
if batch_size1 < max_batch_size:
|
||||
last_image1 = image1[-1].unsqueeze(0).repeat(max_batch_size - batch_size1, 1, 1, 1)
|
||||
image1 = torch.cat([image1, last_image1], dim=0)
|
||||
if batch_size2 < max_batch_size:
|
||||
last_image2 = image2[-1].unsqueeze(0).repeat(max_batch_size - batch_size2, 1, 1, 1)
|
||||
image2 = torch.cat([image2, last_image2], dim=0)
|
||||
|
||||
# Get original dimensions
|
||||
h1, w1 = image1.shape[1:3]
|
||||
h2, w2 = image2.shape[1:3]
|
||||
|
||||
# If direction is auto, determine the best direction based on image dimensions
|
||||
if direction == 'auto':
|
||||
horizontal_ratio = (w1 + w2) / max(h1, h2)
|
||||
vertical_ratio = max(w1, w2) / (h1 + h2)
|
||||
direction = 'right' if abs(horizontal_ratio - 1) <= abs(vertical_ratio - 1) else 'down'
|
||||
|
||||
# Match image sizes if requested
|
||||
if match_image_size:
|
||||
if direction in ['right', 'left', 'auto']:
|
||||
target_height = max(h1, h2)
|
||||
if h1 < target_height:
|
||||
scale = target_height / h1
|
||||
new_width = int(w1 * scale)
|
||||
image1 = torch.nn.functional.interpolate(
|
||||
image1.permute(0, 3, 1, 2),
|
||||
size=(target_height, new_width),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).permute(0, 2, 3, 1)
|
||||
if h2 < target_height:
|
||||
scale = target_height / h2
|
||||
new_width = int(w2 * scale)
|
||||
image2 = torch.nn.functional.interpolate(
|
||||
image2.permute(0, 3, 1, 2),
|
||||
size=(target_height, new_width),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).permute(0, 2, 3, 1)
|
||||
else: # up, down
|
||||
target_width = max(w1, w2)
|
||||
if w1 < target_width:
|
||||
scale = target_width / w1
|
||||
new_height = int(h1 * scale)
|
||||
image1 = torch.nn.functional.interpolate(
|
||||
image1.permute(0, 3, 1, 2),
|
||||
size=(new_height, target_width),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).permute(0, 2, 3, 1)
|
||||
if w2 < target_width:
|
||||
scale = target_width / w2
|
||||
new_height = int(h2 * scale)
|
||||
image2 = torch.nn.functional.interpolate(
|
||||
image2.permute(0, 3, 1, 2),
|
||||
size=(new_height, target_width),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).permute(0, 2, 3, 1)
|
||||
|
||||
# Update dimensions after scaling
|
||||
h1, w1 = image1.shape[1:3]
|
||||
h2, w2 = image2.shape[1:3]
|
||||
|
||||
# Calculate final dimensions
|
||||
if direction in ['right', 'left']:
|
||||
final_height = max(h1, h2)
|
||||
final_width = w1 + w2
|
||||
else: # up, down
|
||||
final_height = h1 + h2
|
||||
final_width = max(w1, w2)
|
||||
|
||||
# Check if we need to scale down
|
||||
if max(final_height, final_width) > max_size:
|
||||
scale = max_size / max(final_height, final_width)
|
||||
new_h1 = int(h1 * scale)
|
||||
new_w1 = int(w1 * scale)
|
||||
new_h2 = int(h2 * scale)
|
||||
new_w2 = int(w2 * scale)
|
||||
image1 = torch.nn.functional.interpolate(
|
||||
image1.permute(0, 3, 1, 2),
|
||||
size=(new_h1, new_w1),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).permute(0, 2, 3, 1)
|
||||
image2 = torch.nn.functional.interpolate(
|
||||
image2.permute(0, 3, 1, 2),
|
||||
size=(new_h2, new_w2),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).permute(0, 2, 3, 1)
|
||||
h1, w1 = image1.shape[1:3]
|
||||
h2, w2 = image2.shape[1:3]
|
||||
if direction in ['right', 'left']:
|
||||
final_height = max(h1, h2)
|
||||
final_width = w1 + w2
|
||||
else:
|
||||
final_height = h1 + h2
|
||||
final_width = max(w1, w2)
|
||||
|
||||
# Ensure both images have the same number of channels
|
||||
channels_image1 = image1.shape[-1]
|
||||
channels_image2 = image2.shape[-1]
|
||||
if channels_image1 != channels_image2:
|
||||
if channels_image1 < channels_image2:
|
||||
alpha_channel = torch.ones((*image1.shape[:-1], channels_image2 - channels_image1), device=image1.device)
|
||||
image1 = torch.cat((image1, alpha_channel), dim=-1)
|
||||
else:
|
||||
alpha_channel = torch.ones((*image2.shape[:-1], channels_image1 - channels_image2), device=image2.device)
|
||||
image2 = torch.cat((image2, alpha_channel), dim=-1)
|
||||
|
||||
# 创建输出张量，batch维度与输入一致
|
||||
batch_size = image1.shape[0]
|
||||
if background_color == "transparent":
|
||||
output = torch.zeros((batch_size, final_height, final_width, image1.shape[-1]), dtype=image1.dtype, device=image1.device)
|
||||
else:
|
||||
color_value = 1.0 if background_color == "white" else 0.0 if background_color == "black" else 0.5
|
||||
output = torch.full((batch_size, final_height, final_width, image1.shape[-1]), color_value, dtype=image1.dtype, device=image1.device)
|
||||
|
||||
# 计算放置位置
|
||||
if direction == 'right':
|
||||
x1 = 0
|
||||
x2 = w1
|
||||
y1 = (final_height - h1) // 2
|
||||
y2 = (final_height - h2) // 2
|
||||
elif direction == 'left':
|
||||
x1 = w2
|
||||
x2 = 0
|
||||
y1 = (final_height - h1) // 2
|
||||
y2 = (final_height - h2) // 2
|
||||
elif direction == 'down':
|
||||
x1 = (final_width - w1) // 2
|
||||
x2 = (final_width - w2) // 2
|
||||
y1 = 0
|
||||
y2 = h1
|
||||
else: # up
|
||||
x1 = (final_width - w1) // 2
|
||||
x2 = (final_width - w2) // 2
|
||||
y1 = h2
|
||||
y2 = 0
|
||||
|
||||
# 批量放置图片
|
||||
output[:, y1:y1+h1, x1:x1+w1] = image1
|
||||
output[:, y2:y2+h2, x2:x2+w2] = image2
|
||||
|
||||
return (output,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageConcatenate_UTK": ImageConcatenate_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageConcatenate_UTK": "Image Concatenate",
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
"""
|
||||
Image Concatenate Multi Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Concatenates multiple images in various directions and layouts.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
class ImageConcatenateMulti_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"direction": (
|
||||
[ 'right',
|
||||
'down',
|
||||
'left',
|
||||
'up',
|
||||
'auto',
|
||||
],
|
||||
{
|
||||
"default": 'auto'
|
||||
}),
|
||||
"match_image_size": ("BOOLEAN", {"default": True}),
|
||||
"max_size": ("INT", {"default": 4096, "min": 64, "max": 8192, "step": 64}),
|
||||
"background_color": (["black", "white", "gray", "transparent"], {"default": "black"}),
|
||||
"grid_size": (["auto", "1x1", "2x2", "3x3", "4x4"], {"default": "auto"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "concatenate_multi"
|
||||
|
||||
def concatenate_multi(self, images, direction, match_image_size, max_size, background_color, grid_size):
|
||||
if len(images.shape) != 4:
|
||||
raise ValueError("输入必须是4D张量 [batch, height, width, channels]")
|
||||
|
||||
batch_size = images.shape[0]
|
||||
|
||||
# 处理网格布局
|
||||
if grid_size != "auto":
|
||||
rows, cols = map(int, grid_size.split("x"))
|
||||
if batch_size > rows * cols:
|
||||
raise ValueError(f"图像数量 ({batch_size}) 超过网格大小 ({grid_size})")
|
||||
# 填充到网格大小
|
||||
if batch_size < rows * cols:
|
||||
padding = torch.zeros((rows * cols - batch_size, *images.shape[1:]), dtype=images.dtype, device=images.device)
|
||||
images = torch.cat([images, padding], dim=0)
|
||||
batch_size = rows * cols
|
||||
else:
|
||||
# 自动计算网格大小
|
||||
if direction in ['right', 'left', 'auto']:
|
||||
rows = 1
|
||||
cols = batch_size
|
||||
else: # up, down
|
||||
rows = batch_size
|
||||
cols = 1
|
||||
|
||||
# 获取所有图像的尺寸
|
||||
heights = []
|
||||
widths = []
|
||||
for i in range(batch_size):
|
||||
h, w = images[i].shape[:2]
|
||||
heights.append(h)
|
||||
widths.append(w)
|
||||
|
||||
# 如果方向是auto，计算最佳方向
|
||||
if direction == 'auto':
|
||||
# 计算水平和垂直拼接的宽高比
|
||||
total_width = sum(widths)
|
||||
max_height = max(heights)
|
||||
horizontal_ratio = total_width / max_height
|
||||
|
||||
total_height = sum(heights)
|
||||
max_width = max(widths)
|
||||
vertical_ratio = max_width / total_height
|
||||
|
||||
# 选择更接近1:1的方向
|
||||
direction = 'right' if abs(horizontal_ratio - 1) <= abs(vertical_ratio - 1) else 'down'
|
||||
|
||||
# 如果需要匹配图像尺寸
|
||||
if match_image_size:
|
||||
if direction in ['right', 'left', 'auto']:
|
||||
# 匹配高度
|
||||
target_height = max(heights)
|
||||
for i in range(batch_size):
|
||||
if heights[i] < target_height:
|
||||
scale = target_height / heights[i]
|
||||
new_width = int(widths[i] * scale)
|
||||
images[i] = torch.nn.functional.interpolate(
|
||||
images[i].unsqueeze(0).permute(0, 3, 1, 2),
|
||||
size=(target_height, new_width),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).permute(0, 2, 3, 1).squeeze(0)
|
||||
else: # up, down
|
||||
# 匹配宽度
|
||||
target_width = max(widths)
|
||||
for i in range(batch_size):
|
||||
if widths[i] < target_width:
|
||||
scale = target_width / widths[i]
|
||||
new_height = int(heights[i] * scale)
|
||||
images[i] = torch.nn.functional.interpolate(
|
||||
images[i].unsqueeze(0).permute(0, 3, 1, 2),
|
||||
size=(new_height, target_width),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).permute(0, 2, 3, 1).squeeze(0)
|
||||
|
||||
# 更新尺寸
|
||||
heights = []
|
||||
widths = []
|
||||
for i in range(batch_size):
|
||||
h, w = images[i].shape[:2]
|
||||
heights.append(h)
|
||||
widths.append(w)
|
||||
|
||||
# 计算最终输出尺寸
|
||||
if direction in ['right', 'left', 'auto']:
|
||||
final_height = max(heights)
|
||||
final_width = sum(widths)
|
||||
else: # up, down
|
||||
final_height = sum(heights)
|
||||
final_width = max(widths)
|
||||
|
||||
# 检查是否需要缩放
|
||||
if max(final_height, final_width) > max_size:
|
||||
scale = max_size / max(final_height, final_width)
|
||||
for i in range(batch_size):
|
||||
new_height = int(heights[i] * scale)
|
||||
new_width = int(widths[i] * scale)
|
||||
images[i] = torch.nn.functional.interpolate(
|
||||
images[i].unsqueeze(0).permute(0, 3, 1, 2),
|
||||
size=(new_height, new_width),
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
).permute(0, 2, 3, 1).squeeze(0)
|
||||
|
||||
# 更新最终尺寸
|
||||
heights = []
|
||||
widths = []
|
||||
for i in range(batch_size):
|
||||
h, w = images[i].shape[:2]
|
||||
heights.append(h)
|
||||
widths.append(w)
|
||||
|
||||
if direction in ['right', 'left', 'auto']:
|
||||
final_height = max(heights)
|
||||
final_width = sum(widths)
|
||||
else: # up, down
|
||||
final_height = sum(heights)
|
||||
final_width = max(widths)
|
||||
|
||||
# 创建输出张量
|
||||
if background_color == "transparent":
|
||||
output = torch.zeros((1, final_height, final_width, images.shape[-1]), dtype=images.dtype, device=images.device)
|
||||
else:
|
||||
color_value = 1.0 if background_color == "white" else 0.0 if background_color == "black" else 0.5
|
||||
output = torch.full((1, final_height, final_width, images.shape[-1]), color_value, dtype=images.dtype, device=images.device)
|
||||
|
||||
# 放置图像
|
||||
if direction in ['right', 'left', 'auto']:
|
||||
x_offset = 0
|
||||
for i in range(batch_size):
|
||||
h, w = images[i].shape[:2]
|
||||
y_offset = (final_height - h) // 2
|
||||
if direction == 'left':
|
||||
x_offset = final_width - sum(widths[i:])
|
||||
output[0, y_offset:y_offset+h, x_offset:x_offset+w] = images[i]
|
||||
if direction != 'left':
|
||||
x_offset += w
|
||||
else: # up, down
|
||||
y_offset = 0
|
||||
for i in range(batch_size):
|
||||
h, w = images[i].shape[:2]
|
||||
x_offset = (final_width - w) // 2
|
||||
if direction == 'up':
|
||||
y_offset = final_height - sum(heights[i:])
|
||||
output[0, y_offset:y_offset+h, x_offset:x_offset+w] = images[i]
|
||||
if direction != 'up':
|
||||
y_offset += h
|
||||
|
||||
return (output,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageConcatenateMulti_UTK": ImageConcatenateMulti_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageConcatenateMulti_UTK": "Image Concatenate Multi (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
"""
|
||||
Image Mask Scale As Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Scales images and masks to match the dimensions of a reference image.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from ...common_utils import log, tensor2pil, pil2tensor, image2mask
|
||||
|
||||
def fit_resize_image(image, target_width, target_height, fit_mode, resize_sampler, background_color="#000000"):
|
||||
"""Resize image according to fit mode"""
|
||||
if fit_mode == 'letterbox':
|
||||
# Calculate scaling factor to fit within target dimensions
|
||||
scale = min(target_width / image.width, target_height / image.height)
|
||||
new_width = int(image.width * scale)
|
||||
new_height = int(image.height * scale)
|
||||
|
||||
# Resize image
|
||||
resized = image.resize((new_width, new_height), resize_sampler)
|
||||
|
||||
# Create new image with target dimensions and paste resized image
|
||||
if image.mode == 'RGB':
|
||||
result = Image.new('RGB', (target_width, target_height), background_color)
|
||||
else:
|
||||
result = Image.new('L', (target_width, target_height), 0)
|
||||
paste_x = (target_width - new_width) // 2
|
||||
paste_y = (target_height - new_height) // 2
|
||||
result.paste(resized, (paste_x, paste_y))
|
||||
return result
|
||||
|
||||
elif fit_mode == 'crop':
|
||||
# Calculate scaling factor to cover target dimensions
|
||||
scale = max(target_width / image.width, target_height / image.height)
|
||||
new_width = int(image.width * scale)
|
||||
new_height = int(image.height * scale)
|
||||
|
||||
# Resize image
|
||||
resized = image.resize((new_width, new_height), resize_sampler)
|
||||
|
||||
# Crop to target dimensions
|
||||
crop_x = (new_width - target_width) // 2
|
||||
crop_y = (new_height - target_height) // 2
|
||||
return resized.crop((crop_x, crop_y, crop_x + target_width, crop_y + target_height))
|
||||
|
||||
else: # fill
|
||||
# Simple resize to target dimensions
|
||||
return image.resize((target_width, target_height), resize_sampler)
|
||||
|
||||
class ImageMaskScaleAs_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
fit_mode = ['letterbox', 'crop', 'fill']
|
||||
method_mode = ['lanczos', 'bicubic', 'hamming', 'bilinear', 'box', 'nearest']
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"scale_as": ("IMAGE",),
|
||||
"fit": (fit_mode,),
|
||||
"method": (method_mode,),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",), #
|
||||
"mask": ("MASK",), #
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "BOX", "INT", "INT")
|
||||
RETURN_NAMES = ("image", "mask", "original_size", "width", "height",)
|
||||
FUNCTION = 'image_mask_scale_as'
|
||||
|
||||
def image_mask_scale_as(self, scale_as, fit, method,
|
||||
image=None, mask = None,
|
||||
):
|
||||
if scale_as.shape[0] > 0:
|
||||
_asimage = tensor2pil(scale_as[0])
|
||||
else:
|
||||
_asimage = tensor2pil(scale_as)
|
||||
target_width, target_height = _asimage.size
|
||||
_mask = Image.new('L', size=_asimage.size, color='black')
|
||||
_image = Image.new('RGB', size=_asimage.size, color='black')
|
||||
orig_width = 4
|
||||
orig_height = 4
|
||||
resize_sampler = Image.LANCZOS
|
||||
if method == "bicubic":
|
||||
resize_sampler = Image.BICUBIC
|
||||
elif method == "hamming":
|
||||
resize_sampler = Image.HAMMING
|
||||
elif method == "bilinear":
|
||||
resize_sampler = Image.BILINEAR
|
||||
elif method == "box":
|
||||
resize_sampler = Image.BOX
|
||||
elif method == "nearest":
|
||||
resize_sampler = Image.NEAREST
|
||||
|
||||
ret_images = []
|
||||
ret_masks = []
|
||||
|
||||
if image is not None:
|
||||
for i in image:
|
||||
i = torch.unsqueeze(i, 0)
|
||||
_image = tensor2pil(i).convert('RGB')
|
||||
orig_width, orig_height = _image.size
|
||||
_image = fit_resize_image(_image, target_width, target_height, fit, resize_sampler)
|
||||
ret_images.append(pil2tensor(_image))
|
||||
if mask is not None:
|
||||
if mask.dim() == 2:
|
||||
mask = torch.unsqueeze(mask, 0)
|
||||
for m in mask:
|
||||
m = torch.unsqueeze(m, 0)
|
||||
_mask = tensor2pil(m).convert('L')
|
||||
orig_width, orig_height = _mask.size
|
||||
_mask = fit_resize_image(_mask, target_width, target_height, fit, resize_sampler).convert('L')
|
||||
ret_masks.append(image2mask(_mask))
|
||||
if len(ret_images) > 0 and len(ret_masks) >0:
|
||||
log(f"ImageMaskScaleAs_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,)
|
||||
elif len(ret_images) > 0 and len(ret_masks) == 0:
|
||||
log(f"ImageMaskScaleAs_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,)
|
||||
elif len(ret_images) == 0 and len(ret_masks) > 0:
|
||||
log(f"ImageMaskScaleAs_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,)
|
||||
else:
|
||||
log(f"Error: ImageMaskScaleAs_UTK skipped, because the available image or mask is not found.", message_type='error')
|
||||
return (None, None, [orig_width, orig_height], 0, 0,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageMaskScaleAs_UTK": ImageMaskScaleAs_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageMaskScaleAs_UTK": "Image Mask Scale As (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
"""
|
||||
Image Pad For Outpaint Masked Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Expands image with padding and generates corresponding mask for outpainting.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
MAX_RESOLUTION = 8192
|
||||
|
||||
class ImagePadForOutpaintMasked_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
color_options = ["gray", "white", "black", "red", "green", "blue", "yellow", "cyan", "magenta"]
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"data_mode": (["pixel", "percent"], {"default": "pixel"}),
|
||||
"left": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"top": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"right": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"bottom": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1}),
|
||||
"feathering": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
|
||||
"background_color": (color_options, {"default": "gray"}),
|
||||
},
|
||||
"optional": {
|
||||
"mask": ("MASK",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK")
|
||||
FUNCTION = "expand_image"
|
||||
|
||||
def expand_image(self, image, data_mode, left, top, right, bottom, feathering, background_color, mask=None):
|
||||
B, H, W, C = image.size()
|
||||
# 处理 pad 参数
|
||||
if data_mode == "percent":
|
||||
left = int(W * left / 100)
|
||||
right = int(W * right / 100)
|
||||
top = int(H * top / 100)
|
||||
bottom = int(H * bottom / 100)
|
||||
# 背景色映射
|
||||
color_map = {
|
||||
"gray": [0.5, 0.5, 0.5],
|
||||
"white": [1.0, 1.0, 1.0],
|
||||
"black": [0.0, 0.0, 0.0],
|
||||
"red": [1.0, 0.0, 0.0],
|
||||
"green": [0.0, 1.0, 0.0],
|
||||
"blue": [0.0, 0.0, 1.0],
|
||||
"yellow": [1.0, 1.0, 0.0],
|
||||
"cyan": [0.0, 1.0, 1.0],
|
||||
"magenta": [1.0, 0.0, 1.0],
|
||||
}
|
||||
bg_rgb = color_map.get(background_color, [0.5, 0.5, 0.5])
|
||||
# 新图像
|
||||
new_image = torch.ones((B, H + top + bottom, W + left + right, C), dtype=torch.float32)
|
||||
for i in range(C):
|
||||
new_image[:, :, :, i] = bg_rgb[i]
|
||||
new_image[:, top:top + H, left:left + W, :] = image
|
||||
# 掩码逻辑与原实现一致
|
||||
if mask is not None:
|
||||
if torch.allclose(mask, torch.zeros_like(mask)):
|
||||
print("Warning: The incoming mask is fully black. Handling it as None.")
|
||||
mask = None
|
||||
if mask is None:
|
||||
new_mask = torch.ones((B, H + top + bottom, W + left + right), dtype=torch.float32)
|
||||
t = torch.zeros((B, H, W), dtype=torch.float32)
|
||||
else:
|
||||
mask = F.pad(mask, (left, right, top, bottom), mode='constant', value=0)
|
||||
mask = 1 - mask
|
||||
t = torch.zeros_like(mask)
|
||||
if feathering > 0 and feathering * 2 < H and feathering * 2 < W:
|
||||
for i in range(H):
|
||||
for j in range(W):
|
||||
dt = i if top != 0 else H
|
||||
db = H - i if bottom != 0 else H
|
||||
dl = j if left != 0 else W
|
||||
dr = W - j if right != 0 else W
|
||||
d = min(dt, db, dl, dr)
|
||||
if d >= feathering:
|
||||
continue
|
||||
v = (feathering - d) / feathering
|
||||
if mask is None:
|
||||
t[:, i, j] = v * v
|
||||
else:
|
||||
t[:, top + i, left + j] = v * v
|
||||
if mask is None:
|
||||
new_mask[:, top:top + H, left:left + W] = t
|
||||
return (new_image, new_mask,)
|
||||
else:
|
||||
return (new_image, mask,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImagePadForOutpaintMasked_UTK": ImagePadForOutpaintMasked_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImagePadForOutpaintMasked_UTK": "Image Pad For Outpaint Masked (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
"""
|
||||
Image Ratio Detector Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Detects and analyzes image aspect ratios and dimensions.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import math
|
||||
from ..common_utils import log
|
||||
|
||||
class ImageRatioDetector_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"image": ("IMAGE",)}}
|
||||
|
||||
RETURN_TYPES = ("STRING", "INT", "INT", "STRING")
|
||||
RETURN_NAMES = ("ratio_str", "width", "height", "approx_ratio_str")
|
||||
FUNCTION = "detect"
|
||||
|
||||
def detect(self, image):
|
||||
if hasattr(image, 'dim') and image.dim() == 4:
|
||||
img = image[0]
|
||||
else:
|
||||
img = image
|
||||
shape = img.shape
|
||||
if len(shape) == 3:
|
||||
if shape[0] <= 4:
|
||||
_, h, w = shape
|
||||
else:
|
||||
h, w, _ = shape
|
||||
elif len(shape) == 2:
|
||||
h, w = shape
|
||||
else:
|
||||
return "?", 0, 0, "N/A"
|
||||
h = int(h)
|
||||
w = int(w)
|
||||
if w == 0 or h == 0:
|
||||
ratio_str = "0:0"
|
||||
approx_ratio_str = "N/A"
|
||||
return ratio_str, w, h, approx_ratio_str
|
||||
gcd = math.gcd(w, h)
|
||||
ratio_str = f"{w//gcd}:{h//gcd}"
|
||||
std_ratios = {
|
||||
"1:1": 1.0,
|
||||
"16:9": 16/9,
|
||||
"4:3": 4/3,
|
||||
"3:2": 3/2,
|
||||
"2:3": 2/3,
|
||||
"3:4": 3/4,
|
||||
"9:16": 9/16,
|
||||
"5:4": 5/4,
|
||||
"7:5": 7/5,
|
||||
"21:9": 21/9,
|
||||
"5:3": 5/3,
|
||||
"3:1": 3/1,
|
||||
"1:2": 1/2,
|
||||
"2:1": 2/1,
|
||||
"1:1.85": 1/1.85,
|
||||
"1:2.35": 1/2.35,
|
||||
}
|
||||
wh_ratio = float(w) / float(h)
|
||||
approx_ratio_str = min(std_ratios.keys(), key=lambda k: abs(std_ratios[k] - wh_ratio))
|
||||
return ratio_str, w, h, approx_ratio_str
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageRatioDetector_UTK": ImageRatioDetector_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageRatioDetector_UTK": "Image Ratio Detector",
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
Image Remove Alpha Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Removes alpha channel from RGBA images with optional background filling.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from ...common_utils import log, tensor2pil, pil2tensor
|
||||
|
||||
class ImageRemoveAlpha_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"RGBA_image": ("IMAGE", ), #
|
||||
"fill_background": ("BOOLEAN", {"default": False}),
|
||||
"background_color": ("STRING", {"default": "#000000"}),
|
||||
},
|
||||
"optional": {
|
||||
"mask": ("MASK",), #
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", )
|
||||
RETURN_NAMES = ("RGB_image", )
|
||||
FUNCTION = 'image_remove_alpha'
|
||||
|
||||
def image_remove_alpha(self, RGBA_image, fill_background, background_color, mask=None):
|
||||
|
||||
ret_images = []
|
||||
|
||||
for index, img in enumerate(RGBA_image):
|
||||
_image = tensor2pil(img)
|
||||
|
||||
if fill_background:
|
||||
if mask is not None:
|
||||
m = mask[index].unsqueeze(0) if index < len(mask) else mask[-1].unsqueeze(0)
|
||||
alpha = tensor2pil(m).convert('L')
|
||||
elif _image.mode == "RGBA":
|
||||
alpha = _image.split()[-1]
|
||||
else:
|
||||
log(f"Error: ImageRemoveAlpha_UTK skipped, because the input image is not RGBA and mask is None.",
|
||||
message_type='error')
|
||||
return (RGBA_image,)
|
||||
ret_image = Image.new('RGB', size=_image.size, color=background_color)
|
||||
ret_image.paste(_image, mask=alpha)
|
||||
ret_images.append(pil2tensor(ret_image))
|
||||
|
||||
else:
|
||||
ret_images.append(pil2tensor(tensor2pil(img).convert('RGB')))
|
||||
|
||||
log(f"ImageRemoveAlpha_UTK Processed {len(ret_images)} image(s).", message_type='finish')
|
||||
return (torch.cat(ret_images, dim=0), )
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageRemoveAlpha_UTK": ImageRemoveAlpha_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageRemoveAlpha_UTK": "Image Remove Alpha (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
"""
|
||||
Image Scale By Aspect Ratio Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Scales images to specific aspect ratios with various fitting modes.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
import math
|
||||
from ...common_utils import log, tensor2pil, pil2tensor, image2mask
|
||||
|
||||
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):
|
||||
"""Resize image according to fit mode"""
|
||||
if fit_mode == 'letterbox':
|
||||
# Calculate scaling factor to fit within target dimensions
|
||||
scale = min(target_width / image.width, target_height / image.height)
|
||||
new_width = int(image.width * scale)
|
||||
new_height = int(image.height * scale)
|
||||
|
||||
# Resize image
|
||||
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))
|
||||
paste_x = (target_width - new_width) // 2
|
||||
paste_y = (target_height - new_height) // 2
|
||||
result.paste(resized, (paste_x, paste_y))
|
||||
return result
|
||||
|
||||
elif fit_mode == 'crop':
|
||||
# Calculate scaling factor to cover target dimensions
|
||||
scale = max(target_width / image.width, target_height / image.height)
|
||||
new_width = int(image.width * scale)
|
||||
new_height = int(image.height * scale)
|
||||
|
||||
# Resize image
|
||||
resized = image.resize((new_width, new_height), resize_sampler)
|
||||
|
||||
# Crop to target dimensions
|
||||
crop_x = (new_width - target_width) // 2
|
||||
crop_y = (new_height - target_height) // 2
|
||||
return resized.crop((crop_x, crop_y, crop_x + target_width, crop_y + target_height))
|
||||
|
||||
else: # fill
|
||||
# Simple resize to target dimensions
|
||||
return image.resize((target_width, target_height), resize_sampler)
|
||||
|
||||
class ImageScaleByAspectRatio_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
ratio_list = ['original', 'custom', '1:1', '3:2', '4:3', '16:9', '2:3', '3:4', '9:16']
|
||||
fit_mode = ['letterbox', 'crop', 'fill']
|
||||
method_mode = ['lanczos', 'bicubic', 'hamming', 'bilinear', 'box', 'nearest']
|
||||
multiple_list = ['8', '16', '32', '64', '128', '256', '512', 'None']
|
||||
|
||||
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}),
|
||||
"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}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",), #
|
||||
"mask": ("MASK",), #
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "BOX", "INT", "INT",)
|
||||
RETURN_NAMES = ("image", "mask", "original_size", "width", "height",)
|
||||
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,
|
||||
):
|
||||
orig_images = []
|
||||
orig_masks = []
|
||||
orig_width = 0
|
||||
orig_height = 0
|
||||
target_width = 0
|
||||
target_height = 0
|
||||
ratio = 1.0
|
||||
ret_images = []
|
||||
ret_masks = []
|
||||
if image is not None:
|
||||
for i in image:
|
||||
i = torch.unsqueeze(i, 0)
|
||||
orig_images.append(i)
|
||||
orig_width, orig_height = tensor2pil(orig_images[0]).size
|
||||
if mask is not None:
|
||||
if mask.dim() == 2:
|
||||
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 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')
|
||||
return (None, None, None, 0, 0,)
|
||||
|
||||
if aspect_ratio == 'original':
|
||||
ratio = orig_width / orig_height
|
||||
elif aspect_ratio == 'custom':
|
||||
ratio = proportional_width / proportional_height
|
||||
else:
|
||||
s = aspect_ratio.split(":")
|
||||
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
|
||||
else:
|
||||
target_width = orig_width
|
||||
target_height = int(target_width / ratio)
|
||||
else:
|
||||
if scale_to_longest_side:
|
||||
target_height = longest_side
|
||||
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)
|
||||
|
||||
if round_to_multiple != 'None':
|
||||
multiple = int(round_to_multiple)
|
||||
target_width = num_round_up_to_multiple(target_width, multiple)
|
||||
target_height = num_round_up_to_multiple(target_height, multiple)
|
||||
|
||||
_mask = Image.new('L', size=(target_width, target_height), color='black')
|
||||
_image = Image.new('RGB', size=(target_width, target_height), color='black')
|
||||
|
||||
resize_sampler = Image.LANCZOS
|
||||
if method == "bicubic":
|
||||
resize_sampler = Image.BICUBIC
|
||||
elif method == "hamming":
|
||||
resize_sampler = Image.HAMMING
|
||||
elif method == "bilinear":
|
||||
resize_sampler = Image.BILINEAR
|
||||
elif method == "box":
|
||||
resize_sampler = Image.BOX
|
||||
elif method == "nearest":
|
||||
resize_sampler = Image.NEAREST
|
||||
|
||||
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)
|
||||
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')
|
||||
ret_masks.append(image2mask(_mask))
|
||||
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,)
|
||||
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,)
|
||||
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,)
|
||||
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,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageScaleByAspectRatio_UTK": ImageScaleByAspectRatio_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageScaleByAspectRatio_UTK": "Image Scale By Aspect Ratio (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
"""
|
||||
Image Scale Restore Node
|
||||
~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Restores images to original size or scales them with specified parameters.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from ...common_utils import log, tensor2pil, pil2tensor, image2mask
|
||||
|
||||
class ImageScaleRestore_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
method_mode = ['lanczos', 'bicubic', 'hamming', 'bilinear', 'box', 'nearest']
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", ), #
|
||||
"scale": ("FLOAT", {"default": 1, "min": 0.01, "max": 100, "step": 0.01}),
|
||||
"method": (method_mode,),
|
||||
"scale_by_longest_side": ("BOOLEAN", {"default": False}), # 是否按长边缩放
|
||||
"longest_side": ("INT", {"default": 1024, "min": 4, "max": 999999, "step": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"mask": ("MASK",), #
|
||||
"original_size": ("BOX",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "BOX", "INT", "INT")
|
||||
RETURN_NAMES = ("image", "mask", "original_size", "width", "height",)
|
||||
FUNCTION = 'image_scale_restore'
|
||||
|
||||
def image_scale_restore(self, image, scale, method,
|
||||
scale_by_longest_side, longest_side,
|
||||
mask = None, original_size = None
|
||||
):
|
||||
|
||||
l_images = []
|
||||
l_masks = []
|
||||
ret_images = []
|
||||
ret_masks = []
|
||||
for l in image:
|
||||
l_images.append(torch.unsqueeze(l, 0))
|
||||
m = tensor2pil(l)
|
||||
if m.mode == 'RGBA':
|
||||
l_masks.append(m.split()[-1])
|
||||
|
||||
if mask is not None:
|
||||
if mask.dim() == 2:
|
||||
mask = torch.unsqueeze(mask, 0)
|
||||
l_masks = []
|
||||
for m in mask:
|
||||
l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L'))
|
||||
|
||||
max_batch = max(len(l_images), len(l_masks))
|
||||
|
||||
orig_width, orig_height = tensor2pil(l_images[0]).size
|
||||
if original_size is not None:
|
||||
target_width = original_size[0]
|
||||
target_height = original_size[1]
|
||||
else:
|
||||
target_width = int(orig_width * scale)
|
||||
target_height = int(orig_height * scale)
|
||||
if scale_by_longest_side:
|
||||
if orig_width > orig_height:
|
||||
target_width = longest_side
|
||||
target_height = int(target_width * orig_height / orig_width)
|
||||
else:
|
||||
target_height = longest_side
|
||||
target_width = int(target_height * orig_width / orig_height)
|
||||
if target_width < 4:
|
||||
target_width = 4
|
||||
if target_height < 4:
|
||||
target_height = 4
|
||||
resize_sampler = Image.LANCZOS
|
||||
if method == "bicubic":
|
||||
resize_sampler = Image.BICUBIC
|
||||
elif method == "hamming":
|
||||
resize_sampler = Image.HAMMING
|
||||
elif method == "bilinear":
|
||||
resize_sampler = Image.BILINEAR
|
||||
elif method == "box":
|
||||
resize_sampler = Image.BOX
|
||||
elif method == "nearest":
|
||||
resize_sampler = Image.NEAREST
|
||||
|
||||
for i in range(max_batch):
|
||||
|
||||
_image = l_images[i] if i < len(l_images) else l_images[-1]
|
||||
|
||||
_canvas = tensor2pil(_image).convert('RGB')
|
||||
ret_image = _canvas.resize((target_width, target_height), resize_sampler)
|
||||
ret_mask = Image.new('L', size=ret_image.size, color='white')
|
||||
if mask is not None:
|
||||
_mask = l_masks[i] if i < len(l_masks) else l_masks[-1]
|
||||
ret_mask = _mask.resize((target_width, target_height), resize_sampler)
|
||||
|
||||
ret_images.append(pil2tensor(ret_image))
|
||||
ret_masks.append(image2mask(ret_mask))
|
||||
|
||||
log(f"ImageScaleRestore_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,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImageScaleRestore_UTK": ImageScaleRestore_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageScaleRestore_UTK": "Image Scale Restore (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
"""
|
||||
Imitation Hue Node
|
||||
~~~~~~~~~~~~~~~~~
|
||||
|
||||
Performs color transfer and imitation between images with skin protection.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
class ImitationHueNode_UTK:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"imitation_image": ("IMAGE",),
|
||||
"target_image": ("IMAGE",),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 1.0, "step": 0.1}),
|
||||
"skin_protection": ("FLOAT", {"default": 0.2, "min": 0, "max": 1.0, "step": 0.1}),
|
||||
"auto_brightness": ("BOOLEAN", {"default": True}),
|
||||
"brightness_range": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.1}),
|
||||
"auto_contrast": ("BOOLEAN", {"default": False}),
|
||||
"contrast_range": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.1}),
|
||||
"auto_saturation": ("BOOLEAN", {"default": False}),
|
||||
"saturation_range": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.1}),
|
||||
"auto_tone": ("BOOLEAN", {"default": False}),
|
||||
"tone_strength": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.1}),
|
||||
},
|
||||
"optional": {
|
||||
"mask": ("MASK", {"default": None}),
|
||||
},
|
||||
}
|
||||
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "imitation_hue"
|
||||
|
||||
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
|
||||
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,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ImitationHueNode_UTK": ImitationHueNode_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImitationHueNode_UTK": "Imitation Hue Node (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
"""
|
||||
Purge VRAM Node
|
||||
~~~~~~~~~~~~~~
|
||||
|
||||
Purges GPU memory and optionally unloads models.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch.cuda
|
||||
import gc
|
||||
from ...common_utils import log
|
||||
|
||||
def clear_memory():
|
||||
"""Clear GPU memory"""
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
gc.collect()
|
||||
|
||||
class PurgeVRAM_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"anything": ("*", {}),
|
||||
"purge_cache": ("BOOLEAN", {"default": True}),
|
||||
"purge_models": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
"optional": {
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("*",)
|
||||
RETURN_NAMES = ("anything",)
|
||||
FUNCTION = "purge_vram"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
def purge_vram(self, anything, purge_cache, purge_models):
|
||||
clear_memory()
|
||||
if purge_models:
|
||||
try:
|
||||
import comfy.model_management
|
||||
comfy.model_management.unload_all_models()
|
||||
comfy.model_management.soft_empty_cache()
|
||||
except ImportError:
|
||||
log("ComfyUI model management not available", message_type="warning")
|
||||
log("VRAM purged successfully", message_type="finish")
|
||||
return (anything,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"PurgeVRAM_UTK": PurgeVRAM_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PurgeVRAM_UTK": "Purge VRAM (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
"""
|
||||
Restore Crop Box Node
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Restores cropped images back to their original background.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from ...common_utils import log, tensor2pil, pil2tensor, image2mask
|
||||
|
||||
class RestoreCropBox_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"background_image": ("IMAGE", ),
|
||||
"croped_image": ("IMAGE",),
|
||||
"invert_mask": ("BOOLEAN", {"default": False}), # 反转mask#
|
||||
"crop_box": ("BOX",),
|
||||
},
|
||||
"optional": {
|
||||
"croped_mask": ("MASK",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK", )
|
||||
RETURN_NAMES = ("image", "mask", )
|
||||
FUNCTION = 'restore_crop_box'
|
||||
|
||||
def restore_crop_box(self, background_image, croped_image, invert_mask, crop_box,
|
||||
croped_mask=None
|
||||
):
|
||||
|
||||
b_images = []
|
||||
l_images = []
|
||||
l_masks = []
|
||||
ret_images = []
|
||||
ret_masks = []
|
||||
for b in background_image:
|
||||
b_images.append(torch.unsqueeze(b, 0))
|
||||
for l in croped_image:
|
||||
l_images.append(torch.unsqueeze(l, 0))
|
||||
m = tensor2pil(l)
|
||||
if m.mode == 'RGBA':
|
||||
l_masks.append(m.split()[-1])
|
||||
else:
|
||||
l_masks.append(Image.new('L', size=m.size, color='white'))
|
||||
if croped_mask is not None:
|
||||
if croped_mask.dim() == 2:
|
||||
croped_mask = torch.unsqueeze(croped_mask, 0)
|
||||
l_masks = []
|
||||
for m in croped_mask:
|
||||
if invert_mask:
|
||||
m = 1 - m
|
||||
l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L'))
|
||||
|
||||
max_batch = max(len(b_images), len(l_images), len(l_masks))
|
||||
for i in range(max_batch):
|
||||
background_image = b_images[i] if i < len(b_images) else b_images[-1]
|
||||
croped_image = l_images[i] if i < len(l_images) else l_images[-1]
|
||||
_mask = l_masks[i] if i < len(l_masks) else l_masks[-1]
|
||||
|
||||
_canvas = tensor2pil(background_image).convert('RGB')
|
||||
_layer = tensor2pil(croped_image).convert('RGB')
|
||||
|
||||
ret_mask = Image.new('L', size=_canvas.size, color='black')
|
||||
_canvas.paste(_layer, box=tuple(crop_box), mask=_mask)
|
||||
ret_mask.paste(_mask, box=tuple(crop_box))
|
||||
ret_images.append(pil2tensor(_canvas))
|
||||
ret_masks.append(image2mask(ret_mask))
|
||||
|
||||
log(f"RestoreCropBox_UTK Processed {len(ret_images)} image(s).", message_type='finish')
|
||||
return (torch.cat(ret_images, dim=0), torch.cat(ret_masks, dim=0),)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"RestoreCropBox_UTK": RestoreCropBox_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"RestoreCropBox_UTK": "Restore Crop Box (UTK)",
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,7 +1,17 @@
|
||||
"""
|
||||
Mask Operations Node
|
||||
~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Performs logical operations on masks.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
class MaskAnd_UTK:
|
||||
CATEGORY = "UniversalToolkit/Mask"
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"mask1": ("MASK",), "mask2": ("MASK",)}}
|
||||
@@ -15,7 +25,7 @@ class MaskAnd_UTK:
|
||||
return (mask1 * mask2,)
|
||||
|
||||
class MaskSub_UTK:
|
||||
CATEGORY = "UniversalToolkit/Mask"
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"mask1": ("MASK",), "mask2": ("MASK",)}}
|
||||
@@ -29,7 +39,7 @@ class MaskSub_UTK:
|
||||
return (torch.clamp(mask1 - mask2, 0, 1),)
|
||||
|
||||
class MaskAdd_UTK:
|
||||
CATEGORY = "UniversalToolkit/Mask"
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"mask1": ("MASK",), "mask2": ("MASK",)}}
|
||||
@@ -40,4 +50,17 @@ class MaskAdd_UTK:
|
||||
# 逐像素加法，支持batch，结果裁剪到[0,1]
|
||||
if mask1.shape != mask2.shape:
|
||||
raise ValueError("输入的两个MASK尺寸不一致")
|
||||
return (torch.clamp(mask1 + mask2, 0, 1),)
|
||||
return (torch.clamp(mask1 + mask2, 0, 1),)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"MaskAnd_UTK": MaskAnd_UTK,
|
||||
"MaskSub_UTK": MaskSub_UTK,
|
||||
"MaskAdd_UTK": MaskAdd_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"MaskAnd_UTK": "Mask And (UTK)",
|
||||
"MaskSub_UTK": "Mask Sub (UTK)",
|
||||
"MaskAdd_UTK": "Mask Add (UTK)",
|
||||
}
|
||||
@@ -1,106 +1,18 @@
|
||||
"""
|
||||
Fill Masked Area Node
|
||||
~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
Fills masked areas in images using various algorithms.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
import cv2
|
||||
from scipy.ndimage import binary_erosion, gaussian_filter
|
||||
|
||||
class Show_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"input": ("STRING", "INT", "FLOAT", "LIST", "MASK", "IMAGE", "LATENT")}}
|
||||
RETURN_TYPES = ("STRING", "INT", "FLOAT", "LIST", "MASK", "IMAGE", "LATENT")
|
||||
RETURN_NAMES = ("string", "int", "float", "list", "mask", "image", "latent")
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
|
||||
def show(self, input):
|
||||
outs = [None] * 7
|
||||
if isinstance(input, str):
|
||||
outs[0] = input
|
||||
elif isinstance(input, int):
|
||||
outs[1] = input
|
||||
elif isinstance(input, float):
|
||||
outs[2] = input
|
||||
elif isinstance(input, list):
|
||||
outs[3] = input
|
||||
elif hasattr(input, "shape") and len(input.shape) == 4 and input.shape[1] == 1:
|
||||
outs[4] = input # MASK
|
||||
elif hasattr(input, "shape") and len(input.shape) == 4 and input.shape[1] == 3:
|
||||
outs[5] = input # IMAGE
|
||||
elif isinstance(input, dict) and "samples" in input:
|
||||
outs[6] = input # LATENT
|
||||
return tuple(outs)
|
||||
|
||||
class ShowInt_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"int_val": ("INT",)}}
|
||||
RETURN_TYPES = ("INT",)
|
||||
RETURN_NAMES = ("int_val",)
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
def show(self, int_val=None):
|
||||
if int_val is None:
|
||||
int_val = 0
|
||||
return (int_val,)
|
||||
|
||||
class ShowFloat_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"float_val": ("FLOAT",)}}
|
||||
RETURN_TYPES = ("FLOAT",)
|
||||
RETURN_NAMES = ("float_val",)
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
def show(self, float_val=None):
|
||||
if float_val is None:
|
||||
float_val = 0.0
|
||||
return (float_val,)
|
||||
|
||||
class ShowList_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"list_val": ("LIST",)}}
|
||||
RETURN_TYPES = ("LIST",)
|
||||
RETURN_NAMES = ("list_val",)
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
def show(self, list_val=None):
|
||||
if list_val is None:
|
||||
list_val = []
|
||||
return (list_val,)
|
||||
|
||||
class ShowText_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"text": ("STRING",)}}
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("text",)
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
def show(self, text=None):
|
||||
if text is None:
|
||||
text = ""
|
||||
return (text,)
|
||||
|
||||
class PreviewMask_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"mask": ("MASK",)}}
|
||||
RETURN_TYPES = ("MASK",)
|
||||
RETURN_NAMES = ("mask",)
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
def show(self, mask=None):
|
||||
if mask is None:
|
||||
mask = torch.zeros([1, 1, 64, 64], dtype=torch.float32)
|
||||
return (mask,)
|
||||
|
||||
# mask二值化，阈值0.5
|
||||
def mask_floor(mask):
|
||||
return (mask > 0.5).astype(np.float32)
|
||||
@@ -201,5 +113,11 @@ class FillMaskedArea_UTK:
|
||||
out = result * (1 - alpha[..., None]) + filled * alpha[..., None]
|
||||
return np.clip(out, 0, 1)
|
||||
|
||||
class ImageAndMaskPreview_UTK:
|
||||
pass # This class is now implemented in image_nodes_utk.py
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FillMaskedArea_UTK": FillMaskedArea_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FillMaskedArea_UTK": "Fill Masked Area (UTK)",
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
"""
|
||||
Show Nodes
|
||||
~~~~~~~~~
|
||||
|
||||
Display and preview nodes for various data types.
|
||||
|
||||
:copyright: (c) 2024 by May
|
||||
:license: MIT, see LICENSE for more details.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
class Show_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"input": ("STRING", "INT", "FLOAT", "LIST", "MASK", "IMAGE", "LATENT")}}
|
||||
RETURN_TYPES = ("STRING", "INT", "FLOAT", "LIST", "MASK", "IMAGE", "LATENT")
|
||||
RETURN_NAMES = ("string", "int", "float", "list", "mask", "image", "latent")
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
|
||||
def show(self, input):
|
||||
outs = [None] * 7
|
||||
if isinstance(input, str):
|
||||
outs[0] = input
|
||||
elif isinstance(input, int):
|
||||
outs[1] = input
|
||||
elif isinstance(input, float):
|
||||
outs[2] = input
|
||||
elif isinstance(input, list):
|
||||
outs[3] = input
|
||||
elif hasattr(input, "shape") and len(input.shape) == 4 and input.shape[1] == 1:
|
||||
outs[4] = input # MASK
|
||||
elif hasattr(input, "shape") and len(input.shape) == 4 and input.shape[1] == 3:
|
||||
outs[5] = input # IMAGE
|
||||
elif isinstance(input, dict) and "samples" in input:
|
||||
outs[6] = input # LATENT
|
||||
return tuple(outs)
|
||||
|
||||
class ShowInt_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"int_val": ("INT",)}}
|
||||
RETURN_TYPES = ("INT",)
|
||||
RETURN_NAMES = ("int_val",)
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
def show(self, int_val=None):
|
||||
if int_val is None:
|
||||
int_val = 0
|
||||
return (int_val,)
|
||||
|
||||
class ShowFloat_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"float_val": ("FLOAT",)}}
|
||||
RETURN_TYPES = ("FLOAT",)
|
||||
RETURN_NAMES = ("float_val",)
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
def show(self, float_val=None):
|
||||
if float_val is None:
|
||||
float_val = 0.0
|
||||
return (float_val,)
|
||||
|
||||
class ShowList_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"list_val": ("LIST",)}}
|
||||
RETURN_TYPES = ("LIST",)
|
||||
RETURN_NAMES = ("list_val",)
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
def show(self, list_val=None):
|
||||
if list_val is None:
|
||||
list_val = []
|
||||
return (list_val,)
|
||||
|
||||
class ShowText_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"text": ("STRING",)}}
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("text",)
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
def show(self, text=None):
|
||||
if text is None:
|
||||
text = ""
|
||||
return (text,)
|
||||
|
||||
class PreviewMask_UTK:
|
||||
CATEGORY = "UniversalToolkit"
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"mask": ("MASK",)}}
|
||||
RETURN_TYPES = ("MASK",)
|
||||
RETURN_NAMES = ("mask",)
|
||||
FUNCTION = "show"
|
||||
IS_PREVIEW = True
|
||||
def show(self, mask=None):
|
||||
if mask is None:
|
||||
mask = torch.zeros([1, 1, 64, 64], dtype=torch.float32)
|
||||
return (mask,)
|
||||
|
||||
# Node mappings
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Show_UTK": Show_UTK,
|
||||
"ShowInt_UTK": ShowInt_UTK,
|
||||
"ShowFloat_UTK": ShowFloat_UTK,
|
||||
"ShowList_UTK": ShowList_UTK,
|
||||
"ShowText_UTK": ShowText_UTK,
|
||||
"PreviewMask_UTK": PreviewMask_UTK,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"Show_UTK": "Show (UTK)",
|
||||
"ShowInt_UTK": "Show Int (UTK)",
|
||||
"ShowFloat_UTK": "Show Float (UTK)",
|
||||
"ShowList_UTK": "Show List (UTK)",
|
||||
"ShowText_UTK": "Show Text (UTK)",
|
||||
"PreviewMask_UTK": "Preview Mask (UTK)",
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user