| from comfy_execution.graph_utils import GraphBuilder |
| from .tools import VariantSupport |
|
|
| @VariantSupport() |
| class TestAccumulateNode: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "to_add": ("*",), |
| }, |
| "optional": { |
| "accumulation": ("ACCUMULATION",), |
| }, |
| } |
|
|
| RETURN_TYPES = ("ACCUMULATION",) |
| FUNCTION = "accumulate" |
|
|
| CATEGORY = "Testing/Lists" |
|
|
| def accumulate(self, to_add, accumulation = None): |
| if accumulation is None: |
| value = [to_add] |
| else: |
| value = accumulation["accum"] + [to_add] |
| return ({"accum": value},) |
|
|
| @VariantSupport() |
| class TestAccumulationHeadNode: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "accumulation": ("ACCUMULATION",), |
| }, |
| } |
|
|
| RETURN_TYPES = ("ACCUMULATION", "*",) |
| FUNCTION = "accumulation_head" |
|
|
| CATEGORY = "Testing/Lists" |
|
|
| def accumulation_head(self, accumulation): |
| accum = accumulation["accum"] |
| if len(accum) == 0: |
| return (accumulation, None) |
| else: |
| return ({"accum": accum[1:]}, accum[0]) |
|
|
| class TestAccumulationTailNode: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "accumulation": ("ACCUMULATION",), |
| }, |
| } |
|
|
| RETURN_TYPES = ("ACCUMULATION", "*",) |
| FUNCTION = "accumulation_tail" |
|
|
| CATEGORY = "Testing/Lists" |
|
|
| def accumulation_tail(self, accumulation): |
| accum = accumulation["accum"] |
| if len(accum) == 0: |
| return (None, accumulation) |
| else: |
| return ({"accum": accum[:-1]}, accum[-1]) |
|
|
| @VariantSupport() |
| class TestAccumulationToListNode: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "accumulation": ("ACCUMULATION",), |
| }, |
| } |
|
|
| RETURN_TYPES = ("*",) |
| OUTPUT_IS_LIST = (True,) |
|
|
| FUNCTION = "accumulation_to_list" |
|
|
| CATEGORY = "Testing/Lists" |
|
|
| def accumulation_to_list(self, accumulation): |
| return (accumulation["accum"],) |
|
|
| @VariantSupport() |
| class TestListToAccumulationNode: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "list": ("*",), |
| }, |
| } |
|
|
| RETURN_TYPES = ("ACCUMULATION",) |
| INPUT_IS_LIST = (True,) |
|
|
| FUNCTION = "list_to_accumulation" |
|
|
| CATEGORY = "Testing/Lists" |
|
|
| def list_to_accumulation(self, list): |
| return ({"accum": list},) |
|
|
| @VariantSupport() |
| class TestAccumulationGetLengthNode: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "accumulation": ("ACCUMULATION",), |
| }, |
| } |
|
|
| RETURN_TYPES = ("INT",) |
|
|
| FUNCTION = "accumlength" |
|
|
| CATEGORY = "Testing/Lists" |
|
|
| def accumlength(self, accumulation): |
| return (len(accumulation['accum']),) |
| |
| @VariantSupport() |
| class TestAccumulationGetItemNode: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "accumulation": ("ACCUMULATION",), |
| "index": ("INT", {"default":0, "step":1}) |
| }, |
| } |
|
|
| RETURN_TYPES = ("*",) |
|
|
| FUNCTION = "get_item" |
|
|
| CATEGORY = "Testing/Lists" |
|
|
| def get_item(self, accumulation, index): |
| return (accumulation['accum'][index],) |
| |
| @VariantSupport() |
| class TestAccumulationSetItemNode: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "accumulation": ("ACCUMULATION",), |
| "index": ("INT", {"default":0, "step":1}), |
| "value": ("*",), |
| }, |
| } |
|
|
| RETURN_TYPES = ("ACCUMULATION",) |
|
|
| FUNCTION = "set_item" |
|
|
| CATEGORY = "Testing/Lists" |
|
|
| def set_item(self, accumulation, index, value): |
| new_accum = accumulation['accum'][:] |
| new_accum[index] = value |
| return ({"accum": new_accum},) |
|
|
| class TestIntMathOperation: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "a": ("INT", {"default": 0, "min": -0xffffffffffffffff, "max": 0xffffffffffffffff, "step": 1}), |
| "b": ("INT", {"default": 0, "min": -0xffffffffffffffff, "max": 0xffffffffffffffff, "step": 1}), |
| "operation": (["add", "subtract", "multiply", "divide", "modulo", "power"],), |
| }, |
| } |
|
|
| RETURN_TYPES = ("INT",) |
| FUNCTION = "int_math_operation" |
|
|
| CATEGORY = "Testing/Logic" |
|
|
| def int_math_operation(self, a, b, operation): |
| if operation == "add": |
| return (a + b,) |
| elif operation == "subtract": |
| return (a - b,) |
| elif operation == "multiply": |
| return (a * b,) |
| elif operation == "divide": |
| return (a // b,) |
| elif operation == "modulo": |
| return (a % b,) |
| elif operation == "power": |
| return (a ** b,) |
|
|
|
|
| from .flow_control import NUM_FLOW_SOCKETS |
| @VariantSupport() |
| class TestForLoopOpen: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "remaining": ("INT", {"default": 1, "min": 0, "max": 100000, "step": 1}), |
| }, |
| "optional": { |
| f"initial_value{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"] + [f"value{i}" for i in range(1, NUM_FLOW_SOCKETS)]) |
| FUNCTION = "for_loop_open" |
|
|
| CATEGORY = "Testing/Flow" |
|
|
| def for_loop_open(self, remaining, **kwargs): |
| graph = GraphBuilder() |
| if "initial_value0" in kwargs: |
| remaining = kwargs["initial_value0"] |
| while_open = graph.node("TestWhileLoopOpen", condition=remaining, initial_value0=remaining, **{(f"initial_value{i}"): kwargs.get(f"initial_value{i}", None) for i in range(1, NUM_FLOW_SOCKETS)}) |
| outputs = [kwargs.get(f"initial_value{i}", None) for i in range(1, NUM_FLOW_SOCKETS)] |
| return { |
| "result": tuple(["stub", remaining] + outputs), |
| "expand": graph.finalize(), |
| } |
|
|
| @VariantSupport() |
| class TestForLoopClose: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "flow_control": ("FLOW_CONTROL", {"rawLink": True}), |
| }, |
| "optional": { |
| f"initial_value{i}": ("*",{"rawLink": True}) for i in range(1, NUM_FLOW_SOCKETS) |
| }, |
| } |
|
|
| RETURN_TYPES = tuple(["*"] * (NUM_FLOW_SOCKETS-1)) |
| RETURN_NAMES = tuple([f"value{i}" for i in range(1, NUM_FLOW_SOCKETS)]) |
| FUNCTION = "for_loop_close" |
|
|
| CATEGORY = "Testing/Flow" |
|
|
| def for_loop_close(self, flow_control, **kwargs): |
| graph = GraphBuilder() |
| while_open = flow_control[0] |
| sub = graph.node("TestIntMathOperation", operation="subtract", a=[while_open,1], b=1) |
| cond = graph.node("TestToBoolNode", value=sub.out(0)) |
| input_values = {f"initial_value{i}": kwargs.get(f"initial_value{i}", None) for i in range(1, NUM_FLOW_SOCKETS)} |
| while_close = graph.node("TestWhileLoopClose", |
| flow_control=flow_control, |
| condition=cond.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(), |
| } |
|
|
| NUM_LIST_SOCKETS = 10 |
| @VariantSupport() |
| class TestMakeListNode: |
| def __init__(self): |
| pass |
|
|
| @classmethod |
| def INPUT_TYPES(cls): |
| return { |
| "required": { |
| "value1": ("*",), |
| }, |
| "optional": { |
| f"value{i}": ("*",) for i in range(1, NUM_LIST_SOCKETS) |
| }, |
| } |
|
|
| RETURN_TYPES = ("*",) |
| FUNCTION = "make_list" |
| OUTPUT_IS_LIST = (True,) |
|
|
| CATEGORY = "Testing/Lists" |
|
|
| def make_list(self, **kwargs): |
| result = [] |
| for i in range(NUM_LIST_SOCKETS): |
| if f"value{i}" in kwargs: |
| result.append(kwargs[f"value{i}"]) |
| return (result,) |
|
|
| UTILITY_NODE_CLASS_MAPPINGS = { |
| "TestAccumulateNode": TestAccumulateNode, |
| "TestAccumulationHeadNode": TestAccumulationHeadNode, |
| "TestAccumulationTailNode": TestAccumulationTailNode, |
| "TestAccumulationToListNode": TestAccumulationToListNode, |
| "TestListToAccumulationNode": TestListToAccumulationNode, |
| "TestAccumulationGetLengthNode": TestAccumulationGetLengthNode, |
| "TestAccumulationGetItemNode": TestAccumulationGetItemNode, |
| "TestAccumulationSetItemNode": TestAccumulationSetItemNode, |
| "TestForLoopOpen": TestForLoopOpen, |
| "TestForLoopClose": TestForLoopClose, |
| "TestIntMathOperation": TestIntMathOperation, |
| "TestMakeListNode": TestMakeListNode, |
| } |
| UTILITY_NODE_DISPLAY_NAME_MAPPINGS = { |
| "TestAccumulateNode": "Accumulate", |
| "TestAccumulationHeadNode": "Accumulation Head", |
| "TestAccumulationTailNode": "Accumulation Tail", |
| "TestAccumulationToListNode": "Accumulation to List", |
| "TestListToAccumulationNode": "List to Accumulation", |
| "TestAccumulationGetLengthNode": "Accumulation Get Length", |
| "TestAccumulationGetItemNode": "Accumulation Get Item", |
| "TestAccumulationSetItemNode": "Accumulation Set Item", |
| "TestForLoopOpen": "For Loop Open", |
| "TestForLoopClose": "For Loop Close", |
| "TestIntMathOperation": "Int Math Operation", |
| "TestMakeListNode": "Make List", |
| } |
|
|
