From f1dc587c74ab5eb4e8143b9aa9ba078ca771f082 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Tue, 16 Jun 2026 04:40:10 -0700 Subject: [PATCH] =?UTF-8?q?[feat]=20mini-fastvideo=20phase=202:=20omni/MoT?= =?UTF-8?q?=20=E2=80=94=20Cosmos3=20+=20canonical=20vllm-omni=20(BAGEL/lan?= =?UTF-8?q?ce)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 | 33 ++++++--- mini_fastvideo/examples.py | 24 +++++- mini_fastvideo/models/__init__.py | 29 +++++++- mini_fastvideo/models/bagel/__init__.py | 7 ++ mini_fastvideo/models/bagel/card.py | 85 ++++++++++++++++++++++ mini_fastvideo/models/bagel/program.py | 31 ++++++++ mini_fastvideo/models/cosmos3/__init__.py | 7 ++ mini_fastvideo/models/cosmos3/card.py | 89 +++++++++++++++++++++++ mini_fastvideo/models/cosmos3/program.py | 32 ++++++++ mini_fastvideo/models/omni/__init__.py | 7 ++ mini_fastvideo/models/omni/pack.py | 45 ++++++++++++ mini_fastvideo/tests/test_omni.py | 83 +++++++++++++++++++++ 12 files changed, 457 insertions(+), 15 deletions(-) create mode 100644 mini_fastvideo/models/bagel/__init__.py create mode 100644 mini_fastvideo/models/bagel/card.py create mode 100644 mini_fastvideo/models/bagel/program.py create mode 100644 mini_fastvideo/models/cosmos3/__init__.py create mode 100644 mini_fastvideo/models/cosmos3/card.py create mode 100644 mini_fastvideo/models/cosmos3/program.py create mode 100644 mini_fastvideo/models/omni/__init__.py create mode 100644 mini_fastvideo/models/omni/pack.py create mode 100644 mini_fastvideo/tests/test_omni.py diff --git a/mini_fastvideo/README.md b/mini_fastvideo/README.md index 204364589..e977a6321 100644 --- a/mini_fastvideo/README.md +++ b/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) diff --git a/mini_fastvideo/examples.py b/mini_fastvideo/examples.py index 582194039..9a46ddba0 100644 --- a/mini_fastvideo/examples.py +++ b/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") diff --git a/mini_fastvideo/models/__init__.py b/mini_fastvideo/models/__init__.py index ce25894a4..98d731c78 100644 --- a/mini_fastvideo/models/__init__.py +++ b/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 diff --git a/mini_fastvideo/models/bagel/__init__.py b/mini_fastvideo/models/bagel/__init__.py new file mode 100644 index 000000000..95fe87387 --- /dev/null +++ b/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"] diff --git a/mini_fastvideo/models/bagel/card.py b/mini_fastvideo/models/bagel/card.py new file mode 100644 index 000000000..a752ff461 --- /dev/null +++ b/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() diff --git a/mini_fastvideo/models/bagel/program.py b/mini_fastvideo/models/bagel/program.py new file mode 100644 index 000000000..f3ba83cc7 --- /dev/null +++ b/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() diff --git a/mini_fastvideo/models/cosmos3/__init__.py b/mini_fastvideo/models/cosmos3/__init__.py new file mode 100644 index 000000000..ad4695ddb --- /dev/null +++ b/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"] diff --git a/mini_fastvideo/models/cosmos3/card.py b/mini_fastvideo/models/cosmos3/card.py new file mode 100644 index 000000000..cf5316756 --- /dev/null +++ b/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() diff --git a/mini_fastvideo/models/cosmos3/program.py b/mini_fastvideo/models/cosmos3/program.py new file mode 100644 index 000000000..b79d59215 --- /dev/null +++ b/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() diff --git a/mini_fastvideo/models/omni/__init__.py b/mini_fastvideo/models/omni/__init__.py new file mode 100644 index 000000000..352b0652e --- /dev/null +++ b/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"] diff --git a/mini_fastvideo/models/omni/pack.py b/mini_fastvideo/models/omni/pack.py new file mode 100644 index 000000000..a2863052f --- /dev/null +++ b/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 diff --git a/mini_fastvideo/tests/test_omni.py b/mini_fastvideo/tests/test_omni.py new file mode 100644 index 000000000..e2ac91b82 --- /dev/null +++ b/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)