jammyfu
71fa5bbe83
- Introduced MaskSwitch node for conditional switching between two masks. - Enhanced NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS to include MaskSwitch. - Updated README files to document the new MaskSwitch and Web Image Loader features, including usage instructions and examples. - Improved error handling in MaskSwitch to return a default empty mask when no inputs are provided.
134 lines
4.0 KiB
Python
134 lines
4.0 KiB
Python
import torch
|
|
|
|
class ImageSwitch:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"use_first": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"image_1": ("IMAGE",),
|
|
"image_2": ("IMAGE",),
|
|
},
|
|
"_meta": {
|
|
"preferred_width": 300,
|
|
"maintain_dimensions": True
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "switch_image"
|
|
CATEGORY = "🎨Painting👓Coder/🔄Switch"
|
|
|
|
def switch_image(self, use_first, image_1=None, image_2=None):
|
|
try:
|
|
# 创建空白图像作为默认值
|
|
empty_image = torch.ones((1, 1024, 1024, 3))
|
|
|
|
# 如果两个输入都为空，返回空白图像
|
|
if image_1 is None and image_2 is None:
|
|
return (empty_image,)
|
|
|
|
# 根据use_first选择输出
|
|
if use_first:
|
|
return (image_1,) if image_1 is not None else (empty_image,)
|
|
else:
|
|
return (image_2,) if image_2 is not None else (empty_image,)
|
|
|
|
except Exception as e:
|
|
print(f"Error in ImageSwitch: {str(e)}")
|
|
return (torch.ones((1, 512, 512, 3)),)
|
|
|
|
class TextSwitch:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"use_first": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"text_1": ("STRING", {"multiline": True}),
|
|
"text_2": ("STRING", {"multiline": True}),
|
|
},
|
|
"_meta": {
|
|
"preferred_width": 300,
|
|
"maintain_dimensions": True
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text",)
|
|
FUNCTION = "switch_text"
|
|
CATEGORY = "🎨Painting👓Coder/🔄Switch"
|
|
|
|
def switch_text(self, use_first, text_1=None, text_2=None):
|
|
try:
|
|
# 如果两个输入都为空，返回空字符串
|
|
if text_1 is None and text_2 is None:
|
|
return ("",)
|
|
|
|
# 根据use_first选择输出
|
|
if use_first:
|
|
return (text_1,) if text_1 is not None else ("",)
|
|
else:
|
|
return (text_2,) if text_2 is not None else ("",)
|
|
|
|
except Exception as e:
|
|
print(f"Error in TextSwitch: {str(e)}")
|
|
return ("",)
|
|
|
|
class MaskSwitch:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"use_first": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"mask_1": ("MASK",),
|
|
"mask_2": ("MASK",),
|
|
},
|
|
"_meta": {
|
|
"preferred_width": 300,
|
|
"maintain_dimensions": True
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
RETURN_NAMES = ("mask",)
|
|
FUNCTION = "switch_mask"
|
|
CATEGORY = "🎨Painting👓Coder/🔄Switch"
|
|
|
|
def switch_mask(self, use_first, mask_1=None, mask_2=None):
|
|
try:
|
|
# 创建空白掩码作为默认值
|
|
empty_mask = torch.zeros((1, 1024, 1024))
|
|
|
|
# 如果两个输入都为空，返回空白掩码
|
|
if mask_1 is None and mask_2 is None:
|
|
return (empty_mask,)
|
|
|
|
# 根据use_first选择输出
|
|
if use_first:
|
|
return (mask_1,) if mask_1 is not None else (empty_mask,)
|
|
else:
|
|
return (mask_2,) if mask_2 is not None else (empty_mask,)
|
|
|
|
except Exception as e:
|
|
print(f"Error in MaskSwitch: {str(e)}")
|
|
return (torch.zeros((1, 512, 512)),)
|
|
|
|
# 添加到 ComfyUI 节点注册
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageSwitch": ImageSwitch,
|
|
"TextSwitch": TextSwitch,
|
|
"MaskSwitch": MaskSwitch
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImageSwitch": "Image Switch 🔄",
|
|
"TextSwitch": "Text Switch 🔄",
|
|
"MaskSwitch": "Mask Switch 🔄"
|
|
} |