[feat] mini-fastvideo phase 2: omni/MoT — Cosmos3 + canonical vllm-omni (BAGEL/lance)

One resident MoT instance runs BOTH an ar_decode loop and a diffusion_denoise loop on shared weights
(the §16 claim no DAG-of-engines can express), with both loops runtime-visible: the scheduler prices
ar_token AND diffusion_step WorkUnits — unlike vllm-omni's opaque DIFFUSION stage the scheduler never
sees inside.

- ARDecodeLoop: token decode until EOS/max_tokens, paged text-KV — the omni AR pathway (loop/§5).
- ToyMoTDiT: one module exposing an und pathway (ar_forward) AND a gen pathway (denoise __call__);
  ToyTokenizer. Binding both loops to one instance = shared weights, no duplication.
- models/cosmos3/: tokenize → reason(ar_decode) → pack(tokens→conditioning) → diffusion_denoise →
  vae_decode; sound_vae declared optional_for non-t2vs (the lazy-component P8 fix, not an env-var hack).
- models/bagel/: the canonical vllm-omni model — generate_text(ar_decode) → generate_image(diffusion),
  text+image outputs, both loops step-scheduled.
- The diffusion loop is WanDenoiseLoop reused (one loop definition bound to the MoT module).
- build_omni_engine() + an omni worked example; 7 omni tests (shared-instance, both-kinds-scheduled,
  interleave parity across loop types, lazy sound_vae). 61 tests pass (pytest + zero-dep runner).

mini_fastvideo/README.md

@@ -17,9 +17,7 @@ point of the (recipe, runtime) separation.
 |---|---|
 | **Inference** | Wan2.1-1.3B (T2V, bidirectional diffusion) · LTX2.3 (two-stage distilled: 8-step base → upsample → 3-step refine) · Wan-causal (chunked/streaming rollout, the self-forcing student) |
 | **Training** (all on Wan2.1-1.3B) | finetuning (flow-match) · DMD2 distillation · DiffusionNFT (likelihood-free RL) · self-forcing (causal, on Wan-causal) |
-
-Phase 2 (designed-for, not yet built): a Cosmos3 omni card and a canonical vllm-omni-style MoT model.
-The spine is already omni-ready — see *Omni readiness* below.
+| **Omni / MoT** (phase 2) | Cosmos3 (reasoner `ar_decode` → joint `diffusion_denoise`) · BAGEL/lance (`generate_text` → `generate_image`) — each ONE resident MoT instance running both loop types on shared weights |
 
 ## The one invariant (design_v3 §1)
 
