diff --git a/comfy_extras/nodes_execution_control.py b/comfy_extras/nodes_execution_control.py new file mode 100644 index 000000000..3609b8257 --- /dev/null +++ b/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, +} diff --git a/comfy_extras/nodes_loop.py b/comfy_extras/nodes_loop.py deleted file mode 100644 index dbf5d50ab..000000000 --- a/comfy_extras/nodes_loop.py +++ /dev/null @@ -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, -} diff --git a/execution.py b/execution.py index cab38b470..a8e04333f 100644 --- a/execution.py +++ b/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 not in outputs or outputs[input_unique_id][input_data[1]] == [None]: - return None + 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 + 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] diff --git a/nodes.py b/nodes.py index 39d692eec..9390eef47 100644 --- a/nodes.py +++ b/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() diff --git a/worklist_execution.py b/worklist_execution.py index a9d74ded6..6fa62cce3 100644 --- a/worklist_execution.py +++ b/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