Acly-comfyui-tooling-nodes/nodes.py

from PIL import Image
from typing import Any
import numpy as np
import base64
import torch
from io import BytesIO
from server import PromptServer, BinaryEventTypes


class LoadImageBase64:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"image": ("STRING", {"multiline": False})}}

    RETURN_TYPES = ("IMAGE", "MASK")
    CATEGORY = "external_tooling"
    FUNCTION = "load_image"

    def load_image(self, image):
        imgdata = base64.b64decode(image)
        img = Image.open(BytesIO(imgdata))

        if "A" in img.getbands():
            mask = np.array(img.getchannel("A")).astype(np.float32) / 255.0
            mask = 1.0 - torch.from_numpy(mask)
        else:
            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")

        img = img.convert("RGB")
        img = np.array(img).astype(np.float32) / 255.0
        img = torch.from_numpy(img)[None,]

        return (img, mask)


class LoadMaskBase64:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"mask": ("STRING", {"multiline": False})}}

    RETURN_TYPES = ("MASK",)
    CATEGORY = "external_tooling"
    FUNCTION = "load_mask"

    def load_mask(self, mask):
        imgdata = base64.b64decode(mask)
        img = Image.open(BytesIO(imgdata))
        img = np.array(img).astype(np.float32) / 255.0
        img = torch.from_numpy(img)
        if img.dim() == 3:  # RGB(A) input, use red channel
            img = img[:, :, 0]
        return (img,)


class SendImageWebSocket:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"images": ("IMAGE",)}}

    RETURN_TYPES = ()
    FUNCTION = "send_images"
    OUTPUT_NODE = True
    CATEGORY = "external_tooling"

    def send_images(self, images):
        results = []
        for tensor in images:
            array = 255.0 * tensor.cpu().numpy()
            image = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))

            server = PromptServer.instance
            server.send_sync(
                BinaryEventTypes.UNENCODED_PREVIEW_IMAGE,
                ["PNG", image, None],
                server.client_id,
            )
            results.append(
                # Could put some kind of ID here, but for now just match them by index
                {"source": "websocket", "content-type": "image/png", "type": "output"}
            )

        return {"ui": {"images": results}}


class CropImage:
    """Deprecated, ComfyUI has an ImageCrop node now which does the same."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "x": (
                    "INT",
                    {"default": 0, "min": 0, "max": 8192, "step": 1},
                ),
                "y": (
                    "INT",
                    {"default": 0, "min": 0, "max": 8192, "step": 1},
                ),
                "width": (
                    "INT",
                    {"default": 512, "min": 1, "max": 8192, "step": 1},
                ),
                "height": (
                    "INT",
                    {"default": 512, "min": 1, "max": 8192, "step": 1},
                ),
            }
        }

    CATEGORY = "external_tooling"
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "crop"

    def crop(self, image, x, y, width, height):
        out = image[:, y : y + height, x : x + width, :]
        return (out,)


class ApplyMaskToImage:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "mask": ("MASK",),
            }
        }

    CATEGORY = "external_tooling"
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "apply_mask"

    def apply_mask(self, image: torch.Tensor, mask: torch.Tensor):
        # Move the channel to the second dimension for processing
        out = image.movedim(-1, 1)

        # Check if the images are RGB, and if so, add an alpha channel initialized to 1
        if out.shape[1] == 3:  # Assuming RGB images
            out = torch.cat([out, torch.ones_like(out[:, :1, :, :])], dim=1)

        # Ensure masks are unsqueezed to match the alpha channel dimension if needed
        if mask.ndim == 2:
            mask = mask.unsqueeze(0)  # Add a batch dimension to masks
        # For single mask, expand it to match size of image batch size.
        if mask.shape[0] == 1:
            mask = mask.repeat(out.shape[0], 1, 1)

        assert mask.ndim == 3, f"Mask should have shape [B, H, W]. {mask.shape}"
        assert out.ndim == 4, f"Image should have shsape [B, C, H, W]. {out.shape}"
        assert out.shape[-2:] == mask.shape[-2:], f"{out.shape[-2:]} != {mask.shape[-2:]}"
        assert out.shape[0] == mask.shape[0], f"{out.shape[0]} != {mask.shape[0]}"
        # Apply each mask in the batch to its corresponding image's alpha channel
        for i in range(out.shape[0]):
            out[i, 3, :, :] = mask[i]

        # Move the channel back to its original dimension
        out = out.movedim(1, -1)

        return (out,)


class ListWrapper:
    content: list[dict[str, Any]]

    def __init__(self, initial: list | None = None):
        self.content = initial or []


class ListAppend:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {},
            "optional": {
                "list": ("LIST",),
                "image": ("IMAGE",),
                "mask": ("MASK",),
                "conditioning": ("CONDITIONING",),
            },
        }

    CATEGORY = "external_tooling"
    RETURN_TYPES = ("LIST",)
    FUNCTION = "append"

    def append(self, list: ListWrapper | None = None, **kwargs):
        if list is None:
            list = ListWrapper()
        list.content.append(kwargs)
        return (list,)


class ListElement:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "list": ("LIST",),
                "index": ("INT", {"default": 0, "min": 0}),
            }
        }

    CATEGORY = "external_tooling"
    RETURN_TYPES = ("IMAGE", "MASK", "CONDITIONING")
    FUNCTION = "element"

    def element(self, list: ListWrapper, index: int):
        elem = list.content[index]
        return (elem.get("image"), elem.get("mask"), elem.get("conditioning"))