@@ -93,15 +91,28 @@ Each `ComponentSpec` records the real `load_id` (e.g.
    scheduler (LTX) inside the loop's step body.
 The loops, policies, scheduler, caches, parity gates, and training methods are unchanged.
 
-## Omni readiness (phase 2)
+## Omni / MoT (phase 2 — implemented)
 
-The spine already carries the minimal concessions so Cosmos3 / vllm-omni MoT is additive, not a
-refactor: `LoopKind`/`WorkUnitKind` include `ar_decode`/`ar_token`/`chunk_step`; `ModelInstance` runs
-multiple loops against shared components; `LoopSpec.shared_weight_components` + `extension_schema`;
-`LoopState.extension` for packed factored sequences; `paged_kv` cache class; the interleave gate
-generalizes to any loop. Adding an omni card means: a `Cosmos3PackedSequence` state extension, an
-`ar_decode` loop bound to the same `transformer` as the `diffusion_denoise` loop, and a program
-`ar_decode(reasoner) → pack → diffusion_denoise(joint) → decoders`.
+The omni-ready spine paid off: adding omni was **additive, not a refactor**. Two omni cards now ship,
+each as ONE resident MoT instance whose `transformer` binds *both* an `ar_decode` loop and a
+`diffusion_denoise` loop (`shared_weight_components=["transformer"]` on both) — the §16 claim no
+DAG-of-engines can express, made native:
+
+- **`models/cosmos3/`** — `tokenize → reason (ar_decode) → pack → diffusion_denoise → vae_decode`.
+  The reasoner upsamples the prompt on the und pathway; its tokens condition the joint denoise on the
+  same weights. `sound_vae` is declared `optional_for` non-t2vs (the lazy-component P8 fix).
+- **`models/bagel/`** — the canonical vllm-omni model: `generate_text (ar_decode) → generate_image
+  (diffusion_denoise)`. Unlike vllm-omni's opaque `DIFFUSION` stage, **both loops are runtime-visible**:
+  the scheduler prices `ar_token` *and* `diffusion_step` WorkUnits (verified in `test_omni.py`).
+
+The diffusion loop is literally `WanDenoiseLoop` (the same step body the engine serves for Wan), bound
+to the MoT module — one loop definition, reused. The interleave gate generalizes across loop types
+(`test_omni.py::test_omni_interleave_parity_holds_across_loop_types`). Run it:
+`python3 -c "from mini_fastvideo.models import build_omni_engine; ..."` or see `mini_fastvideo/examples.py`.
+
+Still future work (declared on the spine, not yet built): the packed factored-sequence
+(`[text|vision|action|sound]`) `LoopState.extension`, joint multi-modality denoise with per-modality
+CFG, and world-model `chunk_rollout` for action-conditioned omni.
 
 ## Deliberately out of scope (per design_v3 non-goals)
 

mini_fastvideo/examples.py

@@ -8,10 +8,10 @@ from __future__ import annotations
 
 import numpy as np
 
-from .models import build_default_engine
+from .models import build_default_engine, build_omni_engine
 from .models.wan21 import build_wan21_card
 from .parity import assert_interleave_parity
-from .request import DiffusionParams, OutputSpec, TaskType, make_request
+from .request import DiffusionParams, OutputSpec, SamplingParams, TaskType, make_request
 from .training import build_diffusion_nft
 
 
@@ -60,9 +60,26 @@ def example_d_rl_rollout() -> None:
           f"(K samples encode the prompt once — the 24× reduction)")
 
 
+def example_g_omni_mot() -> None:
+    print("\n(g) Omni / MoT (§16): ONE resident instance runs AR + diffusion loops on shared weights")
+    eng = build_omni_engine()
+    o = eng.run(make_request(TaskType.T2V, "cosmos3-vfm", "a phoenix",
+                             sampling=SamplingParams(max_tokens=6, seed=1),
+                             diffusion=DiffusionParams(num_steps=4, seed=1)))
+    print(f"    Cosmos3 (reason→joint denoise): text={o.artifacts['text'].text} "
+          f"video={o.artifacts['video'].frames.shape}")
+    o2 = eng.run(make_request(TaskType.T2I, "bagel-mot", "a teapot",
+                              sampling=SamplingParams(max_tokens=6, seed=2),
+                              diffusion=DiffusionParams(num_steps=4, seed=2)))
+    print(f"    BAGEL (generate_text→generate_image): text={o2.artifacts['text'].text} "
+          f"image={o2.artifacts['image'].tensor.shape}")
+    print(f"    scheduler priced BOTH WorkUnit kinds (runtime-visible, not one opaque stage): "
+          f"{dict(eng.admission.metrics.by_kind)}")
+
+
 def main() -> None:
     print("=" * 78)
-    print("mini-fastvideo — worked examples (design_v3 §15). One runtime, many loops.")
+    print("mini-fastvideo — worked examples (design_v3 §15-16). One runtime, many loops.")
     print("=" * 78)
     eng = build_default_engine()
     print(f"registered (recipe, runtime) cards: {list(eng._registry)}")
@@ -71,6 +88,7 @@ def main() -> None:
     example_c_causal_streaming(eng)
     example_c2_interleave_gate(eng)
     example_d_rl_rollout()
