|
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", |
|
} |
|
|
