ComfyUI icon indicating copy to clipboard operation
ComfyUI copied to clipboard

How to extend LoadImage class? or any node class?

Open myusf01 opened this issue 2 years ago • 6 comments

I'm looking to create a custom image loader with width and height outputs.

Duplicating this LoadImage class and then adding width and height output is easy. But I don't want it to do all the process again. So how can we extend this class to add these outputs? https://github.com/comfyanonymous/ComfyUI/blob/f2a7cc912186c89fda9580f36da28c7fc382ea26/nodes.py#L1303

myusf01 avatar Aug 22 '23 15:08 myusf01

And when I try this code as a custom node I can't see 'choose file to upload' button and preview of the image. But when I add image size outputs to main class, every thing is fine so that I thought extending LoadImage class may work.

class CustomImageLoad:
    def __init__(self):
        pass
    
    @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), )},
                }

    CATEGORY = "image"

    RETURN_TYPES = ("IMAGE", "MASK","INT","INT")
    RETURN_NAMES = ("IMAGE", "MASK","WIDTH","HEIGHT")
    FUNCTION = "load_image"
    def load_image(self, image):
        image_path = folder_paths.get_annotated_filepath(image)
        i = Image.open(image_path)
        i = ImageOps.exif_transpose(i)
        image = i.convert("RGB")
        image = np.array(image).astype(np.float32) / 255.0
        image = torch.from_numpy(image)[None,]
        width, height = i.size
        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")
        return (image, mask, width, height)

    @classmethod
    def IS_CHANGED(s, image):
        image_path = folder_paths.get_annotated_filepath(image)
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()

    @classmethod
    def VALIDATE_INPUTS(s, image):
        if not folder_paths.exists_annotated_filepath(image):
            return "Invalid image file: {}".format(image)

        return True

myusf01 avatar Aug 22 '23 15:08 myusf01

The upload button comes from uploadImage.js my solution was to use the following inside my __init__.py

NODE_CLASS_MAPPINGS = {
    "TacoAnimatedLoader": TacoAnimatedLoader,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TacoAnimatedLoader": "Taco Animated Image Loader"
}

# This will copy the required javascript into the correct location so it gets loaded for decorating
def update_javascript():
    extensions_folder = os.path.join(os.path.dirname(os.path.realpath(__main__.__file__)),
                                     "web" + os.sep + "extensions" + os.sep + "ComfyUI-TacoNodes")
    javascript_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "js")

    if not os.path.exists(extensions_folder):
        print("Creating frontend extension folder: " + extensions_folder)
        os.mkdir(extensions_folder)

    result = filecmp.dircmp(javascript_folder, extensions_folder)

    if result.left_only or result.diff_files:
        print('Update to javascripts files detected')
        file_list = list(result.left_only)
        file_list.extend(x for x in result.diff_files if x not in file_list)

        for file in file_list:
            print(f'Copying {file} to extensions folder')
            src_file = os.path.join(javascript_folder, file)
            dst_file = os.path.join(extensions_folder, file)
            if os.path.exists(dst_file):
                os.remove(dst_file)
            shutil.copy(src_file, dst_file)


update_javascript()

then write a file js/someJavascript.js

import {app} from "/scripts/app.js"

app.registerExtension({
	name: "Comfy.UploadImage",
	async beforeRegisterNodeDef(nodeType, nodeData, app) {
		if (nodeData.name === "TacoAnimatedLoader") {
			nodeData.input.required.upload = ["IMAGEUPLOAD"];
		}
	},
});

YOUR-WORST-TACO avatar Aug 24 '23 05:08 YOUR-WORST-TACO

Oh this is awesome. Thanks a lot!

myusf01 avatar Aug 24 '23 12:08 myusf01

if you are interested in full examples I have my latest work here: https://github.com/YOUR-WORST-TACO/ComfyUI-TacoNodes

it includes a couple examples of me extending the ImageLoader and I rewrote the logic for IMAGEUPLOAD so I can extend or change it as needed

YOUR-WORST-TACO avatar Aug 24 '23 17:08 YOUR-WORST-TACO

Also take a look at this https://github.com/comfyanonymous/ComfyUI/pull/1273 rather than manually copying your web assets. cc @YOUR-WORST-TACO

M1kep avatar Aug 24 '23 22:08 M1kep

OOOOOOO, I like that a lot, thank you!

YOUR-WORST-TACO avatar Aug 24 '23 22:08 YOUR-WORST-TACO

This code do not work anymore.

if (nodeData.name === "MyNode") {
	nodeData.input.required.upload = ["IMAGEUPLOAD"];
}

Image

console.log(nodeData.input.required.upload) ----> undefined

KLL535 avatar Mar 09 '25 20:03 KLL535

This code do not work anymore.

You can now just add the image_upload flag to get the upload button and preview:

    @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, allow_batch: False })},
                }

The input name also doesn't need to be "image" anymore.

christian-byrne avatar Mar 10 '25 01:03 christian-byrne