+    example_g_omni_mot()
     print("\n" + "=" * 78)
     print("All examples ran on CPU with numpy toy components. The architecture (cards, driven")
     print("loops, scheduler, caches, parity, training-on-shared-loops) is real; the neural")

mini_fastvideo/models/__init__.py

@@ -7,6 +7,8 @@ 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 .ltx2 import build_ltx2_card, build_ltx2_program
 from .wan21 import build_wan21_card, build_wan_t2v_program
 from .wan_causal import build_wan_causal_card, build_wan_causal_program
@@ -15,7 +17,9 @@ __all__ = [
     "build_wan21_card", "build_wan_t2v_program",
     "build_ltx2_card", "build_ltx2_program",
     "build_wan_causal_card", "build_wan_causal_program",
-    "build_default_engine",
+    "build_cosmos3_card", "build_cosmos3_program",
+    "build_bagel_card", "build_bagel_program",
+    "build_default_engine", "build_omni_engine",
 ]
 
 _BUILDERS = [
@@ -24,6 +28,12 @@ _BUILDERS = [
     (build_wan_causal_card, build_wan_causal_program),
 ]
 
+# Phase-2 omni cards (MoT, one resident instance running AR + diffusion loops on shared weights)
+_OMNI_BUILDERS = [
+    (build_cosmos3_card, build_cosmos3_program),
+    (build_bagel_card, build_bagel_program),
+]
+
 
 def build_default_engine(engine: Any = None) -> Any:
     """Register Wan2.1, LTX2.3, and Wan-causal onto one engine (one resident instance each)."""
@@ -36,3 +46,20 @@ def build_default_engine(engine: Any = None) -> Any:
         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

mini_fastvideo/models/bagel/__init__.py

@@ -0,0 +1,7 @@
+"""BAGEL/lance — canonical vllm-omni MoT (AR generate_text + diffusion generate_image). Phase 2."""
+from __future__ import annotations
+
+from .card import build_bagel_card
+from .program import build_bagel_program
+
+__all__ = ["build_bagel_card", "build_bagel_program"]

mini_fastvideo/models/bagel/card.py

@@ -0,0 +1,85 @@
+"""BAGEL/lance-style MoT ModelCard — the canonical vllm-omni omni model (design_v3 §1, §16, §19).
+
+vllm-omni's ``bagel_single_stage``/``lance`` prove one resident MoT instance can run both AR
+``generate_text`` and diffusion ``generate_image`` on co-resident experts in a single request — BUT
+they bury that interleaving inside one opaque ``DIFFUSION`` stage the scheduler never sees inside
+(``max_num_running_reqs=1``). This card expresses the same shared-weight MoT, but makes both loops
+**runtime-visible, step-scheduled, and batchable**: ``generate_text`` (ar_decode) and
+``generate_image`` (diffusion_denoise) both bind the one resident ``transformer``, and every AR token
+and denoise step is a WorkUnit the scheduler can interleave and price. That visibility is the
+differentiation the §1 thesis names.
+"""
+from __future__ import annotations
+
+from ..._enums import Capability, ConsistencyLevel, LoopKind, WorkUnitKind
+from ...card import (
+    CacheContract,
+    CapabilityMatrix,
+    ComponentSpec,
+    CostModel,
+    LoopSpec,
+    ModelCard,
+    ParallelismContract,
+    ParitySpec,
+    PrecisionContract,
+    RecipeSpec,
+)
+from ...loop.policies import ClassicCFG, FlowShiftPolicy, NoRouting, PrecisionPolicy
+from ...parallel import ParallelPlan
+from ..backend import ToyMoTDiT, ToyTokenizer, ToyVAE, _seed_from
+from ..omni import ARDecodeLoop
+from ..wan21.loop import WanDenoiseLoop
+
+
+def build_bagel_card(model_id: str = "bagel-mot") -> ModelCard:
+    seed = _seed_from(model_id)
+    ar_cost = CostModel(kind=WorkUnitKind.AR_TOKEN, base_seconds=5e-5, per_unit_seconds=1e-7)
+    dn_cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7)
+    cfg, flow = ClassicCFG(), FlowShiftPolicy(shift=3.0)
+    precision, expert = PrecisionPolicy(), NoRouting("transformer")
+
+    def text_factory():
+        return ARDecodeLoop(loop_id="generate_text", transformer_id="transformer", cost=ar_cost, max_tokens=6)
+
+    def image_factory():
+        return WanDenoiseLoop(loop_id="generate_image", cfg=cfg, flow_shift=flow,
+                              precision=precision, expert=expert, cost=dn_cost)
+
+    components = {
+        "tokenizer": ComponentSpec("tokenizer", kind="tokenizer", factory=lambda inst: ToyTokenizer(),
+                                   required_for={"reason", "t2i"}),
+        "transformer": ComponentSpec(
+            "transformer", kind="dit",
+            load_id="vllm_omni.diffusion.models.bagel:BagelTransformer",
+            factory=lambda inst: ToyMoTDiT(seed=seed),
+            resident_for=["generate_text", "generate_image"],   # one resident copy for BOTH loops
+            required_for={"reason", "t2i"}),
+        "vae": ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(), required_for={"t2i"}),
+    }
+    loops = {
+        "generate_text": LoopSpec("generate_text", kind=LoopKind.AR_DECODE,
+                                  work_unit_kind=WorkUnitKind.AR_TOKEN, step_cost_model=ar_cost,
+                                  shared_weight_components=["transformer"], cache_policy=["paged_kv"],
+                                  loop_factory=text_factory),
+        "generate_image": LoopSpec("generate_image", kind=LoopKind.DIFFUSION_DENOISE,
+                                   work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=dn_cost,
+                                   shared_weight_components=["transformer"], cache_policy=["feature"],
+                                   loop_factory=image_factory),
+    }
+    card = ModelCard(
+        model_id=model_id, family="bagel", components=components, loops=loops,
+        capabilities=CapabilityMatrix.of(Capability.TEXT_TO_IMAGE, Capability.REASONING_TEXT,
+                                         Capability.VAE_DECODE),
+        recipe=RecipeSpec(method="base", assumes_loop="generate_image",
+                          assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
+        parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
+        caches={
+            "feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True),
+            "paged_kv": CacheContract("paged_kv", max_bytes=1 << 24, block_bytes=1 << 12,
+                                      reuse_across_requests=False),
+        },
+        precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
+        parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
+                                        default_plan=ParallelPlan.single()),
+    )
+    return card.validate()

