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.
110 lines
4.6 KiB
Python
110 lines
4.6 KiB
Python
"""Model definitions — concrete (recipe, runtime) cards (design_v3 §4).
|
|
|
|
Phase 1: Wan2.1-1.3B (T2V), LTX2.3 (2-stage distilled), Wan-causal (self-forcing student).
|
|
``build_default_engine`` loads all three onto one engine (one resident instance per card).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from .bagel import build_bagel_card, build_bagel_program
|
|
from .cosmos3 import build_cosmos3_card, build_cosmos3_program
|
|
from .image_video import (
|
|
build_flux_t2i_card,
|
|
build_flux_t2i_program,
|
|
build_t2i_then_i2v_workflow,
|
|
build_wan_i2v_card,
|
|
build_wan_i2v_program,
|
|
)
|
|
from .ltx2 import build_ltx2_card, build_ltx2_program
|
|
from .qwen_omni import build_qwen_omni_card, build_qwen_omni_program
|
|
from .unified import build_unified_card, build_unified_program
|
|
from .wan21 import build_wan21_card, build_wan_t2v_program
|
|
from .wan_causal import build_wan_causal_card, build_wan_causal_program
|
|
|
|
__all__ = [
|
|
"build_wan21_card", "build_wan_t2v_program",
|
|
"build_ltx2_card", "build_ltx2_program",
|
|
"build_wan_causal_card", "build_wan_causal_program",
|
|
"build_cosmos3_card", "build_cosmos3_program",
|
|
"build_bagel_card", "build_bagel_program",
|
|
"build_qwen_omni_card", "build_qwen_omni_program",
|
|
"build_unified_card", "build_unified_program",
|
|
"build_flux_t2i_card", "build_flux_t2i_program", "build_wan_i2v_card", "build_wan_i2v_program",
|
|
"build_t2i_then_i2v_workflow",
|
|
"build_default_engine", "build_omni_engine", "build_unified_engine", "build_image_video_engine",
|
|
]
|
|
|
|
_BUILDERS = [
|
|
(build_wan21_card, build_wan_t2v_program),
|
|
(build_ltx2_card, build_ltx2_program),
|
|
(build_wan_causal_card, build_wan_causal_program),
|
|
]
|
|
|
|
# Phase-2 omni cards: MoT shared-weight (Cosmos3/BAGEL) + the cascaded thinker→talker→vocoder
|
|
# (Qwen-Omni, three disjoint experts / three loop types in one request).
|
|
_OMNI_BUILDERS = [
|
|
(build_cosmos3_card, build_cosmos3_program),
|
|
(build_bagel_card, build_bagel_program),
|
|
(build_qwen_omni_card, build_qwen_omni_program),
|
|
]
|
|
|
|
|
|
def build_default_engine(engine: Any = None) -> Any:
|
|
"""Register Wan2.1, LTX2.3, and Wan-causal onto one engine (one resident instance each)."""
|
|
from ..cache import CacheManager
|
|
from ..card import load_card
|
|
from ..runtime import Engine
|
|
eng = engine if engine is not None else Engine()
|
|
for build_card, build_program in _BUILDERS:
|
|
card = build_card()
|
|
inst = load_card(card, cache_manager=CacheManager.from_card(card))
|
|
eng.register(card.model_id, inst, build_program())
|
|
return eng
|
|
|
|
|
|
def build_omni_engine(engine: Any = None) -> Any:
|
|
"""Register the phase-2 omni cards (Cosmos3 + BAGEL/lance) onto one engine.
|
|
|
|
Each is ONE resident MoT instance whose ``transformer`` runs BOTH an ar_decode loop and a
|
|
diffusion_denoise loop (shared weights) — the §16 'true omni/MoT serving' claim made native.
|
|
"""
|
|
from ..cache import CacheManager
|
|
from ..card import load_card
|
|
from ..runtime import Engine
|
|
eng = engine if engine is not None else Engine()
|
|
for build_card, build_program in _OMNI_BUILDERS:
|
|
card = build_card()
|
|
inst = load_card(card, cache_manager=CacheManager.from_card(card))
|
|
eng.register(card.model_id, inst, build_program())
|
|
return eng
|
|
|
|
|
|
def build_image_video_engine(engine: Any = None) -> Any:
|
|
"""Register the T2I (``flux-t2i``) and I2V (``wan-i2v``) cards on one engine — two *separate*
|
|
models chained by ``build_t2i_then_i2v_workflow`` (design_v3 §13 cross-model composition)."""
|
|
from ..cache import CacheManager
|
|
from ..card import load_card
|
|
from ..runtime import Engine
|
|
eng = engine if engine is not None else Engine()
|
|
for build_card, build_program in [(build_flux_t2i_card, build_flux_t2i_program),
|
|
(build_wan_i2v_card, build_wan_i2v_program)]:
|
|
card = build_card()
|
|
inst = load_card(card, cache_manager=CacheManager.from_card(card))
|
|
eng.register(card.model_id, inst, build_program())
|
|
return eng
|
|
|
|
|
|
def build_unified_engine(engine: Any = None) -> Any:
|
|
"""Register the unified LM+generator card (UniRL/PromptRL): a prompt-refiner ``llm`` expert + a
|
|
flow ``transformer`` expert on the SAME engine, each driving its own loop. Two separate experts,
|
|
one request, both trainable under one RL reward (the §10 joint-RL stress test)."""
|
|
from ..cache import CacheManager
|
|
from ..card import load_card
|
|
from ..runtime import Engine
|
|
eng = engine if engine is not None else Engine()
|
|
card = build_unified_card()
|
|
inst = load_card(card, cache_manager=CacheManager.from_card(card))
|
|
eng.register(card.model_id, inst, build_unified_program())
|
|
return eng
|