From d7b3b0f8c11c6261d0d8b859ea98f2d818b7e67d Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Tue, 22 Aug 2023 19:41:49 -0400
Subject: [PATCH] Don't hardcode node names for image upload widget.

---
 nodes.py                           | 4 ++--
 web/extensions/core/uploadImage.js | 3 ++-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/nodes.py b/nodes.py
index 27a329c6..b2f224ea 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1306,7 +1306,7 @@ class LoadImage:
         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": (sorted(files), {"image_upload": True})},
                 }
 
     CATEGORY = "image"
@@ -1349,7 +1349,7 @@ class LoadImageMask:
         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": (sorted(files), {"image_upload": True}),
                      "channel": (s._color_channels, ), }
                 }
 
diff --git a/web/extensions/core/uploadImage.js b/web/extensions/core/uploadImage.js
index f50473ae..fda83f8c 100644
--- a/web/extensions/core/uploadImage.js
+++ b/web/extensions/core/uploadImage.js
@@ -5,7 +5,8 @@ import { app } from "../../scripts/app.js";
 app.registerExtension({
 	name: "Comfy.UploadImage",
 	async beforeRegisterNodeDef(nodeType, nodeData, app) {
-		if (nodeData.name === "LoadImage" || nodeData.name === "LoadImageMask") {
+		console.log(nodeData);
+		if (nodeData?.input?.required?.image?.[1]?.image_upload === true) {
 			nodeData.input.required.upload = ["IMAGEUPLOAD"];
 		}
 	},