How to extend LoadImage class? or any node class?
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
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
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"];
}
},
});
Oh this is awesome. Thanks a lot!
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
Also take a look at this https://github.com/comfyanonymous/ComfyUI/pull/1273 rather than manually copying your web assets. cc @YOUR-WORST-TACO
OOOOOOO, I like that a lot, thank you!
This code do not work anymore.
if (nodeData.name === "MyNode") {
nodeData.input.required.upload = ["IMAGEUPLOAD"];
}
console.log(nodeData.input.required.upload) ----> undefined
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.