Shaman
68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
"""Build a concrete workflow graph from a template + request parameters."""
|
|
import copy
|
|
import secrets
|
|
from typing import Any
|
|
|
|
from .errors import InvalidRequest
|
|
|
|
SEED_MAX = 2**63 - 1
|
|
|
|
|
|
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 build_workflow(
|
|
template: dict[str, Any],
|
|
mapping: dict[str, dict[str, str]],
|
|
*,
|
|
prompt: str,
|
|
seed: int,
|
|
size: str | None = None,
|
|
image_name: str | None = None,
|
|
negative: str = "",
|
|
) -> dict[str, Any]:
|
|
"""Deep-copy the template and inject request values per the node mapping."""
|
|
graph = copy.deepcopy(template)
|
|
|
|
def put(key: str, value: Any) -> None:
|
|
spec = mapping.get(key)
|
|
if spec is None:
|
|
return
|
|
node = graph.get(spec["node"])
|
|
if node is None:
|
|
raise InvalidRequest(f"mapping key '{key}' points to missing node '{spec['node']}'")
|
|
node["inputs"][spec["field"]] = value
|
|
|
|
put("prompt", prompt)
|
|
put("negative", negative)
|
|
put("seed", seed)
|
|
if image_name is not None:
|
|
put("image", image_name)
|
|
if size is not None:
|
|
spec = mapping.get("size")
|
|
if spec is not None:
|
|
width, height = parse_size(size)
|
|
node = graph.get(spec["node"])
|
|
if node is None:
|
|
raise InvalidRequest(f"mapping key 'size' points to missing node '{spec['node']}'")
|
|
node["inputs"][spec["width_field"]] = width
|
|
node["inputs"][spec["height_field"]] = height
|
|
return graph
|