Dumitru Picalau
26df88216f
Chat round (explain on selection, page starters, reactions, streamed notes, thinking), negatives on node fixed values and combines, Pack builder quality fixes from 31 real-model eval rounds, and small-model support (scoped tools, checked tool arguments on local servers, two-phase steps, fixed prefix shown in Test connection).
366 lines
16 KiB
Python
366 lines
16 KiB
Python
"""WP_Context — runs a module list and emits a PipelineContext payload."""
|
|
|
|
import contextlib
|
|
|
|
from comfy_api.latest import io # pyright: ignore[reportMissingImports]
|
|
|
|
from engine.db.connection import get_connection, resolve_db_path
|
|
from engine.db.migrations import migrate
|
|
from engine.db.repositories import ModuleNotFound, ModuleRepository
|
|
from engine.modules.snapshot import ref_seed_uuids, walk_transitive_refs
|
|
from engine.pipeline import FRAME0_CTX_KEY, PipelineEngine
|
|
from engine.seed_derive import effective_chain_seed
|
|
from wp_nodes.types import (
|
|
FRAME0_KEY,
|
|
ContextModulesInput,
|
|
PipelineContext,
|
|
build_payload,
|
|
deserialize_node_input,
|
|
frame0_payload,
|
|
)
|
|
|
|
#: DB files this process already migrated from the node path. The API routes
|
|
#: migrate at startup; this only covers a graph run that reaches the library
|
|
#: before them (or without them), and keeps the sha256 pass over every
|
|
#: migration file off each run.
|
|
_MIGRATED: set[str] = set()
|
|
|
|
|
|
def _migrate_once(conn) -> None:
|
|
key = str(resolve_db_path())
|
|
if key in _MIGRATED:
|
|
return
|
|
migrate(conn)
|
|
_MIGRATED.add(key)
|
|
|
|
|
|
def _input_ctx(upstream_ctx: dict, upstream_internals: dict, catalog: dict) -> dict:
|
|
"""The ctx a run starts from: the upstream variables, the catalog, and
|
|
the upstream cross-node internals."""
|
|
ctx: dict = dict(upstream_ctx)
|
|
# Inject catalog ONCE at the top of execute, before pipeline
|
|
# run. Spec §2.6 — never re-inject mid-run.
|
|
ctx["__wp_catalog__"] = catalog
|
|
# Merge cross-node internals from the upstream payload —
|
|
# `__wp_picks__` so a downstream constraint can read source
|
|
# picks recorded in an earlier Context, `__wp_constraints__`
|
|
# so a constraint registered upstream still applies to a
|
|
# downstream target. Copy in defensively so this node's
|
|
# mutations don't corrupt the upstream payload (frozen
|
|
# dataclass implies the dict references are still shared).
|
|
for key, value in upstream_internals.items():
|
|
if isinstance(value, list):
|
|
ctx[key] = list(value)
|
|
elif isinstance(value, dict):
|
|
ctx[key] = dict(value)
|
|
elif isinstance(value, set):
|
|
# Defensive copy of any set-valued cross-node internal so
|
|
# upstream payload's set isn't mutated by this node's run.
|
|
# (SP3 swapped the constraint bookkeeping from a set to
|
|
# the `__wp_constraint_hits__` dict — handled by the dict
|
|
# branch above — but keep this for any future set keys.)
|
|
ctx[key] = set(value)
|
|
else:
|
|
ctx[key] = value
|
|
return ctx
|
|
|
|
|
|
def _expand_catalog_via_live_db(catalog: dict, modules: list | None = None) -> dict:
|
|
"""Per-issue-#2: the embed-bundle endpoint no longer walks
|
|
transitive ``@{}`` refs — workflow JSON only carries what the
|
|
user explicitly picked. At graph-run time we fill in any nested
|
|
wildcards by querying the live library on the executing machine.
|
|
|
|
Strategy: feed the picked wildcards' uuids — PLUS every uuid referenced by
|
|
any module on any surface (see ``_ref_seed_uuids``) — into
|
|
``walk_transitive_refs`` with a fetch callback that returns the
|
|
embedded snapshot when available (so picked entries keep their
|
|
saved payload — drift detection still works) and falls back to
|
|
the live DB when a referenced uuid is NOT embedded.
|
|
|
|
Failure modes:
|
|
- DB not reachable / migrations not yet run → return embedded
|
|
catalog as-is. Resolver will emit "Unknown wildcard ref"
|
|
warnings for unresolved nested refs (existing lenient behavior).
|
|
- Referenced uuid not in DB → walker records ``missing_target``;
|
|
resolver still emits the warning at run time.
|
|
"""
|
|
seeds = ref_seed_uuids(modules)
|
|
if not catalog and not seeds:
|
|
return catalog
|
|
|
|
try:
|
|
conn = get_connection()
|
|
except Exception:
|
|
# No DB access from this graph run (e.g. embedded ComfyUI
|
|
# without the SPA stack) — return what we have.
|
|
return catalog
|
|
# One connection per run, closed when the walk is done (it used to stay
|
|
# open: one leaked connection per WP_Context execution).
|
|
with contextlib.closing(conn):
|
|
try:
|
|
_migrate_once(conn)
|
|
except Exception:
|
|
return catalog
|
|
return _walk_live(catalog, seeds, ModuleRepository(conn))
|
|
|
|
|
|
def _walk_live(catalog: dict, seeds, repo: ModuleRepository) -> dict:
|
|
"""The walk itself, on the caller's open repository."""
|
|
|
|
def _fetch(uuid: str) -> dict | None:
|
|
embedded = catalog.get(uuid)
|
|
if embedded is not None:
|
|
# Re-shape embedded SnapshotEntry → repo row shape so
|
|
# walk_transitive_refs's downstream logic sees a
|
|
# consistent dict.
|
|
return {
|
|
"id": uuid,
|
|
"type": embedded["type"],
|
|
"name": embedded.get("name", ""),
|
|
"payload": embedded["payload"],
|
|
"payload_hash": embedded.get("payload_hash", ""),
|
|
}
|
|
try:
|
|
return repo.get(uuid)
|
|
except ModuleNotFound:
|
|
return None
|
|
except Exception:
|
|
return None
|
|
|
|
# Picked wildcards first so they keep `source: user`; ref-only targets
|
|
# follow and land as deps. `walk_transitive_refs` skips a uuid it has
|
|
# already snapshotted, so the overlap between the two is harmless.
|
|
roots = list(catalog.keys()) + [u for u in sorted(seeds) if u not in catalog]
|
|
walk = walk_transitive_refs(roots, fetch_module=_fetch)
|
|
_stamp_pool_provenance(walk.snapshots, catalog, repo)
|
|
return walk.snapshots
|
|
|
|
|
|
def _stamp_pool_provenance(
|
|
snapshots: dict, embedded: dict, repo: ModuleRepository,
|
|
) -> None:
|
|
"""Record, per catalog entry, WHERE its option pool came from.
|
|
|
|
Runtime-only bookkeeping consumed by ``ref_option_pool`` when a filtered
|
|
nested ref matches nothing. Without it that warning can only say "matched
|
|
no options", which is true and useless: the user cannot tell whether the
|
|
filter is wrong, or whether it is right but this node is holding a stale
|
|
copy of a pool the library has since changed. That ambiguity is what made
|
|
the nested-ref bug so hard to read from the Debug panel.
|
|
|
|
Two fields:
|
|
* ``pool_origin`` — ``"node"`` when the entry came from THIS node's
|
|
embedded snapshot (which shadows the library), ``"library"`` when it
|
|
was fetched because no module here holds that uuid.
|
|
* ``library_option_count`` — only for ``node`` entries whose live
|
|
library row now has a DIFFERENT number of options. Absent when they
|
|
agree, so its presence alone means "this node's copy has drifted".
|
|
|
|
Mutates in place. Never raises: this is diagnostic garnish, and a DB
|
|
hiccup here must not take a graph run down with it. The extra reads are
|
|
one per embedded wildcard, at graph-run time, on a connection already open.
|
|
|
|
The keys ride on the runtime catalog only — `__wp_catalog__` is stripped at
|
|
the socket boundary and the workflow-embedded snapshot is built elsewhere,
|
|
so neither reaches saved JSON.
|
|
"""
|
|
for uuid, entry in snapshots.items():
|
|
if not isinstance(entry, dict):
|
|
continue
|
|
from_node = uuid in embedded
|
|
entry["pool_origin"] = "node" if from_node else "library"
|
|
if not from_node:
|
|
continue
|
|
try:
|
|
live = repo.get(uuid)
|
|
except Exception:
|
|
continue
|
|
live_options = (live.get("payload") or {}).get("options")
|
|
own_options = (entry.get("payload") or {}).get("options")
|
|
if not isinstance(live_options, list) or not isinstance(own_options, list):
|
|
continue
|
|
if len(live_options) != len(own_options):
|
|
entry["library_option_count"] = len(live_options)
|
|
|
|
|
|
def _node_id(cls: type) -> str:
|
|
"""This node's graph id when ComfyUI supplied it (tests call `execute`
|
|
directly and get "")."""
|
|
hidden = getattr(cls, "hidden", None)
|
|
uid = getattr(hidden, "unique_id", None) if hidden is not None else None
|
|
return str(uid) if uid is not None else ""
|
|
|
|
|
|
class WPContext(io.ComfyNode):
|
|
"""Context node: runs an ordered list of modules, emits PipelineContext."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="WP_Context",
|
|
display_name="WP Context",
|
|
category="wildcard-pipeline",
|
|
inputs=[
|
|
PipelineContext.Input(
|
|
"upstream",
|
|
optional=True,
|
|
tooltip=(
|
|
"Optional upstream Context to extend. Variables pass "
|
|
"through; this node's modules add or override bindings "
|
|
"on top. Leave unconnected to start fresh."
|
|
),
|
|
),
|
|
io.Int.Input(
|
|
"seed",
|
|
default=0,
|
|
min=0,
|
|
max=0xFFFFFFFFFFFFFFFF,
|
|
control_after_generate=True,
|
|
tooltip=(
|
|
"Starting point for randomisation. Change for a "
|
|
"different roll; keep the same to reproduce. "
|
|
"Right-click → Convert widget to input to drive from "
|
|
"another seed node."
|
|
),
|
|
),
|
|
ContextModulesInput.Input("wp_modules", socketless=True),
|
|
],
|
|
outputs=[PipelineContext.Output("context")],
|
|
# The node's graph id, so WP Debug can say which Context ran what.
|
|
hidden=[io.Hidden.unique_id],
|
|
not_idempotent=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, seed, wp_modules, upstream=None):
|
|
upstream_ctx: dict = upstream.context if upstream is not None else {}
|
|
upstream_debug: dict = upstream.debug if upstream is not None else {}
|
|
upstream_internals: dict = (
|
|
upstream.internals if upstream is not None else {}
|
|
)
|
|
|
|
# Unified-list model: one `modules` array holds every kind.
|
|
# The catalog for `@{}` ref resolution is the wildcard subset,
|
|
# synthesised on the fly. Then we expand that catalog with
|
|
# any transitive nested wildcards by hitting the live library
|
|
# — picker no longer auto-walks at pick time (issue #2).
|
|
module_list, catalog, _pick_order = deserialize_node_input(wp_modules)
|
|
# `module_list` (not just the wildcard catalog) so a `@{}` ref written
|
|
# on a derivation action — or in its per-instance override — pulls its
|
|
# target in too. See `ref_seed_uuids`.
|
|
catalog = _expand_catalog_via_live_db(catalog, module_list)
|
|
|
|
ctx = _input_ctx(upstream_ctx, upstream_internals, catalog)
|
|
|
|
# Compute the actual chain seed for THIS iteration. When no
|
|
# WP_ContextLoop is upstream, `upstream_internals` lacks both
|
|
# keys → helper returns `widget_seed` unchanged (backwards-compat).
|
|
# When ContextLoop is upstream with override_seed=true, the
|
|
# override replaces the widget seed; loop_index>0 XORs in a
|
|
# stable hash shift so each iteration walks a distinct seed.
|
|
loop_index = int(upstream_internals.get("__wp_loop_index__", 0))
|
|
chain_seed = effective_chain_seed(
|
|
widget_seed=int(seed),
|
|
seed_override=upstream_internals.get("__wp_seed_override__"),
|
|
loop_index=loop_index,
|
|
)
|
|
# Hold base seed: the CONSTANT frame-0 seed so a held module resolves
|
|
# the SAME base-frame value on every iteration. When the loop drives
|
|
# seeds the per-iteration override varies, and effective_chain_seed
|
|
# returns an override verbatim (ignoring loop_index), so loop_index=0
|
|
# alone can't make it constant — read the loop's derived series head
|
|
# (frame 0's seed) instead. Falls back to the widget seed when no loop
|
|
# drives seeds (override off).
|
|
_loop_seeds = upstream_internals.get("__wp_loop_seeds__")
|
|
if isinstance(_loop_seeds, list) and _loop_seeds:
|
|
hold_seed = int(_loop_seeds[0])
|
|
else:
|
|
hold_seed = effective_chain_seed(
|
|
widget_seed=int(seed),
|
|
seed_override=None,
|
|
loop_index=0,
|
|
)
|
|
# Sweep mode with "hold other picks": every frame rolls the
|
|
# non-swept modules on the frame-0 seed, so only the swept
|
|
# wildcards (pinned per frame) change across the grid.
|
|
if upstream_internals.get("__wp_sweep_hold__") is True:
|
|
chain_seed = hold_seed
|
|
# WP Debug detail: per-module explanations and the nested @{} picks.
|
|
# Both are `__`-prefixed, so neither reaches the public context.
|
|
ctx["__wp_explain__"] = True
|
|
ctx["__wp_ref_log__"] = []
|
|
# Inside a loop, frame 0's upstream payload rides along as
|
|
# `__wp_frame0__`. The engine's hold base pass runs on it, so a held
|
|
# module resolves against frame 0's inputs ($iteration and anything
|
|
# upstream derived from it), not this frame's.
|
|
frame0 = upstream_internals.get(FRAME0_KEY)
|
|
if loop_index != 0 and isinstance(frame0, dict):
|
|
frame0_ctx = _input_ctx(
|
|
frame0.get("context") or {}, frame0.get("internals") or {}, catalog,
|
|
)
|
|
frame0_ctx["__wp_explain__"] = True
|
|
frame0_ctx["__wp_ref_log__"] = []
|
|
ctx[FRAME0_CTX_KEY] = frame0_ctx
|
|
ctx = PipelineEngine().run(
|
|
module_list,
|
|
ctx=ctx,
|
|
seed=chain_seed,
|
|
hold_seed=hold_seed,
|
|
loop_index=loop_index,
|
|
)
|
|
|
|
# Hand downstream nodes THIS node's frame-0 output, from the base pass.
|
|
base_ctx = ctx.get("__wp_hold_base_ctx__")
|
|
if loop_index != 0 and isinstance(base_ctx, dict):
|
|
ctx[FRAME0_KEY] = frame0_payload(base_ctx)
|
|
payload = build_payload(
|
|
ctx,
|
|
upstream_debug=upstream_debug,
|
|
seed=chain_seed,
|
|
node_id=_node_id(cls),
|
|
)
|
|
# Emit two seed-tracking values via the UI payload so the
|
|
# frontend `executed` listener (widgets/context.ts) gets
|
|
# authoritative state — works whether the seed was supplied
|
|
# by this node's local widget OR by an upstream node feeding
|
|
# the socket (the local widget value would be stale in that
|
|
# case, but the engine always sees the real value here).
|
|
#
|
|
# - `seed`: the chain seed actually used this run.
|
|
# - `module_seeds`: { module_id → effective_seed_used }
|
|
# for every module that ran. Effective = `locked_seed`
|
|
# when locked, chain seed otherwise. Wrapped in a list
|
|
# because ComfyUI's UI payload convention is value-as-array.
|
|
# Key by per-instance `_uid` when present, falling back to the
|
|
# library `id`. Without this, sibling rows (same library entry
|
|
# instantiated twice in one Context) collapsed into a single
|
|
# map entry — last sibling overwrote the others — and the
|
|
# lock-toggle defaulted every sibling's locked_seed to the
|
|
# last-executed sibling's roll, not its own. Pre-`_uid`
|
|
# entries still get keyed by `id`, matching the old behaviour
|
|
# for unmigrated workflows.
|
|
module_seeds: dict[str, int] = {}
|
|
for entry in ctx.get("__wp_trace__", []):
|
|
uid = entry.get("_uid")
|
|
mid = entry.get("id")
|
|
key = uid if isinstance(uid, str) and uid else mid
|
|
es = entry.get("seed")
|
|
if isinstance(key, str) and key and isinstance(es, int):
|
|
module_seeds[key] = es
|
|
# The Loop's full derived per-iteration series (present only when the
|
|
# Loop drives seeds via override). Captured by the frontend so the
|
|
# per-frame seed lock pins the seed each frame ACTUALLY rolled THIS
|
|
# run — robust to control_after_generate=randomize, where the live
|
|
# seed widget has already rotated to the NEXT run's seed.
|
|
_loop_seeds = upstream_internals.get("__wp_loop_seeds__")
|
|
loop_seeds = [int(s) for s in _loop_seeds] if isinstance(_loop_seeds, list) else []
|
|
return io.NodeOutput(
|
|
payload,
|
|
ui={
|
|
"seed": [int(seed)],
|
|
"module_seeds": [module_seeds],
|
|
"loop_seeds": [loop_seeds],
|
|
},
|
|
)
|