From e12fb88b1b84e354872c0d761544558479bcfad2 Mon Sep 17 00:00:00 2001 From: missionfloyd Date: Tue, 11 Apr 2023 16:49:39 -0600 Subject: [PATCH 1/3] Image/mask conversion nodes --- nodes.py | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/nodes.py b/nodes.py index 14a73bcd..ecd931d6 100644 --- a/nodes.py +++ b/nodes.py @@ -1059,6 +1059,43 @@ class ImagePadForOutpaint: return (new_image, mask) +class ImageToMask: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "channel": (["red", "green", "blue"],), + } + } + + CATEGORY = "image" + + RETURN_TYPES = ("MASK",) + FUNCTION = "image_to_mask" + + def image_to_mask(self, image, channel): + channels = ["red", "green", "blue"] + mask = torch.select(image[0], 2, channels.index(channel)) + return (mask,) + +class MaskToImage: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + } + } + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "mask_to_image" + + def mask_to_image(self, mask): + result = mask[None, :, :, None].expand(-1, -1, -1, 3) + return (result,) NODE_CLASS_MAPPINGS = { "KSampler": KSampler, @@ -1102,6 +1139,8 @@ NODE_CLASS_MAPPINGS = { "unCLIPCheckpointLoader": unCLIPCheckpointLoader, "CheckpointLoader": CheckpointLoader, "DiffusersLoader": DiffusersLoader, + "ImageToMask": ImageToMask, + "MaskToImage": MaskToImage, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -1147,6 +1186,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ImageUpscaleWithModel": "Upscale Image (using Model)", "ImageInvert": "Invert Image", "ImagePadForOutpaint": "Pad Image for Outpainting", + "ImageToMask": "Convert Image to Mask", + "MaskToImage": "Convert Mask to Image", # _for_testing "VAEDecodeTiled": "VAE Decode (Tiled)", "VAEEncodeTiled": "VAE Encode (Tiled)", From e1d289c1ec6894e15af0b57b6630b853341c61fa Mon Sep 17 00:00:00 2001 From: missionfloyd Date: Tue, 11 Apr 2023 20:26:24 -0600 Subject: [PATCH 2/3] use slice instead of torch.select() --- nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index ecd931d6..815631f5 100644 --- a/nodes.py +++ b/nodes.py @@ -1076,7 +1076,7 @@ class ImageToMask: def image_to_mask(self, image, channel): channels = ["red", "green", "blue"] - mask = torch.select(image[0], 2, channels.index(channel)) + mask = image[0, :, :, channels.index(channel)] return (mask,) class MaskToImage: From 9371924e654128258cc82419e83c2a788a32e2be Mon Sep 17 00:00:00 2001 From: missionfloyd Date: Thu, 13 Apr 2023 03:11:17 -0600 Subject: [PATCH 3/3] Move mask conversion to separate file --- comfy_extras/nodes_mask_conversion.py | 54 +++++++++++++++++++++++++++ nodes.py | 42 +-------------------- 2 files changed, 55 insertions(+), 41 deletions(-) create mode 100644 comfy_extras/nodes_mask_conversion.py diff --git a/comfy_extras/nodes_mask_conversion.py b/comfy_extras/nodes_mask_conversion.py new file mode 100644 index 00000000..04dcbd0d --- /dev/null +++ b/comfy_extras/nodes_mask_conversion.py @@ -0,0 +1,54 @@ +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image + +import comfy.utils + +class ImageToMask: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "channel": (["red", "green", "blue"],), + } + } + + CATEGORY = "image" + + RETURN_TYPES = ("MASK",) + FUNCTION = "image_to_mask" + + def image_to_mask(self, image, channel): + channels = ["red", "green", "blue"] + mask = image[0, :, :, channels.index(channel)] + return (mask,) + +class MaskToImage: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + } + } + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "mask_to_image" + + def mask_to_image(self, mask): + result = mask[None, :, :, None].expand(-1, -1, -1, 3) + return (result,) + +NODE_CLASS_MAPPINGS = { + "ImageToMask": ImageToMask, + "MaskToImage": MaskToImage, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageToMask": "Convert Image to Mask", + "MaskToImage": "Convert Mask to Image", +} diff --git a/nodes.py b/nodes.py index 325e3ba6..3ed9cf49 100644 --- a/nodes.py +++ b/nodes.py @@ -1061,43 +1061,6 @@ class ImagePadForOutpaint: return (new_image, mask) -class ImageToMask: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "image": ("IMAGE",), - "channel": (["red", "green", "blue"],), - } - } - - CATEGORY = "image" - - RETURN_TYPES = ("MASK",) - FUNCTION = "image_to_mask" - - def image_to_mask(self, image, channel): - channels = ["red", "green", "blue"] - mask = image[0, :, :, channels.index(channel)] - return (mask,) - -class MaskToImage: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "mask": ("MASK",), - } - } - - CATEGORY = "image" - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "mask_to_image" - - def mask_to_image(self, mask): - result = mask[None, :, :, None].expand(-1, -1, -1, 3) - return (result,) NODE_CLASS_MAPPINGS = { "KSampler": KSampler, @@ -1141,8 +1104,6 @@ NODE_CLASS_MAPPINGS = { "unCLIPCheckpointLoader": unCLIPCheckpointLoader, "CheckpointLoader": CheckpointLoader, "DiffusersLoader": DiffusersLoader, - "ImageToMask": ImageToMask, - "MaskToImage": MaskToImage, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -1188,8 +1149,6 @@ NODE_DISPLAY_NAME_MAPPINGS = { "ImageUpscaleWithModel": "Upscale Image (using Model)", "ImageInvert": "Invert Image", "ImagePadForOutpaint": "Pad Image for Outpainting", - "ImageToMask": "Convert Image to Mask", - "MaskToImage": "Convert Mask to Image", # _for_testing "VAEDecodeTiled": "VAE Decode (Tiled)", "VAEEncodeTiled": "VAE Encode (Tiled)", @@ -1233,3 +1192,4 @@ def init_custom_nodes(): load_custom_nodes() load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_upscale_model.py")) load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_post_processing.py")) + load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask_conversion.py"))