diff --git a/custom_nodes/execution-inversion-demo-comfyui/__init__.py b/custom_nodes/execution-inversion-demo-comfyui/__init__.py index 00356fc41..d8e6b875f 100644 --- a/custom_nodes/execution-inversion-demo-comfyui/__init__.py +++ b/custom_nodes/execution-inversion-demo-comfyui/__init__.py @@ -1,6 +1,7 @@ from .nodes import GENERAL_NODE_CLASS_MAPPINGS, GENERAL_NODE_DISPLAY_NAME_MAPPINGS from .components import setup_js, COMPONENT_NODE_CLASS_MAPPINGS, COMPONENT_NODE_DISPLAY_NAME_MAPPINGS from .flow_control import FLOW_CONTROL_NODE_CLASS_MAPPINGS, FLOW_CONTROL_NODE_DISPLAY_NAME_MAPPINGS +from .utility_nodes import UTILITY_NODE_CLASS_MAPPINGS, UTILITY_NODE_DISPLAY_NAME_MAPPINGS # NODE_CLASS_MAPPINGS = GENERAL_NODE_CLASS_MAPPINGS.update(COMPONENT_NODE_CLASS_MAPPINGS) # NODE_DISPLAY_NAME_MAPPINGS = GENERAL_NODE_DISPLAY_NAME_MAPPINGS.update(COMPONENT_NODE_DISPLAY_NAME_MAPPINGS) @@ -9,11 +10,13 @@ NODE_CLASS_MAPPINGS = {} NODE_CLASS_MAPPINGS.update(GENERAL_NODE_CLASS_MAPPINGS) NODE_CLASS_MAPPINGS.update(COMPONENT_NODE_CLASS_MAPPINGS) NODE_CLASS_MAPPINGS.update(FLOW_CONTROL_NODE_CLASS_MAPPINGS) +NODE_CLASS_MAPPINGS.update(UTILITY_NODE_CLASS_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS.update(GENERAL_NODE_DISPLAY_NAME_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(COMPONENT_NODE_DISPLAY_NAME_MAPPINGS) NODE_DISPLAY_NAME_MAPPINGS.update(FLOW_CONTROL_NODE_DISPLAY_NAME_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(UTILITY_NODE_DISPLAY_NAME_MAPPINGS) setup_js() diff --git a/custom_nodes/execution-inversion-demo-comfyui/flow_control.py b/custom_nodes/execution-inversion-demo-comfyui/flow_control.py index 7876ae488..d1eef313a 100644 --- a/custom_nodes/execution-inversion-demo-comfyui/flow_control.py +++ b/custom_nodes/execution-inversion-demo-comfyui/flow_control.py @@ -19,6 +19,7 @@ class WhileLoopOpen: return inputs RETURN_TYPES = tuple(["FLOW_CONTROL"] + ["*"] * NUM_FLOW_SOCKETS) + RETURN_NAMES = tuple(["FLOW_CONTROL"] + ["value%d" % i for i in range(NUM_FLOW_SOCKETS)]) FUNCTION = "while_loop_open" CATEGORY = "Flow Control" @@ -37,7 +38,7 @@ class WhileLoopClose: def INPUT_TYPES(cls): inputs = { "required": { - "flow_control": ("FLOW_CONTROL",), + "flow_control": ("FLOW_CONTROL", {"raw_link": True}), "condition": ("INT", {"default": 0, "min": 0, "max": 1, "step": 1}), }, "optional": { @@ -52,6 +53,7 @@ class WhileLoopClose: return inputs RETURN_TYPES = tuple(["*"] * NUM_FLOW_SOCKETS) + RETURN_NAMES = tuple(["value%d" % i for i in range(NUM_FLOW_SOCKETS)]) FUNCTION = "while_loop_close" CATEGORY = "Flow Control" @@ -92,7 +94,7 @@ class WhileLoopClose: self.explore_dependencies(unique_id, dynprompt, upstream) contained = {} - open_node = this_node["inputs"]["flow_control"][0] + open_node = flow_control[0] self.collect_contained(open_node, upstream, contained) contained[unique_id] = True contained[open_node] = True diff --git a/custom_nodes/execution-inversion-demo-comfyui/utility_nodes.py b/custom_nodes/execution-inversion-demo-comfyui/utility_nodes.py new file mode 100644 index 000000000..7ac74a846 --- /dev/null +++ b/custom_nodes/execution-inversion-demo-comfyui/utility_nodes.py @@ -0,0 +1,246 @@ +import torch +from comfy.graph_utils import GraphBuilder + +class AccumulateNode: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "to_add": ("*",), + }, + "optional": { + "accumulation": ("ACCUMULATION",), + }, + } + + RETURN_TYPES = ("ACCUMULATION",) + FUNCTION = "accumulate" + + CATEGORY = "InversionDemo Nodes" + + def accumulate(self, to_add, accumulation = None): + if accumulation is None: + value = [to_add] + else: + value = accumulation["accum"] + [to_add] + return ({"accum": value},) + +class AccumulationHeadNode: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "accumulation": ("ACCUMULATION",), + }, + } + + RETURN_TYPES = ("ACCUMULATION", "*",) + FUNCTION = "accumulation_head" + + CATEGORY = "InversionDemo Nodes" + + def accumulation_head(self, accumulation): + accum = accumulation["accum"] + if len(accum) == 0: + return (accumulation, None) + else: + return ({"accum": accum[1:]}, accum[0]) + +class AccumulationTailNode: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "accumulation": ("ACCUMULATION",), + }, + } + + RETURN_TYPES = ("ACCUMULATION", "*",) + FUNCTION = "accumulation_tail" + + CATEGORY = "InversionDemo Nodes" + + def accumulation_tail(self, accumulation): + accum = accumulation["accum"] + if len(accum) == 0: + return (None, accumulation) + else: + return ({"accum": accum[:-1]}, accum[-1]) + +class AccumulationToListNode: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "accumulation": ("ACCUMULATION",), + }, + } + + RETURN_TYPES = ("*",) + OUTPUT_IS_LIST = (True,) + + FUNCTION = "accumulation_to_list" + + CATEGORY = "InversionDemo Nodes" + + def accumulation_to_list(self, accumulation): + return (accumulation["accum"],) + +class ListToAccumulationNode: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "list": ("*",), + }, + } + + RETURN_TYPES = ("*",) + INPUT_IS_LIST = (True,) + + FUNCTION = "list_to_accumulation" + + CATEGORY = "InversionDemo Nodes" + + def accumulation_to_list(self, list): + return ({"accum": list},) + +class IsTruthyNode: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "value": ("*",), + }, + } + + RETURN_TYPES = ("INT",) + FUNCTION = "is_truthy" + + CATEGORY = "InversionDemo Nodes" + + def is_truthy(self, value): + if isinstance(value, torch.Tensor): + if value.max().item() == 0 and value.min().item() == 0: + return (0,) + else: + return (1,) + try: + return (int(bool(value)),) + except: + # Can't convert it? Well then it's something or other. I dunno, I'm not a Python programmer. + return (1,) + +from .flow_control import NUM_FLOW_SOCKETS +class ForLoopOpen: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "remaining": ("INT", {"default": 1, "min": 0, "max": 100000, "step": 1}), + }, + "optional": { + "initial_value%d" % i: ("*",) for i in range(1, NUM_FLOW_SOCKETS) + }, + "hidden": { + "initial_value0": ("*",) + } + } + + RETURN_TYPES = tuple(["FLOW_CONTROL", "INT",] + ["*"] * (NUM_FLOW_SOCKETS-1)) + RETURN_NAMES = tuple(["flow_control", "remaining"] + ["value%d" % i for i in range(1, NUM_FLOW_SOCKETS)]) + FUNCTION = "for_loop_open" + + CATEGORY = "Flow Control" + + def for_loop_open(self, remaining, **kwargs): + graph = GraphBuilder() + if "initial_value0" in kwargs: + remaining = kwargs["initial_value0"] + while_open = graph.node("WhileLoopOpen", condition=remaining, initial_value0=remaining, **{("initial_value%d" % i): kwargs.get("initial_value%d" % i, None) for i in range(1, NUM_FLOW_SOCKETS)}) + outputs = [kwargs.get("initial_value%d" % i, None) for i in range(1, NUM_FLOW_SOCKETS)] + return { + "result": tuple(["stub", remaining] + outputs), + "expand": graph.finalize(), + } + +class ForLoopClose: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "flow_control": ("FLOW_CONTROL", {"raw_link": True}), + "old_remaining": ("INT", {"default": 1, "min": 0, "max": 100000, "step": 1}), + }, + "optional": { + "initial_value%d" % i: ("*",{"raw_link": True}) for i in range(1, NUM_FLOW_SOCKETS) + }, + } + + RETURN_TYPES = tuple(["*"] * (NUM_FLOW_SOCKETS-1)) + RETURN_NAMES = tuple(["value%d" % i for i in range(1, NUM_FLOW_SOCKETS)]) + FUNCTION = "for_loop_close" + + CATEGORY = "Flow Control" + + def for_loop_close(self, flow_control, old_remaining, **kwargs): + graph = GraphBuilder() + while_open = flow_control[0] + # TODO - Requires WAS-ns. Will definitely want to solve before merging + sub = graph.node("Number Operation", operation="subtraction", number_a=[while_open,1], number_b=1) + input_values = {("initial_value%d" % i): kwargs.get("initial_value%d" % i, None) for i in range(1, NUM_FLOW_SOCKETS)} + while_close = graph.node("WhileLoopClose", + flow_control=flow_control, + condition=sub.out(0), + initial_value0=sub.out(0), + **input_values) + return { + "result": tuple([while_close.out(i) for i in range(1, NUM_FLOW_SOCKETS)]), + "expand": graph.finalize(), + } + + +UTILITY_NODE_CLASS_MAPPINGS = { + "AccumulateNode": AccumulateNode, + "AccumulationHeadNode": AccumulationHeadNode, + "AccumulationTailNode": AccumulationTailNode, + "AccumulationToListNode": AccumulationToListNode, + "ListToAccumulationNode": ListToAccumulationNode, + "IsTruthyNode": IsTruthyNode, + "ForLoopOpen": ForLoopOpen, + "ForLoopClose": ForLoopClose, +} +UTILITY_NODE_DISPLAY_NAME_MAPPINGS = { + "AccumulateNode": "Accumulate", + "AccumulationHeadNode": "Accumulation Head", + "AccumulationTailNode": "Accumulation Tail", + "AccumulationToListNode": "Accumulation to List", + "ListToAccumulationNode": "List to Accumulation", + "IsTruthyNode": "Is Truthy", + "ForLoopOpen": "For Loop Open", + "ForLoopClose": "For Loop Close", +}