Shaman
105 lines
3.8 KiB
Python
105 lines
3.8 KiB
Python
"""Build a concrete workflow graph from a template + request values.
|
|
|
|
Values are injected through placeholder tokens written directly inside the
|
|
workflow graph (e.g. "{prompt}", "{seed}"), so the graph itself declares which
|
|
node input receives which request field. The set of recognised tokens is fixed
|
|
and owned by the adapter, not by the model.
|
|
"""
|
|
import copy
|
|
import re
|
|
import secrets
|
|
from typing import Any
|
|
|
|
from .errors import InvalidRequest
|
|
|
|
SEED_MAX = 2**63 - 1
|
|
|
|
# Fixed token vocabulary owned by the adapter. A model's workflow may only use
|
|
# these tokens; anything else is rejected. Types drive typed vs string injection.
|
|
TOKEN_TYPES: dict[str, str] = {
|
|
"prompt": "str",
|
|
"negative": "str",
|
|
"seed": "int",
|
|
"width": "int",
|
|
"height": "int",
|
|
"image": "str",
|
|
}
|
|
KNOWN_TOKENS = frozenset(TOKEN_TYPES)
|
|
|
|
_TOKEN = re.compile(r"\{(\w+)\}")
|
|
_WHOLE = re.compile(r"^\{(\w+)\}$")
|
|
|
|
|
|
def parse_size(size: str) -> tuple[int, int]:
|
|
"""Parse OpenAI `WxH` size strings (e.g. '1024x1024')."""
|
|
try:
|
|
w_str, h_str = size.lower().split("x", 1)
|
|
width, height = int(w_str), int(h_str)
|
|
except ValueError as e:
|
|
raise InvalidRequest(f"invalid size '{size}', expected format 'WIDTHxHEIGHT'") from e
|
|
if width <= 0 or height <= 0:
|
|
raise InvalidRequest(f"invalid size '{size}', dimensions must be positive")
|
|
return width, height
|
|
|
|
|
|
def new_seed(seed: int | None) -> int:
|
|
if seed is None:
|
|
return secrets.randbelow(SEED_MAX) + 1
|
|
if not 0 <= seed <= SEED_MAX:
|
|
raise InvalidRequest(f"seed must be between 0 and {SEED_MAX}")
|
|
return seed
|
|
|
|
|
|
def extract_tokens(graph: dict[str, Any]) -> tuple[set[str], list[str]]:
|
|
"""Token names referenced by the graph.
|
|
|
|
Returns (known_tokens, unknown_token_errors). A token is any `{name}` inside
|
|
a string node input. Unknown names are reported so startup can reject the
|
|
workflow before it ever reaches a request.
|
|
"""
|
|
declared: set[str] = set()
|
|
unknown: set[str] = set()
|
|
for node_id, node in graph.items():
|
|
for field, value in node.get("inputs", {}).items():
|
|
if isinstance(value, str):
|
|
for name in _TOKEN.findall(value):
|
|
if name in KNOWN_TOKENS:
|
|
declared.add(name)
|
|
else:
|
|
unknown.add(f"node {node_id}: input '{field}' uses unknown token '{{{name}}}'")
|
|
return declared, sorted(unknown)
|
|
|
|
|
|
def build_workflow(template: dict[str, Any], values: dict[str, Any]) -> dict[str, Any]:
|
|
"""Deep-copy the template and substitute placeholder tokens with values.
|
|
|
|
Injection is type-preserving: a value that is exactly one token (e.g.
|
|
"{seed}") is replaced by the typed object (int for seed/width/height, str
|
|
otherwise); a token embedded in surrounding text is string-interpolated.
|
|
A referenced token that is unrecognized or missing from ``values`` rejects
|
|
the build, so a malformed graph never reaches ComfyUI.
|
|
"""
|
|
graph = copy.deepcopy(template)
|
|
for node in graph.values():
|
|
inputs = node.get("inputs", {})
|
|
for field, value in inputs.items():
|
|
if isinstance(value, str):
|
|
inputs[field] = _substitute(value, values)
|
|
return graph
|
|
|
|
|
|
def _substitute(value: str, values: dict[str, Any]) -> Any:
|
|
for name in _TOKEN.findall(value):
|
|
if name not in KNOWN_TOKENS:
|
|
raise InvalidRequest(f"workflow uses unrecognized token '{{{name}}}'")
|
|
if name not in values:
|
|
raise InvalidRequest(f"no value supplied for token '{{{name}}}'")
|
|
whole = _WHOLE.fullmatch(value)
|
|
if whole:
|
|
return _coerce(whole.group(1), values[whole.group(1)])
|
|
return _TOKEN.sub(lambda m: str(values[m.group(1)]), value)
|
|
|
|
|
|
def _coerce(name: str, value: Any) -> Any:
|
|
return int(value) if TOKEN_TYPES[name] == "int" else str(value)
|