feature/refactor: Loop system and Worklist execution model

comfy_extras/nodes_execution_control.py

@@ -0,0 +1,191 @@
+class LoopControl:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                        "loop_condition": ("LOOP_CONDITION", ),
+                        "initial_input": ("*", ),
+                        "loopback_input": ("*", ),
+                    },
+                }
+
+    RETURN_TYPES = ("*", )
+    FUNCTION = "doit"
+
+    CATEGORY = "execution"
+    def doit(s, **kwargs):
+        if 'loopback_input' not in kwargs or kwargs['loopback_input'] is None:
+            current = kwargs['initial_input']
+        else:
+            current = kwargs['loopback_input']
+
+        result = kwargs['loop_condition'].get_next(kwargs['initial_input'], current)
+        if result is None:
+            return None
+        else:
+            return (result, )
+
+
+class CounterCondition:
+    def __init__(self, value):
+        self.max = value
+        self.current = 0
+
+    def get_next(self, initial_value, value):
+        print(f"CounterCondition: {self.current}/{self.max}")
+
+        self.current += 1
+        if self.current == 1:
+            return initial_value
+        elif self.current <= self.max:
+            return value
+        else:
+            return None
+
+
+class LoopCounterCondition:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                        "count": ("INT", {"default": 1, "min": 0, "max": 9999999, "step": 1}),
+                        "trigger": (["A", "B"], )
+                    },
+                }
+
+    RETURN_TYPES = ("LOOP_CONDITION", )
+    FUNCTION = "doit"
+
+    CATEGORY = "execution"
+    def doit(s, count, trigger):
+        return (CounterCondition(count), )
+
+
+# To facilitate the use of multiple inputs as loopback inputs, InputZip and InputUnzip are provided.
+class InputZip:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                        "input1": ("*", ),
+                        "input2": ("*", ),
+                    },
+                }
+
+    RETURN_TYPES = ("*", )
+    FUNCTION = "doit"
+
+    CATEGORY = "execution"
+    def doit(s, input1, input2):
+        return ((input1, input2), )
+
+
+class InputUnzip:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                        "zipped_input": ("*", ),
+                    },
+                }
+
+    RETURN_TYPES = ("*", "*", )
+    FUNCTION = "doit"
+
+    CATEGORY = "execution"
+    def doit(s, zipped_input):
+        input1, input2 = zipped_input
+        return (input1, input2, )
+
+
+class ExecutionBlocker:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                        "input": ("*", ),
+                        "signal": ("*", ),
+                     },
+                }
+
+    RETURN_TYPES = ("*", )
+    FUNCTION = "doit"
+
+    CATEGORY = "execution"
+    def doit(s, input, signal):
+        return input
+
+
+class ExecutionOneOf:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                        "input1": ("*", ),
+            },
+            "optional": {
+                        "input2": ("*", ),
+                        "input3": ("*", ),
+                        "input4": ("*", ),
+                        "input5": ("*", ),
+                     },
+               }
+
+    RETURN_TYPES = ("*", )
+    FUNCTION = "doit"
+
+    CATEGORY = "execution"
+    def doit(s, **kwargs):
+        if 'input1' in kwargs and kwargs['input1'] is not None:
+            return (kwargs['input1'], )
+        elif 'input2' in kwargs and kwargs['input2'] is not None:
+            return (kwargs['input2'], )
+        elif 'input3' in kwargs and kwargs['input3'] is not None:
+            return (kwargs['input3'],)
+        elif 'input4' in kwargs and kwargs['input4'] is not None:
+            return (kwargs['input4'],)
+        elif 'input5' in kwargs and kwargs['input5'] is not None:
+            return (kwargs['input5'],)
+        else:
+            return None
+
+
+class ExecutionSwitch:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                        "select": ("INT", {"default": 1, "min": 0, "max": 5}),
+                        "input1": ("*", ),
+                    },
+                "optional": {
+                        "input2_opt": ("*", ),
+                        "input3_opt": ("*", ),
+                        "input4_opt": ("*", ),
+                        "input5_opt": ("*", ),
+                    }
+                }
+
+    RETURN_TYPES = ("*", "*", "*", "*", "*", )
+    FUNCTION = "doit"
+
+    CATEGORY = "execution"
+
+    def doit(s, select, input1, input2_opt=None, input3_opt=None, input4_opt=None, input5_opt=None):
+        if select == 1:
+            return input1, None, None, None, None
+        elif select == 2:
+            return None, input2_opt, None, None, None
+        elif select == 3:
+            return None, None, input3_opt, None, None
+        elif select == 4:
+            return None, None, None, input4_opt, None
+        elif select == 5:
+            return None, None, None, None, input5_opt
+        else:
+            return None, None, None, None, None
+
+
+NODE_CLASS_MAPPINGS = {
+    "ExecutionSwitch": ExecutionSwitch,
+    "ExecutionBlocker": ExecutionBlocker,
+    "ExecutionOneOf": ExecutionOneOf,
+    "LoopControl": LoopControl,
+    "LoopCounterCondition": LoopCounterCondition,
+    "InputZip": InputZip,
+    "InputUnzip": InputUnzip,
+}