mini_fastvideo/models/bagel/program.py

@@ -0,0 +1,31 @@
+"""BAGEL/lance program (design_v3 §16): one resident MoT instance, two runtime-visible loops.
+
+    tokenize → generate_text (ar_decode) → pack(text→cond) → generate_image (diffusion) → vae_decode
+
+generate_text and generate_image hit the SAME resident weights; their steps are WorkUnits the
+scheduler interleaves — the shared-weight MoT made native, not a monolith bypassing the abstraction.
+"""
+from __future__ import annotations
+
+from ...program import ComponentNode, ModelLoopNode, Program, ProgramKind
+from ..omni import emit_text_node, pack_cond_from_tokens, tokenize_node, vae_decode_node
+
+
+def build_bagel_program() -> Program:
+    return Program(
+        program_id="bagel.mot", kind=ProgramKind.INLINE,
+        nodes=[
+            ComponentNode("tokenize", fn=tokenize_node, writes=("prompt_tokens",)),
+            ModelLoopNode("generate_text", loop_id="generate_text", output_slot="gen_text_out",
+                          reads=("prompt_tokens",), writes=("gen_text_out",)),
+            ComponentNode("emit_text", fn=emit_text_node("gen_text_out", "text"),
+                          reads=("gen_text_out",), writes=("text",)),
+            ComponentNode("pack", fn=pack_cond_from_tokens("gen_text_out"),
+                          reads=("gen_text_out",), writes=("text_embeds", "neg_text_embeds")),
+            ModelLoopNode("generate_image", loop_id="generate_image", output_slot="gen_image_out",
+                          reads=("text_embeds",), writes=("gen_image_out",)),
+            ComponentNode("vae_decode", fn=vae_decode_node("gen_image_out", "image"),
+                          reads=("gen_image_out",), writes=("image",)),
+        ],
+        output_artifacts={"text": "text", "image": "image"},
+    ).validate()

