369 lines
13 KiB
Python
369 lines
13 KiB
Python
"""
|
|
Simple For Loop nodes for ComfyUI.
|
|
Clean implementation without the duplicate-index bug from Easy-Use.
|
|
|
|
Nodes:
|
|
- Simple For Loop Start
|
|
- Simple For Loop End
|
|
- SimpleWhileOpen (internal)
|
|
- SimpleWhileClose (internal)
|
|
"""
|
|
|
|
from comfy_execution.graph_utils import GraphBuilder, is_link
|
|
from comfy_execution.graph import ExecutionBlocker
|
|
|
|
try:
|
|
from nodes import NODE_CLASS_MAPPINGS as ALL_NODE_CLASS_MAPPINGS
|
|
except ImportError:
|
|
ALL_NODE_CLASS_MAPPINGS = {}
|
|
|
|
MAX_FLOW_NUM = 5 # value0=index, value1-4=passthrough
|
|
any_type = "*"
|
|
|
|
_LOOP_INITIAL_DONE = set()
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Internal While Loop nodes
|
|
# ─────────────────────────────────────────────
|
|
|
|
class SimpleWhileOpen:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = {
|
|
"required": {
|
|
"condition": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {},
|
|
}
|
|
for i in range(MAX_FLOW_NUM):
|
|
inputs["optional"]["initial_value%d" % i] = (any_type,)
|
|
return inputs
|
|
|
|
RETURN_TYPES = tuple(["FLOW_CONTROL"] + [any_type] * MAX_FLOW_NUM)
|
|
RETURN_NAMES = tuple(["flow"] + ["value%d" % i for i in range(MAX_FLOW_NUM)])
|
|
FUNCTION = "while_loop_open"
|
|
CATEGORY = "comfyui-sortlist/flow"
|
|
|
|
def while_loop_open(self, condition, **kwargs):
|
|
values = []
|
|
for i in range(MAX_FLOW_NUM):
|
|
if condition:
|
|
values.append(kwargs.get("initial_value%d" % i, None))
|
|
else:
|
|
values.append(ExecutionBlocker(None))
|
|
return tuple(["stub"] + values)
|
|
|
|
|
|
class SimpleWhileClose:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = {
|
|
"required": {
|
|
"flow": ("FLOW_CONTROL", {"rawLink": True}),
|
|
"condition": ("BOOLEAN", {}),
|
|
},
|
|
"optional": {},
|
|
"hidden": {
|
|
"dynprompt": "DYNPROMPT",
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
for i in range(MAX_FLOW_NUM):
|
|
inputs["optional"]["initial_value%d" % i] = (any_type,)
|
|
return inputs
|
|
|
|
RETURN_TYPES = tuple([any_type] * MAX_FLOW_NUM)
|
|
RETURN_NAMES = tuple(["value%d" % i for i in range(MAX_FLOW_NUM)])
|
|
FUNCTION = "while_loop_close"
|
|
CATEGORY = "comfyui-sortlist/flow"
|
|
|
|
def explore_dependencies(self, node_id, dynprompt, upstream, parent_ids):
|
|
node_info = dynprompt.get_node(node_id)
|
|
if "inputs" not in node_info:
|
|
return
|
|
for k, v in node_info["inputs"].items():
|
|
if is_link(v):
|
|
parent_id = v[0]
|
|
display_id = dynprompt.get_display_node_id(parent_id)
|
|
display_node = dynprompt.get_node(display_id)
|
|
class_type = display_node["class_type"]
|
|
if class_type not in ["SimpleForLoopEnd", "SimpleWhileClose"]:
|
|
parent_ids.append(display_id)
|
|
if parent_id not in upstream:
|
|
upstream[parent_id] = []
|
|
self.explore_dependencies(parent_id, dynprompt, upstream, parent_ids)
|
|
upstream[parent_id].append(node_id)
|
|
|
|
def explore_output_nodes(self, dynprompt, upstream, output_nodes, parent_ids):
|
|
for parent_id in upstream:
|
|
display_id = dynprompt.get_display_node_id(parent_id)
|
|
for output_id in output_nodes:
|
|
id = output_nodes[output_id][0]
|
|
if id in parent_ids and display_id == id and output_id not in upstream[parent_id]:
|
|
if '.' in parent_id:
|
|
arr = parent_id.split('.')
|
|
arr[len(arr) - 1] = output_id
|
|
upstream[parent_id].append('.'.join(arr))
|
|
else:
|
|
upstream[parent_id].append(output_id)
|
|
|
|
def collect_contained(self, node_id, upstream, contained):
|
|
if node_id not in upstream:
|
|
return
|
|
for child_id in upstream[node_id]:
|
|
if child_id not in contained:
|
|
contained[child_id] = True
|
|
self.collect_contained(child_id, upstream, contained)
|
|
|
|
def while_loop_close(self, flow, condition, dynprompt=None, unique_id=None, **kwargs):
|
|
if not condition:
|
|
# Loop done — pass final values through
|
|
values = []
|
|
for i in range(MAX_FLOW_NUM):
|
|
values.append(kwargs.get("initial_value%d" % i, None))
|
|
return tuple(values)
|
|
|
|
# Continue looping — clone the subgraph for next iteration
|
|
upstream = {}
|
|
parent_ids = []
|
|
self.explore_dependencies(unique_id, dynprompt, upstream, parent_ids)
|
|
parent_ids = list(set(parent_ids))
|
|
|
|
# Find output nodes
|
|
prompts = dynprompt.get_original_prompt()
|
|
output_nodes = {}
|
|
for id in prompts:
|
|
node = prompts[id]
|
|
if "inputs" not in node:
|
|
continue
|
|
class_type = node["class_type"]
|
|
class_def = ALL_NODE_CLASS_MAPPINGS.get(class_type, None)
|
|
if class_def is not None and hasattr(class_def, 'OUTPUT_NODE') and class_def.OUTPUT_NODE:
|
|
for k, v in node['inputs'].items():
|
|
if is_link(v):
|
|
output_nodes[id] = v
|
|
|
|
graph = GraphBuilder()
|
|
self.explore_output_nodes(dynprompt, upstream, output_nodes, parent_ids)
|
|
contained = {}
|
|
open_node = flow[0]
|
|
self.collect_contained(open_node, upstream, contained)
|
|
contained[unique_id] = True
|
|
contained[open_node] = True
|
|
|
|
# Clone all contained nodes into the expand graph
|
|
for node_id in contained:
|
|
original_node = dynprompt.get_node(node_id)
|
|
node = graph.node(
|
|
original_node["class_type"],
|
|
"Recurse" if node_id == unique_id else node_id,
|
|
)
|
|
node.set_override_display_id(node_id)
|
|
|
|
for node_id in contained:
|
|
original_node = dynprompt.get_node(node_id)
|
|
node = graph.lookup_node("Recurse" if node_id == unique_id else node_id)
|
|
for k, v in original_node["inputs"].items():
|
|
if is_link(v) and v[0] in contained:
|
|
parent = graph.lookup_node(v[0])
|
|
node.set_input(k, parent.out(v[1]))
|
|
else:
|
|
node.set_input(k, v)
|
|
|
|
# Feed updated values back into the open node
|
|
new_open = graph.lookup_node(open_node)
|
|
for i in range(MAX_FLOW_NUM):
|
|
key = "initial_value%d" % i
|
|
new_open.set_input(key, kwargs.get(key, None))
|
|
|
|
my_clone = graph.lookup_node("Recurse")
|
|
result = map(lambda x: my_clone.out(x), range(MAX_FLOW_NUM))
|
|
return {
|
|
"result": tuple(result),
|
|
"expand": graph.finalize(),
|
|
}
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# User-facing For Loop nodes
|
|
# ─────────────────────────────────────────────
|
|
|
|
class SimpleForLoopStart:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"total": ("INT", {"default": 1, "min": 1, "max": 100000, "step": 1}),
|
|
},
|
|
"optional": {
|
|
"initial_value%d" % i: (any_type,) for i in range(1, MAX_FLOW_NUM)
|
|
},
|
|
"hidden": {
|
|
"initial_value0": (any_type,),
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = tuple(["FLOW_CONTROL", "INT"] + [any_type] * (MAX_FLOW_NUM - 1))
|
|
RETURN_NAMES = tuple(["flow", "index"] + ["value%d" % i for i in range(1, MAX_FLOW_NUM)])
|
|
FUNCTION = "for_loop_start"
|
|
CATEGORY = "comfyui-sortlist/flow"
|
|
|
|
def for_loop_start(self, total, unique_id=None, **kwargs):
|
|
global _LOOP_INITIAL_DONE
|
|
graph = GraphBuilder()
|
|
|
|
i = 1
|
|
is_recursion = "initial_value0" in kwargs
|
|
if is_recursion:
|
|
i = kwargs["initial_value0"]
|
|
|
|
uid = str(unique_id) if unique_id else "default"
|
|
|
|
# ── Dedup fix ──
|
|
# ComfyUI calls for_loop_start TWICE with same uid:
|
|
# Call 1 = pre-evaluation → BLOCK
|
|
# Call 2 = real call → create expand graph
|
|
if not is_recursion:
|
|
if uid not in _LOOP_INITIAL_DONE:
|
|
_LOOP_INITIAL_DONE.add(uid)
|
|
outputs = [ExecutionBlocker(None)] * (MAX_FLOW_NUM - 1)
|
|
return tuple(["stub", ExecutionBlocker(None)] + outputs)
|
|
else:
|
|
_LOOP_INITIAL_DONE.discard(uid)
|
|
else:
|
|
base_uid = uid.split(".")[-1] if "." in uid else uid
|
|
_LOOP_INITIAL_DONE.discard(base_uid)
|
|
|
|
# Create expand graph with SimpleWhileOpen
|
|
initial_values = {
|
|
"initial_value%d" % num: kwargs.get("initial_value%d" % num, None)
|
|
for num in range(1, MAX_FLOW_NUM)
|
|
}
|
|
graph.node("SimpleWhileOpen", condition=total, initial_value0=i, **initial_values)
|
|
|
|
outputs = [kwargs.get("initial_value%d" % num, None) for num in range(1, MAX_FLOW_NUM)]
|
|
return {
|
|
"result": tuple(["stub", i] + outputs),
|
|
"expand": graph.finalize(),
|
|
}
|
|
|
|
|
|
class _LoopAdd:
|
|
"""Internal helper: a + b (integers)"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"a": ("INT", {"default": 0}),
|
|
"b": ("INT", {"default": 1}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("INT",)
|
|
FUNCTION = "add"
|
|
CATEGORY = "comfyui-sortlist/flow"
|
|
|
|
def add(self, a, b):
|
|
return (a + b,)
|
|
|
|
|
|
class _LoopLessThan:
|
|
"""Internal helper: a < b → boolean"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"a": ("INT", {"default": 0}),
|
|
"b": ("INT", {"default": 1}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("BOOLEAN",)
|
|
FUNCTION = "compare"
|
|
CATEGORY = "comfyui-sortlist/flow"
|
|
|
|
def compare(self, a, b):
|
|
return (a <= b,)
|
|
|
|
|
|
class SimpleForLoopEnd:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"flow": ("FLOW_CONTROL", {"rawLink": True}),
|
|
},
|
|
"optional": {
|
|
"initial_value%d" % i: (any_type, {"rawLink": True})
|
|
for i in range(1, MAX_FLOW_NUM)
|
|
},
|
|
"hidden": {
|
|
"dynprompt": "DYNPROMPT",
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = tuple([any_type] * (MAX_FLOW_NUM - 1))
|
|
RETURN_NAMES = tuple(["value%d" % i for i in range(1, MAX_FLOW_NUM)])
|
|
FUNCTION = "for_loop_end"
|
|
CATEGORY = "comfyui-sortlist/flow"
|
|
|
|
def for_loop_end(self, flow, dynprompt=None, unique_id=None, **kwargs):
|
|
graph = GraphBuilder()
|
|
while_open = flow[0]
|
|
|
|
# Get total from the original ForLoopStart node
|
|
forstart_node = dynprompt.get_node(while_open)
|
|
total = None
|
|
if forstart_node["class_type"] == "SimpleForLoopStart":
|
|
total = forstart_node["inputs"]["total"]
|
|
|
|
if total is None:
|
|
raise ValueError("Simple For Loop End must be connected to Simple For Loop Start")
|
|
|
|
# Increment index: i + 1 (using self-contained helper)
|
|
sub = graph.node("_LoopAdd", a=[while_open, 1], b=1)
|
|
# Check condition: (i + 1) < total
|
|
cond = graph.node("_LoopLessThan", a=sub.out(0), b=total)
|
|
|
|
input_values = {
|
|
"initial_value%d" % i: kwargs.get("initial_value%d" % i, None)
|
|
for i in range(1, MAX_FLOW_NUM)
|
|
}
|
|
|
|
while_close = graph.node(
|
|
"SimpleWhileClose",
|
|
flow=flow,
|
|
condition=cond.out(0),
|
|
initial_value0=sub.out(0),
|
|
**input_values,
|
|
)
|
|
|
|
return {
|
|
"result": tuple([while_close.out(i) for i in range(1, MAX_FLOW_NUM)]),
|
|
"expand": graph.finalize(),
|
|
}
|
|
|
|
|
|
# ─────────────────────────────────────────────
|
|
# Registration
|
|
# ─────────────────────────────────────────────
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SimpleForLoopStart": SimpleForLoopStart,
|
|
"SimpleForLoopEnd": SimpleForLoopEnd,
|
|
"SimpleWhileOpen": SimpleWhileOpen,
|
|
"SimpleWhileClose": SimpleWhileClose,
|
|
"_LoopAdd": _LoopAdd,
|
|
"_LoopLessThan": _LoopLessThan,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SimpleForLoopStart": "Simple For Loop Start A100",
|
|
"SimpleForLoopEnd": "Simple For Loop End A100",
|
|
"SimpleWhileOpen": "Simple While Open",
|
|
"SimpleWhileClose": "Simple While Close",
|
|
"_LoopAdd": "_Loop Add (internal)",
|
|
"_LoopLessThan": "_Loop Less Than (internal)",
|
|
}
|