execution.py

@@ -20,10 +20,16 @@ def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_da
         if isinstance(input_data, list):
             input_unique_id = input_data[0]
             output_index = input_data[1]
-            if input_unique_id not in outputs:
-                return None
-            obj = outputs[input_unique_id][output_index]
-            input_data_all[x] = obj
+            if input_unique_id in outputs and len(outputs[input_unique_id]) == 0:
+                input_data_all[x] = []
+            else:
+                if class_def.__name__ != "LoopControl" and class_def.__name__ != "ExecutionOneOf":
+                    if input_unique_id not in outputs or outputs[input_unique_id][input_data[1]] == [None]:
+                        return None
+
+                if input_unique_id in outputs and outputs[input_unique_id][input_data[1]] != [None]:
+                    obj = outputs[input_unique_id][output_index]
+                    input_data_all[x] = obj
         else:
             if ("required" in valid_inputs and x in valid_inputs["required"]) or ("optional" in valid_inputs and x in valid_inputs["optional"]):
                 input_data_all[x] = [input_data]
@@ -46,16 +52,20 @@ def map_node_over_list(obj, input_data_all, func, allow_interrupt=False):
     if hasattr(obj, "INPUT_IS_LIST"):
         input_is_list = obj.INPUT_IS_LIST
 
-    if len(input_data_all) == 0:
-        max_len_input = 0
-    else:
+    if input_data_all is not None and len(input_data_all) > 0:
         max_len_input = max([len(x) for x in input_data_all.values()])
-
+    else:
+        max_len_input = 0
+
     # get a slice of inputs, repeat last input when list isn't long enough
     def slice_dict(d, i):
         d_new = dict()
-        for k,v in d.items():
-            d_new[k] = v[i if len(v) > i else -1]
+
+        for k, v in d.items():
+            if not v:
+                continue
+            else:
+                d_new[k] = v[i if len(v) > i else -1]
         return d_new
 
     results = []
@@ -71,15 +81,24 @@ def slice_dict(d, i):
         for i in range(max_len_input):
             if allow_interrupt:
                 nodes.before_node_execution()
-            results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
+
+            params = slice_dict(input_data_all, i)
+
+            if params is None:
+                return None
+
+            results.append(getattr(obj, func)(**params))
     return results
 
+
 def get_output_data(obj, input_data_all):
-
     results = []
     uis = []
     return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)
 
+    if return_values is None:
+        return None, None
+
     for r in return_values:
         if isinstance(r, dict):
             if 'ui' in r:
@@ -90,7 +109,7 @@ def get_output_data(obj, input_data_all):
             results.append(r)
 
     output = []
-    if len(results) > 0:
+    if len(results) > 0 and results[0] is not None:
         # check which outputs need concatenating
         output_is_list = [False] * len(results[0])
         if hasattr(obj, "OUTPUT_IS_LIST"):
@@ -103,7 +122,7 @@ def get_output_data(obj, input_data_all):
             else:
                 output.append([o[i] for o in results])
 
-    ui = dict()    
+    ui = dict()
     if len(uis) > 0:
         ui = {k: [y for x in uis for y in x[k]] for k in uis[0].keys()}
     return output, ui
@@ -210,7 +229,7 @@ def recursive_will_execute(prompt, outputs, current_item):
 
     return will_execute + [unique_id]
 