mini_fastvideo/models/cosmos3/__init__.py

@@ -0,0 +1,7 @@
+"""Cosmos3 — omni MoT (reasoner + joint denoise on one resident instance). Phase 2."""
+from __future__ import annotations
+
+from .card import build_cosmos3_card
+from .program import build_cosmos3_program
+
+__all__ = ["build_cosmos3_card", "build_cosmos3_program"]

mini_fastvideo/models/cosmos3/card.py

@@ -0,0 +1,89 @@
+"""Cosmos3 omni ModelCard (design_v3 §4.1, §15b) — the forcing function for omni-native serving.
+
+ONE resident MoT ``transformer`` is bound to BOTH the ``ar_decode`` loop (the reasoner, und pathway)
+and the ``diffusion_denoise`` loop (the joint generation pathway) — ``shared_weight_components`` on
+both. A request runs the reasoner (prompt upsampling), packs its tokens into conditioning, then the
+joint denoise — all on the same resident weights, with every AR token and denoise step a
+runtime-visible WorkUnit (the differentiation vs vllm-omni's opaque DIFFUSION stage). The diffusion
+loop is literally ``WanDenoiseLoop`` — the same step body the engine serves for Wan — bound to the
+MoT module. ``sound_vae`` is declared ``optional_for`` non-t2vs tasks (the lazy-component fix, P8).
+"""
+from __future__ import annotations
+
+from ..._enums import Capability, ConsistencyLevel, LoopKind, WorkUnitKind
+from ...card import (
+    CacheContract,
+    CapabilityMatrix,
+    ComponentSpec,
+    CostModel,
+    LoopSpec,
+    ModelCard,
+    ParallelismContract,
+    ParitySpec,
+    PrecisionContract,
+    RecipeSpec,
+)
+from ...loop.policies import ClassicCFG, FlowShiftPolicy, NoRouting, PrecisionPolicy
+from ...parallel import ParallelPlan
+from ..backend import ToyMoTDiT, ToyTokenizer, ToyVAE, _seed_from
+from ..omni import ARDecodeLoop
+from ..wan21.loop import WanDenoiseLoop
+
+
+def build_cosmos3_card(model_id: str = "cosmos3-vfm") -> ModelCard:
+    seed = _seed_from(model_id)
+    ar_cost = CostModel(kind=WorkUnitKind.AR_TOKEN, base_seconds=5e-5, per_unit_seconds=1e-7)
+    dn_cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7)
+    cfg, flow = ClassicCFG(), FlowShiftPolicy(shift=5.0)
+    precision, expert = PrecisionPolicy(), NoRouting("transformer")
+
+    def ar_factory():
+        return ARDecodeLoop(loop_id="ar_decode", transformer_id="transformer", cost=ar_cost, max_tokens=6)
+
+    def dn_factory():
+        return WanDenoiseLoop(loop_id="diffusion_denoise", cfg=cfg, flow_shift=flow,
+                              precision=precision, expert=expert, cost=dn_cost)
+
+    components = {
+        "tokenizer": ComponentSpec("tokenizer", kind="tokenizer",
+                                   factory=lambda inst: ToyTokenizer(),
+                                   required_for={"reason", "t2v", "t2vs"}),
+        "transformer": ComponentSpec(
+            "transformer", kind="dit",
+            load_id="fastvideo.models.dits.cosmos3:Cosmos3VFMTransformer",
+            factory=lambda inst: ToyMoTDiT(seed=seed),
+            resident_for=["ar_decode", "diffusion_denoise"],     # never unload mid-request
+            required_for={"reason", "t2v", "t2vs"}),
+        "vae": ComponentSpec("vae", kind="vae", factory=lambda inst: ToyVAE(),
+                             required_for={"t2v", "t2vs"}),
+        "sound_vae": ComponentSpec("sound_vae", kind="audio_vae",
+                                   load_id="fastvideo.models.audio:Cosmos3SoundVAE",
+                                   optional_for={"reason", "t2v"}, required_for={"t2vs"}),
+    }
+    loops = {
+        "ar_decode": LoopSpec("ar_decode", kind=LoopKind.AR_DECODE, work_unit_kind=WorkUnitKind.AR_TOKEN,
+                              step_cost_model=ar_cost, shared_weight_components=["transformer"],
+                              cache_policy=["paged_kv"], loop_factory=ar_factory),
+        "diffusion_denoise": LoopSpec("diffusion_denoise", kind=LoopKind.DIFFUSION_DENOISE,
+                                      work_unit_kind=WorkUnitKind.DIFFUSION_STEP, step_cost_model=dn_cost,
+                                      shared_weight_components=["transformer"], cache_policy=["feature"],
+                                      loop_factory=dn_factory),
+    }
+    card = ModelCard(
+        model_id=model_id, family="cosmos3", components=components, loops=loops,
+        capabilities=CapabilityMatrix.of(
+            Capability.TEXT_TO_VIDEO, Capability.REASONING_TEXT, Capability.TEXT_TO_VIDEO_SOUND,
+            Capability.VAE_DECODE, Capability.POLICY_ROLLOUT),
+        recipe=RecipeSpec(method="base", assumes_loop="diffusion_denoise",
+                          assumes_precision="float32", consistency_required=ConsistencyLevel.C1),
+        parity=ParitySpec(consistency_levels=[ConsistencyLevel.C1], interleave_required=True),
+        caches={
+            "feature": CacheContract("feature", max_bytes=1 << 24, reuse_across_requests=True),
+            "paged_kv": CacheContract("paged_kv", max_bytes=1 << 24, block_bytes=1 << 12,
+                                      reuse_across_requests=False),
+        },
+        precision=PrecisionContract(default_dtype="float32", training_precision="float32"),
+        parallelism=ParallelismContract(valid_plans=[ParallelPlan.single()],
+                                        default_plan=ParallelPlan.single()),
+    )
+    return card.validate()

