qnloft
/
ComfyUI
mirror of https://github.com/comfyanonymous/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor: move nodes_mask_convertion nodes to nodes_mask.

This commit is contained in:
comfyanonymous 2023-04-14 00:21:01 -04:00
parent e1db7a2038
commit f48f0872e2
3 changed files with 31 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
comfy_extras/nodes_mask.py
View File

@ -59,7 +59,7 @@ class LatentCompositeMasked:
class MaskToImage:
@classmethod
def INPUT_TYPES(cls):
def INPUT_TYPES(s):
return {
"required": {
"mask": ("MASK",),
@ -69,13 +69,31 @@ class MaskToImage:
CATEGORY = "mask"
RETURN_TYPES = ("IMAGE",)
FUNCTION = "mask_to_image"
FUNCTION = "convert"
def mask_to_image(self, mask):
result = mask[None, :, :, None].expand(-1, -1, -1, 3)
return (result,)
def convert(self, mask):
image = torch.cat([torch.reshape(mask.clone(), [1, mask.shape[0], mask.shape[1], 1,])] * 3, 3)
class ImageToMask:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"channel": (["red", "green", "blue"],),
}
}
return (image,)
CATEGORY = "mask"
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 SolidMask:
@classmethod
@ -231,6 +249,7 @@ class FeatherMask:
NODE_CLASS_MAPPINGS = {
"LatentCompositeMasked": LatentCompositeMasked,
"MaskToImage": MaskToImage,
"ImageToMask": ImageToMask,
"SolidMask": SolidMask,
"InvertMask": InvertMask,
"CropMask": CropMask,
@ -238,3 +257,7 @@ NODE_CLASS_MAPPINGS = {
"FeatherMask": FeatherMask,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ImageToMask": "Convert Image to Mask",
"MaskToImage": "Convert Mask to Image",
}

54
comfy_extras/nodes_mask_conversion.py
View File

@ -1,54 +0,0 @@
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",
}

1
nodes.py
View File

@ -1193,4 +1193,3 @@ def init_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.py"))
load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask_conversion.py"))