From fce46389fc6437ec648866e4df3ab1a8808d8b3f Mon Sep 17 00:00:00 2001
From: "Dr.Lt.Data" <dr.lt.data@gmail.com>
Date: Sun, 26 Mar 2023 22:05:07 +0900
Subject: [PATCH] make more clean code for img-send

---
 nodes.py           | 15 ++++++++-------
 server.py          |  2 ++
 web/scripts/app.js |  2 +-
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/nodes.py b/nodes.py
index f347da195..d676c97e9 100644
--- a/nodes.py
+++ b/nodes.py
@@ -727,10 +727,11 @@ class SaveImage:
         return {"required": 
                     {"images": ("IMAGE", ),
                      "filename_prefix": ("STRING", {"default": "ComfyUI"}) },
-                "widget": {"send to img": ("IMAGESEND", )},
                 "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                 }
 
+    WIDGET_TYPES = {"send to img": ("IMAGESEND", )}
+
     RETURN_TYPES = ()
     FUNCTION = "save_images"
 
@@ -808,9 +809,9 @@ class LoadImage:
     def INPUT_TYPES(s):
         if not os.path.exists(s.input_dir):
             os.makedirs(s.input_dir)
-        return {"required": {"image": (sorted(os.listdir(s.input_dir)), ) }, 
-                "widget": {"recv img": (["disable", "enable"], ) }
-                }
+        return {"required": {"image": (sorted(os.listdir(s.input_dir)), ) }}
+
+    WIDGET_TYPES = {"recv img": (["disable", "enable"], )}
 
     CATEGORY = "image"
 
@@ -843,9 +844,9 @@ class LoadImageMask:
     def INPUT_TYPES(s):
         return {"required":
                     {"image": (sorted(os.listdir(s.input_dir)), ),
-                    "channel": (["alpha", "red", "green", "blue"], ),},
-                "widget": {"recv img": (["disable", "enable"], ) }
-                }
+                    "channel": (["alpha", "red", "green", "blue"], ),}}
+
+    WIDGET_TYPES = { "recv img": (["disable", "enable"], ) }
 
     CATEGORY = "image"
 
diff --git a/server.py b/server.py
index 6a640fcc7..0a83295c3 100644
--- a/server.py
+++ b/server.py
@@ -163,6 +163,8 @@ class PromptServer():
                 info = {}
                 info['input'] = obj_class.INPUT_TYPES()
                 info['output'] = obj_class.RETURN_TYPES
+                if hasattr(obj_class, 'WIDGET_TYPES'):
+                    info['widget'] = obj_class.WIDGET_TYPES
                 info['name'] = x #TODO
                 info['description'] = ''
                 info['category'] = 'sd'
diff --git a/web/scripts/app.js b/web/scripts/app.js
index c84859535..cdac43a5e 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -565,7 +565,7 @@ class ComfyApp {
 			const nodeData = defs[nodeId];
 			const node = Object.assign(
 				function ComfyNode() {
-					const inputs = Object.assign(nodeData["input"]["required"], nodeData["input"]["widget"]);
+					const inputs = Object.assign(nodeData["input"]["required"], nodeData["widget"]);
 					const config = { minWidth: 1, minHeight: 1 };
 					for (const inputName in inputs) {
 						const inputData = inputs[inputName];