mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-11 06:40:48 +08:00
6d09dd70f8
Additionally, if `VALIDATE_INPUTS` takes an argument named `input_types`, that variable will be a dictionary of the socket type of all incoming connections. If that argument exists, normal socket type validation will not occur. This removes the last hurdle for enabling variant types entirely from custom nodes, so I've removed that command-line option. I've added appropriate unit tests for these changes.
365 lines
9.9 KiB
Python
365 lines
9.9 KiB
Python
from comfy.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": {
|
|
"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 = "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, **{("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(),
|
|
}
|
|
|
|
@VariantSupport()
|
|
class TestForLoopClose:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"flow_control": ("FLOW_CONTROL", {"rawLink": True}),
|
|
},
|
|
"optional": {
|
|
"initial_value%d" % i: ("*",{"rawLink": 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 = "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 = {("initial_value%d" % i): kwargs.get("initial_value%d" % 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": {
|
|
"value%d" % 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 "value%d" % i in kwargs:
|
|
result.append(kwargs["value%d" % 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",
|
|
}
|