Jason HokuandClaude Fable 5.1
9ef138bf72
workflow_sweep_runner feeds jobs through ComfyUI's own queue (windowed feeder that
mirrors server.py post_prompt in-process via ComfyBridge), tracks them in history,
pauses on interrupt/queue-clear/3 errors and resumes by fingerprint. USCGSweepCollect
saves each job's image into the USCG session and appends manifest items. Routes
/config_tester/sweep/{parse,dry_run,start,status,pause,resume,cancel}; the Generator
node starts a sweep for workflow configs and returns the dashboard without waiting.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
970 lines
42 KiB
Python
970 lines
42 KiB
Python
"""
|
||
Workflow Sweep: sweep any ComfyUI API-format workflow through ComfyUI's own executor.
|
||
|
||
Pure helpers (no torch / server at import time) for:
|
||
- parsing an API-format workflow and describing it for the Builder UI,
|
||
- suggesting sweep bindings (which literal inputs to turn into axes),
|
||
- expanding bindings into concrete jobs (Cartesian product, cache-friendly order),
|
||
- substituting values into the template and appending the collector node,
|
||
- fingerprinting jobs (resume key) and building manifest items.
|
||
|
||
The runner (`workflow_sweep_runner.py`) feeds the resulting prompts to ComfyUI's queue and the
|
||
`USCGSweepCollect` node (`workflow_sweep_collect.py`) turns each job's image into a manifest item.
|
||
|
||
IMPORTANT: no network code, no subprocess, no eval/exec here (registry yara scan, brief rule 2).
|
||
Workflow JSON is only ever request-body data; nothing in this module reads a user-supplied path.
|
||
"""
|
||
|
||
import copy
|
||
import hashlib
|
||
import itertools
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
import random
|
||
|
||
|
||
COLLECTOR_CLASS = "USCGSweepCollect"
|
||
MAX_JOBS = 50000
|
||
|
||
# input name -> manifest field (class-agnostic; see plan section 4.4)
|
||
ROLE_HINTS = {
|
||
"ckpt_name": "model",
|
||
"unet_name": "model",
|
||
"seed": "seed",
|
||
"noise_seed": "seed",
|
||
"steps": "steps",
|
||
"cfg": "cfg",
|
||
"sampler_name": "sampler",
|
||
"scheduler": "scheduler",
|
||
"denoise": "denoise",
|
||
"vae_name": "vae",
|
||
"clip_name": "text_encoders",
|
||
"clip_name1": "text_encoders",
|
||
"clip_name2": "text_encoders",
|
||
"clip_name3": "text_encoders",
|
||
"clip_name4": "text_encoders",
|
||
}
|
||
# input name -> folder_paths type used for `folder/` syntax inside a literal axis
|
||
FOLDER_TYPES = {
|
||
"ckpt_name": ["checkpoints"],
|
||
"unet_name": ["diffusion_models", "unet"],
|
||
"vae_name": ["vae"],
|
||
"clip_name": ["text_encoders", "clip"],
|
||
"clip_name1": ["text_encoders", "clip"],
|
||
"clip_name2": ["text_encoders", "clip"],
|
||
"clip_name3": ["text_encoders", "clip"],
|
||
"clip_name4": ["text_encoders", "clip"],
|
||
"lora_name": ["loras"],
|
||
}
|
||
# manifest fields a literal axis may feed
|
||
AXIS_FIELDS = ("model", "sampler", "scheduler", "steps", "cfg", "seed", "denoise", "vae")
|
||
# primitive-string-ish input names a prompt text can hide behind
|
||
_TEXT_INPUTS = ("value", "string", "text", "prompt", "multiline_text")
|
||
# text key in a "text link" hop chain, tried in this order when following links
|
||
_FOLLOW_INPUTS = ("source", "on_false", "on_true", "text", "value", "string", "prompt")
|
||
_RES_RE = re.compile(r"^\s*(\d+)\s*[xX×]\s*(\d+)\s*$")
|
||
|
||
|
||
# --------------------------------------------------------------------------- parsing
|
||
|
||
def is_ui_format(obj):
|
||
"""True for the editor (UI) workflow format, which has `nodes` + `links` lists."""
|
||
return isinstance(obj, dict) and isinstance(obj.get("nodes"), list) and isinstance(obj.get("links"), list)
|
||
|
||
|
||
def is_link(value):
|
||
"""API-format link: [node_id, output_index]."""
|
||
return (
|
||
isinstance(value, (list, tuple)) and len(value) == 2
|
||
and isinstance(value[0], (str, int)) and not isinstance(value[0], bool)
|
||
and isinstance(value[1], int) and not isinstance(value[1], bool)
|
||
)
|
||
|
||
|
||
def parse_api_workflow(data):
|
||
"""
|
||
Parse an API-format workflow (ComfyUI "Export (API)"), from a JSON string or an object.
|
||
|
||
Returns the prompt dict {node_id(str): {class_type, inputs, _meta?}} unchanged.
|
||
Raises ValueError with a readable message otherwise.
|
||
"""
|
||
if isinstance(data, (str, bytes)):
|
||
text = data.decode("utf-8", "replace") if isinstance(data, bytes) else data
|
||
if not text.strip():
|
||
raise ValueError("The workflow JSON is empty.")
|
||
try:
|
||
data = json.loads(text)
|
||
except json.JSONDecodeError as e:
|
||
raise ValueError(f"The workflow is not valid JSON: {e}")
|
||
if is_ui_format(data):
|
||
raise ValueError(
|
||
"This is an editor-format workflow (it has 'nodes' and 'links'). "
|
||
"In ComfyUI use Workflow > Export (API) and load that file instead."
|
||
)
|
||
if isinstance(data, dict) and isinstance(data.get("prompt"), dict) and not any(
|
||
isinstance(v, dict) and "class_type" in v for v in data.values()):
|
||
data = data["prompt"] # a saved /prompt request body
|
||
if not isinstance(data, dict) or not data:
|
||
raise ValueError("An API-format workflow is a JSON object of numbered nodes.")
|
||
prompt = {}
|
||
for nid, node in data.items():
|
||
if not isinstance(node, dict) or "class_type" not in node:
|
||
raise ValueError(f"Node '{nid}' has no class_type; this does not look like an API-format workflow.")
|
||
if not isinstance(node.get("inputs", {}), dict):
|
||
raise ValueError(f"Node '{nid}' has malformed inputs.")
|
||
node.setdefault("inputs", {})
|
||
prompt[str(nid)] = node
|
||
return prompt
|
||
|
||
|
||
def node_title(prompt, node_id):
|
||
node = prompt.get(str(node_id)) or {}
|
||
return (node.get("_meta") or {}).get("title") or node.get("class_type") or str(node_id)
|
||
|
||
|
||
def _node_sort_key(nid):
|
||
s = str(nid)
|
||
return (0, int(s), s) if s.isdigit() else (1, 0, s)
|
||
|
||
|
||
def comfy_class_info(class_type):
|
||
"""
|
||
Describe a node class from ComfyUI's registry (V1 and V3 alike), or None when ComfyUI is not
|
||
importable (tests). Returns {missing, return_types, output_node, inputs:{name:{type,options,required}}}.
|
||
"""
|
||
try:
|
||
import nodes # ComfyUI
|
||
cls = nodes.NODE_CLASS_MAPPINGS.get(class_type)
|
||
except Exception:
|
||
return None
|
||
if cls is None:
|
||
return {"missing": True, "return_types": [], "output_node": False, "inputs": {}}
|
||
inputs = {}
|
||
try:
|
||
spec_all = cls.INPUT_TYPES() or {}
|
||
except Exception:
|
||
spec_all = {}
|
||
for section in ("required", "optional"):
|
||
for name, spec in (spec_all.get(section) or {}).items():
|
||
t = spec[0] if isinstance(spec, (list, tuple)) and spec else spec
|
||
options = None
|
||
if isinstance(t, (list, tuple)):
|
||
options, t = list(t), "COMBO"
|
||
elif t == "COMBO" and isinstance(spec, (list, tuple)) and len(spec) > 1 and isinstance(spec[1], dict):
|
||
options = list(spec[1].get("options") or [])
|
||
inputs[name] = {"type": t if isinstance(t, str) else str(t), "options": options,
|
||
"required": section == "required"}
|
||
try:
|
||
return_types = [str(x) for x in (getattr(cls, "RETURN_TYPES", ()) or ())]
|
||
except Exception:
|
||
return_types = []
|
||
try:
|
||
output_node = bool(getattr(cls, "OUTPUT_NODE", False))
|
||
except Exception:
|
||
output_node = False
|
||
return {"missing": False, "return_types": return_types, "output_node": output_node, "inputs": inputs}
|
||
|
||
|
||
def _info(class_info, class_type):
|
||
if class_info is None:
|
||
return None
|
||
try:
|
||
return class_info(class_type)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
# --------------------------------------------------------------------------- graph helpers
|
||
|
||
def ancestors_ordered(prompt, node_id):
|
||
"""Breadth-first list of node ids upstream of node_id, starting with node_id itself.
|
||
Closest nodes come first, so 'first ancestor wins' picks the node nearest the output."""
|
||
start = str(node_id)
|
||
if start not in prompt:
|
||
return []
|
||
seen, order, queue = {start}, [], [start]
|
||
while queue:
|
||
cur = queue.pop(0)
|
||
order.append(cur)
|
||
for value in (prompt[cur].get("inputs") or {}).values():
|
||
if is_link(value):
|
||
up = str(value[0])
|
||
if up in prompt and up not in seen:
|
||
seen.add(up)
|
||
queue.append(up)
|
||
return order
|
||
|
||
|
||
def graph_descendants(prompt, restrict_to=None):
|
||
"""{node_id: number of nodes that (transitively) depend on it}, optionally restricted to a node set."""
|
||
nodes_ = set(restrict_to) if restrict_to is not None else set(prompt)
|
||
downstream = {n: set() for n in nodes_}
|
||
for nid in nodes_:
|
||
for value in (prompt[nid].get("inputs") or {}).values():
|
||
if is_link(value) and str(value[0]) in nodes_:
|
||
downstream[str(value[0])].add(nid)
|
||
memo = {}
|
||
|
||
def reach(n, stack=()):
|
||
if n in memo:
|
||
return memo[n]
|
||
out = set()
|
||
for d in downstream.get(n, ()):
|
||
if d in stack:
|
||
continue
|
||
out.add(d)
|
||
out |= reach(d, stack + (n,))
|
||
memo[n] = out
|
||
return out
|
||
|
||
return {n: len(reach(n)) for n in nodes_}
|
||
|
||
|
||
def find_image_sources(prompt, class_info=None):
|
||
"""
|
||
Every IMAGE socket a collector could be spliced onto, best candidate first:
|
||
the feeder of the first SaveImage, then PreviewImage feeders, then any node whose
|
||
class advertises an IMAGE output (needs ComfyUI's registry via class_info).
|
||
Returns [{node_id, output_index, type, picked_from?}].
|
||
"""
|
||
sources, seen = [], set()
|
||
|
||
def add(nid, idx, picked_from=None):
|
||
key = (str(nid), int(idx))
|
||
if key in seen or str(nid) not in prompt:
|
||
return
|
||
seen.add(key)
|
||
entry = {"node_id": key[0], "output_index": key[1], "type": "IMAGE"}
|
||
if picked_from is not None:
|
||
entry["picked_from"] = str(picked_from)
|
||
sources.append(entry)
|
||
|
||
def feeders(class_pred):
|
||
for nid in sorted(prompt, key=_node_sort_key):
|
||
ct = prompt[nid].get("class_type", "")
|
||
if not class_pred(ct):
|
||
continue
|
||
for iname, value in (prompt[nid].get("inputs") or {}).items():
|
||
if iname in ("images", "image") and is_link(value):
|
||
add(value[0], value[1], picked_from=nid)
|
||
|
||
feeders(lambda ct: ct == "SaveImage" or ct.startswith("SaveImage"))
|
||
feeders(lambda ct: ct == "PreviewImage" or ct.startswith("PreviewImage"))
|
||
for nid in sorted(prompt, key=_node_sort_key):
|
||
info = _info(class_info, prompt[nid].get("class_type"))
|
||
if info and not info.get("missing"):
|
||
for idx, rt in enumerate(info.get("return_types") or []):
|
||
if rt == "IMAGE":
|
||
add(nid, idx)
|
||
return sources
|
||
|
||
|
||
def find_output_nodes(prompt, class_info=None):
|
||
"""Output nodes of the workflow (SaveImage-like or OUTPUT_NODE classes) with their image source."""
|
||
outs = []
|
||
for nid in sorted(prompt, key=_node_sort_key):
|
||
ct = prompt[nid].get("class_type", "")
|
||
info = _info(class_info, ct)
|
||
is_out = bool(info and info.get("output_node")) or ct.startswith(("SaveImage", "PreviewImage", "PreviewAny", "SaveVideo"))
|
||
if not is_out:
|
||
continue
|
||
src = None
|
||
for iname in ("images", "image", "source", "video"):
|
||
v = (prompt[nid].get("inputs") or {}).get(iname)
|
||
if is_link(v):
|
||
src = [str(v[0]), int(v[1])]
|
||
break
|
||
outs.append({"node_id": nid, "class_type": ct, "title": node_title(prompt, nid), "source": src})
|
||
return outs
|
||
|
||
|
||
def default_output(prompt, class_info=None):
|
||
"""The socket that should feed the collector: {node_id, output_index, picked_from?} or None."""
|
||
sources = find_image_sources(prompt, class_info)
|
||
return sources[0] if sources else None
|
||
|
||
|
||
def describe_nodes(prompt, class_info=None):
|
||
"""Node table for the Builder: one entry per node with typed inputs (links flagged)."""
|
||
out = []
|
||
for nid in sorted(prompt, key=_node_sort_key):
|
||
node = prompt[nid]
|
||
ct = node.get("class_type", "")
|
||
info = _info(class_info, ct)
|
||
specs = (info or {}).get("inputs") or {}
|
||
inputs = []
|
||
for iname, value in (node.get("inputs") or {}).items():
|
||
spec = specs.get(iname) or {}
|
||
row = {"name": iname, "value": None if is_link(value) else value, "is_link": is_link(value),
|
||
"link": [str(value[0]), int(value[1])] if is_link(value) else None,
|
||
"type": spec.get("type") or _guess_type(value), "options": spec.get("options")}
|
||
inputs.append(row)
|
||
out.append({"id": nid, "class_type": ct, "title": node_title(prompt, nid),
|
||
"missing_class": bool(info and info.get("missing")), "inputs": inputs})
|
||
return out
|
||
|
||
|
||
def _guess_type(value):
|
||
if isinstance(value, bool):
|
||
return "BOOLEAN"
|
||
if isinstance(value, int):
|
||
return "INT"
|
||
if isinstance(value, float):
|
||
return "FLOAT"
|
||
if isinstance(value, str):
|
||
return "STRING"
|
||
return "UNKNOWN"
|
||
|
||
|
||
# --------------------------------------------------------------------------- suggestions
|
||
|
||
def _follow_text_source(prompt, node_id, input_name, max_hops=6):
|
||
"""CLIPTextEncode.text is often a LINK (MOSS: 30.text <- 34 <- 35 <- 26.value). Find the first
|
||
node upstream that holds a literal string; returns {node_id, input} or None."""
|
||
node = prompt.get(str(node_id))
|
||
if not node:
|
||
return None
|
||
v = (node.get("inputs") or {}).get(input_name)
|
||
if isinstance(v, str):
|
||
return {"node_id": str(node_id), "input": input_name}
|
||
if not is_link(v):
|
||
return None
|
||
seen = {str(node_id)}
|
||
queue = [(str(v[0]), 1)]
|
||
while queue:
|
||
cur, depth = queue.pop(0)
|
||
if cur in seen or cur not in prompt or depth > max_hops:
|
||
continue
|
||
seen.add(cur)
|
||
inputs = prompt[cur].get("inputs") or {}
|
||
for key in _TEXT_INPUTS:
|
||
if isinstance(inputs.get(key), str):
|
||
return {"node_id": cur, "input": key}
|
||
for key in _FOLLOW_INPUTS:
|
||
link = inputs.get(key)
|
||
if is_link(link):
|
||
queue.append((str(link[0]), depth + 1))
|
||
for key, link in inputs.items():
|
||
if key not in _FOLLOW_INPUTS and is_link(link) and key not in ("clip", "model", "vae"):
|
||
queue.append((str(link[0]), depth + 1))
|
||
return None
|
||
|
||
|
||
def find_prompt_targets(prompt, output_node_id):
|
||
"""Locate the positive / negative text inputs feeding the sampler(s) upstream of the output.
|
||
Class-agnostic: any node with link inputs named `positive` and `negative` (KSampler, CFGGuider...).
|
||
Returns {"positive": {node_id,input}|None, "negative": ...}."""
|
||
result = {"positive": None, "negative": None}
|
||
for nid in ancestors_ordered(prompt, output_node_id):
|
||
inputs = prompt[nid].get("inputs") or {}
|
||
for role in ("positive", "negative"):
|
||
if result[role] is not None or not is_link(inputs.get(role)):
|
||
continue
|
||
enc = str(inputs[role][0])
|
||
enc_inputs = (prompt.get(enc) or {}).get("inputs") or {}
|
||
for key in ("text", "prompt", "text_g", "value"):
|
||
if key in enc_inputs:
|
||
target = _follow_text_source(prompt, enc, key)
|
||
if target:
|
||
result[role] = target
|
||
break
|
||
return result
|
||
|
||
|
||
def _lora_slot_nodes(prompt, anc):
|
||
slots = []
|
||
for nid in anc:
|
||
inputs = prompt[nid].get("inputs") or {}
|
||
if "lora_name" in inputs and "strength_model" in inputs:
|
||
slots.append(nid)
|
||
return slots
|
||
|
||
|
||
def suggest_bindings(prompt, output, class_info=None):
|
||
"""Binding suggestions for the Builder, restricted to nodes upstream of the output.
|
||
Returns [{kind, field, targets, reason}] (targets use {node_id, input})."""
|
||
out_id = str(output["node_id"] if isinstance(output, dict) else output[0])
|
||
anc = ancestors_ordered(prompt, out_id)
|
||
suggestions = []
|
||
targets = find_prompt_targets(prompt, out_id)
|
||
for role in ("positive", "negative"):
|
||
if targets[role]:
|
||
suggestions.append({"kind": role, "field": None, "targets": [targets[role]],
|
||
"reason": f"{role} prompt text feeding the sampler"})
|
||
for nid in anc: # nearest ancestor first, so the first candidate of a field is the first-ancestor-wins pick
|
||
for iname, value in (prompt[nid].get("inputs") or {}).items():
|
||
field = ROLE_HINTS.get(iname)
|
||
if not field or field == "text_encoders" or is_link(value) or isinstance(value, (list, dict)):
|
||
continue
|
||
suggestions.append({"kind": "literal", "field": field,
|
||
"targets": [{"node_id": nid, "input": iname}],
|
||
"reason": f"{node_title(prompt, nid)}.{iname}"})
|
||
slots = _lora_slot_nodes(prompt, anc)
|
||
if slots:
|
||
suggestions.append({"kind": "lora", "field": "lora", "slots": [{"node_id": s} for s in slots],
|
||
"targets": [], "reason": "LoRA loader node(s) upstream of the output"})
|
||
for nid in anc:
|
||
inputs = prompt[nid].get("inputs") or {}
|
||
ct = prompt[nid].get("class_type", "")
|
||
if "width" in inputs and "height" in inputs and "Latent" in ct and "Empty" in ct \
|
||
and not is_link(inputs["width"]) and not is_link(inputs["height"]):
|
||
suggestions.append({"kind": "resolution", "field": None, "targets": [],
|
||
"width": {"node_id": nid, "input": "width"},
|
||
"height": {"node_id": nid, "input": "height"},
|
||
"reason": f"{node_title(prompt, nid)} width x height"})
|
||
break
|
||
return suggestions
|
||
|
||
|
||
# --------------------------------------------------------------------------- template -> known fields
|
||
|
||
def _field_sources(prompt, output_id):
|
||
"""Static (field, node_id, input) triples, first ancestor wins, plus encoder/clip_type/lora sources."""
|
||
anc = ancestors_ordered(prompt, output_id)
|
||
sources = {}
|
||
encoders = []
|
||
for nid in anc:
|
||
ct = prompt[nid].get("class_type", "")
|
||
for iname, value in (prompt[nid].get("inputs") or {}).items():
|
||
field = ROLE_HINTS.get(iname)
|
||
if field == "text_encoders":
|
||
encoders.append((nid, iname))
|
||
elif field and field not in sources:
|
||
sources[field] = (nid, iname)
|
||
elif iname == "type" and "CLIPLoader" in ct and "clip_type" not in sources:
|
||
sources["clip_type"] = (nid, iname)
|
||
prompt_targets = find_prompt_targets(prompt, output_id)
|
||
for role in ("positive", "negative"):
|
||
if prompt_targets[role]:
|
||
sources[role] = (prompt_targets[role]["node_id"], prompt_targets[role]["input"])
|
||
return sources, encoders, _lora_slot_nodes(prompt, anc)
|
||
|
||
|
||
def _literal(prompt, nid, iname):
|
||
v = ((prompt.get(nid) or {}).get("inputs") or {}).get(iname)
|
||
return None if is_link(v) or isinstance(v, (dict,)) else v
|
||
|
||
|
||
def known_fields(prompt, sources, encoders, lora_nodes):
|
||
"""Manifest fields read from a (substituted) prompt using precomputed sources."""
|
||
fields = {}
|
||
for field, (nid, iname) in sources.items():
|
||
v = _literal(prompt, nid, iname)
|
||
if v is not None:
|
||
fields[field] = v
|
||
enc = [_literal(prompt, nid, iname) for nid, iname in encoders]
|
||
enc = [e for e in enc if isinstance(e, str) and e]
|
||
if enc:
|
||
fields["text_encoders"] = enc
|
||
parts = []
|
||
for nid in lora_nodes:
|
||
inputs = prompt[nid].get("inputs") or {}
|
||
name = inputs.get("lora_name")
|
||
ms, cs = inputs.get("strength_model"), inputs.get("strength_clip", inputs.get("strength_model"))
|
||
if isinstance(name, str) and name and not is_link(ms) and not is_link(cs):
|
||
try:
|
||
if float(ms) == 0 and float(cs) == 0:
|
||
continue
|
||
parts.append(f"{name}:{float(ms):g}:{float(cs):g}")
|
||
except (TypeError, ValueError):
|
||
parts.append(name)
|
||
fields["lora"] = " + ".join(parts) if parts else "None"
|
||
fields["lora_expanded"] = fields["lora"]
|
||
if fields.get("denoise") == 1.0:
|
||
fields["denoise"] = 1
|
||
return fields
|
||
|
||
|
||
def known_fields_from_template(prompt, output_id):
|
||
sources, encoders, lora_nodes = _field_sources(prompt, str(output_id))
|
||
return known_fields(prompt, sources, encoders, lora_nodes)
|
||
|
||
|
||
# --------------------------------------------------------------------------- coercion, prompts, fingerprints
|
||
|
||
def coerce_like(value, template_value):
|
||
"""Coerce a sweep value to the type of the template literal it replaces
|
||
(int stays int, integral float -> int for INT inputs, "true"/"false" -> bool, combos as str)."""
|
||
if template_value is None or is_link(template_value):
|
||
return value
|
||
try:
|
||
if isinstance(template_value, bool):
|
||
if isinstance(value, str):
|
||
low = value.strip().lower()
|
||
if low in ("true", "1", "yes", "on"):
|
||
return True
|
||
if low in ("false", "0", "no", "off"):
|
||
return False
|
||
return bool(value)
|
||
if isinstance(template_value, int):
|
||
f = float(value)
|
||
return int(f) if f == int(f) else f
|
||
if isinstance(template_value, float):
|
||
return float(value)
|
||
if isinstance(template_value, str):
|
||
return str(value)
|
||
except (TypeError, ValueError):
|
||
return value
|
||
return value
|
||
|
||
|
||
def template_hash(prompt, output, run_other_outputs=False):
|
||
payload = json.dumps({"p": prompt, "o": [str(output[0]), int(output[1])], "r": bool(run_other_outputs)},
|
||
sort_keys=True, separators=(",", ":"), default=str)
|
||
return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:16]
|
||
|
||
|
||
def _strip_meta(prompt):
|
||
return {nid: {k: v for k, v in node.items() if k != "_meta"} for nid, node in prompt.items()}
|
||
|
||
|
||
def fingerprint(substituted_prompt, output):
|
||
"""Resume key of one job: sha256 of the substituted template (no _meta, no collector) + output socket."""
|
||
payload = json.dumps({"p": _strip_meta(substituted_prompt), "o": [str(output[0]), int(output[1])]},
|
||
sort_keys=True, separators=(",", ":"), default=str)
|
||
return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:20]
|
||
|
||
|
||
def apply_subs(template, subs):
|
||
"""Deepcopy the template and set each (node_id, input, value)."""
|
||
prompt = copy.deepcopy(template)
|
||
for nid, iname, value in subs:
|
||
nid = str(nid)
|
||
if nid not in prompt:
|
||
raise ValueError(f"Binding targets node {nid}, which is not in the workflow.")
|
||
prompt[nid].setdefault("inputs", {})[iname] = value
|
||
return prompt
|
||
|
||
|
||
def collector_node_id(prompt):
|
||
"""Numeric string one above the highest numeric id (keeps ComfyUI's history 'load' path happy)."""
|
||
nums = [int(k) for k in prompt if str(k).isdigit()]
|
||
cid = (max(nums) + 1) if nums else 1
|
||
while str(cid) in prompt:
|
||
cid += 1
|
||
return str(cid)
|
||
|
||
|
||
def build_job_prompt(template, subs, output, collector_inputs):
|
||
"""
|
||
Substitute `subs` into a copy of the template and append the collector node.
|
||
Nothing else is touched. Returns (prompt, collector_id).
|
||
"""
|
||
prompt = apply_subs(template, subs)
|
||
cid = collector_node_id(prompt)
|
||
inputs = {"images": [str(output[0]), int(output[1])]}
|
||
inputs.update(collector_inputs or {})
|
||
prompt[cid] = {"class_type": COLLECTOR_CLASS, "inputs": inputs, "_meta": {"title": "USCG Sweep Collect"}}
|
||
return prompt, cid
|
||
|
||
|
||
# --------------------------------------------------------------------------- expansion
|
||
|
||
def _cu():
|
||
try:
|
||
from . import config_utils
|
||
except ImportError:
|
||
import config_utils
|
||
return config_utils
|
||
|
||
|
||
def _parse_prompts(raw, fallback):
|
||
"""Same rules as config_utils.expand_configs per-config prompt handling."""
|
||
cu = _cu()
|
||
if isinstance(raw, (list, tuple)):
|
||
vals = cu.parse_prompt_input_nested(json.dumps(raw))
|
||
elif isinstance(raw, str) and raw.strip():
|
||
vals = cu.parse_prompt_input_nested(raw)
|
||
else:
|
||
vals = list(fallback or [])
|
||
return [v for v in vals if isinstance(v, str)]
|
||
|
||
|
||
def _expand_folder_value(value, iname):
|
||
if not (isinstance(value, str) and value.replace("\\", "/").endswith("/")):
|
||
return [value]
|
||
cu = _cu()
|
||
for folder_type in FOLDER_TYPES.get(iname, []):
|
||
try:
|
||
found = cu.get_files_from_folder(value, folder_type)
|
||
except Exception:
|
||
found = []
|
||
if found:
|
||
return found
|
||
return []
|
||
|
||
|
||
def _dedupe(values):
|
||
seen, out = set(), []
|
||
for v in values:
|
||
key = json.dumps(v, sort_keys=True, default=str)
|
||
if key not in seen:
|
||
seen.add(key)
|
||
out.append(v)
|
||
return out
|
||
|
||
|
||
def _is_loader_class(class_info, class_type):
|
||
info = _info(class_info, class_type)
|
||
if info and info.get("return_types"):
|
||
return any(t in ("MODEL", "CLIP", "VAE") for t in info["return_types"]) and not info.get("inputs", {}).get("model") \
|
||
and not info.get("inputs", {}).get("clip")
|
||
return "Loader" in class_type or "Checkpoint" in class_type
|
||
|
||
|
||
class _Axis:
|
||
__slots__ = ("label", "kind", "field", "options", "nodes", "binding_id", "loader")
|
||
|
||
def __init__(self, label, kind, field, options, nodes, binding_id, loader=False):
|
||
self.label, self.kind, self.field = label, kind, field
|
||
self.options = options # list of {display, subs, fields}
|
||
self.nodes = nodes # node ids the axis writes to
|
||
self.binding_id = binding_id
|
||
self.loader = loader
|
||
|
||
|
||
def _label_for(prompt, nid, iname, title_counts):
|
||
title = node_title(prompt, nid)
|
||
return f"{title}.{iname} #{nid}" if title_counts.get(title, 0) > 1 else f"{title}.{iname}"
|
||
|
||
|
||
def _as_target(t):
|
||
if not isinstance(t, dict) or "node_id" not in t or "input" not in t:
|
||
raise ValueError(f"A binding target needs node_id and input, got {t!r}.")
|
||
return str(t["node_id"]), str(t["input"])
|
||
|
||
|
||
def _build_axes(prompt, bindings, entry, anc_set, extra_seeds, pos_prompts, neg_prompts, class_info, warnings):
|
||
cu = _cu()
|
||
title_counts = {}
|
||
for nid in prompt:
|
||
t = node_title(prompt, nid)
|
||
title_counts[t] = title_counts.get(t, 0) + 1
|
||
axes = []
|
||
written = {} # (node, input) -> binding id, to catch two bindings writing the same input
|
||
|
||
def claim(nid, iname, bid, override_link):
|
||
if nid not in prompt:
|
||
raise ValueError(f"Binding {bid}: node {nid} is not in the workflow.")
|
||
cur = (prompt[nid].get("inputs") or {}).get(iname)
|
||
if is_link(cur) and not override_link:
|
||
raise ValueError(
|
||
f"Binding {bid}: {node_title(prompt, nid)}.{iname} is linked to another node. "
|
||
"Tick 'override link' to replace the link with the swept value.")
|
||
if (nid, iname) in written:
|
||
raise ValueError(f"Bindings {written[(nid, iname)]} and {bid} both write {node_title(prompt, nid)}.{iname}.")
|
||
written[(nid, iname)] = bid
|
||
if nid not in anc_set:
|
||
warnings.append(f"Binding {bid} ({node_title(prompt, nid)}.{iname}) is not upstream of the chosen output; "
|
||
"it has no effect (the image would come from cache).")
|
||
|
||
for bi, b in enumerate(bindings):
|
||
if not isinstance(b, dict):
|
||
continue
|
||
bid = str(b.get("id") or f"b{bi + 1}")
|
||
kind = b.get("kind", "literal")
|
||
override = bool(b.get("override_link"))
|
||
|
||
if kind == "literal":
|
||
targets = [_as_target(t) for t in (b.get("targets") or [])]
|
||
if not targets:
|
||
raise ValueError(f"Binding {bid} has no targets.")
|
||
for nid, iname in targets:
|
||
claim(nid, iname, bid, override)
|
||
nid0, in0 = targets[0]
|
||
tmpl = (prompt[nid0].get("inputs") or {}).get(in0)
|
||
if "field" in b:
|
||
field = b.get("field") or None
|
||
else:
|
||
field = ROLE_HINTS.get(in0)
|
||
field = field if field in AXIS_FIELDS else None
|
||
if field is not None and field not in AXIS_FIELDS:
|
||
raise ValueError(f"Binding {bid}: unknown manifest field '{field}'.")
|
||
raw_values = b.get("values")
|
||
raw_values = raw_values if isinstance(raw_values, list) else ([] if raw_values is None else [raw_values])
|
||
values = []
|
||
for v in raw_values:
|
||
values.extend(_expand_folder_value(v, in0))
|
||
values = [coerce_like(v, tmpl) for v in values]
|
||
if field == "seed":
|
||
values = values + [coerce_like(s, tmpl) for s in (extra_seeds or [])]
|
||
values = _dedupe(values)
|
||
if not values:
|
||
raise ValueError(f"Binding {bid} ({node_title(prompt, nid0)}.{in0}) has no values.")
|
||
options = [{"display": v, "subs": [(n, i, v) for n, i in targets],
|
||
"fields": ({field: v} if field else {})} for v in values]
|
||
loader = _is_loader_class(class_info, prompt[nid0].get("class_type", ""))
|
||
axes.append(_Axis(_label_for(prompt, nid0, in0, title_counts), "literal", field, options,
|
||
[n for n, _ in targets], bid, loader))
|
||
|
||
elif kind in ("positive", "negative"):
|
||
targets = [_as_target(t) for t in (b.get("targets") or [])]
|
||
if not targets:
|
||
raise ValueError(f"Binding {bid} ({kind}) has no targets.")
|
||
for nid, iname in targets:
|
||
claim(nid, iname, bid, override)
|
||
fallback = pos_prompts if kind == "positive" else neg_prompts
|
||
prompts = _dedupe(_parse_prompts(entry.get(kind, ""), fallback))
|
||
if not prompts or all(not p.strip() for p in prompts):
|
||
raise ValueError(
|
||
f"Binding {bid} swaps the {kind} prompt but there are no {kind} prompts. "
|
||
"Add them in this config's Prompts section (or the global prompts).")
|
||
options = [{"display": p, "subs": [(n, i, p) for n, i in targets], "fields": {kind: p}} for p in prompts]
|
||
axes.append(_Axis(_label_for(prompt, targets[0][0], targets[0][1], title_counts), kind, kind, options,
|
||
[n for n, _ in targets], bid))
|
||
|
||
elif kind == "lora":
|
||
slots = [str(s["node_id"] if isinstance(s, dict) else s) for s in (b.get("slots") or [])]
|
||
if not slots:
|
||
raise ValueError(f"Binding {bid} (LoRA) has no slots.")
|
||
for nid in slots:
|
||
if nid not in prompt:
|
||
raise ValueError(f"Binding {bid}: LoRA node {nid} is not in the workflow.")
|
||
inputs = prompt[nid].get("inputs") or {}
|
||
for iname in ("lora_name", "strength_model", "strength_clip"):
|
||
if iname in inputs:
|
||
claim(nid, iname, bid, override)
|
||
raw = b.get("values")
|
||
raw = raw if isinstance(raw, list) else ([] if raw is None else [raw])
|
||
lora_strings = _dedupe(cu.expand_lora_stack([str(v) for v in raw])) if raw else []
|
||
if not lora_strings:
|
||
raise ValueError(f"Binding {bid} (LoRA) has no values.")
|
||
options = []
|
||
for s in lora_strings:
|
||
parts = [] if s == "None" else cu.parse_lora_definition(s)
|
||
if len(parts) > len(slots):
|
||
raise ValueError(
|
||
f"Binding {bid}: '{s}' stacks {len(parts)} LoRAs but only {len(slots)} LoRA slot(s) are bound.")
|
||
subs = []
|
||
for si, nid in enumerate(slots):
|
||
inputs = prompt[nid].get("inputs") or {}
|
||
if si < len(parts):
|
||
name, ms, cs = parts[si]
|
||
subs.append((nid, "lora_name", name))
|
||
subs.append((nid, "strength_model", ms))
|
||
if "strength_clip" in inputs:
|
||
subs.append((nid, "strength_clip", cs))
|
||
else:
|
||
subs.append((nid, "strength_model", 0.0))
|
||
if "strength_clip" in inputs:
|
||
subs.append((nid, "strength_clip", 0.0))
|
||
options.append({"display": s, "subs": subs, "fields": {"lora": s, "lora_expanded": s}})
|
||
axes.append(_Axis(_label_for(prompt, slots[0], "lora_name", title_counts), "lora", "lora", options, slots, bid, True))
|
||
|
||
elif kind == "resolution":
|
||
try:
|
||
wt, ht = _as_target(b.get("width")), _as_target(b.get("height"))
|
||
except (ValueError, TypeError):
|
||
raise ValueError(f"Binding {bid} (resolution) needs width and height targets.")
|
||
claim(wt[0], wt[1], bid, override)
|
||
claim(ht[0], ht[1], bid, override)
|
||
raw = b.get("values")
|
||
raw = raw if isinstance(raw, list) else ([] if raw is None else [raw])
|
||
options = []
|
||
for v in _dedupe([str(x) for x in raw]):
|
||
m = _RES_RE.match(v)
|
||
if not m:
|
||
raise ValueError(f"Binding {bid}: '{v}' is not a WxH resolution like 1024x1024.")
|
||
w, h = int(m.group(1)), int(m.group(2))
|
||
options.append({"display": f"{w}x{h}", "subs": [(wt[0], wt[1], w), (ht[0], ht[1], h)],
|
||
"fields": {"width": w, "height": h}})
|
||
if not options:
|
||
raise ValueError(f"Binding {bid} (resolution) has no values.")
|
||
axes.append(_Axis(_label_for(prompt, wt[0], wt[1], title_counts), "resolution", None, options,
|
||
[wt[0], ht[0]], bid))
|
||
else:
|
||
raise ValueError(f"Binding {bid}: unknown kind '{kind}'.")
|
||
return axes
|
||
|
||
|
||
def _pair_prompt_axes(axes):
|
||
"""Prompt pairing rule of expand_configs: both lists > 1 and equal length -> 1-to-1, else product."""
|
||
pos = [a for a in axes if a.kind == "positive"]
|
||
neg = [a for a in axes if a.kind == "negative"]
|
||
if len(pos) == 1 and len(neg) == 1 and len(pos[0].options) > 1 and len(pos[0].options) == len(neg[0].options):
|
||
p, n = pos[0], neg[0]
|
||
merged = [{"display": a["display"], "subs": a["subs"] + b["subs"], "fields": {**a["fields"], **b["fields"]},
|
||
"extra_params": {n.label: b["display"]}}
|
||
for a, b in zip(p.options, n.options)]
|
||
pair = _Axis(p.label, "prompt_pair", None, merged, p.nodes + n.nodes, p.binding_id, False)
|
||
out = [a for a in axes if a is not p and a is not n]
|
||
out.append(pair)
|
||
return out
|
||
return axes
|
||
|
||
|
||
def order_axes(prompt, axes, anc_set):
|
||
"""Loader axes outermost, prompts next, sampler params innermost: sort by how many nodes of the
|
||
sub-graph depend on the axis' target nodes (descending), loaders first on ties. itertools.product
|
||
then keeps loader values adjacent, which is what ComfyUI's output cache needs."""
|
||
desc = graph_descendants(prompt, anc_set)
|
||
def key(item):
|
||
i, a = item
|
||
d = max([desc.get(n, 0) for n in a.nodes] or [0])
|
||
return (-d, 0 if a.loader else 1, i)
|
||
return [a for _, a in sorted(enumerate(axes), key=key)]
|
||
|
||
|
||
def normalize_output(entry, prompt, class_info=None):
|
||
wo = entry.get("workflow_output")
|
||
if isinstance(wo, dict) and wo.get("node_id") is not None:
|
||
nid, idx = str(wo["node_id"]), int(wo.get("output_index", 0) or 0)
|
||
else:
|
||
d = default_output(prompt, class_info)
|
||
if not d:
|
||
raise ValueError("No image output found. Pick the image socket to collect in the workflow section.")
|
||
nid, idx = d["node_id"], d["output_index"]
|
||
if nid not in prompt:
|
||
raise ValueError(f"The chosen output node {nid} is not in the workflow.")
|
||
return nid, idx
|
||
|
||
|
||
def plan_workflow_entry(entry, extra_seeds=None, pos_prompts=None, neg_prompts=None, class_info=None):
|
||
"""
|
||
Expand one `model_type: "workflow"` config into jobs.
|
||
|
||
Returns {"jobs": [conf, ...], "axes": [{label, count, order}], "warnings": [...],
|
||
"template": {"hash", "prompt", "output", "run_other_outputs"}}.
|
||
Each job conf is flat (4.3 in the plan) and carries wf_subs; the template is NOT copied into confs.
|
||
"""
|
||
prompt = parse_api_workflow(entry.get("workflow"))
|
||
out_id, out_idx = normalize_output(entry, prompt, class_info)
|
||
run_other = bool(entry.get("workflow_run_other_outputs", False))
|
||
warnings = []
|
||
anc = ancestors_ordered(prompt, out_id)
|
||
anc_set = set(anc)
|
||
bindings = entry.get("workflow_bindings") or []
|
||
|
||
axes = _build_axes(prompt, bindings, entry, anc_set, extra_seeds, pos_prompts, neg_prompts, class_info, warnings)
|
||
if extra_seeds and not any(a.field == "seed" for a in axes):
|
||
seed_inputs = [f"{node_title(prompt, n)}.{i}" for n in anc for i, v in (prompt[n].get("inputs") or {}).items()
|
||
if ROLE_HINTS.get(i) == "seed" and not is_link(v)]
|
||
warnings.append("Random extra seeds were ignored: no axis is bound to a seed input. Detected seed inputs: "
|
||
+ (", ".join(seed_inputs) or "none") + ".")
|
||
axes = order_axes(prompt, _pair_prompt_axes(axes), anc_set)
|
||
|
||
total = 1
|
||
for a in axes:
|
||
total *= len(a.options)
|
||
if total > MAX_JOBS:
|
||
raise ValueError(f"This sweep would create more than {MAX_JOBS} jobs; narrow the axes.")
|
||
|
||
sources, encoders, lora_nodes = _field_sources(prompt, out_id)
|
||
thash = template_hash(prompt, (out_id, out_idx), run_other)
|
||
source_name = str(entry.get("workflow_source") or "")
|
||
base = {"model_type": "workflow", "media_type": "image", "wf_template_hash": thash,
|
||
"wf_source": source_name, "wf_output": f"{out_id}:{out_idx}"}
|
||
|
||
jobs = []
|
||
for combo in itertools.product(*[range(len(a.options)) for a in axes]):
|
||
subs, fields, params = [], {}, {}
|
||
for a, oi in zip(axes, combo):
|
||
opt = a.options[oi]
|
||
subs.extend(opt["subs"])
|
||
fields.update(opt["fields"])
|
||
params[a.label] = opt["display"]
|
||
params.update(opt.get("extra_params") or {})
|
||
substituted = apply_subs(prompt, subs)
|
||
conf = dict(base)
|
||
conf.update(known_fields(substituted, sources, encoders, lora_nodes))
|
||
conf.update(fields)
|
||
conf["positive"] = conf.get("positive", "") if isinstance(conf.get("positive", ""), str) else ""
|
||
conf["negative"] = conf.get("negative", "") if isinstance(conf.get("negative", ""), str) else ""
|
||
conf["wf_subs"] = [[n, i, v] for n, i, v in subs]
|
||
conf["wf_params"] = params
|
||
conf["wf_fingerprint"] = fingerprint(substituted, (out_id, out_idx))
|
||
jobs.append(conf)
|
||
|
||
return {
|
||
"jobs": jobs,
|
||
"axes": [{"label": a.label, "count": len(a.options), "order": i} for i, a in enumerate(axes)],
|
||
"warnings": warnings,
|
||
"template": {"hash": thash, "prompt": prompt, "output": [out_id, out_idx], "run_other_outputs": run_other},
|
||
}
|
||
|
||
|
||
def expand_workflow_entry(entry, extra_seeds=None, pos_prompts=None, neg_prompts=None):
|
||
"""`config_utils.expand_configs` hook: the expanded job confs of one workflow entry."""
|
||
return plan_workflow_entry(entry, extra_seeds, pos_prompts, neg_prompts)["jobs"]
|
||
|
||
|
||
def plan_workflow_configs(configs, extra_seeds=None, pos_prompts=None, neg_prompts=None, class_info=None):
|
||
"""Plan every workflow entry of a configs list (Builder order). Raises ValueError on anything that is
|
||
not a workflow entry. Returns {"jobs", "templates": {hash: {...}}, "axes", "warnings"}."""
|
||
jobs, templates, axes, warnings = [], {}, [], []
|
||
for ci, entry in enumerate(configs):
|
||
if not isinstance(entry, dict) or entry.get("model_type") != "workflow":
|
||
raise ValueError("Workflow sweeps cannot be mixed with standard configs in one Generator run. "
|
||
"Use one Builder for the workflow sweep and another for standard configs.")
|
||
plan = plan_workflow_entry(entry, extra_seeds, pos_prompts, neg_prompts, class_info)
|
||
jobs.extend(plan["jobs"])
|
||
templates[plan["template"]["hash"]] = plan["template"]
|
||
axes.append({"config": entry.get("name", f"Config {ci + 1}"), "axes": plan["axes"]})
|
||
warnings.extend(plan["warnings"])
|
||
return {"jobs": jobs, "templates": templates, "axes": axes, "warnings": warnings}
|
||
|
||
|
||
# --------------------------------------------------------------------------- session / manifest helpers
|
||
|
||
def session_paths(session_name):
|
||
"""Same layout as generation_orchestrator.setup_session_directories (replicated: the orchestrator
|
||
is not imported here). The name is sanitised with the same rule as every other route."""
|
||
import folder_paths
|
||
session = _cu().sanitize_session_name(str(session_name or ""))
|
||
base = os.path.join(folder_paths.get_output_directory(), "benchmarks", session)
|
||
return {"session": session, "base": base, "images": os.path.join(base, "images"),
|
||
"manifest": os.path.join(base, "manifest.json"),
|
||
"sweeps": os.path.join(base, "workflow_sweeps")}
|
||
|
||
|
||
def read_manifest_items(manifest_path):
|
||
try:
|
||
with open(manifest_path, "r") as f:
|
||
data = json.load(f)
|
||
return data if isinstance(data, dict) else {"items": []}
|
||
except (OSError, ValueError):
|
||
return {"items": []}
|
||
|
||
|
||
def build_collect_item(conf, width, height, duration, batch_idx, gen_index, create_image_metadata=None):
|
||
"""Manifest item for one collected image (without id/file; the collector adds those).
|
||
Uses image_generation.create_image_metadata so the field set matches standard runs."""
|
||
if create_image_metadata is None:
|
||
try:
|
||
from .image_generation import create_image_metadata
|
||
except ImportError:
|
||
from image_generation import create_image_metadata
|
||
conf = dict(conf or {})
|
||
seed = conf.get("seed", 0)
|
||
seed = 0 if seed is None else seed
|
||
meta = create_image_metadata(conf, int(width), int(height), duration, seed, batch_idx,
|
||
conf.get("positive", ""), conf.get("negative", ""), gen_index=gen_index)
|
||
if meta.get("denoise") == 1.0:
|
||
meta["denoise"] = 1
|
||
meta.setdefault("model_type", "workflow")
|
||
meta.update({"favorited": False, "rejected": False, "note": ""})
|
||
return meta
|
||
|
||
|
||
def new_item_id():
|
||
return int(time.time() * 100000) + random.randint(0, 1000)
|
||
|
||
|
||
def parse_for_ui(api_json, class_info=comfy_class_info):
|
||
"""Response body of POST /config_tester/sweep/parse. Raises ValueError for unusable input."""
|
||
prompt = parse_api_workflow(api_json)
|
||
sources = find_image_sources(prompt, class_info)
|
||
default = sources[0] if sources else None
|
||
suggestions = suggest_bindings(prompt, default, class_info) if default else []
|
||
warnings = []
|
||
nodes_ = describe_nodes(prompt, class_info)
|
||
missing = sorted({n["class_type"] for n in nodes_ if n["missing_class"]})
|
||
if missing:
|
||
warnings.append("These node types are not installed in this ComfyUI: " + ", ".join(missing))
|
||
if not sources:
|
||
warnings.append("No IMAGE output was found; pick the image socket to collect manually.")
|
||
return {"ok": True, "nodes": nodes_, "outputs": find_output_nodes(prompt, class_info),
|
||
"image_sources": sources, "suggestions": suggestions, "warnings": warnings, "errors": []}
|