From ab77967ae68110485e4d040f0863263e146a02b4 Mon Sep 17 00:00:00 2001 From: Chris Date: Mon, 11 Sep 2023 14:55:34 +1000 Subject: [PATCH] decorator for UI outputs --- comfy_extras/ui_decorator.py | 41 ++++++++++++++++++++++++++++++++++++ 1 file changed, 41 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..c5bd339a1 --- /dev/null +++ b/comfy_extras/ui_decorator.py @@ -0,0 +1,41 @@ +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 decorator + +