mini_fastvideo/models/cosmos3/program.py

@@ -0,0 +1,32 @@
+"""Cosmos3 omni program (design_v3 §15b):
+
+    tokenize → reason (ar_decode) → pack(tokens→cond) → diffusion_denoise → vae_decode
+
+The reasoner and the denoiser bind the SAME resident ``transformer``; the scheduler is the mode
+multiplexer. This is the workload no DAG-of-engines can express (splitting them doubles the weights
+and severs the shared state).
+"""
+from __future__ import annotations
+
+from ...program import ComponentNode, ModelLoopNode, Program, ProgramKind
+from ..omni import emit_text_node, pack_cond_from_tokens, tokenize_node, vae_decode_node
+
+
+def build_cosmos3_program() -> Program:
+    return Program(
+        program_id="cosmos3.omni", kind=ProgramKind.INLINE,
+        nodes=[
+            ComponentNode("tokenize", fn=tokenize_node, writes=("prompt_tokens",)),
+            ModelLoopNode("reason", loop_id="ar_decode", output_slot="reason_out",
+                          reads=("prompt_tokens",), writes=("reason_out",)),
+            ComponentNode("emit_text", fn=emit_text_node("reason_out", "reason_text"),
+                          reads=("reason_out",), writes=("reason_text",)),
+            ComponentNode("pack", fn=pack_cond_from_tokens("reason_out"),
+                          reads=("reason_out",), writes=("text_embeds", "neg_text_embeds")),
+            ModelLoopNode("denoise", loop_id="diffusion_denoise", output_slot="denoise_out",
+                          reads=("text_embeds",), writes=("denoise_out",)),
+            ComponentNode("vae_decode", fn=vae_decode_node("denoise_out", "video"),
+                          reads=("denoise_out",), writes=("video",)),
+        ],
+        output_artifacts={"video": "video", "text": "reason_text", "latents": "denoise_out"},
+    ).validate()

