500 lines
18 KiB
Python
500 lines
18 KiB
Python
"""
|
|
Utilities for loading and processing ComfyUI workflows.
|
|
Supports both UI format (saved normally via Ctrl+S) and API format.
|
|
"""
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
from comfy_execution.graph_utils import GraphBuilder
|
|
|
|
log = logging.getLogger("ComfyUI-Workflow-Functions")
|
|
|
|
_UUID_RE = re.compile(
|
|
r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.I
|
|
)
|
|
|
|
WFF_FUNC_INPUT = "WFF_FunctionInput"
|
|
WFF_FUNC_OUTPUT = "WFF_FunctionOutput"
|
|
MAX_SLOTS = 8
|
|
PLACEHOLDER = "[select workflow]"
|
|
|
|
|
|
def _workflows_dir() -> str:
|
|
import folder_paths
|
|
return os.path.join(folder_paths.base_path, "user", "default", "workflows")
|
|
|
|
|
|
def list_workflow_files() -> list[str]:
|
|
d = _workflows_dir()
|
|
files = []
|
|
if os.path.isdir(d):
|
|
for root, _, fnames in os.walk(d):
|
|
for f in fnames:
|
|
if f.endswith(".json"):
|
|
rel = os.path.relpath(os.path.join(root, f), d)
|
|
files.append(rel.replace(os.sep, "/"))
|
|
return [PLACEHOLDER] + sorted(files)
|
|
|
|
|
|
def load_workflow(filename: str) -> dict:
|
|
path = os.path.join(_workflows_dir(), filename)
|
|
with open(path, encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def is_ui_format(data: dict) -> bool:
|
|
return "nodes" in data
|
|
|
|
|
|
def _sort_key(node_id: str):
|
|
try:
|
|
return (0, int(node_id))
|
|
except (ValueError, TypeError):
|
|
return (1, node_id)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# I/O discovery
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def get_workflow_io(data: dict) -> tuple[list[dict], list[dict]]:
|
|
"""
|
|
Return (inputs, outputs) sorted by node id.
|
|
Each entry: {"node_id": str, "slot_name": str}
|
|
"""
|
|
if is_ui_format(data):
|
|
return _get_workflow_io_ui(data)
|
|
return _get_workflow_io_api(data)
|
|
|
|
|
|
def _get_workflow_io_ui(data: dict) -> tuple[list[dict], list[dict]]:
|
|
inputs, outputs = [], []
|
|
for node in data.get("nodes", []):
|
|
ntype = _node_class_type(node)
|
|
nid = str(node.get("id"))
|
|
widgets = node.get("widgets_values") or []
|
|
slot_name = widgets[0] if widgets else nid
|
|
if ntype == WFF_FUNC_INPUT:
|
|
inputs.append({"node_id": nid, "slot_name": slot_name})
|
|
elif ntype == WFF_FUNC_OUTPUT:
|
|
outputs.append({"node_id": nid, "slot_name": slot_name})
|
|
inputs.sort(key=lambda x: _sort_key(x["node_id"]))
|
|
outputs.sort(key=lambda x: _sort_key(x["node_id"]))
|
|
return inputs, outputs
|
|
|
|
|
|
def _get_workflow_io_api(data: dict) -> tuple[list[dict], list[dict]]:
|
|
inputs, outputs = [], []
|
|
for nid, node in data.items():
|
|
if nid.startswith("_"):
|
|
continue
|
|
ct = node.get("class_type", "")
|
|
slot = node.get("inputs", {}).get("slot_name", nid)
|
|
if ct == WFF_FUNC_INPUT:
|
|
inputs.append({"node_id": nid, "slot_name": slot})
|
|
elif ct == WFF_FUNC_OUTPUT:
|
|
outputs.append({"node_id": nid, "slot_name": slot})
|
|
inputs.sort(key=lambda x: _sort_key(x["node_id"]))
|
|
outputs.sort(key=lambda x: _sort_key(x["node_id"]))
|
|
return inputs, outputs
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Widget value resolution (UI format only)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _node_class_type(node: dict) -> str | None:
|
|
"""
|
|
Resolve the ComfyUI class type for a UI-format node.
|
|
|
|
In standard LiteGraph format the class type lives in node["type"].
|
|
In newer ComfyUI builds the type field holds a UUID. For regular nodes
|
|
the real class name is in node["properties"]["Node name for S&R"].
|
|
For group nodes the UUID itself IS the registered class key in
|
|
NODE_CLASS_MAPPINGS, so we fall back to the UUID when the property is absent.
|
|
"""
|
|
ct = node.get("type")
|
|
if not ct:
|
|
return None
|
|
if not _UUID_RE.match(str(ct)):
|
|
return ct # plain class name, use as-is
|
|
# UUID: prefer the human-readable S&R name, fall back to UUID (group nodes)
|
|
return (node.get("properties") or {}).get("Node name for S&R") or ct
|
|
|
|
|
|
def _parse_link(lnk) -> tuple[str, str, int, str, int]:
|
|
"""
|
|
Normalise a link entry to (link_id, src_node_id, src_slot, dst_node_id, dst_slot).
|
|
Handles both the classic array format [id, src, src_slot, dst, dst_slot, type]
|
|
and the newer object format {"id":…, "origin_id":…, "origin_slot":…, …}.
|
|
"""
|
|
if isinstance(lnk, dict):
|
|
return (
|
|
str(lnk.get("id", "")),
|
|
str(lnk.get("origin_id", "")),
|
|
int(lnk.get("origin_slot", 0)),
|
|
str(lnk.get("target_id", "")),
|
|
int(lnk.get("target_slot", 0)),
|
|
)
|
|
return (str(lnk[0]), str(lnk[1]), int(lnk[2]), str(lnk[3]), int(lnk[4]))
|
|
|
|
|
|
def _get_widget_values(class_type: str, linked_names: set, widgets_values: list) -> dict:
|
|
"""
|
|
Return {widget_name: value} by aligning widgets_values with the full ordered
|
|
widget list for the node class (including linked inputs and hidden widgets).
|
|
|
|
widgets_values (from the UI-format JSON) stores ALL widget values in UI order,
|
|
even for inputs that are wired as links. ComfyUI also inserts a hidden
|
|
'control_after_generate' widget after every INT input named 'seed'.
|
|
We must account for both when computing positional offsets.
|
|
"""
|
|
import nodes as _comfy_nodes
|
|
cls = _comfy_nodes.NODE_CLASS_MAPPINGS.get(class_type)
|
|
if cls is None:
|
|
return {}
|
|
try:
|
|
input_types = cls.INPUT_TYPES()
|
|
except Exception:
|
|
return {}
|
|
|
|
# Build the full ordered widget list exactly as ComfyUI saves it,
|
|
# including linked inputs (for positional tracking) and hidden widgets.
|
|
full_order: list[tuple[str, bool]] = [] # (name, should_skip)
|
|
for category in ("required", "optional"):
|
|
for name, type_def in input_types.get(category, {}).items():
|
|
if not isinstance(type_def, (list, tuple)):
|
|
continue
|
|
input_type = type_def[0]
|
|
opts = type_def[1] if len(type_def) > 1 else {}
|
|
if isinstance(opts, dict) and opts.get("forceInput"):
|
|
continue
|
|
if isinstance(input_type, list) or input_type in ("INT", "FLOAT", "STRING", "BOOLEAN"):
|
|
skip = name in linked_names
|
|
full_order.append((name, skip))
|
|
# ComfyUI inserts control_after_generate after seed INT inputs
|
|
if input_type == "INT" and name == "seed":
|
|
full_order.append(("control_after_generate", True))
|
|
|
|
result = {}
|
|
for i, (name, skip) in enumerate(full_order):
|
|
if i >= len(widgets_values):
|
|
break
|
|
if not skip:
|
|
result[name] = widgets_values[i]
|
|
return result
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Subgraph expansion
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def build_expansion(data: dict, outer_inputs: dict):
|
|
"""
|
|
Build a GraphBuilder subgraph from an inner workflow.
|
|
Accepts both UI format and API format.
|
|
|
|
outer_inputs: {"wff_in_0": value, "wff_in_1": value, ...}
|
|
Returns (output_refs, graph).
|
|
"""
|
|
if is_ui_format(data):
|
|
return _build_expansion_ui(data, outer_inputs)
|
|
return _build_expansion_api(data, outer_inputs)
|
|
|
|
|
|
def _build_subgraph_defs(data: dict) -> dict:
|
|
"""Extract subgraph (group node) definitions from workflow data, keyed by UUID."""
|
|
result = {}
|
|
for sg in ((data.get("definitions") or {}).get("subgraphs") or []):
|
|
if isinstance(sg, dict) and sg.get("id"):
|
|
result[sg["id"]] = sg
|
|
return result
|
|
|
|
|
|
def _expand_subgraph(outer_node: dict, sg_def: dict, outer_link_map: dict,
|
|
outer_resolve_fn, graph, id_prefix: str) -> list:
|
|
"""
|
|
Expand a subgraph (group node) definition inline into the GraphBuilder.
|
|
|
|
The subgraph uses two virtual boundary nodes:
|
|
inputNode (id = inputNode.id, e.g. -10): outputs carry each outer input value
|
|
outputNode (id = outputNode.id, e.g. -20): inputs receive the subgraph outputs
|
|
|
|
Returns a list of NodeRefs corresponding to sg_def['outputs'].
|
|
"""
|
|
import nodes as _comfy_nodes
|
|
|
|
sg_nodes = sg_def.get("nodes") or []
|
|
sg_links = sg_def.get("links") or []
|
|
sg_inputs_def = sg_def.get("inputs") or []
|
|
sg_outputs_def = sg_def.get("outputs") or []
|
|
input_node_id = str((sg_def.get("inputNode") or {}).get("id", -10))
|
|
|
|
# Parse inner links.
|
|
parsed = [_parse_link(lnk) for lnk in sg_links]
|
|
sg_link_map = {lid: (src, ss) for lid, src, ss, _, _ in parsed}
|
|
|
|
# Build outer_values_by_slot: {subgraph_input_slot_index → resolved value}
|
|
sg_input_name_to_slot = {inp.get("name"): i for i, inp in enumerate(sg_inputs_def)}
|
|
outer_values_by_slot: dict[int, object] = {}
|
|
for outer_inp in (outer_node.get("inputs") or []):
|
|
name = outer_inp.get("name")
|
|
link_id = outer_inp.get("link")
|
|
if name and link_id is not None:
|
|
slot_idx = sg_input_name_to_slot.get(name)
|
|
if slot_idx is not None:
|
|
src = outer_link_map.get(str(link_id))
|
|
if src:
|
|
outer_values_by_slot[slot_idx] = outer_resolve_fn(src[0], src[1])
|
|
|
|
# Pass 1: register inner nodes (with prefixed IDs to avoid conflicts).
|
|
sg_refs: dict = {}
|
|
for node in sg_nodes:
|
|
nid = str(node.get("id"))
|
|
ct = _node_class_type(node)
|
|
if not ct:
|
|
log.warning("WFF sg: node id=%s has no class_type, skipping", nid)
|
|
continue
|
|
if ct not in _comfy_nodes.NODE_CLASS_MAPPINGS:
|
|
log.warning("WFF sg: node id=%s type=%r not in NODE_CLASS_MAPPINGS, skipping", nid, ct)
|
|
continue
|
|
sg_refs[nid] = graph.node(ct, id=f"{id_prefix}_{nid}")
|
|
|
|
def resolve_sg_link(link_id: str):
|
|
src = sg_link_map.get(str(link_id))
|
|
if src is None:
|
|
return None
|
|
src_id, src_slot = str(src[0]), int(src[1])
|
|
if src_id == input_node_id:
|
|
return outer_values_by_slot.get(src_slot)
|
|
ref = sg_refs.get(src_id)
|
|
if ref is None:
|
|
log.warning("WFF sg: link %s src node %s not in sg_refs", link_id, src_id)
|
|
return None
|
|
return ref.out(src_slot)
|
|
|
|
# Pass 2: wire inner nodes.
|
|
for node in sg_nodes:
|
|
nid = str(node.get("id"))
|
|
gb_node = sg_refs.get(nid)
|
|
if gb_node is None:
|
|
continue
|
|
ct = _node_class_type(node)
|
|
node_inputs = node.get("inputs") or []
|
|
widgets_values = node.get("widgets_values") or []
|
|
|
|
linked_names: set = set()
|
|
for inp in node_inputs:
|
|
link_id = inp.get("link")
|
|
name = inp.get("name")
|
|
if link_id is not None and name:
|
|
resolved = resolve_sg_link(str(link_id))
|
|
if resolved is not None:
|
|
linked_names.add(name)
|
|
gb_node.set_input(name, resolved)
|
|
else:
|
|
src = sg_link_map.get(str(link_id))
|
|
if src and str(src[0]) == input_node_id:
|
|
pass # outer not wired — fall through to inner widget default
|
|
else:
|
|
linked_names.add(name)
|
|
log.warning("WFF sg: node %s (%s) input %r UNRESOLVED", nid, ct, name)
|
|
|
|
for wname, val in _get_widget_values(ct, linked_names, widgets_values).items():
|
|
gb_node.set_input(wname, val)
|
|
|
|
# Collect output refs from the subgraph's virtual outputNode.
|
|
output_refs = []
|
|
for out_def in sg_outputs_def:
|
|
link_ids = out_def.get("linkIds") or []
|
|
if not link_ids:
|
|
log.warning("WFF sg: output %r has no linkIds", out_def.get("name"))
|
|
output_refs.append(None)
|
|
continue
|
|
src = sg_link_map.get(str(link_ids[0]))
|
|
if src is None:
|
|
log.warning("WFF sg: output link %s not in sg_link_map", link_ids[0])
|
|
output_refs.append(None)
|
|
continue
|
|
src_id, src_slot = str(src[0]), int(src[1])
|
|
ref = sg_refs.get(src_id)
|
|
if ref is None:
|
|
log.warning("WFF sg: output src node %s not in sg_refs", src_id)
|
|
output_refs.append(None)
|
|
continue
|
|
output_refs.append(ref.out(src_slot))
|
|
|
|
return output_refs
|
|
|
|
|
|
def _build_expansion_ui(data: dict, outer_inputs: dict):
|
|
nodes_list = data.get("nodes", [])
|
|
links_list = data.get("links", [])
|
|
|
|
subgraph_defs = _build_subgraph_defs(data)
|
|
|
|
parsed_links = [_parse_link(lnk) for lnk in links_list]
|
|
link_map = {lid: (src, ss) for lid, src, ss, _, _ in parsed_links}
|
|
dst_to_src: dict[str, tuple[str, int]] = {}
|
|
for _, src, ss, dst, _ in parsed_links:
|
|
dst_to_src.setdefault(dst, (src, ss))
|
|
|
|
inputs_info, outputs_info = _get_workflow_io_ui(data)
|
|
|
|
fi_node_ids = {inp["node_id"] for inp in inputs_info}
|
|
fo_node_ids = {out["node_id"] for out in outputs_info}
|
|
func_node_ids = fi_node_ids | fo_node_ids
|
|
|
|
fi_value = {inp["node_id"]: outer_inputs.get(f"wff_in_{i}") for i, inp in enumerate(inputs_info)}
|
|
|
|
graph = GraphBuilder()
|
|
node_refs: dict = {}
|
|
subgraph_outputs: dict = {} # nid → [output_ref, ...]
|
|
|
|
def resolve_link(src_node_id: str, src_slot: int):
|
|
if src_node_id in fi_value:
|
|
return fi_value[src_node_id]
|
|
if src_node_id in subgraph_outputs:
|
|
refs = subgraph_outputs[src_node_id]
|
|
return refs[src_slot] if src_slot < len(refs) else None
|
|
ref = node_refs.get(src_node_id)
|
|
if ref is None:
|
|
log.warning("WFF: node %s slot %d unresolvable", src_node_id, src_slot)
|
|
return ref.out(src_slot) if ref is not None else None
|
|
|
|
import nodes as _comfy_nodes
|
|
|
|
# Pass 1: register regular inner nodes (skip WFF boundary and subgraph nodes).
|
|
for node in nodes_list:
|
|
nid = str(node["id"])
|
|
if nid in func_node_ids:
|
|
continue
|
|
ct = _node_class_type(node)
|
|
if not ct:
|
|
log.warning("WFF: node id=%s has no resolvable class_type, skipping", nid)
|
|
continue
|
|
if ct in subgraph_defs:
|
|
continue # expanded separately below
|
|
if ct not in _comfy_nodes.NODE_CLASS_MAPPINGS:
|
|
log.warning("WFF: node id=%s type=%r not in NODE_CLASS_MAPPINGS, skipping", nid, ct)
|
|
continue
|
|
node_refs[nid] = graph.node(ct, id=nid)
|
|
|
|
# Subgraph expansion pass (runs after pass 1 so node_refs is populated).
|
|
for node in nodes_list:
|
|
nid = str(node["id"])
|
|
if nid in func_node_ids:
|
|
continue
|
|
ct = _node_class_type(node)
|
|
if not ct or ct not in subgraph_defs:
|
|
continue
|
|
refs = _expand_subgraph(
|
|
outer_node=node,
|
|
sg_def=subgraph_defs[ct],
|
|
outer_link_map=link_map,
|
|
outer_resolve_fn=resolve_link,
|
|
graph=graph,
|
|
id_prefix=nid,
|
|
)
|
|
subgraph_outputs[nid] = refs
|
|
|
|
# Pass 2: wire inputs for regular inner nodes.
|
|
for node in nodes_list:
|
|
nid = str(node["id"])
|
|
if nid in func_node_ids:
|
|
continue
|
|
gb_node = node_refs.get(nid)
|
|
if gb_node is None:
|
|
continue # subgraph nodes are not in node_refs
|
|
|
|
ct = _node_class_type(node)
|
|
node_inputs = node.get("inputs") or []
|
|
widgets_values = node.get("widgets_values") or []
|
|
|
|
linked_names: set = set()
|
|
for inp in node_inputs:
|
|
link_id = inp.get("link")
|
|
name = inp.get("name")
|
|
if link_id is not None and name:
|
|
linked_names.add(name)
|
|
src = link_map.get(str(link_id))
|
|
if src:
|
|
resolved = resolve_link(src[0], src[1])
|
|
if resolved is not None:
|
|
gb_node.set_input(name, resolved)
|
|
else:
|
|
log.warning("WFF: node %s (%s) input %r unresolved", nid, ct, name)
|
|
else:
|
|
log.warning("WFF: node %s (%s) input %r: link %s not in link_map",
|
|
nid, ct, name, link_id)
|
|
|
|
for wname, val in _get_widget_values(ct, linked_names, widgets_values).items():
|
|
gb_node.set_input(wname, val)
|
|
|
|
# Collect output refs from FunctionOutput nodes.
|
|
output_refs = []
|
|
for out in outputs_info:
|
|
src = dst_to_src.get(out["node_id"])
|
|
if src is None:
|
|
log.warning("WFF: FunctionOutput node=%s has no incoming link", out["node_id"])
|
|
ref = resolve_link(src[0], src[1]) if src else None
|
|
output_refs.append(ref)
|
|
|
|
return output_refs, graph
|
|
|
|
|
|
def _build_expansion_api(data: dict, outer_inputs: dict):
|
|
inputs_info, outputs_info = _get_workflow_io_api(data)
|
|
|
|
fi_value = {inp["node_id"]: outer_inputs.get(f"wff_in_{i}") for i, inp in enumerate(inputs_info)}
|
|
func_node_ids = {inp["node_id"] for inp in inputs_info} | {out["node_id"] for out in outputs_info}
|
|
|
|
graph = GraphBuilder()
|
|
node_refs: dict = {}
|
|
|
|
def resolve_link(link_val):
|
|
src_id = str(link_val[0])
|
|
src_slot = int(link_val[1])
|
|
if src_id in fi_value:
|
|
return fi_value[src_id]
|
|
ref = node_refs.get(src_id)
|
|
if ref is None:
|
|
log.warning("WFF api: node %s slot %d unresolvable", src_id, src_slot)
|
|
return ref.out(src_slot) if ref is not None else None
|
|
|
|
for nid, node in data.items():
|
|
if nid.startswith("_") or nid in func_node_ids:
|
|
continue
|
|
ct = node.get("class_type")
|
|
if ct is None:
|
|
continue
|
|
node_refs[nid] = graph.node(ct, id=nid)
|
|
|
|
for nid, node in data.items():
|
|
if nid.startswith("_") or nid in func_node_ids:
|
|
continue
|
|
gb_node = node_refs.get(nid)
|
|
if gb_node is None:
|
|
continue
|
|
for inp_name, inp_val in node.get("inputs", {}).items():
|
|
if isinstance(inp_val, list) and len(inp_val) == 2 and isinstance(inp_val[0], (str, int)):
|
|
resolved = resolve_link(inp_val)
|
|
if resolved is not None:
|
|
gb_node.set_input(inp_name, resolved)
|
|
else:
|
|
log.warning("WFF api: node %s input %r unresolved link %s", nid, inp_name, inp_val)
|
|
else:
|
|
gb_node.set_input(inp_name, inp_val)
|
|
|
|
output_refs = []
|
|
for out in outputs_info:
|
|
inp_dict = data[out["node_id"]].get("inputs", {})
|
|
val = inp_dict.get("value")
|
|
if isinstance(val, list) and len(val) == 2 and isinstance(val[0], (str, int)):
|
|
output_refs.append(resolve_link(val))
|
|
else:
|
|
output_refs.append(val)
|
|
|
|
return output_refs, graph
|