Chaoses-Ib / ComfyScript

A Python frontend and library for ComfyUI
MIT License
383 stars 21 forks source link

Is it possible to edit nodes and have it be reflected in ComfyScript? #63

Closed pranaykoppula closed 1 month ago

pranaykoppula commented 1 month ago

For example, in nodes.py in the main ComfyUI directory, the code for the Load Image class is as follows:

class LoadImage:
    @classmethod
    def INPUT_TYPES(s):
        input_dir = folder_paths.get_input_directory()
        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
        return {"required":
                    {"image": (sorted(files), {"image_upload": True})},
                }

    CATEGORY = "image"

    RETURN_TYPES = ("IMAGE", "MASK")
    FUNCTION = "load_image"
    def load_image(self, image):
        image_path = folder_paths.get_annotated_filepath(image)

        img = node_helpers.pillow(Image.open, image_path)

        output_images = []
        output_masks = []
        w, h = None, None

        excluded_formats = ['MPO']

        for i in ImageSequence.Iterator(img):
            i = node_helpers.pillow(ImageOps.exif_transpose, i)

            if i.mode == 'I':
                i = i.point(lambda i: i * (1 / 255))
            image = i.convert("RGB")

            if len(output_images) == 0:
                w = image.size[0]
                h = image.size[1]

            if image.size[0] != w or image.size[1] != h:
                continue

            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[None,]
            if 'A' in i.getbands():
                mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
                mask = 1. - torch.from_numpy(mask)
            else:
                mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
            output_images.append(image)
            output_masks.append(mask.unsqueeze(0))

        if len(output_images) > 1 and img.format not in excluded_formats:
            output_image = torch.cat(output_images, dim=0)
            output_mask = torch.cat(output_masks, dim=0)
        else:
            output_image = output_images[0]
            output_mask = output_masks[0]

        return (output_image, output_mask)

In load_image, I want to change

image_path = folder_paths.get_annotated_filepath(image)

 img = node_helpers.pillow(Image.open, image_path)

into this:

img = image

because I wanna pass the PIL Image object directly rather than giving it a path. I want to make similar changes to the SaveImage class.

This is possible to do and integrate in the ComfyUI-to-Python-Extension because it just imports the original ComfyUI nodes as is. But that extension hasn't been updated in a year, and the same thing doesn't seem to work in ComfyScript

I imagine this is because of the way the nodes are being imported. Is there a way around this?

Chaoses-Ib commented 1 month ago

In ComfyScript's real mode, nodes are also just imported, but with a wrapped __new__ to handle the arguments. You can modify it directly:

from comfy_script.runtime.real import *
load()
from comfy_script.runtime.real.nodes import *

print(getattr(LoadImage, LoadImage.FUNCTION))
# <function LoadImage.load_image at 0x0000021B17F70160>

def my_load_image(self, image):
    print('LoadImage is called')
setattr(LoadImage, LoadImage.FUNCTION, my_load_image)

LoadImage('test.png')
# LoadImage is called

By the way, there are already built-in nodes installed by ComfyScript to pass PIL images directly:

def PILToImage(
    images: PilImage
) -> Image

def PILToMask(
    images: PilImage
) -> Image

def ImageToPIL(
    images: Image
) -> PilImage
pranaykoppula commented 1 month ago

Thank you! This was very helpful