69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Filename: dynamic_mask_input.py
|
|
# Developer: jammyfu
|
|
# Category: 🎨Painting👓Coder/🖼️Image
|
|
|
|
import torch
|
|
from ..text.i18n import I18n
|
|
|
|
class DynamicMaskCombiner:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"mask_1": ("MASK",), # 添加一个初始掩码输入
|
|
},
|
|
"optional": {},
|
|
"_meta": {
|
|
"preferred_width": 300,
|
|
"maintain_dimensions": True
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
RETURN_NAMES = ("masks",)
|
|
OUTPUT_IS_LIST = (True,)
|
|
FUNCTION = "combine_masks"
|
|
CATEGORY = "🎨Painting👓Coder/🖼️Image"
|
|
|
|
def combine_masks(self, **kwargs):
|
|
try:
|
|
masks = []
|
|
# 获取所有掩码输入键
|
|
mask_inputs = sorted(
|
|
[k for k in kwargs.keys() if k.startswith('mask_')],
|
|
key=lambda x: int(x.split('_')[1])
|
|
)
|
|
|
|
# 只处理非None的输入（已连接的输入）
|
|
for key in mask_inputs:
|
|
input_mask = kwargs.get(key)
|
|
|
|
# 跳过未连接的输入（None值）
|
|
if input_mask is None:
|
|
continue
|
|
|
|
# 处理已连接的输入
|
|
if isinstance(input_mask, torch.Tensor):
|
|
masks.append(input_mask)
|
|
|
|
# 如果没有有效的掩码输入，返回空掩码
|
|
if not masks:
|
|
empty_mask = torch.zeros((1, 512, 512))
|
|
return ([empty_mask],)
|
|
|
|
return (masks,)
|
|
|
|
except Exception as e:
|
|
print(f"Error in DynamicMaskCombiner: {str(e)}")
|
|
empty_mask = torch.zeros((1, 512, 512))
|
|
return ([empty_mask],)
|
|
|
|
# 更新节点映射
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PaintingCoder::DynamicMaskCombiner": DynamicMaskCombiner,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PaintingCoder::DynamicMaskCombiner": "Dynamic Mask Input 🎭",
|
|
} |