Support execution control feature.

ExecutionSwitch, ExecutionBlocker, ExecutionOneOf

comfy_extras/nodes_execution_control.py

@@ -0,0 +1,179 @@
+class LoopControl:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                        "loop_condition": ("LOOP_CONDITION", ),
+                        "initial_input": ("*", ),
+                        "loopback_input": ("*", ),
+                    },
+                }
+
+    RETURN_TYPES = ("*", )
+    FUNCTION = "doit"
+
+    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']
+
+        return (kwargs['loop_condition'].get_next(kwargs['initial_input'], current), )
+
+
+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"
+
+    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"
+
+    def doit(s, input1, input2):
+        return ((input1, input2), )
+
+
+class InputUnzip:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                        "zipped_input": ("*", ),
+                    },
+                }
+
+    RETURN_TYPES = ("*", "*", )
+    FUNCTION = "doit"
+
+    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"
+
+    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"
+
+    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"
+
+    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,
+}

comfy_extras/nodes_loop.py

@@ -1,94 +0,0 @@
-class LoopControl:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-                        "loop_condition": ("LOOP_CONDITION", ),
-                        "initial_input": ("*", ),
-                        "loopback_input": ("*", ),
-                    },
-                }
-
-    RETURN_TYPES = ("*", )
-    FUNCTION = "doit"
-
-    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']
-
-        return (kwargs['loop_condition'].get_next(kwargs['initial_input'], current), )
-
-
-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"
-
-    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"
-
-    def doit(s, input1, input2):
-        return ((input1, input2), )
-
-
-class InputUnzip:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-                        "zipped_input": ("*", ),
-                    },
-                }
-
-    RETURN_TYPES = ("*", "*", )
-    FUNCTION = "doit"
-
-    def doit(s, zipped_input):
-        input1, input2 = zipped_input
-        return (input1, input2, )
-
-
-NODE_CLASS_MAPPINGS = {
-    "LoopControl": LoopControl,
-    "LoopCounterCondition": LoopCounterCondition,
-    "InputZip": InputZip,
-    "InputUnzip": InputUnzip,
-}

execution.py

@@ -21,13 +21,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 class_def.__name__ != "LoopControl":
+            if input_unique_id in outputs and len(outputs[input_unique_id]) == 0:
-                if input_unique_id not in outputs or outputs[input_unique_id][input_data[1]] == [None]:
+                input_data_all[x] = []
-                    return None
+            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]:
+                if input_unique_id in outputs and outputs[input_unique_id][input_data[1]] != [None]:
-                obj = outputs[input_unique_id][output_index]
+                    obj = outputs[input_unique_id][output_index]
-                input_data_all[x] = obj
+                    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]

nodes.py

@@ -1459,5 +1459,5 @@ def init_custom_nodes():
     load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_post_processing.py"))
     load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_mask.py"))
     load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_rebatch.py"))
-    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_loop.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()

worklist_execution.py

@@ -86,6 +86,19 @@ def is_incomplete_input_slots(class_def, inputs, outputs):
     if len(required_inputs - inputs.keys()) > 0:
         return True
 
+    # "ExecutionOneof" node is a special node that allows only one of the multiple execution paths to be reached and passed through.
+    if class_def.__name__ == "ExecutionOneOf":
+        for x in inputs:
+            input_data = inputs[x]
+
+            if isinstance(input_data, list):
+                input_unique_id = input_data[0]
+                if input_unique_id in outputs and outputs[input_unique_id][input_data[1]] != [None]:
+                    return False
+
+        return True
+
+    # The "LoopControl" is a special node that can be executed even without loopback_input.
     if class_def.__name__ == "LoopControl":
         inputs = {
                     'loop_condition': inputs['loop_condition'],
@@ -97,7 +110,7 @@ def is_incomplete_input_slots(class_def, inputs, outputs):
 
         if isinstance(input_data, list):
             input_unique_id = input_data[0]
-            if input_unique_id not in outputs:
+            if input_unique_id not in outputs or outputs[input_unique_id][input_data[1]] == [None]:
                 return True
 
     return False