mini_fastvideo/models/omni/__init__.py

@@ -0,0 +1,7 @@
+"""Shared omni building blocks (design_v3 §4): the AR decode loop + pack/emit program nodes."""
+from __future__ import annotations
+
+from .ar_loop import ARDecodeLoop
+from .pack import emit_text_node, pack_cond_from_tokens, tokenize_node, vae_decode_node
+
+__all__ = ["ARDecodeLoop", "tokenize_node", "pack_cond_from_tokens", "emit_text_node", "vae_decode_node"]

mini_fastvideo/models/omni/pack.py

@@ -0,0 +1,45 @@
+"""Omni program nodes shared by Cosmos3 and the bagel/lance MoT (design_v3 §4.1, §15b).
+
+The ``pack`` node is the seam where the AR (understanding) pathway's tokens become the
+conditioning the diffusion (generation) pathway consumes — "the AR reasoner ... used as prompt
+upsampling before diffusion in the same request". Both pathways are the SAME resident MoT
+module, so this is weight-sharing, not a hand-off between two engines.
+"""
+from __future__ import annotations
+
+from ..backend import ToyTokenizer
+
+
+def tokenize_node(instance, slots, request, ctx) -> None:
+    tok = instance.component("tokenizer") if instance.has_component("tokenizer") else ToyTokenizer()
+    slots["prompt_tokens"] = tok.encode(request.prompt())
+
+
+def _tokens_of(slot_value) -> list:
+    if isinstance(slot_value, dict):
+        return slot_value.get("tokens", [])
+    return list(slot_value or [])
+
+
+def pack_cond_from_tokens(token_slot: str):
+    """Build a ComponentNode fn that packs the und-pathway tokens into the gen-pathway conditioning."""
+    def fn(instance, slots, request, ctx) -> None:
+        mot = instance.component("transformer")          # the SHARED MoT module (und + gen pathways)
+        toks = _tokens_of(slots.get(token_slot))
+        slots["text_embeds"] = mot.reasoner_embed(toks)
+        slots["neg_text_embeds"] = mot.reasoner_embed([0])
+    return fn
+
+
+def emit_text_node(token_slot: str, out_slot: str = "text"):
+    def fn(instance, slots, request, ctx) -> None:
+        out = slots.get(token_slot, {})
+        slots[out_slot] = out.get("text", "") if isinstance(out, dict) else str(out)
+    return fn
+
+
+def vae_decode_node(latent_slot: str, out_slot: str):
+    def fn(instance, slots, request, ctx) -> None:
+        latents = slots[latent_slot]["latents"] if isinstance(slots.get(latent_slot), dict) else None
+        slots[out_slot] = instance.component("vae").decode(latents) if latents is not None else None
+    return fn

mini_fastvideo/tests/test_omni.py