-def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item):
+def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item, object_storage):
     unique_id = current_item
     inputs = prompt[unique_id]['inputs']
     class_type = prompt[unique_id]['class_type']
@@ -219,6 +238,12 @@ def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item
     is_changed_old = ''
     is_changed = ''
     to_delete = False
+
+    obj = object_storage.get((unique_id, class_type), None)
+    if obj is None:
+        obj = class_def()
+        object_storage[(unique_id, class_type)] = obj
+
     if hasattr(class_def, 'IS_CHANGED'):
         if unique_id in old_prompt and 'is_changed' in old_prompt[unique_id]:
             is_changed_old = old_prompt[unique_id]['is_changed']
@@ -227,7 +252,7 @@ def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item
             if input_data_all is not None:
                 try:
                     #is_changed = class_def.IS_CHANGED(**input_data_all)
-                    is_changed = map_node_over_list(class_def, input_data_all, "IS_CHANGED")
+                    is_changed = map_node_over_list(obj, input_data_all, "IS_CHANGED")
                     prompt[unique_id]['is_changed'] = is_changed
                 except:
                     to_delete = True
@@ -250,7 +275,7 @@ def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item
                     input_unique_id = input_data[0]
                     output_index = input_data[1]
                     if input_unique_id in outputs:
-                        to_delete = recursive_output_delete_if_changed(prompt, old_prompt, outputs, input_unique_id)
+                        to_delete = recursive_output_delete_if_changed(prompt, old_prompt, outputs, input_unique_id, object_storage)
                     else:
                         to_delete = True
                     if to_delete:
@@ -346,7 +371,7 @@ def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
                 del d
 
             for x in prompt:
-                recursive_output_delete_if_changed(prompt, self.old_prompt, self.outputs, x)
+                recursive_output_delete_if_changed(prompt, self.old_prompt, self.outputs, x, self.object_storage)
 
             current_outputs = set(self.outputs.keys())
             for x in list(self.outputs_ui.keys()):
@@ -382,7 +407,12 @@ def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
 
 
 
-def validate_inputs(prompt, item, validated):
+def validate_inputs(prompt, item, validated, visited=set()):
+    if item in visited:
+        return (True, [], item)
+    else:
+        visited.add(item)
+
     unique_id = item
     if unique_id in validated:
         return validated[unique_id]
@@ -431,7 +461,7 @@ def validate_inputs(prompt, item, validated):
             o_id = val[0]
             o_class_type = prompt[o_id]['class_type']
             r = nodes.NODE_CLASS_MAPPINGS[o_class_type].RETURN_TYPES
-            if r[val[1]] != type_input:
+            if r[val[1]] != type_input and r[val[1]] != '*' and type_input != '*':
                 received_type = r[val[1]]
                 details = f"{x}, {received_type} != {type_input}"
                 error = {
@@ -448,7 +478,7 @@ def validate_inputs(prompt, item, validated):
                 errors.append(error)
                 continue
             try:
-                r = validate_inputs(prompt, o_id, validated)
+                r = validate_inputs(prompt, o_id, validated, visited)
                 if r[0] is False:
                     # `r` will be set in `validated[o_id]` already
                     valid = False

main.py

@@ -46,6 +46,7 @@ def execute_script(script_path):
 # Main code
 import asyncio
 import itertools
+import os
 import shutil
 import threading
 import gc
@@ -67,6 +68,8 @@ def execute_script(script_path):
 import yaml
 
 import execution
+import worklist_execution
+import folder_paths
 import server
 from server import BinaryEventTypes
 from nodes import init_custom_nodes
@@ -84,7 +87,7 @@ def cuda_malloc_warning():
             print("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n")
 
 def prompt_worker(q, server):
-    e = execution.PromptExecutor(server)
+    e = worklist_execution.PromptExecutor(server)
     while True:
         item, item_id = q.get()
         execution_start_time = time.perf_counter()

nodes.py

@@ -1697,4 +1697,5 @@ def init_custom_nodes():
     load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_tomesd.py"))
     load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_clip_sdxl.py"))
     load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_canny.py"))
+    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_execution_control.py"))
     load_custom_nodes()