From c8e01b0d8bdf3cbfb959d6a678304f9d7c6405b4 Mon Sep 17 00:00:00 2001
From: Chris <chris.goringe@gmail.com>
Date: Mon, 11 Sep 2023 17:53:44 +1000
Subject: [PATCH] Merge branch 'node-ui-enhancements'

---
 comfy_extras/ui_decorator.py | 42 ++++++++++++++++++++++++++++++++++++
 server.py                    |  1 +
 web/scripts/app.js           |  1 +
 3 files changed, 44 insertions(+)
 create mode 100644 comfy_extras/ui_decorator.py

diff --git a/comfy_extras/ui_decorator.py b/comfy_extras/ui_decorator.py
new file mode 100644
index 000000000..a9090ebca
--- /dev/null
+++ b/comfy_extras/ui_decorator.py
@@ -0,0 +1,42 @@
+def ui_signal(signals:str|list[str]):
+    """
+    Return a decorator for Node classes.
+    @param signals - a list of strings that name the signals to be sent to the UI. 
+    (For convenience, a string gets converted to a list of length 1)
+
+    The decorator performs the following:
+    The class has OUTPUT_NODE set to True.
+    The class UI_OUTPUT is appended (or created) with a comma separated list of these signals
+    The class FUNCTION is wrapped such that the last len(signals) are removed, and added to the
+    ui dictionary using signals as keys.
+
+    So ui_signals(["first","second"]) will wrap a function returning (something, somethingelse, first_signal, second_signal)
+    and will return { "ui": {"first":first_signal, "second":second_signal}, "result":(something, somethingelse) }
+    """
+    signals:iter = [signals,] if isinstance(signals,str) else signals
+    def decorator(clazz):
+        internal_function_name = getattr(clazz,'FUNCTION')
+        if internal_function_name=='_ui_signal_decorated_function':
+            raise Exception("Can't nest ui_signal decorators")
+        def _ui_signal_decorated_function(self, **kwargs):
+            returns = getattr(self,internal_function_name)(**kwargs)
+            returns_tuple = returns['result']    if isinstance(returns,dict) else returns
+            returns_ui    = returns.get('ui',{}) if isinstance(returns,dict) else {}
+
+            popped_returns = returns_tuple[-len(signals):]
+            returns_tuple  = returns_tuple[:-len(signals)]
+
+            for i,key in enumerate(signals):
+                returns_ui['key'] = popped_returns[i]
+
+            return { "ui":returns_ui, "result": returns_tuple }
+        clazz._ui_signal_decorated_function = _ui_signal_decorated_function
+        clazz.FUNCTION = '_ui_signal_decorated_function'
+        clazz.OUTPUT_NODE = True
+        clazz.UI_OUTPUT = clazz.UI_OUTPUT+"," if hasattr(clazz, 'UI_OUTPUT') else ""
+        clazz.UI_OUTPUT += ",".join(signals)
+        return clazz
+
+    return decorator
+        
+            
diff --git a/server.py b/server.py
index d04060499..1029b42fe 100644
--- a/server.py
+++ b/server.py
@@ -399,6 +399,7 @@ class PromptServer():
             info['name'] = node_class
             info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[node_class] if node_class in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else node_class
             info['description'] = obj_class.DESCRIPTION if hasattr(obj_class,'DESCRIPTION') else ''
+            info['ui_output'] = obj_class.UI_OUTPUT if hasattr(obj_class, 'UI_OUTPUT') else ''
             info['category'] = 'sd'
             if hasattr(obj_class, 'OUTPUT_NODE') and obj_class.OUTPUT_NODE == True:
                 info['output_node'] = True
diff --git a/web/scripts/app.js b/web/scripts/app.js
index 6dd1f3edd..99e9d065e 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -1277,6 +1277,7 @@ export class ComfyApp {
 				}
 			);
 			node.prototype.comfyClass = nodeData.name;
+			node.prototype.ui_output = nodeData.ui_output;
 
 			this.#addNodeContextMenuHandler(node);
 			this.#addDrawBackgroundHandler(node, app);