xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2025-12-09 05:54:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

[V3] convert nodes_mask.py to V3 schema (#10669)

Browse Source 
* convert nodes_mask.py to V3 schema

* set "Preview Mask" as display name for MaskPreview
This commit is contained in:
Alexander Piskun 2025-12-06 06:24:10 +02:00 committed by GitHub
parent 117bf3f2bd
commit 913f86b727
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 263 additions and 245 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

508
comfy_extras/nodes_mask.py
View File

@ -3,11 +3,10 @@ import scipy.ndimage
import torch import torch
import comfy.utils import comfy.utils
import node_helpers import node_helpers
import folder_paths from typing_extensions import override
import random from comfy_api.latest import ComfyExtension, IO, UI
import nodes import nodes
from nodes import MAX_RESOLUTION
def composite(destination, source, x, y, mask = None, multiplier = 8, resize_source = False): def composite(destination, source, x, y, mask = None, multiplier = 8, resize_source = False):
source = source.to(destination.device) source = source.to(destination.device)
@ -46,202 +45,213 @@ def composite(destination, source, x, y, mask = None, multiplier = 8, resize_sou
destination[..., top:bottom, left:right] = source_portion + destination_portion destination[..., top:bottom, left:right] = source_portion + destination_portion
return destination return destination
class LatentCompositeMasked: class LatentCompositeMasked(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="LatentCompositeMasked",
"destination": ("LATENT",), category="latent",
"source": ("LATENT",), inputs=[
"x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}), IO.Latent.Input("destination"),
"y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}), IO.Latent.Input("source"),
"resize_source": ("BOOLEAN", {"default": False}), IO.Int.Input("x", default=0, min=0, max=nodes.MAX_RESOLUTION, step=8),
}, IO.Int.Input("y", default=0, min=0, max=nodes.MAX_RESOLUTION, step=8),
"optional": { IO.Boolean.Input("resize_source", default=False),
"mask": ("MASK",), IO.Mask.Input("mask", optional=True),
} ],
} outputs=[IO.Latent.Output()],
RETURN_TYPES = ("LATENT",) )
FUNCTION = "composite"
CATEGORY = "latent" @classmethod
def execute(cls, destination, source, x, y, resize_source, mask = None) -> IO.NodeOutput:
def composite(self, destination, source, x, y, resize_source, mask = None):
output = destination.copy() output = destination.copy()
destination = destination["samples"].clone() destination = destination["samples"].clone()
source = source["samples"] source = source["samples"]
output["samples"] = composite(destination, source, x, y, mask, 8, resize_source) output["samples"] = composite(destination, source, x, y, mask, 8, resize_source)
return (output,) return IO.NodeOutput(output)
class ImageCompositeMasked: composite = execute # TODO: remove
class ImageCompositeMasked(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="ImageCompositeMasked",
"destination": ("IMAGE",), category="image",
"source": ("IMAGE",), inputs=[
"x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), IO.Image.Input("destination"),
"y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), IO.Image.Input("source"),
"resize_source": ("BOOLEAN", {"default": False}), IO.Int.Input("x", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
}, IO.Int.Input("y", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
"optional": { IO.Boolean.Input("resize_source", default=False),
"mask": ("MASK",), IO.Mask.Input("mask", optional=True),
} ],
} outputs=[IO.Image.Output()],
RETURN_TYPES = ("IMAGE",) )
FUNCTION = "composite"
CATEGORY = "image" @classmethod
def execute(cls, destination, source, x, y, resize_source, mask = None) -> IO.NodeOutput:
def composite(self, destination, source, x, y, resize_source, mask = None):
destination, source = node_helpers.image_alpha_fix(destination, source) destination, source = node_helpers.image_alpha_fix(destination, source)
destination = destination.clone().movedim(-1, 1) destination = destination.clone().movedim(-1, 1)
output = composite(destination, source.movedim(-1, 1), x, y, mask, 1, resize_source).movedim(1, -1) output = composite(destination, source.movedim(-1, 1), x, y, mask, 1, resize_source).movedim(1, -1)
return (output,) return IO.NodeOutput(output)
class MaskToImage: composite = execute # TODO: remove
class MaskToImage(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="MaskToImage",
"mask": ("MASK",), display_name="Convert Mask to Image",
} category="mask",
} inputs=[
IO.Mask.Input("mask"),
],
outputs=[IO.Image.Output()],
)
CATEGORY = "mask" @classmethod
def execute(cls, mask) -> IO.NodeOutput:
RETURN_TYPES = ("IMAGE",)
FUNCTION = "mask_to_image"
def mask_to_image(self, mask):
result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
return (result,) return IO.NodeOutput(result)
class ImageToMask: mask_to_image = execute # TODO: remove
class ImageToMask(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="ImageToMask",
"image": ("IMAGE",), display_name="Convert Image to Mask",
"channel": (["red", "green", "blue", "alpha"],), category="mask",
} inputs=[
} IO.Image.Input("image"),
IO.Combo.Input("channel", options=["red", "green", "blue", "alpha"]),
],
outputs=[IO.Mask.Output()],
)
CATEGORY = "mask" @classmethod
def execute(cls, image, channel) -> IO.NodeOutput:
RETURN_TYPES = ("MASK",)
FUNCTION = "image_to_mask"
def image_to_mask(self, image, channel):
channels = ["red", "green", "blue", "alpha"] channels = ["red", "green", "blue", "alpha"]
mask = image[:, :, :, channels.index(channel)] mask = image[:, :, :, channels.index(channel)]
return (mask,) return IO.NodeOutput(mask)
class ImageColorToMask: image_to_mask = execute # TODO: remove
class ImageColorToMask(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="ImageColorToMask",
"image": ("IMAGE",), category="mask",
"color": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFF, "step": 1, "display": "color"}), inputs=[
} IO.Image.Input("image"),
} IO.Int.Input("color", default=0, min=0, max=0xFFFFFF, step=1, display_mode=IO.NumberDisplay.number),
],
outputs=[IO.Mask.Output()],
)
CATEGORY = "mask" @classmethod
def execute(cls, image, color) -> IO.NodeOutput:
RETURN_TYPES = ("MASK",)
FUNCTION = "image_to_mask"
def image_to_mask(self, image, color):
temp = (torch.clamp(image, 0, 1.0) * 255.0).round().to(torch.int) temp = (torch.clamp(image, 0, 1.0) * 255.0).round().to(torch.int)
temp = torch.bitwise_left_shift(temp[:,:,:,0], 16) + torch.bitwise_left_shift(temp[:,:,:,1], 8) + temp[:,:,:,2] temp = torch.bitwise_left_shift(temp[:,:,:,0], 16) + torch.bitwise_left_shift(temp[:,:,:,1], 8) + temp[:,:,:,2]
mask = torch.where(temp == color, 1.0, 0).float() mask = torch.where(temp == color, 1.0, 0).float()
return (mask,) return IO.NodeOutput(mask)
class SolidMask: image_to_mask = execute # TODO: remove
class SolidMask(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(cls): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="SolidMask",
"value": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), category="mask",
"width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), inputs=[
"height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), IO.Float.Input("value", default=1.0, min=0.0, max=1.0, step=0.01),
} IO.Int.Input("width", default=512, min=1, max=nodes.MAX_RESOLUTION, step=1),
} IO.Int.Input("height", default=512, min=1, max=nodes.MAX_RESOLUTION, step=1),
],
outputs=[IO.Mask.Output()],
)
CATEGORY = "mask" @classmethod
def execute(cls, value, width, height) -> IO.NodeOutput:
RETURN_TYPES = ("MASK",)
FUNCTION = "solid"
def solid(self, value, width, height):
out = torch.full((1, height, width), value, dtype=torch.float32, device="cpu") out = torch.full((1, height, width), value, dtype=torch.float32, device="cpu")
return (out,) return IO.NodeOutput(out)
class InvertMask: solid = execute # TODO: remove
class InvertMask(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(cls): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="InvertMask",
"mask": ("MASK",), category="mask",
} inputs=[
} IO.Mask.Input("mask"),
],
outputs=[IO.Mask.Output()],
)
CATEGORY = "mask" @classmethod
def execute(cls, mask) -> IO.NodeOutput:
RETURN_TYPES = ("MASK",)
FUNCTION = "invert"
def invert(self, mask):
out = 1.0 - mask out = 1.0 - mask
return (out,) return IO.NodeOutput(out)
class CropMask: invert = execute # TODO: remove
class CropMask(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(cls): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="CropMask",
"mask": ("MASK",), category="mask",
"x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), inputs=[
"y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), IO.Mask.Input("mask"),
"width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), IO.Int.Input("x", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
"height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}), IO.Int.Input("y", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
} IO.Int.Input("width", default=512, min=1, max=nodes.MAX_RESOLUTION, step=1),
} IO.Int.Input("height", default=512, min=1, max=nodes.MAX_RESOLUTION, step=1),
],
outputs=[IO.Mask.Output()],
)
CATEGORY = "mask" @classmethod
def execute(cls, mask, x, y, width, height) -> IO.NodeOutput:
RETURN_TYPES = ("MASK",)
FUNCTION = "crop"
def crop(self, mask, x, y, width, height):
mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1])) mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
out = mask[:, y:y + height, x:x + width] out = mask[:, y:y + height, x:x + width]
return (out,) return IO.NodeOutput(out)
class MaskComposite: crop = execute # TODO: remove
class MaskComposite(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(cls): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="MaskComposite",
"destination": ("MASK",), category="mask",
"source": ("MASK",), inputs=[
"x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), IO.Mask.Input("destination"),
"y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), IO.Mask.Input("source"),
"operation": (["multiply", "add", "subtract", "and", "or", "xor"],), IO.Int.Input("x", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
} IO.Int.Input("y", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
} IO.Combo.Input("operation", options=["multiply", "add", "subtract", "and", "or", "xor"]),
],
outputs=[IO.Mask.Output()],
)
CATEGORY = "mask" @classmethod
def execute(cls, destination, source, x, y, operation) -> IO.NodeOutput:
RETURN_TYPES = ("MASK",)
FUNCTION = "combine"
def combine(self, destination, source, x, y, operation):
output = destination.reshape((-1, destination.shape[-2], destination.shape[-1])).clone() output = destination.reshape((-1, destination.shape[-2], destination.shape[-1])).clone()
source = source.reshape((-1, source.shape[-2], source.shape[-1])) source = source.reshape((-1, source.shape[-2], source.shape[-1]))
@ -267,28 +277,29 @@ class MaskComposite:
output = torch.clamp(output, 0.0, 1.0) output = torch.clamp(output, 0.0, 1.0)
return (output,) return IO.NodeOutput(output)
class FeatherMask: combine = execute # TODO: remove
class FeatherMask(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(cls): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="FeatherMask",
"mask": ("MASK",), category="mask",
"left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), inputs=[
"top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), IO.Mask.Input("mask"),
"right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), IO.Int.Input("left", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
"bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}), IO.Int.Input("top", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
} IO.Int.Input("right", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
} IO.Int.Input("bottom", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
],
outputs=[IO.Mask.Output()],
)
CATEGORY = "mask" @classmethod
def execute(cls, mask, left, top, right, bottom) -> IO.NodeOutput:
RETURN_TYPES = ("MASK",)
FUNCTION = "feather"
def feather(self, mask, left, top, right, bottom):
output = mask.reshape((-1, mask.shape[-2], mask.shape[-1])).clone() output = mask.reshape((-1, mask.shape[-2], mask.shape[-1])).clone()
left = min(left, output.shape[-1]) left = min(left, output.shape[-1])
@ -312,26 +323,28 @@ class FeatherMask:
feather_rate = (y + 1) / bottom feather_rate = (y + 1) / bottom
output[:, -y, :] *= feather_rate output[:, -y, :] *= feather_rate
return (output,) return IO.NodeOutput(output)
class GrowMask: feather = execute # TODO: remove
class GrowMask(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(cls): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="GrowMask",
"mask": ("MASK",), display_name="Grow Mask",
"expand": ("INT", {"default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1}), category="mask",
"tapered_corners": ("BOOLEAN", {"default": True}), inputs=[
}, IO.Mask.Input("mask"),
} IO.Int.Input("expand", default=0, min=-nodes.MAX_RESOLUTION, max=nodes.MAX_RESOLUTION, step=1),
IO.Boolean.Input("tapered_corners", default=True),
],
outputs=[IO.Mask.Output()],
)
CATEGORY = "mask" @classmethod
def execute(cls, mask, expand, tapered_corners) -> IO.NodeOutput:
RETURN_TYPES = ("MASK",)
FUNCTION = "expand_mask"
def expand_mask(self, mask, expand, tapered_corners):
c = 0 if tapered_corners else 1 c = 0 if tapered_corners else 1
kernel = np.array([[c, 1, c], kernel = np.array([[c, 1, c],
[1, 1, 1], [1, 1, 1],
@ -347,69 +360,74 @@ class GrowMask:
output = scipy.ndimage.grey_dilation(output, footprint=kernel) output = scipy.ndimage.grey_dilation(output, footprint=kernel)
output = torch.from_numpy(output) output = torch.from_numpy(output)
out.append(output) out.append(output)
return (torch.stack(out, dim=0),) return IO.NodeOutput(torch.stack(out, dim=0))
class ThresholdMask: expand_mask = execute # TODO: remove
class ThresholdMask(IO.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return { return IO.Schema(
"required": { node_id="ThresholdMask",
"mask": ("MASK",), category="mask",
"value": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), inputs=[
} IO.Mask.Input("mask"),
} IO.Float.Input("value", default=0.5, min=0.0, max=1.0, step=0.01),
],
outputs=[IO.Mask.Output()],
)
CATEGORY = "mask" @classmethod
def execute(cls, mask, value) -> IO.NodeOutput:
RETURN_TYPES = ("MASK",)
FUNCTION = "image_to_mask"
def image_to_mask(self, mask, value):
mask = (mask > value).float() mask = (mask > value).float()
return (mask,) return IO.NodeOutput(mask)
image_to_mask = execute # TODO: remove
# Mask Preview - original implement from # Mask Preview - original implement from
# https://github.com/cubiq/ComfyUI_essentials/blob/9d9f4bedfc9f0321c19faf71855e228c93bd0dc9/mask.py#L81 # https://github.com/cubiq/ComfyUI_essentials/blob/9d9f4bedfc9f0321c19faf71855e228c93bd0dc9/mask.py#L81
# upstream requested in https://github.com/Kosinkadink/rfcs/blob/main/rfcs/0000-corenodes.md#preview-nodes # upstream requested in https://github.com/Kosinkadink/rfcs/blob/main/rfcs/0000-corenodes.md#preview-nodes
class MaskPreview(nodes.SaveImage): class MaskPreview(IO.ComfyNode):
def __init__(self): @classmethod
self.output_dir = folder_paths.get_temp_directory() def define_schema(cls):
self.type = "temp" return IO.Schema(
self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5)) node_id="MaskPreview",
self.compress_level = 4 display_name="Preview Mask",
category="mask",
description="Saves the input images to your ComfyUI output directory.",
inputs=[
IO.Mask.Input("mask"),
],
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
is_output_node=True,
)
@classmethod @classmethod
def INPUT_TYPES(s): def execute(cls, mask, filename_prefix="ComfyUI") -> IO.NodeOutput:
return { return IO.NodeOutput(ui=UI.PreviewMask(mask))
"required": {"mask": ("MASK",), },
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
}
FUNCTION = "execute"
CATEGORY = "mask"
def execute(self, mask, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
preview = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
return self.save_images(preview, filename_prefix, prompt, extra_pnginfo)
NODE_CLASS_MAPPINGS = { class MaskExtension(ComfyExtension):
"LatentCompositeMasked": LatentCompositeMasked, @override
"ImageCompositeMasked": ImageCompositeMasked, async def get_node_list(self) -> list[type[IO.ComfyNode]]:
"MaskToImage": MaskToImage, return [
"ImageToMask": ImageToMask, LatentCompositeMasked,
"ImageColorToMask": ImageColorToMask, ImageCompositeMasked,
"SolidMask": SolidMask, MaskToImage,
"InvertMask": InvertMask, ImageToMask,
"CropMask": CropMask, ImageColorToMask,
"MaskComposite": MaskComposite, SolidMask,
"FeatherMask": FeatherMask, InvertMask,
"GrowMask": GrowMask, CropMask,
"ThresholdMask": ThresholdMask, MaskComposite,
"MaskPreview": MaskPreview FeatherMask,
} GrowMask,
ThresholdMask,
MaskPreview,
]
NODE_DISPLAY_NAME_MAPPINGS = {
"ImageToMask": "Convert Image to Mask", async def comfy_entrypoint() -> MaskExtension:
"MaskToImage": "Convert Mask to Image", return MaskExtension()
}