DumiFlex
e41ae4ded0
Wildcard steps list their options with weights and why each was cut, the pool chain, locked seed, separator and matched value; fixed values, combines and node templates show what they read; constraint steps link the steps they re-weighted; Model Info, Toggles, Injector and Context Loop rows carry their own facts.
405 lines
18 KiB
Python
405 lines
18 KiB
Python
"""WP_ContextLoop — optional loop-head node for WP_Context chains.
|
||
|
||
Emits N PIPELINE_CONTEXT outputs (`is_list=True`) so ComfyUI auto-iterates
|
||
the entire downstream chain N times in a single run. Each emitted payload
|
||
carries `__wp_loop_index__` + (optionally) `__wp_seed_override__` in its
|
||
`internals` dict; WP_Context downstream picks these up via
|
||
`engine.seed_derive.effective_chain_seed` and varies per-iteration without
|
||
disturbing locked-seed modules (whose precedence runs at the per-module
|
||
handler level, untouched by this node).
|
||
"""
|
||
|
||
import json
|
||
from typing import Any
|
||
|
||
from comfy_api.latest import io # pyright: ignore[reportMissingImports]
|
||
|
||
from engine.seed_derive import apply_seed_locks, derive_loop_seeds
|
||
from engine.sweep import HOLD_KEY, PIN_KEY, parse_sweep, sweep_frames, sweep_total
|
||
from wp_nodes.types import (
|
||
FRAME0_KEY,
|
||
ContextLoopConfigInput,
|
||
ContextLoopWidgetInput,
|
||
ContextPayload,
|
||
PipelineContext,
|
||
)
|
||
|
||
_STRATEGIES = {"sequential", "hash_index", "prime_stride"}
|
||
|
||
#: Most frames a loop emits (the `count` widget's max).
|
||
_MAX_FRAMES = 999
|
||
|
||
|
||
def _first(values: Any, default: Any = None) -> Any:
|
||
"""A widget's value under ``is_input_list`` (ComfyUI wraps it in a list)."""
|
||
if isinstance(values, list):
|
||
return values[0] if values else default
|
||
return default if values is None else values
|
||
|
||
|
||
def _wired_values(values: Any) -> list[str]:
|
||
"""The ``values`` input's entries: every list entry, one per line, blank
|
||
lines dropped. Empty when nothing is wired."""
|
||
if values is None:
|
||
return []
|
||
out: list[str] = []
|
||
for entry in values if isinstance(values, list) else [values]:
|
||
if entry is None:
|
||
continue
|
||
out.extend(line.strip() for line in str(entry).splitlines() if line.strip())
|
||
return out
|
||
|
||
|
||
def _parse_config(raw: str) -> dict[str, Any]:
|
||
"""Recovery-friendly JSON parse for the DOM widget value.
|
||
|
||
Unknown keys ignored; missing keys backfilled from defaults; type
|
||
mismatches collapse to defaults. Mirrors the TS-side recovery path
|
||
so a workflow with a malformed widget value still loads instead of
|
||
crashing the run.
|
||
"""
|
||
defaults: dict[str, Any] = {
|
||
"strategy": "hash_index",
|
||
"override_seed": False,
|
||
"iteration_var_name": "iteration",
|
||
"bypass": False,
|
||
"iteration_internal": True,
|
||
"total_internal": True,
|
||
"seed_locks": {},
|
||
"bypass_frames": [],
|
||
"sweep": parse_sweep(None),
|
||
}
|
||
if not raw or not isinstance(raw, str):
|
||
return defaults
|
||
try:
|
||
parsed = json.loads(raw)
|
||
except (TypeError, ValueError):
|
||
return defaults
|
||
if not isinstance(parsed, dict):
|
||
return defaults
|
||
out = dict(defaults)
|
||
if parsed.get("strategy") in _STRATEGIES:
|
||
out["strategy"] = parsed["strategy"]
|
||
out["override_seed"] = bool(parsed.get("override_seed", False))
|
||
name = parsed.get("iteration_var_name", "iteration")
|
||
if isinstance(name, str) and name.strip():
|
||
out["iteration_var_name"] = name.strip()
|
||
out["bypass"] = bool(parsed.get("bypass", False))
|
||
out["iteration_internal"] = bool(parsed.get("iteration_internal", True))
|
||
out["total_internal"] = bool(parsed.get("total_internal", True))
|
||
locks_raw = parsed.get("seed_locks", {})
|
||
locks: dict[int, int] = {}
|
||
if isinstance(locks_raw, dict):
|
||
for k, v in locks_raw.items():
|
||
try:
|
||
locks[int(k)] = int(v)
|
||
except (TypeError, ValueError):
|
||
continue
|
||
out["seed_locks"] = locks
|
||
frames_raw = parsed.get("bypass_frames", [])
|
||
frames: set[int] = set()
|
||
if isinstance(frames_raw, list):
|
||
for x in frames_raw:
|
||
if isinstance(x, bool): # bool is an int subclass — exclude
|
||
continue
|
||
if isinstance(x, int):
|
||
n = x
|
||
elif isinstance(x, float) and x.is_integer():
|
||
n = int(x)
|
||
else:
|
||
continue
|
||
if n >= 0:
|
||
frames.add(n)
|
||
out["bypass_frames"] = sorted(frames)
|
||
out["sweep"] = parse_sweep(parsed.get("sweep"))
|
||
return out
|
||
|
||
|
||
class WPContextLoop(io.ComfyNode):
|
||
"""Optional loop head — emits N PIPELINE_CONTEXT outputs in one run."""
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
return io.Schema(
|
||
node_id="WP_ContextLoop",
|
||
display_name="WP Context Loop",
|
||
category="wildcard-pipeline",
|
||
inputs=[
|
||
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."
|
||
),
|
||
),
|
||
io.Int.Input(
|
||
"count",
|
||
default=1,
|
||
min=1,
|
||
max=999,
|
||
tooltip=(
|
||
"How many iterations the series runs (1–999). Set to 1 "
|
||
"to disable looping without removing the node."
|
||
),
|
||
),
|
||
# Widget-only TYPE (`WP_CONTEXT_LOOP_WIDGET`), distinct
|
||
# from the wire-payload type (`WP_CONTEXT_LOOP_CONFIG`)
|
||
# used on the output below + on WP_SeedList's
|
||
# `loop_config` socket. Separating the two types lets the
|
||
# `getCustomWidgets` factory register against the widget
|
||
# type only — the wire type has NO factory, so ComfyUI
|
||
# renders consumers (SeedList) as plain sockets and
|
||
# Loop's own input stays a DOM widget. The previous
|
||
# single-type design forced an `inputName` gate in
|
||
# main.ts that ended up hiding SeedList's socket
|
||
# entirely.
|
||
io.String.Input(
|
||
"values",
|
||
optional=True,
|
||
force_input=True,
|
||
tooltip=(
|
||
"Optional text from another node: one frame per line (or per "
|
||
"list entry), each writing its line to $iteration_value. Sets "
|
||
"the frame count; with a sweep on, every value runs every "
|
||
"combination."
|
||
),
|
||
),
|
||
ContextLoopWidgetInput.Input(
|
||
"wp_context_loop_config",
|
||
socketless=True,
|
||
default="{}",
|
||
),
|
||
],
|
||
outputs=[
|
||
# `is_output_list=True` is the V3 contract for emitting a
|
||
# list that ComfyUI auto-iterates downstream. Not `is_list`
|
||
# (that was wrong — stub accepted any kwarg, real API
|
||
# rejects). See `comfy_api/latest/_io.py:Output.__init__`.
|
||
PipelineContext.Output("context", is_output_list=True),
|
||
# Single-value side output: the resolved config payload
|
||
# so sibling helpers (WP_SeedList etc.) can mirror the
|
||
# loop's count / strategy / base_seed without the user
|
||
# double-wiring widgets. Plain (non-list) — receiving
|
||
# nodes get one dict per graph run regardless of the
|
||
# context-fan-out happening on the first output.
|
||
ContextLoopConfigInput.Output("loop_config"),
|
||
],
|
||
# A list wired into `values` must arrive in ONE call (one loop),
|
||
# not run the whole loop once per entry.
|
||
is_input_list=True,
|
||
not_idempotent=False,
|
||
# The node id stamps the Loop's trace rows so WP Debug groups
|
||
# them under this node.
|
||
hidden=[io.Hidden.unique_id],
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, seed, count, wp_context_loop_config, values=None):
|
||
seed = _first(seed, 0)
|
||
count = _first(count, 1)
|
||
cfg = _parse_config(_first(wp_context_loop_config, "{}"))
|
||
wired_values = _wired_values(values)
|
||
override_seed = bool(cfg["override_seed"])
|
||
strategy = cfg["strategy"]
|
||
iteration_var_name = cfg["iteration_var_name"]
|
||
bypass = bool(cfg["bypass"])
|
||
iteration_internal = bool(cfg["iteration_internal"])
|
||
total_internal = bool(cfg["total_internal"])
|
||
seed_locks: dict[int, int] = cfg["seed_locks"] # type: ignore[assignment]
|
||
bypass_frames: list[int] = cfg["bypass_frames"] # type: ignore[assignment]
|
||
|
||
# Resolve which frames to emit + the seed series they draw from.
|
||
#
|
||
# - Whole-loop `bypass` (existing behavior) collapses to a single
|
||
# run while preserving the list-shape output contract.
|
||
# - Per-frame bypass (`bypass_frames`, 0-based) skips individual
|
||
# iterations WITHOUT renumbering the kept ones: each kept frame
|
||
# keeps its original index as `__wp_loop_index__`, so it draws
|
||
# `derived[idx]` and its per-frame overrides still key correctly.
|
||
# Out-of-range indices are ignored; an all-bypassed config floors
|
||
# to a full run so the `is_output_list` socket is never empty.
|
||
# The seed series is derived for the FULL count (not the reduced
|
||
# one) so kept frame N always draws `derived[N]`.
|
||
has_override = override_seed
|
||
configured_count = max(1, int(count))
|
||
# Sweep mode: one frame per combination of the chosen wildcards'
|
||
# options, capped at the sweep limit. The combination count replaces
|
||
# `count` (the widget keeps `count` in step, but the config is the
|
||
# authority so an API-queued workflow still sweeps correctly).
|
||
sweep = cfg["sweep"]
|
||
sweep_pins: list[dict[str, str]] = []
|
||
if sweep["enabled"] and sweep["axes"] and not bypass:
|
||
sweep_pins = sweep_frames(sweep["axes"], sweep["limit"])
|
||
if sweep_pins:
|
||
configured_count = len(sweep_pins)
|
||
# A wired value list sets the frame count: every value runs every
|
||
# sweep combination (the value is the outer, slower axis).
|
||
combos = max(1, len(sweep_pins))
|
||
if wired_values and not bypass:
|
||
wired_values = wired_values[: max(1, _MAX_FRAMES // combos)]
|
||
configured_count = len(wired_values) * combos
|
||
if bypass:
|
||
total_count = 1
|
||
kept = [0]
|
||
else:
|
||
in_range_bypassed = {
|
||
i for i in bypass_frames if 0 <= i < configured_count
|
||
}
|
||
kept = [i for i in range(configured_count) if i not in in_range_bypassed]
|
||
if not kept:
|
||
kept = list(range(configured_count))
|
||
total_count = configured_count
|
||
derived = (
|
||
apply_seed_locks(
|
||
derive_loop_seeds(int(seed), total_count, strategy),
|
||
seed_locks,
|
||
)
|
||
if has_override else None
|
||
)
|
||
# Build the loop_config payload once so the same dict drives both
|
||
# the per-iteration internals AND the side output. `count` is the
|
||
# configured/effective total (so WP_SeedList + the seed modal
|
||
# re-derive the identical full series); `active_count` is how many
|
||
# frames actually emit after per-frame bypass. `base_seed` carries
|
||
# the user's `seed` input verbatim so WP_SeedList can re-derive the
|
||
# same series locally.
|
||
loop_config_payload: dict[str, object] = {
|
||
"count": total_count,
|
||
"active_count": len(kept),
|
||
# The 0-based indices this loop actually emits (== range(count)
|
||
# for a normal run). A wired WP_SeedList mirrors this so it emits
|
||
# ONE seed per kept frame — otherwise it would derive the full
|
||
# `count` seeds and the sampler would run more times than the
|
||
# loop has prompts.
|
||
"kept_indices": list(kept),
|
||
"strategy": strategy,
|
||
"base_seed": int(seed),
|
||
"override_seed": has_override,
|
||
}
|
||
if sweep_pins:
|
||
# Combinations before the limit, so a UI can say "64 of 120".
|
||
loop_config_payload["sweep_total"] = sweep_total(sweep["axes"])
|
||
if wired_values:
|
||
loop_config_payload["values_count"] = len(wired_values)
|
||
|
||
# Build the internal-flags map once — same across iterations
|
||
# (per-iteration values differ but their internal-ness doesn't).
|
||
# Engine reads `__wp_internal_flags__` from the payload internals
|
||
# at downstream merge time; PromptAssembler then strips internal
|
||
# keys at render so they never leak into prompts.
|
||
total_var_name = f"{iteration_var_name}_total"
|
||
value_var_name = f"{iteration_var_name}_value"
|
||
internal_flags: dict[str, bool] = {}
|
||
if iteration_internal:
|
||
internal_flags[iteration_var_name] = True
|
||
if total_internal:
|
||
internal_flags[total_var_name] = True
|
||
|
||
def frame(idx: int) -> tuple[dict[str, str], dict[str, object]]:
|
||
# User-facing iteration counter is the frame's ORIGINAL 1-based
|
||
# number (bypass keeps identity — frame 5 stays "5 of 18", not
|
||
# renumbered). The internal `__wp_loop_index__` below stays the
|
||
# 0-based original index because `derived[idx]` indexes the
|
||
# full series + `effective_chain_seed` math assumes zero-origin
|
||
# (loop_index=0 means "no XOR shift"). `total` is the configured
|
||
# count so bypassing a frame never rewrites the other frames'
|
||
# "of N" text (which would bust their downstream caches).
|
||
context_vars = {
|
||
iteration_var_name: str(idx + 1),
|
||
total_var_name: str(total_count),
|
||
}
|
||
if wired_values:
|
||
# Bypass collapses to frame 0, which reads the first value.
|
||
pick = min(idx // combos, len(wired_values) - 1)
|
||
context_vars[value_var_name] = wired_values[pick]
|
||
internals: dict[str, object] = {
|
||
"__wp_loop_index__": idx,
|
||
"__wp_seed_override__": (derived[idx] if has_override else None),
|
||
}
|
||
# Only stamp the derived series when override is active —
|
||
# otherwise the values would be meaningless (downstream uses
|
||
# widget seeds). WP_Debug renders this only when present.
|
||
if has_override:
|
||
internals["__wp_loop_seeds__"] = list(derived or [])
|
||
if sweep_pins:
|
||
internals[PIN_KEY] = dict(sweep_pins[idx % combos])
|
||
internals[HOLD_KEY] = bool(sweep["hold_others"])
|
||
if internal_flags:
|
||
# Merge by-value so successive iterations don't share
|
||
# the same dict reference; downstream WP_Context copies
|
||
# internals defensively but explicit copy here is safer.
|
||
internals["__wp_internal_flags__"] = dict(internal_flags)
|
||
return context_vars, internals
|
||
|
||
hidden = getattr(cls, "hidden", None)
|
||
node_id = getattr(hidden, "unique_id", None) if hidden is not None else None
|
||
|
||
def frame_trace(idx: int, context_vars: dict[str, str]) -> list[dict[str, Any]]:
|
||
# One WP Debug trace row per variable the Loop writes, so Debug
|
||
# names the Loop as their writer and shows the frame it ran.
|
||
detail: dict[str, Any] = {"frame": idx + 1, "of": total_count}
|
||
if has_override and derived:
|
||
detail["seed"] = int(derived[idx])
|
||
if sweep_pins:
|
||
detail["combo"] = idx % combos + 1
|
||
detail["combos"] = combos
|
||
if wired_values:
|
||
detail["value_line"] = min(idx // combos, len(wired_values) - 1) + 1
|
||
detail["values"] = len(wired_values)
|
||
if bypass:
|
||
detail["bypass"] = True
|
||
if has_override:
|
||
detail["strategy"] = strategy
|
||
if len(kept) < total_count:
|
||
detail["skipped"] = total_count - len(kept)
|
||
if sweep_pins:
|
||
# This frame's sweep pins: module uid → option id.
|
||
detail["pins"] = dict(sweep_pins[idx % combos])
|
||
detail["hold_others"] = bool(sweep["hold_others"])
|
||
rows: list[dict[str, Any]] = []
|
||
for name, value in context_vars.items():
|
||
row: dict[str, Any] = {
|
||
"node": "WP_ContextLoop",
|
||
"binding": name,
|
||
"internal": bool(internal_flags.get(name)),
|
||
"type": "value" if name == value_var_name else (
|
||
"total" if name == total_var_name else "iteration"
|
||
),
|
||
"value": value,
|
||
"detail": detail,
|
||
}
|
||
if node_id is not None:
|
||
row["node_id"] = str(node_id)
|
||
rows.append(row)
|
||
return rows
|
||
|
||
payloads = []
|
||
for idx in kept:
|
||
context_vars, internals = frame(idx)
|
||
if idx != 0:
|
||
# Frame 0's payload rides along so a held module downstream
|
||
# resolves against frame 0's inputs on every frame (its
|
||
# $iteration, and whatever a Context derives from it).
|
||
ctx0, int0 = frame(0)
|
||
internals[FRAME0_KEY] = {"context": ctx0, "internals": int0}
|
||
payloads.append(ContextPayload(
|
||
context=context_vars,
|
||
debug={"__wp_trace__": frame_trace(idx, context_vars)},
|
||
internals=internals,
|
||
))
|
||
# Surface the per-iteration series to the node UI so the frontend can
|
||
# offer "lock previous" in the seed modal (captured in context_loop.ts).
|
||
# Only meaningful when the Loop drives the seeds (override on); an empty
|
||
# list otherwise greys the button out per-frame.
|
||
ui_seeds = [int(s) for s in derived] if (has_override and derived) else []
|
||
ui: dict[str, list[Any]] = {"loop_seeds": [ui_seeds]}
|
||
if wired_values:
|
||
# The widget shows the frame count the wire set, after a run.
|
||
ui["loop_values"] = [{"count": len(wired_values), "frames": total_count}]
|
||
return io.NodeOutput(payloads, loop_config_payload, ui=ui)
|