DumiFlex
3a7671b467
lookup_var reads $name.bool from the toggles' on/off table; example 2's Var → Bool reads $quick_preview.bool since the Quick toggle's Value (12) now fills $quick_preview. Tests, docs note and release-notes line.
312 lines
13 KiB
Python
312 lines
13 KiB
Python
"""The workflows in `example_workflows/` run on an EMPTY library.
|
|
|
|
ComfyUI lists them under Browse Templates. Every list they use is embedded on
|
|
their WP Context nodes, so this test runs each one's WP nodes through the
|
|
engine with a fresh database and fails on a payload a handler rejects, a
|
|
warning (an unknown `@{}` reference, a missing variable…) or an empty prompt.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from engine.modules import dispatcher
|
|
from engine.modules.snapshot import payload_hash
|
|
from wp_nodes.assembler_node import WPPromptAssembler
|
|
from wp_nodes.context_loop import WPContextLoop
|
|
from wp_nodes.context_node import WPContext
|
|
from wp_nodes.injector_node import WPContextInjector
|
|
from wp_nodes.prompt_cleaner import WPPromptCleaner
|
|
from wp_nodes.toggles_node import WPToggles
|
|
from wp_nodes.var_to_bool import WPVarToBool
|
|
from wp_nodes.var_to_float import WPVarToFloat
|
|
from wp_nodes.var_to_int import WPVarToInt
|
|
|
|
EXAMPLES = sorted((Path(__file__).parent.parent / "example_workflows").glob("*.json"))
|
|
BYPASS, MUTE = 4, 2
|
|
VAR_NODES = {"WP_VarToInt": WPVarToInt, "WP_VarToFloat": WPVarToFloat, "WP_VarToBool": WPVarToBool}
|
|
# Every WP node, frontend-only ones included, shows up in at least one example.
|
|
ALL_WP_NODES = {
|
|
"WP_Context", "WP_ContextLoop", "WP_ContextInjector", "WP_Toggles", "WP_ModelInfo",
|
|
"WP_PromptAssembler", "WP_PromptCleaner", "WP_Debug", "WP_SeedList", "WP_ImageFilter",
|
|
"WP_LabelGrid", "WP_VarToInt", "WP_VarToFloat", "WP_VarToBool", "WP_VarToChoice",
|
|
"WP_Note", "WP_Bookmark", "WP_NodeGroup", "WP_ModeRelay",
|
|
}
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _empty_library(tmp_path, monkeypatch):
|
|
monkeypatch.setenv("WP_DB_PATH", os.fspath(tmp_path / "wp.db"))
|
|
|
|
|
|
def _json_value(values: list[Any], key: str) -> str:
|
|
"""The widget value that is a JSON object holding `key`."""
|
|
for v in values:
|
|
if isinstance(v, str) and v.startswith("{"):
|
|
try:
|
|
if key in json.loads(v):
|
|
return v
|
|
except ValueError:
|
|
continue
|
|
raise AssertionError(f"no JSON widget value with {key!r}")
|
|
|
|
|
|
class Graph:
|
|
"""Runs the WP nodes of a saved (UI-format) workflow, frame by frame."""
|
|
|
|
def __init__(self, wf: dict[str, Any], toggles_on: bool = False, unmute: bool = False):
|
|
self.toggles_on = toggles_on
|
|
# Muted groups are the modes a Pick one toggle switched off; with
|
|
# `unmute` they run too, so every mode gets checked.
|
|
self.unmute = unmute
|
|
self.nodes = {n["id"]: n for n in wf["nodes"]}
|
|
self.links = {link[0]: link for link in wf["links"]}
|
|
self.frame: Any = None
|
|
self.cache: dict[tuple[int, int], Any] = {}
|
|
self.warnings: list[dict[str, Any]] = []
|
|
|
|
def source(self, node: dict[str, Any], name: str) -> tuple[dict[str, Any], int] | None:
|
|
for inp in node.get("inputs") or []:
|
|
if inp["name"] == name and inp.get("link") is not None:
|
|
link = self.links[inp["link"]]
|
|
return self.nodes[link[1]], link[2]
|
|
return None
|
|
|
|
def input(self, node: dict[str, Any], name: str) -> Any:
|
|
src = self.source(node, name)
|
|
return None if src is None else self.output(*src)
|
|
|
|
def output(self, node: dict[str, Any], slot: int) -> Any:
|
|
key = (node["id"], slot)
|
|
if key not in self.cache:
|
|
self.cache[key] = self._run(node, slot)
|
|
return self.cache[key]
|
|
|
|
def _run(self, node: dict[str, Any], slot: int) -> Any:
|
|
t, wv, mode = node["type"], node.get("widgets_values") or [], node.get("mode", 0)
|
|
if mode == MUTE and self.unmute:
|
|
mode = 0
|
|
if mode == MUTE:
|
|
return None
|
|
if mode == BYPASS:
|
|
# A bypassed context node hands its upstream on; anything else
|
|
# gives nothing and the receiving widget keeps its own value.
|
|
return self.input(node, "upstream") if t.startswith("WP_Context") else None
|
|
if t == "WP_ContextLoop":
|
|
return self.frame if slot == 0 else None
|
|
if t == "WP_ModelInfo":
|
|
# No model here: it hands its upstream on, as with no model wired.
|
|
return self.input(node, "upstream")
|
|
if t == "WP_Context":
|
|
out = WPContext.execute(
|
|
wv[0], _json_value(wv, "modules"), upstream=self.input(node, "upstream"),
|
|
)
|
|
ctx = out.values[0]
|
|
self.warnings += ctx.debug.get("__wp_warnings__", [])
|
|
return ctx
|
|
if t == "WP_ContextInjector":
|
|
slots = {
|
|
i["name"]: self.input(node, i["name"])
|
|
for i in node.get("inputs") or [] if i["name"].startswith("input_")
|
|
}
|
|
return WPContextInjector.execute(
|
|
_json_value(wv, "rows"), upstream=self.input(node, "upstream"), **slots,
|
|
).values[0]
|
|
if t == "WP_Toggles":
|
|
cfg = json.loads(_json_value(wv, "rows"))
|
|
if self.toggles_on:
|
|
for row in cfg["rows"]:
|
|
row["on"] = True
|
|
return WPToggles.execute(
|
|
upstream=self.input(node, "upstream"), wp_toggles=json.dumps(cfg),
|
|
).values[slot]
|
|
if t == "WP_PromptAssembler":
|
|
ctx = self.input(node, "context")
|
|
return WPPromptAssembler.execute(ctx, wv[0], wv[1]).values[slot]
|
|
if t == "WP_PromptCleaner":
|
|
return WPPromptCleaner.execute(
|
|
self.input(node, "prompt"), _json_value(wv, "mode"),
|
|
negative=self.input(node, "negative"),
|
|
).values[slot]
|
|
if t in VAR_NODES:
|
|
return VAR_NODES[t].execute(self.input(node, "context"), wv[0], wv[1], wv[2]).values[0]
|
|
if t.startswith("Primitive"):
|
|
return wv[0]
|
|
return None # Model loaders, samplers, images: not ours to run here.
|
|
|
|
def loops(self) -> list[dict[str, Any]]:
|
|
return [n for n in self.nodes.values() if n["type"] == "WP_ContextLoop"]
|
|
|
|
def frames(self, loop: dict[str, Any] | None = None) -> list[Any]:
|
|
loops = self.loops()
|
|
if not loops:
|
|
return [None]
|
|
wv = (loop or loops[0])["widgets_values"]
|
|
return WPContextLoop.execute([wv[0]], [wv[2]], [_json_value(wv, "strategy")]).values[0]
|
|
|
|
def downstream(self, start: dict[str, Any]) -> set[int]:
|
|
"""Ids of every node fed, directly or not, by `start`."""
|
|
seen: set[int] = set()
|
|
todo = [start["id"]]
|
|
while todo:
|
|
nid = todo.pop()
|
|
for link in self.links.values():
|
|
if link[1] == nid and link[3] not in seen:
|
|
seen.add(link[3])
|
|
todo.append(link[3])
|
|
return seen
|
|
|
|
def per_loop(self) -> list[tuple[dict[str, Any] | None, list[dict[str, Any]]]]:
|
|
"""Each loop (or None) with the active Prompt Assemblers it feeds."""
|
|
asms = [
|
|
n for n in self.nodes.values() if n["type"] == "WP_PromptAssembler"
|
|
and (n.get("mode", 0) == 0 or (self.unmute and n.get("mode") == MUTE))
|
|
]
|
|
loops = self.loops()
|
|
if len(loops) <= 1:
|
|
return [(loops[0] if loops else None, asms)]
|
|
return [(lp, [a for a in asms if a["id"] in self.downstream(lp)]) for lp in loops]
|
|
|
|
def run_frames(
|
|
self, loop: dict[str, Any] | None, node: dict[str, Any], slot: int = 0,
|
|
) -> list[Any]:
|
|
out = []
|
|
for frame in self.frames(loop) if loop else [None]:
|
|
self.frame, self.cache = frame, {}
|
|
out.append(self.output(node, slot))
|
|
return out
|
|
|
|
|
|
def _load(path: Path) -> dict[str, Any]:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def test_examples_exist():
|
|
assert len(EXAMPLES) >= 3
|
|
|
|
|
|
@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem)
|
|
def test_embedded_rows_are_engine_valid(path: Path):
|
|
for node in _load(path)["nodes"]:
|
|
if node["type"] != "WP_Context":
|
|
continue
|
|
value = json.loads(_json_value(node["widgets_values"], "modules"))
|
|
uids = {b["_uid"] for b in value.get("bundles", [])}
|
|
for b in value.get("bundles", []):
|
|
assert b.get("example") is True, f"bundle {b['name']} is not flagged example"
|
|
assert b.get("parent_uid") in (None, *uids)
|
|
for row in value["modules"]:
|
|
handler = dispatcher.get_handler(row["type"])
|
|
assert handler is not None, row["type"]
|
|
handler.validate_payload(row["payload"])
|
|
assert row["payload_hash"] == payload_hash(row["payload"]), row["meta"]["name"]
|
|
assert row.get("example") is True, f"row {row['meta']['name']} is not flagged example"
|
|
|
|
|
|
@pytest.mark.parametrize("path", EXAMPLES, ids=lambda p: p.stem)
|
|
def test_runs_on_an_empty_library_without_warnings(path: Path):
|
|
wf = _load(path)
|
|
graph = Graph(wf, unmute=True)
|
|
groups = graph.per_loop()
|
|
assert any(asms for _, asms in groups)
|
|
for loop, assemblers in groups:
|
|
assert assemblers, f"loop {loop and loop['id']} feeds no Prompt Assembler"
|
|
for frame in graph.frames(loop) if loop else [None]:
|
|
graph.frame, graph.cache = frame, {}
|
|
for asm in assemblers:
|
|
assert graph.output(asm, 0).strip(), f"{asm.get('title')} printed nothing"
|
|
for node in wf["nodes"]:
|
|
if node["type"] == "WP_PromptCleaner" or node["type"] in VAR_NODES:
|
|
graph.output(node, 0)
|
|
bad = [w for w in graph.warnings if w.get("severity") in ("warn", "warning", "error")]
|
|
assert not bad, bad
|
|
|
|
|
|
def test_lora_toggles_swap_in_the_trigger_words():
|
|
path = next(p for p in EXAMPLES if p.stem.startswith("1 "))
|
|
wf = _load(path)
|
|
asm = next(n for n in wf["nodes"] if n["type"] == "WP_PromptAssembler")
|
|
off = Graph(wf).output(asm, 0)
|
|
on = Graph(wf, toggles_on=True).output(asm, 0)
|
|
words = [n["widgets_values"][0] for n in wf["nodes"] if n["type"].startswith("PrimitiveString")]
|
|
assert len(words) == 3
|
|
assert not any(w in off for w in words)
|
|
assert all(w in on for w in words), on
|
|
|
|
|
|
def test_every_wp_node_is_in_an_example():
|
|
seen = {n["type"] for p in EXAMPLES for n in _load(p)["nodes"]}
|
|
assert ALL_WP_NODES <= seen, sorted(ALL_WP_NODES - seen)
|
|
|
|
|
|
def _sets_and_stories() -> tuple[Graph, dict[str, tuple[dict[str, Any], dict[str, Any]]]]:
|
|
"""The Sets and stories workflow: each group's loop, prompt and labels."""
|
|
path = next(p for p in EXAMPLES if p.stem.startswith("4 "))
|
|
wf = _load(path)
|
|
graph = Graph(wf, unmute=True)
|
|
titled = {n["title"]: n for n in wf["nodes"] if n.get("title")}
|
|
out = {}
|
|
toggles = titled["WP Toggles · pick a mode"]
|
|
rows = json.loads(toggles["widgets_values"][0])["rows"]
|
|
modes = {g for r in rows for g in r.get("groups", [])}
|
|
for group in wf["groups"]:
|
|
if group["title"] not in modes:
|
|
continue # the parts inside a mode row (Lists, Prompt, Image…)
|
|
x, y, w, h = group["bounding"]
|
|
inside = [
|
|
n for n in wf["nodes"] if x <= n["pos"][0] <= x + w and y <= n["pos"][1] <= y + h
|
|
]
|
|
loop = next(n for n in inside if n["type"] == "WP_ContextLoop")
|
|
asm = next(n for n in inside if n["type"] == "WP_PromptAssembler" and not n.get("title"))
|
|
out[group["title"]] = (loop, asm)
|
|
assert titled["WP Toggles · pick a mode"]
|
|
return graph, out
|
|
|
|
|
|
def _subject(prompt: str) -> str:
|
|
found = re.search(r"\ban? \d+-year-old [^,]+", prompt)
|
|
assert found, prompt
|
|
return found.group(0)
|
|
|
|
|
|
def test_sets_and_stories_modes_hold_what_they_say():
|
|
graph, groups = _sets_and_stories()
|
|
loop, asm = groups["Every frame different"]
|
|
random_people = {_subject(p) for p in graph.run_frames(loop, asm)}
|
|
assert len(random_people) > 3, random_people
|
|
|
|
loop, asm = groups["Photo shoot"]
|
|
shots = graph.run_frames(loop, asm)
|
|
assert len({_subject(p) for p in shots}) == 1, shots
|
|
order = ["wide shot", "full body", "cowboy shot", "upper body",
|
|
"portrait", "close-up"]
|
|
assert [next(s for s in order if f", {s}, " in p) for p in shots] == order
|
|
|
|
loop, asm = groups["Storyboard"]
|
|
beats = graph.run_frames(loop, asm)
|
|
assert len({_subject(p) for p in beats}) == 1, beats
|
|
assert len(set(beats)) == len(beats)
|
|
|
|
|
|
@pytest.mark.parametrize("pick,quick", [("Quick", True), ("Full", False)])
|
|
def test_quick_preview_switch_follows_the_pick(pick: str, quick: bool):
|
|
"""Example 2: the Quick toggle has a Value (its steps), so Var → Bool reads its switch
|
|
through `$quick_preview.bool`, not the value."""
|
|
path = next(p for p in EXAMPLES if p.stem.startswith("2 "))
|
|
wf = _load(path)
|
|
tog = next(n for n in wf["nodes"] if n["type"] == "WP_Toggles")
|
|
wv = tog["widgets_values"]
|
|
i = wv.index(_json_value(wv, "rows"))
|
|
cfg = json.loads(wv[i])
|
|
for row in cfg["rows"]:
|
|
row["on"] = row.get("name") == pick
|
|
wv[i] = json.dumps(cfg)
|
|
node = next(n for n in wf["nodes"] if n["type"] == "WP_VarToBool")
|
|
assert Graph(wf).output(node, 0) is quick
|