Add composable world-model pipelines and Cosmos3 export

CHANGELOG.md

@@ -7,35 +7,113 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Compositional world-model pipelines
+
+#### Added
+
+- A runtime-agnostic `PipelineBuilder` / `PipelineManifest` /
+  `PipelinePackage` framework for heterogeneous world models. It validates
+  typed graph ports, input-source closure, dataflow, recurrent state,
+  registered roles/strategies/transforms, runtime capabilities, safe assets,
+  and atomic `pipeline.json` persistence.
+- Executable pipeline schema 1.1: versioned model profiles, input semantics,
+  registered generated-input programs, explicit state lifecycle,
+  scheduler/sampling/stopping controls, transform parameters, and per-component
+  dtype/EP hints.
+- `build_world_model()` and `--features world-model`, with a complete
+  `cosmos3_omni` implementation: Qwen3-VL Reasoner, unified MoT diffusion
+  Generator, Wan video VAE, optional full/decoder-only Cosmos3 AVAE, Sound
+  projection, and domain-aware Action projection.
+- Complete `cosmos3_edge` world-model composition for
+  `nvidia/Cosmos3-Edge` and `Cosmos3-Edge-Policy-DROID`: Edge
+  Nemotron/SigLIP Reasoner, shared MoT Generator, Wan VAE, and Action head.
+- Exact Mobius implementations for `Cosmos3OmniTransformer`,
+  `AutoencoderKLWan`, and `Cosmos3AVAEAudioTokenizer`.
+- Single-frame (image) mode for the Wan video VAE: the exported encoder accepts
+  `frames = 1` and the decoder accepts `latent_frames = 1`, matching upstream
+  diffusers' chunk-0 behaviour, which enables native Cosmos3 text-to-image.
+
+#### Changed
+
+- The original fixed four-output world-model API is now accurately named
+  `LatentDynamicsTask` / `LatentDynamicsConfig` /
+  `MLPLatentDynamicsModel`. The original `WorldModel*` names remain aliases.
+
 ### NVIDIA Cosmos 3 Edge vision-language model (`cosmos3_edge`)
 
 #### Added
 
 - Support for the **full `cosmos3_edge` vision-language model**
   (`nvidia/Cosmos3-Edge`, `Cosmos3EdgeForConditionalGeneration`) as a 3-model
-  onnxruntime-genai split (`decoder` + `vision_encoder` + `embedding`):
+  split (`decoder` + `vision_encoder` + `embedding`):
   - **decoder**: grouped-query-attention text reasoner with a **non-gated
     squared-ReLU FFN** (`hidden_act="relu2"`, `up_proj → relu2 → down_proj`)
-    and 3D multimodal RoPE (`mrope_section=[24, 20, 20]`); takes
-    `inputs_embeds`.
-  - **vision_encoder**: SigLIP vision tower + a new
-    `Cosmos3EdgeMultiModalProjector` (pre-shuffle `LayerNorm` → 2×2
-    pixel-shuffle → `linear_fc1` → GELU → `linear_fc2`).
-  - **embedding**: token embedding + image-feature fusion at
-    `image_token_id=19`.
+    and **interleaved** 3D multimodal RoPE (`mrope_section=[24, 20, 20]`,
+    axis per frequency channel `i % 3`); takes `inputs_embeds` and
+    `position_ids [3, batch, seq]`.
+  - **vision_encoder**: **variable-resolution** SigLIP2 tower
+    (`Cosmos3EdgeVisionTower`) + `Cosmos3EdgePatchMerger` (pre-shuffle
+    `LayerNorm` → 2×2 pixel-shuffle → `linear_fc1` → GELU → `linear_fc2`).
+    Consumes the processor's packed, block-major, channel-last patches
+    (`pixel_values [total_patches, patch*patch*3]`) plus `grid_thw [3]`, and
+    resamples the learned 16×16 position grid to the image's own patch grid
+    with an exact antialiased-bilinear filter. **The same graph serves images
+    (`grid_t = 1`) and videos (`grid_t = num_frames`)**, matching
+    `Cosmos3EdgeModel.get_video_features`, which delegates to
+    `get_image_features`.
+  - **embedding**: token embedding + **two** feature streams —
+    `image_features` scattered at `image_token_id=19` and `video_features` at
+    `video_token_id=18` — mirroring the reference's two `masked_scatter`
+    calls. Either stream may be empty.
   `preprocess_weights` routes the single HF checkpoint to the three
