SolitaryThinker
b254d1affe
Two more pipelines stress-testing the design, plus a BAGEL-placement note in designv4. Both fit with no new runtime primitive (109 tests pass). Pipeline 1 — cross-model T2I→I2V (program/workflow.py, models/image_video/): - Realizes ProgramKind.WORKFLOW as a thin multi-instance orchestrator ABOVE the engine (the hot path stays single-instance). A Program composes one model's loops; a Workflow chains full engine.run calls across distinct cards, threading artifacts. (LTX-2 already covers same-card multi-stage; cross-model — FLUX→Wan — is the new capability the single-instance runner can't express.) - flux-t2i (text→image) and wan-i2v (text+image→video) cards; the I2V program folds the conditioning image into text_embeds so WanDenoiseLoop is unchanged. - Each model keeps its own interleave-parity guarantee (crossing instances is a Workflow boundary, not a loop step). 7 tests incl. video-depends-on-image. Pipeline 2 — N-way joint RL (training/methods/joint_multi_rl.py, models/multi_expert/): - JointMultiExpertRL generalizes UnifiedRLMethod (N=2) to N refiner LMs + a generator: one reward → one group advantage → N token-PG updates + 1 FlowGRPO PPO update, N+1 independent WeightSyncPlans. Proves the substrate was already N-ready (card holds N components/loops; per-component weight-sync; dict grad targets) — only the method body looped over two; now it loops over a list. - credit="per_expert" learns all N cleanly; credit="shared" (faithful to UniRL) works but is noisier — the honest multi-agent credit-assignment result, a reward-shaping choice, not a substrate limit. 6 tests (N=1,3,4; prompt-only). Fix — flow_sde_ml_velocity (loop/sampler.py): the toy FlowGRPO generator update targeted the velocity the model already produced (a no-op once guidance_scale=1 was set for the C2 identity; the unified generator moved only on ~1e-7 noise). The correct PG surrogate targets the max-likelihood velocity of the realized sample — nonzero at ratio==1. Both UniRL and N-way generators now learn for real; the C2 ratio==1 identity still holds (measured before the update). BAGEL: MoT/shared-weight (one transformer on both loops), same row as Cosmos3; real BAGEL's co-resident experts are expressible via the expert-routing policy (partial sharing) — captured in designv4 §2.3.
62 lines
3.1 KiB
Python
62 lines
3.1 KiB
Python
"""Workflow — multi-MODEL composition above the engine (design_v3 §13; ``ProgramKind.WORKFLOW``).
|
|
|
|
A ``Program`` composes the loops of ONE resident model instance (the omni/MoT default — every
|
|
``ModelLoopNode`` resolves on the program's single instance). A ``Workflow`` composes *across model
|
|
instances*: each stage is a full ``engine.run`` on a (possibly different) registered card, with typed
|
|
artifacts threaded stage→stage. This is the realistic "T2I model **followed by** I2V model" pipeline
|
|
(FLUX→Wan): two distinct cards, two distinct weight sets, chained.
|
|
|
|
The split mirrors the doc's two layers exactly:
|
|
* **Program** — single-model, step-interleaved loop composition (the hot path; unchanged).
|
|
* **Workflow** — multi-model orchestration; a thin layer that calls ``engine.run`` per stage and
|
|
passes artifacts forward. No change to the engine's single-instance program runner — so cross-
|
|
model chaining costs nothing on the per-step scheduling path. The engine already holds every
|
|
instance in its registry; the Workflow just selects which card each stage runs on.
|
|
|
|
Why not put cross-model in ``Program``? Because a step-interleaved program shares one resident
|
|
instance, one cache scope, one parity gate. Crossing instances is a *boundary* (different weights,
|
|
different caches, an artifact hand-off) — exactly a Workflow edge, not a loop step. Keeping it out of
|
|
the program runner is what preserves the interleave-parity guarantee for the single-model case.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Callable
|
|
|
|
|
|
@dataclass
|
|
class WorkflowStage:
|
|
"""One model invocation. ``make_request(state)`` builds this stage's Request from the accumulated
|
|
``state`` (the initial inputs plus every prior stage's artifacts, keyed ``"<stage_label>:<name>"``
|
|
and also under ``"prev"`` for the immediately-preceding stage's artifact dict)."""
|
|
model_id: str
|
|
make_request: Callable[[dict[str, Any]], Any]
|
|
label: str = ""
|
|
|
|
|
|
@dataclass
|
|
class Workflow:
|
|
workflow_id: str
|
|
stages: list[WorkflowStage] = field(default_factory=list)
|
|
|
|
def run(self, engine: Any, **initial: Any) -> Any:
|
|
"""Execute the stages in order on ``engine``; return the final stage's Output.
|
|
|
|
Each stage's artifacts are merged into ``state`` so downstream stages can read them
|
|
(e.g. the I2V stage reads the T2I stage's ``image`` artifact)."""
|
|
if not self.stages:
|
|
raise ValueError(f"workflow {self.workflow_id!r} has no stages")
|
|
state: dict[str, Any] = dict(initial)
|
|
out = None
|
|
for i, stage in enumerate(self.stages):
|
|
label = stage.label or stage.model_id
|
|
req = stage.make_request(state)
|
|
if req.model_id != stage.model_id:
|
|
raise ValueError(f"workflow {self.workflow_id!r} stage {label!r}: request model "
|
|
f"{req.model_id!r} != stage model {stage.model_id!r}")
|
|
out = engine.run(req)
|
|
state["prev"] = out.artifacts
|
|
for name, art in out.artifacts.items():
|
|
state[f"{label}:{name}"] = art
|
|
return out
|