@@ -0,0 +1,83 @@
+"""Phase-2 omni: one resident MoT instance runs AR + diffusion loops on shared weights.
+
+This is the §16 claim — true omni/MoT serving: one resident model, many loop types, scheduled at
+step granularity, in one request — made native (not a DAG doubling weights, not a monolith bypassing
+the abstraction). Cosmos3 = reason → joint denoise; BAGEL/lance = generate_text → generate_image.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from mini_fastvideo._enums import WorkUnitKind
+from mini_fastvideo.models import build_omni_engine
+from mini_fastvideo.models.cosmos3 import build_cosmos3_card
+from mini_fastvideo.parity import assert_interleave_parity
+from mini_fastvideo.request import DiffusionParams, SamplingParams, TaskType, make_request
+
+
+def _eng():
+    return build_omni_engine()
+
+
+def _omni_req(model_id, task, prompt, seed, max_tokens=6, steps=4):
+    return make_request(task, model_id, prompt,
+                        sampling=SamplingParams(max_tokens=max_tokens, seed=seed),
+                        diffusion=DiffusionParams(num_steps=steps, seed=seed))
+
+
+def test_cosmos3_reasoner_then_joint_denoise_one_request():
+    eng = _eng()
+    out = eng.run(_omni_req("cosmos3-vfm", TaskType.T2V, "a phoenix", 1))
+    assert out.artifacts["video"].frames is not None       # diffusion pathway produced video
+    assert out.artifacts["text"].text.startswith("tok:")   # und pathway produced reasoned tokens
+    assert out.metrics["tokens"] > 0 and out.metrics["denoise_steps"] > 0
+
+
+def test_bagel_generate_text_then_image_one_request():
+    eng = _eng()
+    out = eng.run(_omni_req("bagel-mot", TaskType.T2I, "a teapot", 2))
+    assert out.artifacts["text"].text.startswith("tok:")   # AR generate_text
+    assert out.artifacts["image"].tensor is not None       # diffusion generate_image
+    assert out.metrics["tokens"] > 0 and out.metrics["denoise_steps"] > 0
+
+
+def test_mot_one_resident_instance_shared_by_both_loops():
+    eng = _eng()
+    inst = eng._registry["cosmos3-vfm"][0]
+    t = inst.component("transformer")
+    assert inst.component("transformer") is t              # exactly one resident copy
+    # both loops declare they bind the same component (no weight duplication, no DAG split)
+    assert "transformer" in inst.card.loops["ar_decode"].shared_weight_components
+    assert "transformer" in inst.card.loops["diffusion_denoise"].shared_weight_components
+    # and the one module exposes BOTH pathways
+    assert hasattr(t, "ar_forward") and callable(t)
+
+
+def test_omni_scheduler_sees_both_work_unit_kinds():
+    eng = _eng()
+    eng.run(_omni_req("cosmos3-vfm", TaskType.T2V, "a comet", 3))
+    by_kind = eng.admission.metrics.by_kind
+    # the scheduler priced BOTH AR tokens and denoise steps — runtime-visible, not one opaque stage
+    assert by_kind[WorkUnitKind.AR_TOKEN.value] > 0
+    assert by_kind[WorkUnitKind.DIFFUSION_STEP.value] > 0
+
+
+def test_omni_interleave_parity_holds_across_loop_types():
+    eng = _eng()
+    reqs = [_omni_req("cosmos3-vfm", TaskType.T2V, "alpha", 11),
+            _omni_req("cosmos3-vfm", TaskType.T2V, "beta", 22)]
+    assert not assert_interleave_parity(eng, reqs)         # ar+denoise interleave is still bit-identical
+
+
+def test_cosmos3_sound_vae_is_optional_declared_not_loaded():
+    # the lazy-component fix (P8): sound_vae is declared optional_for non-t2vs, never an env-var hack
+    card = build_cosmos3_card()
+    assert "t2v" in card.components["sound_vae"].optional_for
+    assert "t2vs" in card.components["sound_vae"].required_for
+
+
+def test_omni_determinism():
+    eng = _eng()
+    a = eng.run(_omni_req("bagel-mot", TaskType.T2I, "x", 5)).artifacts["image"].tensor
+    b = eng.run(_omni_req("bagel-mot", TaskType.T2I, "x", 5)).artifacts["image"].tensor
+    assert np.array_equal(a, b)