-  sub-models: `model.visual.*` / `model.projector.*` → vision (with SigLIP
-  `mlp.fc1/fc2` → `up_proj/down_proj`), `embed_tokens` → embedding, the
-  top-level text tower (`layers.*` / `norm` / `lm_head`) → decoder (renaming
+  sub-models: `model.visual.*` / `model.projector.*` → vision (only the
+  `model.` prefix and the SigLIP `mlp.fc1/fc2` → `up_proj/down_proj` naming
+  differ), `embed_tokens` → embedding, the top-level text tower
+  (`layers.*` / `norm` / `lm_head`) → decoder (renaming
   `self_attn.to_{q,k,v,out}` → `{q,k,v,o}_proj`), and drops the
   generator-tower `k_norm_und_for_gen` key-norm. Built via a new
   `Cosmos3EdgeVLTask` (`cosmos3-edge-vl`). The decoder-only text reasoner
   remains available as `cosmos3_edge_text`.
-- **L1 graph-build tested only.** NVIDIA does not publish modeling code for
-  `cosmos3_edge` (not in `transformers`, no remote-code module), so the exact
-  pixel-shuffle ordering and numerical parity are unverifiable; L4/L5 parity
-  is deferred. The `cosmos3_omni` variants (`Cosmos3-Nano`/`-Super`) are
-  two-tower diffusion world models tracked separately.
+- Video understanding for `cosmos3_edge`: per-frame vision spans with
+  timestamps, the `video_token_id=18` placeholder stream, and a
+  `vision_understanding` world-model manifest block describing the token ids,
+  per-frame token expansion, feature routing, the full packed-patch
+  preprocessing contract (bicubic `smart_resize` to a multiple of 32 inside the
+  processor's pixel-area bounds, `1/255` rescale, RGB conversion, mean/std
+  `0.5` — the first three are processor class defaults absent from the shipped
+  `preprocessor_config.json`), and the interleaved M-RoPE axis assignment
+  including the per-frame `grid_t = 1` video index rule.
+- **Numerically verified against the published reference.** The Reasoner
+  (vision tower, merger projector, image/video token fusion and decoder
+  logits) is compared against `tests/_cosmos3_edge_reference.py`, a PyTorch
+  transcription of `transformers`' `modular_cosmos3_edge.py`
+  (`models/cosmos3_edge`, commit `e8ea728`; cross-checked with vLLM's
+  `cosmos3_edge.py`), at tiny scale
+  (`tests/cosmos3_edge_vision_test.py`) and with the real checkpoint
+  (`tests/cosmos3_edge_integration_test.py`). The Cosmos3-Edge
+  Generator/Action/Sound towers that share the same checkpoint remain
+  proprietary rectified-flow components with no published reference, so their
+  numerics stay unverifiable.
+- The complete Cosmos3-Edge world-model pipeline is also available through
+  `build_world_model()` / `--features world-model`.
+
+#### Fixed
+
+- `cosmos3_edge` image understanding produced uncorrelated vision features
+  (Pearson r ≈ 0.001 against the reference on a 256×256 image). Three
+  independent root causes:
+  - the patch embedding was reshaped into a `Conv2d` `[out, C, kH, kW]`
+    kernel, but the checkpoint ships an `nn.Linear` over **channel-last**
+    `(patch_h, patch_w, channel)` values;
+  - the vision graph assumed a fixed 256×256 square input in raster patch
+    order instead of the processor's variable-resolution, **block-major**
+    packed patches with resampled position embeddings;
+  - the merger projector concatenated the merged block as
+    `(hidden, merge, merge)` instead of `(merge_h, merge_w, hidden)`.
+- `cosmos3_edge` used Qwen-style **chunked** M-RoPE rather than Cosmos'
+  **interleaved** M-RoPE. The two agree exactly on text tokens (`|Δcos| = 0`)
+  but differ by up to 1.95 on visual tokens, which is why text-only output was
+  correct while image understanding was not.
 
 ### Cargo-style `--features` build option
 
@@ -45,8 +123,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   single Rust/cargo-style option. Accepts a comma-separated list and may be
   repeated (`--features fp8-kv-cache,static-cache` or `--features fp8-kv-cache
   --features static-cache`). Available features: `static-cache`, `fp8-kv-cache`,
-  `prune-prefill-prefix`, `text-only`. Unknown feature names are rejected with an error
-  listing the valid set.
+  `prune-prefill-prefix`, `text-only`, `world-model`. Unknown feature names are
+  rejected with an error listing the valid set.
 
 #### Changed
 

README.md

@@ -117,12 +117,19 @@ mobius build --model openai/whisper-tiny output_dir/
 ```
 
 Build-mode toggles use the cargo-style `--features` option. Available features
-are `static-cache`, `fp8-kv-cache`, `prune-prefill-prefix`, and `text-only`. Pass them
-as a comma-separated list or repeat the option:
+are `static-cache`, `fp8-kv-cache`, `prune-prefill-prefix`, `text-only`, and
+`world-model`. Pass them as a comma-separated list or repeat the option:
 
 ```sh
 mobius build --model meta-llama/Llama-3.2-1B output_dir/ \
-      --features static-cache,prune-prefill-prefix --max-seq-len 2048
+    --features static-cache,prune-prefill-prefix --max-seq-len 2048
+
+# Export a complete heterogeneous world-model pipeline
+mobius build --model nvidia/Cosmos3-Nano cosmos3_onnx/ \
+    --features world-model
+
+mobius build --model nvidia/Cosmos3-Edge cosmos3_edge_onnx/ \
+    --features world-model
 ```
 
 See the [CLI Reference](https://onnxruntime.github.io/mobius/cli_reference.html) for all subcommands and flags.

docs/api/build_world_model.md

@@ -0,0 +1,34 @@
+# `build_world_model()`
+
+Export a supported world-model checkpoint as a `PipelinePackage`.
+
+```python
+from mobius import build_world_model
+
+package = build_world_model(
+    "nvidia/Cosmos3-Nano",
+    load_weights=True,
+    execution_provider="cuda",
+)
+```
+
+Save the returned package with `package.save(output_directory)`.
+
+## Supported models
+
+| `model_type` | Checkpoints |
+|---|---|
+| `cosmos3_omni` | Qwen3-VL-based Cosmos3-Nano, Cosmos3-Super, Policy-DROID, Text2Image, and Image2Video variants whose public component configs match the supported architecture |
+| `cosmos3_edge` | `nvidia/Cosmos3-Edge` |
+
+`nvidia/Cosmos3-Edge-Policy-DROID` is detected automatically despite its
+different top-level model type.
+
+## Options
+
+`load_weights=False` builds and validates the complete graph topology without
+downloading tensor payloads. Small configuration files, runtime assets, and
+safetensors header metadata may still be downloaded.
+
+Use `dtype="f32"` for CPU inference or the checkpoint's native BF16 dtype for
+CUDA inference.

docs/api/index.md

@@ -7,6 +7,7 @@ Public API documentation for `mobius`.
 - [`build()`](build.md) — Build from a HuggingFace model ID
 - [`build_from_module()`](build_from_module.md) — Build from a module instance
 - [`build_from_gguf()`](build_from_gguf.md) — Build from a GGUF file
+- [`build_world_model()`](build_world_model.md) — Build a complete compositional world model
 - [`apply_weights()`](apply_weights.md) — Apply weights to a built model
 
 ## Core Classes
@@ -22,6 +23,7 @@ Public API documentation for `mobius`.
 build
 build_from_module
 build_from_gguf
+build_world_model
 apply_weights
 model_package
 base_model_config

docs/cli_reference.md

@@ -177,6 +177,7 @@ option. Pass a comma-separated list (and/or repeat the flag):
 --features fp8-kv-cache,static-cache
 --features prune-prefill-prefix
 --features text-only
+--features world-model
 ```
 
 Available features:
@@ -187,6 +188,7 @@ Available features:
 | `fp8-kv-cache` | Store the `GroupQueryAttention` KV cache as `FLOAT8E4M3FN` (per-tensor E4M3), halving KV-cache memory. Requires a GQA build (e.g. `--ep cuda --dtype f16`) and an ORT runtime with the FP8 KV-cache kernel (SM89+). Pair with `--kv-cache-scale-file` for calibrated scales. |
 | `prune-prefill-prefix` | Emit logits shaped `[B, 1, vocab]` by selecting the final token before the LM head. Gemma 4 also prunes its KV-sharing layer suffix and per-layer inputs to reduce prefill compute. |
 | `text-only` | Export the text backbone of a multimodal checkpoint as a standalone decoder-only LLM (see below). |
+| `world-model` | Export a supported world-model package containing ONNX components and `pipeline.json`. |
 
 The legacy boolean flags `--static-cache`, `--fp8-kv-cache`, and
 `--text-only` have been removed in favor of `--features`.
@@ -200,6 +202,12 @@ mobius build --model Qwen/Qwen2.5-0.5B output/ \
 
 mobius build --model meta-llama/Llama-3.2-1B output/ \
     --features prune-prefill-prefix
+
+mobius build --model nvidia/Cosmos3-Nano output/cosmos3/ \
+    --features world-model
+
+mobius build --model nvidia/Cosmos3-Edge output/cosmos3-edge/ \
+    --features world-model
 ```
 
 ### Static Cache (`--features static-cache`)

docs/index.md

@@ -14,6 +14,7 @@ getting-started
 cli_reference
 module-architecture
 model-catalog
+world-models
 models/index
 ```
 

docs/world-models.md

@@ -0,0 +1,92 @@
+# Export world models
+
+Mobius exports a world model as a directory containing multiple ONNX models
+and a `pipeline.json` file. Use a compatible runtime, such as
+`onnx-world-model`, to run the package.
+
+## Quick start
+
+```bash
+# Cosmos3 Edge
+mobius build --model nvidia/Cosmos3-Edge output/cosmos3-edge \
+    --features world-model
+
+# Cosmos3 Omni
+mobius build --model nvidia/Cosmos3-Nano output/cosmos3-nano \
+    --features world-model
+```
+
+Python API:
+
+```python
+from mobius import build_world_model
+
+package = build_world_model("nvidia/Cosmos3-Edge")
+package.save("output/cosmos3-edge")
+```
+
+## Output
+
+The components depend on the checkpoint. A Cosmos3 Edge package contains:
+
+```text
+cosmos3-edge/
+├── pipeline.json
+├── reasoner_decoder/model.onnx
+├── reasoner_embedding/model.onnx
+├── reasoner_vision_encoder/model.onnx
+├── generator/model.onnx
+├── video_encoder/model.onnx
+├── video_decoder/model.onnx
+├── tokenizer.json
+└── scheduler/scheduler_config.json
+```
+
+`pipeline.json` tells the runtime how to execute the models, including
+generated inputs, recurrent state, sampling, scheduling, and component
+dtype/EP requirements.
+
+## Run the package
+
+```python
+from onnx_world_model import Pipeline
+
+pipeline = Pipeline("output/cosmos3-edge")
+session = pipeline.create_session()
+```
+
+See the `onnx-world-model` documentation for text, image/video, and action
+inference examples.
+
+## Supported checkpoints
+
+| Family | Examples |
+|---|---|
+| Cosmos3 Edge | `nvidia/Cosmos3-Edge`, `nvidia/Cosmos3-Edge-Policy-DROID` |
+| Cosmos3 Omni | Cosmos3-Nano, Cosmos3-Super, Text2Image, Image2Video, and Policy-DROID variants |
+
+Some variants omit optional components such as the vision or audio encoder.
+Mobius includes only the components present in the checkpoint.
+
+## Options
+
+- Use `load_weights=False` to build graph structure without downloading model
+  weights.
+- Use `--dtype f32` for CPU inference. The native Cosmos3 transformer dtype is
+  BF16 and normally targets CUDA.
+
+## Cosmos3 Edge image and video input
+
+`reasoner_vision_encoder` accepts packed image or video patches:
+
+| Input | Shape |
+|---|---|
+| `pixel_values` | `[total_patches, patch_dim]` |
+| `grid_thw` | `[3]` (`frames`, `grid_height`, `grid_width`) |
+
+Use the checkpoint's Cosmos3 Edge image/video processor to resize, normalize,
+and patchify media. Route the encoder output to `image_features` for images or
+`video_features` for videos.
+
+The Edge Reasoner vision, fusion, and decoder outputs are numerically verified
+against the published Transformers implementation with the real checkpoint.