diff --git a/CHANGELOG.md b/CHANGELOG.md index 29b8d268d..88e51a4a7 100644 --- a/CHANGELOG.md +++ b/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-lm-head`, `text-only`. Unknown feature names are rejected with an error - listing the valid set. + `prune-lm-head`, `text-only`, `world-model`. Unknown feature names are + rejected with an error listing the valid set. #### Changed diff --git a/README.md b/README.md index f9f2320b6..7b96607fa 100644 --- a/README.md +++ b/README.md @@ -113,12 +113,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-lm-head`, and `text-only`. Pass them -as a comma-separated list or repeat the option: +are `static-cache`, `fp8-kv-cache`, `prune-lm-head`, `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-lm-head --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. diff --git a/docs/api/build_world_model.md b/docs/api/build_world_model.md new file mode 100644 index 000000000..44ed25559 --- /dev/null +++ b/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. diff --git a/docs/api/index.md b/docs/api/index.md index ff1246334..8d1892a28 100644 --- a/docs/api/index.md +++ b/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 diff --git a/docs/cli_reference.md b/docs/cli_reference.md index 989608091..874e6331f 100644 --- a/docs/cli_reference.md +++ b/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-lm-head --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-lm-head` | Select the final hidden-state position before the LM-head projection and emit logits shaped `[B, 1, vocab]`. Supported by models using the base `CausalLMModel.forward()` path; unsupported custom forwards fail explicitly. Use only when the downstream workflow does not need per-token logits. | | `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-lm-head + +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`) diff --git a/docs/index.md b/docs/index.md index 103db3353..226c158b2 100644 --- a/docs/index.md +++ b/docs/index.md @@ -14,6 +14,7 @@ getting-started cli_reference module-architecture model-catalog +world-models models/index ``` diff --git a/docs/world-models.md b/docs/world-models.md new file mode 100644 index 000000000..b76744a86 --- /dev/null +++ b/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. diff --git a/src/mobius/__init__.py b/src/mobius/__init__.py index 0962f505a..068d3ed39 100644 --- a/src/mobius/__init__.py +++ b/src/mobius/__init__.py @@ -9,6 +9,14 @@ "BaseModelConfig", "CausalLMConfig", "CausalLMTask", + "Cosmos3AVAEAudioDecoderOnlyTokenizer", + "Cosmos3AVAEAudioDecoderTask", + "Cosmos3AVAEAudioTokenizer", + "Cosmos3AVAEAudioTokenizerTask", + "Cosmos3AudioConfig", + "Cosmos3OmniGeneratorConfig", + "Cosmos3OmniGeneratorModel", + "Cosmos3OmniGeneratorTask", "DepthAnythingConfig", "EncoderConfig", "EpCapabilities", @@ -16,30 +24,56 @@ "Gemma3nConfig", "Gemma4AudioConfig", "Gemma4Config", + "GeneratedInputRule", + "LatentDynamicsConfig", "MambaConfig", "MllamaConfig", "ModelPackage", "ModelRegistration", "ModelRegistry", "ModelTask", + "MLPLatentDynamicsModel", "MLPWorldModel", "MMSConfig", "OPSET_VERSION", + "PipelineBuilder", + "PipelineComponent", + "PipelineConnection", + "PipelineInput", + "PipelineManifest", + "PipelineOutput", + "PipelinePackage", + "PipelinePort", + "PipelineProfile", + "PipelineState", + "PipelineStage", + "PipelineValidationError", "Sam2Config", "SegformerConfig", "VisionConfig", "VisionLanguageConfig", "WhisperConfig", + "WorldModelBuilderRegistry", + "WorldModelBuildConfig", "WorldModelConfig", + "WorldModelGenerationConfig", + "WorldModelPipelineConfig", "WorldModelTask", + "WanVAEConfig", + "WanVAETask", + "AutoencoderKLWanModel", + "LatentDynamicsTask", "YolosConfig", "apply_weights", "build", "build_context", + "build_cosmos3_edge_world_model", + "build_cosmos3_world_model", "build_diffusers_pipeline", "build_from_gguf", "build_from_module", "build_from_nemo", + "build_world_model", "components", "ep_capabilities", "ep_registry", @@ -48,8 +82,15 @@ "models", "optimize_model", "register_ep", + "register_generated_input", + "register_phase", + "register_role", + "register_strategy", + "register_state", + "register_transform", "registry", "tasks", + "world_model_registry", ] __version__ = "0.1.0" @@ -65,12 +106,15 @@ AudioConfig, BaseModelConfig, CausalLMConfig, + Cosmos3AudioConfig, + Cosmos3OmniGeneratorConfig, DepthAnythingConfig, EncoderConfig, Gemma2Config, Gemma3nConfig, Gemma4AudioConfig, Gemma4Config, + LatentDynamicsConfig, MambaConfig, MllamaConfig, MMSConfig, @@ -78,22 +122,72 @@ SegformerConfig, VisionConfig, VisionLanguageConfig, + WanVAEConfig, WhisperConfig, WorldModelConfig, YolosConfig, ) from mobius._constants import OPSET_VERSION +from mobius._cosmos3_edge_world_model import build_cosmos3_edge_world_model +from mobius._cosmos3_world_model import build_cosmos3_world_model from mobius._diffusers_builder import build_diffusers_pipeline from mobius._execution_providers import EpCapabilities, ep_registry, get_ep, register_ep from mobius._model_package import ModelPackage from mobius._optimizations import optimize_model +from mobius._pipeline import ( + GeneratedInputRule, + PipelineBuilder, + PipelineComponent, + PipelineConnection, + PipelineInput, + PipelineManifest, + PipelineOutput, + PipelinePackage, + PipelinePort, + PipelineProfile, + PipelineStage, + PipelineState, + PipelineValidationError, + register_generated_input, + register_phase, + register_role, + register_state, + register_strategy, + register_transform, +) from mobius._registry import ( ModelRegistration, ModelRegistry, registry, ) from mobius._weight_loading import apply_weights +from mobius._world_model_builder import ( + WorldModelBuilderRegistry, + build_world_model, + world_model_registry, +) +from mobius._world_model_config import ( + WorldModelBuildConfig, + WorldModelGenerationConfig, + WorldModelPipelineConfig, +) from mobius.integrations.gguf import build_from_gguf from mobius.integrations.nemo import build_from_nemo -from mobius.models import MLPWorldModel -from mobius.tasks import CausalLMTask, ModelTask, WorldModelTask +from mobius.models import ( + AutoencoderKLWanModel, + Cosmos3AVAEAudioDecoderOnlyTokenizer, + Cosmos3AVAEAudioTokenizer, + Cosmos3OmniGeneratorModel, + MLPLatentDynamicsModel, + MLPWorldModel, +) +from mobius.tasks import ( + CausalLMTask, + Cosmos3AVAEAudioDecoderTask, + Cosmos3AVAEAudioTokenizerTask, + Cosmos3OmniGeneratorTask, + LatentDynamicsTask, + ModelTask, + WanVAETask, + WorldModelTask, +) diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py index 821490e8f..8b5d6a430 100644 --- a/src/mobius/__main__.py +++ b/src/mobius/__main__.py @@ -41,6 +41,7 @@ "fp8-kv-cache": "fp8_kv_cache", "prune-lm-head": "prune_lm_head", "text-only": "text_only", + "world-model": "world_model", } @@ -217,6 +218,39 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: "--component selects a component of a diffusers pipeline." ) + if args.world_model: + incompatible = [ + feature + for feature, enabled in ( + ("static-cache", args.static_cache), + ("fp8-kv-cache", args.fp8_kv_cache), + ("prune-lm-head", args.prune_lm_head), + ("text-only", args.text_only), + ) + if enabled + ] + if incompatible: + names = ", ".join(incompatible) + raise SystemExit( + "Error: --features world-model cannot be combined with " + f"decoder-specific feature(s): {names}." + ) + if args.task is not None: + raise SystemExit( + "Error: --features world-model cannot be combined with --task. " + "A world-model builder composes multiple tasks." + ) + if args.component is not None: + raise SystemExit( + "Error: --features world-model cannot be combined with --component. " + "A complete pipeline must be saved atomically with its manifest." + ) + if args.runtime is not None: + raise SystemExit( + "Error: --features world-model emits a runtime-agnostic pipeline " + "manifest and cannot be combined with --runtime." + ) + load_weights = not args.no_weights task: str | ModelTask | None = args.task @@ -252,6 +286,19 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: component_filter = args.component execution_provider = args.execution_provider + if args.world_model: + from mobius._world_model_builder import build_world_model + + source = args.config or args.model + pkg = build_world_model( + source, + dtype=dtype_override, + load_weights=load_weights, + execution_provider=execution_provider, + ) + _save_package(pkg, output_dir, args, optimize, component_filter=None) + return + # Auto-detect diffusers pipelines. Skipped when the text-only feature is set: # that flag only applies to transformers decoder exports, so we let the # central build() validation reject a diffusers/unsupported repo rather @@ -266,6 +313,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask: args.model, dtype=dtype_override, load_weights=load_weights, + execution_provider=execution_provider, ) _save_package(pkg, output_dir, args, optimize, component_filter) return diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py index afb241a16..21a3c6617 100644 --- a/src/mobius/_configs/__init__.py +++ b/src/mobius/_configs/__init__.py @@ -67,6 +67,8 @@ _shallow_fields, _shared_expert_size, ) +from mobius._configs._cosmos3_audio import Cosmos3AudioConfig +from mobius._configs._cosmos3_omni_generator import Cosmos3OmniGeneratorConfig from mobius._configs._quantization import QuantizationConfig from mobius._configs._sub_configs import ( AudioConfig, @@ -79,7 +81,8 @@ TTSConfig, VisionConfig, ) -from mobius._configs._world_model import WorldModelConfig +from mobius._configs._wan_vae import WanVAEConfig +from mobius._configs._world_model import LatentDynamicsConfig, WorldModelConfig __all__ = [ "DEFAULT_INT", @@ -89,6 +92,8 @@ "BaseModelConfig", "CausalLMConfig", "CodePredictorConfig", + "Cosmos3AudioConfig", + "Cosmos3OmniGeneratorConfig", "CodecDecoderConfig", "CodecEncoderConfig", "DepthAnythingConfig", @@ -102,6 +107,7 @@ "JambaConfig", "JetMoeConfig", "LongcatFlashConfig", + "LatentDynamicsConfig", "Mamba2Config", "MambaConfig", "MllamaConfig", @@ -117,6 +123,7 @@ "TTSConfig", "VisionConfig", "VisionLanguageConfig", + "WanVAEConfig", "WhisperConfig", "WorldModelConfig", "YolosConfig", diff --git a/src/mobius/_configs/_cosmos3_audio.py b/src/mobius/_configs/_cosmos3_audio.py new file mode 100644 index 000000000..98d7f65b0 --- /dev/null +++ b/src/mobius/_configs/_cosmos3_audio.py @@ -0,0 +1,490 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Configuration for the Cosmos3 AVAE audio tokenizer (``sound_tokenizer``). + +Mirrors the ``Cosmos3AVAEAudioTokenizer`` config surface from +``diffusers.models.autoencoders.autoencoder_cosmos3_audio`` as published for +``nvidia/Cosmos3-Nano`` / ``nvidia/Cosmos3-Super`` (``sound_tokenizer/config.json``). + +The shipped checkpoint configuration is:: + + model_type = "autoencoder_v2" sampling_rate = 48000 + enc_type = "spec_convnext" dec_type = "oobleck" + bottleneck_type = "vae" activation = "snakebeta" + hop_size = 1920 vocoder_input_dim = 64 + stereo = True dec_out_channels = 2 + +Only that configuration is supported; every other variant raises +:class:`NotImplementedError` from :meth:`Cosmos3AudioConfig.validate`, exactly +like the upstream ``__init__`` guards. + +Encoder presence is **not** a config property +--------------------------------------------- + +``nvidia/Cosmos3-Nano`` and ``nvidia/Cosmos3-Super`` ship full +encoder+decoder AVAE weights (249 tensors: 67 ``encoder.*`` + 182 +``decoder.*``), while ``nvidia/Cosmos3-Super-Text2Image`` ships decoder-only +weights (182 tensors, zero ``encoder.*``). All three ``sound_tokenizer/ +config.json`` files are **byte-identical**: every ``enc_*`` field is present +and none of them carries an ``encoder_enabled`` key. + +Build-time config therefore cannot reveal encoder absence, which is why +:func:`state_dict_has_encoder` / :meth:`Cosmos3AudioConfig.from_diffusers` +(``weight_names=``) / :meth:`Cosmos3AudioConfig.with_encoder_from_state_dict` +exist, and why the model and task layers expose separate decoder-only and +encoder-decoder paths rather than one flag-driven class. + +Two invariants tie the encoder and decoder together and are checked here: + +* ``enc_latent_dim == 2 * vocoder_input_dim`` — the encoder emits + ``[mean, scale]`` moments that the VAE bottleneck splits in half. +* ``prod(enc_strides) * enc_hop_length == prod(dec_strides) == hop_size`` — the + waveform/latent compression factor must agree between the STFT front-end, + the ConvNeXt downsampling stack and the Oobleck upsampling stack. +""" + +from __future__ import annotations + +import dataclasses +import math +from collections.abc import Iterable, Mapping, Sequence + +from mobius._configs._base import BaseModelConfig + +__all__ = ["Cosmos3AudioConfig", "state_dict_has_encoder"] + +_ENCODER_WEIGHT_PREFIX = "encoder." + + +def state_dict_has_encoder(weight_names: Iterable[str]) -> bool: + """Return ``True`` when a checkpoint carries Cosmos3 AVAE encoder weights. + + The upstream ``Cosmos3AVAEAudioTokenizer._fix_state_dict_keys_on_load`` + drops the encoder whenever no ``encoder.*`` key is present, because AVAE + checkpoints are sometimes shipped decoder-only for sound generation. Mobius + needs the same signal *before* graph construction so it never emits + encoder initializers that no weight can fill. + + Args: + weight_names: Iterable of checkpoint tensor names (or any mapping whose + iteration yields them, such as a ``state_dict``). + + Returns: + ``True`` if at least one name starts with ``"encoder."``. + """ + return any(str(name).startswith(_ENCODER_WEIGHT_PREFIX) for name in weight_names) + + +def _as_int_tuple(values: Sequence[int] | int, field: str) -> tuple[int, ...]: + """Coerce a config list/tuple of positive ints into a tuple.""" + if isinstance(values, int): + values = (values,) + result = tuple(int(v) for v in values) + if not result: + raise ValueError(f"{field} must contain at least one entry") + if any(v <= 0 for v in result): + raise ValueError(f"{field} must contain only positive integers, got {result}") + return result + + +@dataclasses.dataclass +class Cosmos3AudioConfig(BaseModelConfig): + """Configuration for ``Cosmos3AVAEAudioTokenizer``. + + Field names and defaults are copied verbatim from the upstream diffusers + ``__init__`` signature so a ``sound_tokenizer/config.json`` can be splatted + in without renaming. + + Attributes: + model_type: AVAE variant; only ``"autoencoder_v2"`` is supported. + sampling_rate: Waveform sample rate in Hz (48000 for Cosmos3). + vocoder_input_dim: Latent channel count consumed by the decoder; equals + the transformer ``sound_dim``. + dec_dim: Base decoder channel count. + dec_c_mults: Decoder channel multipliers (low → high resolution). + dec_strides: Decoder strides; the decoder consumes them reversed. + dec_out_channels: Output waveform channels (2 = stereo). + stereo: Whether audio is stereo; doubles the encoder input channels. + use_wav_as_input: Whether the encoder consumes raw waveforms. + normalize_volume: Whether ``encode`` peak-normalizes before encoding. + hop_size: Waveform→latent compression factor; defaults to + ``prod(dec_strides)`` when ``None``. + input_channels: Per-channel encoder input count before ``stereo`` + doubling. + enc_type: Encoder type; only ``"spec_convnext"`` is supported. + enc_dim: Base encoder channel count. + enc_intermediate_dim: Unused upstream (ConvNeXt blocks use + ``input_dim * 4``); retained for config fidelity. + enc_num_layers: Unused upstream (depth derives from ``enc_num_blocks``); + retained for config fidelity. + enc_num_blocks: ConvNeXt blocks per encoder downsampling stage. + enc_n_fft: STFT size of the encoder spectrogram front-end. + enc_hop_length: STFT hop length of the encoder front-end. + enc_latent_dim: Encoder output channels (``2 * vocoder_input_dim``). + enc_c_mults: Encoder channel multipliers per stage. + enc_strides: Encoder downsampling strides per stage. + enc_identity_init: Zero-init flag for the ConvNeXt residual 1x1 conv. + Training-time only; it does not change the exported graph. + enc_use_snake: Whether ConvNeXt blocks use SnakeBeta (else GELU). + dec_type: Decoder type; only ``"oobleck"`` is supported. + dec_use_snake: Whether the decoder uses SnakeBeta; must be ``True``. + dec_final_tanh: Vestigial decoder flag; must be ``False``. + dec_anti_aliasing: Decoder anti-aliasing flag; must be ``False``. + dec_use_nearest_upsample: Decoder upsample flag; must be ``False``. + dec_use_tanh_at_final: Decoder final-tanh flag; must be ``False``. + bottleneck_type: Bottleneck type; only ``"vae"`` is supported. + bottleneck: Optional bottleneck dict whose ``"type"`` must be ``"vae"``. + activation: Activation family; only ``"snakebeta"`` is supported. + snake_logscale: Whether SnakeBeta parameters are log-scaled; must be + ``True``. + anti_aliasing: Global anti-aliasing flag; must be ``False``. + use_cuda_kernel: Fused-CUDA-kernel flag; must be ``False``. + causal: Whether convolutions are causal; must be ``False``. + padding_mode: Convolution padding mode; only ``"zeros"`` is supported. + latent_mean: Latent normalization mean; must be ``None`` (upstream does + not implement latent normalization). + latent_std: Latent normalization std; must be ``None``. + encoder_enabled: Whether the encoder exists in this checkpoint. Set to + ``False`` for decoder-only AVAE weights so no encoder initializer + is ever created. + """ + + model_type: str = "autoencoder_v2" + sampling_rate: int = 48000 + vocoder_input_dim: int = 64 + dec_dim: int = 320 + dec_c_mults: tuple[int, ...] = (1, 2, 4, 8, 16) + dec_strides: tuple[int, ...] = (2, 4, 5, 6, 8) + dec_out_channels: int = 2 + stereo: bool = True + use_wav_as_input: bool = True + normalize_volume: bool = True + hop_size: int | None = None + input_channels: int = 1 + enc_type: str = "spec_convnext" + enc_dim: int = 192 + enc_intermediate_dim: int = 768 + enc_num_layers: int = 12 + enc_num_blocks: int = 2 + enc_n_fft: int = 64 + enc_hop_length: int = 16 + enc_latent_dim: int = 128 + enc_c_mults: tuple[int, ...] = (1, 2, 4) + enc_strides: tuple[int, ...] = (4, 5, 6) + enc_identity_init: bool = False + enc_use_snake: bool = True + dec_type: str = "oobleck" + dec_use_snake: bool = True + dec_final_tanh: bool = False + dec_anti_aliasing: bool = False + dec_use_nearest_upsample: bool = False + dec_use_tanh_at_final: bool = False + bottleneck_type: str = "vae" + bottleneck: dict | None = None + activation: str = "snakebeta" + snake_logscale: bool = True + anti_aliasing: bool = False + use_cuda_kernel: bool = False + causal: bool = False + padding_mode: str = "zeros" + latent_mean: float | list[float] | None = None + latent_std: float | list[float] | None = None + encoder_enabled: bool = True + + def __post_init__(self) -> None: + """Normalize sequence fields and resolve the default ``hop_size``.""" + self.dec_c_mults = _as_int_tuple(self.dec_c_mults, "dec_c_mults") + self.dec_strides = _as_int_tuple(self.dec_strides, "dec_strides") + self.enc_c_mults = _as_int_tuple(self.enc_c_mults, "enc_c_mults") + self.enc_strides = _as_int_tuple(self.enc_strides, "enc_strides") + if self.hop_size is None: + self.hop_size = math.prod(self.dec_strides) + else: + self.hop_size = int(self.hop_size) + + # -- Derived geometry -------------------------------------------------- + + @property + def encoder_input_channels(self) -> int: + """Waveform channels the encoder expects (``input_channels`` x stereo).""" + return self.input_channels * (2 if self.stereo else 1) + + @property + def stft_num_bins(self) -> int: + """One-sided STFT bin count (``n_fft // 2 + 1``).""" + return self.enc_n_fft // 2 + 1 + + @property + def stft_pad_left(self) -> int: + """Left zero-pad applied before the (``center=False``) STFT.""" + return (self.enc_n_fft - self.enc_hop_length) // 2 + + @property + def stft_pad_right(self) -> int: + """Right zero-pad applied before the (``center=False``) STFT.""" + return (self.enc_n_fft - self.enc_hop_length) - self.stft_pad_left + + @property + def spectrogram_channels(self) -> int: + """Channels of the packed real/imaginary spectrogram fed to ``layers.0``. + + ``(n_fft + 2)`` equals ``2 * stft_num_bins`` (real bins followed by + imaginary bins), multiplied by the number of waveform channels. + """ + return (self.enc_n_fft + 2) * self.encoder_input_channels + + @property + def latent_channels(self) -> int: + """Latent channels after the VAE bottleneck split (``vocoder_input_dim``).""" + return self.vocoder_input_dim + + @property + def moments_channels(self) -> int: + """Channels of the un-split posterior moments (``enc_latent_dim``).""" + return self.enc_latent_dim + + @property + def audio_channels(self) -> int: + """Waveform channels produced by the decoder (``dec_out_channels``).""" + return self.dec_out_channels + + @property + def decoder_upsampling_ratios(self) -> tuple[int, ...]: + """Decoder strides in application order (``reversed(dec_strides)``).""" + return tuple(reversed(self.dec_strides)) + + @property + def decoder_channel_multiples(self) -> tuple[int, ...]: + """``[1] + dec_c_mults`` — the multiplier table the decoder indexes.""" + return (1, *self.dec_c_mults) + + @property + def decoder_upsample_factor(self) -> int: + """Latent frames → waveform samples ratio (``prod(dec_strides)``).""" + return math.prod(self.dec_strides) + + @property + def encoder_downsample_factor(self) -> int: + """Waveform samples → latent frames ratio for the spec-ConvNeXt encoder.""" + return math.prod(self.enc_strides) * self.enc_hop_length + + @property + def resolved_hop_size(self) -> int: + """The effective ``hop_size`` (never ``None`` after ``__post_init__``).""" + assert self.hop_size is not None + return self.hop_size + + # -- Validation -------------------------------------------------------- + + def validate(self) -> None: + """Validate that this config matches the supported Cosmos3 AVAE variant. + + Raises: + NotImplementedError: For any architecture variant the upstream + ``Cosmos3AVAEAudioTokenizer.__init__`` also rejects. + ValueError: For structurally inconsistent dimensions (channel + multiplier / stride length mismatch, latent-dim mismatch, or a + ``hop_size`` that disagrees with the stride products). + """ + if self.model_type != "autoencoder_v2": + raise NotImplementedError( + f"Cosmos3 AVAE model type {self.model_type!r} is not supported." + ) + if not self.use_wav_as_input: + raise NotImplementedError("Cosmos3 AVAE tokenizer only supports waveform input.") + if self.enc_type != "spec_convnext": + raise NotImplementedError( + f"Cosmos3 AVAE encoder type {self.enc_type!r} is not supported." + ) + if self.bottleneck is not None: + declared = self.bottleneck.get("type", self.bottleneck_type) + if declared != "vae": + raise NotImplementedError( + "Cosmos3 AVAE tokenizer only supports the VAE bottleneck, got " + f"bottleneck={{'type': {declared!r}}}." + ) + if self.bottleneck_type != "vae": + raise NotImplementedError( + "Cosmos3 AVAE tokenizer only supports the VAE bottleneck, got " + f"bottleneck_type={self.bottleneck_type!r}." + ) + if self.dec_type != "oobleck": + raise NotImplementedError( + f"Cosmos3 AVAE decoder type {self.dec_type!r} is not supported." + ) + if ( + not self.dec_use_snake + or self.dec_final_tanh + or self.dec_anti_aliasing + or self.dec_use_nearest_upsample + or self.dec_use_tanh_at_final + ): + raise NotImplementedError( + "Cosmos3 AVAE decoder only supports the shipped Oobleck decoder configuration " + "(dec_use_snake=True and every dec_* toggle False)." + ) + if ( + self.activation != "snakebeta" + or not self.snake_logscale + or self.anti_aliasing + or self.use_cuda_kernel + ): + raise NotImplementedError( + "Cosmos3 AVAE tokenizer only supports the shipped SnakeBeta configuration " + "(activation='snakebeta', snake_logscale=True, anti_aliasing=False, " + "use_cuda_kernel=False)." + ) + if self.causal: + raise NotImplementedError( + "Cosmos3 AVAE causal audio encoder is not supported yet." + ) + if self.padding_mode != "zeros": + raise NotImplementedError( + f"Cosmos3 AVAE only supports padding_mode='zeros', got {self.padding_mode!r}." + ) + if self.latent_mean is not None or self.latent_std is not None: + raise NotImplementedError( + "Cosmos3 AVAE tokenizer does not apply latent normalization; " + "`latent_mean`/`latent_std` must be None." + ) + + if len(self.enc_c_mults) != len(self.enc_strides): + raise ValueError( + "`enc_c_mults` and `enc_strides` must have the same length, got " + f"{len(self.enc_c_mults)} and {len(self.enc_strides)}." + ) + if len(self.dec_c_mults) != len(self.dec_strides): + raise ValueError( + "`dec_c_mults` and `dec_strides` must have the same length, got " + f"{len(self.dec_c_mults)} and {len(self.dec_strides)}." + ) + if self.enc_latent_dim != 2 * self.vocoder_input_dim: + raise ValueError( + "Cosmos3 AVAE VAE bottleneck splits the encoder output into mean/scale, so " + f"enc_latent_dim must be 2 * vocoder_input_dim; got enc_latent_dim=" + f"{self.enc_latent_dim} and vocoder_input_dim={self.vocoder_input_dim}." + ) + if self.enc_n_fft <= 0 or self.enc_n_fft % 2 != 0: + raise ValueError( + f"enc_n_fft must be a positive even number, got {self.enc_n_fft}." + ) + if not 0 < self.enc_hop_length <= self.enc_n_fft: + raise ValueError( + "enc_hop_length must satisfy 0 < enc_hop_length <= enc_n_fft; got " + f"enc_hop_length={self.enc_hop_length}, enc_n_fft={self.enc_n_fft}." + ) + for name in ( + "sampling_rate", + "enc_dim", + "dec_dim", + "vocoder_input_dim", + "enc_num_blocks", + ): + if getattr(self, name) <= 0: + raise ValueError(f"{name} must be positive, got {getattr(self, name)}.") + if self.input_channels <= 0: + raise ValueError(f"input_channels must be positive, got {self.input_channels}.") + if self.dec_out_channels != self.encoder_input_channels: + raise ValueError( + "Cosmos3 AVAE round-trips audio, so dec_out_channels must equal " + f"input_channels * (2 if stereo else 1); got dec_out_channels=" + f"{self.dec_out_channels} and encoder_input_channels=" + f"{self.encoder_input_channels}." + ) + if self.decoder_upsample_factor != self.resolved_hop_size: + raise ValueError( + "hop_size must equal prod(dec_strides); got hop_size=" + f"{self.resolved_hop_size} and prod(dec_strides)=" + f"{self.decoder_upsample_factor}." + ) + if self.encoder_downsample_factor != self.resolved_hop_size: + raise ValueError( + "The encoder compression factor prod(enc_strides) * enc_hop_length must equal " + f"hop_size; got {self.encoder_downsample_factor} != {self.resolved_hop_size}." + ) + + # -- Construction ------------------------------------------------------ + + @classmethod + def from_diffusers( + cls, + config: Mapping[str, object] | object, + *, + encoder_enabled: bool | None = None, + weight_names: Iterable[str] | None = None, + ) -> Cosmos3AudioConfig: + """Create a config from a diffusers ``sound_tokenizer/config.json`` dict. + + Unknown/private diffusers bookkeeping keys (``_class_name``, + ``_diffusers_version``, ...) are ignored. + + .. warning:: + The published ``sound_tokenizer/config.json`` files **cannot** tell + you whether an encoder exists. ``nvidia/Cosmos3-Nano``, + ``nvidia/Cosmos3-Super`` (both full) and + ``nvidia/Cosmos3-Super-Text2Image`` (decoder-only) all ship the + *byte-identical* config: every ``enc_*`` field is present and there + is no ``encoder_enabled`` key. Encoder presence is only observable + from the checkpoint, so pass ``weight_names`` (preferred) or + ``encoder_enabled``. Leaving both unset assumes a full checkpoint + and will emit encoder initializers that decoder-only weights cannot + fill. + + Args: + config: Parsed ``sound_tokenizer/config.json`` mapping, or any + object exposing ``to_dict()`` / ``items()``. + encoder_enabled: Explicit override for encoder presence. + weight_names: Checkpoint tensor names (or a ``state_dict``) from + which encoder presence is detected. Mutually exclusive with + ``encoder_enabled``. + + Returns: + A validated :class:`Cosmos3AudioConfig`. + + Raises: + ValueError: If both ``encoder_enabled`` and ``weight_names`` are given. + """ + if encoder_enabled is not None and weight_names is not None: + raise ValueError( + "Pass either `encoder_enabled` or `weight_names`, not both — they are two " + "ways to answer the same question." + ) + if hasattr(config, "to_dict"): + data = dict(config.to_dict()) # type: ignore[attr-defined] + elif isinstance(config, Mapping): + data = dict(config) + else: + data = dict(config.items()) # type: ignore[attr-defined] + + known = {field.name for field in dataclasses.fields(cls)} + kwargs = {key: value for key, value in data.items() if key in known} + if weight_names is not None: + kwargs["encoder_enabled"] = state_dict_has_encoder(weight_names) + elif encoder_enabled is not None: + kwargs["encoder_enabled"] = encoder_enabled + resolved = cls(**kwargs) # type: ignore[arg-type] + resolved.validate() + return resolved + + def with_encoder_from_state_dict( + self, + weight_names: Iterable[str], + ) -> Cosmos3AudioConfig: + """Return a copy whose ``encoder_enabled`` reflects the checkpoint. + + This is the mobius equivalent of the upstream + ``_fix_state_dict_keys_on_load`` hook: decoder-only AVAE checkpoints + must never build an encoder graph, because every encoder initializer + would be left without weight data. + + Args: + weight_names: Checkpoint tensor names (or a ``state_dict``). + + Returns: + ``self`` when the flag already matches, otherwise a new config. + """ + has_encoder = state_dict_has_encoder(weight_names) + if has_encoder == self.encoder_enabled: + return self + return dataclasses.replace(self, encoder_enabled=has_encoder) diff --git a/src/mobius/_configs/_cosmos3_omni_generator.py b/src/mobius/_configs/_cosmos3_omni_generator.py new file mode 100644 index 000000000..d076d9e70 --- /dev/null +++ b/src/mobius/_configs/_cosmos3_omni_generator.py @@ -0,0 +1,488 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Configuration for the NVIDIA Cosmos3-Omni unified MoT transformer. + +Describes ``Cosmos3OmniTransformer`` — the *unified* Mixture-of-Transformers +(MoT) backbone shipped in the ``transformer/`` folder of the diffusers +checkpoints ``nvidia/Cosmos3-Nano`` and ``nvidia/Cosmos3-Super``. This is the +model that carries **both** towers in one stack of layers: + +* the **understanding** ("und" / Reasoner) expert — causal self-attention over + the text/understanding prefix, using ``self_attn.to_{q,k,v,out}``, + ``mlp.*``, ``input_layernorm``, ``post_attention_layernorm`` and ``norm``; +* the **generation** ("gen" / Generator) expert — a rectified-flow diffusion + branch whose tokens attend non-causally over *understanding + generation* + keys/values, using ``self_attn.add_{q,k,v}_proj`` / ``to_add_out``, + ``mlp_moe_gen.*``, ``input_layernorm_moe_gen``, + ``post_attention_layernorm_moe_gen`` and ``norm_moe_gen``. + +On top of the shared backbone the checkpoint carries per-modality projection +heads: vision latents (``proj_in`` / ``proj_out``), an optional Sound head +(``audio_proj_in`` / ``audio_proj_out`` / ``audio_modality_embed``, gated by +``sound_gen``) and an optional Action head (``action_proj_in`` / +``action_proj_out`` / ``action_modality_embed``, gated by ``action_gen``, +implemented as per-embodiment-domain ``DomainAwareLinear`` layers). + +Architecture reference: ``huggingface/diffusers`` +``src/diffusers/models/transformers/transformer_cosmos3.py`` (Apache-2.0, +Copyright 2025 The NVIDIA Team and The HuggingFace Team) and the public +``nvidia/Cosmos3-Nano`` ``transformer/config.json``. +""" + +from __future__ import annotations + +import dataclasses +from typing import Any + +import onnx_ir as ir + +from mobius._configs._base import BaseModelConfig + +# Feed-forward variants that the Cosmos3 backbone is defined for. ``silu`` +# selects the gated (SwiGLU) ``gate_proj``/``up_proj``/``down_proj`` MLP used +# by the Qwen3 dense backbone (Cosmos3-Nano / -Super); ``relu2`` selects the +# non-gated squared-ReLU ``up_proj``/``down_proj`` MLP used by the Nemotron +# backbone. Anything else is an architecture mismatch. +SUPPORTED_HIDDEN_ACTS: tuple[str, ...] = ("silu", "relu2") + +# Compute dtypes the graph may be built in. ``bfloat16`` is the dtype of the +# published weights; ``float32`` is used by the unit tests and ``float16`` +# is accepted for EPs without bf16 kernels. +SUPPORTED_DTYPES: tuple[ir.DataType, ...] = ( + ir.DataType.FLOAT, + ir.DataType.FLOAT16, + ir.DataType.BFLOAT16, +) + +_DTYPE_BY_NAME: dict[str, ir.DataType] = { + "bfloat16": ir.DataType.BFLOAT16, + "float16": ir.DataType.FLOAT16, + "float32": ir.DataType.FLOAT, + "fp16": ir.DataType.FLOAT16, + "fp32": ir.DataType.FLOAT, +} + + +def _as_dict(config: Any) -> dict[str, Any]: + """Normalize a diffusers ``FrozenDict`` / config object into a plain dict.""" + if hasattr(config, "to_dict"): + return dict(config.items()) + return dict(config) + + +def _resolve_diffusers_dtype(value: Any) -> ir.DataType: + """Map the ``dtype`` string in ``transformer/config.json`` to an IR dtype.""" + if value is None: + return ir.DataType.BFLOAT16 + if isinstance(value, ir.DataType): + return value + name = str(value).removeprefix("torch.") + if name not in _DTYPE_BY_NAME: + raise ValueError( + f"Unsupported Cosmos3-Omni dtype {value!r}. Supported: {sorted(_DTYPE_BY_NAME)}" + ) + return _DTYPE_BY_NAME[name] + + +@dataclasses.dataclass +class Cosmos3OmniGeneratorConfig(BaseModelConfig): + """Architecture configuration for ``Cosmos3OmniTransformer``. + + Field names follow the published ``transformer/config.json`` so that + :meth:`from_diffusers` is a near-identity mapping. Fields that only drive + *host-side* preprocessing (position-id construction, patchify/unpatchify, + FPS modulation) are still parsed and validated here so the exported graph + can be paired with a correct pre/post-processor, but they are documented + as such — the ONNX graph itself consumes already-packed tokens and + already-computed mRoPE position IDs. + """ + + # --- Backbone ----------------------------------------------------------- + #: RMSNorm epsilon shared by every norm in the backbone. + rms_norm_eps: float = 1e-6 + #: Rotary base. Cosmos3 uses a very large theta (5e6) for long video. + rope_theta: float = 5_000_000.0 + #: mRoPE channel budget ``(T, H, W)``. These are *channel counts*, and + #: their sum must equal ``head_dim // 2`` (the rotary dimension). The + #: layout is interleaved, not chunked — see + #: :class:`~mobius.models.cosmos3_omni_generator.Cosmos3OmniRotaryEmbedding`. + rope_axes_dim: tuple[int, int, int] = (24, 20, 20) + #: Bias on every attention projection (both experts). ``False`` upstream. + attention_bias: bool = False + #: Attention dropout. Must be ``0.0`` — the exported graph is inference-only. + attention_dropout: float = 0.0 + #: Whether the understanding expert applies per-head QK RMSNorm + #: (``self_attn.norm_q`` / ``self_attn.norm_k``). + qk_norm_for_text: bool = True + #: Whether the generation expert applies per-head QK RMSNorm + #: (``self_attn.norm_added_q`` / ``self_attn.norm_added_k``). The upstream + #: module always constructs these, so ``False`` is an architecture mismatch. + qk_norm_for_diffusion: bool = True + #: Adds a *separate* ``self_attn.k_norm_und_for_gen`` RMSNorm applied to the + #: understanding keys that the generation pathway consumes. Upstream only + #: instantiates it when ``use_und_k_norm_for_gen and not qk_norm_for_text`` + #: — see :attr:`has_und_k_norm_for_gen`. + use_und_k_norm_for_gen: bool = False + + # --- Vision (diffusion) head ------------------------------------------- + #: Channel count of a single video-VAE latent (pre-patchify). + latent_channel: int = 48 + #: Spatial patch size applied to the latent grid (host-side patchify). + latent_patch_size: int = 2 + #: Width of ``proj_in`` / ``proj_out``. Must equal + #: ``latent_channel * latent_patch_size ** 2``. + patch_latent_dim: int = 192 + #: Multiplier applied to the raw timesteps before the sinusoidal + #: projection (``timesteps * timestep_scale``). + timestep_scale: float = 0.001 + #: Width of the sinusoidal timestep projection feeding ``time_embedder``. + #: Upstream hardcodes ``Timesteps(num_channels=256, ...)``; it is fixed by + #: the published ``time_embedder.linear_1.weight`` shape ``[hidden, 256]``. + time_proj_channels: int = 256 + + # --- Sound head (optional) --------------------------------------------- + #: Enables ``audio_proj_in`` / ``audio_proj_out`` / ``audio_modality_embed``. + sound_gen: bool = False + #: Channel count of a sound latent frame. Required when ``sound_gen``. + sound_dim: int | None = None + #: Sound latent frame rate (host-side position-id construction only). + sound_latent_fps: float = 25.0 + #: Sound temporal compression (host-side packing only). + temporal_compression_factor_sound: int = 1 + + # --- Action head (optional) -------------------------------------------- + #: Enables ``action_proj_in`` / ``action_proj_out`` / ``action_modality_embed``. + action_gen: bool = False + #: Action vector width. Required when ``action_gen``. + action_dim: int | None = None + #: Upper bound on ``action_dim`` across embodiments (metadata only). + max_action_dim: int | None = None + #: Number of embodiment domains indexed by the ``DomainAwareLinear`` heads. + num_embodiment_domains: int = 32 + + # --- Host-side (documented, not consumed by the graph) ------------------ + #: Reference FPS used by the host when building temporal position IDs. + base_fps: int = 24 + #: Whether the host scales temporal position IDs by the clip FPS. + enable_fps_modulation: bool = True + #: Whether the host restarts H/W position IDs per modality segment. + unified_3d_mrope_reset_spatial_ids: bool = True + #: Temporal offset the host inserts between modality segments. + unified_3d_mrope_temporal_modality_margin: int = 15_000 + #: Maximum position id the host may emit. + max_position_embeddings: int = 262_144 + + # --- Architecture assertions (parsed so mismatches fail loudly) --------- + #: Must be ``"unified_3d_mrope"``. + position_embedding_type: str = "unified_3d_mrope" + #: Must be ``"two_way"`` (causal und pathway + non-causal gen pathway). + joint_attn_implementation: str = "two_way" + #: Must be ``True`` — the MoT (per-expert weights) structure. + use_moe: bool = True + #: Must be ``False`` — temporal-causal video attention is not implemented. + video_temporal_causal: bool = False + #: HuggingFace-style model type tag used by downstream tooling. + model_type: str = "cosmos3_omni_generator" + + # ------------------------------------------------------------------ + # Derived properties + # ------------------------------------------------------------------ + + @property + def rotary_dim(self) -> int: + """Number of rotary frequency channels (``head_dim // 2``).""" + return self.head_dim // 2 + + @property + def num_key_value_groups(self) -> int: + """Query heads per key/value head (GQA group size).""" + return self.num_attention_heads // self.num_key_value_heads + + @property + def is_gated_mlp(self) -> bool: + """``True`` when the feed-forward is SwiGLU (``gate/up/down_proj``).""" + return self.hidden_act == "silu" + + @property + def has_und_k_norm_for_gen(self) -> bool: + """``True`` when ``self_attn.k_norm_und_for_gen`` exists. + + Mirrors upstream exactly: the extra norm is only instantiated when + ``use_und_k_norm_for_gen`` is set *and* the understanding pathway has + no QK norm of its own. When ``qk_norm_for_text`` is ``True`` the + understanding keys are already normalized, so the flag is inert. + """ + return self.use_und_k_norm_for_gen and not self.qk_norm_for_text + + @property + def attention_out_size(self) -> int: + """Flattened attention output width (``num_attention_heads * head_dim``).""" + return self.num_attention_heads * self.head_dim + + @property + def key_value_size(self) -> int: + """Flattened key/value width (``num_key_value_heads * head_dim``).""" + return self.num_key_value_heads * self.head_dim + + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ + + def validate(self) -> None: + """Validate every shape relationship the graph builder relies on. + + Raises: + ValueError: If any field is missing, non-positive, or inconsistent + with the rest of the architecture. + """ + self._validate_backbone() + self._validate_rope() + self._validate_vision_head() + self._validate_sound_head() + self._validate_action_head() + self._validate_unsupported_variants() + + def _validate_backbone(self) -> None: + for name in ( + "vocab_size", + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + ): + value = getattr(self, name) + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + if self.num_attention_heads % self.num_key_value_heads != 0: + raise ValueError( + "num_attention_heads must be divisible by num_key_value_heads " + f"(got {self.num_attention_heads} and {self.num_key_value_heads})" + ) + if self.head_dim % 2 != 0: + raise ValueError( + f"head_dim must be even for rotary embedding, got {self.head_dim}" + ) + if self.hidden_act not in SUPPORTED_HIDDEN_ACTS: + raise ValueError( + f"hidden_act must be one of {SUPPORTED_HIDDEN_ACTS}, got {self.hidden_act!r}" + ) + if self.rms_norm_eps <= 0: + raise ValueError(f"rms_norm_eps must be positive, got {self.rms_norm_eps}") + if self.dtype not in SUPPORTED_DTYPES: + raise ValueError( + f"dtype must be one of {[d.name for d in SUPPORTED_DTYPES]}, got {self.dtype}" + ) + # Any non-zero dropout is a training-time config, not an inference graph. + if self.attention_dropout: + raise ValueError( + "attention_dropout must be 0.0 for an inference graph, got " + f"{self.attention_dropout}" + ) + + def _validate_rope(self) -> None: + if len(self.rope_axes_dim) != 3: + raise ValueError( + f"rope_axes_dim must have exactly 3 (T, H, W) entries, got {self.rope_axes_dim!r}" + ) + if any( + not isinstance(d, int) or isinstance(d, bool) or d <= 0 for d in self.rope_axes_dim + ): + raise ValueError( + f"rope_axes_dim entries must be positive integers, got {self.rope_axes_dim!r}" + ) + if sum(self.rope_axes_dim) != self.rotary_dim: + raise ValueError( + "sum(rope_axes_dim) must equal head_dim // 2 " + f"(got sum={sum(self.rope_axes_dim)}, head_dim // 2={self.rotary_dim})" + ) + if self.rope_theta <= 0: + raise ValueError(f"rope_theta must be positive, got {self.rope_theta}") + if self.max_position_embeddings <= 0: + raise ValueError( + f"max_position_embeddings must be positive, got {self.max_position_embeddings}" + ) + + def _validate_vision_head(self) -> None: + if self.latent_channel <= 0: + raise ValueError(f"latent_channel must be positive, got {self.latent_channel}") + if self.latent_patch_size <= 0: + raise ValueError( + f"latent_patch_size must be positive, got {self.latent_patch_size}" + ) + expected = self.latent_channel * self.latent_patch_size**2 + if self.patch_latent_dim != expected: + raise ValueError( + "patch_latent_dim must equal latent_channel * latent_patch_size ** 2 " + f"(got {self.patch_latent_dim}, expected {expected})" + ) + if self.timestep_scale <= 0: + raise ValueError(f"timestep_scale must be positive, got {self.timestep_scale}") + if self.time_proj_channels <= 0 or self.time_proj_channels % 2 != 0: + raise ValueError( + "time_proj_channels must be a positive even number, got " + f"{self.time_proj_channels}" + ) + + def _validate_sound_head(self) -> None: + if not self.sound_gen: + return + if self.sound_dim is None or self.sound_dim <= 0: + raise ValueError( + f"sound_dim must be a positive integer when sound_gen=True, got {self.sound_dim!r}" + ) + if self.temporal_compression_factor_sound <= 0: + raise ValueError( + "temporal_compression_factor_sound must be positive, got " + f"{self.temporal_compression_factor_sound}" + ) + if self.sound_latent_fps <= 0: + raise ValueError(f"sound_latent_fps must be positive, got {self.sound_latent_fps}") + + def _validate_action_head(self) -> None: + if not self.action_gen: + return + if self.action_dim is None or self.action_dim <= 0: + raise ValueError( + "action_dim must be a positive integer when action_gen=True, got " + f"{self.action_dim!r}" + ) + if self.num_embodiment_domains <= 0: + raise ValueError( + "num_embodiment_domains must be positive when action_gen=True, got " + f"{self.num_embodiment_domains}" + ) + if self.max_action_dim is not None and self.max_action_dim < self.action_dim: + raise ValueError( + f"max_action_dim ({self.max_action_dim}) must be >= action_dim ({self.action_dim})" + ) + + def _validate_unsupported_variants(self) -> None: + if self.position_embedding_type != "unified_3d_mrope": + raise ValueError( + "Only position_embedding_type='unified_3d_mrope' is supported, got " + f"{self.position_embedding_type!r}" + ) + if self.joint_attn_implementation != "two_way": + raise ValueError( + "Only joint_attn_implementation='two_way' is supported, got " + f"{self.joint_attn_implementation!r}" + ) + if not self.use_moe: + raise ValueError( + "use_moe=False is not a Cosmos3-Omni MoT checkpoint — the generation " + "expert weights (mlp_moe_gen, add_*_proj, ...) would be absent." + ) + if self.video_temporal_causal: + raise ValueError( + "video_temporal_causal=True is not supported: the generation pathway is " + "built as fully non-causal attention over understanding + generation K/V." + ) + if not self.qk_norm_for_diffusion: + raise ValueError( + "qk_norm_for_diffusion=False is not supported: the published checkpoint " + "always carries self_attn.norm_added_q / norm_added_k." + ) + + # ------------------------------------------------------------------ + # Parsing + # ------------------------------------------------------------------ + + @classmethod + def from_diffusers(cls, config: Any) -> Cosmos3OmniGeneratorConfig: + """Build a config from a diffusers ``transformer/config.json`` dict. + + ``rope_axes_dim`` is absent from the published config; upstream derives + it from ``rope_scaling["mrope_section"]`` (falling back to + ``[24, 20, 20]``), which this reproduces exactly. + + Args: + config: Parsed ``transformer/config.json`` (a plain ``dict`` or any + mapping-like diffusers config object). + + Returns: + A validated :class:`Cosmos3OmniGeneratorConfig`. + + Raises: + ValueError: If the parsed architecture fails :meth:`validate`. + """ + raw = _as_dict(config) + rope_scaling = raw.get("rope_scaling") or {} + rope_axes_dim = raw.get("rope_axes_dim") + if rope_axes_dim is None: + rope_axes_dim = rope_scaling.get("mrope_section", [24, 20, 20]) + attention_bias = bool(raw.get("attention_bias", False)) + + parsed = cls( + # BaseModelConfig fields + vocab_size=int(raw.get("vocab_size", 151936)), + hidden_size=int(raw.get("hidden_size", 4096)), + intermediate_size=int(raw.get("intermediate_size", 12288)), + num_hidden_layers=int(raw.get("num_hidden_layers", 36)), + num_attention_heads=int(raw.get("num_attention_heads", 32)), + num_key_value_heads=int(raw.get("num_key_value_heads", 8)), + head_dim=int(raw.get("head_dim", 128)), + hidden_act=raw.get("hidden_act", "silu"), + attn_qkv_bias=attention_bias, + attn_o_bias=attention_bias, + dtype=_resolve_diffusers_dtype(raw.get("dtype")), + # Backbone + rms_norm_eps=float(raw.get("rms_norm_eps", 1e-6)), + rope_theta=float(raw.get("rope_theta", 5_000_000.0)), + rope_axes_dim=tuple(int(d) for d in rope_axes_dim), # type: ignore[arg-type] + attention_bias=attention_bias, + attention_dropout=float(raw.get("attention_dropout", 0.0)), + qk_norm_for_text=bool(raw.get("qk_norm_for_text", True)), + qk_norm_for_diffusion=bool(raw.get("qk_norm_for_diffusion", True)), + use_und_k_norm_for_gen=bool(raw.get("use_und_k_norm_for_gen", False)), + # Vision head + latent_channel=int(raw.get("latent_channel", 48)), + latent_patch_size=int(raw.get("latent_patch_size", 2)), + patch_latent_dim=int(raw.get("patch_latent_dim", 192)), + timestep_scale=float(raw.get("timestep_scale", 0.001)), + # Sound head + sound_gen=bool(raw.get("sound_gen", False)), + sound_dim=_optional_int(raw.get("sound_dim")), + sound_latent_fps=float(raw.get("sound_latent_fps", 25.0)), + temporal_compression_factor_sound=int( + raw.get("temporal_compression_factor_sound", 1) + ), + # Action head + action_gen=bool(raw.get("action_gen", False)), + action_dim=_optional_int(raw.get("action_dim")), + max_action_dim=_optional_int(raw.get("max_action_dim")), + num_embodiment_domains=int(raw.get("num_embodiment_domains", 32)), + # Host-side + base_fps=int(raw.get("base_fps", 24)), + enable_fps_modulation=bool(raw.get("enable_fps_modulation", True)), + unified_3d_mrope_reset_spatial_ids=bool( + raw.get("unified_3d_mrope_reset_spatial_ids", True) + ), + unified_3d_mrope_temporal_modality_margin=int( + raw.get("unified_3d_mrope_temporal_modality_margin", 15_000) + ), + max_position_embeddings=int(raw.get("max_position_embeddings", 262_144)), + # Architecture assertions + position_embedding_type=str( + raw.get("position_embedding_type", "unified_3d_mrope") + ), + joint_attn_implementation=str(raw.get("joint_attn_implementation", "two_way")), + use_moe=bool(raw.get("use_moe", True)), + video_temporal_causal=bool(raw.get("video_temporal_causal", False)), + ) + parsed.validate() + return parsed + + +def _optional_int(value: Any) -> int | None: + """Return ``int(value)`` or ``None`` when the field is absent/null.""" + return None if value is None else int(value) + + +__all__ = [ + "SUPPORTED_DTYPES", + "SUPPORTED_HIDDEN_ACTS", + "Cosmos3OmniGeneratorConfig", +] diff --git a/src/mobius/_configs/_extractors.py b/src/mobius/_configs/_extractors.py index ac4c6e41c..bcbfd4e74 100644 --- a/src/mobius/_configs/_extractors.py +++ b/src/mobius/_configs/_extractors.py @@ -188,6 +188,7 @@ def extract_vision_config(config, parent_config, model_type: str) -> dict: "fullatt_block_indexes", "window_size", "mrope_section", + "mrope_interleaved", "image_crop_size", ): val = fields.get(shared) diff --git a/src/mobius/_configs/_sub_configs.py b/src/mobius/_configs/_sub_configs.py index 79dfddc5f..97edf108c 100644 --- a/src/mobius/_configs/_sub_configs.py +++ b/src/mobius/_configs/_sub_configs.py @@ -78,6 +78,22 @@ class VisionConfig: # (HF: projector_config.merger_intermediate_size). ``None`` means the model # does not use a Cosmos-style merger projector. projector_intermediate_size: int | None = None + # Cosmos3-Edge: apply the merger LayerNorm *after* the spatial shuffle + # (HF: projector_config.use_postshuffle_norm). + use_postshuffle_norm: bool = False + # SigLIP2-style learned position-embedding reference grid size, in patches + # (HF: vision_config.num_patches, e.g. 256 -> 16x16). Variable-resolution + # towers resample this grid per image instead of using a fixed image_size. + num_patches: int | None = None + # Cosmos3-Edge video placeholder token (HF: video_token_id, default 18). + video_token_id: int | None = None + # Cosmos3-Edge vision span delimiters (HF: vision_start/end_token_id). + vision_start_token_id: int | None = None + vision_end_token_id: int | None = None + # ``True`` when the text decoder's 3D M-RoPE assigns axes to *interleaved* + # frequency channels (T, H, W, T, H, W, ...) rather than contiguous + # ``mrope_section`` chunks. ``None`` leaves the decoder default alone. + mrope_interleaved: bool | None = None # CLIP-style feature extraction: which ``hidden_states`` index to output # (HuggingFace convention, e.g. -2 for Phi-3.5-Vision). ``None`` means use # the final hidden state (all layers + post_layernorm). diff --git a/src/mobius/_configs/_wan_vae.py b/src/mobius/_configs/_wan_vae.py new file mode 100644 index 000000000..350470845 --- /dev/null +++ b/src/mobius/_configs/_wan_vae.py @@ -0,0 +1,427 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Configuration for the diffusers ``AutoencoderKLWan`` 3D causal video VAE. + +``AutoencoderKLWan`` is the video autoencoder shipped with the Wan 2.1 / Wan 2.2 +family and re-used verbatim as the video VAE of NVIDIA's Cosmos3 models +(``nvidia/Cosmos3-Nano/vae``, which declares +``_name_or_path = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"``). + +The public ``vae/config.json`` for that checkpoint is:: + + { + "_class_name": "AutoencoderKLWan", + "attn_scales": [], "base_dim": 160, "clip_output": false, + "decoder_base_dim": 256, "dim_mult": [1, 2, 4, 4], "dropout": 0.0, + "in_channels": 12, "is_residual": true, + "latents_mean": [...48 floats...], "latents_std": [...48 floats...], + "num_res_blocks": 2, "out_channels": 12, "patch_size": 2, + "scale_factor_spatial": 16, "scale_factor_temporal": 4, + "temperal_downsample": [false, true, true], "z_dim": 48 + } + +Two upstream naming quirks are preserved here: + +``temperal_downsample`` + Upstream misspells "temporal". :meth:`WanVAEConfig.from_diffusers` accepts + the misspelled JSON key (and the corrected spelling as a fallback), but the + dataclass exposes the correctly spelled :attr:`WanVAEConfig.temporal_downsample`. + +``clip_output`` + Present in the public JSON but **not** a parameter of upstream + ``AutoencoderKLWan.__init__`` — ``ConfigMixin`` drops it as an unused kwarg, + so diffusers always clamps the decoded video to ``[-1, 1]`` regardless of its + value. It is parsed here for round-trip fidelity only; see + :attr:`WanVAEConfig.clip_output`. +""" + +from __future__ import annotations + +import dataclasses +import itertools +from collections.abc import Sequence +from typing import Any + +import onnx_ir as ir + +__all__ = ["WanVAEConfig"] + +#: ``latents_mean`` from Wan 2.1's ``AutoencoderKLWan`` signature (z_dim=16). +#: Used only when a config omits the field; real checkpoints always provide it. +_WAN21_LATENTS_MEAN: tuple[float, ...] = ( + -0.7571, + -0.7089, + -0.9113, + 0.1075, + -0.1745, + 0.9653, + -0.1517, + 1.5508, + 0.4134, + -0.0715, + 0.5517, + -0.3632, + -0.1922, + -0.9497, + 0.2503, + -0.2921, +) + +#: ``latents_std`` from Wan 2.1's ``AutoencoderKLWan`` signature (z_dim=16). +_WAN21_LATENTS_STD: tuple[float, ...] = ( + 2.8184, + 1.4541, + 2.3275, + 2.6558, + 1.2196, + 1.7708, + 2.6052, + 2.0743, + 3.2687, + 2.1526, + 2.8652, + 1.5579, + 1.6382, + 1.1253, + 2.8251, + 1.9160, +) + + +def _as_dict(config: Any) -> dict[str, Any]: + """Normalise a diffusers config (dict, ``FrozenDict`` or object) to a dict.""" + if isinstance(config, dict): + return dict(config) + if hasattr(config, "items"): + return dict(config.items()) + if hasattr(config, "to_dict"): + return dict(config.to_dict()) + return dict(vars(config)) + + +def _resolve_dtype(raw: Any) -> ir.DataType: + """Map a diffusers ``dtype``/``torch_dtype`` field onto an ``ir.DataType``. + + Returns ``ir.DataType.FLOAT`` for ``None`` and for the ``"auto"`` sentinel, + matching how diffusers materialises weights when no dtype is requested. + """ + if raw is None or raw == "auto": + return ir.DataType.FLOAT + if isinstance(raw, ir.DataType): + return raw + name = str(raw).replace("torch.", "") + mapping = { + "float32": ir.DataType.FLOAT, + "float": ir.DataType.FLOAT, + "float16": ir.DataType.FLOAT16, + "half": ir.DataType.FLOAT16, + "bfloat16": ir.DataType.BFLOAT16, + } + if name not in mapping: + raise ValueError( + f"Unsupported dtype {raw!r} for AutoencoderKLWan; " + f"expected one of {sorted(mapping)}" + ) + return mapping[name] + + +@dataclasses.dataclass(frozen=True) +class WanVAEConfig: + """Architecture configuration for ``AutoencoderKLWan``. + + Attributes mirror the upstream ``AutoencoderKLWan.__init__`` signature + one-for-one, except that ``temperal_downsample`` is exposed under the + corrected spelling :attr:`temporal_downsample`. + + Attributes: + base_dim: Encoder base channel width (``160`` for Wan 2.2 / Cosmos3). + decoder_base_dim: Decoder base channel width. ``None`` means "same as + :attr:`base_dim`"; :meth:`__post_init__` resolves it eagerly so the + model never has to re-apply the default. + z_dim: Latent channel count (``48``). The encoder emits ``2 * z_dim`` + channels (mean ‖ logvar) which ``quant_conv`` maps 1:1. + dim_mult: Per-stage channel multipliers applied to the base dim. + num_res_blocks: Residual blocks per encoder stage. Decoder stages use + ``num_res_blocks + 1`` blocks, matching upstream. + attn_scales: Spatial scales at which the *non-residual* (Wan 2.1) + encoder inserts an attention block. Empty for Wan 2.2 / Cosmos3. + temporal_downsample: One flag per down/up-sampling stage + (``len(dim_mult) - 1`` entries) selecting 3D (spatio-temporal) + instead of 2D (spatial-only) resampling. Parsed from the upstream + misspelled key ``temperal_downsample``. + dropout: Upstream dropout probability. Inference-only ONNX export + always evaluates dropout in eval mode, so any value is a no-op; + retained for config fidelity. + latents_mean: Per-channel latent mean used by the *pipeline* to + normalise / denormalise latents (length ``z_dim``). + latents_std: Per-channel latent standard deviation (length ``z_dim``). + is_residual: ``True`` selects the Wan 2.2 residual down/up blocks with + ``AvgDown3D`` / ``DupUp3D`` shortcuts; ``False`` selects the flat + Wan 2.1 blocks. + in_channels: Encoder input channels **after** patchification, i.e. + ``3 * patch_size ** 2`` (``12`` for Cosmos3). + out_channels: Decoder output channels **before** unpatchification + (``12`` for Cosmos3, which unpatchifies to 3-channel RGB video). + patch_size: Spatial patch size folded into the channel dim before the + encoder and unfolded after the decoder. ``None`` disables it. + scale_factor_temporal: Total temporal compression ratio (``4``). + scale_factor_spatial: Total spatial compression ratio (``16``), which + includes the ``patch_size`` factor. + clip_output: Parsed from ``vae/config.json`` for round-trip fidelity. + Upstream ``AutoencoderKLWan.__init__`` does **not** accept it, so + diffusers silently drops it and *always* applies + ``torch.clamp(out, -1.0, 1.0)`` in ``_decode``. Mobius mirrors + diffusers and clamps unconditionally, so this flag does not change + the exported graph. + dtype: Element type used for the exported graph inputs/outputs. + """ + + base_dim: int = 96 + decoder_base_dim: int | None = None + z_dim: int = 16 + dim_mult: tuple[int, ...] = (1, 2, 4, 4) + num_res_blocks: int = 2 + attn_scales: tuple[float, ...] = () + temporal_downsample: tuple[bool, ...] = (False, True, True) + dropout: float = 0.0 + latents_mean: tuple[float, ...] = _WAN21_LATENTS_MEAN + latents_std: tuple[float, ...] = _WAN21_LATENTS_STD + is_residual: bool = False + in_channels: int = 3 + out_channels: int = 3 + patch_size: int | None = None + scale_factor_temporal: int = 4 + scale_factor_spatial: int = 8 + clip_output: bool = False + dtype: ir.DataType = ir.DataType.FLOAT + + def __post_init__(self) -> None: + """Coerce sequence fields to tuples, resolve defaults and validate.""" + object.__setattr__(self, "dim_mult", tuple(int(v) for v in self.dim_mult)) + object.__setattr__(self, "attn_scales", tuple(float(v) for v in self.attn_scales)) + object.__setattr__( + self, "temporal_downsample", tuple(bool(v) for v in self.temporal_downsample) + ) + object.__setattr__(self, "latents_mean", tuple(float(v) for v in self.latents_mean)) + object.__setattr__(self, "latents_std", tuple(float(v) for v in self.latents_std)) + if self.decoder_base_dim is None: + object.__setattr__(self, "decoder_base_dim", self.base_dim) + object.__setattr__(self, "dtype", _resolve_dtype(self.dtype)) + self.validate() + + # ------------------------------------------------------------------ + # Derived properties + # ------------------------------------------------------------------ + + @property + def temporal_upsample(self) -> tuple[bool, ...]: + """Decoder-side flags: :attr:`temporal_downsample` reversed.""" + return tuple(reversed(self.temporal_downsample)) + + @property + def temperal_downsample(self) -> tuple[bool, ...]: + """Upstream-spelled alias of :attr:`temporal_downsample`. + + Kept so code written against the diffusers field name keeps working. + """ + return self.temporal_downsample + + @property + def video_channels(self) -> int: + """Channel count of the un-patchified video tensor (3 for RGB). + + ``in_channels`` counts *patchified* channels, so the pixel-space video + has ``in_channels / patch_size ** 2`` channels. + """ + if self.patch_size is None: + return self.in_channels + return self.in_channels // (self.patch_size * self.patch_size) + + @property + def decoded_video_channels(self) -> int: + """Channel count of the decoded, un-patchified video tensor.""" + if self.patch_size is None: + return self.out_channels + return self.out_channels // (self.patch_size * self.patch_size) + + @property + def encoder_dims(self) -> tuple[int, ...]: + """Encoder stage widths ``[base, base * m0, base * m1, ...]``.""" + return tuple(self.base_dim * u for u in (1, *self.dim_mult)) + + @property + def decoder_dims(self) -> tuple[int, ...]: + """Decoder stage widths ``[dbase * m[-1], dbase * m[::-1]...]``.""" + assert self.decoder_base_dim is not None + base = self.decoder_base_dim + return tuple(base * u for u in (self.dim_mult[-1], *reversed(self.dim_mult))) + + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ + + def validate(self) -> None: + """Check the parsed fields for internal consistency. + + Raises: + ValueError: If any field is out of range or inconsistent with the + other fields (stage counts, compression ratios, latent + statistics length, or ``AvgDown3D``/``DupUp3D`` divisibility). + """ + if self.base_dim <= 0: + raise ValueError(f"base_dim must be positive, got {self.base_dim}") + assert self.decoder_base_dim is not None + if self.decoder_base_dim <= 0: + raise ValueError(f"decoder_base_dim must be positive, got {self.decoder_base_dim}") + if self.z_dim <= 0: + raise ValueError(f"z_dim must be positive, got {self.z_dim}") + if not self.dim_mult: + raise ValueError("dim_mult must contain at least one stage") + if self.num_res_blocks <= 0: + raise ValueError(f"num_res_blocks must be positive, got {self.num_res_blocks}") + if len(self.temporal_downsample) != len(self.dim_mult) - 1: + raise ValueError( + "temperal_downsample must have len(dim_mult) - 1 = " + f"{len(self.dim_mult) - 1} entries, got {len(self.temporal_downsample)}" + ) + if self.in_channels <= 0 or self.out_channels <= 0: + raise ValueError( + f"in_channels/out_channels must be positive, got " + f"{self.in_channels}/{self.out_channels}" + ) + if len(self.latents_mean) != self.z_dim: + raise ValueError( + f"latents_mean must have z_dim = {self.z_dim} entries, " + f"got {len(self.latents_mean)}" + ) + if len(self.latents_std) != self.z_dim: + raise ValueError( + f"latents_std must have z_dim = {self.z_dim} entries, " + f"got {len(self.latents_std)}" + ) + if any(std == 0 for std in self.latents_std): + raise ValueError("latents_std entries must be non-zero (used as a divisor)") + self._validate_patch_size() + self._validate_scale_factors() + if self.is_residual: + self._validate_residual_shortcuts() + + def _validate_patch_size(self) -> None: + if self.patch_size is None: + return + if self.patch_size < 1: + raise ValueError(f"patch_size must be >= 1, got {self.patch_size}") + fold = self.patch_size * self.patch_size + if self.in_channels % fold != 0: + raise ValueError( + f"in_channels ({self.in_channels}) must be divisible by " + f"patch_size ** 2 ({fold}); in_channels counts patchified channels" + ) + if self.out_channels % fold != 0: + raise ValueError( + f"out_channels ({self.out_channels}) must be divisible by " + f"patch_size ** 2 ({fold})" + ) + + def _validate_scale_factors(self) -> None: + patch = self.patch_size or 1 + expected_spatial = (2 ** (len(self.dim_mult) - 1)) * patch + if self.scale_factor_spatial != expected_spatial: + raise ValueError( + f"scale_factor_spatial ({self.scale_factor_spatial}) does not match the " + f"architecture: 2 ** (len(dim_mult) - 1) * patch_size = {expected_spatial}" + ) + expected_temporal = 2 ** sum(self.temporal_downsample) + if self.scale_factor_temporal != expected_temporal: + raise ValueError( + f"scale_factor_temporal ({self.scale_factor_temporal}) does not match the " + f"architecture: 2 ** sum(temperal_downsample) = {expected_temporal}" + ) + + def _validate_residual_shortcuts(self) -> None: + """Check ``AvgDown3D`` / ``DupUp3D`` channel divisibility per stage. + + Upstream asserts ``in_channels * factor % out_channels == 0`` (down) and + ``out_channels * factor % in_channels == 0`` (up); failing that, the + residual shortcut cannot be expressed as a grouped mean / repeat. + """ + dims = self.encoder_dims + last = len(self.dim_mult) - 1 + for i, (in_dim, out_dim) in enumerate(itertools.pairwise(dims)): + down = i != last + factor_t = 2 if (down and self.temporal_downsample[i]) else 1 + factor_s = 2 if down else 1 + factor = factor_t * factor_s * factor_s + if in_dim * factor % out_dim != 0: + raise ValueError( + f"encoder stage {i}: AvgDown3D requires in_dim * factor " + f"({in_dim} * {factor}) to be divisible by out_dim ({out_dim})" + ) + up_dims = self.decoder_dims + temporal_upsample = self.temporal_upsample + for i, (in_dim, out_dim) in enumerate(itertools.pairwise(up_dims)): + if i == last: + continue + factor_t = 2 if temporal_upsample[i] else 1 + factor = factor_t * 4 + if out_dim * factor % in_dim != 0: + raise ValueError( + f"decoder stage {i}: DupUp3D requires out_dim * factor " + f"({out_dim} * {factor}) to be divisible by in_dim ({in_dim})" + ) + + # ------------------------------------------------------------------ + # Parsing + # ------------------------------------------------------------------ + + @classmethod + def from_diffusers(cls, config: Any) -> WanVAEConfig: + """Build a :class:`WanVAEConfig` from a diffusers ``vae/config.json``. + + Args: + config: Parsed ``vae/config.json`` as a dict, a diffusers + ``FrozenDict``, or any object exposing ``items()``/``to_dict()``. + + Returns: + A validated :class:`WanVAEConfig`. + + Raises: + ValueError: If ``_class_name`` is present but is not + ``"AutoencoderKLWan"``, or if the parsed fields are inconsistent. + """ + raw = _as_dict(config) + class_name = raw.get("_class_name") + if class_name is not None and class_name != "AutoencoderKLWan": + raise ValueError( + f"WanVAEConfig expects _class_name 'AutoencoderKLWan', got {class_name!r}" + ) + # Upstream misspells "temporal"; accept the corrected spelling as a fallback + # so hand-written configs work too. + temporal_downsample: Sequence[Any] = raw.get( + "temperal_downsample", raw.get("temporal_downsample", (False, True, True)) + ) + return cls( + base_dim=int(raw.get("base_dim", 96)), + decoder_base_dim=( + int(raw["decoder_base_dim"]) + if raw.get("decoder_base_dim") is not None + else None + ), + z_dim=int(raw.get("z_dim", 16)), + dim_mult=tuple(raw.get("dim_mult", (1, 2, 4, 4))), + num_res_blocks=int(raw.get("num_res_blocks", 2)), + attn_scales=tuple(raw.get("attn_scales", ())), + temporal_downsample=tuple(temporal_downsample), + dropout=float(raw.get("dropout", 0.0)), + latents_mean=tuple(raw.get("latents_mean", _WAN21_LATENTS_MEAN)), + latents_std=tuple(raw.get("latents_std", _WAN21_LATENTS_STD)), + is_residual=bool(raw.get("is_residual", False)), + in_channels=int(raw.get("in_channels", 3)), + out_channels=int(raw.get("out_channels", 3)), + patch_size=(int(raw["patch_size"]) if raw.get("patch_size") is not None else None), + scale_factor_temporal=int(raw.get("scale_factor_temporal") or 4), + scale_factor_spatial=int(raw.get("scale_factor_spatial") or 8), + clip_output=bool(raw.get("clip_output", False)), + dtype=raw.get("dtype", raw.get("torch_dtype")), + ) diff --git a/src/mobius/_configs/_world_model.py b/src/mobius/_configs/_world_model.py index 17eea7d2e..8080799f5 100644 --- a/src/mobius/_configs/_world_model.py +++ b/src/mobius/_configs/_world_model.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Configuration for directly declared world models.""" +"""Configuration for directly declared latent-dynamics models.""" from __future__ import annotations @@ -12,12 +12,12 @@ @dataclasses.dataclass -class WorldModelConfig(BaseModelConfig): - """Configuration shared by single-step world-model graphs. +class LatentDynamicsConfig(BaseModelConfig): + """Configuration shared by single-step latent-dynamics graphs. The three shapes exclude the leading batch dimension. The default - :class:`~mobius.models.MLPWorldModel` flattens each value internally, while - custom modules may preserve their original ranks. + :class:`~mobius.models.MLPLatentDynamicsModel` flattens each value + internally, while custom modules may preserve their original ranks. """ observation_shape: tuple[int, ...] = (1,) @@ -44,7 +44,7 @@ def state_size(self) -> int: return math.prod(self.state_shape) def validate(self) -> None: - """Validate dimensions required by the world-model task and reference model.""" + """Validate dimensions required by the dynamics task and reference model.""" for name, shape in ( ("observation_shape", self.observation_shape), ("action_shape", self.action_shape), @@ -62,3 +62,9 @@ def validate(self) -> None: raise ValueError("num_hidden_layers must be positive") if self.hidden_act is None: raise ValueError("hidden_act must be set") + + +# Backward compatibility for the original, overly broad name. A full world +# model is a pipeline of heterogeneous components; this config describes only +# one state-transition component within such a pipeline. +WorldModelConfig = LatentDynamicsConfig diff --git a/src/mobius/_configs/per_model/_cosmos3_edge_vision.py b/src/mobius/_configs/per_model/_cosmos3_edge_vision.py index bd04c31ca..75952a650 100644 --- a/src/mobius/_configs/per_model/_cosmos3_edge_vision.py +++ b/src/mobius/_configs/per_model/_cosmos3_edge_vision.py @@ -4,16 +4,23 @@ """Vision extractor hook for NVIDIA Cosmos3-Edge. Cosmos3-Edge (``cosmos3_edge`` / ``Cosmos3EdgeForConditionalGeneration``) pairs -a SigLIP-style vision encoder (``cosmos3_edge_vision``) with a pixel-shuffle -merger projector (``cosmos3_edge_projector``). Two config quirks need bridging -into :class:`~mobius._configs.VisionConfig`: +a **variable-resolution** SigLIP2 vision tower (``cosmos3_edge_vision``) with a +pixel-shuffle merger projector (``cosmos3_edge_projector``). Several config +quirks need bridging into :class:`~mobius._configs.VisionConfig`: -- The vision config declares ``num_patches`` (256) instead of ``image_size``. - The standard :class:`~mobius.components.PatchEmbedding` derives the patch - count from ``image_size // patch_size``, so we reconstruct - ``image_size = sqrt(num_patches) * patch_size`` (16 * 16 = 256). -- The projector's intermediate width lives in a sibling ``projector_config`` - (``merger_intermediate_size``), not in ``vision_config``. +- The vision config declares ``num_patches`` (256), the size of the *learned + position-embedding reference grid*, not a fixed input resolution. The tower + resamples that 16x16 grid to each image's own patch grid, so ``image_size`` + is only a nominal value (``sqrt(num_patches) * patch_size``) retained for + metadata — it does not constrain the exported graph. +- The projector's intermediate width and norm placement live in a sibling + ``projector_config`` (``merger_intermediate_size``, + ``use_postshuffle_norm``), not in ``vision_config``. +- Cosmos3-Edge only supports ``temporal_patch_size=1`` (each video frame is an + independent patch run), unlike the Qwen-VL default of 2. +- The text decoder uses **interleaved** 3D M-RoPE (T, H, W, T, H, W, ...) + rather than the contiguous ``mrope_section`` chunking used by Qwen-VL. HF + reference: ``Cosmos3EdgeTextRotaryEmbedding.compute_default_rope_parameters``. """ from __future__ import annotations @@ -22,28 +29,39 @@ from mobius._configs._extractors import register_vision_hook +_DEFAULT_IMAGE_TOKEN_ID = 19 +_DEFAULT_VIDEO_TOKEN_ID = 18 +_DEFAULT_VISION_START_TOKEN_ID = 20 +_DEFAULT_VISION_END_TOKEN_ID = 21 -@register_vision_hook("cosmos3_edge", "cosmos3_edge_vision") + +@register_vision_hook("cosmos3_edge", "cosmos3_edge_text", "cosmos3_edge_vision") def _cosmos3_edge_vision(config, parent_config, model_type: str, fields: dict): vision_source = parent_config or config hf_vision = getattr(vision_source, "vision_config", None) or getattr( config, "vision_config", None ) - # Reconstruct image_size from num_patches (256 -> 16x16 grid -> 256 px). - if fields.get("image_size") is None and hf_vision is not None: + if hf_vision is not None: num_patches = getattr(hf_vision, "num_patches", None) patch_size = getattr(hf_vision, "patch_size", None) or fields.get("patch_size") - if num_patches is not None and patch_size is not None: + if num_patches is not None: grid = math.isqrt(num_patches) if grid * grid != num_patches: raise ValueError( "Cosmos3-Edge vision num_patches must form a square grid, " f"got {num_patches}" ) - fields["image_size"] = grid * patch_size + fields["num_patches"] = num_patches + # Nominal resolution of the learned position grid (16 * 16 = 256). + if fields.get("image_size") is None and patch_size is not None: + fields["image_size"] = grid * patch_size + fields["hidden_act"] = getattr(hf_vision, "hidden_act", None) + + # Cosmos3-Edge patchifies one frame at a time (``temporal_patch_size=1``). + fields["temporal_patch_size"] = 1 - # Pixel-shuffle projector intermediate size from projector_config. + # Pixel-shuffle projector settings from projector_config. projector_cfg = getattr(vision_source, "projector_config", None) or getattr( config, "projector_config", None ) @@ -54,6 +72,7 @@ def _cosmos3_edge_vision(config, parent_config, model_type: str, fields: dict): else lambda k, d=None: getattr(projector_cfg, k, d) ) fields["projector_intermediate_size"] = get("merger_intermediate_size") + fields["use_postshuffle_norm"] = bool(get("use_postshuffle_norm", False)) # out_hidden_size drives the projector output (text hidden size). if fields.get("out_hidden_size") is None: fields["out_hidden_size"] = get("out_hidden_size") @@ -61,7 +80,23 @@ def _cosmos3_edge_vision(config, parent_config, model_type: str, fields: dict): if merge is not None: fields["spatial_merge_size"] = merge - # Cosmos3-Edge places image_token_id at the top level (default 19). + # Cosmos3-Edge places the multimodal token ids at the top level. if fields.get("image_token_id") is None: - fields["image_token_id"] = getattr(vision_source, "image_token_id", 19) + fields["image_token_id"] = getattr( + vision_source, "image_token_id", _DEFAULT_IMAGE_TOKEN_ID + ) + fields["video_token_id"] = getattr( + vision_source, "video_token_id", _DEFAULT_VIDEO_TOKEN_ID + ) + fields["vision_start_token_id"] = getattr( + vision_source, "vision_start_token_id", _DEFAULT_VISION_START_TOKEN_ID + ) + fields["vision_end_token_id"] = getattr( + vision_source, "vision_end_token_id", _DEFAULT_VISION_END_TOKEN_ID + ) + + # Interleaved 3D M-RoPE — the Edge checkpoint stores ``mrope_section`` under + # ``rope_parameters`` without an explicit interleave flag, but + # ``Cosmos3EdgeTextRotaryEmbedding`` builds its axis masks from ``i % 3``. + fields["mrope_interleaved"] = True return None diff --git a/src/mobius/_cosmos3_edge_world_model.py b/src/mobius/_cosmos3_edge_world_model.py new file mode 100644 index 000000000..2e0da6813 --- /dev/null +++ b/src/mobius/_cosmos3_edge_world_model.py @@ -0,0 +1,322 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Complete world-model exporter for NVIDIA Cosmos3-Edge checkpoints.""" + +from __future__ import annotations + +__all__ = ["build_cosmos3_edge_world_model"] + +import json +from collections.abc import Mapping +from typing import Any + +import onnx_ir as ir + +from mobius._configs.per_model import _cosmos3_edge_vision # noqa: F401 +from mobius._cosmos3_world_model import ( + _REASONER_NAMES, + _apply_checkpoint_weights, + _build_components, + _collect_assets, + _compose_pipeline, +) +from mobius._diffusers_checkpoint import ( + component_class, + load_checkpoint_json, + load_optional_checkpoint_json, + resolve_checkpoint_file, +) +from mobius._pipeline import PipelinePackage +from mobius._world_model_config import ( + WorldModelBuildConfig, + WorldModelGenerationConfig, + WorldModelPipelineConfig, +) +from mobius.models.cosmos import Cosmos3EdgeVLModel + +# Official Edge generation recipes use 50 steps for I2V and a mode-specific +# flow schedule. Action recipes use a separate flow schedule. +_SCHEDULER_MODE_OVERRIDES: dict[str, Any] = { + "image_to_video": { + "flow_shift": 3.0, + "use_karras_sigmas": False, + "num_inference_steps": 50, + "guidance_scale": 5.0, + }, + "action": { + "flow_shift": 10.0, + "use_karras_sigmas": False, + "num_inference_steps": 30, + "guidance_scale": 1.0, + }, +} +_DEFAULT_INFERENCE_STEPS = 50 + +# Cosmos3-Edge Reasoner vision/video-understanding contract. The exported +# ``reasoner_vision_encoder`` graph consumes the *pre-patchified* tensor that +# ``Cosmos3EdgeImageProcessor`` / ``Cosmos3EdgeVideoProcessor`` produce, so the +# manifest has to spell out that preprocessing for a runtime to reproduce it. +_IMAGE_PROCESSOR_ASSET = "preprocessor_config.json" +_VIDEO_PROCESSOR_ASSET = "video_preprocessor_config.json" + + +def _vision_understanding_metadata( + root_config: Mapping[str, Any], + reasoner_config: Any, +) -> dict[str, Any]: + """Describe the Reasoner's packed image/video input contract.""" + vision = getattr(reasoner_config, "vision", None) + patch_size = getattr(vision, "patch_size", None) or 16 + merge_size = getattr(vision, "spatial_merge_size", None) or 2 + channels = getattr(vision, "in_channels", None) or 3 + temporal_patch_size = getattr(vision, "temporal_patch_size", None) or 1 + return { + "encoder": _REASONER_NAMES["vision_encoder"], + "invocation": "once_per_visual_item", + "tokens": { + "image": root_config.get("image_token_id", 19), + "video": root_config.get("video_token_id", 18), + "vision_start": root_config.get("vision_start_token_id", 20), + "vision_end": root_config.get("vision_end_token_id", 21), + }, + "token_expansion": { + "image": ( + "<|vision_start|> + grid_h*grid_w/merge_size**2 image tokens + <|vision_end|>" + ), + "video": ( + "one '' timestamp followed by " + "<|vision_start|> + grid_h*grid_w/merge_size**2 video tokens + " + "<|vision_end|>, repeated per sampled frame" + ), + }, + "routing": { + "image": f"{_REASONER_NAMES['embedding']}.image_features", + "video": f"{_REASONER_NAMES['embedding']}.video_features", + }, + "presence": {"video": "video_understanding"}, + "preprocessing": { + "image_processor_asset": _IMAGE_PROCESSOR_ASSET, + "video_processor_asset": _VIDEO_PROCESSOR_ASSET, + # ``size`` in the shipped processor assets holds pixel *areas* + # (shortest_edge/longest_edge), not edge lengths. + "resize": "smart_resize_area_bounded_multiple_of_patch_times_merge", + "alignment": patch_size * merge_size, + # Class defaults of Cosmos3EdgeImageProcessor / + # Cosmos3EdgeVideoProcessor; they are NOT serialised into the + # shipped *_preprocessor_config.json assets. + "resample": "bicubic", + "convert_rgb": True, + "rescale_factor": 1 / 255, + "normalize": {"mean": [0.5, 0.5, 0.5], "std": [0.5, 0.5, 0.5]}, + "video_frame_sampling": {"fps": 2, "min_frames": 4, "max_frames": 768}, + "patchify": { + "layout": "time_major_block_major", + "patch_value_order": "patch_height_patch_width_channel", + "patch_size": patch_size, + "merge_size": merge_size, + "temporal_patch_size": temporal_patch_size, + "patch_dim": patch_size * patch_size * channels * temporal_patch_size, + }, + }, + "grid_thw": { + "layout": "t_h_w", + "units": "patches", + "source": "image_grid_thw[i] / video_grid_thw[i]", + "note": "grid_h and grid_w must be multiples of merge_size", + }, + "position_ids": { + "mrope": "interleaved", + "mrope_section": list(getattr(reasoner_config, "mrope_section", None) or []), + "axis_assignment": ( + "channel i uses height when i%3==1 and i<3*mrope_section[1], " + "width when i%3==2 and i<3*mrope_section[2], temporal otherwise" + ), + "video_index_rule": ( + "expand video_grid_thw to one row per frame and set grid_t=1, so " + "every frame is an independent visual span for position indexing" + ), + }, + } + + +def _edge_text_model_type(root_config: Mapping[str, Any]) -> str | None: + text_config = root_config.get("text_config") + if isinstance(text_config, Mapping): + model_type = text_config.get("model_type") + return model_type if isinstance(model_type, str) else None + return None + + +def build_cosmos3_edge_world_model( + model_id: str, + *, + dtype: str | ir.DataType | None = None, + load_weights: bool = True, + execution_provider: str = "default", + trace_optimization: bool = False, + **_options: Any, +) -> PipelinePackage: + """Build the complete Cosmos3-Edge Reasoner/Generator/VAE/Action package. + + Both ``nvidia/Cosmos3-Edge`` and the historically mislabeled + ``nvidia/Cosmos3-Edge-Policy-DROID`` use the Edge text/vision architecture. + The latter advertises top-level ``model_type="cosmos3_omni"``; dispatch is + therefore based on ``text_config.model_type="cosmos3_edge_text"``. + """ + build_config = WorldModelBuildConfig( + dtype=dtype, + load_weights=load_weights, + execution_provider=execution_provider, + trace_optimization=trace_optimization, + ) + root_config, _ = load_checkpoint_json(model_id, "config.json") + if _edge_text_model_type(root_config) != "cosmos3_edge_text": + raise ValueError( + f"{model_id!r} is not a Cosmos3-Edge checkpoint: expected " + "text_config.model_type='cosmos3_edge_text'." + ) + + pipeline_index, _ = load_checkpoint_json(model_id, "model_index.json") + transformer_class = component_class(pipeline_index, "transformer") + vae_class = component_class(pipeline_index, "vae") + sound_class = component_class(pipeline_index, "sound_tokenizer") + if transformer_class != "Cosmos3OmniTransformer": + raise ValueError(f"Unsupported Cosmos3-Edge transformer class {transformer_class!r}") + if vae_class != "AutoencoderKLWan": + raise ValueError(f"Unsupported Cosmos3-Edge VAE class {vae_class!r}") + if sound_class is not None: + raise ValueError( + "Cosmos3-Edge Sound generation is not supported by the public architecture; " + f"unexpected sound tokenizer {sound_class!r}." + ) + + transformer_config, _ = load_checkpoint_json(model_id, "transformer/config.json") + vae_config, _ = load_checkpoint_json(model_id, "vae/config.json") + scheduler_config, _ = load_checkpoint_json(model_id, "scheduler/scheduler_config.json") + generation_config = WorldModelGenerationConfig.from_generation_config( + load_optional_checkpoint_json(model_id, "generation_config.json"), + default_inference_steps=_DEFAULT_INFERENCE_STEPS, + scheduler_mode_overrides=_SCHEDULER_MODE_OVERRIDES, + ) + + ( + reasoner_package, + reasoner_module, + generator_package, + generator_module, + vae_package, + vae_module, + audio_package, + audio_module, + ) = _build_components( + model_id, + build_config=build_config, + pipeline_index=pipeline_index, + transformer_config_dict=transformer_config, + vae_config_dict=vae_config, + audio_config_dict=None, + audio_weight_names=None, + has_reasoner_vision=True, + reasoner_module_class=Cosmos3EdgeVLModel, + reasoner_task="cosmos3-edge-vl", + ) + assert audio_package is None and audio_module is None + + if build_config.load_weights: + _apply_checkpoint_weights( + model_id, + reasoner_package=reasoner_package, + reasoner_module=reasoner_module, + generator_package=generator_package, + generator_module=generator_module, + vae_package=vae_package, + vae_module=vae_module, + audio_package=None, + audio_module=None, + ) + + assets = _collect_assets(model_id, has_sound_tokenizer=False) + negative_prompt_asset = "assets/negative_prompt.json" + i2v_prompt: dict[str, Any] = { + "positive": "json_or_text", + "negative_default": ("asset" if negative_prompt_asset in assets else "empty"), + "add_resolution_template": False, + "add_duration_template": False, + "use_system_prompt": False, + } + if negative_prompt_asset in assets: + i2v_prompt["negative_asset"] = negative_prompt_asset + text_config = root_config.get("text_config") + eos_token_id = ( + text_config.get("eos_token_id") + if isinstance(text_config, Mapping) + else root_config.get("eos_token_id") + ) + vision_start_token_id = root_config.get("vision_start_token_id") + if not isinstance(eos_token_id, int) or not isinstance(vision_start_token_id, int): + raise TypeError( + "Cosmos3-Edge requires integer text_config.eos_token_id and " + "vision_start_token_id for generator prompt packing." + ) + policy: dict[str, Any] | None = None + checkpoint_path = resolve_checkpoint_file(model_id, "checkpoint.json", required=False) + if checkpoint_path is not None: + with open(checkpoint_path, encoding="utf-8") as handle: + checkpoint = json.load(handle) + if isinstance(checkpoint, Mapping) and isinstance(checkpoint.get("policy"), dict): + policy = dict(checkpoint["policy"]) + return _compose_pipeline( + pipeline_config=WorldModelPipelineConfig( + model_id=model_id, + model_type="cosmos3_edge", + build=build_config, + generation=generation_config, + extra_metadata={ + "edge": { + "checkpoint_model_type": root_config.get("model_type"), + "policy": policy, + }, + "generation_recipes": { + "image_to_video": { + "conditioning": { + "modality": "image", + "encoder_stage": "encode_video", + "conditioned_latent_frames": [0], + }, + "prompt": i2v_prompt, + "height": 480, + "width": 832, + "frames": 121, + "fps": 24.0, + }, + }, + "generator_prompt": { + "chat": { + "add_generation_prompt": True, + "add_vision_id": False, + "enable_thinking": True, + }, + "suffix_token_ids": [ + eos_token_id, + vision_start_token_id, + ], + }, + "vision_understanding": _vision_understanding_metadata( + root_config, reasoner_package.config + ), + }, + ), + reasoner_package=reasoner_package, + generator_package=generator_package, + vae_package=vae_package, + audio_package=None, + generator_config=generator_module.config, + vae_config=vae_module.config, + scheduler_config=scheduler_config, + assets=assets, + reasoner_architecture="cosmos3_edge", + default_action_domain=( + policy.get("domain_name", "no_action") if policy is not None else "no_action" + ), + ) diff --git a/src/mobius/_cosmos3_edge_world_model_test.py b/src/mobius/_cosmos3_edge_world_model_test.py new file mode 100644 index 000000000..1f74ba57c --- /dev/null +++ b/src/mobius/_cosmos3_edge_world_model_test.py @@ -0,0 +1,240 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +from types import SimpleNamespace +from unittest import mock + +import onnx_ir as ir +import pytest + +from mobius._cosmos3_edge_world_model import ( + _edge_text_model_type, + build_cosmos3_edge_world_model, +) +from mobius._cosmos3_world_model import build_cosmos3_world_model +from mobius._world_model_builder import world_model_registry +from mobius.models.cosmos import Cosmos3EdgeVLModel + + +def test_edge_model_type_is_registered() -> None: + assert "cosmos3_edge" in world_model_registry.model_types() + + +def test_edge_text_model_type_uses_nested_config() -> None: + assert ( + _edge_text_model_type( + {"model_type": "cosmos3_omni", "text_config": {"model_type": "cosmos3_edge_text"}} + ) + == "cosmos3_edge_text" + ) + assert _edge_text_model_type({"model_type": "cosmos3_edge"}) is None + + +def test_omni_dispatch_delegates_mislabeled_edge_checkpoint(tmp_path) -> None: + (tmp_path / "model_index.json").write_text("{}", encoding="utf-8") + (tmp_path / "config.json").write_text( + json.dumps( + { + "model_type": "cosmos3_omni", + "text_config": {"model_type": "cosmos3_edge_text"}, + } + ), + encoding="utf-8", + ) + package = object() + + with mock.patch( + "mobius._cosmos3_edge_world_model.build_cosmos3_edge_world_model", + return_value=package, + ) as edge_builder: + result = build_cosmos3_world_model( + str(tmp_path), + dtype="bf16", + load_weights=False, + execution_provider="cuda", + trace_optimization=True, + custom=True, + ) + + assert result is package + edge_builder.assert_called_once_with( + str(tmp_path), + dtype="bf16", + load_weights=False, + execution_provider="cuda", + trace_optimization=True, + custom=True, + ) + + +def test_edge_builder_uses_edge_reasoner_and_shared_generator_pipeline() -> None: + loaded = { + "config.json": ( + { + "model_type": "cosmos3_edge", + "text_config": { + "model_type": "cosmos3_edge_text", + "eos_token_id": 11, + }, + "vision_start_token_id": 20, + }, + "config.json", + ), + "model_index.json": ( + { + "transformer": ["diffusers", "Cosmos3OmniTransformer"], + "vae": ["diffusers", "AutoencoderKLWan"], + "sound_tokenizer": [None, None], + }, + "model_index.json", + ), + "transformer/config.json": ({"hidden_size": 8}, "transformer/config.json"), + "vae/config.json": ({"z_dim": 2}, "vae/config.json"), + "scheduler/scheduler_config.json": ( + {"prediction_type": "flow_prediction"}, + "scheduler/scheduler_config.json", + ), + } + reasoner_package = SimpleNamespace( + config=SimpleNamespace( + vision=SimpleNamespace( + patch_size=16, + spatial_merge_size=2, + in_channels=3, + temporal_patch_size=1, + ), + mrope_section=[24, 20, 20], + ) + ) + reasoner_module = mock.sentinel.reasoner_module + generator_package = mock.sentinel.generator_package + generator_module = SimpleNamespace(config=mock.sentinel.generator_config) + vae_package = mock.sentinel.vae_package + vae_module = SimpleNamespace(config=mock.sentinel.vae_config) + package = mock.sentinel.pipeline_package + + with ( + mock.patch( + "mobius._cosmos3_edge_world_model.load_checkpoint_json", + side_effect=lambda _model_id, filename: loaded[filename], + ), + mock.patch( + "mobius._cosmos3_edge_world_model._build_components", + return_value=( + reasoner_package, + reasoner_module, + generator_package, + generator_module, + vae_package, + vae_module, + None, + None, + ), + ) as build_components, + mock.patch( + "mobius._cosmos3_edge_world_model._collect_assets", + return_value={ + "assets/negative_prompt.json": ( + "cached/assets/negative_prompt.json", + False, + ) + }, + ), + mock.patch( + "mobius._cosmos3_edge_world_model.load_optional_checkpoint_json", + return_value={}, + ), + mock.patch( + "mobius._cosmos3_edge_world_model.resolve_checkpoint_file", + return_value=None, + ), + mock.patch( + "mobius._cosmos3_edge_world_model._compose_pipeline", + return_value=package, + ) as compose, + ): + result = build_cosmos3_edge_world_model( + "nvidia/Cosmos3-Edge", + dtype="bf16", + load_weights=False, + ) + + assert result is package + assert build_components.call_args.kwargs["reasoner_module_class"] is Cosmos3EdgeVLModel + assert build_components.call_args.kwargs["reasoner_task"] == "cosmos3-edge-vl" + build_config = build_components.call_args.kwargs["build_config"] + assert build_config.resolved_dtype() is ir.DataType.BFLOAT16 + assert build_config.load_weights is False + pipeline_config = compose.call_args.kwargs["pipeline_config"] + assert pipeline_config.model_type == "cosmos3_edge" + assert pipeline_config.model_id == "nvidia/Cosmos3-Edge" + assert pipeline_config.build is build_config + assert compose.call_args.kwargs["reasoner_architecture"] == "cosmos3_edge" + assert pipeline_config.extra_metadata["edge"]["checkpoint_model_type"] == "cosmos3_edge" + assert pipeline_config.generation.default_inference_steps == 50 + assert pipeline_config.generation.scheduler_mode_overrides_manifest() == { + "image_to_video": { + "flow_shift": 3.0, + "use_karras_sigmas": False, + "num_inference_steps": 50, + "guidance_scale": 5.0, + }, + "action": { + "flow_shift": 10.0, + "use_karras_sigmas": False, + "num_inference_steps": 30, + "guidance_scale": 1.0, + }, + } + i2v = pipeline_config.extra_metadata["generation_recipes"]["image_to_video"] + assert i2v["conditioning"]["conditioned_latent_frames"] == [0] + assert i2v["prompt"]["negative_asset"] == "assets/negative_prompt.json" + assert i2v["prompt"]["negative_default"] == "asset" + assert (i2v["width"], i2v["height"], i2v["frames"]) == (832, 480, 121) + assert pipeline_config.extra_metadata["generator_prompt"] == { + "chat": { + "add_generation_prompt": True, + "add_vision_id": False, + "enable_thinking": True, + }, + "suffix_token_ids": [11, 20], + } + vision = pipeline_config.extra_metadata["vision_understanding"] + assert vision["encoder"] == "reasoner_vision_encoder" + assert vision["tokens"] == { + "image": 19, + "video": 18, + "vision_start": 20, + "vision_end": 21, + } + assert vision["routing"]["video"] == "reasoner_embedding.video_features" + preprocessing = vision["preprocessing"] + assert preprocessing["patchify"] == { + "layout": "time_major_block_major", + "patch_value_order": "patch_height_patch_width_channel", + "patch_size": 16, + "merge_size": 2, + "temporal_patch_size": 1, + "patch_dim": 16 * 16 * 3, + } + assert preprocessing["alignment"] == 32 + # Class defaults that the shipped preprocessor_config.json assets omit. + assert preprocessing["resample"] == "bicubic" + assert preprocessing["rescale_factor"] == pytest.approx(1 / 255) + assert preprocessing["convert_rgb"] is True + assert vision["position_ids"]["mrope"] == "interleaved" + assert vision["position_ids"]["mrope_section"] == [24, 20, 20] + assert "grid_t=1" in vision["position_ids"]["video_index_rule"] + + +def test_edge_builder_rejects_non_edge_checkpoint(tmp_path) -> None: + (tmp_path / "config.json").write_text( + json.dumps({"model_type": "cosmos3_omni"}), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="not a Cosmos3-Edge checkpoint"): + build_cosmos3_edge_world_model(str(tmp_path), load_weights=False) diff --git a/src/mobius/_cosmos3_world_model.py b/src/mobius/_cosmos3_world_model.py new file mode 100644 index 000000000..56e559c77 --- /dev/null +++ b/src/mobius/_cosmos3_world_model.py @@ -0,0 +1,1266 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Complete compositional exporter for NVIDIA Cosmos3-Omni world models. + +The public Cosmos3-Omni checkpoint combines several independently executable +neural systems: + +* a Qwen3-VL Reasoner (decoder, vision encoder, and embedding graphs); +* the unified MoT rectified-flow transformer, including optional Sound and + Action heads; +* a Wan video VAE (encoder and decoder); +* an optional Cosmos3 AVAE sound tokenizer (full or decoder-only). + +This module builds those graphs through their normal Mobius tasks and composes +them into a :class:`mobius.PipelinePackage`. The topology manifest stays +runtime-agnostic: schedulers, tokenizers, and processors are copied as opaque +assets, while transform edges declare the capabilities a runtime must supply. +""" + +from __future__ import annotations + +__all__ = ["build_cosmos3_world_model"] + +import dataclasses +from collections.abc import Iterable, Mapping +from typing import Any + +import onnx_ir as ir + +from mobius._builder import build as build_model +from mobius._builder import build_from_module +from mobius._configs import ( + Cosmos3AudioConfig, + Cosmos3OmniGeneratorConfig, + WanVAEConfig, +) +from mobius._diffusers_checkpoint import ( + component_class, + component_weight_names, + load_checkpoint_json, + load_component_weights, + load_optional_checkpoint_json, + resolve_assets, +) +from mobius._model_package import ModelPackage +from mobius._pipeline import ( + PipelineBuilder, + PipelinePackage, + register_transform, +) +from mobius._weight_loading import iter_weight_shards +from mobius._world_model_config import ( + WorldModelBuildConfig, + WorldModelGenerationConfig, + WorldModelPipelineConfig, +) +from mobius.models.cosmos3_audio import create_cosmos3_avae_audio_tokenizer +from mobius.models.cosmos3_omni import Cosmos3OmniReasonerModel +from mobius.models.cosmos3_omni_generator import Cosmos3OmniGeneratorModel +from mobius.models.wan_vae import AutoencoderKLWanModel +from mobius.tasks._cosmos3_audio import select_cosmos3_audio_task +from mobius.tasks._wan_vae import WanVAETask + +_VIDEO_DIFFUSION_FINALIZE = "video_diffusion_finalize" +register_transform( + _VIDEO_DIFFUSION_FINALIZE, + description=( + "Advance the final scheduler step and restore packed diffusion tokens " + "to a video decoder's latent tensor layout" + ), + capabilities=( + "iterative_scheduler", + "tensor_patchify", + "tensor_reshape", + "tensor_cast", + ), + required_parameters=( + "scheduler_asset", + "state", + "spatial_patch_size", + "latent_channels", + "input_layout", + "output_layout", + "source_dtype", + "target_dtype", + ), + allowed_parameters=( + "scheduler_asset", + "state", + "spatial_patch_size", + "latent_channels", + "input_layout", + "output_layout", + "source_dtype", + "target_dtype", + ), +) +_AUDIO_DIFFUSION_FINALIZE = "audio_diffusion_finalize" +register_transform( + _AUDIO_DIFFUSION_FINALIZE, + description=( + "Advance the final scheduler step and restore packed sound tokens " + "to an audio decoder's latent tensor layout" + ), + capabilities=( + "iterative_scheduler", + "tensor_reshape", + "tensor_cast", + ), + required_parameters=( + "scheduler_asset", + "state", + "input_layout", + "output_layout", + "source_dtype", + "target_dtype", + ), + allowed_parameters=( + "scheduler_asset", + "state", + "input_layout", + "output_layout", + "source_dtype", + "target_dtype", + ), +) + +_REASONER_NAMES = { + "decoder": "reasoner_decoder", + "vision_encoder": "reasoner_vision_encoder", + "embedding": "reasoner_embedding", +} +# Presence key for the optional video-understanding feature stream on the +# Reasoner embedding graph (Cosmos3-Edge routes video frames through the same +# vision encoder but scatters them at the video placeholder id). +_VIDEO_UNDERSTANDING = "video_understanding" +_GENERATOR_NAME = "generator" +_VIDEO_ENCODER_NAME = "video_encoder" +_VIDEO_DECODER_NAME = "video_decoder" +_AUDIO_ENCODER_NAME = "audio_encoder" +_AUDIO_DECODER_NAME = "audio_decoder" + +# Public diffusers Cosmos3 action-domain contract. Domain IDs select the +# DomainAwareLinear bank; raw widths remove zero-padding from action outputs. +_ACTION_DOMAIN_IDS: dict[str, int] = { + "no_action": 0, + "av": 1, + "camera_pose": 2, + "hand_pose": 3, + "pusht": 4, + "libero": 5, + "umi": 6, + "bridge_orig_lerobot": 7, + "droid_lerobot": 8, + "robomind-franka": 8, + "galbot": 9, + "robomind-franka-dual": 12, + "robomind-ur": 13, + "agibotworld": 15, + "agibot_gear_gripper": 15, + "agibot_gear_gripper_ext": 15, + "fractal": 20, +} +_ACTION_RAW_DIMS: dict[str, int] = { + "no_action": 0, + "av": 9, + "camera_pose": 9, + "hand_pose": 57, + "pusht": 2, + "umi": 10, + "bridge_orig_lerobot": 10, + "droid_lerobot": 10, + "robomind-franka": 10, + "galbot": 30, + "robomind-franka-dual": 20, + "robomind-ur": 10, + "agibotworld": 29, + "agibot_gear_gripper": 29, + "agibot_gear_gripper_ext": 29, + "fractal": 10, +} + +_ASSET_CANDIDATES: tuple[tuple[str, bool], ...] = ( + ("config.json", True), + ("model_index.json", True), + ("transformer/config.json", True), + ("vae/config.json", True), + ("vision_encoder/config.json", False), + ("scheduler/scheduler_config.json", True), + ("tokenizer.json", False), + ("tokenizer_config.json", False), + ("special_tokens_map.json", False), + ("vocab.json", False), + ("merges.txt", False), + ("chat_template.json", False), + ("chat_template.jinja", False), + ("generation_config.json", False), + ("checkpoint.json", False), + ("assets/negative_prompt.json", False), + ("preprocessor_config.json", False), + ("video_preprocessor_config.json", False), + ("processor_config.json", False), + ("text_tokenizer/tokenizer.json", False), + ("text_tokenizer/tokenizer_config.json", False), + ("text_tokenizer/vocab.json", False), + ("text_tokenizer/merges.txt", False), +) + + +def _build_components( + model_id: str, + *, + build_config: WorldModelBuildConfig, + pipeline_index: Mapping[str, Any], + transformer_config_dict: Mapping[str, Any], + vae_config_dict: Mapping[str, Any], + audio_config_dict: Mapping[str, Any] | None, + audio_weight_names: Iterable[str] | None, + has_reasoner_vision: bool, + reasoner_module_class: type[Any] = Cosmos3OmniReasonerModel, + reasoner_task: str | None = None, +) -> tuple[ + ModelPackage, + Any, + ModelPackage, + Cosmos3OmniGeneratorModel, + ModelPackage, + AutoencoderKLWanModel, + ModelPackage | None, + Any | None, +]: + """Build all component graphs while retaining modules for weight routing.""" + reasoner_package = build_model( + model_id, + task=reasoner_task, + module_class=reasoner_module_class, + dtype=build_config.dtype, + load_weights=False, + execution_provider=build_config.execution_provider, + trace_optimization=build_config.trace_optimization, + ) + if not has_reasoner_vision: + reasoner_package.pop("vision_encoder", None) + expected_reasoner = {"decoder", "embedding"} + if has_reasoner_vision: + expected_reasoner.add("vision_encoder") + if set(reasoner_package) != expected_reasoner: + raise ValueError( + "Cosmos3 Reasoner components do not match the checkpoint layout; " + f"got {sorted(reasoner_package)}" + ) + reasoner_module = reasoner_module_class(reasoner_package.config) + + generator_config = Cosmos3OmniGeneratorConfig.from_diffusers(transformer_config_dict) + if resolved_dtype := build_config.resolved_dtype(): + generator_config = dataclasses.replace(generator_config, dtype=resolved_dtype) + generator_config.validate() + generator_module = Cosmos3OmniGeneratorModel(generator_config) + generator_package = build_from_module( + generator_module, + generator_config, + task="cosmos3-omni-generator", + execution_provider=build_config.execution_provider, + trace_optimization=build_config.trace_optimization, + ) + + vae_config = WanVAEConfig.from_diffusers(vae_config_dict) + vae_module = AutoencoderKLWanModel(vae_config) + vae_package = build_from_module( + vae_module, + vae_config, + task=WanVAETask(), + execution_provider=build_config.execution_provider, + trace_optimization=build_config.trace_optimization, + ) + + audio_package: ModelPackage | None = None + audio_module: Any | None = None + sound_class = component_class(pipeline_index, "sound_tokenizer") + if generator_config.sound_gen and sound_class is None: + raise ValueError( + "The transformer enables Sound generation, but model_index.json has no " + "sound_tokenizer component." + ) + if sound_class is not None: + if sound_class != "Cosmos3AVAEAudioTokenizer": + raise ValueError(f"Unsupported Cosmos3 sound tokenizer class {sound_class!r}") + if audio_config_dict is None or audio_weight_names is None: + raise ValueError("Sound tokenizer config and weight metadata are required") + audio_config = Cosmos3AudioConfig.from_diffusers( + audio_config_dict, + weight_names=audio_weight_names, + ) + audio_module = create_cosmos3_avae_audio_tokenizer(audio_config) + audio_task = select_cosmos3_audio_task(audio_config)() + audio_package = build_from_module( + audio_module, + audio_config, + task=audio_task, + execution_provider=build_config.execution_provider, + trace_optimization=build_config.trace_optimization, + ) + + return ( + reasoner_package, + reasoner_module, + generator_package, + generator_module, + vae_package, + vae_module, + audio_package, + audio_module, + ) + + +def _apply_checkpoint_weights( + model_id: str, + *, + reasoner_package: ModelPackage, + reasoner_module: Any, + generator_package: ModelPackage, + generator_module: Cosmos3OmniGeneratorModel, + vae_package: ModelPackage, + vae_module: AutoencoderKLWanModel, + audio_package: ModelPackage | None, + audio_module: Any | None, +) -> None: + """Stream the shared checkpoint once, routing each shard to both towers.""" + if reasoner_module.config.tie_word_embeddings: + raise ValueError( + "Streaming a tied-embedding Cosmos3 Reasoner is not supported because " + "the shared embedding and LM-head tensors may live in different shards." + ) + generator_seen: set[str] = set() + for shard in iter_weight_shards(model_id): + reasoner_package.apply_weights_partial(reasoner_module.preprocess_weights(shard)) + generator_weights = generator_module.preprocess_weight_shard(shard) + generator_seen.update(generator_weights) + generator_package.apply_weights_partial(generator_weights) + + missing_generator = sorted(generator_module.expected_checkpoint_keys() - generator_seen) + if missing_generator: + raise ValueError( + "Cosmos3 unified checkpoint is missing Generator weights required by the " + f"configured architecture: {missing_generator[:16]}" + ) + reasoner_package.finalize_weights() + generator_package.finalize_weights() + + vae_weights = load_component_weights(model_id, "vae") + vae_package.apply_weights(vae_module.preprocess_weights(vae_weights)) + + if audio_package is not None and audio_module is not None: + audio_weights = load_component_weights(model_id, "sound_tokenizer") + audio_package.apply_weights(audio_module.preprocess_weights(audio_weights)) + + reasoner_package.validate_weights() + generator_package.validate_weights() + vae_package.validate_weights() + if audio_package is not None: + audio_package.validate_weights() + + +def _connect_same_named_ports( + builder: PipelineBuilder, + source_name: str, + source_model: ir.Model, + target_name: str, + target_model: ir.Model, + connected_targets: set[str], +) -> None: + outputs = {value.name for value in source_model.graph.outputs} + for value in target_model.graph.inputs: + if value.name in outputs: + connection = builder.connect( + f"{source_name}.{value.name}", + f"{target_name}.{value.name}", + ) + connected_targets.add(connection.target.qualified) + + +def _component_metadata(dtype: ir.DataType, **values: Any) -> dict[str, Any]: + return {"dtype": dtype.name, **values} + + +def _compose_pipeline( + *, + pipeline_config: WorldModelPipelineConfig, + reasoner_package: ModelPackage, + generator_package: ModelPackage, + vae_package: ModelPackage, + audio_package: ModelPackage | None, + generator_config: Cosmos3OmniGeneratorConfig, + vae_config: WanVAEConfig, + scheduler_config: Mapping[str, Any], + assets: Mapping[str, tuple[str, bool]], + reasoner_architecture: str = "qwen3_vl", + default_action_domain: str = "no_action", +) -> PipelinePackage: + """Compose already-built component graphs into the complete topology. + + *pipeline_config* carries the model-agnostic identity, build, and runtime + settings; the remaining arguments are the Cosmos3-specific packages, + architecture configs, and action-domain default. + """ + model_id = pipeline_config.model_id + build_config = pipeline_config.build + generation = pipeline_config.generation + if generator_config.sound_gen and audio_package is None: + raise ValueError( + "A transformer with sound_gen=True requires a sound-tokenizer package." + ) + if vae_config.z_dim != generator_config.latent_channel: + raise ValueError( + "Cosmos3 video VAE latent width must match the Generator vision head: " + f"{vae_config.z_dim} != {generator_config.latent_channel}." + ) + if audio_package is not None: + audio_config = audio_package.config + if audio_config.vocoder_input_dim != generator_config.sound_dim: + raise ValueError( + "Cosmos3 AVAE latent width must match the Generator Sound head: " + f"{audio_config.vocoder_input_dim} != {generator_config.sound_dim}." + ) + if generator_config.action_gen: + assert generator_config.action_dim is not None + if max(_ACTION_DOMAIN_IDS.values()) >= generator_config.num_embodiment_domains: + raise ValueError("Cosmos3 action domain metadata exceeds num_embodiment_domains.") + oversized = { + name: width + for name, width in _ACTION_RAW_DIMS.items() + if width > generator_config.action_dim + } + if oversized: + raise ValueError(f"Cosmos3 raw action widths exceed action_dim: {oversized}.") + if default_action_domain not in _ACTION_DOMAIN_IDS: + raise ValueError( + f"Unknown default Cosmos3 action domain {default_action_domain!r}." + ) + builder = PipelineBuilder() + component_configs: dict[str, object] = {} + models: dict[str, ir.Model] = {} + + for key, name in _REASONER_NAMES.items(): + if key not in reasoner_package: + continue + model = reasoner_package[key] + role = {"decoder": "decoder", "vision_encoder": "encoder", "embedding": "embedding"}[ + key + ] + builder.add_model( + name, + model, + role=role, + source=model_id, + config=_component_metadata( + reasoner_package.config.dtype, + subsystem="reasoner", + architecture=reasoner_architecture, + ), + preferred_execution_providers=build_config.preferred_execution_providers( + reasoner_package.config.dtype, + ), + parameter_dtype=reasoner_package.config.dtype.name, + ) + models[name] = model + component_configs[name] = reasoner_package.config + + generator = generator_package["model"] + builder.add_model( + _GENERATOR_NAME, + generator, + role="dynamics", + source=model_id, + config=_component_metadata( + generator_config.dtype, + subsystem="generator", + sound_gen=generator_config.sound_gen, + action_gen=generator_config.action_gen, + patch_latent_dim=generator_config.patch_latent_dim, + ), + preferred_execution_providers=build_config.preferred_execution_providers( + generator_config.dtype, + ), + parameter_dtype=generator_config.dtype.name, + ) + models[_GENERATOR_NAME] = generator + component_configs[_GENERATOR_NAME] = generator_config + + video_names = {"encoder": _VIDEO_ENCODER_NAME, "decoder": _VIDEO_DECODER_NAME} + for key, name in video_names.items(): + model = vae_package[key] + builder.add_model( + name, + model, + role=key, + presence="video_conditioning" if key == "encoder" else None, + source=model_id, + config=_component_metadata( + vae_config.dtype, + subsystem="video_vae", + spatial_compression=vae_config.scale_factor_spatial, + temporal_compression=vae_config.scale_factor_temporal, + ), + preferred_execution_providers=build_config.preferred_execution_providers( + vae_config.dtype, + ), + parameter_dtype=vae_config.dtype.name, + ) + models[name] = model + component_configs[name] = vae_config + + audio_config: Cosmos3AudioConfig | None = None + if audio_package is not None: + audio_config = audio_package.config + audio_names = {"decoder": _AUDIO_DECODER_NAME} + if "encoder" in audio_package: + audio_names["encoder"] = _AUDIO_ENCODER_NAME + for key, name in audio_names.items(): + model = audio_package[key] + builder.add_model( + name, + model, + role=key, + presence="audio_conditioning" if key == "encoder" else None, + source=model_id, + config=_component_metadata( + audio_config.dtype, + subsystem="sound_tokenizer", + sample_rate=audio_config.sampling_rate, + hop_size=audio_config.resolved_hop_size, + ), + preferred_execution_providers=build_config.preferred_execution_providers( + audio_config.dtype, + ), + parameter_dtype=audio_config.dtype.name, + ) + models[name] = model + component_configs[name] = audio_config + + initial_targets: set[str] = set() + recurrent_targets: set[str] = set() + if _REASONER_NAMES["vision_encoder"] in models: + _connect_same_named_ports( + builder, + _REASONER_NAMES["vision_encoder"], + models[_REASONER_NAMES["vision_encoder"]], + _REASONER_NAMES["embedding"], + models[_REASONER_NAMES["embedding"]], + initial_targets, + ) + _connect_same_named_ports( + builder, + _REASONER_NAMES["embedding"], + models[_REASONER_NAMES["embedding"]], + _REASONER_NAMES["decoder"], + models[_REASONER_NAMES["decoder"]], + initial_targets, + ) + + reasoner_decoder = models[_REASONER_NAMES["decoder"]] + decoder_inputs = {value.name for value in reasoner_decoder.graph.inputs} + state_specs: list[dict[str, Any]] = [] + for output in reasoner_decoder.graph.outputs: + if not output.name.startswith("present."): + continue + cache_input = f"past_key_values.{output.name.removeprefix('present.')}" + if cache_input in decoder_inputs: + connection = builder.connect( + f"{_REASONER_NAMES['decoder']}.{output.name}", + f"{_REASONER_NAMES['decoder']}.{cache_input}", + recurrent=True, + ) + recurrent_targets.add(connection.target.qualified) + cache_suffix = output.name.removeprefix("present.").replace(".", "_") + state_specs.append( + { + "name": f"reasoner_kv_{cache_suffix}", + "kind": "kv_cache", + "input": connection.target, + "output": connection.source, + "lifetime": "sequence", + "release_after": "reasoner_decode", + "sequence_axis": 2, + "metadata": { + "update": "append", + "initializer": "empty_tensor", + }, + } + ) + + iterative_state_inputs = ["vision_tokens"] + for output_name, input_name in ( + ("vision_pred", "vision_tokens"), + ("sound_pred", "sound_tokens"), + ("action_pred", "action_tokens"), + ): + if output_name not in {value.name for value in generator.graph.outputs}: + continue + state_name = input_name.removesuffix("_tokens") + "_state" + connection = builder.connect( + f"{_GENERATOR_NAME}.{output_name}", + f"{_GENERATOR_NAME}.{input_name}", + recurrent=True, + transform="scheduler_step", + parameters={ + "scheduler_asset": "scheduler/scheduler_config.json", + "stage": "world_generation", + "state": state_name, + "timestep_input": f"{_GENERATOR_NAME}.{input_name.removesuffix('_tokens')}_timesteps", + }, + ) + recurrent_targets.add(connection.target.qualified) + state_kind = "action_state" if input_name == "action_tokens" else "diffusion_latent" + release_after = { + "vision_tokens": "decode_video", + "sound_tokens": "decode_audio", + "action_tokens": "world_generation", + }[input_name] + state_specs.append( + { + "name": state_name, + "kind": state_kind, + "input": connection.target, + "output": connection.source, + "lifetime": "request", + "release_after": release_after, + "sequence_axis": 0, + "metadata": { + "update": "scheduler_step", + "scheduler_asset": "scheduler/scheduler_config.json", + }, + } + ) + if input_name != "vision_tokens": + iterative_state_inputs.append(input_name) + + video_final = builder.connect( + f"{_GENERATOR_NAME}.vision_pred", + f"{_VIDEO_DECODER_NAME}.latent", + transform=_VIDEO_DIFFUSION_FINALIZE, + context=(f"{_GENERATOR_NAME}.vision_tokens",), + parameters={ + "scheduler_asset": "scheduler/scheduler_config.json", + "state": "vision_state", + "spatial_patch_size": generator_config.latent_patch_size, + "latent_channels": generator_config.latent_channel, + "input_layout": "packed_tokens", + "output_layout": "BCTHW", + "source_dtype": generator_config.dtype.name, + "target_dtype": vae_config.dtype.name, + }, + ) + initial_targets.add(video_final.target.qualified) + + if audio_package is not None: + if "sound_pred" not in {value.name for value in generator.graph.outputs}: + raise ValueError( + "A sound tokenizer is present, but the transformer has no Sound output." + ) + audio_final = builder.connect( + f"{_GENERATOR_NAME}.sound_pred", + f"{_AUDIO_DECODER_NAME}.latents", + transform=_AUDIO_DIFFUSION_FINALIZE, + context=(f"{_GENERATOR_NAME}.sound_tokens",), + parameters={ + "scheduler_asset": "scheduler/scheduler_config.json", + "state": "sound_state", + "input_layout": "TC", + "output_layout": "BCT", + "source_dtype": generator_config.dtype.name, + "target_dtype": audio_config.dtype.name, + }, + ) + initial_targets.add(audio_final.target.qualified) + + generated_names = { + "attention_mask", + "position_ids", + "cache_position", + "text_indexes", + "und_len", + "vision_sequence_indexes", + "vision_timesteps", + "vision_timestep_token_indexes", + "vision_mse_loss_indexes", + "sound_sequence_indexes", + "sound_timesteps", + "sound_timestep_token_indexes", + "sound_mse_loss_indexes", + "action_domain_ids", + "action_sequence_indexes", + "action_timesteps", + "action_timestep_token_indexes", + "action_mse_loss_indexes", + "action_pred_domain_ids", + } + + def generated_program( + component_name: str, + input_name: str, + ) -> tuple[str, dict[str, Any], str]: + if input_name == "attention_mask": + return ( + "causal_attention_mask", + { + "sequence_input": f"{component_name}.inputs_embeds", + "past_state": [ + state["name"] for state in state_specs if state["kind"] == "kv_cache" + ], + "visible_value": 1, + "masked_value": 0, + }, + "attention.causal_mask", + ) + if input_name == "position_ids": + source = ( + f"{_GENERATOR_NAME}.input_ids" + if component_name == _GENERATOR_NAME + else f"{_REASONER_NAMES['embedding']}.input_ids" + ) + sections = ( + generator_config.rope_axes_dim + if component_name == _GENERATOR_NAME + else ( + getattr(reasoner_package.config, "mrope_section", None) + or generator_config.rope_axes_dim + ) + ) + parameters: dict[str, Any] = { + "source": source, + "axes": 3, + "mrope_sections": list(sections), + } + if component_name == _GENERATOR_NAME: + parameters.update( + { + "temporal_margin": ( + generator_config.unified_3d_mrope_temporal_modality_margin + ), + "reset_spatial": (generator_config.unified_3d_mrope_reset_spatial_ids), + } + ) + else: + parameters["past_state"] = [ + state["name"] for state in state_specs if state["kind"] == "kv_cache" + ] + return ("multimodal_position_ids", parameters, "position.multimodal") + if input_name.endswith("_timesteps"): + modality = input_name.removesuffix("_timesteps") + return ( + "scheduler_timesteps", + {"stage": "world_generation", "modality": modality}, + f"diffusion.{modality}.timesteps", + ) + if input_name in {"action_domain_ids", "action_pred_domain_ids"}: + return ( + "action_domain_ids", + { + "domain_input": "action_domain", + "default": default_action_domain, + "domain_map": _ACTION_DOMAIN_IDS, + "padded_dimension": generator_config.action_dim, + }, + "action.domain_ids", + ) + modality = input_name.split("_", 1)[0] + source = ( + f"{_GENERATOR_NAME}.input_ids" + if modality in {"text", "und"} + else f"{_GENERATOR_NAME}.{modality}_tokens" + ) + return ( + "packed_sequence_layout", + { + "modality": modality, + "source": source, + "layout": "flat_token_rows", + "understanding_prefix": input_name == "und_len", + "index_kind": input_name, + }, + f"packing.{input_name}", + ) + + def external_semantic(component_name: str, input_name: str) -> str: + if input_name == "input_ids": + return "text.token_ids" + if input_name == "pixel_values": + return "vision.pixel_values" + if input_name == "grid_thw": + return "vision.grid_thw" + if input_name == "image_features": + return "vision.image_features" + if input_name == "video_features": + return "vision.video_features" + if input_name == "audio": + return "audio.waveform" + if input_name == "sample": + return "video.frames" + if input_name.endswith("_tokens") and component_name == _GENERATOR_NAME: + return f"diffusion.initial_{input_name.removesuffix('_tokens')}_latent" + return f"tensor.{component_name}.{input_name}" + + for component_name, model in models.items(): + for value in model.graph.inputs: + endpoint = f"{component_name}.{value.name}" + if endpoint in initial_targets: + continue + if endpoint in recurrent_targets: + if component_name == _GENERATOR_NAME: + alias = f"initial_{value.name}".replace(".", "_") + builder.declare_external( + endpoint, + alias=alias, + semantic=external_semantic(component_name, value.name), + required=True, + ) + else: + cache_axis = value.shape[2] + cache_axis_name = getattr(cache_axis, "value", None) + if not isinstance(cache_axis_name, str): + raise ValueError( + f"KV-cache input {endpoint!r} must expose a symbolic " + "sequence axis at index 2." + ) + builder.declare_generated( + endpoint, + generator="empty_tensor", + parameters={ + "dynamic_axes": {cache_axis_name: 0}, + "fill": 0, + }, + semantic=( + "kv_cache.key" if value.name.endswith(".key") else "kv_cache.value" + ), + ) + continue + if value.name.startswith(("past_key_values.", "key_cache.", "value_cache.")): + builder.declare_stateful(endpoint, semantic=f"state.{value.name}") + elif value.name in generated_names: + generator, parameters, semantic = generated_program( + component_name, + value.name, + ) + builder.declare_generated( + endpoint, + generator=generator, + parameters=parameters, + semantic=semantic, + ) + else: + alias = f"{component_name}_{value.name}".replace(".", "_") + presence = None + required = True + if component_name == _VIDEO_ENCODER_NAME: + presence = "video_conditioning" + required = False + elif component_name == _AUDIO_ENCODER_NAME: + presence = "audio_conditioning" + required = False + elif ( + component_name == _REASONER_NAMES["embedding"] + and value.name == "video_features" + ): + # Video frames reuse the Reasoner vision encoder; the host + # runs it once per visual item and routes the projected + # features to the placeholder id of that item's modality. + presence = _VIDEO_UNDERSTANDING + required = False + elif component_name == _REASONER_NAMES["vision_encoder"]: + presence = None + required = True + builder.declare_external( + endpoint, + alias=alias, + semantic=external_semantic(component_name, value.name), + required=required, + presence=presence, + ) + + reasoner_prompt_components = [_REASONER_NAMES["embedding"]] + if _REASONER_NAMES["vision_encoder"] in models: + reasoner_prompt_components.insert(0, _REASONER_NAMES["vision_encoder"]) + builder.add_stage( + "reasoner_prompt", + "single_pass", + reasoner_prompt_components, + run_on="prefill", + ) + builder.add_stage( + "reasoner_decode", + "autoregressive", + [_REASONER_NAMES["embedding"], _REASONER_NAMES["decoder"]], + run_on="decode", + options={ + "tokenizer_asset": ( + "tokenizer.json" + if "tokenizer.json" in assets + else "text_tokenizer/tokenizer.json" + ), + "sampling": generation.sampling_manifest(), + "stop": generation.stop_manifest( + max_sequence_length=getattr( + reasoner_package.config, + "max_position_embeddings", + None, + ), + ), + "max_tokens": generation.max_tokens_manifest( + limit=getattr( + reasoner_package.config, + "max_position_embeddings", + None, + ), + ), + "state_names": [ + state["name"] for state in state_specs if state["kind"] == "kv_cache" + ], + }, + ) + builder.add_stage( + "world_generation", + "iterative", + [_GENERATOR_NAME], + run_on="step", + capabilities=("classifier_free_guidance", "conditioned_diffusion"), + options={ + "scheduler": { + "kind": scheduler_config.get("_class_name"), + "config_asset": "scheduler/scheduler_config.json", + "overrideable": [ + "num_inference_steps", + "guidance_scale", + "flow_shift", + "use_karras_sigmas", + ], + "mode_overrides": generation.scheduler_mode_overrides_manifest(), + }, + "guidance": { + "kind": "classifier_free", + "conditioning_input": f"{_GENERATOR_NAME}.input_ids", + "scale_option": "guidance_scale", + "default_scale": 1.0, + "combine": "unconditional + scale * (conditional - unconditional)", + }, + "conditioning": { + "vision": { + "encoder_stage": "encode_video", + "encoder_input": f"{_VIDEO_ENCODER_NAME}.sample", + "encoder_output": f"{_VIDEO_ENCODER_NAME}.latent", + "state": "vision_state", + "conditioned_latent_frames_option": "vision_conditioned_latent_frames", + "default_conditioned_latent_frames": [], + "preprocessing": { + "resize": "bilinear", + "normalize": { + "mean": [0.5, 0.5, 0.5], + "std": [0.5, 0.5, 0.5], + }, + }, + "packing": { + "spatial_patch_size": generator_config.latent_patch_size, + "temporal_patch_size": 1, + "input_layout": "BCTHW", + "output_layout": "NC", + "channel_order": "patch_height_patch_width_channel", + }, + }, + }, + "default_steps": generation.default_inference_steps, + "timestep": { + "generator": "scheduler_timesteps", + "scale": generator_config.timestep_scale, + }, + "prediction_type": scheduler_config.get("prediction_type", "flow_prediction"), + "state_inputs": [f"{_GENERATOR_NAME}.{name}" for name in iterative_state_inputs], + "packed_modalities": True, + }, + ) + builder.add_stage( + "encode_video", + "on_demand", + [_VIDEO_ENCODER_NAME], + run_on="on_demand", + options={"presence": "video_conditioning"}, + ) + builder.add_stage( + "decode_video", + "single_pass", + [_VIDEO_DECODER_NAME], + run_on="finalize", + ) + + if audio_package is not None: + if "encoder" in audio_package: + builder.add_stage( + "encode_audio", + "on_demand", + [_AUDIO_ENCODER_NAME], + run_on="on_demand", + options={"presence": "audio_conditioning"}, + ) + builder.add_stage( + "decode_audio", + "single_pass", + [_AUDIO_DECODER_NAME], + run_on="finalize", + ) + + for state in state_specs: + builder.add_state(**state) + + builder.add_public_output(f"{_REASONER_NAMES['decoder']}.logits", alias="logits") + builder.add_public_output(f"{_VIDEO_DECODER_NAME}.sample", alias="video") + builder.add_public_output(f"{_GENERATOR_NAME}.vision_pred", alias="vision_velocity") + builder.add_public_output(f"{_VIDEO_ENCODER_NAME}.latent", alias="encoded_video_latent") + if generator_config.action_gen: + builder.add_public_output( + f"{_GENERATOR_NAME}.action_pred", + alias="action_velocity", + ) + builder.add_public_state_output("action_state", alias="action") + if audio_package is not None: + builder.add_public_output(f"{_AUDIO_DECODER_NAME}.waveform", alias="sound") + if "encoder" in audio_package: + builder.add_public_output( + f"{_AUDIO_ENCODER_NAME}.latent_mean", + alias="encoded_audio_latent", + ) + + for destination, (source, required) in assets.items(): + builder.add_asset(destination, source, required=required) + builder.set_profile(pipeline_config.profile_name, pipeline_config.profile_version) + for key, value in pipeline_config.manifest_metadata().items(): + builder.set_metadata(key, value) + builder.set_metadata( + "modalities", + { + "vision": True, + "sound": generator_config.sound_gen, + "action": generator_config.action_gen, + }, + ) + builder.set_metadata( + "packing", + { + "generator_boundary": "packed_tokens", + "latent_patch_size": generator_config.latent_patch_size, + "patch_latent_dim": generator_config.patch_latent_dim, + }, + ) + builder.set_metadata( + "conditioning_handoffs", + { + "video": { + "from": f"{_VIDEO_ENCODER_NAME}.latent", + "to": f"{_GENERATOR_NAME}.vision_tokens", + "transform": "patchify", + "parameters": { + "spatial_patch_size": generator_config.latent_patch_size, + "temporal_patch_size": 1, + "input_layout": "BCTHW", + "output_layout": "NC", + "channel_order": "patch_height_patch_width_channel", + }, + "optional": True, + }, + "audio": ( + { + "from": f"{_AUDIO_ENCODER_NAME}.latent_mean", + "to": f"{_GENERATOR_NAME}.sound_tokens", + "transform": "reshape", + "parameters": { + "input_layout": "BCT", + "output_layout": "TC", + }, + "optional": True, + } + if audio_package is not None and "encoder" in audio_package + else None + ), + }, + ) + builder.set_metadata( + "shared_parameters", + { + "understanding_expert": [ + _REASONER_NAMES["decoder"], + _GENERATOR_NAME, + ] + }, + ) + builder.set_metadata( + "prompt_inputs", + { + "reasoner": f"{_REASONER_NAMES['embedding']}.input_ids", + "generator": f"{_GENERATOR_NAME}.input_ids", + "relationship": "same_tokenized_prompt_repacked_for_generator", + }, + ) + if generator_config.action_gen: + builder.set_metadata( + "action", + { + "modes": ["policy", "forward_dynamics", "inverse_dynamics"], + "domain_ids": _ACTION_DOMAIN_IDS, + "raw_dimensions": _ACTION_RAW_DIMS, + "padded_dimension": generator_config.action_dim, + "input_padding": { + "side": "right", + "value": 0.0, + "target_dimension": generator_config.action_dim, + }, + "output_slicing": "raw_dimensions", + "clipping": None, + "resolution_tiers": [256, 480, 704, 720], + }, + ) + for key, value in pipeline_config.extra_metadata.items(): + builder.set_metadata(key, value) + + return builder.build( + config=reasoner_package.config, + component_configs=component_configs, + ) + + +def _collect_assets( + model_id: str, + *, + has_sound_tokenizer: bool, +) -> dict[str, tuple[str, bool]]: + """Resolve the runtime assets a Cosmos3 package must ship.""" + candidates = list(_ASSET_CANDIDATES) + if has_sound_tokenizer: + candidates.append(("sound_tokenizer/config.json", True)) + assets = resolve_assets(model_id, candidates) + tokenizer_paths = ("tokenizer.json", "text_tokenizer/tokenizer.json") + available_tokenizers = [path for path in tokenizer_paths if path in assets] + if not available_tokenizers: + raise FileNotFoundError( + "Cosmos3 checkpoint must provide tokenizer.json either at the root " + "or under text_tokenizer/." + ) + preferred = available_tokenizers[0] + assets[preferred] = (assets[preferred][0], True) + return assets + + +def build_cosmos3_world_model( + model_id: str, + *, + dtype: str | ir.DataType | None = None, + load_weights: bool = True, + execution_provider: str = "default", + trace_optimization: bool = False, + **_options: Any, +) -> PipelinePackage: + """Build the complete neural Cosmos3-Omni world-model package.""" + build_config = WorldModelBuildConfig( + dtype=dtype, + load_weights=load_weights, + execution_provider=execution_provider, + trace_optimization=trace_optimization, + ) + pipeline_index, _ = load_checkpoint_json(model_id, "model_index.json") + root_config, _ = load_checkpoint_json(model_id, "config.json") + text_config = root_config.get("text_config") or {} + if ( + isinstance(text_config, Mapping) + and text_config.get("model_type") == "cosmos3_edge_text" + ): + from mobius._cosmos3_edge_world_model import build_cosmos3_edge_world_model + + return build_cosmos3_edge_world_model( + model_id, + dtype=dtype, + load_weights=load_weights, + execution_provider=execution_provider, + trace_optimization=trace_optimization, + **_options, + ) + transformer_class = component_class(pipeline_index, "transformer") + vae_class = component_class(pipeline_index, "vae") + if transformer_class != "Cosmos3OmniTransformer": + raise ValueError(f"Unsupported Cosmos3 transformer class {transformer_class!r}") + if vae_class != "AutoencoderKLWan": + raise ValueError(f"Unsupported Cosmos3 VAE class {vae_class!r}") + + transformer_config_dict, _ = load_checkpoint_json(model_id, "transformer/config.json") + vae_config_dict, _ = load_checkpoint_json(model_id, "vae/config.json") + scheduler_config, _ = load_checkpoint_json(model_id, "scheduler/scheduler_config.json") + generation_config = WorldModelGenerationConfig.from_generation_config( + load_optional_checkpoint_json(model_id, "generation_config.json"), + # Distilled 4-step checkpoints advertise their step budget only in the + # repository name. + default_inference_steps=4 if "4Step" in model_id else 35, + ) + checkpoint_config = load_optional_checkpoint_json(model_id, "checkpoint.json") + policy_config = checkpoint_config.get("policy") + default_action_domain = ( + policy_config.get("domain_name", "no_action") + if isinstance(policy_config, Mapping) + else "no_action" + ) + + has_sound = component_class(pipeline_index, "sound_tokenizer") is not None + has_reasoner_vision = component_class(pipeline_index, "vision_encoder") is not None + audio_config_dict: dict[str, Any] | None = None + audio_weight_names: set[str] | None = None + if has_sound: + audio_config_dict, _ = load_checkpoint_json(model_id, "sound_tokenizer/config.json") + audio_weight_names = component_weight_names(model_id, "sound_tokenizer") + + ( + reasoner_package, + reasoner_module, + generator_package, + generator_module, + vae_package, + vae_module, + audio_package, + audio_module, + ) = _build_components( + model_id, + build_config=build_config, + pipeline_index=pipeline_index, + transformer_config_dict=transformer_config_dict, + vae_config_dict=vae_config_dict, + audio_config_dict=audio_config_dict, + audio_weight_names=audio_weight_names, + has_reasoner_vision=has_reasoner_vision, + ) + + if build_config.load_weights: + _apply_checkpoint_weights( + model_id, + reasoner_package=reasoner_package, + reasoner_module=reasoner_module, + generator_package=generator_package, + generator_module=generator_module, + vae_package=vae_package, + vae_module=vae_module, + audio_package=audio_package, + audio_module=audio_module, + ) + + assets = _collect_assets(model_id, has_sound_tokenizer=has_sound) + return _compose_pipeline( + pipeline_config=WorldModelPipelineConfig( + model_id=model_id, + model_type="cosmos3_omni", + build=build_config, + generation=generation_config, + ), + reasoner_package=reasoner_package, + generator_package=generator_package, + vae_package=vae_package, + audio_package=audio_package, + generator_config=generator_module.config, + vae_config=vae_module.config, + scheduler_config=scheduler_config, + assets=assets, + default_action_domain=default_action_domain, + ) diff --git a/src/mobius/_cosmos3_world_model_test.py b/src/mobius/_cosmos3_world_model_test.py new file mode 100644 index 000000000..3867ad5b6 --- /dev/null +++ b/src/mobius/_cosmos3_world_model_test.py @@ -0,0 +1,661 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import dataclasses +import json +from types import SimpleNamespace +from unittest import mock + +import onnx_ir as ir +import pytest +import safetensors.torch +import torch + +from mobius._configs import Cosmos3AudioConfig, Cosmos3OmniGeneratorConfig, WanVAEConfig +from mobius._cosmos3_world_model import _compose_pipeline, build_cosmos3_world_model +from mobius._diffusers_checkpoint import component_weight_names +from mobius._model_package import ModelPackage +from mobius._world_model_config import ( + WorldModelBuildConfig, + WorldModelGenerationConfig, + WorldModelPipelineConfig, +) + + +def _pipeline_config( + model_id: str = "example/world", + **overrides, +) -> WorldModelPipelineConfig: + return WorldModelPipelineConfig( + model_id=model_id, + model_type=overrides.pop("model_type", "cosmos3_omni"), + **overrides, + ) + + +def _value( + name: str, + dtype: ir.DataType, + shape: list[int | str], +) -> ir.Value: + return ir.Value(name=name, type=ir.TensorType(dtype), shape=ir.Shape(shape)) + + +def _model( + inputs: dict[str, tuple[ir.DataType, list[int | str]]], + outputs: dict[str, tuple[ir.DataType, list[int | str]]], +) -> ir.Model: + input_values = [_value(name, dtype, shape) for name, (dtype, shape) in inputs.items()] + nodes: list[ir.Node] = [] + output_values: list[ir.Value] = [] + for name, (dtype, shape) in outputs.items(): + node = ir.Node("", "Identity", inputs=[input_values[0]], num_outputs=1) + output = node.outputs[0] + output.name = name + output.type = ir.TensorType(dtype) + output.shape = ir.Shape(shape) + nodes.append(node) + output_values.append(output) + graph = ir.Graph( + input_values, + output_values, + nodes=nodes, + name="component", + opset_imports={"": 24}, + ) + return ir.Model(graph, ir_version=10) + + +def _generator_config(*, sound: bool = True, action: bool = True): + return Cosmos3OmniGeneratorConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=8, + hidden_act="silu", + rope_axes_dim=(2, 1, 1), + latent_channel=16, + latent_patch_size=2, + patch_latent_dim=64, + sound_gen=sound, + sound_dim=4 if sound else None, + action_gen=action, + action_dim=64 if action else None, + max_action_dim=64 if action else None, + num_embodiment_domains=32, + dtype=ir.DataType.FLOAT, + ) + + +def _packages(*, sound: bool = True, action: bool = True, audio_encoder: bool = True): + f32 = ir.DataType.FLOAT + i64 = ir.DataType.INT64 + reasoner_config = SimpleNamespace(dtype=f32) + reasoner = ModelPackage( + { + "vision_encoder": _model( + {"pixel_values": (f32, ["b", 3, 8, 8])}, + {"image_features": (f32, ["b", "n", 16])}, + ), + "embedding": _model( + { + "input_ids": (i64, ["b", "s"]), + "image_features": (f32, ["b", "n", 16]), + }, + {"inputs_embeds": (f32, ["b", "s", 16])}, + ), + "decoder": _model( + { + "inputs_embeds": (f32, ["b", "s", 16]), + "attention_mask": (i64, ["b", "s"]), + "position_ids": (i64, [3, "b", "s"]), + "past_key_values.0.key": (f32, ["b", 1, "p", 8]), + "past_key_values.0.value": (f32, ["b", 1, "p", 8]), + }, + { + "logits": (f32, ["b", "s", 32]), + "present.0.key": (f32, ["b", 1, "total", 8]), + "present.0.value": (f32, ["b", 1, "total", 8]), + }, + ), + }, + config=reasoner_config, + ) + + generator_inputs = { + "input_ids": (i64, ["text"]), + "position_ids": (i64, [3, "sequence"]), + "vision_tokens": (f32, ["vision", 64]), + "vision_timesteps": (f32, ["vision"]), + } + generator_outputs = {"vision_pred": (f32, ["vision", 64])} + if sound: + generator_inputs["sound_tokens"] = (f32, ["sound", 4]) + generator_inputs["sound_timesteps"] = (f32, ["sound"]) + generator_outputs["sound_pred"] = (f32, ["sound", 4]) + if action: + generator_inputs["action_tokens"] = (f32, ["action", 64]) + generator_inputs["action_timesteps"] = (f32, ["action"]) + generator_outputs["action_pred"] = (f32, ["action", 64]) + generator = ModelPackage( + {"model": _model(generator_inputs, generator_outputs)}, + config=_generator_config(sound=sound, action=action), + ) + + vae_config = WanVAEConfig() + vae = ModelPackage( + { + "encoder": _model( + {"sample": (f32, ["b", 3, "t", "h", "w"])}, + {"latent": (f32, ["b", 16, "lt", "lh", "lw"])}, + ), + "decoder": _model( + {"latent": (f32, ["b", 16, "lt", "lh", "lw"])}, + {"sample": (f32, ["b", 3, "t", "h", "w"])}, + ), + }, + config=vae_config, + ) + + audio = None + if sound: + audio_config = Cosmos3AudioConfig( + dec_dim=4, + dec_c_mults=(1,), + dec_strides=(2,), + enc_dim=4, + enc_intermediate_dim=16, + enc_num_layers=1, + enc_num_blocks=1, + enc_n_fft=4, + enc_hop_length=2, + enc_c_mults=(1,), + enc_strides=(1,), + enc_latent_dim=8, + vocoder_input_dim=4, + hop_size=2, + encoder_enabled=audio_encoder, + ) + audio_models = { + "decoder": _model( + {"latents": (f32, ["b", 4, "t"])}, + {"waveform": (f32, ["b", 2, "samples"])}, + ) + } + if audio_encoder: + audio_models["encoder"] = _model( + {"audio": (f32, ["b", 2, "samples"])}, + {"latent_mean": (f32, ["b", 4, "t"])}, + ) + audio = ModelPackage(audio_models, config=audio_config) + return reasoner, generator, vae, audio + + +def _runtime_assets(tmp_path): + tokenizer = tmp_path / "tokenizer.json" + tokenizer.write_text("{}", encoding="utf-8") + scheduler_dir = tmp_path / "scheduler" + scheduler_dir.mkdir(exist_ok=True) + scheduler = scheduler_dir / "scheduler_config.json" + scheduler.write_text("{}", encoding="utf-8") + return { + "tokenizer.json": (str(tokenizer), True), + "scheduler/scheduler_config.json": (str(scheduler), True), + } + + +def test_compose_complete_cosmos3_pipeline(tmp_path) -> None: + reasoner, generator, vae, audio = _packages() + config = generator.config + package = _compose_pipeline( + pipeline_config=_pipeline_config("nvidia/Cosmos3-Nano"), + reasoner_package=reasoner, + generator_package=generator, + vae_package=vae, + audio_package=audio, + generator_config=config, + vae_config=vae.config, + scheduler_config={"prediction_type": "flow_prediction"}, + assets=_runtime_assets(tmp_path), + ) + + assert set(package) == { + "reasoner_decoder", + "reasoner_vision_encoder", + "reasoner_embedding", + "generator", + "video_encoder", + "video_decoder", + "audio_encoder", + "audio_decoder", + } + recurrent = [edge for edge in package.manifest.connections if edge.recurrent] + assert {edge.target.qualified for edge in recurrent} >= { + "generator.vision_tokens", + "generator.sound_tokens", + "generator.action_tokens", + "reasoner_decoder.past_key_values.0.key", + "reasoner_decoder.past_key_values.0.value", + } + assert "iterative_scheduler" in package.manifest.required_capabilities + assert {output.name for output in package.manifest.outputs} >= { + "logits", + "video", + "sound", + "action", + } + assert package.manifest.metadata["profile"] == "world-model" + assert package.manifest.metadata["action"]["domain_ids"]["droid_lerobot"] == 8 + assert package.manifest.metadata["action"]["raw_dimensions"]["hand_pose"] == 57 + assert package.manifest.profile is not None + assert package.manifest.profile.name == "cosmos3-omni" + assert {state.kind for state in package.manifest.states} >= { + "kv_cache", + "diffusion_latent", + "action_state", + } + position_input = next( + value + for value in package.manifest.inputs + if value.port.qualified == "generator.position_ids" + ) + assert position_input.generator is not None + assert position_input.generator.kind == "multimodal_position_ids" + cache_input = next( + value + for value in package.manifest.inputs + if value.port.qualified == "reasoner_decoder.past_key_values.0.key" + ) + assert cache_input.generator is not None + assert cache_input.generator.parameters["dynamic_axes"] == {"p": 0} + assert package.manifest.component("generator").parameter_dtype == "FLOAT" + assert package.manifest.component("generator").preferred_execution_providers[0] == ("cuda") + + +def test_compose_without_optional_sound_or_action(tmp_path) -> None: + reasoner, generator, vae, audio = _packages(sound=False, action=False) + package = _compose_pipeline( + pipeline_config=_pipeline_config(), + reasoner_package=reasoner, + generator_package=generator, + vae_package=vae, + audio_package=audio, + generator_config=generator.config, + vae_config=vae.config, + scheduler_config={}, + assets=_runtime_assets(tmp_path), + ) + + assert "audio_decoder" not in package + assert {output.name for output in package.manifest.outputs}.isdisjoint({"sound", "action"}) + + +def test_sound_head_requires_sound_tokenizer() -> None: + reasoner, generator, vae, _audio = _packages(sound=True) + + with pytest.raises(ValueError, match="sound_gen=True"): + _compose_pipeline( + pipeline_config=_pipeline_config(), + reasoner_package=reasoner, + generator_package=generator, + vae_package=vae, + audio_package=None, + generator_config=generator.config, + vae_config=vae.config, + scheduler_config={}, + assets={}, + ) + + +def test_local_audio_weight_metadata_detects_encoder(tmp_path) -> None: + component = tmp_path / "sound_tokenizer" + component.mkdir() + safetensors.torch.save_file( + { + "encoder.layers.0.weight": torch.zeros(1), + "decoder.conv1.weight": torch.zeros(1), + }, + str(component / "diffusion_pytorch_model.safetensors"), + ) + + names = component_weight_names(str(tmp_path), "sound_tokenizer") + + assert names == {"encoder.layers.0.weight", "decoder.conv1.weight"} + + +def test_decoder_only_audio_package_is_supported(tmp_path) -> None: + reasoner, generator, vae, audio = _packages(audio_encoder=False) + package = _compose_pipeline( + pipeline_config=_pipeline_config(), + reasoner_package=reasoner, + generator_package=generator, + vae_package=vae, + audio_package=audio, + generator_config=generator.config, + vae_config=vae.config, + scheduler_config={}, + assets=_runtime_assets(tmp_path), + ) + + assert "audio_encoder" not in package + assert "audio_decoder" in package + + +def test_reasoner_without_standalone_vision_tower_is_supported(tmp_path) -> None: + reasoner, generator, vae, audio = _packages(sound=False, action=False) + del reasoner["vision_encoder"] + package = _compose_pipeline( + pipeline_config=_pipeline_config("example/distilled-world"), + reasoner_package=reasoner, + generator_package=generator, + vae_package=vae, + audio_package=audio, + generator_config=generator.config, + vae_config=vae.config, + scheduler_config={}, + assets=_runtime_assets(tmp_path), + ) + + assert "reasoner_vision_encoder" not in package + prompt_stage = next( + stage for stage in package.manifest.stages if stage.name == "reasoner_prompt" + ) + assert prompt_stage.components == ("reasoner_embedding",) + + +def test_generator_and_video_vae_latent_width_must_match() -> None: + reasoner, generator, vae, audio = _packages(sound=False, action=False) + incompatible = dataclasses.replace( + generator.config, + latent_channel=8, + patch_latent_dim=32, + ) + + with pytest.raises(ValueError, match="latent width"): + _compose_pipeline( + pipeline_config=_pipeline_config(), + reasoner_package=reasoner, + generator_package=generator, + vae_package=vae, + audio_package=audio, + generator_config=incompatible, + vae_config=vae.config, + scheduler_config={}, + assets={}, + ) + + +def test_manifest_carries_generation_and_pipeline_config(tmp_path) -> None: + reasoner, generator, vae, audio = _packages(sound=False, action=False) + generation = WorldModelGenerationConfig.from_generation_config( + { + "do_sample": True, + "temperature": 0.6, + "top_k": 20, + "top_p": 0.95, + "repetition_penalty": 1.05, + "max_new_tokens": 512, + "eos_token_id": 151645, + }, + default_inference_steps=4, + scheduler_mode_overrides={"action": {"flow_shift": 10.0}}, + ) + package = _compose_pipeline( + pipeline_config=WorldModelPipelineConfig( + model_id="example/world-4Step", + model_type="cosmos3_edge", + build=WorldModelBuildConfig(execution_provider="cuda"), + generation=generation, + extra_metadata={"edge": {"policy": None}}, + ), + reasoner_package=reasoner, + generator_package=generator, + vae_package=vae, + audio_package=audio, + generator_config=generator.config, + vae_config=vae.config, + scheduler_config={"_class_name": "FlowMatchEulerDiscreteScheduler"}, + assets=_runtime_assets(tmp_path), + reasoner_architecture="cosmos3_edge", + ) + + decode = next( + stage for stage in package.manifest.stages if stage.name == "reasoner_decode" + ) + assert decode.options["sampling"] == { + "do_sample": True, + "temperature": 0.6, + "top_k": 20, + "top_p": 0.95, + "repetition_penalty": 1.05, + } + assert decode.options["stop"]["eos_token_ids"] == [151645] + assert decode.options["max_tokens"] == { + "default": 512, + "required_override": False, + "limit": None, + } + generation_stage = next( + stage for stage in package.manifest.stages if stage.name == "world_generation" + ) + assert generation_stage.options["default_steps"] == 4 + assert generation_stage.options["scheduler"]["mode_overrides"] == { + "action": {"flow_shift": 10.0} + } + assert generation_stage.options["guidance"] == { + "kind": "classifier_free", + "conditioning_input": "generator.input_ids", + "scale_option": "guidance_scale", + "default_scale": 1.0, + "combine": "unconditional + scale * (conditional - unconditional)", + } + assert set(generation_stage.capabilities) == { + "classifier_free_guidance", + "conditioned_diffusion", + "loop_carried_state", + } + assert generation_stage.options["conditioning"]["vision"] == { + "encoder_stage": "encode_video", + "encoder_input": "video_encoder.sample", + "encoder_output": "video_encoder.latent", + "state": "vision_state", + "conditioned_latent_frames_option": "vision_conditioned_latent_frames", + "default_conditioned_latent_frames": [], + "preprocessing": { + "resize": "bilinear", + "normalize": { + "mean": [0.5, 0.5, 0.5], + "std": [0.5, 0.5, 0.5], + }, + }, + "packing": { + "spatial_patch_size": 2, + "temporal_patch_size": 1, + "input_layout": "BCTHW", + "output_layout": "NC", + "channel_order": "patch_height_patch_width_channel", + }, + } + assert generation_stage.options["scheduler"]["overrideable"] == [ + "num_inference_steps", + "guidance_scale", + "flow_shift", + "use_karras_sigmas", + ] + video_handoff = package.manifest.metadata["conditioning_handoffs"]["video"] + assert video_handoff["parameters"]["channel_order"] == ("patch_height_patch_width_channel") + assert package.manifest.metadata["profile"] == "world-model" + assert package.manifest.metadata["model_type"] == "cosmos3_edge" + assert package.manifest.metadata["source"] == "example/world-4Step" + assert package.manifest.metadata["edge"] == {"policy": None} + assert package.manifest.profile is not None + assert package.manifest.profile.name == "cosmos3-edge" + assert package.manifest.component("generator").preferred_execution_providers == ("cuda",) + + +def test_missing_generation_config_requires_runtime_token_budget(tmp_path) -> None: + reasoner, generator, vae, audio = _packages(sound=False, action=False) + package = _compose_pipeline( + pipeline_config=_pipeline_config(), + reasoner_package=reasoner, + generator_package=generator, + vae_package=vae, + audio_package=audio, + generator_config=generator.config, + vae_config=vae.config, + scheduler_config={}, + assets=_runtime_assets(tmp_path), + ) + + decode = next( + stage for stage in package.manifest.stages if stage.name == "reasoner_decode" + ) + assert decode.options["sampling"] == { + "do_sample": False, + "temperature": 1.0, + "top_k": 50, + "top_p": 1.0, + "repetition_penalty": 1.0, + } + assert decode.options["stop"] == { + "kind": "token_ids", + "eos_token_ids": [], + "max_sequence_length": None, + } + assert decode.options["max_tokens"]["required_override"] is True + + +def test_build_cosmos3_world_model_threads_shared_configs(tmp_path) -> None: + (tmp_path / "model_index.json").write_text( + json.dumps( + { + "transformer": ["diffusers", "Cosmos3OmniTransformer"], + "vae": ["diffusers", "AutoencoderKLWan"], + "sound_tokenizer": [None, None], + "vision_encoder": ["transformers", "Qwen3VLVisionModel"], + } + ), + encoding="utf-8", + ) + (tmp_path / "config.json").write_text( + json.dumps({"model_type": "cosmos3_omni"}), + encoding="utf-8", + ) + for name in ("transformer", "vae"): + (tmp_path / name).mkdir() + (tmp_path / name / "config.json").write_text("{}", encoding="utf-8") + (tmp_path / "scheduler").mkdir() + (tmp_path / "scheduler" / "scheduler_config.json").write_text("{}", encoding="utf-8") + (tmp_path / "generation_config.json").write_text( + json.dumps({"max_new_tokens": 256, "eos_token_id": 11}), + encoding="utf-8", + ) + (tmp_path / "checkpoint.json").write_text( + json.dumps({"policy": {"domain_name": "droid_lerobot"}}), + encoding="utf-8", + ) + generator_module = SimpleNamespace(config=mock.sentinel.generator_config) + vae_module = SimpleNamespace(config=mock.sentinel.vae_config) + package = mock.sentinel.pipeline_package + + with ( + mock.patch( + "mobius._cosmos3_world_model._build_components", + return_value=( + mock.sentinel.reasoner_package, + mock.sentinel.reasoner_module, + mock.sentinel.generator_package, + generator_module, + mock.sentinel.vae_package, + vae_module, + None, + None, + ), + ) as build_components, + mock.patch( + "mobius._cosmos3_world_model._collect_assets", + return_value={}, + ), + mock.patch( + "mobius._cosmos3_world_model._compose_pipeline", + return_value=package, + ) as compose, + ): + result = build_cosmos3_world_model( + str(tmp_path), + dtype="f16", + load_weights=False, + execution_provider="cuda", + ) + + assert result is package + build_config = build_components.call_args.kwargs["build_config"] + assert build_config.resolved_dtype() is ir.DataType.FLOAT16 + assert build_config.execution_provider == "cuda" + assert build_config.load_weights is False + pipeline_config = compose.call_args.kwargs["pipeline_config"] + assert pipeline_config.model_type == "cosmos3_omni" + assert pipeline_config.model_id == str(tmp_path) + assert pipeline_config.build is build_config + assert pipeline_config.generation.max_new_tokens == 256 + assert pipeline_config.generation.eos_token_ids == (11,) + assert pipeline_config.generation.default_inference_steps == 35 + assert pipeline_config.generation.scheduler_mode_overrides_manifest() == {} + assert compose.call_args.kwargs["default_action_domain"] == "droid_lerobot" + + +def test_build_cosmos3_world_model_uses_distilled_step_budget(tmp_path) -> None: + checkpoint = tmp_path / "Cosmos3-Omni-4Step" + checkpoint.mkdir() + (checkpoint / "model_index.json").write_text( + json.dumps( + { + "transformer": ["diffusers", "Cosmos3OmniTransformer"], + "vae": ["diffusers", "AutoencoderKLWan"], + } + ), + encoding="utf-8", + ) + (checkpoint / "config.json").write_text( + json.dumps({"model_type": "cosmos3_omni"}), + encoding="utf-8", + ) + for name in ("transformer", "vae"): + (checkpoint / name).mkdir() + (checkpoint / name / "config.json").write_text("{}", encoding="utf-8") + (checkpoint / "scheduler").mkdir() + (checkpoint / "scheduler" / "scheduler_config.json").write_text("{}", encoding="utf-8") + + with ( + mock.patch( + "mobius._cosmos3_world_model._build_components", + return_value=( + mock.sentinel.reasoner_package, + mock.sentinel.reasoner_module, + mock.sentinel.generator_package, + SimpleNamespace(config=mock.sentinel.generator_config), + mock.sentinel.vae_package, + SimpleNamespace(config=mock.sentinel.vae_config), + None, + None, + ), + ), + mock.patch( + "mobius._cosmos3_world_model._collect_assets", + return_value={}, + ), + mock.patch( + "mobius._cosmos3_world_model._compose_pipeline", + return_value=mock.sentinel.pipeline_package, + ) as compose, + ): + build_cosmos3_world_model(str(checkpoint), load_weights=False) + + pipeline_config = compose.call_args.kwargs["pipeline_config"] + assert pipeline_config.generation.default_inference_steps == 4 + assert compose.call_args.kwargs["default_action_domain"] == "no_action" diff --git a/src/mobius/_diffusers_builder.py b/src/mobius/_diffusers_builder.py index f95ee9e20..1c88de094 100644 --- a/src/mobius/_diffusers_builder.py +++ b/src/mobius/_diffusers_builder.py @@ -42,6 +42,7 @@ def _init_diffusers_class_map() -> None: if _DIFFUSERS_CLASS_MAP: return + from mobius._configs import Cosmos3OmniGeneratorConfig, WanVAEConfig from mobius._diffusers_configs import ( CLIPTextConfig, CogVideoXConfig, @@ -54,6 +55,7 @@ def _init_diffusers_class_map() -> None: from mobius.models.cogvideox import ( CogVideoXTransformer3DModel, ) + from mobius.models.cosmos3_omni_generator import Cosmos3OmniGeneratorModel from mobius.models.dit import DiTConfig, DiTTransformer2DModel from mobius.models.flux_sd3 import ( FluxConfig, @@ -67,6 +69,7 @@ def _init_diffusers_class_map() -> None: from mobius.models.unet import UNet2DConditionModel from mobius.models.vae import AutoencoderKLModel from mobius.models.video_vae import VideoAutoencoderModel, VideoVAEConfig + from mobius.models.wan_vae import AutoencoderKLWanModel _DIFFUSERS_CLASS_MAP.update( { @@ -104,6 +107,16 @@ def _init_diffusers_class_map() -> None: CogVideoXConfig, "video-denoising", ), + "Cosmos3OmniTransformer": ( + Cosmos3OmniGeneratorModel, + Cosmos3OmniGeneratorConfig, + "cosmos3-omni-generator", + ), + "AutoencoderKLWan": ( + AutoencoderKLWanModel, + WanVAEConfig, + "wan-vae", + ), } ) @@ -236,6 +249,8 @@ def build_diffusers_pipeline( dtype: str | ir.DataType | None = None, load_weights: bool = True, unet_loras: dict | None = None, + execution_provider: str = "default", + trace_optimization: bool = False, ) -> ModelPackage: """Build ONNX models for all supported components in a diffusers pipeline. @@ -255,6 +270,9 @@ def build_diffusers_pipeline( inferred from the file); at inference a ``lora_gate.{name}`` scalar input switches/blends it. Requires ``load_weights=True`` to apply the adapter weights. + execution_provider: Target execution provider for component-specific + optimization and lowering. + trace_optimization: Whether to log each component optimization stage. Returns: A :class:`ModelPackage` containing the built component model(s). @@ -318,7 +336,13 @@ def build_diffusers_pipeline( model_module = module_class(config) - sub_pkg = build_from_module(model_module, config, task_name) + sub_pkg = build_from_module( + model_module, + config, + task_name, + execution_provider=execution_provider, + trace_optimization=trace_optimization, + ) # Flatten sub-package into the top-level package if len(sub_pkg) == 1 and "model" in sub_pkg: diff --git a/src/mobius/_diffusers_builder_test.py b/src/mobius/_diffusers_builder_test.py index 8f1cf5dfd..3e42349a1 100644 --- a/src/mobius/_diffusers_builder_test.py +++ b/src/mobius/_diffusers_builder_test.py @@ -59,6 +59,8 @@ def test_populates_expected_classes(self): "AutoencoderKLQwenImage", "AutoencoderKLCogVideoX", "CogVideoXTransformer3DModel", + "Cosmos3OmniTransformer", + "AutoencoderKLWan", } assert expected_keys == set(_DIFFUSERS_CLASS_MAP.keys()) @@ -83,6 +85,8 @@ def test_task_names_are_valid(self): "qwen-image-vae", "video-denoising", "feature-extraction", + "cosmos3-omni-generator", + "wan-vae", } for class_name, (_, _, task_name) in _DIFFUSERS_CLASS_MAP.items(): assert task_name in valid_tasks, f"Unknown task '{task_name}' for {class_name}" @@ -346,7 +350,7 @@ def test_multiple_components_built( ) mock_load_config.return_value = {} - def fake_build(module, config, task_name): + def fake_build(module, config, task_name, **_kwargs): graph = ir.Graph([], [], nodes=[], name="g") return ModelPackage({"model": ir.Model(graph, ir_version=10)}) @@ -400,6 +404,33 @@ def test_dtype_ir_datatype_passthrough( # Verify build_from_module was called (ir.DataType accepted without error) mock_build_from_module.assert_called_once() + @patch("mobius._diffusers_builder.build_from_module") + @patch( + "mobius._diffusers_builder._load_diffusers_component_config", + ) + @patch( + "mobius._diffusers_builder._load_diffusers_pipeline_index", + ) + def test_execution_provider_forwarded_to_component_build( + self, + mock_load_index, + mock_load_config, + mock_build_from_module, + ): + self._mock_build_for_vae(mock_load_index, mock_load_config, mock_build_from_module) + + build_diffusers_pipeline( + "fake/vae-model", + load_weights=False, + execution_provider="cuda", + trace_optimization=True, + ) + + assert mock_build_from_module.call_args.kwargs == { + "execution_provider": "cuda", + "trace_optimization": True, + } + # ── build_diffusers_pipeline weight loading ────────────────────────────── @@ -504,7 +535,7 @@ def test_preprocess_weights_called_when_available( # The module class will have preprocess_weights set by AutoencoderKLModel # We patch it at the module instance level via build_from_module's first arg - def capture_build(module, config, task_name): + def capture_build(module, config, task_name, **_kwargs): module.preprocess_weights = lambda sd: processed_weights return ModelPackage({"model": model}) diff --git a/src/mobius/_diffusers_checkpoint.py b/src/mobius/_diffusers_checkpoint.py new file mode 100644 index 000000000..91a8f8acd --- /dev/null +++ b/src/mobius/_diffusers_checkpoint.py @@ -0,0 +1,267 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Read-only access to diffusers-style multi-component checkpoints. + +A diffusers checkpoint is a directory (local or on the Hub) with a +``model_index.json`` at the root and one subdirectory per component, each +holding its own ``config.json`` and safetensors shards. World-model exporters +need to inspect such a checkpoint before building anything: which components +exist, what class each one declares, which tensors a component ships, and +which auxiliary files must travel with the exported package. + +Every helper here is model-agnostic — it resolves and reads files without +interpreting their contents. Local paths are resolved with traversal +protection so a malicious index or component name cannot read outside the +checkpoint directory. +""" + +from __future__ import annotations + +__all__ = [ + "component_class", + "component_shard_paths", + "component_weight_names", + "load_checkpoint_json", + "load_component_weights", + "load_optional_checkpoint_json", + "resolve_assets", + "resolve_checkpoint_file", +] + +import json +import pathlib +from collections.abc import Iterable, Mapping +from typing import Any + +import safetensors.torch +from huggingface_hub import HfApi, hf_hub_download +from huggingface_hub.utils import EntryNotFoundError +from safetensors import safe_open + +from mobius._diffusers_builder import _download_diffusers_component_weights +from mobius._weight_loading import _dequantize_fp8_weights + +#: Weight-index basenames used by diffusers and transformers components. +_INDEX_NAMES: tuple[str, ...] = ( + "diffusion_pytorch_model.safetensors.index.json", + "model.safetensors.index.json", +) +#: Single-shard weight filenames used by diffusers and transformers components. +_SINGLE_NAMES: tuple[str, ...] = ( + "diffusion_pytorch_model.safetensors", + "model.safetensors", +) + + +def _local_component_directory( + model_dir: pathlib.Path, + component: str, +) -> pathlib.Path: + """Resolve a component directory without allowing checkpoint-root escape.""" + root = model_dir.resolve() + component_dir = (root / component).resolve() + try: + component_dir.relative_to(root) + except ValueError as error: + raise ValueError(f"Unsafe component path {component!r}") from error + return component_dir + + +def resolve_checkpoint_file( + model_id: str, + filename: str, + *, + required: bool = True, +) -> str | None: + """Resolve one local-or-Hub checkpoint file without interpreting it. + + Args: + model_id: Local checkpoint directory or Hub repository id. + filename: ``/``-separated path relative to the checkpoint root. + required: Whether a missing file is an error. + + Returns: + An absolute local path, or ``None`` when an optional file is absent. + + Raises: + ValueError: If *filename* resolves outside a local checkpoint. + FileNotFoundError: If a required local file is missing. + """ + root = pathlib.Path(model_id) + if root.is_dir(): + path = (root / pathlib.PurePosixPath(filename)).resolve() + try: + path.relative_to(root.resolve()) + except ValueError as error: + raise ValueError( + f"Checkpoint file escapes model directory: {filename!r}" + ) from error + if path.is_file(): + return str(path) + if required: + raise FileNotFoundError(f"Required checkpoint file not found: {path}") + return None + try: + return hf_hub_download(repo_id=model_id, filename=filename) + except EntryNotFoundError: + if required: + raise + return None + + +def load_checkpoint_json(model_id: str, filename: str) -> tuple[dict[str, Any], str]: + """Load a required checkpoint JSON object and return it with its path.""" + path = resolve_checkpoint_file(model_id, filename) + assert path is not None + with open(path, encoding="utf-8") as handle: + value = json.load(handle) + if not isinstance(value, dict): + raise TypeError(f"{filename!r} must contain a JSON object") + return value, path + + +def load_optional_checkpoint_json(model_id: str, filename: str) -> dict[str, Any]: + """Load an optional checkpoint JSON object, or ``{}`` when it is absent.""" + path = resolve_checkpoint_file(model_id, filename, required=False) + if path is None: + return {} + with open(path, encoding="utf-8") as handle: + value = json.load(handle) + if not isinstance(value, dict): + raise TypeError(f"{filename!r} must contain a JSON object") + return value + + +def component_class( + pipeline_index: Mapping[str, Any], + component: str, +) -> str | None: + """Return the class a ``model_index.json`` declares for *component*. + + Diffusers records each component as a ``[library, class_name]`` pair and + marks an absent component as ``[null, null]``. + """ + info = pipeline_index.get(component) + if info in (None, [None, None]): + return None + if not isinstance(info, list) or len(info) != 2 or not isinstance(info[1], str): + raise ValueError(f"Invalid model_index.json entry for {component!r}: {info!r}") + return info[1] + + +def component_weight_names(model_id: str, component: str) -> set[str]: + """Read only safetensors metadata to determine component graph shape. + + Tensor names alone reveal which optional towers a component ships, so no + tensor data is downloaded or read here. + """ + root = pathlib.Path(model_id) + if root.is_dir(): + component_dir = _local_component_directory(root, component) + for name in _INDEX_NAMES: + path = component_dir / name + if path.is_file(): + with path.open(encoding="utf-8") as handle: + index = json.load(handle) + return set(index["weight_map"]) + for name in _SINGLE_NAMES: + path = component_dir / name + if path.is_file(): + with safe_open(str(path), framework="pt", device="cpu") as file: + return set(file.keys()) + raise FileNotFoundError( + f"No safetensors checkpoint found for component {component!r} in {model_id!r}" + ) + + for name in _INDEX_NAMES: + filename = f"{component}/{name}" + path = resolve_checkpoint_file(model_id, filename, required=False) + if path is not None: + with open(path, encoding="utf-8") as handle: + index = json.load(handle) + return set(index["weight_map"]) + api = HfApi() + for name in _SINGLE_NAMES: + filename = f"{component}/{name}" + try: + metadata = api.parse_safetensors_file_metadata(model_id, filename) + except EntryNotFoundError: + continue + return set(metadata.tensors) + raise FileNotFoundError( + f"No safetensors checkpoint found for component {component!r} in {model_id!r}" + ) + + +def component_shard_paths( + model_dir: pathlib.Path, + component: str, +) -> list[pathlib.Path]: + """Resolve local component shards with traversal protection.""" + component_dir = _local_component_directory(model_dir, component) + + for basename in ( + "diffusion_pytorch_model", + "model", + ): + index_path = component_dir / f"{basename}.safetensors.index.json" + if not index_path.is_file(): + continue + with index_path.open(encoding="utf-8") as handle: + index = json.load(handle) + paths: list[pathlib.Path] = [] + for filename in sorted(set(index["weight_map"].values())): + relative = pathlib.PurePosixPath(str(filename).replace("\\", "/")) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError(f"Unsafe component weight filename: {filename!r}") + path = (component_dir / relative).resolve() + try: + path.relative_to(component_dir) + except ValueError as error: + raise ValueError( + f"Component weight filename escapes its directory: {filename!r}" + ) from error + if not path.is_file(): + raise FileNotFoundError(path) + paths.append(path) + return paths + + for basename in _SINGLE_NAMES: + path = component_dir / basename + if path.is_file(): + return [path] + raise FileNotFoundError( + f"No safetensors checkpoint found for component {component!r} in {model_dir}" + ) + + +def load_component_weights(model_id: str, component: str) -> dict[str, Any]: + """Load one diffusers component, supporting Hub and local directories.""" + root = pathlib.Path(model_id) + if not root.is_dir(): + return _download_diffusers_component_weights(model_id, component) + + state_dict: dict[str, Any] = {} + for path in component_shard_paths(root, component): + state_dict.update(safetensors.torch.load_file(str(path))) + return _dequantize_fp8_weights(state_dict) + + +def resolve_assets( + model_id: str, + candidates: Iterable[tuple[str, bool]], +) -> dict[str, tuple[str, bool]]: + """Resolve runtime asset candidates to ``destination -> (source, required)``. + + Candidates are ``(relative path, required)`` pairs. Required files must + exist; optional files are skipped when absent. Contents are never read: + an asset is copied into the exported package verbatim. The caller owns the + candidate list, so which files a model family ships stays with that family. + """ + assets: dict[str, tuple[str, bool]] = {} + for destination, required in candidates: + source = resolve_checkpoint_file(model_id, destination, required=required) + if source is not None: + assets[destination] = (source, required) + return assets diff --git a/src/mobius/_diffusers_checkpoint_test.py b/src/mobius/_diffusers_checkpoint_test.py new file mode 100644 index 000000000..a6f8e5b50 --- /dev/null +++ b/src/mobius/_diffusers_checkpoint_test.py @@ -0,0 +1,227 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +import pathlib + +import pytest +import safetensors.torch +import torch + +from mobius._diffusers_checkpoint import ( + component_class, + component_shard_paths, + component_weight_names, + load_checkpoint_json, + load_component_weights, + load_optional_checkpoint_json, + resolve_assets, + resolve_checkpoint_file, +) + + +def _write(path: pathlib.Path, value: object) -> pathlib.Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value), encoding="utf-8") + return path + + +def test_resolve_file_returns_local_path(tmp_path) -> None: + _write(tmp_path / "scheduler" / "scheduler_config.json", {"_class_name": "Flow"}) + + path = resolve_checkpoint_file(str(tmp_path), "scheduler/scheduler_config.json") + + assert path is not None + assert pathlib.Path(path) == (tmp_path / "scheduler" / "scheduler_config.json").resolve() + + +def test_resolve_file_rejects_paths_escaping_the_checkpoint(tmp_path) -> None: + _write(tmp_path.parent / "secret.json", {"token": "value"}) + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + + with pytest.raises(ValueError, match="escapes model directory"): + resolve_checkpoint_file(str(tmp_path), "../secret.json") + + +def test_resolve_file_reports_missing_required_file(tmp_path) -> None: + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + + with pytest.raises(FileNotFoundError, match="Required checkpoint file not found"): + resolve_checkpoint_file(str(tmp_path), "model_index.json") + + +def test_resolve_file_skips_missing_optional_file(tmp_path) -> None: + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + + assert resolve_checkpoint_file(str(tmp_path), "tokenizer.json", required=False) is None + + +def test_load_json_returns_object_and_path(tmp_path) -> None: + _write(tmp_path / "config.json", {"model_type": "example"}) + + value, path = load_checkpoint_json(str(tmp_path), "config.json") + + assert value == {"model_type": "example"} + assert pathlib.Path(path).name == "config.json" + + +def test_load_json_rejects_non_object_documents(tmp_path) -> None: + _write(tmp_path / "config.json", [1, 2, 3]) + + with pytest.raises(TypeError, match="must contain a JSON object"): + load_checkpoint_json(str(tmp_path), "config.json") + + +def test_optional_json_defaults_to_empty_mapping(tmp_path) -> None: + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + + assert load_optional_checkpoint_json(str(tmp_path), "generation_config.json") == {} + + +def test_optional_json_reads_present_file(tmp_path) -> None: + _write(tmp_path / "generation_config.json", {"top_k": 20}) + + assert load_optional_checkpoint_json(str(tmp_path), "generation_config.json") == { + "top_k": 20 + } + + +def test_component_class_reads_model_index_entries() -> None: + index = { + "transformer": ["diffusers", "Cosmos3OmniTransformer"], + "sound_tokenizer": [None, None], + } + + assert component_class(index, "transformer") == "Cosmos3OmniTransformer" + assert component_class(index, "sound_tokenizer") is None + assert component_class(index, "absent") is None + + +def test_component_class_rejects_malformed_entries() -> None: + with pytest.raises(ValueError, match=r"Invalid model_index\.json entry"): + component_class({"vae": "AutoencoderKLWan"}, "vae") + + +def test_component_weight_names_reads_single_shard_metadata(tmp_path) -> None: + component = tmp_path / "vae" + component.mkdir() + safetensors.torch.save_file( + {"encoder.conv.weight": torch.zeros(1), "decoder.conv.weight": torch.zeros(1)}, + str(component / "diffusion_pytorch_model.safetensors"), + ) + + names = component_weight_names(str(tmp_path), "vae") + + assert names == {"encoder.conv.weight", "decoder.conv.weight"} + + +def test_component_weight_names_reads_sharded_index(tmp_path) -> None: + _write( + tmp_path / "transformer" / "model.safetensors.index.json", + {"weight_map": {"a.weight": "model-00001.safetensors"}}, + ) + + assert component_weight_names(str(tmp_path), "transformer") == {"a.weight"} + + +def test_component_weight_names_requires_a_checkpoint(tmp_path) -> None: + (tmp_path / "vae").mkdir() + + with pytest.raises(FileNotFoundError, match="No safetensors checkpoint"): + component_weight_names(str(tmp_path), "vae") + + +def test_component_weight_names_rejects_escaping_component_names(tmp_path) -> None: + with pytest.raises(ValueError, match="Unsafe component path"): + component_weight_names(str(tmp_path), "../elsewhere") + + +def test_component_shard_paths_rejects_escaping_shard_entries(tmp_path) -> None: + _write( + tmp_path / "vae" / "diffusion_pytorch_model.safetensors.index.json", + {"weight_map": {"a.weight": "../outside.safetensors"}}, + ) + + with pytest.raises(ValueError, match="Unsafe component weight filename"): + component_shard_paths(tmp_path, "vae") + + +def test_component_shard_paths_rejects_escaping_component_names(tmp_path) -> None: + with pytest.raises(ValueError, match="Unsafe component path"): + component_shard_paths(tmp_path, "../elsewhere") + + +def test_component_shard_paths_orders_index_shards(tmp_path) -> None: + component = tmp_path / "vae" + component.mkdir() + for name in ("model-00002.safetensors", "model-00001.safetensors"): + safetensors.torch.save_file({"w": torch.zeros(1)}, str(component / name)) + _write( + component / "diffusion_pytorch_model.safetensors.index.json", + { + "weight_map": { + "b.weight": "model-00002.safetensors", + "a.weight": "model-00001.safetensors", + } + }, + ) + + paths = component_shard_paths(tmp_path, "vae") + + assert [path.name for path in paths] == [ + "model-00001.safetensors", + "model-00002.safetensors", + ] + + +def test_load_component_weights_merges_local_shards(tmp_path) -> None: + component = tmp_path / "vae" + component.mkdir() + safetensors.torch.save_file( + {"a.weight": torch.ones(2)}, + str(component / "model-00001.safetensors"), + ) + safetensors.torch.save_file( + {"b.weight": torch.zeros(2)}, + str(component / "model-00002.safetensors"), + ) + _write( + component / "model.safetensors.index.json", + { + "weight_map": { + "a.weight": "model-00001.safetensors", + "b.weight": "model-00002.safetensors", + } + }, + ) + + weights = load_component_weights(str(tmp_path), "vae") + + assert set(weights) == {"a.weight", "b.weight"} + + +def test_resolve_assets_skips_absent_optional_candidates(tmp_path) -> None: + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + _write(tmp_path / "text_tokenizer" / "tokenizer.json", {}) + + assets = resolve_assets( + str(tmp_path), + ( + ("config.json", True), + ("tokenizer.json", False), + ("text_tokenizer/tokenizer.json", False), + ), + ) + + assert set(assets) == {"config.json", "text_tokenizer/tokenizer.json"} + assert assets["config.json"][1] is True + assert assets["text_tokenizer/tokenizer.json"][1] is False + + +def test_resolve_assets_propagates_missing_required_candidate(tmp_path) -> None: + (tmp_path / "config.json").write_text("{}", encoding="utf-8") + + with pytest.raises(FileNotFoundError): + resolve_assets(str(tmp_path), (("model_index.json", True),)) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 00317f41f..f4e655841 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -151,6 +151,11 @@ def save( else: ir.save(model, path, external_data="model.onnx.data", callback=callback) + def validate_weights(self) -> None: + """Raise if any component initializer has no assigned tensor data.""" + for name, model in self.data.items(): + _check_weights(name, model) + @classmethod def load(cls, directory: str) -> ModelPackage: """Load all ``.onnx`` files from a directory into a package. @@ -208,6 +213,26 @@ def apply_weights( named component (with the prefix stripped). Unmatched weights are applied to all components. """ + applied = self.apply_weights_partial(state_dict, prefix_map=prefix_map) + _log_weight_mapping(state_dict, applied) + self.finalize_weights() + + def apply_weights_partial( + self, + state_dict: dict[str, torch.Tensor], + prefix_map: dict[str, str] | None = None, + ) -> set[str]: + """Apply one weight shard without folding partially populated graphs. + + This is the streaming counterpart of :meth:`apply_weights`. Call it + once for each checkpoint shard, then call :meth:`finalize_weights` + after the final shard. It intentionally does not log unmapped weights: + a shard shared by heterogeneous pipeline components is expected to + contain tensors that do not belong to every component. + + Returns: + Original state-dict names that matched graph initializers. + """ applied: set[str] = set() if len(self.data) == 1: @@ -249,8 +274,10 @@ def apply_weights( for model in self.data.values(): applied |= _apply_weights_to_model(model, unmatched) - _log_weight_mapping(state_dict, applied) + return applied + def finalize_weights(self) -> None: + """Fold initializer-only subgraphs after all weight shards are applied.""" # Fold constants now that weights have been loaded. # PackQKV emits Concat(w_q, w_k, w_v) in the graph; those nodes can only # be constant-folded once the weight tensors carry their const_value. diff --git a/src/mobius/_model_package_test.py b/src/mobius/_model_package_test.py index 12072f2e2..38d329c9f 100644 --- a/src/mobius/_model_package_test.py +++ b/src/mobius/_model_package_test.py @@ -6,8 +6,10 @@ from __future__ import annotations import logging +from unittest import mock import onnx_ir as ir +import pytest import torch from mobius._builder import build_from_module @@ -157,6 +159,31 @@ def test_multi_component_with_prefix_map(self): ) assert model1.graph.initializers[init_name].const_value is not None + def test_partial_weights_defer_folding_until_finalize(self): + config = make_config() + pkg = build_from_module(CausalLMModel(config), config) + model = pkg["model"] + init_name = next(iter(model.graph.initializers)) + shape = list(model.graph.initializers[init_name].shape) + + with mock.patch( + "mobius._model_package.fold_initializers_after_weights" + ) as fold_initializers: + applied = pkg.apply_weights_partial({init_name: torch.ones(shape)}) + fold_initializers.assert_not_called() + + pkg.finalize_weights() + + assert applied == {init_name} + fold_initializers.assert_called_once_with(model) + + def test_validate_weights_rejects_unassigned_initializers(self): + config = make_config() + pkg = build_from_module(CausalLMModel(config), config) + + with pytest.raises(ValueError, match="without weights"): + pkg.validate_weights() + class TestBuildPackageFromModule: def test_returns_model_package(self): diff --git a/src/mobius/_pipeline.py b/src/mobius/_pipeline.py new file mode 100644 index 000000000..b62f1331a --- /dev/null +++ b/src/mobius/_pipeline.py @@ -0,0 +1,3427 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Runtime-agnostic, model-agnostic compositional pipeline core. + +This module describes how already-built ONNX graphs are wired and how a +runtime executes them. It is a runtime-neutral contract: + +- it never runs, traces, or optimizes a graph; +- it serializes typed input semantics, generated-input programs, state + lifecycle, scheduler/sampling controls, transform parameters, assets, and + execution-provider hints; +- it never executes a :attr:`PipelineConnection.transform` — a transform is a + registered *kind* that a runtime resolves, and the core only records it and + its validated parameters/capabilities. + +Model-specific exporters compile HuggingFace/Diffusers semantics into this +standard contract. A runtime implements the registered programs and strategies +instead of inferring behavior from tensor names. + +Layers +------ + +``PipelineComponent`` + One :class:`onnx_ir.Model` plus a validated role/phase and optional + presence, capability, source, and config metadata. Its typed graph ports + are derived from the graph signature. +``PipelineConnection`` + ``component.output -> component.input``. Fan-out is allowed; an input may + have one initializer and one recurrent update. A recurrent connection is + loop-carried state scoped to a looping stage. A registered transform kind + may adapt the edge and declare additional context ports/capabilities. +``PipelineStage`` + A validated strategy kind over a set of components, with JSON-safe options + and capabilities. +``PipelineManifest`` + The serializable, deterministic executable contract. +``PipelinePackage`` + A :class:`~mobius._model_package.ModelPackage` that also carries the + manifest and per-component configs, and persists ``pipeline.json``. It can + ship opaque runtime assets (tokenizers, scheduler configs) whose + destinations are recorded and whose contents are never interpreted. +``PipelineBuilder`` + Convenience composition front-end with full structural validation on + :meth:`PipelineBuilder.build`. + +Example:: + + builder = PipelineBuilder() + builder.add_model("encoder", encoder_model, role="encoder") + builder.add_model("decoder", decoder_model, role="decoder") + builder.connect("encoder.hidden", "decoder.encoder_hidden") + builder.declare_external("encoder.pixel_values") + builder.declare_generated("decoder.position_ids", generator="zeros") + builder.add_stage("encode", "single_pass", ["encoder"]) + builder.add_stage("generate", "autoregressive", ["decoder"]) + builder.add_public_output("decoder.logits") + pkg = builder.build() + pkg.save("/out") +""" + +from __future__ import annotations + +__all__ = [ + "DEFAULT_PHASE", + "InputSource", + "GeneratedInputDefinition", + "GeneratedInputRule", + "LOOP_CARRIED_STATE_CAPABILITY", + "PIPELINE_FILENAME", + "PIPELINE_SCHEMA_VERSION", + "PhaseDefinition", + "PipelineAsset", + "PipelineBuilder", + "PipelineComponent", + "PipelineConnection", + "PipelineInput", + "PipelineManifest", + "PipelineOutput", + "PipelinePackage", + "PipelinePort", + "PipelineProfile", + "PipelineState", + "PipelineStage", + "PipelineValidationError", + "RoleDefinition", + "StrategyDefinition", + "StateDefinition", + "TensorSpec", + "TransformDefinition", + "phase_definition", + "generated_input_definition", + "register_phase", + "register_generated_input", + "register_role", + "register_strategy", + "register_state", + "register_transform", + "role_definition", + "strategy_definition", + "state_definition", + "transform_definition", +] + +import dataclasses +import json +import math +import os +import shutil +import tempfile +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from typing import Any, TypeAlias + +import onnx_ir as ir + +from mobius._model_package import ModelPackage + +JSONValue: TypeAlias = ( + "bool | int | float | str | list[JSONValue] | dict[str, JSONValue] | None" +) + +#: Schema version of :class:`PipelineManifest`. ``major.minor``; a mismatching +#: *major* version is a hard failure on load. +PIPELINE_SCHEMA_VERSION = "1.1" + +#: Name of the manifest file written next to the component graphs. +PIPELINE_FILENAME = "pipeline.json" + +#: Separator between a component name and a graph port name in an endpoint. +ENDPOINT_SEPARATOR = "." + +#: Phase meaning "no phase restriction". +DEFAULT_PHASE = "always" + +#: Capability contributed by a stage that carries loop state across iterations. +LOOP_CARRIED_STATE_CAPABILITY = "loop_carried_state" + + +class PipelineValidationError(ValueError): + """Raised when a pipeline topology is structurally invalid.""" + + +# --------------------------------------------------------------------------- +# Name validation +# --------------------------------------------------------------------------- + +_INVALID_NAME_CHARS = frozenset('/\\:*?"<>|\0' + ENDPOINT_SEPARATOR) +_RESERVED_NAMES = frozenset( + ["con", "prn", "aux", "nul"] + + [f"com{i}" for i in range(1, 10)] + + [f"lpt{i}" for i in range(1, 10)] +) + + +def _validate_component_name(name: str) -> str: + """Validate that *name* is usable as a single, safe directory name.""" + if not isinstance(name, str) or not name or name.strip() != name: + raise PipelineValidationError( + f"Component name {name!r} must be a non-blank string without " + "surrounding whitespace." + ) + if name in {".", ".."} or ".." in name: + raise PipelineValidationError(f"Component name {name!r} must not contain '..'.") + bad = sorted(_INVALID_NAME_CHARS & set(name)) + if bad: + chars = ", ".join(repr(c) for c in bad) + raise PipelineValidationError( + f"Component name {name!r} must not contain {chars} because it is used " + "as a directory name and as an endpoint prefix." + ) + if any(ord(c) < 32 for c in name): + raise PipelineValidationError( + f"Component name {name!r} must not contain control characters." + ) + if name.split(".")[0].lower() in _RESERVED_NAMES: + raise PipelineValidationError( + f"Component name {name!r} is a reserved filesystem name." + ) + return name + + +def _validate_token(name: str, what: str) -> str: + """Validate a short identifier-like token (stage/role/strategy/phase name).""" + if not isinstance(name, str) or not name or name.strip() != name: + raise PipelineValidationError( + f"{what} name {name!r} must be a non-blank string without surrounding whitespace." + ) + if any(c in name for c in "/\\\0"): + raise PipelineValidationError( + f"{what} name {name!r} must not contain path separators." + ) + return name + + +def _validate_port_name(name: str, component: str) -> str: + """Validate a graph port name (ONNX value names may contain dots).""" + if not isinstance(name, str) or not name or name.strip() != name: + raise PipelineValidationError( + f"Port name {name!r} on component {component!r} must be a non-blank string." + ) + return name + + +def _validate_asset_path(path: str) -> str: + """Validate a ``/``-separated relative destination inside a package directory. + + Rejects anything that could escape the package directory or collide with a + platform-reserved name: absolute paths, drive letters, UNC prefixes, + backslashes, ``.``/``..`` segments, empty segments, and control characters. + """ + if not isinstance(path, str) or not path or path.strip() != path: + raise PipelineValidationError( + f"Asset destination {path!r} must be a non-blank string without " + "surrounding whitespace." + ) + if "\\" in path: + raise PipelineValidationError( + f"Asset destination {path!r} must use '/' separators, not backslashes." + ) + if path.startswith(("/", "~")) or os.path.isabs(path): + raise PipelineValidationError( + f"Asset destination {path!r} must be relative to the package directory." + ) + if len(path) >= 2 and path[1] == ":": + raise PipelineValidationError( + f"Asset destination {path!r} must not contain a drive letter." + ) + if path.endswith("/"): + raise PipelineValidationError( + f"Asset destination {path!r} must name a file, not a directory." + ) + for segment in path.split("/"): + if not segment: + raise PipelineValidationError( + f"Asset destination {path!r} must not contain empty path segments." + ) + if segment in {".", ".."}: + raise PipelineValidationError( + f"Asset destination {path!r} must not contain '.' or '..' segments." + ) + if segment.strip() != segment: + raise PipelineValidationError( + f"Asset destination {path!r} must not have padded path segments." + ) + if any(ord(c) < 32 for c in segment) or any(c in segment for c in ':*?"<>|'): + raise PipelineValidationError( + f"Asset destination {path!r} contains characters that are unsafe in a " + "file path." + ) + if ( + segment.lower() in _RESERVED_NAMES + or segment.split(".")[0].lower() in _RESERVED_NAMES + ): + raise PipelineValidationError( + f"Asset destination {path!r} uses reserved filesystem name {segment!r}." + ) + return path + + +# --------------------------------------------------------------------------- +# JSON safety +# --------------------------------------------------------------------------- + + +def _ensure_json_value(value: Any, context: str) -> JSONValue: + """Return a normalized, JSON-safe copy of *value* or raise.""" + if value is None or isinstance(value, (str, bool)): + return value + if isinstance(value, int): + return int(value) + if isinstance(value, float): + if not math.isfinite(value): + raise PipelineValidationError( + f"{context} must be JSON-safe; {value!r} is not a finite number." + ) + return float(value) + if isinstance(value, Mapping): + result: dict[str, JSONValue] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise PipelineValidationError( + f"{context} must be JSON-safe; mapping key {key!r} is not a string." + ) + result[key] = _ensure_json_value(item, f"{context}[{key!r}]") + return result + if isinstance(value, (list, tuple)): + return [_ensure_json_value(item, f"{context}[{i}]") for i, item in enumerate(value)] + raise PipelineValidationError( + f"{context} must be JSON-safe; got value of type {type(value).__name__}." + ) + + +def _ensure_json_mapping( + value: Mapping[str, Any] | None, context: str +) -> dict[str, JSONValue]: + if value is None: + return {} + normalized = _ensure_json_value(value, context) + if not isinstance(normalized, dict): + raise PipelineValidationError(f"{context} must be a mapping.") + return normalized + + +def _string_tuple(values: Iterable[str] | None, context: str) -> tuple[str, ...]: + """Normalize an iterable of capability-like strings: deduped and sorted.""" + if values is None: + return () + if isinstance(values, str): + raise PipelineValidationError( + f"{context} must be a sequence of strings, not a string." + ) + out = set() + for value in values: + if not isinstance(value, str) or not value.strip(): + raise PipelineValidationError(f"{context} entries must be non-blank strings.") + out.add(value) + return tuple(sorted(out)) + + +def _ordered_string_tuple( + values: Iterable[str] | None, + context: str, +) -> tuple[str, ...]: + """Normalize strings while preserving preference/declaration order.""" + if values is None: + return () + if isinstance(values, str): + raise PipelineValidationError( + f"{context} must be a sequence of strings, not a string." + ) + result: list[str] = [] + seen: set[str] = set() + for value in values: + if not isinstance(value, str) or not value.strip(): + raise PipelineValidationError(f"{context} entries must be non-blank strings.") + if value not in seen: + result.append(value) + seen.add(value) + return tuple(result) + + +# --------------------------------------------------------------------------- +# Open registries with closed validation +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class RoleDefinition: + """A registered component role.""" + + name: str + description: str = "" + + +@dataclasses.dataclass(frozen=True) +class PhaseDefinition: + """A registered ``run_on`` phase.""" + + name: str + description: str = "" + + +@dataclasses.dataclass(frozen=True) +class StrategyDefinition: + """A registered stage strategy kind. + + Attributes: + name: The strategy identifier used in manifests. + description: Human-readable summary. + loop_carried_state: Whether stages of this kind may own recurrent + (loop-carried) connections. + """ + + name: str + description: str = "" + loop_carried_state: bool = False + required_options: tuple[str, ...] = () + allowed_options: tuple[str, ...] | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "required_options", + _string_tuple(self.required_options, f"Strategy {self.name!r} required options"), + ) + if self.allowed_options is not None: + allowed = _string_tuple( + self.allowed_options, f"Strategy {self.name!r} allowed options" + ) + if not set(self.required_options) <= set(allowed): + raise PipelineValidationError( + f"Strategy {self.name!r} required options must be allowed." + ) + object.__setattr__(self, "allowed_options", allowed) + + +@dataclasses.dataclass(frozen=True) +class TransformDefinition: + """A registered edge transform kind. + + A transform is applied *by the runtime* on a connection. The core neither + defines nor executes the computation; it only records that the edge is + mediated and what the runtime must be able to do. Because a transform may + legitimately change dtype, rank, and shape (a VAE posterior parameterization + sampled and patchified into generator tokens, a scheduler turning a noise + prediction into the next latent), direct port compatibility is *not* + assumed for transformed edges. + + Attributes: + name: The transform identifier used in manifests. + description: Human-readable summary. + capabilities: Capabilities the executing runtime must provide for this + transform. A manifest using the transform must list them in + :attr:`PipelineManifest.required_capabilities`. + """ + + name: str + description: str = "" + capabilities: tuple[str, ...] = () + required_parameters: tuple[str, ...] = () + allowed_parameters: tuple[str, ...] | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "capabilities", + _string_tuple(self.capabilities, f"Transform {self.name!r} capabilities"), + ) + object.__setattr__( + self, + "required_parameters", + _string_tuple( + self.required_parameters, + f"Transform {self.name!r} required parameters", + ), + ) + if self.allowed_parameters is not None: + allowed = _string_tuple( + self.allowed_parameters, + f"Transform {self.name!r} allowed parameters", + ) + if not set(self.required_parameters) <= set(allowed): + raise PipelineValidationError( + f"Transform {self.name!r} required parameters must be allowed." + ) + object.__setattr__(self, "allowed_parameters", allowed) + + +@dataclasses.dataclass(frozen=True) +class GeneratedInputDefinition: + """A registered runtime program that materializes a graph input.""" + + name: str + description: str = "" + capabilities: tuple[str, ...] = () + required_parameters: tuple[str, ...] = () + allowed_parameters: tuple[str, ...] | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "capabilities", + _string_tuple(self.capabilities, f"Generator {self.name!r} capabilities"), + ) + object.__setattr__( + self, + "required_parameters", + _string_tuple( + self.required_parameters, + f"Generator {self.name!r} required parameters", + ), + ) + if self.allowed_parameters is not None: + allowed = _string_tuple( + self.allowed_parameters, + f"Generator {self.name!r} allowed parameters", + ) + if not set(self.required_parameters) <= set(allowed): + raise PipelineValidationError( + f"Generator {self.name!r} required parameters must be allowed." + ) + object.__setattr__(self, "allowed_parameters", allowed) + + +@dataclasses.dataclass(frozen=True) +class StateDefinition: + """A registered recurrent-state semantic.""" + + name: str + description: str = "" + + +_ROLES: dict[str, RoleDefinition] = {} +_PHASES: dict[str, PhaseDefinition] = {} +_STRATEGIES: dict[str, StrategyDefinition] = {} +_TRANSFORMS: dict[str, TransformDefinition] = {} +_GENERATED_INPUTS: dict[str, GeneratedInputDefinition] = {} +_STATES: dict[str, StateDefinition] = {} + + +def _register( + registry: dict[str, Any], + definition: Any, + what: str, +) -> Any: + """Idempotently register *definition*, rejecting conflicting redefinition.""" + _validate_token(definition.name, what) + existing = registry.get(definition.name) + if existing is not None: + if existing != definition: + raise PipelineValidationError( + f"{what} {definition.name!r} is already registered with a different " + f"definition ({existing!r} != {definition!r})." + ) + return existing + registry[definition.name] = definition + return definition + + +def register_role(name: str, *, description: str = "") -> RoleDefinition: + """Register a component role. + + Registration is idempotent for an identical definition and raises for a + conflicting redefinition of the same name. + """ + return _register(_ROLES, RoleDefinition(name, description), "Role") + + +def register_phase(name: str, *, description: str = "") -> PhaseDefinition: + """Register a ``run_on`` phase (idempotent for an identical definition).""" + return _register(_PHASES, PhaseDefinition(name, description), "Phase") + + +def register_strategy( + name: str, + *, + description: str = "", + loop_carried_state: bool = False, + required_options: Iterable[str] | None = None, + allowed_options: Iterable[str] | None = None, +) -> StrategyDefinition: + """Register a stage strategy kind (idempotent for an identical definition). + + Args: + name: Strategy identifier. + description: Human-readable summary. + loop_carried_state: Whether stages of this kind may own recurrent + connections. + """ + definition = StrategyDefinition( + name, + description, + loop_carried_state, + tuple(required_options or ()), + tuple(allowed_options) if allowed_options is not None else None, + ) + return _register(_STRATEGIES, definition, "Strategy") + + +def register_transform( + name: str, + *, + description: str = "", + capabilities: Iterable[str] | None = None, + required_parameters: Iterable[str] | None = None, + allowed_parameters: Iterable[str] | None = None, +) -> TransformDefinition: + """Register an edge transform kind (idempotent for an identical definition). + + Args: + name: Transform identifier used by :attr:`PipelineConnection.transform`. + description: Human-readable summary. + capabilities: Capabilities the executing runtime must provide. Any + manifest that uses the transform must list them in its + ``required_capabilities``; :class:`PipelineBuilder` adds them + automatically. + """ + definition = TransformDefinition( + name, + description, + tuple(capabilities or ()), + tuple(required_parameters or ()), + tuple(allowed_parameters) if allowed_parameters is not None else None, + ) + return _register(_TRANSFORMS, definition, "Transform") + + +def register_generated_input( + name: str, + *, + description: str = "", + capabilities: Iterable[str] | None = None, + required_parameters: Iterable[str] | None = None, + allowed_parameters: Iterable[str] | None = None, +) -> GeneratedInputDefinition: + """Register a runtime input-generation program.""" + definition = GeneratedInputDefinition( + name, + description, + tuple(capabilities or ()), + tuple(required_parameters or ()), + tuple(allowed_parameters) if allowed_parameters is not None else None, + ) + return _register(_GENERATED_INPUTS, definition, "Generated input") + + +def register_state(name: str, *, description: str = "") -> StateDefinition: + """Register a recurrent-state semantic.""" + return _register(_STATES, StateDefinition(name, description), "State") + + +def role_definition(name: str) -> RoleDefinition: + """Return a registered role or raise :class:`PipelineValidationError`.""" + return _lookup(_ROLES, name, "role") + + +def phase_definition(name: str) -> PhaseDefinition: + """Return a registered phase or raise :class:`PipelineValidationError`.""" + return _lookup(_PHASES, name, "phase") + + +def strategy_definition(name: str) -> StrategyDefinition: + """Return a registered strategy or raise :class:`PipelineValidationError`.""" + return _lookup(_STRATEGIES, name, "strategy") + + +def transform_definition(name: str) -> TransformDefinition: + """Return a registered transform or raise :class:`PipelineValidationError`.""" + return _lookup(_TRANSFORMS, name, "transform") + + +def generated_input_definition(name: str) -> GeneratedInputDefinition: + """Return a registered generated-input program.""" + return _lookup(_GENERATED_INPUTS, name, "generated_input") + + +def state_definition(name: str) -> StateDefinition: + """Return a registered recurrent-state semantic.""" + return _lookup(_STATES, name, "state") + + +def _validate_registered_parameters( + *, + kind: str, + parameters: Mapping[str, Any], + required: Iterable[str], + allowed: Iterable[str] | None, + context: str, +) -> dict[str, JSONValue]: + normalized = _ensure_json_mapping(parameters, f"{context} parameters") + missing = sorted(set(required) - set(normalized)) + if missing: + raise PipelineValidationError( + f"{context} {kind!r} is missing required parameter(s) {missing}." + ) + if allowed is not None: + unknown = sorted(set(normalized) - set(allowed)) + if unknown: + raise PipelineValidationError( + f"{context} {kind!r} has unknown parameter(s) {unknown}." + ) + return normalized + + +def _lookup(registry: dict[str, Any], name: str, what: str) -> Any: + definition = registry.get(name) + if definition is None: + known = ", ".join(sorted(registry)) or "" + raise PipelineValidationError( + f"Unknown {what} {name!r}. Known {what}s: {known}. " + f"Use register_{what}() to extend the registry." + ) + return definition + + +# Built-in, deliberately model-agnostic vocabulary. +for _role, _role_doc in ( + ("encoder", "Maps raw or embedded observations into a latent representation"), + ("decoder", "Maps latent representations back into outputs or tokens"), + ("embedding", "Turns discrete ids into dense vectors"), + ("projector", "Adapts one representation space into another"), + ("dynamics", "Advances latent state given an action or control signal"), + ("observation", "Consumes or produces environment observations"), + ("action", "Produces or consumes action representations"), + ("policy", "Maps state to an action distribution"), + ("value", "Estimates a scalar value of a state"), + ("reward", "Estimates a scalar reward"), + ("sampler", "Turns scores into concrete selections"), + ("transform", "Pure tensor reshaping/normalization graph"), + ("generic", "Unclassified component"), +): + register_role(_role, description=_role_doc) + +for _phase, _phase_doc in ( + (DEFAULT_PHASE, "No phase restriction"), + ("init", "Runs once when the pipeline is created"), + ("warmup", "Runs once before the first real request"), + ("prefill", "Runs on the initial/context pass"), + ("decode", "Runs on each incremental pass"), + ("step", "Runs on each iteration of a loop"), + ("refine", "Runs on refinement iterations"), + ("finalize", "Runs once after the loop terminates"), + ("on_demand", "Runs only when its presence condition holds"), +): + register_phase(_phase, description=_phase_doc) + +for _strategy, _strategy_doc, _loop, _options in ( + ("single_pass", "Run every component exactly once, in dependency order", False, ()), + ( + "autoregressive", + "Repeat until a stop condition, feeding outputs back", + True, + ("tokenizer_asset", "sampling", "stop", "max_tokens", "state_names"), + ), + ( + "iterative", + "Repeat a fixed or condition-bound number of iterations", + True, + ( + "scheduler", + "guidance", + "conditioning", + "default_steps", + "timestep", + "state_inputs", + "initial_state_inputs", + "prediction_type", + "packed_modalities", + ), + ), + ( + "state_transition", + "Advance a carried state one step per invocation", + True, + ("state_names", "max_steps", "stop"), + ), + ("composite", "Group other stages; ordering delegated to the runtime", False, ("stages",)), + ( + "on_demand", + "Run only when the component presence condition holds", + False, + ("presence",), + ), +): + register_strategy( + _strategy, + description=_strategy_doc, + loop_carried_state=_loop, + allowed_options=_options, + ) + +for _transform, _transform_doc, _transform_caps, _parameters in ( + ("cast", "Convert the tensor element type", ("tensor_cast",), ("to",)), + ( + "reshape", + "Rearrange dimensions without changing the element count", + ("tensor_reshape",), + ("shape", "input_layout", "output_layout"), + ), + ( + "normalize", + "Apply a scale/shift such as a latent scaling factor", + ("tensor_normalize",), + ("mean", "std", "scale", "shift"), + ), + ( + "sample", + "Draw a sample from a distribution parameterization, e.g. VAE moments", + ("stochastic_sampling",), + ("distribution", "seed_input"), + ), + ( + "patchify", + "Fold spatial/temporal axes into a token sequence", + ("tensor_patchify",), + ( + "spatial_patch_size", + "temporal_patch_size", + "input_layout", + "output_layout", + "channel_order", + ), + ), + ( + "unpatchify", + "Unfold a token sequence back into spatial/temporal axes", + ("tensor_patchify",), + ( + "spatial_patch_size", + "temporal_patch_size", + "input_layout", + "output_layout", + "channel_order", + ), + ), + ( + "scheduler_step", + "Advance an iterative solver, e.g. noise prediction to the next latent", + ("iterative_scheduler",), + ("scheduler_asset", "stage", "state", "timestep_input"), + ), + ("concat", "Join several producers along one axis", ("tensor_concat",), ("axis",)), + ( + "slice", + "Select a sub-range of one axis", + ("tensor_slice",), + ("axes", "starts", "ends", "steps"), + ), +): + register_transform( + _transform, + description=_transform_doc, + capabilities=_transform_caps, + allowed_parameters=_parameters, + ) + +for _generator, _generator_doc, _caps, _required, _allowed in ( + ( + "empty_tensor", + "Create an empty tensor matching a declared graph port", + (), + (), + ("shape", "dynamic_axes", "fill"), + ), + ( + "zeros", + "Create a zero-filled tensor", + (), + (), + ("shape", "shape_from", "dtype"), + ), + ( + "causal_attention_mask", + "Build an autoregressive attention mask", + ("attention_mask_program",), + ("sequence_input",), + ("sequence_input", "past_state", "visible_value", "masked_value"), + ), + ( + "multimodal_position_ids", + "Build one- or multi-axis position ids", + ("position_program",), + ("source", "axes"), + ( + "source", + "axes", + "mrope_sections", + "temporal_margin", + "reset_spatial", + "past_state", + ), + ), + ( + "packed_sequence_layout", + "Build packed-token indexes for one modality", + ("packed_sequence_program",), + ("modality",), + ("modality", "source", "layout", "understanding_prefix", "index_kind"), + ), + ( + "scheduler_timesteps", + "Materialize timesteps from an iterative-stage scheduler", + ("iterative_scheduler",), + ("stage",), + ("stage", "modality"), + ), + ( + "action_domain_ids", + "Map an action-domain semantic name to projection-bank ids", + ("action_domain_program",), + ("domain_input",), + ("domain_input", "default", "domain_map", "padded_dimension"), + ), +): + register_generated_input( + _generator, + description=_generator_doc, + capabilities=_caps, + required_parameters=_required, + allowed_parameters=_allowed, + ) + +for _state, _state_doc in ( + ("kv_cache", "Autoregressive key/value cache"), + ("diffusion_latent", "Loop-carried diffusion latent"), + ("action_state", "Loop-carried action trajectory"), + ("recurrent", "Generic recurrent tensor state"), +): + register_state(_state, description=_state_doc) + + +# --------------------------------------------------------------------------- +# Endpoints and typed graph ports +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True, order=True) +class PipelinePort: + """An endpoint: one graph port of one component. + + ``"decoder.logits"`` denotes port ``logits`` of component ``decoder``. + Component names may not contain :data:`ENDPOINT_SEPARATOR`, so the + qualified form is parsed by splitting on the first separator only; ONNX + value names containing dots (``past_key_values.0.key``) round-trip + unchanged. + """ + + component: str + port: str + + def __post_init__(self) -> None: + _validate_component_name(self.component) + _validate_port_name(self.port, self.component) + + @property + def qualified(self) -> str: + """``"component.port"``.""" + return f"{self.component}{ENDPOINT_SEPARATOR}{self.port}" + + def __str__(self) -> str: + return self.qualified + + @classmethod + def parse(cls, endpoint: str | PipelinePort) -> PipelinePort: + """Parse ``"component.port"`` (or pass through a ``PipelinePort``).""" + if isinstance(endpoint, PipelinePort): + return endpoint + if not isinstance(endpoint, str) or ENDPOINT_SEPARATOR not in endpoint: + raise PipelineValidationError( + f"Endpoint {endpoint!r} must have the form " + f"'component{ENDPOINT_SEPARATOR}port'." + ) + component, _, port = endpoint.partition(ENDPOINT_SEPARATOR) + return cls(component, port) + + +@dataclasses.dataclass(frozen=True) +class TensorSpec: + """The declared type of one graph port. + + Attributes: + name: The ONNX value name. + dtype: :class:`onnx_ir.DataType` name, e.g. ``"FLOAT"``. + shape: One entry per dimension. ``int`` for a static dimension, + ``str`` for a named symbolic dimension, ``None`` for an anonymous + dynamic dimension. Symbolic names are *never* compared across + components — only the rank and concrete dimensions are. + """ + + name: str + dtype: str + shape: tuple[int | str | None, ...] + + def __post_init__(self) -> None: + if not isinstance(self.name, str) or not self.name: + raise PipelineValidationError("Tensor spec name must be a non-blank string.") + try: + ir.DataType[self.dtype] + except KeyError as error: + raise PipelineValidationError( + f"Unknown dtype {self.dtype!r} for port {self.name!r}." + ) from error + dims: list[int | str | None] = [] + for dim in self.shape: + if dim is None or isinstance(dim, str): + dims.append(dim) + elif isinstance(dim, int) and not isinstance(dim, bool): + if dim < 0: + raise PipelineValidationError( + f"Port {self.name!r} has negative static dimension {dim}." + ) + dims.append(int(dim)) + else: + raise PipelineValidationError( + f"Port {self.name!r} has invalid dimension {dim!r}; " + "expected int, str, or None." + ) + object.__setattr__(self, "shape", tuple(dims)) + + @property + def rank(self) -> int: + """Number of dimensions.""" + return len(self.shape) + + def to_dict(self) -> dict[str, JSONValue]: + """Return the deterministic serializable form.""" + return {"name": self.name, "dtype": self.dtype, "shape": list(self.shape)} + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> TensorSpec: + """Inverse of :meth:`to_dict`.""" + body = _read_mapping(data, {"name", "dtype", "shape"}, "tensor spec") + shape = body.get("shape", []) + if not isinstance(shape, list): + raise PipelineValidationError("Tensor spec 'shape' must be a list.") + return cls(str(body["name"]), str(body["dtype"]), tuple(shape)) + + @classmethod + def from_value(cls, value: ir.Value, *, component: str, direction: str) -> TensorSpec: + """Derive a spec from an ``ir.Value`` on a component graph signature.""" + name = value.name + if not name: + raise PipelineValidationError( + f"Component {component!r} has an unnamed graph {direction}; " + "every pipeline port must be named." + ) + if value.dtype is None: + raise PipelineValidationError( + f"Component {component!r} {direction} {name!r} has no dtype; " + "pipeline components must have fully typed signatures." + ) + if value.shape is None: + raise PipelineValidationError( + f"Component {component!r} {direction} {name!r} has no shape; " + "pipeline components must have fully typed signatures." + ) + dims: list[int | str | None] = [] + for dim in value.shape: + if isinstance(dim, ir.SymbolicDim): + dims.append(dim.value) + else: + dims.append(int(dim)) + return cls(name, value.dtype.name, tuple(dims)) + + +def _tensor_mismatch(source: TensorSpec, target: TensorSpec) -> str | None: + """Return a human-readable mismatch reason, or ``None`` when compatible. + + Exact dtype and rank are required. Static dimensions are compared only + when *both* sides are concrete ints; symbolic names are never compared + because two graphs may legitimately use different names for the same axis. + """ + if source.dtype != target.dtype: + return f"dtype {source.dtype} != {target.dtype}" + if source.rank != target.rank: + return f"rank {source.rank} != {target.rank}" + for axis, (left, right) in enumerate(zip(source.shape, target.shape)): + left_static = isinstance(left, int) and not isinstance(left, bool) + right_static = isinstance(right, int) and not isinstance(right, bool) + if left_static and right_static and left != right: + return f"dim {axis}: {left} != {right}" + return None + + +# --------------------------------------------------------------------------- +# Components +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class PipelineComponent: + """One :class:`onnx_ir.Model` and its typed graph ports. + + A component is a *pure graph*: it has named, typed inputs and outputs and + no runtime semantics of its own. + + Attributes: + name: Unique component name. Also used as the directory name when the + package is saved, so it must be a safe single path segment. + role: Registered role (see :func:`register_role`). + inputs: Typed graph inputs, in graph order. + outputs: Typed graph outputs, in graph order. + run_on: Registered phase (see :func:`register_phase`). + presence: Optional opaque key naming the condition under which the + component exists at all. The core never evaluates it. + capabilities: Capabilities this component contributes. + source: Optional free-form provenance string (e.g. a model id). + config: JSON-safe, topology-relevant configuration. Runtime concerns + (tokenizers, preprocessing, sampling) do **not** belong here. + metadata: JSON-safe extension bag; unknown keys are preserved verbatim + across serialization round-trips. + model: The wrapped graph. Excluded from equality and serialization — + the manifest is topology only — and ``None`` for a manifest loaded + from disk (the graphs then live in the owning + :class:`PipelinePackage`). + """ + + name: str + role: str + inputs: tuple[TensorSpec, ...] = () + outputs: tuple[TensorSpec, ...] = () + run_on: str = DEFAULT_PHASE + presence: str | None = None + capabilities: tuple[str, ...] = () + preferred_execution_providers: tuple[str, ...] = () + parameter_dtype: str | None = None + source: str | None = None + config: dict[str, JSONValue] = dataclasses.field(default_factory=dict) + metadata: dict[str, JSONValue] = dataclasses.field(default_factory=dict) + model: ir.Model | None = dataclasses.field( + default=None, compare=False, repr=False, hash=False + ) + + def __post_init__(self) -> None: + _validate_component_name(self.name) + role_definition(self.role) + phase_definition(self.run_on) + if self.presence is not None: + _validate_token(self.presence, "Presence key") + if self.source is not None and not isinstance(self.source, str): + raise PipelineValidationError(f"Component {self.name!r} source must be a string.") + object.__setattr__(self, "inputs", tuple(self.inputs)) + object.__setattr__(self, "outputs", tuple(self.outputs)) + object.__setattr__( + self, "capabilities", _string_tuple(self.capabilities, f"{self.name} capabilities") + ) + if self.parameter_dtype is not None: + try: + ir.DataType[self.parameter_dtype] + except KeyError as error: + raise PipelineValidationError( + f"Component {self.name!r} has unknown parameter dtype " + f"{self.parameter_dtype!r}." + ) from error + object.__setattr__( + self, + "preferred_execution_providers", + _ordered_string_tuple( + self.preferred_execution_providers, + f"{self.name} preferred execution providers", + ), + ) + object.__setattr__( + self, "config", _ensure_json_mapping(self.config, f"{self.name} config") + ) + object.__setattr__( + self, "metadata", _ensure_json_mapping(self.metadata, f"{self.name} metadata") + ) + _check_unique((spec.name for spec in self.inputs), f"Component {self.name!r} input") + _check_unique((spec.name for spec in self.outputs), f"Component {self.name!r} output") + + @classmethod + def from_model( + cls, + name: str, + model: ir.Model, + *, + role: str, + run_on: str = DEFAULT_PHASE, + presence: str | None = None, + capabilities: Iterable[str] | None = None, + preferred_execution_providers: Iterable[str] | None = None, + parameter_dtype: str | None = None, + source: str | None = None, + config: Mapping[str, Any] | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> PipelineComponent: + """Derive a component (and its typed ports) from an ``ir.Model``.""" + _validate_component_name(name) + if not isinstance(model, ir.Model): + raise PipelineValidationError( + f"Component {name!r} must wrap an onnx_ir.Model, got {type(model).__name__}." + ) + inputs = tuple( + TensorSpec.from_value(value, component=name, direction="input") + for value in model.graph.inputs + ) + outputs = tuple( + TensorSpec.from_value(value, component=name, direction="output") + for value in model.graph.outputs + ) + return cls( + name=name, + role=role, + inputs=inputs, + outputs=outputs, + run_on=run_on, + presence=presence, + capabilities=_string_tuple(capabilities, f"{name} capabilities"), + preferred_execution_providers=_ordered_string_tuple( + preferred_execution_providers, + f"{name} preferred execution providers", + ), + parameter_dtype=parameter_dtype, + source=source, + config=_ensure_json_mapping(config, f"{name} config"), + metadata=_ensure_json_mapping(metadata, f"{name} metadata"), + model=model, + ) + + def input(self, port: str) -> TensorSpec | None: + """Return the typed input named *port*, if any.""" + return next((spec for spec in self.inputs if spec.name == port), None) + + def output(self, port: str) -> TensorSpec | None: + """Return the typed output named *port*, if any.""" + return next((spec for spec in self.outputs if spec.name == port), None) + + def with_model(self, model: ir.Model | None) -> PipelineComponent: + """Return a copy bound to *model* (topology unchanged).""" + return dataclasses.replace(self, model=model) + + def to_dict(self) -> dict[str, JSONValue]: + """Return the deterministic serializable form (topology only).""" + data: dict[str, JSONValue] = { + "name": self.name, + "role": self.role, + "run_on": self.run_on, + "inputs": [spec.to_dict() for spec in self.inputs], + "outputs": [spec.to_dict() for spec in self.outputs], + } + if self.presence is not None: + data["presence"] = self.presence + if self.capabilities: + data["capabilities"] = list(self.capabilities) + if self.preferred_execution_providers: + data["preferred_execution_providers"] = list(self.preferred_execution_providers) + if self.parameter_dtype is not None: + data["parameter_dtype"] = self.parameter_dtype + if self.source is not None: + data["source"] = self.source + if self.config: + data["config"] = dict(self.config) + if self.metadata: + data["metadata"] = dict(self.metadata) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PipelineComponent: + """Inverse of :meth:`to_dict`.""" + body = _read_mapping( + data, + { + "name", + "role", + "run_on", + "inputs", + "outputs", + "presence", + "capabilities", + "preferred_execution_providers", + "parameter_dtype", + "source", + "config", + "metadata", + }, + "component", + ) + return cls( + name=str(body["name"]), + role=str(body["role"]), + inputs=tuple(TensorSpec.from_dict(s) for s in body.get("inputs", [])), + outputs=tuple(TensorSpec.from_dict(s) for s in body.get("outputs", [])), + run_on=str(body.get("run_on", DEFAULT_PHASE)), + presence=body.get("presence"), + capabilities=tuple(body.get("capabilities", ())), + preferred_execution_providers=tuple(body.get("preferred_execution_providers", ())), + parameter_dtype=body.get("parameter_dtype"), + source=body.get("source"), + config=dict(body.get("config", {})), + metadata=dict(body.get("metadata", {})), + ) + + +# --------------------------------------------------------------------------- +# Connections +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class PipelineConnection: + """A directed edge ``source output -> target input``. + + Fan-out is allowed (one output may feed many inputs). An input may have at + most one initial producer and one recurrent producer, so an encoder can + seed state that a dynamics model subsequently updates. + + Attributes: + source: The producing component output endpoint. + target: The consuming component input endpoint. + recurrent: When ``True`` this is a loop-carried (state) edge: the value + produced by iteration *n* is consumed by iteration *n + 1*. Such + edges are stage-scoped — both endpoints must live in the same + looping stage — and are excluded from the acyclicity check. + transform: Optional *kind* of transform a runtime applies on this edge, + drawn from the transform registry (see :func:`register_transform`). + The core neither defines nor executes transforms. A transform may + legitimately change dtype, rank, and shape — a VAE posterior + parameterization sampled, normalized, and patchified into generator + tokens; a scheduler turning a noise prediction into the next latent + — so direct port compatibility is *not* checked for a transformed + edge. Endpoint existence, the single-producer rule, and the + transform's declared runtime capabilities are still enforced. + context: Additional component inputs/outputs consumed by a transform. + For example, a diffusion scheduler needs both the denoiser output + and the current loop-carried latent. Context ports do not produce + the target on their own and therefore do not participate in the + single-producer rule. + """ + + source: PipelinePort + target: PipelinePort + recurrent: bool = False + transform: str | None = None + context: tuple[PipelinePort, ...] = () + parameters: dict[str, JSONValue] = dataclasses.field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "source", PipelinePort.parse(self.source)) + object.__setattr__(self, "target", PipelinePort.parse(self.target)) + object.__setattr__(self, "recurrent", bool(self.recurrent)) + transform = None + if self.transform is not None: + _validate_token(self.transform, "Transform") + transform = transform_definition(self.transform) + context = tuple(PipelinePort.parse(port) for port in self.context) + if context and self.transform is None: + raise PipelineValidationError( + "Connection context is only valid when a transform is declared." + ) + _check_unique((port.qualified for port in context), "Connection context") + object.__setattr__(self, "context", context) + if transform is None: + if self.parameters: + raise PipelineValidationError( + "Connection parameters are only valid when a transform is declared." + ) + parameters: dict[str, JSONValue] = {} + else: + parameters = _validate_registered_parameters( + kind=transform.name, + parameters=self.parameters, + required=transform.required_parameters, + allowed=transform.allowed_parameters, + context="Transform", + ) + object.__setattr__(self, "parameters", parameters) + + @property + def transform_capabilities(self) -> tuple[str, ...]: + """Runtime capabilities required by this edge's transform, if any.""" + if self.transform is None: + return () + return transform_definition(self.transform).capabilities + + @property + def sort_key(self) -> tuple[str, str, bool, str, tuple[str, ...]]: + """Deterministic ordering key.""" + return ( + self.target.qualified, + self.source.qualified, + self.recurrent, + self.transform or "", + tuple(port.qualified for port in self.context), + ) + + def to_dict(self) -> dict[str, JSONValue]: + """Return the deterministic serializable form.""" + data: dict[str, JSONValue] = { + "source": self.source.qualified, + "target": self.target.qualified, + } + if self.recurrent: + data["recurrent"] = True + if self.transform is not None: + data["transform"] = self.transform + if self.context: + data["context"] = [port.qualified for port in self.context] + if self.parameters: + data["parameters"] = dict(self.parameters) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PipelineConnection: + """Inverse of :meth:`to_dict`.""" + body = _read_mapping( + data, + {"source", "target", "recurrent", "transform", "context", "parameters"}, + "connection", + ) + return cls( + source=PipelinePort.parse(str(body["source"])), + target=PipelinePort.parse(str(body["target"])), + recurrent=bool(body.get("recurrent", False)), + transform=body.get("transform"), + context=tuple(PipelinePort.parse(str(port)) for port in body.get("context", ())), + parameters=dict(body.get("parameters", {})), + ) + + +# --------------------------------------------------------------------------- +# Input sources and public outputs +# --------------------------------------------------------------------------- + + +class InputSource: + """The exhaustive set of ways a component input can be satisfied. + + Every graph input must be classified by exactly one of these. The core + never assumes that an unconnected input is caller-supplied — that must be + declared, so that a forgotten wire is an error rather than a silent + external input. + """ + + #: Produced by another component in the same invocation (a connection). + DATAFLOW = "dataflow" + #: Supplied by the caller of the pipeline. + EXTERNAL = "external" + #: Produced by the runtime harness per invocation (e.g. counters, ids). + GENERATED = "generated" + #: Carried across iterations (a recurrent connection, or runtime-owned state). + STATEFUL = "stateful" + #: Filled from a constant declared in the manifest. + DEFAULTED = "defaulted" + + #: Kinds that can be declared on a :class:`PipelineInput`. + DECLARABLE = (EXTERNAL, GENERATED, STATEFUL, DEFAULTED) + ALL = (DATAFLOW, EXTERNAL, GENERATED, STATEFUL, DEFAULTED) + + +@dataclasses.dataclass(frozen=True) +class PipelineProfile: + """Versioned runtime semantic profile implemented by the exporter.""" + + name: str + version: str + + def __post_init__(self) -> None: + _validate_token(self.name, "Profile") + _parse_version_format(self.version, "Profile") + + def to_dict(self) -> dict[str, JSONValue]: + return {"name": self.name, "version": self.version} + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PipelineProfile: + body = _read_mapping(data, {"name", "version"}, "profile") + return cls(name=str(body["name"]), version=str(body["version"])) + + +@dataclasses.dataclass(frozen=True) +class GeneratedInputRule: + """A concrete registered program for one runtime-generated input.""" + + kind: str + parameters: dict[str, JSONValue] = dataclasses.field(default_factory=dict) + + def __post_init__(self) -> None: + definition = generated_input_definition(self.kind) + parameters = _validate_registered_parameters( + kind=self.kind, + parameters=self.parameters, + required=definition.required_parameters, + allowed=definition.allowed_parameters, + context="Generated input", + ) + object.__setattr__(self, "parameters", parameters) + + @property + def capabilities(self) -> tuple[str, ...]: + return generated_input_definition(self.kind).capabilities + + def to_dict(self) -> dict[str, JSONValue]: + data: dict[str, JSONValue] = {"kind": self.kind} + if self.parameters: + data["parameters"] = dict(self.parameters) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> GeneratedInputRule: + body = _read_mapping(data, {"kind", "parameters"}, "generated input rule") + return cls( + kind=str(body["kind"]), + parameters=dict(body.get("parameters", {})), + ) + + +@dataclasses.dataclass(frozen=True) +class PipelineState: + """Explicit lifecycle contract for one recurrent connection.""" + + name: str + kind: str + input: PipelinePort + output: PipelinePort + lifetime: str + release_after: str + sequence_axis: int | None = None + metadata: dict[str, JSONValue] = dataclasses.field(default_factory=dict) + + def __post_init__(self) -> None: + _validate_token(self.name, "State") + state_definition(self.kind) + object.__setattr__(self, "input", PipelinePort.parse(self.input)) + object.__setattr__(self, "output", PipelinePort.parse(self.output)) + if self.lifetime not in {"iteration", "sequence", "request", "session"}: + raise PipelineValidationError( + f"State {self.name!r} has unsupported lifetime {self.lifetime!r}." + ) + _validate_token(self.release_after, "State release stage") + if self.sequence_axis is not None and self.sequence_axis < 0: + raise PipelineValidationError( + f"State {self.name!r} sequence_axis must be non-negative." + ) + object.__setattr__( + self, + "metadata", + _ensure_json_mapping(self.metadata, f"State {self.name!r} metadata"), + ) + + def to_dict(self) -> dict[str, JSONValue]: + data: dict[str, JSONValue] = { + "name": self.name, + "kind": self.kind, + "input": self.input.qualified, + "output": self.output.qualified, + "lifetime": self.lifetime, + "release_after": self.release_after, + } + if self.sequence_axis is not None: + data["sequence_axis"] = self.sequence_axis + if self.metadata: + data["metadata"] = dict(self.metadata) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PipelineState: + body = _read_mapping( + data, + { + "name", + "kind", + "input", + "output", + "lifetime", + "release_after", + "sequence_axis", + "metadata", + }, + "state", + ) + return cls( + name=str(body["name"]), + kind=str(body["kind"]), + input=PipelinePort.parse(str(body["input"])), + output=PipelinePort.parse(str(body["output"])), + lifetime=str(body["lifetime"]), + release_after=str(body["release_after"]), + sequence_axis=body.get("sequence_axis"), + metadata=dict(body.get("metadata", {})), + ) + + +@dataclasses.dataclass(frozen=True) +class PipelineInput: + """A declared source for a component input that no connection feeds. + + Attributes: + port: The component input endpoint. + kind: One of :data:`InputSource.DECLARABLE`. + value: The constant for ``defaulted`` inputs; must be JSON-safe and + must be ``None`` for every other kind. + alias: Optional pipeline-level name for ``external`` inputs. + """ + + port: PipelinePort + kind: str + value: JSONValue = None + alias: str | None = None + semantic: str | None = None + required: bool = True + presence: str | None = None + generator: GeneratedInputRule | None = None + + def __post_init__(self) -> None: + object.__setattr__(self, "port", PipelinePort.parse(self.port)) + if self.kind not in InputSource.DECLARABLE: + allowed = ", ".join(InputSource.DECLARABLE) + raise PipelineValidationError( + f"Input {self.port.qualified!r} has unknown source kind " + f"{self.kind!r}; expected one of: {allowed}." + ) + if self.kind == InputSource.DEFAULTED: + if self.value is None: + raise PipelineValidationError( + f"Defaulted input {self.port.qualified!r} requires a value; " + "use kind 'external' or 'generated' if there is no constant." + ) + object.__setattr__( + self, + "value", + _ensure_json_value(self.value, f"Default for {self.port.qualified!r}"), + ) + object.__setattr__(self, "required", False) + elif self.value is not None: + raise PipelineValidationError( + f"Input {self.port.qualified!r} of kind {self.kind!r} must not " + "carry a default value." + ) + if self.kind == InputSource.GENERATED: + if self.generator is None: + raise PipelineValidationError( + f"Generated input {self.port.qualified!r} requires a generation rule." + ) + elif self.generator is not None: + raise PipelineValidationError( + f"Only generated inputs may declare a generation rule; " + f"{self.port.qualified!r} is {self.kind!r}." + ) + if self.alias is not None: + if self.kind != InputSource.EXTERNAL: + raise PipelineValidationError( + f"Only external inputs may declare an alias; " + f"{self.port.qualified!r} is {self.kind!r}." + ) + _validate_token(self.alias, "Input alias") + if self.semantic is not None: + _validate_token(self.semantic, "Input semantic") + if self.presence is not None: + _validate_token(self.presence, "Input presence") + object.__setattr__(self, "required", bool(self.required)) + + @property + def name(self) -> str: + """The pipeline-level name of this input.""" + return self.alias or self.port.port + + def to_dict(self) -> dict[str, JSONValue]: + """Return the deterministic serializable form.""" + data: dict[str, JSONValue] = {"port": self.port.qualified, "kind": self.kind} + data["required"] = self.required + if self.kind == InputSource.DEFAULTED: + data["value"] = self.value + if self.alias is not None: + data["alias"] = self.alias + if self.semantic is not None: + data["semantic"] = self.semantic + if self.presence is not None: + data["presence"] = self.presence + if self.generator is not None: + data["generator"] = self.generator.to_dict() + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PipelineInput: + """Inverse of :meth:`to_dict`.""" + body = _read_mapping( + data, + { + "port", + "kind", + "value", + "alias", + "semantic", + "required", + "presence", + "generator", + }, + "input", + ) + return cls( + port=PipelinePort.parse(str(body["port"])), + kind=str(body["kind"]), + value=body.get("value"), + alias=body.get("alias"), + semantic=body.get("semantic"), + required=bool(body.get("required", True)), + presence=body.get("presence"), + generator=( + GeneratedInputRule.from_dict(body["generator"]) + if body.get("generator") is not None + else None + ), + ) + + +@dataclasses.dataclass(frozen=True) +class PipelineOutput: + """A component output or final recurrent state exposed as a result.""" + + port: PipelinePort | None = None + alias: str | None = None + state: str | None = None + + def __post_init__(self) -> None: + if (self.port is None) == (self.state is None): + raise PipelineValidationError( + "Pipeline output must reference exactly one component port or state." + ) + if self.port is not None: + object.__setattr__(self, "port", PipelinePort.parse(self.port)) + if self.state is not None: + _validate_token(self.state, "Output state") + if self.alias is not None: + _validate_token(self.alias, "Output alias") + + @property + def name(self) -> str: + """The pipeline-level name of this output.""" + if self.alias is not None: + return self.alias + if self.port is not None: + return self.port.port + assert self.state is not None + return self.state + + def to_dict(self) -> dict[str, JSONValue]: + """Return the deterministic serializable form.""" + data: dict[str, JSONValue] = {} + if self.port is not None: + data["port"] = self.port.qualified + else: + data["state"] = self.state + if self.alias is not None: + data["alias"] = self.alias + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PipelineOutput: + """Inverse of :meth:`to_dict`.""" + body = _read_mapping(data, {"port", "state", "alias"}, "output") + return cls( + port=( + PipelinePort.parse(str(body["port"])) if body.get("port") is not None else None + ), + state=body.get("state"), + alias=body.get("alias"), + ) + + +@dataclasses.dataclass(frozen=True) +class PipelineAsset: + """An opaque runtime file that ships next to the component graphs. + + Tokenizers, scheduler configs, and processor configs are runtime assets: + the topology core copies them and records *where* they live, and never + opens, parses, or interprets their contents. Only the destination is part + of the manifest — machine-local source paths never are. + + Attributes: + path: ``/``-separated destination relative to the package directory. + required: Whether :meth:`PipelinePackage.load` must find the file. + """ + + path: str + required: bool = True + + def __post_init__(self) -> None: + _validate_asset_path(self.path) + object.__setattr__(self, "required", bool(self.required)) + + def to_dict(self) -> dict[str, JSONValue]: + """Return the deterministic serializable form.""" + data: dict[str, JSONValue] = {"path": self.path} + if not self.required: + data["required"] = False + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PipelineAsset: + """Inverse of :meth:`to_dict`.""" + body = _read_mapping(data, {"path", "required"}, "asset") + return cls(path=str(body["path"]), required=bool(body.get("required", True))) + + +# --------------------------------------------------------------------------- +# Stages +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class PipelineStage: + """A group of components executed under one strategy. + + A stage says what kind of control flow applies to its components. + Strategy-specific ``options`` are validated by the strategy registry and + carry runtime controls such as scheduler assets, sampling, and stopping. + + Attributes: + name: Unique stage name. + kind: Registered strategy (see :func:`register_strategy`). + components: Names of member components (order is declaration order). + run_on: Registered phase. + options: JSON-safe, registry-validated strategy parameters. + capabilities: Capabilities this stage contributes. A stage owning a + recurrent connection must contribute + :data:`LOOP_CARRIED_STATE_CAPABILITY`. + metadata: JSON-safe extension bag preserved across round-trips. + """ + + name: str + kind: str + components: tuple[str, ...] + run_on: str = DEFAULT_PHASE + options: dict[str, JSONValue] = dataclasses.field(default_factory=dict) + capabilities: tuple[str, ...] = () + metadata: dict[str, JSONValue] = dataclasses.field(default_factory=dict) + + def __post_init__(self) -> None: + _validate_token(self.name, "Stage") + strategy = strategy_definition(self.kind) + phase_definition(self.run_on) + components = tuple(self.components) + if not components: + raise PipelineValidationError(f"Stage {self.name!r} must contain a component.") + for component in components: + _validate_component_name(component) + _check_unique(components, f"Stage {self.name!r} component") + object.__setattr__(self, "components", components) + object.__setattr__( + self, + "capabilities", + _string_tuple(self.capabilities, f"Stage {self.name!r} capabilities"), + ) + options = _validate_registered_parameters( + kind=self.kind, + parameters=self.options, + required=strategy.required_options, + allowed=strategy.allowed_options, + context=f"Stage {self.name!r}", + ) + object.__setattr__(self, "options", options) + object.__setattr__( + self, + "metadata", + _ensure_json_mapping(self.metadata, f"Stage {self.name!r} metadata"), + ) + + @property + def strategy(self) -> StrategyDefinition: + """The registered strategy definition for :attr:`kind`.""" + return strategy_definition(self.kind) + + @property + def supports_loop_carried_state(self) -> bool: + """Whether this stage may own recurrent connections.""" + return self.strategy.loop_carried_state + + def to_dict(self) -> dict[str, JSONValue]: + """Return the deterministic serializable form.""" + data: dict[str, JSONValue] = { + "name": self.name, + "kind": self.kind, + "components": list(self.components), + "run_on": self.run_on, + } + if self.options: + data["options"] = dict(self.options) + if self.capabilities: + data["capabilities"] = list(self.capabilities) + if self.metadata: + data["metadata"] = dict(self.metadata) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PipelineStage: + """Inverse of :meth:`to_dict`.""" + body = _read_mapping( + data, + {"name", "kind", "components", "run_on", "options", "capabilities", "metadata"}, + "stage", + ) + return cls( + name=str(body["name"]), + kind=str(body["kind"]), + components=tuple(body.get("components", ())), + run_on=str(body.get("run_on", DEFAULT_PHASE)), + options=dict(body.get("options", {})), + capabilities=tuple(body.get("capabilities", ())), + metadata=dict(body.get("metadata", {})), + ) + + +# --------------------------------------------------------------------------- +# Manifest +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass(frozen=True) +class PipelineManifest: + """The topology of a composed pipeline. + + The manifest is fully validated on construction, so any instance is + structurally sound. It is deterministic: components, connections, and + inputs are canonically ordered, while stages and public outputs keep their + declaration order because that order is meaningful. + + A manifest with :attr:`profile` is an executable contract: every declared + input has semantics, generated inputs name a registered program, recurrent + edges have state lifecycle, iterative/autoregressive stages carry control + parameters, assets are explicit, and components provide dtype/EP hints. + """ + + components: tuple[PipelineComponent, ...] = () + connections: tuple[PipelineConnection, ...] = () + stages: tuple[PipelineStage, ...] = () + inputs: tuple[PipelineInput, ...] = () + outputs: tuple[PipelineOutput, ...] = () + assets: tuple[PipelineAsset, ...] = () + states: tuple[PipelineState, ...] = () + profile: PipelineProfile | None = None + required_capabilities: tuple[str, ...] = () + schema_version: str = PIPELINE_SCHEMA_VERSION + metadata: dict[str, JSONValue] = dataclasses.field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__( + self, "components", tuple(sorted(self.components, key=lambda c: c.name)) + ) + object.__setattr__( + self, "connections", tuple(sorted(self.connections, key=lambda c: c.sort_key)) + ) + object.__setattr__( + self, "inputs", tuple(sorted(self.inputs, key=lambda i: i.port.qualified)) + ) + object.__setattr__(self, "assets", tuple(sorted(self.assets, key=lambda a: a.path))) + object.__setattr__(self, "states", tuple(sorted(self.states, key=lambda s: s.name))) + object.__setattr__(self, "stages", tuple(self.stages)) + object.__setattr__(self, "outputs", tuple(self.outputs)) + object.__setattr__( + self, + "required_capabilities", + _string_tuple(self.required_capabilities, "Required capabilities"), + ) + object.__setattr__( + self, "metadata", _ensure_json_mapping(self.metadata, "Manifest metadata") + ) + _parse_schema_version(self.schema_version) + self.validate() + + # -- Lookup ------------------------------------------------------------ + + @property + def component_names(self) -> tuple[str, ...]: + """Component names in canonical (sorted) order.""" + return tuple(component.name for component in self.components) + + def component(self, name: str) -> PipelineComponent: + """Return the component named *name*.""" + for component in self.components: + if component.name == name: + return component + raise KeyError(name) + + @property + def external_inputs(self) -> tuple[PipelineInput, ...]: + """Inputs the caller must supply.""" + return tuple(i for i in self.inputs if i.kind == InputSource.EXTERNAL) + + def inputs_of_kind(self, kind: str) -> tuple[PipelineInput, ...]: + """Return declared inputs whose source classification is *kind*.""" + return tuple(i for i in self.inputs if i.kind == kind) + + def stage_of(self, component: str) -> tuple[PipelineStage, ...]: + """Return the stages a component belongs to.""" + return tuple(stage for stage in self.stages if component in stage.components) + + def source_of(self, port: PipelinePort | str) -> str: + """Return the :class:`InputSource` classification of an input endpoint.""" + endpoint = PipelinePort.parse(port) + for connection in self.connections: + if connection.target == endpoint and connection.recurrent: + return InputSource.STATEFUL + for connection in self.connections: + if connection.target == endpoint: + return InputSource.DATAFLOW + for declared in self.inputs: + if declared.port == endpoint: + return declared.kind + raise KeyError(endpoint.qualified) + + def initial_source_of(self, port: PipelinePort | str) -> str: + """Return how a possibly recurrent input receives its first value.""" + endpoint = PipelinePort.parse(port) + for connection in self.connections: + if connection.target == endpoint and not connection.recurrent: + return InputSource.DATAFLOW + for declared in self.inputs: + if declared.port == endpoint: + return declared.kind + raise KeyError(endpoint.qualified) + + # -- Validation -------------------------------------------------------- + + def validate(self) -> None: + """Run full structural validation (called on construction).""" + by_name = {component.name: component for component in self.components} + if len(by_name) != len(self.components): + _check_unique((c.name for c in self.components), "Component") + _check_unique_casefold( + (component.name for component in self.components), + "Component", + ) + + self._validate_connections(by_name) + self._validate_stages(by_name) + self._validate_input_sources(by_name) + self._validate_states(by_name) + self._validate_outputs(by_name) + self._validate_assets() + self._validate_runtime_references(by_name) + self._validate_acyclic() + self._validate_capabilities() + self._validate_profile_contract() + + def _endpoint_spec( + self, + by_name: Mapping[str, PipelineComponent], + port: PipelinePort, + direction: str, + context: str, + ) -> TensorSpec: + component = by_name.get(port.component) + if component is None: + known = ", ".join(sorted(by_name)) or "" + raise PipelineValidationError( + f"{context} references unknown component {port.component!r}. " + f"Known components: {known}." + ) + spec = ( + component.input(port.port) if direction == "input" else component.output(port.port) + ) + if spec is None: + available = ", ".join( + s.name + for s in (component.inputs if direction == "input" else component.outputs) + ) + raise PipelineValidationError( + f"{context} references unknown {direction} {port.port!r} on component " + f"{port.component!r}. Available {direction}s: {available or ''}." + ) + return spec + + def _validate_connections(self, by_name: Mapping[str, PipelineComponent]) -> None: + producers: dict[tuple[str, bool], PipelineConnection] = {} + for connection in self.connections: + context = f"Connection {connection.source} -> {connection.target}" + source_spec = self._endpoint_spec(by_name, connection.source, "output", context) + target_spec = self._endpoint_spec(by_name, connection.target, "input", context) + producer_key = (connection.target.qualified, connection.recurrent) + existing = producers.get(producer_key) + if existing is not None: + lifecycle = "recurrent" if connection.recurrent else "initial" + raise PipelineValidationError( + f"Input {connection.target} has more than one {lifecycle} producer " + f"({existing.source} and {connection.source}); an input accepts " + "at most one producer for each lifecycle phase " + "(fan-out on outputs is fine)." + ) + producers[producer_key] = connection + for context_port in connection.context: + component = by_name.get(context_port.component) + if component is None: + raise PipelineValidationError( + f"{context} transform context references unknown component " + f"{context_port.component!r}." + ) + if ( + component.input(context_port.port) is None + and component.output(context_port.port) is None + ): + raise PipelineValidationError( + f"{context} transform context references unknown port " + f"{context_port.qualified!r}." + ) + if connection.transform is not None: + # A registered transform may legitimately change dtype, rank, + # and shape, and the core does not execute it, so port + # compatibility is not assumed here. The transform kind itself + # was validated against the registry on construction, and its + # runtime capabilities are checked in _validate_capabilities. + continue + mismatch = _tensor_mismatch(source_spec, target_spec) + if mismatch is not None: + raise PipelineValidationError( + f"{context} is incompatible: {mismatch} " + f"(source {source_spec.dtype}{list(source_spec.shape)}, " + f"target {target_spec.dtype}{list(target_spec.shape)}). " + "Declare a transform if a runtime adapts this edge." + ) + + def _validate_stages(self, by_name: Mapping[str, PipelineComponent]) -> None: + _check_unique((stage.name for stage in self.stages), "Stage") + staged: set[str] = set() + for stage in self.stages: + for name in stage.components: + component = by_name.get(name) + if component is None: + known = ", ".join(sorted(by_name)) or "" + raise PipelineValidationError( + f"Stage {stage.name!r} references unknown component {name!r}. " + f"Known components: {known}." + ) + if ( + stage.run_on != DEFAULT_PHASE + and component.run_on != DEFAULT_PHASE + and stage.run_on != component.run_on + ): + raise PipelineValidationError( + f"Component {name!r} runs on {component.run_on!r} but stage " + f"{stage.name!r} runs on {stage.run_on!r}; the component could " + "never execute in this stage." + ) + staged.add(name) + if stage.kind == "on_demand" and not any( + by_name[name].presence is not None for name in stage.components + ): + raise PipelineValidationError( + f"On-demand stage {stage.name!r} has no component presence " + "condition, so a runtime cannot decide when to execute it." + ) + missing = sorted(set(by_name) - staged) + if missing: + names = ", ".join(repr(n) for n in missing) + raise PipelineValidationError( + f"Component(s) {names} belong to no declared stage; every component " + "must be reachable through a stage." + ) + + stages_by_component: dict[str, list[PipelineStage]] = {} + for stage in self.stages: + for name in stage.components: + stages_by_component.setdefault(name, []).append(stage) + + for connection in self.connections: + if not connection.recurrent: + continue + source_stages = stages_by_component.get(connection.source.component, []) + target_stages = stages_by_component.get(connection.target.component, []) + shared = [ + stage + for stage in source_stages + if stage in target_stages and stage.supports_loop_carried_state + ] + if not shared: + kinds = ", ".join( + sorted(k for k, v in _STRATEGIES.items() if v.loop_carried_state) + ) + raise PipelineValidationError( + f"Recurrent connection {connection.source} -> {connection.target} " + "must be scoped to a single stage that both components belong to " + f"and whose strategy supports loop-carried state ({kinds})." + ) + without_capability = [ + stage.name + for stage in shared + if LOOP_CARRIED_STATE_CAPABILITY not in stage.capabilities + ] + if len(without_capability) == len(shared): + raise PipelineValidationError( + f"Recurrent connection {connection.source} -> {connection.target} " + f"requires its stage to contribute the " + f"{LOOP_CARRIED_STATE_CAPABILITY!r} capability; " + f"stage(s) {', '.join(repr(s) for s in without_capability)} do not." + ) + + def _validate_input_sources(self, by_name: Mapping[str, PipelineComponent]) -> None: + declared: dict[str, PipelineInput] = {} + for entry in self.inputs: + context = f"Declared input {entry.port}" + self._endpoint_spec(by_name, entry.port, "input", context) + if entry.port.qualified in declared: + raise PipelineValidationError( + f"{context} is declared more than once; each input needs exactly " + "one source classification." + ) + declared[entry.port.qualified] = entry + + aliases = [entry.name for entry in self.external_inputs] + _check_unique(aliases, "External input name") + + initial_connections = { + connection.target.qualified: connection + for connection in self.connections + if not connection.recurrent + } + recurrent_connections = { + connection.target.qualified: connection + for connection in self.connections + if connection.recurrent + } + for component in self.components: + for spec in component.inputs: + endpoint = PipelinePort(component.name, spec.name) + key = endpoint.qualified + initial = initial_connections.get(key) + recurrent = recurrent_connections.get(key) + entry = declared.get(key) + if initial is not None and entry is not None: + raise PipelineValidationError( + f"Input {key} is both initialized by {initial.source} and " + f"declared as {entry.kind!r}; exactly one initial source is allowed." + ) + if initial is None and entry is None: + kinds = ", ".join(InputSource.DECLARABLE) + prefix = "Recurrent input" if recurrent is not None else "Input" + raise PipelineValidationError( + f"{prefix} {key} has no initial source. Connect it, or declare it as one " + f"of: {kinds}. Unconnected inputs are never assumed to be " + "external." + ) + + def _validate_outputs(self, by_name: Mapping[str, PipelineComponent]) -> None: + state_names = {state.name for state in self.states} + for output in self.outputs: + if output.port is not None: + self._endpoint_spec( + by_name, + output.port, + "output", + f"Public output {output.port}", + ) + elif output.state not in state_names: + raise PipelineValidationError( + f"Public output references unknown state {output.state!r}." + ) + _check_unique((output.name for output in self.outputs), "Public output name") + + def _validate_states(self, by_name: Mapping[str, PipelineComponent]) -> None: + _check_unique((state.name for state in self.states), "State") + _check_unique((state.input.qualified for state in self.states), "State input") + stages = {stage.name for stage in self.stages} + recurrent = { + (connection.source.qualified, connection.target.qualified) + for connection in self.connections + if connection.recurrent + } + declared: set[tuple[str, str]] = set() + for state in self.states: + input_spec = self._endpoint_spec( + by_name, + state.input, + "input", + f"State {state.name!r}", + ) + self._endpoint_spec( + by_name, + state.output, + "output", + f"State {state.name!r}", + ) + edge = (state.output.qualified, state.input.qualified) + if edge not in recurrent: + raise PipelineValidationError( + f"State {state.name!r} does not match a recurrent connection " + f"{state.output} -> {state.input}." + ) + if state.release_after not in stages: + raise PipelineValidationError( + f"State {state.name!r} releases after unknown stage " + f"{state.release_after!r}." + ) + if state.sequence_axis is not None and state.sequence_axis >= input_spec.rank: + raise PipelineValidationError( + f"State {state.name!r} sequence_axis {state.sequence_axis} " + f"is outside input rank {input_spec.rank}." + ) + declared.add(edge) + missing = sorted(recurrent - declared) + if self.profile is not None and missing: + raise PipelineValidationError( + "Every recurrent connection requires an explicit state lifecycle; " + f"missing declarations for {missing}." + ) + + def _validate_assets(self) -> None: + _check_unique((asset.path for asset in self.assets), "Asset destination") + _check_unique_casefold( + (asset.path for asset in self.assets), + "Asset destination", + ) + reserved = {PIPELINE_FILENAME, *self.component_file_layout().values()} + reserved |= {f"{path}.data" for path in self.component_file_layout().values()} + reserved_casefold = {path.casefold() for path in reserved} + for asset in self.assets: + if asset.path.casefold() in reserved_casefold: + raise PipelineValidationError( + f"Asset destination {asset.path!r} collides with a file written by " + "the package itself." + ) + + def _validate_runtime_references( + self, + by_name: Mapping[str, PipelineComponent], + ) -> None: + state_names = {state.name for state in self.states} + stage_names = {stage.name for stage in self.stages} + asset_paths = {asset.path for asset in self.assets} + + def validate_port(value: Any, context: str) -> None: + if not isinstance(value, str) or "." not in value: + raise PipelineValidationError( + f"{context} must reference a qualified component port." + ) + port = PipelinePort.parse(value) + component = by_name.get(port.component) + if component is None or ( + component.input(port.port) is None and component.output(port.port) is None + ): + raise PipelineValidationError(f"{context} references unknown port {value!r}.") + + for connection in self.connections: + parameters = connection.parameters + if "state" in parameters and parameters["state"] not in state_names: + raise PipelineValidationError( + f"Connection {connection.source} -> {connection.target} references " + f"unknown state {parameters['state']!r}." + ) + if "stage" in parameters and parameters["stage"] not in stage_names: + raise PipelineValidationError( + f"Connection {connection.source} -> {connection.target} references " + f"unknown stage {parameters['stage']!r}." + ) + if ( + "scheduler_asset" in parameters + and parameters["scheduler_asset"] not in asset_paths + ): + raise PipelineValidationError( + f"Connection {connection.source} -> {connection.target} references " + f"undeclared scheduler asset {parameters['scheduler_asset']!r}." + ) + if "timestep_input" in parameters: + validate_port( + parameters["timestep_input"], + f"Connection {connection.source} timestep_input", + ) + + for entry in self.inputs: + if entry.generator is None: + continue + parameters = entry.generator.parameters + if "stage" in parameters and parameters["stage"] not in stage_names: + raise PipelineValidationError( + f"Generated input {entry.port} references unknown stage " + f"{parameters['stage']!r}." + ) + for key in ("source", "sequence_input"): + if key in parameters: + validate_port(parameters[key], f"Generated input {entry.port} {key}") + if "past_state" in parameters: + values = parameters["past_state"] + if not isinstance(values, list) or not all( + isinstance(value, str) and value in state_names for value in values + ): + raise PipelineValidationError( + f"Generated input {entry.port} references unknown past state." + ) + if "dynamic_axes" in parameters: + dynamic_axes = parameters["dynamic_axes"] + input_spec = self._endpoint_spec( + by_name, + entry.port, + "input", + f"Generated input {entry.port}", + ) + symbolic_dims = {dim for dim in input_spec.shape if isinstance(dim, str)} + if ( + not isinstance(dynamic_axes, dict) + or not set(dynamic_axes) <= symbolic_dims + ): + raise PipelineValidationError( + f"Generated input {entry.port} references unknown dynamic axis." + ) + + def component_file_layout(self) -> dict[str, str]: + """Return ``{component: relative onnx path}`` for the saved layout. + + Mirrors :meth:`ModelPackage.save`: a single-component package is stored + flat as ``model.onnx``; otherwise each component gets its own folder. + """ + names = self.component_names + if len(names) == 1: + return {names[0]: "model.onnx"} + return {name: f"{name}/model.onnx" for name in names} + + @property + def required_assets(self) -> tuple[PipelineAsset, ...]: + """Assets that must be present for the package to load.""" + return tuple(asset for asset in self.assets if asset.required) + + def _validate_acyclic(self) -> None: + """Non-recurrent edges must form a DAG (hence acyclic within a stage).""" + edges: dict[str, set[str]] = {c.name: set() for c in self.components} + for connection in self.connections: + if connection.recurrent: + continue + if connection.source.component == connection.target.component: + raise PipelineValidationError( + f"Connection {connection.source} -> {connection.target} makes " + f"component {connection.source.component!r} depend on itself; " + "mark the edge recurrent if it is loop-carried state." + ) + edges[connection.source.component].add(connection.target.component) + + visiting: set[str] = set() + visited: set[str] = set() + stack: list[str] = [] + + def visit(node: str) -> None: + if node in visited: + return + if node in visiting: + cycle = " -> ".join([*stack[stack.index(node) :], node]) + raise PipelineValidationError( + f"Pipeline has a cycle in non-recurrent connections: {cycle}. " + "Mark loop-carried edges recurrent and scope them to an " + "iterative stage." + ) + visiting.add(node) + stack.append(node) + for successor in sorted(edges[node]): + visit(successor) + stack.pop() + visiting.discard(node) + visited.add(node) + + for name in sorted(edges): + visit(name) + + def _validate_capabilities(self) -> None: + provided = set() + for component in self.components: + provided.update(component.capabilities) + for stage in self.stages: + provided.update(stage.capabilities) + for entry in self.inputs: + if entry.generator is None: + continue + required = entry.generator.capabilities + provided.update(required) + undeclared = sorted(set(required) - set(self.required_capabilities)) + if undeclared: + raise PipelineValidationError( + f"Generated input {entry.port} requires undeclared capabilities " + f"{undeclared}." + ) + + # A transformed edge contributes the capabilities its transform kind + # declares, and the manifest must require them explicitly so that the + # runtime obligation is visible without inspecting every connection. + for connection in self.connections: + required = connection.transform_capabilities + if not required: + continue + provided.update(required) + undeclared = sorted(set(required) - set(self.required_capabilities)) + if undeclared: + names = ", ".join(repr(c) for c in undeclared) + raise PipelineValidationError( + f"Connection {connection.source} -> {connection.target} uses " + f"transform {connection.transform!r}, which requires capabilities " + f"{names}; add them to the manifest's required_capabilities." + ) + + missing = sorted(set(self.required_capabilities) - provided) + if missing: + names = ", ".join(repr(m) for m in missing) + raise PipelineValidationError( + f"Required capabilities {names} are not provided by any component or stage." + ) + + def _validate_profile_contract(self) -> None: + """Require executable semantics when a runtime profile is declared.""" + if self.profile is None: + return + missing_semantics = [ + entry.port.qualified for entry in self.inputs if entry.semantic is None + ] + if missing_semantics: + raise PipelineValidationError( + f"Profile {self.profile.name!r} requires semantic names for every " + f"declared input; missing {missing_semantics}." + ) + missing_ep_hints = [ + component.name + for component in self.components + if not component.preferred_execution_providers + ] + if missing_ep_hints: + raise PipelineValidationError( + f"Profile {self.profile.name!r} requires execution-provider hints " + f"for every component; missing {missing_ep_hints}." + ) + missing_dtypes = [ + component.name + for component in self.components + if component.parameter_dtype is None + ] + if missing_dtypes: + raise PipelineValidationError( + f"Profile {self.profile.name!r} requires parameter dtype for every " + f"component; missing {missing_dtypes}." + ) + for stage in self.stages: + if stage.kind == "autoregressive": + required = {"tokenizer_asset", "sampling", "stop"} + elif stage.kind == "iterative": + required = {"scheduler", "default_steps", "timestep", "state_inputs"} + else: + continue + missing = sorted(required - set(stage.options)) + if missing: + raise PipelineValidationError( + f"Executable profile stage {stage.name!r} ({stage.kind}) is " + f"missing control option(s) {missing}." + ) + asset_paths = {asset.path for asset in self.assets} + referenced_assets: set[str] = set() + for stage in self.stages: + tokenizer_asset = stage.options.get("tokenizer_asset") + if isinstance(tokenizer_asset, str): + referenced_assets.add(tokenizer_asset) + scheduler = stage.options.get("scheduler") + if isinstance(scheduler, dict): + config_asset = scheduler.get("config_asset") + if isinstance(config_asset, str): + referenced_assets.add(config_asset) + missing_assets = sorted(referenced_assets - asset_paths) + if missing_assets: + raise PipelineValidationError( + f"Profile {self.profile.name!r} references undeclared runtime assets " + f"{missing_assets}." + ) + + def validate_models(self, models: Mapping[str, ir.Model]) -> None: + """Check that *models* match the manifest exactly (names and signatures).""" + manifest_names = set(self.component_names) + model_names = set(models) + if manifest_names != model_names: + missing = ", ".join(repr(n) for n in sorted(manifest_names - model_names)) or "-" + extra = ", ".join(repr(n) for n in sorted(model_names - manifest_names)) or "-" + raise PipelineValidationError( + f"Pipeline models do not match the manifest. Missing: {missing}. " + f"Unexpected: {extra}." + ) + for component in self.components: + actual = PipelineComponent.from_model( + component.name, models[component.name], role=component.role + ) + if actual.inputs != component.inputs or actual.outputs != component.outputs: + raise PipelineValidationError( + f"Component {component.name!r} graph signature does not match the " + "manifest; the manifest is out of date with the graph." + ) + + # -- Serialization ----------------------------------------------------- + + def to_dict(self) -> dict[str, JSONValue]: + """Return the deterministic serializable form.""" + data: dict[str, JSONValue] = { + "schema_version": self.schema_version, + "components": [component.to_dict() for component in self.components], + "connections": [connection.to_dict() for connection in self.connections], + "stages": [stage.to_dict() for stage in self.stages], + "inputs": [entry.to_dict() for entry in self.inputs], + "outputs": [output.to_dict() for output in self.outputs], + } + if self.profile is not None: + data["profile"] = self.profile.to_dict() + if self.states: + data["states"] = [state.to_dict() for state in self.states] + if self.assets: + data["assets"] = [asset.to_dict() for asset in self.assets] + if self.required_capabilities: + data["required_capabilities"] = list(self.required_capabilities) + if self.metadata: + data["metadata"] = dict(self.metadata) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> PipelineManifest: + """Inverse of :meth:`to_dict`. + + Unknown top-level keys, unknown roles/strategies/phases, and an unknown + schema *major* version are hard failures. Unknown keys nested inside a + ``metadata`` bag are preserved verbatim. + """ + body = _read_mapping( + data, + { + "schema_version", + "components", + "connections", + "stages", + "inputs", + "outputs", + "assets", + "states", + "profile", + "required_capabilities", + "metadata", + }, + "manifest", + ) + version = str(body.get("schema_version", PIPELINE_SCHEMA_VERSION)) + _parse_schema_version(version) + return cls( + components=tuple( + PipelineComponent.from_dict(c) for c in body.get("components", []) + ), + connections=tuple( + PipelineConnection.from_dict(c) for c in body.get("connections", []) + ), + stages=tuple(PipelineStage.from_dict(s) for s in body.get("stages", [])), + inputs=tuple(PipelineInput.from_dict(i) for i in body.get("inputs", [])), + outputs=tuple(PipelineOutput.from_dict(o) for o in body.get("outputs", [])), + assets=tuple(PipelineAsset.from_dict(a) for a in body.get("assets", [])), + states=tuple(PipelineState.from_dict(s) for s in body.get("states", [])), + profile=( + PipelineProfile.from_dict(body["profile"]) + if body.get("profile") is not None + else None + ), + required_capabilities=tuple(body.get("required_capabilities", ())), + schema_version=version, + metadata=dict(body.get("metadata", {})), + ) + + def to_json(self, *, indent: int | None = 2) -> str: + """Serialize to a deterministic JSON string.""" + return json.dumps(self.to_dict(), indent=indent, sort_keys=False) + + @classmethod + def from_json(cls, text: str) -> PipelineManifest: + """Inverse of :meth:`to_json`.""" + return cls.from_dict(json.loads(text)) + + +def _parse_version_format(version: str, context: str) -> tuple[int, int]: + """Parse a non-negative ``major.minor`` version without compatibility checks.""" + if not isinstance(version, str): + raise PipelineValidationError(f"{context} version {version!r} must be a string.") + parts = version.split(".") + try: + major = int(parts[0]) + minor = int(parts[1]) if len(parts) > 1 else 0 + except (ValueError, IndexError) as error: + raise PipelineValidationError( + f"{context} version {version!r} must have the form 'major.minor'." + ) from error + if len(parts) > 2 or major < 0 or minor < 0: + raise PipelineValidationError( + f"{context} version {version!r} must have the form 'major.minor'." + ) + return major, minor + + +def _parse_schema_version(version: str) -> tuple[int, int]: + """Parse and compatibility-check a pipeline schema version.""" + major, minor = _parse_version_format(version, "Schema") + current_major = int(PIPELINE_SCHEMA_VERSION.split(".")[0]) + if major != current_major: + raise PipelineValidationError( + f"Unsupported pipeline schema version {version!r}; this build understands " + f"major version {current_major}." + ) + return major, minor + + +def _check_unique(names: Iterable[str], what: str) -> None: + seen: set[str] = set() + for name in names: + if name in seen: + raise PipelineValidationError(f"{what} {name!r} is declared more than once.") + seen.add(name) + + +def _check_unique_casefold(names: Iterable[str], what: str) -> None: + """Reject names that collide on case-insensitive filesystems.""" + seen: dict[str, str] = {} + for name in names: + folded = name.casefold() + existing = seen.get(folded) + if existing is not None and existing != name: + raise PipelineValidationError( + f"{what} names {existing!r} and {name!r} collide on " + "case-insensitive filesystems." + ) + seen[folded] = name + + +def _read_mapping(data: Mapping[str, Any], allowed: set[str], what: str) -> dict[str, Any]: + """Return a plain dict of *data*, rejecting unknown keys.""" + if not isinstance(data, Mapping): + raise PipelineValidationError( + f"Expected a mapping for {what}, got {type(data).__name__}." + ) + unknown = sorted(set(data) - allowed) + if unknown: + keys = ", ".join(repr(k) for k in unknown) + raise PipelineValidationError( + f"Unknown key(s) {keys} in {what}. Put forward-compatible extensions in " + "the 'metadata' field." + ) + return dict(data) + + +# --------------------------------------------------------------------------- +# Package +# --------------------------------------------------------------------------- + + +class PipelinePackage(ModelPackage): + """A :class:`ModelPackage` that also carries a :class:`PipelineManifest`. + + The package keeps the ``ModelPackage`` on-disk layout (``model.onnx`` for a + single component, ``{name}/model.onnx`` otherwise) and adds + ``pipeline.json`` describing the topology and the component filenames. + + It may also carry *runtime assets* — tokenizers, scheduler configs, + processor configs — as a mapping of safe relative destination to an + existing local source file. Assets are opaque: they are copied and their + destinations are recorded, and nothing in this module ever reads them. + + Attributes: + manifest: The validated topology. + config: The primary configuration, exactly as on ``ModelPackage``. + component_configs: Optional per-component configuration objects. + assets: Mapping of manifest-declared destination to the local source + file it is copied from. After :meth:`load`, the sources are the + resolved paths inside the loaded directory. + """ + + def __init__( + self, + models: Mapping[str, ir.Model] | None = None, + manifest: PipelineManifest | None = None, + config: object | None = None, + component_configs: Mapping[str, object] | None = None, + assets: Mapping[str, str] | None = None, + ) -> None: + if manifest is None: + manifest = PipelineManifest() + if models is None: + models = { + component.name: component.model + for component in manifest.components + if component.model is not None + } + super().__init__(dict(models), config=config) + self.manifest = manifest + self.component_configs: dict[str, object] = dict(component_configs or {}) + unknown = sorted(set(self.component_configs) - set(manifest.component_names)) + if unknown: + names = ", ".join(repr(n) for n in unknown) + raise PipelineValidationError( + f"Per-component config given for unknown component(s) {names}." + ) + self.assets: dict[str, str] = dict(assets or {}) + self._validate_assets() + manifest.validate_models(self.data) + + def _validate_assets(self) -> None: + """Check asset sources against the manifest's declared destinations.""" + declared = {asset.path: asset for asset in self.manifest.assets} + undeclared = sorted(set(self.assets) - set(declared)) + if undeclared: + names = ", ".join(repr(n) for n in undeclared) + raise PipelineValidationError( + f"Asset source(s) given for undeclared destination(s) {names}; declare " + "them on the manifest so that pipeline.json stays in sync with the " + "saved directory." + ) + missing = sorted( + path + for path, asset in declared.items() + if asset.required and path not in self.assets + ) + if missing: + names = ", ".join(repr(n) for n in missing) + raise PipelineValidationError( + f"Required asset(s) {names} are declared by the manifest but have no " + "source file." + ) + for destination, source in self.assets.items(): + _validate_asset_path(destination) + if not isinstance(source, str) or not source: + raise PipelineValidationError( + f"Asset {destination!r} must map to a local file path." + ) + if not os.path.isfile(source): + raise PipelineValidationError( + f"Asset {destination!r} source {source!r} does not exist." + ) + + def asset_path(self, destination: str) -> str: + """Return the local source path for a declared asset destination.""" + if destination not in self.assets: + known = ", ".join(sorted(self.assets)) or "" + raise KeyError(f"No asset {destination!r}; known assets: {known}.") + return self.assets[destination] + + def __repr__(self) -> str: + names = ", ".join(repr(k) for k in self.data) + return f"PipelinePackage({{{names}}}, stages={len(self.manifest.stages)})" + + def config_for(self, component: str) -> object | None: + """Return the component config, falling back to the primary config.""" + if component in self.component_configs: + return self.component_configs[component] + return self.config + + def component_files(self) -> dict[str, str]: + """Return ``{component: relative onnx path}`` for the saved layout. + + Mirrors :meth:`ModelPackage.save`: a single-component package is stored + flat as ``model.onnx``; otherwise each component gets its own folder. + """ + return self.manifest.component_file_layout() + + def to_dict(self) -> dict[str, JSONValue]: + """Return the ``pipeline.json`` document.""" + return { + "format": "mobius-pipeline", + "schema_version": self.manifest.schema_version, + "manifest": self.manifest.to_dict(), + "component_files": dict(sorted(self.component_files().items())), + } + + def save( + self, + directory: str, + *, + external_data: str = "onnx", + max_shard_size_bytes: int | None = None, + components: Callable[[str], bool] | None = None, + progress_bar: bool = True, + check_weights: bool = True, + ) -> None: + """Save every component, every declared asset, and ``pipeline.json``. + + Assets are validated (safe relative destination, existing source) before + anything is written, then copied one at a time via a temporary file and + an atomic rename, so a reader never observes a half-written asset. + ``pipeline.json`` is written last: its presence marks a complete + directory. + + Raises: + PipelineValidationError: If *components* is given (a partial save + would desynchronize the manifest from the saved graphs), if the + in-memory models no longer match the manifest, or if an asset + destination is unsafe or its source has disappeared. + """ + if components is not None: + raise PipelineValidationError( + "PipelinePackage.save() does not support partial saves: the manifest " + "describes every component, so writing a subset would produce a " + "directory whose pipeline.json references missing graphs. Build a " + "smaller pipeline instead." + ) + self.manifest.validate_models(self.data) + self._validate_assets() + marker = os.path.join(directory, PIPELINE_FILENAME) + if os.path.isfile(marker): + os.remove(marker) + super().save( + directory, + external_data=external_data, + max_shard_size_bytes=max_shard_size_bytes, + progress_bar=progress_bar, + check_weights=check_weights, + ) + self._copy_assets(directory) + document = self.to_dict() + handle, staged = tempfile.mkstemp(dir=directory, prefix=".mobius-pipeline-") + try: + with os.fdopen(handle, "w", encoding="utf-8") as file: + json.dump(document, file, indent=2) + file.write("\n") + os.replace(staged, marker) + except BaseException: + if os.path.exists(staged): + os.remove(staged) + raise + files = document["component_files"] + assert isinstance(files, dict) + for name, relative in files.items(): + expected = os.path.join(directory, *str(relative).split("/")) + if not os.path.isfile(expected): + raise PipelineValidationError( + f"Component {name!r} was not written to {expected!r}; the saved " + "layout does not match pipeline.json." + ) + for asset in self.manifest.assets: + written = os.path.join(directory, *asset.path.split("/")) + if asset.required and not os.path.isfile(written): + raise PipelineValidationError( + f"Required asset {asset.path!r} was not written to {written!r}." + ) + + def _copy_assets(self, directory: str) -> None: + """Copy every declared asset into *directory*, one atomic rename each.""" + root = os.path.abspath(directory) + for destination in sorted(self.assets): + source = self.assets[destination] + target = os.path.abspath(os.path.join(root, *destination.split("/"))) + # Defence in depth: the destination was validated as a safe relative + # path, so the resolved target must still be inside the package. + if os.path.commonpath([root, target]) != root: + raise PipelineValidationError( + f"Asset destination {destination!r} resolves outside the package " + f"directory ({target!r})." + ) + if os.path.abspath(source) == target: + continue + os.makedirs(os.path.dirname(target), exist_ok=True) + handle, staged = tempfile.mkstemp( + dir=os.path.dirname(target), prefix=".mobius-asset-" + ) + os.close(handle) + try: + shutil.copyfile(source, staged) + os.replace(staged, target) + except BaseException: + if os.path.exists(staged): + os.remove(staged) + raise + + @classmethod + def load(cls, directory: str) -> PipelinePackage: + """Load a pipeline directory written by :meth:`save`. + + Raises: + PipelineValidationError: If ``pipeline.json`` is missing, refers to + components that are absent from the directory, or declares a + required asset that is not present. + """ + path = os.path.join(directory, PIPELINE_FILENAME) + if not os.path.isfile(path): + raise PipelineValidationError( + f"{path!r} not found; use ModelPackage.load() for plain model directories." + ) + try: + with open(path, encoding="utf-8") as file: + document = json.load(file) + except (OSError, json.JSONDecodeError) as error: + raise PipelineValidationError( + f"Could not read a valid {PIPELINE_FILENAME!r} from {directory!r}." + ) from error + body = _read_mapping( + document, + {"format", "schema_version", "manifest", "component_files"}, + PIPELINE_FILENAME, + ) + if body.get("format") != "mobius-pipeline": + raise PipelineValidationError( + f"{PIPELINE_FILENAME!r} has unsupported format {body.get('format')!r}." + ) + _parse_schema_version(str(body.get("schema_version", PIPELINE_SCHEMA_VERSION))) + manifest = PipelineManifest.from_dict(body["manifest"]) + files = body.get("component_files") or {} + if not isinstance(files, Mapping): + raise PipelineValidationError("'component_files' must be a mapping.") + expected_files = manifest.component_file_layout() + if dict(files) != expected_files: + raise PipelineValidationError( + "'component_files' must exactly match the safe layout derived from " + "the manifest." + ) + models: dict[str, ir.Model] = {} + for name in manifest.component_names: + model_path = os.path.join(directory, *str(files[name]).split("/")) + if not os.path.isfile(model_path): + raise PipelineValidationError( + f"Component {name!r} file {model_path!r} is missing." + ) + models[name] = ir.load(model_path) + assets: dict[str, str] = {} + for asset in manifest.assets: + resolved = os.path.join(directory, *asset.path.split("/")) + if os.path.isfile(resolved): + assets[asset.path] = resolved + elif asset.required: + raise PipelineValidationError( + f"Required asset {asset.path!r} is missing from {directory!r}." + ) + manifest = dataclasses.replace( + manifest, + components=tuple( + component.with_model(models[component.name]) + for component in manifest.components + ), + ) + return cls(models, manifest, assets=assets) + + +# --------------------------------------------------------------------------- +# Builder +# --------------------------------------------------------------------------- + + +class PipelineBuilder: + """Compose already-built graphs into a validated :class:`PipelinePackage`. + + The builder never builds, traces, or optimizes a graph: every component is + an ``ir.Model`` that already exists. All structural validation happens in + :meth:`build`. + """ + + def __init__(self, *, schema_version: str = PIPELINE_SCHEMA_VERSION) -> None: + _parse_schema_version(schema_version) + self._schema_version = schema_version + self._components: dict[str, PipelineComponent] = {} + self._connections: list[PipelineConnection] = [] + self._inputs: list[PipelineInput] = [] + self._states: list[PipelineState] = [] + self._outputs: list[PipelineOutput] = [] + self._stages: list[PipelineStage] = [] + self._assets: dict[str, str] = {} + self._asset_specs: dict[str, PipelineAsset] = {} + self._required_capabilities: set[str] = set() + self._metadata: dict[str, JSONValue] = {} + self._profile: PipelineProfile | None = None + + # -- Components -------------------------------------------------------- + + def add_model( + self, + name: str, + model: ir.Model, + *, + role: str, + run_on: str = DEFAULT_PHASE, + presence: str | None = None, + capabilities: Iterable[str] | None = None, + preferred_execution_providers: Iterable[str] | None = None, + parameter_dtype: str | None = None, + source: str | None = None, + config: Mapping[str, Any] | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> PipelineComponent: + """Add one graph as a component and derive its typed ports.""" + if name in self._components: + raise PipelineValidationError(f"Component {name!r} is already registered.") + component = PipelineComponent.from_model( + name, + model, + role=role, + run_on=run_on, + presence=presence, + capabilities=capabilities, + preferred_execution_providers=preferred_execution_providers, + parameter_dtype=parameter_dtype, + source=source, + config=config, + metadata=metadata, + ) + self._components[name] = component + return component + + def add_package( + self, + package: Mapping[str, ir.Model], + *, + roles: Mapping[str, str] | Callable[[str], str], + prefix: str | None = None, + run_on: str | Mapping[str, str] = DEFAULT_PHASE, + source: str | None = None, + configs: Mapping[str, Mapping[str, Any]] | None = None, + ) -> tuple[PipelineComponent, ...]: + """Add every model of a :class:`ModelPackage` as a component. + + Component names are derived deterministically: keys are visited in + sorted order and, when *prefix* is given, namespaced as + ``"{prefix}_{key}"`` — so two packages that both contain ``"model"`` + can coexist. + + Roles must be given explicitly (a mapping keyed by *package key*, or a + callable). The core deliberately does not infer a role from a literal + package key such as ``"vision"``: package keys are a storage detail and + carry no guaranteed semantics. + + Args: + package: Mapping of package key to ``ir.Model``. + roles: Mapping ``package key -> role`` or a callable of the key. + prefix: Optional namespace prefix. + run_on: A phase for every component, or a mapping keyed by package + key. + source: Optional provenance recorded on each component. + configs: Optional per-key JSON-safe component configs. + + Returns: + The created components, in sorted key order. + """ + if prefix is not None: + _validate_component_name(prefix) + created: list[PipelineComponent] = [] + for key in sorted(package): + if isinstance(roles, Mapping): + if key not in roles: + known = ", ".join(sorted(roles)) or "" + raise PipelineValidationError( + f"No role declared for package key {key!r}; roles were given " + f"for: {known}. Roles must be explicit." + ) + role = roles[key] + elif callable(roles): + role = roles(key) + else: + raise PipelineValidationError( + "roles must be a mapping of package key to role, or a callable." + ) + name = f"{prefix}_{key}" if prefix else key + phase = run_on[key] if isinstance(run_on, Mapping) else run_on + created.append( + self.add_model( + name, + package[key], + role=role, + run_on=phase, + source=source, + config=(configs or {}).get(key), + metadata={"package_key": key}, + ) + ) + return tuple(created) + + # -- Wiring ------------------------------------------------------------ + + def connect( + self, + source: str | PipelinePort, + target: str | PipelinePort, + *, + recurrent: bool = False, + transform: str | None = None, + context: Iterable[str | PipelinePort] | None = None, + parameters: Mapping[str, Any] | None = None, + ) -> PipelineConnection: + """Wire ``source`` output to ``target`` input. + + Args: + source: Producing endpoint, e.g. ``"encoder.hidden"``. + target: Consuming endpoint, e.g. ``"decoder.encoder_hidden"``. + recurrent: Mark the edge as loop-carried state. The owning stage + must support it and will be given the + :data:`LOOP_CARRIED_STATE_CAPABILITY` capability at build time. + transform: Registered transform kind for the runtime; never + executed here. Its declared capabilities are added to the + manifest's required capabilities, and the edge is exempt from + direct dtype/rank/shape compatibility because a transform may + legitimately change all three. + context: Additional component input/output endpoints consumed by + the transform, such as the current latent for a scheduler step. + parameters: JSON-safe parameters validated by the transform + definition and consumed directly by the runtime. + """ + connection = PipelineConnection( + PipelinePort.parse(source), + PipelinePort.parse(target), + recurrent, + transform, + tuple(PipelinePort.parse(port) for port in (context or ())), + _ensure_json_mapping(parameters, "Transform parameters"), + ) + self._connections.append(connection) + if recurrent: + self._required_capabilities.add(LOOP_CARRIED_STATE_CAPABILITY) + self._required_capabilities.update(connection.transform_capabilities) + return connection + + def declare_external( + self, + port: str | PipelinePort, + *, + alias: str | None = None, + semantic: str | None = None, + required: bool = True, + presence: str | None = None, + ) -> PipelineInput: + """Declare that the caller supplies this input.""" + return self._declare( + port, + InputSource.EXTERNAL, + alias=alias, + semantic=semantic, + required=required, + presence=presence, + ) + + def declare_generated( + self, + port: str | PipelinePort, + *, + generator: str, + parameters: Mapping[str, Any] | None = None, + semantic: str | None = None, + presence: str | None = None, + ) -> PipelineInput: + """Declare that the runtime harness produces this input per invocation.""" + rule = GeneratedInputRule( + generator, + _ensure_json_mapping(parameters, "Generated input parameters"), + ) + self._required_capabilities.update(rule.capabilities) + return self._declare( + port, + InputSource.GENERATED, + semantic=semantic, + presence=presence, + generator=rule, + ) + + def declare_stateful( + self, + port: str | PipelinePort, + *, + semantic: str | None = None, + ) -> PipelineInput: + """Declare that this input is runtime-owned state carried across steps.""" + return self._declare(port, InputSource.STATEFUL, semantic=semantic) + + def declare_default( + self, + port: str | PipelinePort, + value: JSONValue, + *, + semantic: str | None = None, + ) -> PipelineInput: + """Declare a JSON-safe constant for this input.""" + return self._declare( + port, + InputSource.DEFAULTED, + value=value, + semantic=semantic, + required=False, + ) + + def _declare( + self, + port: str | PipelinePort, + kind: str, + *, + value: JSONValue = None, + alias: str | None = None, + semantic: str | None = None, + required: bool = True, + presence: str | None = None, + generator: GeneratedInputRule | None = None, + ) -> PipelineInput: + entry = PipelineInput( + PipelinePort.parse(port), + kind, + value, + alias, + semantic, + required, + presence, + generator, + ) + self._inputs.append(entry) + return entry + + def add_state( + self, + name: str, + *, + kind: str, + input: str | PipelinePort, + output: str | PipelinePort, + lifetime: str, + release_after: str, + sequence_axis: int | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> PipelineState: + """Declare lifecycle semantics for one recurrent connection.""" + state = PipelineState( + name=name, + kind=kind, + input=PipelinePort.parse(input), + output=PipelinePort.parse(output), + lifetime=lifetime, + release_after=release_after, + sequence_axis=sequence_axis, + metadata=_ensure_json_mapping(metadata, f"State {name!r} metadata"), + ) + self._states.append(state) + return state + + def add_public_output( + self, port: str | PipelinePort, *, alias: str | None = None + ) -> PipelineOutput: + """Expose a component output as a pipeline result.""" + output = PipelineOutput(PipelinePort.parse(port), alias) + self._outputs.append(output) + return output + + def add_public_state_output( + self, + state: str, + *, + alias: str | None = None, + ) -> PipelineOutput: + """Expose the final value of a recurrent state as a pipeline result.""" + output = PipelineOutput(state=state, alias=alias) + self._outputs.append(output) + return output + + def add_stage( + self, + name: str, + kind: str, + components: Sequence[str], + *, + run_on: str = DEFAULT_PHASE, + options: Mapping[str, Any] | None = None, + capabilities: Iterable[str] | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> PipelineStage: + """Declare a stage over *components* with a registered strategy.""" + stage = PipelineStage( + name=name, + kind=kind, + components=tuple(components), + run_on=run_on, + options=_ensure_json_mapping(options, f"Stage {name!r} options"), + capabilities=_string_tuple(capabilities, f"Stage {name!r} capabilities"), + metadata=_ensure_json_mapping(metadata, f"Stage {name!r} metadata"), + ) + self._stages.append(stage) + return stage + + def require_capability(self, capability: str) -> None: + """Record a capability the target runtime must provide.""" + _validate_token(capability, "Capability") + self._required_capabilities.add(capability) + + def add_asset( + self, destination: str, source: str, *, required: bool = True + ) -> PipelineAsset: + """Ship an opaque runtime file with the package. + + The file's contents are never read or interpreted here — this only + records that *source* must be copied to *destination* inside the saved + package directory, and only *destination* reaches the manifest. + + Args: + destination: ``/``-separated relative path inside the package, e.g. + ``"tokenizer.json"`` or ``"scheduler/scheduler_config.json"``. + source: Path to an existing local file to copy at save time. + required: Whether :meth:`PipelinePackage.load` must find the file. + """ + _validate_asset_path(destination) + if destination in self._assets: + raise PipelineValidationError( + f"Asset destination {destination!r} is already registered." + ) + if not isinstance(source, str) or not os.path.isfile(source): + raise PipelineValidationError( + f"Asset {destination!r} source {source!r} must be an existing file." + ) + asset = PipelineAsset(destination, required=required) + self._asset_specs[destination] = asset + self._assets[destination] = source + return asset + + def set_metadata(self, key: str, value: JSONValue) -> None: + """Attach a JSON-safe manifest-level metadata entry.""" + _validate_token(key, "Metadata key") + self._metadata[key] = _ensure_json_value(value, f"Metadata {key!r}") + + def set_profile(self, name: str, version: str) -> PipelineProfile: + """Set the versioned runtime profile implemented by this package.""" + profile = PipelineProfile(name, version) + if self._profile is not None and self._profile != profile: + raise PipelineValidationError( + f"Pipeline profile is already set to {self._profile!r}." + ) + self._profile = profile + return profile + + # -- Build ------------------------------------------------------------- + + def build( + self, + *, + config: object | None = None, + component_configs: Mapping[str, object] | None = None, + ) -> PipelinePackage: + """Validate the topology and return the composed package. + + Stages that own a recurrent connection are given the + :data:`LOOP_CARRIED_STATE_CAPABILITY` capability so that the loop-carried + state is visible to a runtime inspecting the manifest. + """ + stages = tuple(self._augment_loop_stages()) + manifest = PipelineManifest( + components=tuple(self._components.values()), + connections=tuple(self._connections), + stages=stages, + inputs=tuple(self._inputs), + outputs=tuple(self._outputs), + assets=tuple(self._asset_specs.values()), + states=tuple(self._states), + profile=self._profile, + required_capabilities=tuple(self._required_capabilities), + schema_version=self._schema_version, + metadata=dict(self._metadata), + ) + models = {name: component.model for name, component in self._components.items()} + missing = [name for name, model in models.items() if model is None] + if missing: + names = ", ".join(repr(n) for n in missing) + raise PipelineValidationError(f"Component(s) {names} have no graph.") + return PipelinePackage( + {name: model for name, model in models.items() if model is not None}, + manifest, + config=config, + component_configs=component_configs, + assets=dict(self._assets), + ) + + def _augment_loop_stages(self) -> Iterator[PipelineStage]: + """Add the loop-carried-state capability to stages owning recurrent edges.""" + looping: set[str] = set() + for connection in self._connections: + if not connection.recurrent: + continue + for stage in self._stages: + if not stage.supports_loop_carried_state: + continue + if ( + connection.source.component in stage.components + and connection.target.component in stage.components + ): + looping.add(stage.name) + for stage in self._stages: + if ( + stage.name in looping + and LOOP_CARRIED_STATE_CAPABILITY not in stage.capabilities + ): + yield dataclasses.replace( + stage, + capabilities=(*stage.capabilities, LOOP_CARRIED_STATE_CAPABILITY), + ) + else: + yield stage diff --git a/src/mobius/_pipeline_test.py b/src/mobius/_pipeline_test.py new file mode 100644 index 000000000..7453f6711 --- /dev/null +++ b/src/mobius/_pipeline_test.py @@ -0,0 +1,1473 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the compositional pipeline core.""" + +from __future__ import annotations + +import dataclasses +import json +import os +from unittest import mock + +import onnx_ir as ir +import pytest + +from mobius._model_package import ModelPackage +from mobius._pipeline import ( + LOOP_CARRIED_STATE_CAPABILITY, + PIPELINE_FILENAME, + GeneratedInputRule, + InputSource, + PipelineAsset, + PipelineBuilder, + PipelineComponent, + PipelineConnection, + PipelineInput, + PipelineManifest, + PipelineOutput, + PipelinePackage, + PipelinePort, + PipelineProfile, + PipelineStage, + PipelineValidationError, + register_generated_input, + register_role, + register_strategy, + register_transform, + transform_definition, +) + +Dim = int | str + + +def _value(name: str, shape: list[Dim], dtype: ir.DataType = ir.DataType.FLOAT) -> ir.Value: + return ir.Value(name=name, type=ir.TensorType(dtype), shape=ir.Shape(shape)) + + +def _make_model( + inputs: dict[str, list[Dim]], + outputs: dict[str, list[Dim]], + *, + dtype: ir.DataType = ir.DataType.FLOAT, + name: str = "g", +) -> ir.Model: + """Build a tiny well-formed ``ir.Model`` with the requested signature.""" + input_values = [_value(n, s, dtype) for n, s in inputs.items()] + nodes = [] + output_values = [] + for out_name, shape in outputs.items(): + node = ir.Node("", "Identity", inputs=[input_values[0]], num_outputs=1) + out = node.outputs[0] + out.name = out_name + out.type = ir.TensorType(dtype) + out.shape = ir.Shape(shape) + nodes.append(node) + output_values.append(out) + graph = ir.Graph( + input_values, output_values, nodes=nodes, name=name, opset_imports={"": 24} + ) + return ir.Model(graph, ir_version=10) + + +def _encoder() -> ir.Model: + return _make_model( + {"pixel_values": [1, 3, 8, 8]}, {"image_features": ["batch", "tokens", 16]} + ) + + +def _decoder() -> ir.Model: + return _make_model( + {"image_features": ["b", "t", 16], "position_ids": ["b", "t"]}, + {"logits": ["b", "t", 32]}, + ) + + +def _simple_pipeline() -> PipelineBuilder: + """Encoder -> decoder, one external input, one generated input.""" + builder = PipelineBuilder() + builder.add_model("encoder", _encoder(), role="encoder") + builder.add_model("decoder", _decoder(), role="decoder") + builder.connect("encoder.image_features", "decoder.image_features") + builder.declare_external("encoder.pixel_values") + builder.declare_generated("decoder.position_ids", generator="zeros") + builder.add_stage("encode", "single_pass", ["encoder"]) + builder.add_stage("generate", "autoregressive", ["decoder"]) + builder.add_public_output("decoder.logits") + return builder + + +def _cosmos_style_pipeline() -> PipelineBuilder: + """VAE moments -> sampled generator tokens, with an iterative generate stage.""" + builder = PipelineBuilder() + builder.add_model( + "vae", + _make_model({"video": [1, 3, 8, 8]}, {"moments": ["b", 8, "h", "w"]}), + role="encoder", + ) + builder.add_model( + "generator", + _make_model({"tokens": ["b", "t", 16]}, {"latent": ["b", "t", 16]}), + role="dynamics", + ) + builder.connect("vae.moments", "generator.tokens", transform="sample") + builder.declare_external("vae.video") + builder.add_stage("encode", "single_pass", ["vae"]) + builder.add_stage("generate", "iterative", ["generator"]) + builder.add_public_output("generator.latent") + return builder + + +class TestPipelinePort: + def test_qualified_round_trip(self): + port = PipelinePort("decoder", "logits") + assert port.qualified == "decoder.logits" + assert PipelinePort.parse("decoder.logits") == port + + def test_dotted_port_name_survives(self): + port = PipelinePort.parse("decoder.past_key_values.0.key") + assert port.component == "decoder" + assert port.port == "past_key_values.0.key" + + def test_parse_requires_separator(self): + with pytest.raises(PipelineValidationError, match=r"component\.port"): + PipelinePort.parse("logits") + + @pytest.mark.parametrize( + "name", ["", " ", "a b ".strip() + "/x", "a\\b", "..", "a..b", "a.b", "nul"] + ) + def test_unsafe_component_names_rejected(self, name): + with pytest.raises(PipelineValidationError): + PipelinePort(name, "x") + + def test_blank_port_rejected(self): + with pytest.raises(PipelineValidationError): + PipelinePort("encoder", "") + + +class TestComposition: + def test_multiple_tiny_models(self): + pkg = _simple_pipeline().build() + assert isinstance(pkg, ModelPackage) + assert sorted(pkg) == ["decoder", "encoder"] + assert pkg.manifest.component_names == ("decoder", "encoder") + assert pkg.manifest.component("encoder").role == "encoder" + assert [o.name for o in pkg.manifest.outputs] == ["logits"] + + def test_component_wraps_its_model(self): + pkg = _simple_pipeline().build() + component = pkg.manifest.component("encoder") + assert component.model is pkg["encoder"] + assert component.inputs[0].name == "pixel_values" + assert component.inputs[0].dtype == "FLOAT" + assert component.outputs[0].shape == ("batch", "tokens", 16) + + def test_input_source_classification(self): + manifest = _simple_pipeline().build().manifest + assert manifest.source_of("decoder.image_features") == InputSource.DATAFLOW + assert manifest.source_of("encoder.pixel_values") == InputSource.EXTERNAL + assert manifest.source_of("decoder.position_ids") == InputSource.GENERATED + assert [i.port.qualified for i in manifest.external_inputs] == ["encoder.pixel_values"] + + def test_deterministic_serialization(self): + first = _simple_pipeline().build().manifest + builder = PipelineBuilder() + # Declare everything in a different order. + builder.add_model("decoder", _decoder(), role="decoder") + builder.add_model("encoder", _encoder(), role="encoder") + builder.declare_generated("decoder.position_ids", generator="zeros") + builder.declare_external("encoder.pixel_values") + builder.connect("encoder.image_features", "decoder.image_features") + builder.add_stage("encode", "single_pass", ["encoder"]) + builder.add_stage("generate", "autoregressive", ["decoder"]) + builder.add_public_output("decoder.logits") + assert builder.build().manifest.to_json() == first.to_json() + + def test_manifest_json_round_trip(self): + manifest = _simple_pipeline().build().manifest + assert PipelineManifest.from_json(manifest.to_json()) == manifest + + def test_builder_does_not_mutate_graphs(self): + model = _encoder() + before = len(model.graph) + builder = PipelineBuilder() + builder.add_model("encoder", model, role="encoder") + builder.declare_external("encoder.pixel_values") + builder.add_stage("encode", "single_pass", ["encoder"]) + builder.add_public_output("encoder.image_features") + builder.build() + assert len(model.graph) == before + + +class TestAddPackage: + def test_multi_model_package_is_namespaced(self): + package = ModelPackage( + { + "model": _make_model({"x": [1, 4]}, {"y": [1, 4]}), + "vision": _make_model({"pixels": [1, 4]}, {"feats": [1, 4]}), + } + ) + builder = PipelineBuilder() + created = builder.add_package( + package, roles={"model": "decoder", "vision": "encoder"}, prefix="left" + ) + assert [c.name for c in created] == ["left_model", "left_vision"] + assert created[0].role == "decoder" + assert created[0].metadata["package_key"] == "model" + + def test_two_packages_can_coexist(self): + def package(): + return ModelPackage({"model": _make_model({"x": [1, 4]}, {"y": [1, 4]})}) + + builder = PipelineBuilder() + builder.add_package(package(), roles={"model": "encoder"}, prefix="left") + builder.add_package(package(), roles={"model": "decoder"}, prefix="right") + builder.connect("left_model.y", "right_model.x") + builder.declare_external("left_model.x") + builder.add_stage("run", "single_pass", ["left_model", "right_model"]) + builder.add_public_output("right_model.y") + pkg = builder.build() + assert pkg.manifest.component_names == ("left_model", "right_model") + + def test_missing_role_rejected(self): + package = ModelPackage({"model": _make_model({"x": [1, 4]}, {"y": [1, 4]})}) + builder = PipelineBuilder() + with pytest.raises(PipelineValidationError, match="No role declared"): + builder.add_package(package, roles={"other": "encoder"}) + + def test_role_callable_ignores_literal_key_semantics(self): + package = ModelPackage( + { + "vision": _make_model({"x": [1, 4]}, {"y": [1, 4]}), + "model": _make_model({"x": [1, 4]}, {"y": [1, 4]}), + } + ) + builder = PipelineBuilder() + created = builder.add_package(package, roles=lambda key: "generic") + assert {c.role for c in created} == {"generic"} + assert [c.name for c in created] == ["model", "vision"] + + def test_duplicate_component_rejected(self): + builder = PipelineBuilder() + builder.add_model("encoder", _encoder(), role="encoder") + with pytest.raises(PipelineValidationError, match="already registered"): + builder.add_model("encoder", _encoder(), role="encoder") + + +class TestConnectionValidation: + def test_fan_out_accepted(self): + builder = PipelineBuilder() + builder.add_model("src", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="decoder") + builder.add_model("b", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="decoder") + builder.connect("src.y", "a.x") + builder.connect("src.y", "b.x") + builder.declare_external("src.x") + builder.add_stage("run", "single_pass", ["src", "a", "b"]) + builder.add_public_output("a.y", alias="a_out") + builder.add_public_output("b.y", alias="b_out") + manifest = builder.build().manifest + targets = [c.target.qualified for c in manifest.connections] + assert sorted(targets) == ["a.x", "b.x"] + + def test_duplicate_producer_rejected(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.add_model("b", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.add_model("c", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="decoder") + builder.connect("a.y", "c.x") + builder.connect("b.y", "c.x") + builder.declare_external("a.x") + builder.declare_external("b.x", alias="b_x") + builder.add_stage("run", "single_pass", ["a", "b", "c"]) + with pytest.raises(PipelineValidationError, match=r"more than one .*producer"): + builder.build() + + def test_unknown_component_endpoint_rejected(self): + builder = _simple_pipeline() + builder.connect("ghost.y", "decoder.position_ids") + with pytest.raises(PipelineValidationError, match="unknown component 'ghost'"): + builder.build() + + def test_unknown_port_rejected(self): + builder = _simple_pipeline() + builder.connect("encoder.nope", "decoder.position_ids") + with pytest.raises(PipelineValidationError, match="unknown output 'nope'"): + builder.build() + + def test_symbolic_dimension_name_mismatch_accepted(self): + builder = PipelineBuilder() + builder.add_model( + "a", _make_model({"x": [1, 4]}, {"y": ["batch", "seq", 8]}), role="encoder" + ) + builder.add_model( + "b", _make_model({"h": ["n", "tokens", 8]}, {"y": [1, 4]}), role="decoder" + ) + builder.connect("a.y", "b.h") + builder.declare_external("a.x") + builder.add_stage("run", "single_pass", ["a", "b"]) + builder.add_public_output("b.y") + assert builder.build().manifest.source_of("b.h") == InputSource.DATAFLOW + + def test_concrete_dimension_mismatch_rejected(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": ["b", 8]}), role="encoder") + builder.add_model("b", _make_model({"h": ["b", 16]}, {"y": [1, 4]}), role="decoder") + builder.connect("a.y", "b.h") + builder.declare_external("a.x") + builder.add_stage("run", "single_pass", ["a", "b"]) + with pytest.raises(PipelineValidationError, match="dim 1: 8 != 16"): + builder.build() + + def test_rank_mismatch_rejected(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": ["b", 8]}), role="encoder") + builder.add_model("b", _make_model({"h": ["b", 1, 8]}, {"y": [1, 4]}), role="decoder") + builder.connect("a.y", "b.h") + builder.declare_external("a.x") + builder.add_stage("run", "single_pass", ["a", "b"]) + with pytest.raises(PipelineValidationError, match="rank 2 != 3"): + builder.build() + + def test_dtype_mismatch_rejected(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": ["b", 8]}), role="encoder") + builder.add_model( + "b", + _make_model({"h": ["b", 8]}, {"y": [1, 4]}, dtype=ir.DataType.INT64), + role="decoder", + ) + builder.connect("a.y", "b.h") + builder.declare_external("a.x") + builder.add_stage("run", "single_pass", ["a", "b"]) + with pytest.raises(PipelineValidationError, match="dtype FLOAT != INT64"): + builder.build() + + def test_registered_transform_allows_rank_and_dtype_change(self): + """A VAE-style edge: FLOAT [b, 8, h, w] moments -> INT64 [b, t] tokens.""" + builder = PipelineBuilder() + builder.add_model( + "vae", + _make_model({"x": [1, 4]}, {"moments": ["b", 8, "h", "w"]}), + role="encoder", + ) + builder.add_model( + "generator", + _make_model({"tokens": ["b", "t"]}, {"y": [1, 4]}, dtype=ir.DataType.INT64), + role="decoder", + ) + builder.connect("vae.moments", "generator.tokens", transform="patchify") + builder.declare_external("vae.x") + builder.add_stage("run", "single_pass", ["vae", "generator"]) + builder.add_public_output("generator.y") + manifest = builder.build().manifest + connection = manifest.connections[0] + assert connection.transform == "patchify" + assert connection.transform_capabilities == ("tensor_patchify",) + assert "tensor_patchify" in manifest.required_capabilities + + def test_unknown_transform_rejected(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": ["b", 8]}), role="encoder") + builder.add_model("b", _make_model({"h": ["b", 16]}, {"y": [1, 4]}), role="decoder") + with pytest.raises(PipelineValidationError, match="Unknown transform 'tile'"): + builder.connect("a.y", "b.h", transform="tile") + + def test_rank_mismatch_rejected_without_transform(self): + """The same edge that a transform makes legal is illegal untransformed.""" + builder = PipelineBuilder() + builder.add_model( + "vae", _make_model({"x": [1, 4]}, {"moments": ["b", 8, "h", "w"]}), role="encoder" + ) + builder.add_model( + "generator", _make_model({"tokens": ["b", "t"]}, {"y": [1, 4]}), role="decoder" + ) + builder.connect("vae.moments", "generator.tokens") + builder.declare_external("vae.x") + builder.add_stage("run", "single_pass", ["vae", "generator"]) + with pytest.raises(PipelineValidationError, match="rank 4 != 2"): + builder.build() + + def test_transform_still_validates_endpoints(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": ["b", 8]}), role="encoder") + builder.add_model("b", _make_model({"h": ["b", 16]}, {"y": [1, 4]}), role="decoder") + builder.connect("a.ghost", "b.h", transform="reshape") + builder.declare_external("a.x") + builder.add_stage("run", "single_pass", ["a", "b"]) + with pytest.raises(PipelineValidationError, match="unknown output 'ghost'"): + builder.build() + + def test_transform_still_enforces_single_producer(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.add_model("b", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.add_model("c", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="decoder") + builder.connect("a.y", "c.x", transform="reshape") + builder.connect("b.y", "c.x", transform="cast") + builder.declare_external("a.x") + builder.declare_external("b.x", alias="b_x") + builder.add_stage("run", "single_pass", ["a", "b", "c"]) + with pytest.raises(PipelineValidationError, match=r"more than one .*producer"): + builder.build() + + def test_transform_capabilities_must_be_declared_in_manifest(self): + manifest = _cosmos_style_pipeline().build().manifest + with pytest.raises(PipelineValidationError, match="requires capabilities"): + dataclasses.replace(manifest, required_capabilities=()) + + def test_recurrent_transform_edge(self): + """A scheduler step feeding the next latent back into the denoiser.""" + model = _make_model( + {"latent": ["b", 16], "timestep": ["b"]}, + {"noise_pred": ["b", 16]}, + ) + builder = PipelineBuilder() + builder.add_model("denoiser", model, role="dynamics", run_on="step") + builder.connect( + "denoiser.noise_pred", + "denoiser.latent", + recurrent=True, + transform="scheduler_step", + ) + builder.declare_generated("denoiser.latent", generator="zeros") + builder.declare_generated( + "denoiser.timestep", + generator="scheduler_timesteps", + parameters={"stage": "denoise"}, + ) + builder.add_stage("denoise", "iterative", ["denoiser"], run_on="step") + builder.add_public_output("denoiser.noise_pred") + manifest = builder.build().manifest + assert manifest.source_of("denoiser.latent") == InputSource.STATEFUL + assert set(manifest.required_capabilities) == { + LOOP_CARRIED_STATE_CAPABILITY, + "iterative_scheduler", + } + + def test_untyped_graph_rejected(self): + model = _make_model({"x": [1, 4]}, {"y": [1, 4]}) + model.graph.inputs[0].shape = None + builder = PipelineBuilder() + with pytest.raises(PipelineValidationError, match="has no shape"): + builder.add_model("a", model, role="encoder") + + +class TestInputSources: + def _pipeline_without_position_ids(self) -> PipelineBuilder: + """Encoder -> decoder where ``decoder.position_ids`` has no source yet.""" + builder = PipelineBuilder() + builder.add_model("encoder", _encoder(), role="encoder") + builder.add_model("decoder", _decoder(), role="decoder") + builder.connect("encoder.image_features", "decoder.image_features") + builder.declare_external("encoder.pixel_values") + builder.add_stage("encode", "single_pass", ["encoder"]) + builder.add_stage("generate", "autoregressive", ["decoder"]) + builder.add_public_output("decoder.logits") + return builder + + def test_missing_input_source_rejected(self): + with pytest.raises(PipelineValidationError, match="has no initial source"): + self._pipeline_without_position_ids().build() + + def test_connected_and_declared_is_conflicting(self): + builder = _simple_pipeline() + builder.declare_external("decoder.image_features") + with pytest.raises(PipelineValidationError, match="exactly one initial source"): + builder.build() + + def test_recurrent_input_requires_initial_source(self): + model = _make_model({"state": [1, 4]}, {"next_state": [1, 4]}) + builder = PipelineBuilder() + builder.add_model("dynamics", model, role="dynamics") + builder.connect("dynamics.next_state", "dynamics.state", recurrent=True) + builder.add_stage("imagine", "state_transition", ["dynamics"]) + builder.add_public_output("dynamics.next_state") + + with pytest.raises(PipelineValidationError, match=r"Recurrent input.*no initial"): + builder.build() + + def test_dataflow_initializer_and_recurrent_update_can_share_input(self): + builder = PipelineBuilder() + builder.add_model( + "encoder", + _make_model({"observation": [1, 4]}, {"state": [1, 4]}), + role="encoder", + ) + builder.add_model( + "dynamics", + _make_model({"state": [1, 4]}, {"next_state": [1, 4]}), + role="dynamics", + ) + builder.connect("encoder.state", "dynamics.state") + builder.connect("dynamics.next_state", "dynamics.state", recurrent=True) + builder.declare_external("encoder.observation") + builder.add_stage("encode", "single_pass", ["encoder"]) + builder.add_stage("imagine", "state_transition", ["dynamics"]) + builder.add_public_output("dynamics.next_state") + + manifest = builder.build().manifest + + assert manifest.source_of("dynamics.state") == InputSource.STATEFUL + assert manifest.initial_source_of("dynamics.state") == InputSource.DATAFLOW + + def test_stateful_classification_does_not_depend_on_source_name_sorting(self): + builder = PipelineBuilder() + builder.add_model( + "aencoder", + _make_model({"observation": [1, 4]}, {"state": [1, 4]}), + role="encoder", + ) + builder.add_model( + "dynamics", + _make_model({"state": [1, 4]}, {"next_state": [1, 4]}), + role="dynamics", + ) + builder.connect("aencoder.state", "dynamics.state") + builder.connect("dynamics.next_state", "dynamics.state", recurrent=True) + builder.declare_external("aencoder.observation") + builder.add_stage("encode", "single_pass", ["aencoder"]) + builder.add_stage("imagine", "state_transition", ["dynamics"]) + builder.add_public_output("dynamics.next_state") + + manifest = builder.build().manifest + + assert manifest.source_of("dynamics.state") == InputSource.STATEFUL + assert manifest.initial_source_of("dynamics.state") == InputSource.DATAFLOW + + def test_double_declaration_rejected(self): + builder = _simple_pipeline() + builder.declare_generated("encoder.pixel_values", generator="zeros") + with pytest.raises(PipelineValidationError, match="declared more than once"): + builder.build() + + def test_defaulted_input_accepted_and_serialized(self): + builder = self._pipeline_without_position_ids() + builder.declare_default("decoder.position_ids", [[0, 1, 2]]) + manifest = builder.build().manifest + assert manifest.source_of("decoder.position_ids") == InputSource.DEFAULTED + entry = manifest.inputs_of_kind(InputSource.DEFAULTED)[0] + assert entry.value == [[0, 1, 2]] + assert PipelineManifest.from_json(manifest.to_json()) == manifest + + def test_defaulted_input_requires_json_safe_value(self): + with pytest.raises(PipelineValidationError, match="JSON-safe"): + PipelineInput(PipelinePort("a", "x"), InputSource.DEFAULTED, value=object()) + + def test_defaulted_input_requires_a_value(self): + with pytest.raises(PipelineValidationError, match="requires a value"): + PipelineInput(PipelinePort("a", "x"), InputSource.DEFAULTED) + + def test_non_default_kind_rejects_value(self): + with pytest.raises(PipelineValidationError, match="must not carry a default"): + PipelineInput(PipelinePort("a", "x"), InputSource.EXTERNAL, value=3) + + def test_unknown_source_kind_rejected(self): + with pytest.raises(PipelineValidationError, match="unknown source kind"): + PipelineInput(PipelinePort("a", "x"), "magic") + + def test_duplicate_external_names_rejected(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.add_model("b", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.declare_external("a.x") + builder.declare_external("b.x") + builder.add_stage("run", "single_pass", ["a", "b"]) + with pytest.raises(PipelineValidationError, match="External input name"): + builder.build() + + def test_alias_only_for_external(self): + with pytest.raises(PipelineValidationError, match="Only external inputs"): + PipelineInput( + PipelinePort("a", "x"), + InputSource.GENERATED, + alias="z", + generator=GeneratedInputRule("zeros"), + ) + + +class TestCycles: + def _loop_builder(self, *, stage_kind: str) -> PipelineBuilder: + model = _make_model( + {"tokens": ["b", "s"], "state": ["b", 8]}, + {"logits": ["b", 32], "new_state": ["b", 8]}, + ) + builder = PipelineBuilder() + builder.add_model("decoder", model, role="dynamics", run_on="decode") + builder.connect("decoder.new_state", "decoder.state", recurrent=True) + builder.declare_stateful("decoder.state") + builder.declare_external("decoder.tokens") + builder.add_stage("loop", stage_kind, ["decoder"], run_on="decode") + builder.add_public_output("decoder.logits") + return builder + + def test_legal_recurrent_cycle(self): + manifest = self._loop_builder(stage_kind="autoregressive").build().manifest + assert manifest.source_of("decoder.state") == InputSource.STATEFUL + stage = manifest.stages[0] + assert LOOP_CARRIED_STATE_CAPABILITY in stage.capabilities + assert LOOP_CARRIED_STATE_CAPABILITY in manifest.required_capabilities + + def test_state_transition_stage_supports_loops(self): + manifest = self._loop_builder(stage_kind="state_transition").build().manifest + assert manifest.connections[0].recurrent + + def test_recurrent_edge_in_single_pass_stage_rejected(self): + builder = self._loop_builder(stage_kind="single_pass") + with pytest.raises(PipelineValidationError, match="loop-carried state"): + builder.build() + + def test_illegal_non_recurrent_cycle(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.add_model("b", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="decoder") + builder.connect("a.y", "b.x") + builder.connect("b.y", "a.x") + builder.add_stage("run", "iterative", ["a", "b"]) + builder.add_public_output("b.y") + with pytest.raises(PipelineValidationError, match="cycle in non-recurrent"): + builder.build() + + def test_non_recurrent_self_edge_rejected(self): + model = _make_model({"x": [1, 4], "s": [1, 4]}, {"y": [1, 4]}) + builder = PipelineBuilder() + builder.add_model("a", model, role="encoder") + builder.connect("a.y", "a.s") + builder.declare_external("a.x") + builder.add_stage("run", "iterative", ["a"]) + with pytest.raises(PipelineValidationError, match="depend on itself"): + builder.build() + + def test_recurrent_edge_across_stages_rejected(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.add_model("b", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="decoder") + builder.connect("b.y", "a.x", recurrent=True) + builder.declare_stateful("a.x") + builder.declare_external("b.x") + builder.add_stage("first", "iterative", ["a"]) + builder.add_stage("second", "iterative", ["b"]) + builder.add_public_output("a.y") + with pytest.raises(PipelineValidationError, match="scoped to a single stage"): + builder.build() + + +class TestStages: + def test_unknown_component_in_stage_rejected(self): + builder = _simple_pipeline() + builder.add_stage("extra", "single_pass", ["ghost"]) + with pytest.raises(PipelineValidationError, match="unknown component 'ghost'"): + builder.build() + + def test_component_without_stage_rejected(self): + builder = PipelineBuilder() + builder.add_model("a", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.declare_external("a.x") + builder.add_public_output("a.y") + with pytest.raises(PipelineValidationError, match="belong to no declared stage"): + builder.build() + + def test_impossible_phase_combination_rejected(self): + builder = PipelineBuilder() + builder.add_model( + "a", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder", run_on="prefill" + ) + builder.declare_external("a.x") + builder.add_stage("late", "single_pass", ["a"], run_on="decode") + with pytest.raises(PipelineValidationError, match="could never execute"): + builder.build() + + def test_on_demand_stage_with_presence(self): + builder = PipelineBuilder() + builder.add_model( + "vision", + _make_model({"pixels": [1, 4]}, {"feats": [1, 4]}), + role="encoder", + presence="has_image", + run_on="on_demand", + ) + builder.declare_external("vision.pixels") + builder.add_stage( + "maybe_vision", + "on_demand", + ["vision"], + run_on="on_demand", + options={"presence": "has_image"}, + ) + builder.add_public_output("vision.feats") + manifest = builder.build().manifest + assert manifest.component("vision").presence == "has_image" + assert manifest.stages[0].options == {"presence": "has_image"} + + def test_on_demand_stage_without_presence_rejected(self): + builder = PipelineBuilder() + builder.add_model( + "vision", + _make_model({"pixels": [1, 4]}, {"feats": [1, 4]}), + role="encoder", + ) + builder.declare_external("vision.pixels") + builder.add_stage("maybe_vision", "on_demand", ["vision"]) + builder.add_public_output("vision.feats") + + with pytest.raises(PipelineValidationError, match="no component presence"): + builder.build() + + def test_composite_stage_kind(self): + builder = _simple_pipeline() + builder.add_stage("everything", "composite", ["encoder", "decoder"]) + manifest = builder.build().manifest + assert [s.kind for s in manifest.stages] == [ + "single_pass", + "autoregressive", + "composite", + ] + + def test_empty_stage_rejected(self): + with pytest.raises(PipelineValidationError, match="must contain a component"): + PipelineStage("s", "single_pass", ()) + + def test_stage_options_must_be_json_safe(self): + with pytest.raises(PipelineValidationError, match="JSON-safe"): + PipelineStage("s", "single_pass", ("a",), options={"f": object()}) + + def test_duplicate_stage_name_rejected(self): + builder = _simple_pipeline() + builder.add_stage("encode", "single_pass", ["decoder"]) + with pytest.raises(PipelineValidationError, match="Stage 'encode'"): + builder.build() + + +class TestRegistries: + def test_unknown_role_rejected(self): + builder = PipelineBuilder() + with pytest.raises(PipelineValidationError, match="Unknown role 'world_sim'"): + builder.add_model("a", _make_model({"x": [1]}, {"y": [1]}), role="world_sim") + + def test_unknown_strategy_rejected(self): + with pytest.raises(PipelineValidationError, match="Unknown strategy 'beam'"): + PipelineStage("s", "beam", ("a",)) + + def test_unknown_phase_rejected(self): + builder = PipelineBuilder() + with pytest.raises(PipelineValidationError, match="Unknown phase"): + builder.add_model( + "a", _make_model({"x": [1]}, {"y": [1]}), role="encoder", run_on="someday" + ) + + def test_registration_is_idempotent(self): + first = register_role("test_role_idempotent", description="d") + second = register_role("test_role_idempotent", description="d") + assert first == second + + def test_transform_registration_is_idempotent(self): + first = register_transform("test_transform_idempotent", capabilities=["cap"]) + second = register_transform("test_transform_idempotent", capabilities=["cap"]) + assert first == second + assert first.capabilities == ("cap",) + + def test_conflicting_transform_registration_rejected(self): + register_transform("test_transform_conflict", capabilities=["a"]) + with pytest.raises(PipelineValidationError, match="already registered"): + register_transform("test_transform_conflict", capabilities=["b"]) + + def test_registered_transform_usable_end_to_end(self): + register_transform( + "test_transform_usable", + description="quantize latents into codebook ids", + capabilities=["vector_quantization"], + ) + builder = PipelineBuilder() + builder.add_model( + "a", _make_model({"x": [1, 4]}, {"y": ["b", "t", 16]}), role="encoder" + ) + builder.add_model( + "b", + _make_model({"ids": ["b", "t"]}, {"y": [1, 4]}, dtype=ir.DataType.INT64), + role="decoder", + ) + builder.connect("a.y", "b.ids", transform="test_transform_usable") + builder.declare_external("a.x") + builder.add_stage("run", "single_pass", ["a", "b"]) + builder.add_public_output("b.y") + manifest = builder.build().manifest + assert manifest.required_capabilities == ("vector_quantization",) + assert PipelineManifest.from_json(manifest.to_json()) == manifest + + def test_builtin_transform_capabilities(self): + assert transform_definition("sample").capabilities == ("stochastic_sampling",) + assert transform_definition("scheduler_step").capabilities == ("iterative_scheduler",) + + def test_unknown_transform_in_manifest_dict_rejected(self): + data = _cosmos_style_pipeline().build().manifest.to_dict() + data["connections"][0]["transform"] = "not_a_transform" + with pytest.raises(PipelineValidationError, match="Unknown transform"): + PipelineManifest.from_dict(data) + + def test_conflicting_registration_rejected(self): + register_strategy("test_strategy_conflict", description="a") + with pytest.raises(PipelineValidationError, match="already registered"): + register_strategy("test_strategy_conflict", description="b") + + def test_registered_role_usable(self): + register_role("test_role_usable", description="a world model rollout head") + builder = PipelineBuilder() + component = builder.add_model( + "a", _make_model({"x": [1]}, {"y": [1]}), role="test_role_usable" + ) + assert component.role == "test_role_usable" + + def test_unknown_role_in_manifest_dict_rejected(self): + data = _simple_pipeline().build().manifest.to_dict() + data["components"][0]["role"] = "not_a_role" + with pytest.raises(PipelineValidationError, match="Unknown role"): + PipelineManifest.from_dict(data) + + def test_unknown_strategy_in_manifest_dict_rejected(self): + data = _simple_pipeline().build().manifest.to_dict() + data["stages"][0]["kind"] = "not_a_strategy" + with pytest.raises(PipelineValidationError, match="Unknown strategy"): + PipelineManifest.from_dict(data) + + def test_generated_input_registration_is_closed_and_parameterized(self): + register_generated_input( + "test_generated_program", + required_parameters=["source"], + allowed_parameters=["source"], + ) + with pytest.raises(PipelineValidationError, match="missing required"): + GeneratedInputRule("test_generated_program") + with pytest.raises(PipelineValidationError, match="unknown parameter"): + GeneratedInputRule( + "test_generated_program", + {"source": "x", "extra": True}, + ) + assert GeneratedInputRule( + "test_generated_program", + {"source": "x"}, + ).parameters == {"source": "x"} + + def test_transform_parameters_are_validated(self): + with pytest.raises(PipelineValidationError, match="unknown parameter"): + PipelineConnection( + PipelinePort("a", "y"), + PipelinePort("b", "x"), + transform="cast", + parameters={"guess": "fp16"}, + ) + + +class TestExecutableProfile: + def _builder(self, tmp_path, *, include_state: bool = True) -> PipelineBuilder: + builder = PipelineBuilder() + builder.add_model( + "decoder", + _make_model( + {"tokens": [1, 1], "state": [1, 4]}, + {"logits": [1, 8], "next_state": [1, 4]}, + ), + role="decoder", + preferred_execution_providers=["cuda", "cpu"], + parameter_dtype="FLOAT", + ) + builder.connect("decoder.next_state", "decoder.state", recurrent=True) + builder.declare_external( + "decoder.tokens", + semantic="text.token_ids", + ) + builder.declare_generated( + "decoder.state", + generator="zeros", + semantic="kv_cache.initial", + ) + builder.add_stage( + "decode", + "autoregressive", + ["decoder"], + options={ + "tokenizer_asset": "tokenizer.json", + "sampling": {"do_sample": False}, + "stop": {"kind": "token_ids", "eos_token_ids": [2]}, + "state_names": ["cache"], + }, + ) + if include_state: + builder.add_state( + "cache", + kind="kv_cache", + input="decoder.state", + output="decoder.next_state", + lifetime="sequence", + release_after="decode", + sequence_axis=1, + ) + builder.add_public_state_output("cache", alias="final_cache") + builder.add_public_output("decoder.logits") + tokenizer = tmp_path / "tokenizer.json" + tokenizer.write_text("{}", encoding="utf-8") + scheduler_dir = tmp_path / "scheduler" + scheduler_dir.mkdir(exist_ok=True) + scheduler = scheduler_dir / "scheduler_config.json" + scheduler.write_text("{}", encoding="utf-8") + builder.add_asset("tokenizer.json", str(tokenizer)) + builder.add_asset("scheduler/scheduler_config.json", str(scheduler)) + builder.set_profile("test-world", "1.0") + return builder + + def test_profile_round_trip_contains_executable_contract(self, tmp_path): + manifest = self._builder(tmp_path).build().manifest + restored = PipelineManifest.from_json(manifest.to_json()) + + assert restored.profile is not None + assert restored.profile.name == "test-world" + assert restored.states[0].kind == "kv_cache" + assert next( + output for output in restored.outputs if output.name == "final_cache" + ).state == ("cache") + assert restored.inputs_of_kind(InputSource.GENERATED)[0].generator == ( + GeneratedInputRule("zeros") + ) + assert restored.component("decoder").preferred_execution_providers == ( + "cuda", + "cpu", + ) + + def test_profile_requires_state_lifecycle(self, tmp_path): + with pytest.raises(PipelineValidationError, match="explicit state lifecycle"): + self._builder(tmp_path, include_state=False).build() + + def test_profile_requires_input_semantics(self, tmp_path): + builder = self._builder(tmp_path) + entry = next(i for i in builder._inputs if i.kind == InputSource.EXTERNAL) + builder._inputs[builder._inputs.index(entry)] = dataclasses.replace( + entry, + semantic=None, + ) + with pytest.raises(PipelineValidationError, match="semantic names"): + builder.build() + + def test_profile_version_is_independent_from_schema_major(self): + assert PipelineProfile("future-runtime", "2.0").version == "2.0" + + def test_generated_program_port_references_are_validated(self): + builder = PipelineBuilder() + builder.add_model( + "decoder", + _make_model({"positions": [1, 1]}, {"logits": [1, 8]}), + role="decoder", + ) + builder.declare_generated( + "decoder.positions", + generator="multimodal_position_ids", + parameters={"source": "ghost.tokens", "axes": 1}, + ) + builder.add_stage("run", "single_pass", ["decoder"]) + builder.add_public_output("decoder.logits") + + with pytest.raises(PipelineValidationError, match="unknown port"): + builder.build() + + def test_generated_empty_tensor_axis_must_exist_on_port(self): + builder = PipelineBuilder() + builder.add_model( + "decoder", + _make_model({"cache": [1, "past"]}, {"logits": [1, 8]}), + role="decoder", + ) + builder.declare_generated( + "decoder.cache", + generator="empty_tensor", + parameters={"dynamic_axes": {"past_sequence_length": 0}}, + ) + builder.add_stage("run", "single_pass", ["decoder"]) + builder.add_public_output("decoder.logits") + + with pytest.raises(PipelineValidationError, match="unknown dynamic axis"): + builder.build() + + +class TestManifestSchema: + def test_unknown_major_version_rejected(self): + data = _simple_pipeline().build().manifest.to_dict() + data["schema_version"] = "99.0" + with pytest.raises(PipelineValidationError, match="Unsupported pipeline schema"): + PipelineManifest.from_dict(data) + + def test_newer_minor_version_accepted(self): + data = _simple_pipeline().build().manifest.to_dict() + data["schema_version"] = "1.99" + assert PipelineManifest.from_dict(data).schema_version == "1.99" + + def test_unknown_top_level_key_rejected(self): + data = _simple_pipeline().build().manifest.to_dict() + data["scheduler"] = {"algorithm": "continuous_batching"} + with pytest.raises(PipelineValidationError, match="Unknown key"): + PipelineManifest.from_dict(data) + + def test_metadata_bag_is_preserved(self): + builder = _simple_pipeline() + builder.set_metadata("provenance", {"tool": "mobius", "extra": [1, 2]}) + manifest = builder.build().manifest + restored = PipelineManifest.from_json(manifest.to_json()) + assert restored.metadata == {"provenance": {"tool": "mobius", "extra": [1, 2]}} + assert restored == manifest + + def test_component_metadata_preserved(self): + data = _simple_pipeline().build().manifest.to_dict() + data["components"][0]["metadata"] = {"future_field": True} + restored = PipelineManifest.from_dict(data) + assert restored.component("decoder").metadata == {"future_field": True} + + def test_public_output_must_exist(self): + builder = _simple_pipeline() + builder.add_public_output("decoder.ghost", alias="ghost") + with pytest.raises(PipelineValidationError, match="unknown output 'ghost'"): + builder.build() + + def test_required_capability_must_be_provided(self): + builder = _simple_pipeline() + builder.require_capability("streaming") + with pytest.raises(PipelineValidationError, match="not provided by any"): + builder.build() + + def test_required_capability_provided_by_component(self): + builder = PipelineBuilder() + builder.add_model( + "a", + _make_model({"x": [1, 4]}, {"y": [1, 4]}), + role="encoder", + capabilities=["streaming"], + ) + builder.declare_external("a.x") + builder.add_stage("run", "single_pass", ["a"]) + builder.add_public_output("a.y") + builder.require_capability("streaming") + assert builder.build().manifest.required_capabilities == ("streaming",) + + def test_manifest_equality_ignores_graph_identity(self): + left = _simple_pipeline().build().manifest + right = _simple_pipeline().build().manifest + assert left == right + assert left.component("encoder").model is not right.component("encoder").model + + +class TestPackagePersistence: + def test_save_load_round_trip(self, tmp_path): + pkg = _simple_pipeline().build() + pkg.save(str(tmp_path), progress_bar=False) + assert (tmp_path / PIPELINE_FILENAME).is_file() + assert (tmp_path / "encoder" / "model.onnx").is_file() + assert (tmp_path / "decoder" / "model.onnx").is_file() + + loaded = PipelinePackage.load(str(tmp_path)) + assert loaded.manifest == pkg.manifest + assert sorted(loaded) == ["decoder", "encoder"] + assert loaded.manifest.component("encoder").model is loaded["encoder"] + + def test_component_filenames_round_trip_exactly(self, tmp_path): + pkg = _simple_pipeline().build() + pkg.save(str(tmp_path), progress_bar=False) + document = json.loads((tmp_path / PIPELINE_FILENAME).read_text(encoding="utf-8")) + assert document["component_files"] == { + "decoder": "decoder/model.onnx", + "encoder": "encoder/model.onnx", + } + for relative in document["component_files"].values(): + assert (tmp_path / relative).is_file() + + def test_single_component_uses_flat_layout(self, tmp_path): + builder = PipelineBuilder() + builder.add_model("solo", _make_model({"x": [1, 4]}, {"y": [1, 4]}), role="encoder") + builder.declare_external("solo.x") + builder.add_stage("run", "single_pass", ["solo"]) + builder.add_public_output("solo.y") + pkg = builder.build() + pkg.save(str(tmp_path), progress_bar=False) + assert (tmp_path / "model.onnx").is_file() + assert pkg.component_files() == {"solo": "model.onnx"} + assert PipelinePackage.load(str(tmp_path)).manifest == pkg.manifest + + def test_partial_save_rejected(self, tmp_path): + pkg = _simple_pipeline().build() + with pytest.raises(PipelineValidationError, match="partial saves"): + pkg.save(str(tmp_path), components=lambda name: name == "encoder") + + def test_load_requires_pipeline_json(self, tmp_path): + ModelPackage({"a": _make_model({"x": [1, 4]}, {"y": [1, 4]})}).save( + str(tmp_path), progress_bar=False + ) + with pytest.raises(PipelineValidationError, match="not found"): + PipelinePackage.load(str(tmp_path)) + + def test_load_rejects_missing_component_file(self, tmp_path): + pkg = _simple_pipeline().build() + pkg.save(str(tmp_path), progress_bar=False) + (tmp_path / "encoder" / "model.onnx").unlink() + with pytest.raises(PipelineValidationError, match="is missing"): + PipelinePackage.load(str(tmp_path)) + + def test_load_rejects_component_path_escape(self, tmp_path): + pkg = _simple_pipeline().build() + pkg.save(str(tmp_path), progress_bar=False) + path = tmp_path / PIPELINE_FILENAME + document = json.loads(path.read_text(encoding="utf-8")) + document["component_files"]["encoder"] = "../outside/model.onnx" + path.write_text(json.dumps(document), encoding="utf-8") + + with pytest.raises(PipelineValidationError, match="safe layout"): + PipelinePackage.load(str(tmp_path)) + + def test_load_wraps_corrupt_manifest_error(self, tmp_path): + (tmp_path / PIPELINE_FILENAME).write_text('{"format":', encoding="utf-8") + + with pytest.raises(PipelineValidationError, match=r"valid 'pipeline\.json'"): + PipelinePackage.load(str(tmp_path)) + + def test_load_rejects_wrong_manifest_format(self, tmp_path): + pkg = _simple_pipeline().build() + pkg.save(str(tmp_path), progress_bar=False) + path = tmp_path / PIPELINE_FILENAME + document = json.loads(path.read_text(encoding="utf-8")) + document["format"] = "other" + path.write_text(json.dumps(document), encoding="utf-8") + + with pytest.raises(PipelineValidationError, match="unsupported format"): + PipelinePackage.load(str(tmp_path)) + + def test_failed_resave_leaves_no_completeness_marker(self, tmp_path): + pkg = _simple_pipeline().build() + pkg.save(str(tmp_path), progress_bar=False) + + with ( + mock.patch("mobius._pipeline.json.dump", side_effect=OSError("disk full")), + pytest.raises(OSError, match="disk full"), + ): + pkg.save(str(tmp_path), progress_bar=False) + + assert not (tmp_path / PIPELINE_FILENAME).exists() + assert not list(tmp_path.glob(".mobius-pipeline-*")) + + def test_desynchronized_models_rejected(self): + pkg = _simple_pipeline().build() + pkg["extra"] = _make_model({"x": [1, 4]}, {"y": [1, 4]}) + with pytest.raises(PipelineValidationError, match="do not match the manifest"): + pkg.save("unused-directory", progress_bar=False) + + def test_package_requires_matching_models(self): + manifest = _simple_pipeline().build().manifest + with pytest.raises(PipelineValidationError, match="do not match the manifest"): + PipelinePackage({"encoder": _encoder()}, manifest) + + +class TestAssets: + def _asset(self, tmp_path, name: str, text: str) -> str: + path = tmp_path / name + path.write_text(text, encoding="utf-8") + return str(path) + + def _pipeline_with_assets(self, tmp_path) -> PipelineBuilder: + builder = _simple_pipeline() + builder.add_asset("tokenizer.json", self._asset(tmp_path, "tokenizer.json", '{"v":1}')) + builder.add_asset( + "scheduler/scheduler_config.json", + self._asset(tmp_path, "scheduler_config.json", '{"steps":30}'), + ) + return builder + + def test_assets_saved_and_recorded_by_destination_only(self, tmp_path): + source_dir = tmp_path / "src" + source_dir.mkdir() + out = tmp_path / "out" + pkg = self._pipeline_with_assets(source_dir).build() + pkg.save(str(out), progress_bar=False) + + assert (out / "tokenizer.json").read_text(encoding="utf-8") == '{"v":1}' + assert (out / "scheduler" / "scheduler_config.json").is_file() + + document = json.loads((out / PIPELINE_FILENAME).read_text(encoding="utf-8")) + assets = document["manifest"]["assets"] + assert assets == [ + {"path": "scheduler/scheduler_config.json"}, + {"path": "tokenizer.json"}, + ] + # No machine-local source path may leak into the manifest. + assert str(source_dir) not in json.dumps(document) + + def test_asset_round_trip_exposes_resolved_paths(self, tmp_path): + source_dir = tmp_path / "src" + source_dir.mkdir() + out = tmp_path / "out" + pkg = self._pipeline_with_assets(source_dir).build() + pkg.save(str(out), progress_bar=False) + + loaded = PipelinePackage.load(str(out)) + assert loaded.manifest == pkg.manifest + assert set(loaded.assets) == {"tokenizer.json", "scheduler/scheduler_config.json"} + resolved = loaded.asset_path("tokenizer.json") + assert os.path.isfile(resolved) + with open(resolved, encoding="utf-8") as file: + assert file.read() == '{"v":1}' + + def test_resaving_a_loaded_package_is_idempotent(self, tmp_path): + source_dir = tmp_path / "src" + source_dir.mkdir() + first = tmp_path / "first" + second = tmp_path / "second" + self._pipeline_with_assets(source_dir).build().save(str(first), progress_bar=False) + loaded = PipelinePackage.load(str(first)) + loaded.save(str(second), progress_bar=False) + assert (second / "tokenizer.json").read_text(encoding="utf-8") == '{"v":1}' + assert PipelinePackage.load(str(second)).manifest == loaded.manifest + + def test_save_in_place_preserves_assets(self, tmp_path): + source_dir = tmp_path / "src" + source_dir.mkdir() + out = tmp_path / "out" + self._pipeline_with_assets(source_dir).build().save(str(out), progress_bar=False) + loaded = PipelinePackage.load(str(out)) + loaded.save(str(out), progress_bar=False) + assert (out / "tokenizer.json").read_text(encoding="utf-8") == '{"v":1}' + + def test_missing_required_asset_fails_load(self, tmp_path): + source_dir = tmp_path / "src" + source_dir.mkdir() + out = tmp_path / "out" + self._pipeline_with_assets(source_dir).build().save(str(out), progress_bar=False) + (out / "tokenizer.json").unlink() + with pytest.raises(PipelineValidationError, match=r"Required asset 'tokenizer\.json'"): + PipelinePackage.load(str(out)) + + def test_optional_asset_may_be_absent(self, tmp_path): + source_dir = tmp_path / "src" + source_dir.mkdir() + out = tmp_path / "out" + builder = _simple_pipeline() + builder.add_asset( + "processor_config.json", + self._asset(source_dir, "processor_config.json", "{}"), + required=False, + ) + builder.build().save(str(out), progress_bar=False) + (out / "processor_config.json").unlink() + loaded = PipelinePackage.load(str(out)) + assert loaded.assets == {} + assert loaded.manifest.required_assets == () + + def test_required_asset_without_source_rejected(self): + manifest = _simple_pipeline().build().manifest + manifest = dataclasses.replace(manifest, assets=(PipelineAsset("tokenizer.json"),)) + with pytest.raises(PipelineValidationError, match="have no source file"): + PipelinePackage( + {name: manifest.component(name).model for name in manifest.component_names}, + manifest, + ) + + def test_undeclared_asset_source_rejected(self, tmp_path): + pkg = _simple_pipeline().build() + with pytest.raises(PipelineValidationError, match="undeclared destination"): + PipelinePackage( + dict(pkg), + pkg.manifest, + assets={"tokenizer.json": self._asset(tmp_path, "t.json", "{}")}, + ) + + def test_missing_source_file_rejected(self, tmp_path): + builder = _simple_pipeline() + with pytest.raises(PipelineValidationError, match="must be an existing file"): + builder.add_asset("tokenizer.json", str(tmp_path / "nope.json")) + + def test_duplicate_asset_destination_rejected(self, tmp_path): + builder = _simple_pipeline() + source = self._asset(tmp_path, "t.json", "{}") + builder.add_asset("tokenizer.json", source) + with pytest.raises(PipelineValidationError, match="already registered"): + builder.add_asset("tokenizer.json", source) + + @pytest.mark.parametrize( + "destination", + [ + "../escape.json", + "a/../../escape.json", + "/etc/passwd", + "C:/Windows/system.ini", + "c:tokenizer.json", + "~/tokenizer.json", + "sub\\tokenizer.json", + "a//b.json", + "./tokenizer.json", + "sub/", + "", + " ", + " tokenizer.json", + "nul.json", + "sub/con", + "bad\x00name.json", + ], + ) + def test_unsafe_asset_destinations_rejected(self, tmp_path, destination): + builder = _simple_pipeline() + source = self._asset(tmp_path, "t.json", "{}") + with pytest.raises(PipelineValidationError): + builder.add_asset(destination, source) + + @pytest.mark.parametrize( + "destination", + ["tokenizer.json", "sub/tokenizer.json", "a/b/c.txt", "chat_template.jinja"], + ) + def test_safe_asset_destinations_accepted(self, tmp_path, destination): + builder = _simple_pipeline() + builder.add_asset(destination, self._asset(tmp_path, "t.json", "{}")) + assert builder.build().manifest.assets[0].path == destination + + def test_asset_cannot_shadow_package_files(self, tmp_path): + builder = _simple_pipeline() + source = self._asset(tmp_path, "t.json", "{}") + builder.add_asset(PIPELINE_FILENAME, source) + with pytest.raises(PipelineValidationError, match="collides with a file written"): + builder.build() + + def test_asset_cannot_shadow_component_graph(self, tmp_path): + builder = _simple_pipeline() + builder.add_asset("encoder/model.onnx", self._asset(tmp_path, "t.json", "{}")) + with pytest.raises(PipelineValidationError, match="collides with a file written"): + builder.build() + + def test_asset_cannot_shadow_component_graph_with_different_case(self, tmp_path): + builder = _simple_pipeline() + builder.add_asset("Encoder/model.onnx", self._asset(tmp_path, "t.json", "{}")) + with pytest.raises(PipelineValidationError, match="collides with a file written"): + builder.build() + + def test_asset_manifest_round_trips_through_json(self, tmp_path): + builder = _simple_pipeline() + builder.add_asset("tokenizer.json", self._asset(tmp_path, "t.json", "{}")) + builder.add_asset("extra.txt", self._asset(tmp_path, "e.txt", "x"), required=False) + manifest = builder.build().manifest + restored = PipelineManifest.from_json(manifest.to_json()) + assert restored == manifest + assert [(a.path, a.required) for a in restored.assets] == [ + ("extra.txt", False), + ("tokenizer.json", True), + ] + + def test_no_temporary_files_left_behind(self, tmp_path): + source_dir = tmp_path / "src" + source_dir.mkdir() + out = tmp_path / "out" + self._pipeline_with_assets(source_dir).build().save(str(out), progress_bar=False) + leftovers = [p.name for p in out.rglob("*") if p.name.startswith(".mobius-asset-")] + assert leftovers == [] + + def test_deleted_source_between_build_and_save_rejected(self, tmp_path): + source_dir = tmp_path / "src" + source_dir.mkdir() + builder = _simple_pipeline() + source = self._asset(source_dir, "tokenizer.json", "{}") + builder.add_asset("tokenizer.json", source) + pkg = builder.build() + os.remove(source) + out = tmp_path / "out" + with pytest.raises(PipelineValidationError, match="does not exist"): + pkg.save(str(out), progress_bar=False) + assert not (out / PIPELINE_FILENAME).exists() + + def test_asset_path_unknown_destination(self): + pkg = _simple_pipeline().build() + with pytest.raises(KeyError): + pkg.asset_path("tokenizer.json") + + +class TestPackageConfigs: + def test_primary_config_preserved(self): + config = object() + pkg = _simple_pipeline().build(config=config) + assert pkg.config is config + assert pkg.config_for("encoder") is config + + def test_per_component_config(self): + primary = object() + encoder_config = object() + pkg = _simple_pipeline().build( + config=primary, component_configs={"encoder": encoder_config} + ) + assert pkg.config_for("encoder") is encoder_config + assert pkg.config_for("decoder") is primary + assert pkg.config is primary + + def test_config_for_unknown_component_rejected(self): + with pytest.raises(PipelineValidationError, match="unknown component"): + _simple_pipeline().build(component_configs={"ghost": object()}) + + def test_component_config_must_be_json_safe(self): + builder = PipelineBuilder() + with pytest.raises(PipelineValidationError, match="JSON-safe"): + builder.add_model( + "a", + _make_model({"x": [1, 4]}, {"y": [1, 4]}), + role="encoder", + config={"bad": object()}, + ) + + def test_component_config_round_trips(self): + builder = PipelineBuilder() + builder.add_model( + "a", + _make_model({"x": [1, 4]}, {"y": [1, 4]}), + role="encoder", + config={"latent_dim": 8}, + source="acme/world-model", + ) + builder.declare_external("a.x") + builder.add_stage("run", "single_pass", ["a"]) + builder.add_public_output("a.y") + manifest = builder.build().manifest + restored = PipelineManifest.from_json(manifest.to_json()) + assert restored.component("a").config == {"latent_dim": 8} + assert restored.component("a").source == "acme/world-model" + + +class TestDataclassSurface: + def test_component_lookup_helpers(self): + component = PipelineComponent.from_model("a", _encoder(), role="encoder") + assert component.input("pixel_values") is not None + assert component.input("missing") is None + assert component.output("image_features") is not None + + def test_connection_serialization(self): + connection = PipelineConnection( + PipelinePort("a", "y"), PipelinePort("b", "x"), recurrent=True, transform="cast" + ) + assert connection.to_dict() == { + "source": "a.y", + "target": "b.x", + "recurrent": True, + "transform": "cast", + } + assert PipelineConnection.from_dict(connection.to_dict()) == connection + + def test_transform_context_round_trips(self): + connection = PipelineConnection( + PipelinePort("denoiser", "velocity"), + PipelinePort("decoder", "latent"), + transform="scheduler_step", + context=(PipelinePort("denoiser", "sample"),), + ) + + assert connection.to_dict()["context"] == ["denoiser.sample"] + assert PipelineConnection.from_dict(connection.to_dict()) == connection + + def test_output_alias_defaults_to_port_name(self): + assert PipelineOutput(PipelinePort("a", "logits")).name == "logits" + assert PipelineOutput(PipelinePort("a", "logits"), "y").name == "y" + + def test_empty_manifest_is_valid(self): + manifest = PipelineManifest() + assert manifest.components == () + assert PipelineManifest.from_dict(manifest.to_dict()) == manifest + + def test_component_names_must_be_portable_across_case_sensitivity(self): + first = PipelineComponent.from_model("Encoder", _encoder(), role="encoder") + second = PipelineComponent.from_model("encoder", _encoder(), role="encoder") + + with pytest.raises(PipelineValidationError, match="case-insensitive"): + PipelineManifest(components=(first, second)) diff --git a/src/mobius/_weight_loading.py b/src/mobius/_weight_loading.py index 33f1653c7..c941b3b0f 100644 --- a/src/mobius/_weight_loading.py +++ b/src/mobius/_weight_loading.py @@ -17,12 +17,14 @@ __all__ = [ "apply_weights", + "iter_weight_shards", ] import concurrent.futures import json import logging import pathlib +from collections.abc import Iterator import onnx_ir as ir import safetensors.torch @@ -302,6 +304,25 @@ def _download_weights(model_id: str) -> dict[str, torch.Tensor]: downloads from HuggingFace Hub. Uses parallel downloads when multiple safetensors shards exist. """ + state_dict: dict[str, torch.Tensor] = {} + for shard in iter_weight_shards(model_id): + state_dict.update(shard) + return _dequantize_fp8_weights(state_dict) + + +def iter_weight_shards(model_id: str) -> Iterator[dict[str, torch.Tensor]]: + """Yield raw checkpoint state dictionaries one safetensors shard at a time. + + Index entries may reference safe relative subpaths, as used by unified + checkpoints whose transformer and vision weights share a top-level index. + Keeping shard boundaries lets composite exporters route a tensor to + multiple component graphs before advancing to the next file. + + FP8 dequantization is intentionally not performed here: a weight and its + ``weight_scale_inv`` may live in different shards. Callers that need + dequantized weights must merge the shards first, as :func:`_download_weights` + does, or implement an index-aware cross-shard scale resolver. + """ paths = _local_weight_paths(pathlib.Path(model_id)) if paths is None: try: @@ -317,9 +338,5 @@ def _download_weights(model_id: str) -> dict[str, torch.Tensor]: paths = _parallel_download(model_id, all_files, desc="safetensors") - state_dict: dict[str, torch.Tensor] = {} for path in tqdm.tqdm(paths, desc="Loading weights"): - state_dict.update(safetensors.torch.load_file(path)) - - state_dict = _dequantize_fp8_weights(state_dict) - return state_dict + yield safetensors.torch.load_file(path) diff --git a/src/mobius/_weight_loading_test.py b/src/mobius/_weight_loading_test.py index 1c7ee2569..6941453b9 100644 --- a/src/mobius/_weight_loading_test.py +++ b/src/mobius/_weight_loading_test.py @@ -32,7 +32,7 @@ from mobius._builder import build_from_module from mobius._model_package import ModelPackage from mobius._testing import make_config -from mobius._weight_loading import _download_weights, apply_weights +from mobius._weight_loading import _download_weights, apply_weights, iter_weight_shards from mobius.models.base import CausalLMModel from mobius.tasks import CausalLMTask, ModelTask @@ -189,6 +189,71 @@ def _unexpected_hub_call(*_args, **_kwargs): assert torch.equal(state_dict["a.weight"], shard_a["a.weight"]) assert torch.equal(state_dict["b.weight"], shard_b["b.weight"]) + def test_local_weight_shards_are_yielded_individually(self, tmp_path, monkeypatch): + shard_a = {"a.weight": torch.ones(1)} + shard_b = {"b.weight": torch.zeros(1)} + safetensors.torch.save_file(shard_a, str(tmp_path / "part-a.safetensors")) + safetensors.torch.save_file(shard_b, str(tmp_path / "nested-part-b.safetensors")) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {}, + "weight_map": { + "a.weight": "part-a.safetensors", + "b.weight": "nested-part-b.safetensors", + }, + } + ) + ) + + def _unexpected_hub_call(*_args, **_kwargs): + raise AssertionError("local checkpoint should not call hf_hub_download") + + monkeypatch.setattr("mobius._weight_loading.hf_hub_download", _unexpected_hub_call) + + shards = list(iter_weight_shards(str(tmp_path))) + + assert [set(shard) for shard in shards] == [{"b.weight"}, {"a.weight"}] + assert torch.equal(shards[0]["b.weight"], shard_b["b.weight"]) + assert torch.equal(shards[1]["a.weight"], shard_a["a.weight"]) + + def test_fp8_scale_in_different_shard_is_applied_after_merge(self, tmp_path, monkeypatch): + weight = torch.full((2, 2), 4.0, dtype=torch.float8_e4m3fn) + scale = torch.tensor(8.0) + safetensors.torch.save_file( + {"layer.weight": weight}, + str(tmp_path / "weight.safetensors"), + ) + safetensors.torch.save_file( + {"layer.weight_scale_inv": scale}, + str(tmp_path / "scale.safetensors"), + ) + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "metadata": {}, + "weight_map": { + "layer.weight": "weight.safetensors", + "layer.weight_scale_inv": "scale.safetensors", + }, + } + ) + ) + + def _unexpected_hub_call(*_args, **_kwargs): + raise AssertionError("local checkpoint should not call hf_hub_download") + + monkeypatch.setattr("mobius._weight_loading.hf_hub_download", _unexpected_hub_call) + + state_dict = _download_weights(str(tmp_path)) + + assert state_dict["layer.weight"].dtype == torch.bfloat16 + torch.testing.assert_close( + state_dict["layer.weight"], + torch.full((2, 2), 32.0, dtype=torch.bfloat16), + ) + assert "layer.weight_scale_inv" not in state_dict + def test_local_directory_without_safetensors_raises_without_hub( self, tmp_path, monkeypatch ): diff --git a/src/mobius/_world_model_builder.py b/src/mobius/_world_model_builder.py new file mode 100644 index 000000000..18bc9f6e1 --- /dev/null +++ b/src/mobius/_world_model_builder.py @@ -0,0 +1,129 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Registry and public build entry point for compositional world models.""" + +from __future__ import annotations + +__all__ = [ + "WorldModelBuilderRegistry", + "build_world_model", + "world_model_registry", +] + +import json +import pathlib +from collections.abc import Callable +from typing import Any + +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import EntryNotFoundError + +from mobius._pipeline import PipelinePackage + +WorldModelBuilder = Callable[..., PipelinePackage] + + +class WorldModelBuilderRegistry: + """Map checkpoint ``model_type`` values to compositional exporters.""" + + def __init__(self) -> None: + self._builders: dict[str, WorldModelBuilder] = {} + + def register(self, model_type: str, builder: WorldModelBuilder) -> None: + """Register one world-model pipeline builder. + + Re-registering the same callable is idempotent. Replacing an existing + builder is rejected so import order cannot silently change export + semantics. + """ + if not model_type: + raise ValueError("World-model model_type must be non-empty") + existing = self._builders.get(model_type) + if existing is builder: + return + if existing is not None: + raise ValueError(f"World-model builder for {model_type!r} is already registered") + self._builders[model_type] = builder + + def get(self, model_type: str) -> WorldModelBuilder: + """Return the registered builder for *model_type*.""" + try: + return self._builders[model_type] + except KeyError as error: + supported = ", ".join(sorted(self._builders)) or "" + raise ValueError( + f"No complete world-model pipeline is registered for model_type " + f"{model_type!r}. Registered types: {supported}." + ) from error + + def model_types(self) -> tuple[str, ...]: + """Return registered model types in deterministic order.""" + return tuple(sorted(self._builders)) + + +def _load_model_type(model_id: str) -> str: + root = pathlib.Path(model_id) + candidates = ("config.json", "model_index.json") + for filename in candidates: + if root.is_dir(): + path = root / filename + if not path.is_file(): + continue + else: + try: + path = pathlib.Path(hf_hub_download(repo_id=model_id, filename=filename)) + except EntryNotFoundError: + continue + with path.open(encoding="utf-8") as handle: + config = json.load(handle) + field = "model_type" if filename == "config.json" else "_class_name" + model_type = config.get(field) + if isinstance(model_type, str) and model_type: + return model_type + raise ValueError( + f"Checkpoint {model_id!r} has neither a non-empty config.json model_type " + "nor a model_index.json _class_name." + ) + + +world_model_registry = WorldModelBuilderRegistry() + + +def build_world_model( + model_id: str, + *, + dtype: Any | None = None, + load_weights: bool = True, + execution_provider: str = "default", + trace_optimization: bool = False, + **options: Any, +) -> PipelinePackage: + """Build every neural component of a registered world-model checkpoint. + + Unlike :func:`mobius.build`, which exports one architecture task, this + entry point dispatches to a model-specific pipeline builder that may + combine reasoners, dynamics or diffusion generators, tokenizers/codecs, + observation decoders, reward models, and action policies. + """ + model_type = _load_model_type(model_id) + builder = world_model_registry.get(model_type) + return builder( + model_id, + dtype=dtype, + load_weights=load_weights, + execution_provider=execution_provider, + trace_optimization=trace_optimization, + **options, + ) + + +def _register_builtin_world_models() -> None: + from mobius._cosmos3_edge_world_model import build_cosmos3_edge_world_model + from mobius._cosmos3_world_model import build_cosmos3_world_model + + world_model_registry.register("cosmos3_edge", build_cosmos3_edge_world_model) + world_model_registry.register("cosmos3_omni", build_cosmos3_world_model) + + +_register_builtin_world_models() diff --git a/src/mobius/_world_model_builder_test.py b/src/mobius/_world_model_builder_test.py new file mode 100644 index 000000000..868cd9301 --- /dev/null +++ b/src/mobius/_world_model_builder_test.py @@ -0,0 +1,90 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import json +from unittest import mock + +import pytest + +from mobius._world_model_builder import ( + WorldModelBuilderRegistry, + build_world_model, + world_model_registry, +) + + +def test_registry_rejects_conflicting_builder() -> None: + registry = WorldModelBuilderRegistry() + + def first(*_args, **_kwargs): + return object() + + def second(*_args, **_kwargs): + return object() + + registry.register("example", first) + registry.register("example", first) + with pytest.raises(ValueError, match="already registered"): + registry.register("example", second) + + +def test_registry_reports_supported_model_types() -> None: + registry = WorldModelBuilderRegistry() + registry.register("zeta", mock.Mock()) + registry.register("alpha", mock.Mock()) + + assert registry.model_types() == ("alpha", "zeta") + with pytest.raises(ValueError, match=r"alpha, zeta"): + registry.get("missing") + + +def test_build_world_model_dispatches_local_checkpoint(tmp_path) -> None: + (tmp_path / "config.json").write_text(json.dumps({"model_type": "test_world"})) + package = object() + builder = mock.Mock(return_value=package) + + with mock.patch.dict(world_model_registry._builders, {"test_world": builder}, clear=True): + result = build_world_model( + str(tmp_path), + dtype="bf16", + load_weights=False, + execution_provider="cuda", + trace_optimization=True, + custom_option=7, + ) + + assert result is package + builder.assert_called_once_with( + str(tmp_path), + dtype="bf16", + load_weights=False, + execution_provider="cuda", + trace_optimization=True, + custom_option=7, + ) + + +def test_build_world_model_rejects_missing_model_type(tmp_path) -> None: + (tmp_path / "config.json").write_text("{}") + + with pytest.raises(ValueError, match="model_type"): + build_world_model(str(tmp_path), load_weights=False) + + +def test_build_world_model_dispatches_pure_diffusers_pipeline(tmp_path) -> None: + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "ExampleWorldPipeline"}) + ) + package = object() + builder = mock.Mock(return_value=package) + + with mock.patch.dict( + world_model_registry._builders, + {"ExampleWorldPipeline": builder}, + clear=True, + ): + result = build_world_model(str(tmp_path), load_weights=False) + + assert result is package diff --git a/src/mobius/_world_model_config.py b/src/mobius/_world_model_config.py new file mode 100644 index 000000000..60aacecd9 --- /dev/null +++ b/src/mobius/_world_model_config.py @@ -0,0 +1,279 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Immutable, model-agnostic configuration for world-model pipeline builders. + +A world-model exporter composes several independently built graphs into one +:class:`mobius.PipelinePackage`. Every such exporter needs the same three +groups of settings: + +* how each component graph is built (dtype, weights, execution provider); +* how a runtime should sample tokens and drive the iterative scheduler; +* which checkpoint the package came from and which profile it implements. + +These dataclasses carry those settings as frozen values so a builder passes a +single object through its private helpers instead of repeating loose keyword +arguments at every call site. Nothing here is specific to one model family: +component names, action domains, scheduler tuning values, and topology stay in +the model-specific exporter modules. +""" + +from __future__ import annotations + +__all__ = [ + "WORLD_MODEL_PROFILE", + "WorldModelBuildConfig", + "WorldModelGenerationConfig", + "WorldModelPipelineConfig", +] + +import dataclasses +from collections.abc import Iterable, Mapping +from types import MappingProxyType +from typing import Any + +import onnx_ir as ir + +from mobius._builder import resolve_dtype + +#: Manifest ``profile`` metadata value shared by every world-model package. +WORLD_MODEL_PROFILE = "world-model" + +#: Manifest profile version emitted by world-model builders. +DEFAULT_PROFILE_VERSION = "1.0" + + +@dataclasses.dataclass(frozen=True) +class WorldModelBuildConfig: + """How each component graph of a world model is built. + + Attributes: + dtype: Requested parameter dtype, as a Mobius dtype string (e.g. + ``"f16"``), an ``ir.DataType``, or ``None`` to keep each + component's checkpoint dtype. + load_weights: Whether checkpoint weights are streamed into the graphs. + execution_provider: Requested execution provider, or ``"default"`` to + let each component advertise a dtype-appropriate preference list. + trace_optimization: Whether graph-optimization tracing is enabled. + """ + + dtype: str | ir.DataType | None = None + load_weights: bool = True + execution_provider: str = "default" + trace_optimization: bool = False + + def __post_init__(self) -> None: + self.validate() + + def validate(self) -> None: + """Reject settings that cannot produce a loadable package.""" + if not isinstance(self.execution_provider, str) or not self.execution_provider: + raise ValueError("execution_provider must be a non-empty string") + # Resolving eagerly turns an unknown dtype string into a ValueError at + # configuration time rather than midway through a multi-graph build. + self.resolved_dtype() + + def resolved_dtype(self) -> ir.DataType | None: + """Return the requested dtype as an ``ir.DataType``, if one was given.""" + if self.dtype is None or isinstance(self.dtype, ir.DataType): + return self.dtype + return resolve_dtype(self.dtype) + + def preferred_execution_providers(self, dtype: ir.DataType) -> tuple[str, ...]: + """Return the provider preference order for a component of *dtype*. + + An explicit request is honoured verbatim. ``"default"`` expands to the + providers that can execute *dtype*: DirectML has no bfloat16 support, + so bfloat16 components omit it. + """ + if self.execution_provider != "default": + return (self.execution_provider,) + if dtype == ir.DataType.BFLOAT16: + return ("cuda", "cpu") + return ("cuda", "dml", "cpu") + + +@dataclasses.dataclass(frozen=True) +class WorldModelGenerationConfig: + """Runtime sampling and scheduling defaults carried into the manifest. + + The token-sampling fields mirror the HuggingFace ``generation_config.json`` + contract; the diffusion fields describe how many scheduler steps a runtime + should take by default and which per-mode scheduler overrides the exporter + recorded. All of them are opaque runtime hints: no graph is built from + them, so they only reach the manifest. + """ + + do_sample: bool = False + temperature: float = 1.0 + top_k: int = 50 + top_p: float = 1.0 + repetition_penalty: float = 1.0 + max_new_tokens: int | None = None + eos_token_ids: tuple[int, ...] = () + default_inference_steps: int = 1 + scheduler_mode_overrides: Mapping[str, Any] = dataclasses.field( + default_factory=dict, + ) + + def __post_init__(self) -> None: + object.__setattr__(self, "eos_token_ids", tuple(self.eos_token_ids)) + object.__setattr__( + self, + "scheduler_mode_overrides", + MappingProxyType(dict(self.scheduler_mode_overrides)), + ) + self.validate() + + def validate(self) -> None: + """Reject values a runtime could not act on.""" + for name in ("temperature", "top_p", "repetition_penalty"): + value = getattr(self, name) + if not isinstance(value, (int, float)) or isinstance(value, bool): + raise TypeError(f"{name} must be a number") + if value < 0: + raise ValueError(f"{name} must not be negative") + if isinstance(self.top_k, bool) or not isinstance(self.top_k, int): + raise TypeError("top_k must be an integer") + if self.top_k < 0: + raise ValueError("top_k must not be negative") + if self.max_new_tokens is not None and self.max_new_tokens <= 0: + raise ValueError("max_new_tokens must be positive when set") + if any( + isinstance(value, bool) or not isinstance(value, int) + for value in self.eos_token_ids + ): + raise TypeError("eos_token_ids must contain integers") + if self.default_inference_steps <= 0: + raise ValueError("default_inference_steps must be positive") + + @classmethod + def from_generation_config( + cls, + values: Mapping[str, Any] | None, + *, + default_inference_steps: int = 1, + scheduler_mode_overrides: Mapping[str, Any] | None = None, + ) -> WorldModelGenerationConfig: + """Parse a HuggingFace ``generation_config.json`` mapping. + + Missing keys — and keys explicitly set to ``null``, which real + checkpoints use to mean "unset" — fall back to the HuggingFace + defaults. ``eos_token_id`` may be a single id or a list; both + normalize to a tuple. + """ + mapping = dict(values or {}) + + def number(key: str, default: float) -> float: + value = mapping.get(key) + return default if value is None else float(value) + + eos: Any = mapping.get("eos_token_id", ()) + if isinstance(eos, int) and not isinstance(eos, bool): + eos = (eos,) + eos_token_ids: Iterable[int] = tuple(eos or ()) + max_new_tokens = mapping.get("max_new_tokens") + top_k = mapping.get("top_k") + return cls( + do_sample=bool(mapping.get("do_sample")), + temperature=number("temperature", 1.0), + top_k=50 if top_k is None else int(top_k), + top_p=number("top_p", 1.0), + repetition_penalty=number("repetition_penalty", 1.0), + max_new_tokens=None if max_new_tokens is None else int(max_new_tokens), + eos_token_ids=tuple(eos_token_ids), + default_inference_steps=default_inference_steps, + scheduler_mode_overrides=scheduler_mode_overrides or {}, + ) + + def sampling_manifest(self) -> dict[str, Any]: + """Return the deterministic ``sampling`` block of a decode stage.""" + return { + "do_sample": self.do_sample, + "temperature": self.temperature, + "top_k": self.top_k, + "top_p": self.top_p, + "repetition_penalty": self.repetition_penalty, + } + + def stop_manifest(self, *, max_sequence_length: int | None) -> dict[str, Any]: + """Return the deterministic ``stop`` block of a decode stage.""" + return { + "kind": "token_ids", + "eos_token_ids": list(self.eos_token_ids), + "max_sequence_length": max_sequence_length, + } + + def max_tokens_manifest(self, *, limit: int | None) -> dict[str, Any]: + """Return the deterministic ``max_tokens`` block of a decode stage. + + A checkpoint without ``max_new_tokens`` cannot bound its own decode + loop, so the runtime must supply the budget. + """ + return { + "default": self.max_new_tokens, + "required_override": self.max_new_tokens is None, + "limit": limit, + } + + def scheduler_mode_overrides_manifest(self) -> dict[str, Any]: + """Return a mutable copy of the per-mode scheduler overrides.""" + return dict(self.scheduler_mode_overrides) + + +@dataclasses.dataclass(frozen=True) +class WorldModelPipelineConfig: + """Identity and runtime contract of one composed world-model package. + + Attributes: + model_id: Checkpoint the package was exported from, recorded as the + manifest ``source``. + model_type: Checkpoint ``model_type`` the package implements, recorded + as manifest ``model_type`` and used to derive the profile name. + build: How the component graphs were built. + generation: Runtime sampling and scheduling defaults. + extra_metadata: Additional manifest metadata entries applied last, so + a model family can attach its own top-level block. + profile_version: Version of the runtime profile implemented here. + """ + + model_id: str + model_type: str + build: WorldModelBuildConfig = dataclasses.field( + default_factory=WorldModelBuildConfig, + ) + generation: WorldModelGenerationConfig = dataclasses.field( + default_factory=WorldModelGenerationConfig, + ) + extra_metadata: Mapping[str, Any] = dataclasses.field(default_factory=dict) + profile_version: str = DEFAULT_PROFILE_VERSION + + def __post_init__(self) -> None: + object.__setattr__( + self, + "extra_metadata", + MappingProxyType(dict(self.extra_metadata)), + ) + self.validate() + + def validate(self) -> None: + """Reject an unidentifiable package.""" + if not isinstance(self.model_id, str) or not self.model_id: + raise ValueError("model_id must be a non-empty string") + if not isinstance(self.model_type, str) or not self.model_type: + raise ValueError("model_type must be a non-empty string") + if not isinstance(self.profile_version, str) or not self.profile_version: + raise ValueError("profile_version must be a non-empty string") + + @property + def profile_name(self) -> str: + """Manifest profile name derived from ``model_type``.""" + return self.model_type.replace("_", "-") + + def manifest_metadata(self) -> dict[str, Any]: + """Return the metadata entries every world-model package declares.""" + return { + "profile": WORLD_MODEL_PROFILE, + "model_type": self.model_type, + "source": self.model_id, + } diff --git a/src/mobius/_world_model_config_test.py b/src/mobius/_world_model_config_test.py new file mode 100644 index 000000000..55fd91cb6 --- /dev/null +++ b/src/mobius/_world_model_config_test.py @@ -0,0 +1,258 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +import dataclasses + +import onnx_ir as ir +import pytest + +from mobius import ( + WorldModelBuildConfig as PublicWorldModelBuildConfig, +) +from mobius import ( + WorldModelGenerationConfig as PublicWorldModelGenerationConfig, +) +from mobius import ( + WorldModelPipelineConfig as PublicWorldModelPipelineConfig, +) +from mobius._world_model_config import ( + WorldModelBuildConfig, + WorldModelGenerationConfig, + WorldModelPipelineConfig, +) + + +def test_world_model_configs_are_public() -> None: + assert PublicWorldModelBuildConfig is WorldModelBuildConfig + assert PublicWorldModelGenerationConfig is WorldModelGenerationConfig + assert PublicWorldModelPipelineConfig is WorldModelPipelineConfig + + +def test_build_config_resolves_dtype_aliases() -> None: + assert WorldModelBuildConfig().resolved_dtype() is None + assert WorldModelBuildConfig(dtype="f16").resolved_dtype() is ir.DataType.FLOAT16 + assert ( + WorldModelBuildConfig(dtype=ir.DataType.BFLOAT16).resolved_dtype() + is ir.DataType.BFLOAT16 + ) + + +def test_build_config_rejects_unknown_settings() -> None: + with pytest.raises(ValueError, match="Unknown dtype"): + WorldModelBuildConfig(dtype="float128") + with pytest.raises(ValueError, match="execution_provider"): + WorldModelBuildConfig(execution_provider="") + + +def test_build_config_is_immutable() -> None: + config = WorldModelBuildConfig(dtype="f32") + + with pytest.raises(dataclasses.FrozenInstanceError): + config.dtype = "f16" # type: ignore[misc] + + +def test_preferred_execution_providers_follow_dtype_support() -> None: + default = WorldModelBuildConfig() + + assert default.preferred_execution_providers(ir.DataType.FLOAT) == ("cuda", "dml", "cpu") + # DirectML has no bfloat16 support, so it is omitted for bfloat16 graphs. + assert default.preferred_execution_providers(ir.DataType.BFLOAT16) == ("cuda", "cpu") + assert WorldModelBuildConfig(execution_provider="cpu").preferred_execution_providers( + ir.DataType.FLOAT16 + ) == ("cpu",) + + +def test_generation_config_parses_huggingface_mapping() -> None: + config = WorldModelGenerationConfig.from_generation_config( + { + "do_sample": True, + "temperature": 0.6, + "top_k": 20, + "top_p": 0.95, + "repetition_penalty": 1.05, + "max_new_tokens": 1024, + "eos_token_id": [151645, 151643], + }, + default_inference_steps=35, + ) + + assert config.do_sample is True + assert config.temperature == pytest.approx(0.6) + assert config.top_k == 20 + assert config.top_p == pytest.approx(0.95) + assert config.repetition_penalty == pytest.approx(1.05) + assert config.max_new_tokens == 1024 + assert config.eos_token_ids == (151645, 151643) + assert config.default_inference_steps == 35 + + +def test_generation_config_normalizes_scalar_and_absent_eos() -> None: + scalar = WorldModelGenerationConfig.from_generation_config({"eos_token_id": 7}) + explicit_null = WorldModelGenerationConfig.from_generation_config({"eos_token_id": None}) + + assert scalar.eos_token_ids == (7,) + assert explicit_null.eos_token_ids == () + assert WorldModelGenerationConfig.from_generation_config({}).eos_token_ids == () + assert WorldModelGenerationConfig.from_generation_config(None).eos_token_ids == () + + +def test_generation_config_defaults_match_huggingface() -> None: + config = WorldModelGenerationConfig.from_generation_config({}) + + assert config.sampling_manifest() == { + "do_sample": False, + "temperature": 1.0, + "top_k": 50, + "top_p": 1.0, + "repetition_penalty": 1.0, + } + + +def test_generation_config_treats_explicit_nulls_as_unset() -> None: + config = WorldModelGenerationConfig.from_generation_config( + { + "do_sample": None, + "temperature": None, + "top_k": None, + "top_p": None, + "repetition_penalty": None, + "max_new_tokens": None, + } + ) + + assert config.sampling_manifest() == { + "do_sample": False, + "temperature": 1.0, + "top_k": 50, + "top_p": 1.0, + "repetition_penalty": 1.0, + } + assert config.max_new_tokens is None + + +def test_generation_manifest_blocks_are_deterministic() -> None: + config = WorldModelGenerationConfig( + do_sample=True, + temperature=0.6, + top_k=20, + top_p=0.95, + repetition_penalty=1.05, + max_new_tokens=64, + eos_token_ids=(2,), + default_inference_steps=4, + scheduler_mode_overrides={"action": {"flow_shift": 10.0}}, + ) + + assert list(config.sampling_manifest()) == [ + "do_sample", + "temperature", + "top_k", + "top_p", + "repetition_penalty", + ] + assert config.stop_manifest(max_sequence_length=128) == { + "kind": "token_ids", + "eos_token_ids": [2], + "max_sequence_length": 128, + } + assert config.max_tokens_manifest(limit=128) == { + "default": 64, + "required_override": False, + "limit": 128, + } + assert config.scheduler_mode_overrides_manifest() == {"action": {"flow_shift": 10.0}} + + +def test_generation_config_requires_runtime_budget_without_max_new_tokens() -> None: + config = WorldModelGenerationConfig() + + assert config.max_tokens_manifest(limit=None) == { + "default": None, + "required_override": True, + "limit": None, + } + + +def test_generation_config_copies_and_freezes_overrides() -> None: + overrides = {"action": {"flow_shift": 10.0}} + config = WorldModelGenerationConfig(scheduler_mode_overrides=overrides) + overrides["action"] = {"flow_shift": 1.0} + + assert config.scheduler_mode_overrides["action"] == {"flow_shift": 10.0} + with pytest.raises(TypeError): + config.scheduler_mode_overrides["extra"] = {} # type: ignore[index] + # The manifest mapping is a mutable copy, so callers cannot corrupt the config. + manifest = config.scheduler_mode_overrides_manifest() + manifest["extra"] = {} + assert "extra" not in config.scheduler_mode_overrides + + +def test_generation_config_rejects_unusable_values() -> None: + with pytest.raises(ValueError, match="temperature"): + WorldModelGenerationConfig(temperature=-1.0) + with pytest.raises(ValueError, match="top_k"): + WorldModelGenerationConfig(top_k=-1) + with pytest.raises(ValueError, match="max_new_tokens"): + WorldModelGenerationConfig(max_new_tokens=0) + with pytest.raises(ValueError, match="default_inference_steps"): + WorldModelGenerationConfig(default_inference_steps=0) + with pytest.raises(TypeError, match="eos_token_ids"): + WorldModelGenerationConfig(eos_token_ids=("2",)) # type: ignore[arg-type] + + +def test_pipeline_config_derives_profile_and_metadata() -> None: + config = WorldModelPipelineConfig( + model_id="nvidia/Cosmos3-Omni", + model_type="cosmos3_omni", + ) + + assert config.profile_name == "cosmos3-omni" + assert config.profile_version == "1.0" + assert config.manifest_metadata() == { + "profile": "world-model", + "model_type": "cosmos3_omni", + "source": "nvidia/Cosmos3-Omni", + } + assert list(config.manifest_metadata()) == ["profile", "model_type", "source"] + + +def test_pipeline_config_holds_build_and_generation_settings() -> None: + build = WorldModelBuildConfig(dtype="bf16", execution_provider="cuda") + generation = WorldModelGenerationConfig(default_inference_steps=50) + config = WorldModelPipelineConfig( + model_id="example/world", + model_type="example_world", + build=build, + generation=generation, + extra_metadata={"edge": {"policy": None}}, + ) + + assert config.build is build + assert config.generation is generation + assert config.extra_metadata["edge"] == {"policy": None} + with pytest.raises(TypeError): + config.extra_metadata["edge"] = {} # type: ignore[index] + + +def test_pipeline_config_requires_identity() -> None: + with pytest.raises(ValueError, match="model_id"): + WorldModelPipelineConfig(model_id="", model_type="example") + with pytest.raises(ValueError, match="model_type"): + WorldModelPipelineConfig(model_id="example/world", model_type="") + with pytest.raises(ValueError, match="profile_version"): + WorldModelPipelineConfig( + model_id="example/world", + model_type="example", + profile_version="", + ) + + +def test_pipeline_config_defaults_are_independent() -> None: + first = WorldModelPipelineConfig(model_id="a/b", model_type="t") + second = WorldModelPipelineConfig(model_id="a/b", model_type="t") + + assert first.build == second.build + assert first.generation == second.generation + assert first.build is not second.build diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 3baeb6e77..92bdc393b 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -47,7 +47,9 @@ "LoRALinear", "MLP", "MLPMultiModalProjector", - "Cosmos3EdgeMultiModalProjector", + "Cosmos3EdgePatchMerger", + "Cosmos3EdgeVisionEmbeddings", + "Cosmos3EdgeVisionTower", "MoELayer", "OffsetRMSNorm", "PatchEmbed", @@ -132,6 +134,15 @@ Conv2dNoBias, ConvTranspose2d, ) +from mobius.components._cosmos3_edge_vision import ( + Cosmos3EdgePatchMerger as Cosmos3EdgePatchMerger, +) +from mobius.components._cosmos3_edge_vision import ( + Cosmos3EdgeVisionEmbeddings as Cosmos3EdgeVisionEmbeddings, +) +from mobius.components._cosmos3_edge_vision import ( + Cosmos3EdgeVisionTower as Cosmos3EdgeVisionTower, +) from mobius.components._decoder import ( DecoderLayer, PostNormDecoderLayer, @@ -170,9 +181,6 @@ SparseMixerGate, TopKGate, ) -from mobius.components._multimodal import ( - Cosmos3EdgeMultiModalProjector as Cosmos3EdgeMultiModalProjector, -) from mobius.components._multimodal import ( Gemma3MultiModalProjector as Gemma3MultiModalProjector, ) diff --git a/src/mobius/components/_cosmos3_edge_vision.py b/src/mobius/components/_cosmos3_edge_vision.py new file mode 100644 index 000000000..6c4f45617 --- /dev/null +++ b/src/mobius/components/_cosmos3_edge_vision.py @@ -0,0 +1,331 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""NVIDIA Cosmos3-Edge SigLIP2 vision tower and pixel-shuffle merger projector. + +Replicates ``Cosmos3EdgeVisionModel`` / ``Cosmos3EdgePatchMerger`` from +``transformers.models.cosmos3_edge.modular_cosmos3_edge`` (cross-checked +against ``vllm/model_executor/models/cosmos3_edge.py``). + +The tower is a **variable-resolution, packed** SigLIP2 encoder, not a +fixed-square ViT: + +- ``pixel_values`` arrives **already patchified** as + ``[total_patches, patch_size**2 * channels * temporal_patch_size]``. The + processor emits patches in **block-major** order (each ``merge x merge`` + block of adjacent patches is contiguous) and stores the values inside a + patch as ``(patch_h, patch_w, channel)`` — channel-**last**. The checkpoint + therefore ships ``embeddings.patch_embedding`` as an ``nn.Linear`` + ``[hidden, patch_dim]``, not a ``Conv2d`` kernel. +- The learned ``num_patches`` (16x16) position table is bilinearly resampled + (antialiased, ``align_corners=False``) to the item's ``(grid_h, grid_w)`` + and then reordered into the same block-major layout. +- Attention is non-causal and runs **independently per frame** + (``cu_seqlens`` delimits every frame), which for a single packed item with a + shared ``(grid_h, grid_w)`` is exactly batched attention over ``grid_t`` + equal-length sequences. + +The merger projector groups ``spatial_merge_size**2`` consecutive +(block-major) patches, applies the pre-shuffle ``LayerNorm`` +(``use_postshuffle_norm=false`` in the public checkpoint), then +``linear_fc1 → GELU → linear_fc2``. +""" + +from __future__ import annotations + +import math + +import numpy as np +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius.components._vision import VisionEncoder, VisionLayerNorm, _VisionLinear + +__all__ = [ + "Cosmos3EdgePatchMerger", + "Cosmos3EdgeVisionEmbeddings", + "Cosmos3EdgeVisionTower", +] + + +def _resample_weights(op: OpBuilder, source: int, target: ir.Value) -> ir.Value: + """Build the 1-D antialiased-bilinear resample matrix ``[target, source]``. + + Reproduces PyTorch's ``F.interpolate(..., mode="bilinear", + align_corners=False, antialias=True)`` separable filter: + + - ``scale = source / target`` (``align_corners=False`` area scale), + - ``support = max(scale, 1)`` — antialiasing only widens the triangle + filter when *downsampling*, + - ``w[i, j] = max(0, 1 - |j + 0.5 - scale * (i + 0.5)| / support)``, + renormalised over the ``source`` window. + + ONNX ``Resize(antialias=1)`` is *not* used: onnxruntime's antialias filter + does not match PyTorch's when downsampling, and CUDA/DML do not implement + it at all. An explicit ``[target, source]`` matrix is exact and portable. + """ + source_f = float(source) + # centers[i] = (i + 0.5) * source / target + target_f = op.Cast(target, to=ir.DataType.FLOAT) + scale = op.Div(op.Constant(value_float=source_f), target_f) + positions = op.Cast( + op.Range( + op.Constant(value_int=0), + op.Reshape(target, op.Constant(value_ints=[])), + op.Constant(value_int=1), + ), + to=ir.DataType.FLOAT, + ) + centers = op.Mul(op.Add(positions, op.Constant(value_float=0.5)), scale) + support = op.Max(scale, op.Constant(value_float=1.0)) + + # Source pixel centres: [0.5, 1.5, ..., source - 0.5] + source_centers = op.Constant( + value=ir.tensor(np.arange(source, dtype=np.float32) + 0.5, name="source_centers") + ) + distance = op.Sub(op.Unsqueeze(source_centers, [0]), op.Unsqueeze(centers, [-1])) + weights = op.Sub( + op.Constant(value_float=1.0), + op.Div(op.Abs(distance), op.Unsqueeze(support, [0])), + ) + weights = op.Relu(weights) # triangle filter: max(0, 1 - |x| / support) + total = op.ReduceSum(weights, op.Constant(value_ints=[-1]), keepdims=1) + return op.Div(weights, total) + + +class Cosmos3EdgeVisionEmbeddings(nn.Module): + """Packed patch embedding + resampled block-major position embedding. + + HF reference: ``Cosmos3EdgeVisionEmbeddings`` (a ``Siglip2VisionEmbeddings`` + subclass whose ``resize_positional_embeddings`` additionally reorders the + resampled grid into the processor's block-major 2x2 layout). + """ + + def __init__( + self, + *, + hidden_size: int, + patch_size: int, + num_channels: int, + num_patches: int, + temporal_patch_size: int = 1, + spatial_merge_size: int = 2, + ): + super().__init__() + grid = math.isqrt(num_patches) + if grid * grid != num_patches: + raise ValueError( + f"num_patches must form a square reference grid, got {num_patches}" + ) + self.position_embedding_size = grid + self.spatial_merge_size = spatial_merge_size + self.hidden_size = hidden_size + self.patch_dim = patch_size * patch_size * num_channels * temporal_patch_size + # nn.Linear over flattened (patch_h, patch_w, channel) values. + self.patch_embedding = _VisionLinear(self.patch_dim, hidden_size) + self.position_embedding = nn.Parameter( + [num_patches, hidden_size], name="position_embedding.weight" + ) + + def _resized_position_embedding( + self, op: OpBuilder, grid_h: ir.Value, grid_w: ir.Value + ) -> ir.Value: + """Resample the learned grid to ``(grid_h, grid_w)``, block-major. + + ``[g, g, D]`` -> bilinear/antialias -> ``[H, W, D]`` -> block-major + ``[H*W, D]`` so element ``k`` matches the processor's patch ``k``. + """ + grid = self.position_embedding_size + dim = self.hidden_size + merge = self.spatial_merge_size + + table = op.Reshape(self.position_embedding, op.Constant(value_ints=[grid, grid * dim])) + # Height pass: [H, g] @ [g, g*D] -> [H, g, D] + height_weights = _resample_weights(op, grid, grid_h) + table = op.MatMul(op.CastLike(height_weights, table), table) + table = op.Reshape( + table, + op.Concat( + op.Reshape(grid_h, op.Constant(value_ints=[1])), + op.Constant(value_ints=[grid, dim]), + axis=0, + ), + ) + # Width pass: [W, g] @ [g, H*D] -> [W, H, D] + table = op.Transpose(table, perm=[1, 0, 2]) # [g, H, D] + table = op.Reshape(table, op.Constant(value_ints=[grid, -1])) + width_weights = _resample_weights(op, grid, grid_w) + table = op.MatMul(op.CastLike(width_weights, table), table) + table = op.Reshape( + table, + op.Concat( + op.Reshape(grid_w, op.Constant(value_ints=[1])), + op.Constant(value_ints=[-1, dim]), + axis=0, + ), + ) + table = op.Transpose(table, perm=[1, 0, 2]) # [H, W, D] + + # Block-major reorder: (H, W) -> (H/m, m, W/m, m) -> (H/m, W/m, m, m) + merge_const = op.Constant(value_ints=[merge]) + blocks_h = op.Div(op.Reshape(grid_h, op.Constant(value_ints=[1])), merge_const) + blocks_w = op.Div(op.Reshape(grid_w, op.Constant(value_ints=[1])), merge_const) + table = op.Reshape( + table, + op.Concat( + blocks_h, + merge_const, + blocks_w, + merge_const, + op.Constant(value_ints=[dim]), + axis=0, + ), + ) + table = op.Transpose(table, perm=[0, 2, 1, 3, 4]) + return op.Reshape(table, op.Constant(value_ints=[-1, dim])) + + def forward( + self, + op: OpBuilder, + pixel_values: ir.Value, + grid_t: ir.Value, + grid_h: ir.Value, + grid_w: ir.Value, + ): + # pixel_values: [T*H*W, patch_dim] -> [T*H*W, hidden] + patch_embeds = self.patch_embedding(op, pixel_values) + # Split frames out so the per-frame position grid broadcasts over T, + # which is what the reference's ``resized.repeat(temporal, 1)`` does. + frame_shape = op.Concat( + op.Reshape(grid_t, op.Constant(value_ints=[1])), + op.Reshape(op.Mul(grid_h, grid_w), op.Constant(value_ints=[1])), + op.Constant(value_ints=[self.hidden_size]), + axis=0, + ) + patch_embeds = op.Reshape(patch_embeds, frame_shape) + position_embeds = self._resized_position_embedding(op, grid_h, grid_w) + # [T, H*W, D] + [H*W, D] -> [T, H*W, D] + return op.Add(patch_embeds, position_embeds) + + +class Cosmos3EdgeVisionTower(nn.Module): + """SigLIP2 encoder stack over packed, per-frame attention sequences. + + ``pixel_values [T*H*W, patch_dim]`` -> ``[T*H*W, vision_hidden]``. The + patches are reshaped to ``[T, H*W, D]`` so the ONNX ``Attention`` op sees + ``T`` independent sequences — the exact semantics of the reference's + per-frame ``cu_seqlens`` (every frame of a packed item shares one + ``(H, W)`` grid, so all sequences have equal length and no mask is needed). + """ + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + num_hidden_layers: int, + num_attention_heads: int, + patch_size: int, + num_channels: int, + num_patches: int, + norm_eps: float = 1e-6, + temporal_patch_size: int = 1, + spatial_merge_size: int = 2, + ): + super().__init__() + self.hidden_size = hidden_size + self.embeddings = Cosmos3EdgeVisionEmbeddings( + hidden_size=hidden_size, + patch_size=patch_size, + num_channels=num_channels, + num_patches=num_patches, + temporal_patch_size=temporal_patch_size, + spatial_merge_size=spatial_merge_size, + ) + self.encoder = VisionEncoder( + num_layers=num_hidden_layers, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_heads=num_attention_heads, + norm_eps=norm_eps, + ) + self.post_layernorm = VisionLayerNorm(hidden_size, eps=norm_eps) + + def forward( + self, + op: OpBuilder, + pixel_values: ir.Value, + grid_t: ir.Value, + grid_h: ir.Value, + grid_w: ir.Value, + ): + # [T, H*W, D]: one non-causal attention sequence per frame. + hidden_states = self.embeddings(op, pixel_values, grid_t, grid_h, grid_w) + hidden_states = self.encoder(op, hidden_states) + hidden_states = self.post_layernorm(op, hidden_states) + # Back to packed rows: [T, H*W, D] -> [T*H*W, D] + return op.Reshape(hidden_states, op.Constant(value_ints=[-1, self.hidden_size])) + + +class Cosmos3EdgePatchMerger(nn.Module): + """Cosmos3-Edge pixel-shuffle merger projector. + + ``LayerNorm -> spatial merge -> linear_fc1 -> GELU -> linear_fc2``. + + The processor already emits patches in block-major order, so merging is a + plain reshape of ``spatial_merge_size**2`` *consecutive* rows — matching + ``Cosmos3EdgePatchMerger.forward``'s + ``x.reshape(-1, spatial_merge_size**2, input_hidden_size)``. + + ``use_postshuffle_norm`` selects whether the ``LayerNorm`` normalises the + raw ``vision_hidden_size`` features (``false``, the public checkpoint) or + the merged ``spatial_merge_size**2 * vision_hidden_size`` vector. + + HF weights (``model.projector.*``): ``norm.{weight,bias}``, + ``linear_fc1.{weight,bias}``, ``linear_fc2.{weight,bias}``. + """ + + def __init__( + self, + *, + vision_hidden_size: int, + text_hidden_size: int, + intermediate_size: int, + spatial_merge_size: int = 2, + use_postshuffle_norm: bool = False, + norm_eps: float = 1e-6, + ): + super().__init__() + if spatial_merge_size <= 0: + raise ValueError(f"spatial_merge_size must be positive, got {spatial_merge_size}") + self.spatial_merge_size = spatial_merge_size + self.vision_hidden_size = vision_hidden_size + self.merged_size = vision_hidden_size * spatial_merge_size * spatial_merge_size + self.use_postshuffle_norm = use_postshuffle_norm + self.norm = VisionLayerNorm( + self.merged_size if use_postshuffle_norm else vision_hidden_size, + eps=norm_eps, + ) + self.linear_fc1 = _VisionLinear(self.merged_size, intermediate_size) + self.linear_fc2 = _VisionLinear(intermediate_size, text_hidden_size) + + def forward(self, op: OpBuilder, vision_features: ir.Value): + # vision_features: [total_patches, vision_hidden] + merged_shape = op.Constant(value_ints=[-1, self.merged_size]) + if self.use_postshuffle_norm: + hidden = self.norm(op, op.Reshape(vision_features, merged_shape)) + else: + grouped = op.Reshape( + vision_features, + op.Constant( + value_ints=[ + -1, + self.spatial_merge_size * self.spatial_merge_size, + self.vision_hidden_size, + ] + ), + ) + hidden = op.Reshape(self.norm(op, grouped), merged_shape) + hidden = self.linear_fc1(op, hidden) + hidden = op.Gelu(hidden) # nn.GELU() -> exact erf GELU + return self.linear_fc2(op, hidden) diff --git a/src/mobius/components/_multimodal.py b/src/mobius/components/_multimodal.py index 9e6c7b9f9..f1108f665 100644 --- a/src/mobius/components/_multimodal.py +++ b/src/mobius/components/_multimodal.py @@ -23,7 +23,6 @@ from mobius.components._common import Linear from mobius.components._rms_norm import RMSNorm -from mobius.components._vision import VisionLayerNorm if TYPE_CHECKING: import onnx_ir as ir @@ -118,90 +117,6 @@ def forward(self, op: OpBuilder, vision_features: ir.Value): return hidden -class Cosmos3EdgeMultiModalProjector(nn.Module): - """Cosmos3-Edge pixel-shuffle merger projector. - - ``LayerNorm → spatial 2x2 pixel-shuffle → Linear(fc1) → GELU → Linear(fc2)`` - - The SigLIP vision encoder emits a fixed ``grid x grid`` patch grid - (``num_patches`` patches, e.g. 16x16 = 256). ``use_postshuffle_norm=false`` - means the ``LayerNorm`` is applied on the raw ``vision_hidden_size`` (1152) - features **before** the spatial merge. The merge concatenates each - ``spatial_merge_size x spatial_merge_size`` block of adjacent patches into a - single ``spatial_merge_size**2 * vision_hidden_size`` (4608) vector, which - ``linear_fc1`` maps to ``intermediate_size`` (11520) and ``linear_fc2`` maps - to ``text_hidden_size`` (2048). - - HF weights (``model.projector.*``): - - ``norm.{weight,bias}`` (pre-shuffle LayerNorm) - - ``linear_fc1.{weight,bias}`` - - ``linear_fc2.{weight,bias}`` - """ - - def __init__( - self, - vision_hidden_size: int, - text_hidden_size: int, - intermediate_size: int, - grid_size: int, - spatial_merge_size: int = 2, - norm_eps: float = 1e-6, - ): - super().__init__() - if grid_size <= 0: - raise ValueError(f"grid_size must be positive, got {grid_size}") - if spatial_merge_size <= 0: - raise ValueError(f"spatial_merge_size must be positive, got {spatial_merge_size}") - if grid_size % spatial_merge_size != 0: - raise ValueError( - f"grid_size ({grid_size}) must be divisible by " - f"spatial_merge_size ({spatial_merge_size})" - ) - self._grid = grid_size - self._ms = spatial_merge_size - self._vision_hidden = vision_hidden_size - merged_dim = vision_hidden_size * spatial_merge_size * spatial_merge_size - # Pre-shuffle LayerNorm over the raw vision hidden size. - self.norm = VisionLayerNorm(vision_hidden_size, eps=norm_eps) - self.linear_fc1 = Linear(merged_dim, intermediate_size, bias=True) - self.linear_fc2 = Linear(intermediate_size, text_hidden_size, bias=True) - - def forward(self, op: OpBuilder, vision_features: ir.Value): - # vision_features: [batch, grid*grid, vision_hidden] - ms = self._ms - g = self._grid - gm = g // ms - d = self._vision_hidden - - # Pre-shuffle LayerNorm (use_postshuffle_norm=false). - x = self.norm(op, vision_features) - - batch = op.Shape(vision_features, start=0, end=1) # dynamic [1] - - # [B, g*g, D] -> [B, g/ms, ms, g/ms, ms, D] - shape_6d = op.Concat( - batch, - op.Constant(value_ints=[gm, ms, gm, ms, d]), - axis=0, - ) - x = op.Reshape(x, shape_6d) - # Group hidden dim outermost per merged block (HF F.unfold ordering): - # [B, g/ms, ms, g/ms, ms, D] -> [B, g/ms, g/ms, D, ms, ms] - x = op.Transpose(x, perm=[0, 1, 3, 5, 2, 4]) - # Flatten to [B, (g/ms)^2, D*ms*ms] - shape_3d = op.Concat( - batch, - op.Constant(value_ints=[gm * gm, d * ms * ms]), - axis=0, - ) - x = op.Reshape(x, shape_3d) - - x = self.linear_fc1(op, x) - x = op.Gelu(x) - x = self.linear_fc2(op, x) - return x - - class LinearMultiModalProjector(nn.Module): """Single linear projection (PaliGemma, Qwen2-Audio). diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 9ca9e2ce3..d421f525f 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -23,8 +23,11 @@ "CogVideoXTransformer3DModel", "CohereCausalLMModel", "ControlNetModel", + "Cosmos3AVAEAudioDecoderOnlyTokenizer", + "Cosmos3AVAEAudioTokenizer", "Cosmos3EdgeTextModel", "Cosmos3EdgeVLModel", + "Cosmos3OmniGeneratorModel", "Cosmos3OmniReasonerModel", "DeepSeekOCR2CausalLMModel", "DeepSeekV3CausalLMModel", @@ -77,6 +80,7 @@ "LayerNormCausalLMModel", "LongcatFlashCausalLMModel", "MPTCausalLMModel", + "MLPLatentDynamicsModel", "Mamba2CausalLMModel", "MambaCausalLMModel", "MiniMaxCausalLMModel", @@ -146,6 +150,7 @@ "remap_diffusers_unet_lora", "ViTModel", "VideoAutoencoderModel", + "AutoencoderKLWanModel", "Wav2Vec2ForCTCModel", "Wav2Vec2Model", "WhisperForConditionalGeneration", @@ -171,7 +176,12 @@ from mobius.models.cohere import CohereCausalLMModel from mobius.models.controlnet import ControlNetModel from mobius.models.cosmos import Cosmos3EdgeTextModel, Cosmos3EdgeVLModel +from mobius.models.cosmos3_audio import ( + Cosmos3AVAEAudioDecoderOnlyTokenizer, + Cosmos3AVAEAudioTokenizer, +) from mobius.models.cosmos3_omni import Cosmos3OmniReasonerModel +from mobius.models.cosmos3_omni_generator import Cosmos3OmniGeneratorModel from mobius.models.ctrl import CTRLCausalLMModel from mobius.models.deepseek import DeepSeekV3CausalLMModel from mobius.models.deepseek_ocr2 import DeepSeekOCR2CausalLMModel @@ -307,9 +317,10 @@ from mobius.models.vae import AutoencoderKLModel from mobius.models.video_vae import VideoAutoencoderModel from mobius.models.vit import ViTModel +from mobius.models.wan_vae import AutoencoderKLWanModel from mobius.models.wav2vec2 import Wav2Vec2Model from mobius.models.wav2vec2_ctc import Wav2Vec2ForCTCModel from mobius.models.whisper import WhisperForConditionalGeneration -from mobius.models.world_model import MLPWorldModel +from mobius.models.world_model import MLPLatentDynamicsModel, MLPWorldModel from mobius.models.xlm import XLMCausalLMModel from mobius.models.zamba2 import Zamba2CausalLMModel diff --git a/src/mobius/models/cosmos.py b/src/mobius/models/cosmos.py index 8b823fec9..32c71a7ec 100644 --- a/src/mobius/models/cosmos.py +++ b/src/mobius/models/cosmos.py @@ -5,8 +5,8 @@ Supports the ``cosmos3_edge`` checkpoint (``nvidia/Cosmos3-Edge``, ``Cosmos3EdgeForConditionalGeneration``) both as a full vision-language model -(:class:`Cosmos3EdgeVLModel`, 3-model onnxruntime-genai split) and as a -standalone text reasoner (:class:`Cosmos3EdgeTextModel`). +(:class:`Cosmos3EdgeVLModel`, 3-model split) and as a standalone text reasoner +(:class:`Cosmos3EdgeTextModel`). Cosmos3-Edge is a LLaVA-style VLM built from three towers: @@ -16,15 +16,27 @@ - **Non-gated feed-forward network** — ``down_proj(relu2(up_proj(x)))`` using a squared-ReLU activation (``hidden_act="relu2"``), rather than the GLU-style gated MLP used by Llama/Qwen. This maps onto :class:`FCMLP`. - - **3D multimodal RoPE** (``mrope_section=[24, 20, 20]``). For text-only - inference the three sections use the same positions, so it reduces to - standard 1D RoPE (same simplification used by the Qwen-VL text decoders). - -- **Vision encoder** — a SigLIP-style patch-embedding + transformer tower - (``model.visual.*``). + - **Interleaved 3D multimodal RoPE** (``mrope_section=[24, 20, 20]``). + Frequency channel ``i`` is driven by the height axis when + ``i % 3 == 1 and i < 3 * mrope_section[1]``, by the width axis when + ``i % 3 == 2 and i < 3 * mrope_section[2]``, and by the temporal axis + otherwise — *not* the contiguous chunking used by Qwen-VL. For text-only + inference the three axes carry the same positions, so it reduces to + standard 1D RoPE. + +- **Vision encoder** — a *variable-resolution* SigLIP2 tower + (:class:`Cosmos3EdgeVisionTower`, ``model.visual.*``) consuming + pre-patchified, packed pixel values plus a ``grid_thw`` triple. - **Merger projector** — a pixel-shuffle projector - (:class:`Cosmos3EdgeMultiModalProjector`, ``model.projector.*``) that merges - each 2x2 patch block and projects to the text hidden size. + (:class:`Cosmos3EdgePatchMerger`, ``model.projector.*``) that merges each + 2x2 patch block and projects to the text hidden size. + +Images and videos share the same tower, projector and token expansion: the +processor emits one ``<|vision_start|> ... <|vision_end|>`` span per image and +one span **per video frame**, with placeholder ids 19 (image) and 18 (video). +:class:`_Cosmos3EdgeEmbeddingModel` therefore scatters two independent feature +streams — ``image_features`` at id 19 and ``video_features`` at id 18 — exactly +mirroring ``Cosmos3EdgeModel.forward``'s two ``masked_scatter`` calls. The HuggingFace weights use ``self_attn.to_{q,k,v,out}`` projection names and place the text tower at the top level (``layers.*``, ``embed_tokens``, @@ -41,18 +53,21 @@ for both paths. .. note:: - NVIDIA does not publish modeling code for ``cosmos3_edge`` (it is not in - ``transformers`` and the repo ships no remote-code module), so exact - pixel-shuffle ordering and numerical parity are unverifiable. These builds - are validated at graph-construction (L1) confidence only. - -The ``cosmos3_omni`` variants (``nvidia/Cosmos3-Nano`` / ``-Super``) are -two-tower diffusion world models exported as diffusers pipelines and are -tracked separately. + The Reasoner architecture is reproduced from the published + ``transformers`` ``cosmos3_edge`` modeling code (cross-checked against + vLLM's ``cosmos3_edge.py``). The *Generator*/Action towers that share the + same checkpoint remain proprietary rectified-flow components and are not + reproduced here. + +The complete ``cosmos3_edge`` world-model package (this Reasoner plus the +shared MoT Generator, Wan VAE, and Action head) is composed by +``build_cosmos3_edge_world_model``. The ``cosmos3_omni`` variants use the same +Generator/VAE building blocks with a Qwen3-VL Reasoner. """ from __future__ import annotations +import re from typing import TYPE_CHECKING from onnxscript import OpBuilder, nn @@ -60,10 +75,10 @@ from mobius._configs import ArchitectureConfig from mobius.components import ( FCMLP, - Cosmos3EdgeMultiModalProjector, + Cosmos3EdgePatchMerger, + Cosmos3EdgeVisionTower, Embedding, Linear, - VisionModel, ) from mobius.models.base import CausalLMModel, TextModel @@ -72,6 +87,30 @@ import torch +_DROPPED_UNIFIED_KEY_RE = re.compile( + "|".join( + ( + r"\.add_q_proj\.", + r"\.add_k_proj\.", + r"\.add_v_proj\.", + r"\.to_add_out\.", + r"\.norm_added_q\.", + r"\.norm_added_k\.", + r"moe_gen", + r"^proj_out\.", + r"^proj_in\.", + r"^time_embedder\.", + r"^audio_proj_out\.", + r"^audio_proj_in\.", + r"^audio_modality_embed$", + r"^action_proj_out\.", + r"^action_proj_in\.", + r"^action_modality_embed$", + ) + ) +) + + def _rename_cosmos_text_key(key: str) -> str: """Rename Cosmos3-Edge attention and QK-norm keys to mobius conventions.""" return ( @@ -198,18 +237,30 @@ def preprocess_weights( class _Cosmos3EdgeVisionEncoderModel(nn.Module): - """Cosmos3-Edge vision encoder: SigLIP tower + pixel-shuffle projector. + """Cosmos3-Edge vision encoder: SigLIP2 tower + pixel-shuffle projector. + + ``pixel_values [total_patches, patch_dim]`` + ``grid_thw [3]`` → + :class:`Cosmos3EdgeVisionTower` → ``[total_patches, vision_hidden]`` → + :class:`Cosmos3EdgePatchMerger` → ``[total_patches / merge², text_hidden]``. - ``pixel_values [B, 3, H, W]`` → SigLIP encoder → ``[B, num_patches, D]`` - → :class:`Cosmos3EdgeMultiModalProjector` → ``[B, num_merged, text_hidden]``. + One call handles exactly one packed visual item: a single image + (``grid_t == 1``) or all frames of one video (``grid_t == num_frames``), + matching ``Cosmos3EdgeModel.get_image_features`` / + ``get_video_features`` (which share the same code path). + + Sub-module attribute names (``visual`` / ``projector``) mirror the + HuggingFace layout so weight mapping is a plain ``model.`` prefix strip. """ def __init__(self, config: ArchitectureConfig): super().__init__() vc = config.vision assert vc is not None, "Cosmos3-Edge requires a VisionConfig" - assert vc.image_size is not None and vc.patch_size is not None assert vc.hidden_size is not None + assert vc.intermediate_size is not None + assert vc.num_hidden_layers is not None + assert vc.num_attention_heads is not None + assert vc.patch_size is not None assert vc.projector_intermediate_size is not None, ( "Cosmos3-Edge projector requires projector_intermediate_size" ) @@ -218,52 +269,81 @@ def __init__(self, config: ArchitectureConfig): "Cosmos3-Edge projector output must match the text hidden size, " f"got {vc.out_hidden_size} != {config.hidden_size}" ) - if vc.image_size % vc.patch_size != 0: - raise ValueError( - f"image_size ({vc.image_size}) must be divisible by " - f"patch_size ({vc.patch_size})" - ) - self.vision_tower = VisionModel(config) - # Fixed square patch grid (image_size // patch_size), e.g. 256 -> 16. - grid_size = vc.image_size // vc.patch_size - self.multi_modal_projector = Cosmos3EdgeMultiModalProjector( + # ``num_patches`` sizes the learned position grid; fall back to the + # nominal image_size//patch_size square when the config omits it. + num_patches = vc.num_patches + if num_patches is None: + assert vc.image_size is not None + if vc.image_size % vc.patch_size != 0: + raise ValueError( + f"image_size ({vc.image_size}) must be divisible by " + f"patch_size ({vc.patch_size})" + ) + num_patches = (vc.image_size // vc.patch_size) ** 2 + merge = vc.spatial_merge_size or 2 + self.patch_size = vc.patch_size + self.temporal_patch_size = vc.temporal_patch_size or 1 + self.num_channels = vc.in_channels + self.spatial_merge_size = merge + self.patch_dim = ( + self.patch_size * self.patch_size * self.num_channels * self.temporal_patch_size + ) + self.visual = Cosmos3EdgeVisionTower( + hidden_size=vc.hidden_size, + intermediate_size=vc.intermediate_size, + num_hidden_layers=vc.num_hidden_layers, + num_attention_heads=vc.num_attention_heads, + patch_size=self.patch_size, + num_channels=self.num_channels, + num_patches=num_patches, + norm_eps=vc.norm_eps, + temporal_patch_size=self.temporal_patch_size, + spatial_merge_size=merge, + ) + self.projector = Cosmos3EdgePatchMerger( vision_hidden_size=vc.hidden_size, text_hidden_size=config.hidden_size, intermediate_size=vc.projector_intermediate_size, - grid_size=grid_size, - spatial_merge_size=vc.spatial_merge_size or 2, + spatial_merge_size=merge, + use_postshuffle_norm=vc.use_postshuffle_norm, norm_eps=vc.norm_eps, ) - def forward(self, op: OpBuilder, pixel_values: ir.Value): - vision_features = self.vision_tower(op, pixel_values) - return self.multi_modal_projector(op, vision_features) + def forward(self, op: OpBuilder, pixel_values: ir.Value, grid_thw: ir.Value): + # grid_thw: [3] int64 = (frames, grid_h, grid_w) for this visual item. + grid_t = op.Gather(grid_thw, op.Constant(value_int=0)) + grid_h = op.Gather(grid_thw, op.Constant(value_int=1)) + grid_w = op.Gather(grid_thw, op.Constant(value_int=2)) + vision_features = self.visual(op, pixel_values, grid_t, grid_h, grid_w) + return self.projector(op, vision_features) def preprocess_weights( self, state_dict: dict[str, torch.Tensor] ) -> dict[str, torch.Tensor]: - # state_dict here is the vision-routed slice: ``model.visual.*`` (SigLIP - # tower) and ``model.projector.*`` (merger projector). + # state_dict here is the vision-routed slice: ``model.visual.*`` (SigLIP2 + # tower) and ``model.projector.*`` (merger projector). The module tree + # mirrors those names, so only the ``model.`` prefix and the SigLIP + # ``mlp.fc1/fc2`` -> ``FCMLP.up_proj/down_proj`` naming differ. renamed: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): - if key.startswith("model.visual."): - new_key = "vision_tower.vision_model." + key[len("model.visual.") :] - # SigLIP MLP: fc1/fc2 -> up_proj/down_proj (FCMLP convention). - new_key = new_key.replace(".mlp.fc1.", ".mlp.up_proj.").replace( - ".mlp.fc2.", ".mlp.down_proj." - ) - renamed[new_key] = value - elif key.startswith("model.projector."): - new_key = "multi_modal_projector." + key[len("model.projector.") :] - renamed[new_key] = value + if not key.startswith(("model.visual.", "model.projector.")): + continue + new_key = key[len("model.") :] + new_key = new_key.replace(".mlp.fc1.", ".mlp.up_proj.").replace( + ".mlp.fc2.", ".mlp.down_proj." + ) + renamed[new_key] = value return renamed class _Cosmos3EdgeEmbeddingModel(nn.Module): - """Cosmos3-Edge embedding: token lookup + image feature fusion. + """Cosmos3-Edge embedding: token lookup + image/video feature fusion. Scatters projected vision features into the text embedding sequence at - ``image_token_id`` (19) positions, matching the LLaVA embedding contract. + ``image_token_id`` (19) and ``video_token_id`` (18) positions. The two + streams are independent, mirroring ``Cosmos3EdgeModel.forward``, which + runs one ``masked_scatter`` per modality. Either stream may be empty + (zero rows) for text-only, image-only or video-only prompts. """ def __init__(self, config: ArchitectureConfig): @@ -273,32 +353,54 @@ def __init__(self, config: ArchitectureConfig): config.vocab_size, config.hidden_size, config.pad_token_id ) self.image_token_id = config.image_token_id or 0 + vision = config.vision + self.video_token_id = ( + vision.video_token_id if vision is not None and vision.video_token_id else None + ) - def forward(self, op: OpBuilder, input_ids: ir.Value, image_features: ir.Value): - text_embeds = self.embed_tokens(op, input_ids) - - image_mask = op.Equal(input_ids, op.Constant(value_int=self.image_token_id)) - image_mask_3d = op.Unsqueeze(image_mask, [-1]) - - mask_int = op.Cast(image_mask, to=7) - cumsum = op.CumSum(mask_int, 1) - indices = op.Sub(cumsum, op.Constant(value_int=1)) - indices = op.Clip(indices, op.Constant(value_int=0)) - - # Pad image_features with one zero row so Gather stays in-bounds for - # text-only input (num_image_tokens == 0); the mask discards it. + def _scatter( + self, + op: OpBuilder, + embeds: ir.Value, + input_ids: ir.Value, + features: ir.Value, + token_id: int, + ) -> ir.Value: + """Replace ``token_id`` positions with rows of ``features`` in order.""" + mask = op.Equal(input_ids, op.Constant(value_int=token_id)) # [B, S] + mask_3d = op.Unsqueeze(mask, [-1]) + + # Running index of the placeholder within the sequence. + cumsum = op.CumSum(op.Cast(mask, to=7), 1) + indices = op.Clip(op.Sub(cumsum, op.Constant(value_int=1)), op.Constant(value_int=0)) + + # Pad the feature table with one zero row so Gather stays in-bounds + # when the modality is absent (num_feature_tokens == 0); the mask + # discards the padded row. pad_row = op.Expand( - op.CastLike(0.0, image_features), + op.CastLike(0.0, features), op.Concat( op.Constant(value_ints=[1]), - op.Shape(image_features, start=1, end=2), + op.Shape(features, start=1, end=2), axis=0, ), ) - padded_features = op.Concat(image_features, pad_row, axis=0) + padded = op.Concat(features, pad_row, axis=0) + gathered = op.Gather(padded, indices, axis=0) + return op.Where(mask_3d, gathered, embeds) - gathered = op.Gather(padded_features, indices, axis=0) - return op.Where(image_mask_3d, gathered, text_embeds) + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + image_features: ir.Value, + video_features: ir.Value | None = None, + ): + embeds = self.embed_tokens(op, input_ids) + embeds = self._scatter(op, embeds, input_ids, image_features, self.image_token_id) + if video_features is not None and self.video_token_id is not None: + embeds = self._scatter(op, embeds, input_ids, video_features, self.video_token_id) + return embeds def preprocess_weights( self, state_dict: dict[str, torch.Tensor] @@ -313,17 +415,21 @@ class Cosmos3EdgeVLModel(nn.Module): ``model_type: cosmos3_edge`` / ``Cosmos3EdgeForConditionalGeneration``. - Builds three ONNX models for onnxruntime-genai deployment: + Builds three ONNX models: - - **decoder**: squared-ReLU GQA text reasoner taking ``inputs_embeds``. - - **vision_encoder**: SigLIP vision tower + pixel-shuffle merger projector. - - **embedding**: token embedding + image feature fusion at - ``image_token_id`` (19). + - **decoder**: squared-ReLU GQA text reasoner taking ``inputs_embeds`` and + interleaved 3D M-RoPE ``position_ids [3, batch, seq]``. + - **vision_encoder**: variable-resolution SigLIP2 tower + pixel-shuffle + merger projector, driven by packed ``pixel_values [total_patches, + patch_dim]`` and ``grid_thw [3]``. The *same* graph serves images + (``grid_t == 1``) and videos (``grid_t == num_frames``). + - **embedding**: token embedding + feature fusion at ``image_token_id`` + (19) and ``video_token_id`` (18). HuggingFace weight layout (single checkpoint, no ``language_model.`` prefix): - - ``model.visual.*`` → vision tower (SigLIP) + - ``model.visual.*`` → vision tower (SigLIP2) - ``model.projector.*`` → merger projector - ``embed_tokens.weight`` → embedding - ``layers.*`` / ``norm.weight`` / ``lm_head.weight`` → decoder @@ -331,10 +437,10 @@ class Cosmos3EdgeVLModel(nn.Module): module docstring) .. note:: - NVIDIA does not publish modeling code for ``cosmos3_edge`` (it is not in - ``transformers`` and the repo ships no remote-code module), so the exact - pixel-shuffle ordering and numerical parity are unverifiable. This build - is validated at graph-construction (L1) confidence only. + The Reasoner is validated numerically against a PyTorch transcription + of the published ``transformers`` ``cosmos3_edge`` modeling code (see + ``tests/_cosmos3_edge_reference.py``). The Generator/Action towers in + the same checkpoint remain proprietary and are not reproduced. """ default_task: str = "cosmos3-edge-vl" @@ -371,6 +477,8 @@ def preprocess_weights( # Generator-tower key-norm is not used by the reasoner (see docstring). if "k_norm_und_for_gen" in key: continue + if _DROPPED_UNIFIED_KEY_RE.search(key): + continue if key.startswith(("model.visual.", "model.projector.")): vision_sd[key] = value elif "embed_tokens" in key: diff --git a/src/mobius/models/cosmos3_audio.py b/src/mobius/models/cosmos3_audio.py new file mode 100644 index 000000000..324145432 --- /dev/null +++ b/src/mobius/models/cosmos3_audio.py @@ -0,0 +1,899 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Cosmos3 AVAE audio tokenizer — direct ONNX graph construction. + +Replicates ``diffusers.models.autoencoders.autoencoder_cosmos3_audio`` +(``Cosmos3AVAEAudioTokenizer``), the ``sound_tokenizer`` component of +``nvidia/Cosmos3-Nano`` / ``nvidia/Cosmos3-Super``. + +Pipeline:: + + waveform [B, C, N] + -> (optional) peak volume normalization + -> (optional) right zero-pad to a multiple of hop_size + -> STFT front-end [B, (n_fft+2)*C, N/hop_length] + -> encoder.layers (1x1 conv -> {ConvNeXt x k, strided conv} x S -> 1x1 conv) + -> moments [B, 2*z, T] + -> bottleneck (VAE): mean, scale = split(moments); std = softplus(scale) + 1e-4 + + latents [B, z, T] + -> decoder.conv1 -> decoder.block.{i} (Snake -> ConvTranspose -> 3x ResidualUnit) + -> decoder.snake1 -> decoder.conv2 -> clamp(-1, 1) + -> waveform [B, 2, T * hop_size] + +Module attribute names mirror the upstream ``nn.Module`` tree exactly +(``encoder.layers.3.weight``, ``decoder.block.0.res_unit2.conv1.weight``, ...) +so ONNX initializer names line up with the HuggingFace checkpoint after the +weight-norm fold performed by :func:`fold_weight_norm`. + +Published checkpoints (all FP32, all using legacy ``weight_g``/``weight_v``):: + + nvidia/Cosmos3-Nano 249 tensors = 67 encoder + 182 decoder + nvidia/Cosmos3-Super 249 tensors = 67 encoder + 182 decoder + nvidia/Cosmos3-Super-Text2Image 182 tensors = 0 encoder + 182 decoder + +After folding (5 weight-norm pairs in the encoder, 37 in the decoder) that is +62 + 145 = 207 initializers for a full build and 145 for a decoder-only build. +Because all three ship the *same* config JSON, encoder presence must be read +from the weights — see :func:`state_dict_has_encoder` and the two tokenizer +classes below. + +Two structural differences from the PyTorch reference, both required by ONNX: + +* ``weight_norm`` is folded offline. PyTorch stores ``weight_g``/``weight_v`` + (or ``parametrizations.weight.original0/original1``); ONNX ``Conv`` takes a + single weight tensor, so :meth:`preprocess_weights` recombines them into + ``.weight``. Everything above the leaf name is unchanged, and no + ``g * v / ||v||`` math is left in the graph. +* Posterior sampling stays out of the graph. The encoder graph emits + ``moments``/``latent_mean``/``latent_std`` and the caller draws the sample, + keeping the ONNX model deterministic. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import numpy as np +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius._configs._cosmos3_audio import Cosmos3AudioConfig, state_dict_has_encoder + +if TYPE_CHECKING: + import torch + +__all__ = [ + "Cosmos3AVAEAudioDecoderOnlyTokenizer", + "Cosmos3AVAEAudioTokenizer", + "Cosmos3AudioConvNeXtBlock", + "Cosmos3AudioDecoder", + "Cosmos3AudioDecoderBlock", + "Cosmos3AudioResidualUnit", + "Cosmos3AudioSpectrogramConvNeXtEncoder", + "Cosmos3AudioVAEBottleneck", + "create_cosmos3_avae_audio_tokenizer", + "fold_weight_norm", + "state_dict_has_encoder", +] + +# Matches ``OobleckDiagonalGaussianDistribution``: std = softplus(scale) + 1e-4. +_POSTERIOR_STD_EPS = 1e-4 +# Matches ``Snake1d``: x + (beta + 1e-9).reciprocal() * sin(alpha * x)^2. +_SNAKE_EPS = 1e-9 +# Matches ``Cosmos3AVAEAudioTokenizer.encode``: x / (|x|.max() + 1e-5) * 0.95. +_VOLUME_EPS = 1e-5 +_VOLUME_PEAK = 0.95 + + +def _scalar(op: OpBuilder, value: float, like: ir.Value) -> ir.Value: + """Build a scalar constant cast to the dtype of *like*.""" + return op.CastLike(op.Constant(value_float=float(value)), like) + + +# --------------------------------------------------------------------------- +# Primitive layers (parameter holders mirroring torch.nn leaf modules) +# --------------------------------------------------------------------------- + + +class _Conv1d(nn.Module): + """``nn.Conv1d`` with symmetric zero padding. + + Parameter names are ``weight`` / ``bias``, matching PyTorch after the + weight-norm fold (``weight_g``/``weight_v`` -> ``weight``). + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + *, + stride: int = 1, + padding: int = 0, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + dtype: ir.DataType = ir.DataType.FLOAT, + ): + super().__init__() + # Conv1d weight layout: (out_channels, in_channels / groups, kernel_size) + self.weight = nn.Parameter([out_channels, in_channels // groups, kernel_size], dtype) + self.bias = nn.Parameter([out_channels], dtype) if bias else None + self._kernel_size = kernel_size + self._stride = stride + self._padding = padding + self._dilation = dilation + self._groups = groups + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Convolve ``(B, C_in, L)`` into ``(B, C_out, L')``.""" + inputs = [x, self.weight] + ([self.bias] if self.bias is not None else []) + return op.Conv( + *inputs, + kernel_shape=[self._kernel_size], + strides=[self._stride], + pads=[self._padding, self._padding], + dilations=[self._dilation], + group=self._groups, + ) + + +class _ConvTranspose1d(nn.Module): + """``nn.ConvTranspose1d`` with symmetric padding and right output padding.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + *, + stride: int = 1, + padding: int = 0, + output_padding: int = 0, + bias: bool = True, + dtype: ir.DataType = ir.DataType.FLOAT, + ): + super().__init__() + # ConvTranspose1d weight layout: (in_channels, out_channels, kernel_size) + self.weight = nn.Parameter([in_channels, out_channels, kernel_size], dtype) + self.bias = nn.Parameter([out_channels], dtype) if bias else None + self._kernel_size = kernel_size + self._stride = stride + self._padding = padding + self._output_padding = output_padding + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Upsample ``(B, C_in, L)`` to ``(B, C_out, L * stride)``. + + With ``kernel_size = 2 * stride``, ``padding = ceil(stride / 2)`` and + ``output_padding = stride % 2`` the ONNX length formula + ``stride * (L - 1) + output_padding + kernel_size - 2 * padding`` + collapses to exactly ``L * stride`` for both even and odd strides. + """ + inputs = [x, self.weight] + ([self.bias] if self.bias is not None else []) + return op.ConvTranspose( + *inputs, + kernel_shape=[self._kernel_size], + strides=[self._stride], + pads=[self._padding, self._padding], + output_padding=[self._output_padding], + dilations=[1], + group=1, + ) + + +class _ConstantPad1d(nn.Module): + """Parameter-free ``nn.ConstantPad1d`` over the temporal axis. + + Kept as its own module so the enclosing ``nn.Sequential`` assigns index + ``0`` here and index ``1`` to the depthwise conv, reproducing the upstream + ``dwconv.1.weight`` parameter names. + """ + + def __init__(self, pad_left: int, pad_right: int): + super().__init__() + # ONNX Pad on rank-3 input: [b0, b1, b2, e0, e1, e2] + self._pads = [0, 0, pad_left, 0, 0, pad_right] + self._needs_pad = pad_left > 0 or pad_right > 0 + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Zero-pad ``(B, C, L)`` to ``(B, C, L + left + right)``.""" + if not self._needs_pad: + return x + return op.Pad(x, op.Constant(value_ints=self._pads)) + + +class _Snake1d(nn.Module): + """SnakeBeta activation ``x + (beta + 1e-9)^-1 * sin(alpha * x)^2``. + + Replicates ``Snake1d`` from the upstream module. ``alpha``/``beta`` keep the + checkpoint's ``(1, C, 1)`` shape so they broadcast over ``(B, C, T)`` + without a reshape, and are exponentiated when ``logscale`` is set. + """ + + def __init__( + self, + hidden_dim: int, + logscale: bool = True, + dtype: ir.DataType = ir.DataType.FLOAT, + ): + super().__init__() + self.alpha = nn.Parameter([1, hidden_dim, 1], dtype) + self.beta = nn.Parameter([1, hidden_dim, 1], dtype) + self._logscale = logscale + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Apply SnakeBeta elementwise to ``(B, C, T)``.""" + alpha = op.Exp(self.alpha) if self._logscale else self.alpha + beta = op.Exp(self.beta) if self._logscale else self.beta + # sin^2(alpha * x) — (B, C, T) broadcast against (1, C, 1) + sin_val = op.Sin(op.Mul(alpha, x)) + sin_sq = op.Mul(sin_val, sin_val) + inv_beta = op.Reciprocal(op.Add(beta, _scalar(op, _SNAKE_EPS, beta))) + return op.Add(x, op.Mul(inv_beta, sin_sq)) + + +class _GELU(nn.Module): + """Exact (erf) GELU, matching ``nn.GELU()`` with default approximation.""" + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Apply GELU elementwise.""" + return op.Gelu(x, approximate="none") + + +class _FP32LayerNorm(nn.Module): + """Weight-only LayerNorm evaluated in float32. + + Mirrors ``diffusers.models.normalization.FP32LayerNorm`` constructed with + ``bias=False``: the input and weight are upcast to float32, normalized over + the last axis, then cast back to the original dtype. For float32 graphs the + casts are elided. + """ + + def __init__( + self, + hidden_dim: int, + eps: float = 1e-5, + dtype: ir.DataType = ir.DataType.FLOAT, + ): + super().__init__() + self.weight = nn.Parameter([hidden_dim], dtype) + self._eps = eps + self._needs_cast = dtype != ir.DataType.FLOAT + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Normalize the last axis of ``(B, T, C)``.""" + if not self._needs_cast: + return op.LayerNormalization(x, self.weight, axis=-1, epsilon=self._eps) + x_fp32 = op.Cast(x, to=ir.DataType.FLOAT) + weight_fp32 = op.Cast(self.weight, to=ir.DataType.FLOAT) + normalized = op.LayerNormalization(x_fp32, weight_fp32, axis=-1, epsilon=self._eps) + return op.CastLike(normalized, x) + + +# --------------------------------------------------------------------------- +# Encoder: spectrogram ConvNeXt +# --------------------------------------------------------------------------- + + +class Cosmos3AudioConvNeXtBlock(nn.Module): + """1-D ConvNeXt block used by the Cosmos3 SpecConvNeXt encoder. + + ``residual + pwconv2(act(pwconv1(norm(dwconv(x)))))`` where ``dwconv`` is a + depthwise k=7 convolution wrapped in an explicit constant pad, and ``norm`` + is a channels-last FP32 LayerNorm. + + HF class: ``Cosmos3AudioConvNeXtBlock``. + """ + + def __init__( + self, + hidden_dim: int, + intermediate_dim: int, + *, + use_snake: bool = True, + snake_logscale: bool = True, + dtype: ir.DataType = ir.DataType.FLOAT, + ): + super().__init__() + # Non-causal variant: symmetric (3, 3) pad around a k=7 depthwise conv, + # so the temporal length is preserved. + self.dwconv = nn.Sequential( + _ConstantPad1d(3, 3), + _Conv1d(hidden_dim, hidden_dim, 7, groups=hidden_dim, dtype=dtype), + ) + self.norm = _FP32LayerNorm(hidden_dim, eps=1e-5, dtype=dtype) + self.pwconv1 = _Conv1d(hidden_dim, intermediate_dim, 1, dtype=dtype) + self.act: nn.Module = ( + _Snake1d(intermediate_dim, snake_logscale, dtype) if use_snake else _GELU() + ) + self.pwconv2 = _Conv1d(intermediate_dim, hidden_dim, 1, dtype=dtype) + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + """Transform ``(B, C, T)`` in place (channel count unchanged).""" + residual = hidden_states + # (B, C, T) -> (B, C, T) + hidden_states = self.dwconv(op, hidden_states) + # LayerNorm is channels-last upstream: (B, C, T) -> (B, T, C) -> (B, C, T) + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) + hidden_states = self.norm(op, hidden_states) + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) + # (B, C, T) -> (B, 4C, T) -> act -> (B, C, T) + hidden_states = self.pwconv1(op, hidden_states) + hidden_states = self.act(op, hidden_states) + hidden_states = self.pwconv2(op, hidden_states) + return op.Add(residual, hidden_states) + + +class Cosmos3AudioSpectrogramConvNeXtEncoder(nn.Module): + """Waveform → latent-moments encoder (STFT front-end + ConvNeXt stack). + + The ``layers`` attribute is a flat ``nn.Sequential`` whose indices match the + checkpoint exactly for ``enc_num_blocks=2`` and three stages:: + + layers.0 1x1 conv, (n_fft+2)*C -> c_mults[0] * enc_dim + layers.1, layers.2 ConvNeXt blocks + layers.3 strided conv, stage 0 -> stage 1 + layers.4, layers.5 ConvNeXt blocks + layers.6 strided conv, stage 1 -> stage 2 + layers.7, layers.8 ConvNeXt blocks + layers.9 strided conv, stage 2 -> stage 2 + layers.10 1x1 conv, c_mults[-1] * enc_dim -> enc_latent_dim + + HF class: ``Cosmos3AudioSpectrogramConvNeXtEncoder``. + + Inputs: ``audio`` ``(B, encoder_input_channels, N)``. + Outputs: ``(B, T, enc_latent_dim)`` with ``T = N / (prod(enc_strides) * enc_hop_length)`` + — channels-last, matching the upstream ``forward`` return. + """ + + def __init__(self, config: Cosmos3AudioConfig): + super().__init__() + config.validate() + self._config = config + dtype = config.dtype + channels = config.enc_dim + multiples = config.enc_c_mults + strides = config.enc_strides + + self.input_channels = config.encoder_input_channels + self.n_fft = config.enc_n_fft + self.hop_length = config.enc_hop_length + + layers: list[nn.Module] = [ + # Packed real/imag spectrogram bins -> first stage width. + _Conv1d( + config.spectrogram_channels, + multiples[0] * channels, + 1, + bias=False, + dtype=dtype, + ) + ] + for index, stride in enumerate(strides): + input_dim = multiples[index] * channels + # The last stage keeps its width (there is no multiples[index + 1]). + output_dim = ( + multiples[index + 1] * channels + if index < len(multiples) - 1 + else multiples[-1] * channels + ) + for _ in range(config.enc_num_blocks): + layers.append( + Cosmos3AudioConvNeXtBlock( + input_dim, + input_dim * 4, + use_snake=config.enc_use_snake, + snake_logscale=config.snake_logscale, + dtype=dtype, + ) + ) + # Downsample by `stride`: k = 2 * stride, pad = ceil(stride / 2). + layers.append( + _Conv1d( + input_dim, + output_dim, + 2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + dtype=dtype, + ) + ) + layers.append( + _Conv1d( + multiples[-1] * channels, config.enc_latent_dim, 1, bias=False, dtype=dtype + ) + ) + self.layers = nn.Sequential(*layers) + + # Periodic Hann window, identical to torch.hann_window(n_fft). + window = np.hanning(self.n_fft + 1)[:-1].astype(np.float32) + self._window = ir.tensor(window, name="stft_hann_window") + + def spectrogram(self, op: OpBuilder, audio: ir.Value) -> ir.Value: + """Compute the packed real/imaginary STFT front-end. + + Replicates ``_spectrogram`` + the channel packing in the upstream + ``forward``. The transform is fixed (no learnable parameters) but it + lives inside ``encoder.forward`` upstream, so it is kept inside the + ONNX graph to preserve the documented ``waveform -> moments`` contract. + + .. note:: + ``STFT`` has no CUDA kernel in onnxruntime; on the CUDA EP this + node falls back to CPU and introduces a host/device copy at the + graph entry. Callers that want a pure-GPU encoder should feed the + spectrogram in directly instead. + + Args: + audio: ``(B, C, N)`` waveform. + + Returns: + ``(B, C * (n_fft + 2), N / hop_length)`` spectrogram, real bins of + each channel followed by that channel's imaginary bins. + """ + config = self._config + channels = self.input_channels + + # (B, C, N) -> (B * C, N): the STFT is applied per waveform channel. + num_samples = op.Shape(audio, start=2, end=3) + flat = audio + if channels > 1: + flat_shape = op.Concat(op.Constant(value_ints=[-1]), num_samples, axis=0) + flat = op.Reshape(flat, flat_shape) + else: + flat = op.Squeeze(flat, op.Constant(value_ints=[1])) + + # torch.stft(center=False) after an explicit symmetric pad. + padded = op.Pad( + flat, + op.Constant(value_ints=[0, config.stft_pad_left, 0, config.stft_pad_right]), + ) + # The reference computes the transform in float32 regardless of dtype. + padded = op.Cast(padded, to=ir.DataType.FLOAT) + + # ONNX STFT wants a trailing "real signal" axis: (B*C, N', 1). + signal = op.Unsqueeze(padded, op.Constant(value_ints=[-1])) + spec = op.STFT( + signal, + op.Constant(value=ir.tensor(np.array(self.hop_length, dtype=np.int64))), + op.Constant(value=self._window), + op.Constant(value=ir.tensor(np.array(self.n_fft, dtype=np.int64))), + onesided=1, + ) + # (B*C, frames, bins, 2) -> (B*C, 2, bins, frames) so that a flatten of + # axes 1-2 yields [real bins ..., imaginary bins ...] like + # torch.cat([real, imaginary], dim=1). + spec = op.Transpose(spec, perm=[0, 3, 2, 1]) + num_frames = op.Shape(spec, start=3, end=4) + packed_shape = op.Concat( + op.Constant(value_ints=[-1, config.enc_n_fft + 2]), num_frames, axis=0 + ) + spec = op.Reshape(spec, packed_shape) + + spec = op.CastLike(spec, audio) + if channels > 1: + # (B*C, n_fft+2, frames) -> (B, C*(n_fft+2), frames) + merged_shape = op.Concat( + op.Constant(value_ints=[-1, config.spectrogram_channels]), num_frames, axis=0 + ) + spec = op.Reshape(spec, merged_shape) + return spec + + def forward(self, op: OpBuilder, audio: ir.Value) -> ir.Value: + """Encode a waveform into channels-last latent moments. + + Args: + audio: ``(B, encoder_input_channels, N)`` waveform. + + Returns: + ``(B, T, enc_latent_dim)`` moments. + """ + spec = self.spectrogram(op, audio) + # (B, spectrogram_channels, T) -> (B, enc_latent_dim, T) + hidden_states = self.layers(op, spec) + # Upstream returns channels-last; the tokenizer transposes it back. + return op.Transpose(hidden_states, perm=[0, 2, 1]) + + +# --------------------------------------------------------------------------- +# Decoder: Oobleck +# --------------------------------------------------------------------------- + + +class Cosmos3AudioResidualUnit(nn.Module): + """Oobleck residual unit: Snake → dilated k=7 conv → Snake → k=1 conv. + + HF class: ``Cosmos3AudioResidualUnit``. + """ + + def __init__( + self, + dimension: int = 16, + dilation: int = 1, + *, + snake_logscale: bool = True, + dtype: ir.DataType = ir.DataType.FLOAT, + ): + super().__init__() + # pad = ((7 - 1) * dilation) // 2 exactly preserves the temporal length, + # so the upstream centre-crop of the skip branch is always a no-op. + pad = ((7 - 1) * dilation) // 2 + self.snake1 = _Snake1d(dimension, snake_logscale, dtype) + self.conv1 = _Conv1d( + dimension, dimension, 7, dilation=dilation, padding=pad, dtype=dtype + ) + self.snake2 = _Snake1d(dimension, snake_logscale, dtype) + self.conv2 = _Conv1d(dimension, dimension, 1, dtype=dtype) + + def forward(self, op: OpBuilder, hidden_state: ir.Value) -> ir.Value: + """Apply the residual unit to ``(B, C, T)``; shape is preserved.""" + output = self.conv1(op, self.snake1(op, hidden_state)) + output = self.conv2(op, self.snake2(op, output)) + return op.Add(hidden_state, output) + + +class Cosmos3AudioDecoderBlock(nn.Module): + """Oobleck decoder block: Snake → transposed conv → 3 dilated residual units. + + HF class: ``Cosmos3AudioDecoderBlock``. + """ + + def __init__( + self, + input_dim: int, + output_dim: int, + stride: int = 1, + output_padding: int = 0, + *, + snake_logscale: bool = True, + dtype: ir.DataType = ir.DataType.FLOAT, + ): + super().__init__() + self.snake1 = _Snake1d(input_dim, snake_logscale, dtype) + self.conv_t1 = _ConvTranspose1d( + input_dim, + output_dim, + 2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + output_padding=output_padding, + dtype=dtype, + ) + self.res_unit1 = Cosmos3AudioResidualUnit( + output_dim, dilation=1, snake_logscale=snake_logscale, dtype=dtype + ) + self.res_unit2 = Cosmos3AudioResidualUnit( + output_dim, dilation=3, snake_logscale=snake_logscale, dtype=dtype + ) + self.res_unit3 = Cosmos3AudioResidualUnit( + output_dim, dilation=9, snake_logscale=snake_logscale, dtype=dtype + ) + + def forward(self, op: OpBuilder, hidden_state: ir.Value) -> ir.Value: + """Upsample ``(B, C_in, T)`` to ``(B, C_out, T * stride)``.""" + hidden_state = self.snake1(op, hidden_state) + hidden_state = self.conv_t1(op, hidden_state) + hidden_state = self.res_unit1(op, hidden_state) + hidden_state = self.res_unit2(op, hidden_state) + return self.res_unit3(op, hidden_state) + + +class Cosmos3AudioDecoder(nn.Module): + """Oobleck-style latent → waveform decoder. + + Layout:: + + conv1 k=7 conv, vocoder_input_dim -> dec_dim * ([1] + dec_c_mults)[-1] + block.{i} Cosmos3AudioDecoderBlock, stride = reversed(dec_strides)[i] + snake1 SnakeBeta at dec_dim + conv2 k=7 conv, dec_dim -> dec_out_channels (no bias) + + HF class: ``Cosmos3AudioDecoder``. + + Inputs: ``(B, vocoder_input_dim, T)``. + Outputs: ``(B, dec_out_channels, T * prod(dec_strides))``. + """ + + def __init__(self, config: Cosmos3AudioConfig): + super().__init__() + config.validate() + self._config = config + dtype = config.dtype + channels = config.dec_dim + strides = config.decoder_upsampling_ratios + multiples = config.decoder_channel_multiples + + self.conv1 = _Conv1d( + config.vocoder_input_dim, channels * multiples[-1], 7, padding=3, dtype=dtype + ) + + # Walk the multiplier table from widest to narrowest as we upsample. + blocks: list[nn.Module] = [] + for stride_index, stride in enumerate(strides): + blocks.append( + Cosmos3AudioDecoderBlock( + input_dim=channels * multiples[len(strides) - stride_index], + output_dim=channels * multiples[len(strides) - stride_index - 1], + stride=stride, + output_padding=stride % 2, + snake_logscale=config.snake_logscale, + dtype=dtype, + ) + ) + self.block = nn.ModuleList(blocks) + + self.snake1 = _Snake1d(channels, config.snake_logscale, dtype) + self.conv2 = _Conv1d( + channels, config.dec_out_channels, 7, padding=3, bias=False, dtype=dtype + ) + + def forward(self, op: OpBuilder, hidden_state: ir.Value) -> ir.Value: + """Decode ``(B, z, T)`` latents into ``(B, audio_channels, T * hop)``.""" + # (B, z, T) -> (B, dec_dim * mults[-1], T) + hidden_state = self.conv1(op, hidden_state) + for layer in self.block: + hidden_state = layer(op, hidden_state) + # (B, dec_dim, T * hop) -> (B, audio_channels, T * hop) + hidden_state = self.snake1(op, hidden_state) + return self.conv2(op, hidden_state) + + +# --------------------------------------------------------------------------- +# VAE bottleneck +# --------------------------------------------------------------------------- + + +class Cosmos3AudioVAEBottleneck(nn.Module): + """Parameter-free diagonal-Gaussian bottleneck (``bottleneck_type="vae"``). + + Replicates ``OobleckDiagonalGaussianDistribution``: the encoder moments are + split in half along the channel axis into ``mean`` and ``scale``, and the + standard deviation is ``softplus(scale) + 1e-4``. + + Sampling (``mean + std * eps``) is deliberately *not* emitted — the ONNX + graph stays deterministic and the caller draws ``eps`` itself. ``mode()`` + is simply ``mean``. + """ + + def __init__(self, latent_channels: int): + super().__init__() + self._latent_channels = latent_channels + + def forward( + self, + op: OpBuilder, + moments: ir.Value, + ) -> tuple[ir.Value, ir.Value]: + """Split moments into ``(mean, std)``. + + Args: + moments: ``(B, 2 * z, T)`` encoder output. + + Returns: + ``(mean, std)``, each ``(B, z, T)``. + """ + split = op.Constant(value_ints=[self._latent_channels, self._latent_channels]) + mean, scale = op.Split(moments, split, axis=1, _outputs=2) + # std = softplus(scale) + 1e-4 + std = op.Add(op.Softplus(scale), _scalar(op, _POSTERIOR_STD_EPS, scale)) + return mean, std + + +# --------------------------------------------------------------------------- +# Weight-norm folding +# --------------------------------------------------------------------------- + + +def fold_weight_norm(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Recombine PyTorch ``weight_norm`` pairs into a single dense weight. + + ``torch.nn.utils.weight_norm(module, dim=0)`` stores a magnitude ``g`` and a + direction ``v``; the effective weight is ``g * v / ||v||`` where the norm is + taken over every axis except axis 0. ONNX ``Conv``/``ConvTranspose`` accept + only the combined tensor, so the fold happens here, at weight-load time. + + Both spellings are handled: + + * legacy — ``.weight_g`` / ``.weight_v`` (what the published + Cosmos3 ``sound_tokenizer`` checkpoint uses); + * ``torch.nn.utils.parametrizations.weight_norm`` — + ``.parametrizations.weight.original0`` / ``...original1``. + + Keys that are already dense are passed through untouched. + + Args: + state_dict: Raw checkpoint tensors. + + Returns: + A new dict with ``.weight`` in place of each ``g``/``v`` pair. + + Raises: + ValueError: If a magnitude tensor has no matching direction tensor. + """ + suffix_pairs = ( + (".weight_g", ".weight_v", ".weight"), + ( + ".parametrizations.weight.original0", + ".parametrizations.weight.original1", + ".weight", + ), + ) + folded: dict[str, torch.Tensor] = {} + + for key, value in state_dict.items(): + matched = False + for g_suffix, v_suffix, target_suffix in suffix_pairs: + if not key.endswith(g_suffix): + continue + prefix = key[: -len(g_suffix)] + v_key = prefix + v_suffix + if v_key not in state_dict: + raise ValueError( + f"Weight-norm magnitude {key!r} has no matching direction tensor {v_key!r}." + ) + direction = state_dict[v_key] + # norm over every axis except axis 0, keeping (out, 1, 1) broadcast shape + reduce_dims = tuple(range(1, direction.dim())) + norm = direction.float().pow(2).sum(dim=reduce_dims, keepdim=True).sqrt() + weight = value.float() * direction.float() / norm + folded[prefix + target_suffix] = weight.to(direction.dtype) + matched = True + break + if matched: + continue + # Direction tensors are consumed alongside their magnitude tensor. + if any(key.endswith(v_suffix) for _, v_suffix, _ in suffix_pairs): + continue + folded[key] = value + + return folded + + +# --------------------------------------------------------------------------- +# Tokenizers +# --------------------------------------------------------------------------- + + +class Cosmos3AVAEAudioDecoderOnlyTokenizer(nn.Module): + """Decoder-only Cosmos3 AVAE sound tokenizer. + + Use this path when the checkpoint ships without ``encoder.*`` weights (the + common case for sound *generation*). No encoder or bottleneck sub-module is + created, so the exported graph can never contain an initializer that has no + weight to fill it. + + HF class: ``Cosmos3AVAEAudioTokenizer`` with ``encoder_enabled=False``. + """ + + #: Whether :meth:`encode` is available on this module. + encoder_available: bool = False + + def __init__(self, config: Cosmos3AudioConfig): + super().__init__() + config.validate() + self.config = config + self.decoder = Cosmos3AudioDecoder(config) + + def decode(self, op: OpBuilder, latents: ir.Value) -> ir.Value: + """Decode sound latents into a waveform. + + Args: + latents: ``(B, vocoder_input_dim, T)`` diffusion-model latents. + + Returns: + ``(B, dec_out_channels, T * hop_size)`` waveform clamped to + ``[-1, 1]``. + """ + audio = self.decoder(op, latents) + return op.Clip(audio, _scalar(op, -1.0, audio), _scalar(op, 1.0, audio)) + + def preprocess_weights( + self, + state_dict: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Fold ``weight_norm`` pairs; module paths are already HF-aligned.""" + return fold_weight_norm(state_dict) + + +class Cosmos3AVAEAudioTokenizer(Cosmos3AVAEAudioDecoderOnlyTokenizer): + """Full encoder + VAE bottleneck + decoder Cosmos3 AVAE sound tokenizer. + + HF class: ``Cosmos3AVAEAudioTokenizer`` with ``encoder_enabled=True``. + """ + + encoder_available: bool = True + + def __init__(self, config: Cosmos3AudioConfig): + super().__init__(config) + self.encoder = Cosmos3AudioSpectrogramConvNeXtEncoder(config) + self.bottleneck = Cosmos3AudioVAEBottleneck(config.latent_channels) + + def normalize_volume(self, op: OpBuilder, waveform: ir.Value) -> ir.Value: + """Peak-normalize a waveform, matching ``encode``'s pre-processing. + + ``x / (|x|.max() + 1e-5) * 0.95`` where the maximum is taken over the + **entire tensor** (batch included), exactly as upstream. Callers that + need per-sample normalization must batch size 1. + + Args: + waveform: ``(B, C, N)`` waveform. + + Returns: + The peak-normalized waveform, same shape. + """ + peak = op.ReduceMax(op.Abs(waveform), keepdims=0) + scaled = op.Div(waveform, op.Add(peak, _scalar(op, _VOLUME_EPS, waveform))) + return op.Mul(scaled, _scalar(op, _VOLUME_PEAK, waveform)) + + def pad_to_hop_size(self, op: OpBuilder, waveform: ir.Value) -> ir.Value: + """Right zero-pad a waveform to a whole multiple of ``hop_size``. + + Mirrors the inference-mode padding in ``encode`` so that the number of + latent frames is exactly ``ceil(N / hop_size)``. + + Args: + waveform: ``(B, C, N)`` waveform. + + Returns: + ``(B, C, ceil(N / hop) * hop)`` waveform. + """ + hop = op.Constant(value=ir.tensor(np.array([self.config.resolved_hop_size], np.int64))) + num_samples = op.Shape(waveform, start=2, end=3) + # padding = (hop - (N % hop)) % hop + remainder = op.Mod(num_samples, hop) + padding = op.Mod(op.Sub(hop, remainder), hop) + pads = op.Concat(op.Constant(value_ints=[0, 0, 0, 0, 0]), padding, axis=0) + return op.Pad(waveform, pads) + + def encode( + self, + op: OpBuilder, + sample: ir.Value, + ) -> tuple[ir.Value, ir.Value, ir.Value]: + """Encode a waveform into deterministic posterior moments. + + Reproduces ``Cosmos3AVAEAudioTokenizer.encode`` in evaluation mode + (``force_pad`` semantics): optional peak normalization, optional + right-padding to ``hop_size``, the STFT ConvNeXt encoder, then the VAE + bottleneck split. No sample is drawn. + + Args: + sample: ``(B, encoder_input_channels, N)`` waveform. + + Returns: + ``(moments, mean, std)`` where ``moments`` is ``(B, 2 * z, T)`` and + ``mean``/``std`` are ``(B, z, T)``. + """ + hidden_states = sample + if self.config.normalize_volume: + hidden_states = self.normalize_volume(op, hidden_states) + hidden_states = self.pad_to_hop_size(op, hidden_states) + + # Upstream `_encode` transposes the channels-last encoder output back to + # channels-first before handing it to the Gaussian distribution. + encoded = self.encoder(op, hidden_states) + moments = op.Transpose(encoded, perm=[0, 2, 1]) + mean, std = self.bottleneck(op, moments) + return moments, mean, std + + +def create_cosmos3_avae_audio_tokenizer( + config: Cosmos3AudioConfig, +) -> Cosmos3AVAEAudioDecoderOnlyTokenizer: + """Instantiate the tokenizer variant that matches ``config.encoder_enabled``. + + Args: + config: A validated :class:`Cosmos3AudioConfig`. Use + :meth:`Cosmos3AudioConfig.with_encoder_from_state_dict` first when + encoder presence must be derived from the checkpoint. + + Returns: + :class:`Cosmos3AVAEAudioTokenizer` when the encoder is enabled, + otherwise :class:`Cosmos3AVAEAudioDecoderOnlyTokenizer`. + """ + if config.encoder_enabled: + return Cosmos3AVAEAudioTokenizer(config) + return Cosmos3AVAEAudioDecoderOnlyTokenizer(config) diff --git a/src/mobius/models/cosmos3_audio_test.py b/src/mobius/models/cosmos3_audio_test.py new file mode 100644 index 000000000..998959f33 --- /dev/null +++ b/src/mobius/models/cosmos3_audio_test.py @@ -0,0 +1,1149 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the Cosmos3 AVAE audio tokenizer (``sound_tokenizer``). + +Coverage: + +* parsing/validating the real published ``sound_tokenizer/config.json``; +* every ``NotImplementedError``/``ValueError`` guard in + :meth:`Cosmos3AudioConfig.validate`; +* the decoder-only graph (no encoder initializers at all) and the full + encoder+decoder package, including exact I/O contracts and initializer names; +* encoder-weight detection from checkpoint key names; +* ``weight_norm`` folding against ``torch.nn.utils.weight_norm``; +* numerical parity with a PyTorch reference transcribed from + ``diffusers.models.autoencoders.autoencoder_cosmos3_audio`` for SnakeBeta, a + residual unit, the whole tiny decoder, and the whole tiny STFT encoder. +""" + +from __future__ import annotations + +import math + +import numpy as np +import onnx_ir as ir +import pytest +import torch +from onnxscript import GraphBuilder, nn +from torch import nn as tnn + +from mobius._configs._cosmos3_audio import Cosmos3AudioConfig, state_dict_has_encoder +from mobius._constants import OPSET_VERSION +from mobius._model_package import ModelPackage +from mobius._testing.ort_inference import OnnxModelSession +from mobius.models.cosmos3_audio import ( + Cosmos3AudioConvNeXtBlock, + Cosmos3AudioResidualUnit, + Cosmos3AVAEAudioDecoderOnlyTokenizer, + Cosmos3AVAEAudioTokenizer, + _Snake1d, + create_cosmos3_avae_audio_tokenizer, + fold_weight_norm, +) +from mobius.tasks._cosmos3_audio import ( + Cosmos3AVAEAudioDecoderTask, + Cosmos3AVAEAudioTokenizerTask, + select_cosmos3_audio_task, +) + +# The published nvidia/Cosmos3-Nano and nvidia/Cosmos3-Super +# ``sound_tokenizer/config.json`` (byte-identical between the two releases). +_REAL_CONFIG: dict = { + "model_type": "autoencoder_v2", + "sampling_rate": 48000, + "stereo": True, + "use_wav_as_input": True, + "normalize_volume": True, + "hop_size": 1920, + "input_channels": 1, + "enc_type": "spec_convnext", + "enc_dim": 192, + "enc_intermediate_dim": 768, + "enc_num_layers": 12, + "enc_num_blocks": 2, + "enc_n_fft": 64, + "enc_hop_length": 16, + "enc_latent_dim": 128, + "enc_c_mults": [1, 2, 4], + "enc_strides": [4, 5, 6], + "enc_identity_init": False, + "enc_use_snake": True, + "dec_type": "oobleck", + "dec_dim": 320, + "dec_c_mults": [1, 2, 4, 8, 16], + "dec_strides": [2, 4, 5, 6, 8], + "dec_use_snake": True, + "dec_final_tanh": False, + "dec_out_channels": 2, + "dec_anti_aliasing": False, + "dec_use_nearest_upsample": False, + "dec_use_tanh_at_final": False, + "bottleneck_type": "vae", + "bottleneck": {"type": "vae"}, + "activation": "snakebeta", + "snake_logscale": True, + "anti_aliasing": False, + "use_cuda_kernel": False, + "causal": False, + "padding_mode": "zeros", + "vocoder_input_dim": 64, + "latent_mean": None, + "latent_std": None, +} + +# Tiny but *structurally faithful* config: two encoder stages, two decoder +# blocks, stereo I/O, and a consistent compression factor of 12 samples/frame +# (prod(dec_strides) == prod(enc_strides) * enc_hop_length == hop_size). +_TINY_KWARGS: dict = { + "vocoder_input_dim": 2, + "dec_dim": 4, + "dec_c_mults": (1, 2), + "dec_strides": (3, 4), + "dec_out_channels": 2, + "hop_size": 12, + "enc_dim": 4, + "enc_num_blocks": 1, + "enc_n_fft": 8, + "enc_hop_length": 2, + "enc_latent_dim": 4, + "enc_c_mults": (1, 2), + "enc_strides": (2, 3), +} + + +def _tiny_config(**overrides) -> Cosmos3AudioConfig: + """Build the shared tiny config, applying *overrides*.""" + config = Cosmos3AudioConfig(**{**_TINY_KWARGS, **overrides}) + config.validate() + return config + + +# --------------------------------------------------------------------------- +# PyTorch reference — a direct transcription of the upstream diffusers module +# --------------------------------------------------------------------------- + + +class _RefSnake1d(tnn.Module): + """Reference ``Snake1d``.""" + + def __init__(self, hidden_dim: int, logscale: bool = True): + super().__init__() + self.alpha = tnn.Parameter(torch.zeros(1, hidden_dim, 1)) + self.beta = tnn.Parameter(torch.zeros(1, hidden_dim, 1)) + self.logscale = logscale + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + alpha = self.alpha if not self.logscale else torch.exp(self.alpha) + beta = self.beta if not self.logscale else torch.exp(self.beta) + return hidden_states + (beta + 1e-9).reciprocal() * torch.sin( + alpha * hidden_states + ).pow(2) + + +class _RefFP32LayerNorm(tnn.LayerNorm): + """Reference ``FP32LayerNorm`` (diffusers).""" + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + origin_dtype = inputs.dtype + return torch.nn.functional.layer_norm( + inputs.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ).to(origin_dtype) + + +class _RefConvNeXtBlock(tnn.Module): + """Reference ``Cosmos3AudioConvNeXtBlock`` (non-causal, SnakeBeta).""" + + def __init__(self, hidden_dim: int, intermediate_dim: int): + super().__init__() + self.dwconv = tnn.Sequential( + tnn.ConstantPad1d((3, 3), 0), + tnn.Conv1d(hidden_dim, hidden_dim, kernel_size=7, groups=hidden_dim), + ) + self.norm = _RefFP32LayerNorm(hidden_dim, eps=1e-5, bias=False) + self.pwconv1 = tnn.Conv1d(hidden_dim, intermediate_dim, kernel_size=1) + self.act = _RefSnake1d(intermediate_dim) + self.pwconv2 = tnn.Conv1d(intermediate_dim, hidden_dim, kernel_size=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + residual = hidden_states + hidden_states = self.dwconv(hidden_states) + hidden_states = self.norm(hidden_states.permute(0, 2, 1)).permute(0, 2, 1) + hidden_states = self.pwconv1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.pwconv2(hidden_states) + return residual + hidden_states + + +class _RefEncoder(tnn.Module): + """Reference ``Cosmos3AudioSpectrogramConvNeXtEncoder`` (weight-norm folded).""" + + def __init__(self, config: Cosmos3AudioConfig): + super().__init__() + self.input_channels = config.encoder_input_channels + self.n_fft = config.enc_n_fft + self.hop_length = config.enc_hop_length + channels = config.enc_dim + multiples = config.enc_c_mults + strides = config.enc_strides + + layers: list[tnn.Module] = [ + tnn.Conv1d(config.spectrogram_channels, multiples[0] * channels, 1, bias=False) + ] + for index, stride in enumerate(strides): + input_dim = multiples[index] * channels + output_dim = ( + multiples[index + 1] * channels + if index < len(multiples) - 1 + else multiples[-1] * channels + ) + for _ in range(config.enc_num_blocks): + layers.append(_RefConvNeXtBlock(input_dim, input_dim * 4)) + layers.append( + tnn.Conv1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ) + ) + layers.append( + tnn.Conv1d(multiples[-1] * channels, config.enc_latent_dim, 1, bias=False) + ) + self.layers = tnn.Sequential(*layers) + + def _spectrogram(self, waveform: torch.Tensor) -> torch.Tensor: + pad_left = (self.n_fft - self.hop_length) // 2 + pad_right = (self.n_fft - self.hop_length) - pad_left + waveform = torch.nn.functional.pad(waveform, (pad_left, pad_right)).float() + window = torch.hann_window(self.n_fft, dtype=waveform.dtype) + return torch.stft( + waveform, + n_fft=self.n_fft, + hop_length=self.hop_length, + win_length=self.n_fft, + window=window, + center=False, + normalized=False, + onesided=True, + return_complex=True, + ) + + def forward(self, audio: torch.Tensor) -> torch.Tensor: + batch_size, num_channels, num_samples = audio.shape + if num_channels > 1: + audio = audio.reshape(batch_size * num_channels, 1, num_samples) + spectrogram = self._spectrogram(audio.squeeze(1)) + real, imaginary = torch.view_as_real(spectrogram).chunk(2, dim=-1) + spectrogram = torch.cat([real, imaginary], dim=1).squeeze(-1) + spectrogram = spectrogram.to(audio.dtype) + if num_channels > 1: + spectrogram = spectrogram.reshape( + batch_size, num_channels * spectrogram.shape[1], spectrogram.shape[2] + ) + return self.layers(spectrogram).transpose(1, 2) + + +class _RefResidualUnit(tnn.Module): + """Reference ``Cosmos3AudioResidualUnit`` (weight-norm folded).""" + + def __init__(self, dimension: int = 16, dilation: int = 1): + super().__init__() + pad = ((7 - 1) * dilation) // 2 + self.snake1 = _RefSnake1d(dimension) + self.conv1 = tnn.Conv1d(dimension, dimension, 7, dilation=dilation, padding=pad) + self.snake2 = _RefSnake1d(dimension) + self.conv2 = tnn.Conv1d(dimension, dimension, 1) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + output_tensor = self.conv1(self.snake1(hidden_state)) + output_tensor = self.conv2(self.snake2(output_tensor)) + padding = (hidden_state.shape[-1] - output_tensor.shape[-1]) // 2 + if padding > 0: + hidden_state = hidden_state[..., padding:-padding] + return hidden_state + output_tensor + + +class _RefDecoderBlock(tnn.Module): + """Reference ``Cosmos3AudioDecoderBlock`` (weight-norm folded).""" + + def __init__(self, input_dim: int, output_dim: int, stride: int, output_padding: int): + super().__init__() + self.snake1 = _RefSnake1d(input_dim) + self.conv_t1 = tnn.ConvTranspose1d( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + output_padding=output_padding, + ) + self.res_unit1 = _RefResidualUnit(output_dim, dilation=1) + self.res_unit2 = _RefResidualUnit(output_dim, dilation=3) + self.res_unit3 = _RefResidualUnit(output_dim, dilation=9) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + hidden_state = self.snake1(hidden_state) + hidden_state = self.conv_t1(hidden_state) + hidden_state = self.res_unit1(hidden_state) + hidden_state = self.res_unit2(hidden_state) + return self.res_unit3(hidden_state) + + +class _RefDecoder(tnn.Module): + """Reference ``Cosmos3AudioDecoder`` (weight-norm folded).""" + + def __init__(self, config: Cosmos3AudioConfig): + super().__init__() + channels = config.dec_dim + strides = config.decoder_upsampling_ratios + multiples = config.decoder_channel_multiples + self.conv1 = tnn.Conv1d( + config.vocoder_input_dim, channels * multiples[-1], kernel_size=7, padding=3 + ) + self.block = tnn.ModuleList( + [ + _RefDecoderBlock( + input_dim=channels * multiples[len(strides) - i], + output_dim=channels * multiples[len(strides) - i - 1], + stride=stride, + output_padding=stride % 2, + ) + for i, stride in enumerate(strides) + ] + ) + self.snake1 = _RefSnake1d(channels) + self.conv2 = tnn.Conv1d( + channels, config.dec_out_channels, kernel_size=7, padding=3, bias=False + ) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + hidden_state = self.conv1(hidden_state) + for layer in self.block: + hidden_state = layer(hidden_state) + hidden_state = self.snake1(hidden_state) + return self.conv2(hidden_state) + + +def _randomize(module: tnn.Module, seed: int) -> dict[str, torch.Tensor]: + """Fill every parameter with small random values and return the state dict. + + Snake ``alpha``/``beta`` are log-scaled, so keeping them near zero keeps + ``exp()`` well conditioned and mirrors a trained checkpoint. + """ + generator = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for name, param in module.named_parameters(): + scale = 0.1 if name.endswith((".alpha", ".beta")) else 0.3 + param.copy_(torch.randn(param.shape, generator=generator) * scale) + return {name: value.detach().clone() for name, value in module.state_dict().items()} + + +def _session(model, state_dict: dict[str, torch.Tensor], config) -> OnnxModelSession: + """Apply weights to a single-model package and open an ORT session. + + Every initializer must be filled — an unfilled one means the reference + state-dict names and the graph initializer names have drifted apart, which + is exactly the failure mode the HF-aligned naming exists to prevent. + """ + package = ModelPackage({"model": model}, config=config) + package.apply_weights(state_dict) + unset = [ + name for name, init in model.graph.initializers.items() if init.const_value is None + ] + assert not unset, f"initializers without weights: {unset}" + return OnnxModelSession(model) + + +def _wrap_module_graph( + module: nn.Module, + name: str, + channels: int, +) -> ir.Model: + """Build a single-module ONNX graph with a ``(B, channels, T)`` contract.""" + module._set_name(name) + graph = ir.Graph([], [], nodes=[], name=name, opset_imports={"": OPSET_VERSION}) + builder = GraphBuilder(graph) + x = builder.input("x", dtype=ir.DataType.FLOAT, shape=["batch", channels, "time"]) + builder.add_output(module(builder.op, x), "y") + return ir.Model(graph, ir_version=11) + + +# --------------------------------------------------------------------------- +# Config parsing and validation +# --------------------------------------------------------------------------- + + +def test_real_config_parses_and_validates(): + """The published sound_tokenizer config round-trips through from_diffusers.""" + config = Cosmos3AudioConfig.from_diffusers(_REAL_CONFIG) + + assert config.model_type == "autoencoder_v2" + assert config.sampling_rate == 48000 + assert config.stereo is True + assert config.input_channels == 1 + assert config.encoder_input_channels == 2 + assert config.dec_out_channels == 2 + assert config.resolved_hop_size == 1920 + assert config.vocoder_input_dim == 64 + assert config.enc_latent_dim == 128 + assert config.enc_type == "spec_convnext" + assert config.dec_type == "oobleck" + assert config.bottleneck_type == "vae" + assert config.activation == "snakebeta" + assert config.snake_logscale is True + assert config.latent_mean is None + assert config.latent_std is None + assert config.encoder_enabled is True + # Tuples, not the lists that came out of JSON. + assert config.enc_c_mults == (1, 2, 4) + assert config.enc_strides == (4, 5, 6) + assert config.dec_c_mults == (1, 2, 4, 8, 16) + assert config.dec_strides == (2, 4, 5, 6, 8) + + +def test_real_config_derived_geometry(): + """Derived shapes agree with the checkpoint tensor shapes.""" + config = Cosmos3AudioConfig.from_diffusers(_REAL_CONFIG) + + # encoder.layers.0.weight is [192, 132, 1] in the published checkpoint. + assert config.spectrogram_channels == 132 + assert config.stft_num_bins == 33 + assert config.stft_pad_left == 24 + assert config.stft_pad_right == 24 + assert config.latent_channels == 64 + assert config.moments_channels == 128 + assert config.audio_channels == 2 + # Encoder and decoder must agree on samples-per-latent-frame. + assert config.decoder_upsample_factor == 1920 + assert config.encoder_downsample_factor == 1920 + assert config.decoder_upsampling_ratios == (8, 6, 5, 4, 2) + assert config.decoder_channel_multiples == (1, 1, 2, 4, 8, 16) + + +def test_hop_size_defaults_to_decoder_stride_product(): + """An absent hop_size falls back to prod(dec_strides).""" + raw = dict(_REAL_CONFIG) + del raw["hop_size"] + config = Cosmos3AudioConfig.from_diffusers(raw) + assert config.resolved_hop_size == 1920 + + +def test_from_diffusers_ignores_unknown_keys(): + """Diffusers bookkeeping keys do not break construction.""" + raw = { + **_REAL_CONFIG, + "_class_name": "Cosmos3AVAEAudioTokenizer", + "_diffusers_version": "1", + } + assert Cosmos3AudioConfig.from_diffusers(raw).sampling_rate == 48000 + + +def test_from_diffusers_encoder_override(): + """The encoder_enabled override reaches the parsed config.""" + config = Cosmos3AudioConfig.from_diffusers(_REAL_CONFIG, encoder_enabled=False) + assert config.encoder_enabled is False + + +@pytest.mark.parametrize( + ("overrides", "expected"), + [ + ({"model_type": "autoencoder_v1"}, NotImplementedError), + ({"use_wav_as_input": False}, NotImplementedError), + ({"enc_type": "convnext"}, NotImplementedError), + ({"dec_type": "hifigan"}, NotImplementedError), + ({"bottleneck_type": "rvq"}, NotImplementedError), + ({"bottleneck": {"type": "rvq"}}, NotImplementedError), + ({"dec_use_snake": False}, NotImplementedError), + ({"dec_final_tanh": True}, NotImplementedError), + ({"dec_anti_aliasing": True}, NotImplementedError), + ({"dec_use_nearest_upsample": True}, NotImplementedError), + ({"dec_use_tanh_at_final": True}, NotImplementedError), + ({"activation": "gelu"}, NotImplementedError), + ({"snake_logscale": False}, NotImplementedError), + ({"anti_aliasing": True}, NotImplementedError), + ({"use_cuda_kernel": True}, NotImplementedError), + ({"causal": True}, NotImplementedError), + ({"padding_mode": "reflect"}, NotImplementedError), + ({"latent_mean": 0.0}, NotImplementedError), + ({"latent_std": [1.0, 2.0]}, NotImplementedError), + ({"enc_c_mults": [1, 2]}, ValueError), + ({"dec_c_mults": [1, 2]}, ValueError), + ({"enc_latent_dim": 100}, ValueError), + ({"enc_n_fft": 63}, ValueError), + ({"enc_hop_length": 128}, ValueError), + ({"input_channels": 0}, ValueError), + ({"dec_dim": 0}, ValueError), + ({"dec_out_channels": 1}, ValueError), + ({"hop_size": 960}, ValueError), + ({"enc_hop_length": 8}, ValueError), + ], +) +def test_config_rejects_unsupported_variants(overrides, expected): + """Every upstream ``__init__`` guard is reproduced by ``validate()``.""" + with pytest.raises(expected): + Cosmos3AudioConfig.from_diffusers({**_REAL_CONFIG, **overrides}) + + +def test_config_rejects_empty_and_non_positive_strides(): + """Sequence fields are validated eagerly in ``__post_init__``.""" + with pytest.raises(ValueError, match="at least one entry"): + Cosmos3AudioConfig(dec_strides=()) + with pytest.raises(ValueError, match="positive integers"): + Cosmos3AudioConfig(enc_strides=(4, 0, 6)) + + +# --------------------------------------------------------------------------- +# Encoder-weight availability +# --------------------------------------------------------------------------- + + +def test_state_dict_has_encoder_detection(): + """Encoder presence is decided by the ``encoder.`` key prefix.""" + assert state_dict_has_encoder(["encoder.layers.0.weight_g", "decoder.conv1.bias"]) + assert not state_dict_has_encoder(["decoder.conv1.bias", "decoder.conv2.weight_g"]) + assert not state_dict_has_encoder([]) + + +def test_with_encoder_from_state_dict_disables_for_decoder_only_weights(): + """A decoder-only checkpoint flips ``encoder_enabled`` off.""" + config = Cosmos3AudioConfig.from_diffusers(_REAL_CONFIG) + assert config.encoder_enabled is True + + decoder_only = config.with_encoder_from_state_dict(["decoder.conv1.weight_g"]) + assert decoder_only.encoder_enabled is False + assert decoder_only is not config + # Everything else is preserved. + assert decoder_only.resolved_hop_size == 1920 + assert decoder_only.dec_c_mults == (1, 2, 4, 8, 16) + + # Idempotent when the flag already matches. + full = config.with_encoder_from_state_dict(["encoder.layers.0.weight_g"]) + assert full is config + + +def test_create_tokenizer_dispatches_on_encoder_enabled(): + """The factory returns the module variant matching the config flag.""" + full = create_cosmos3_avae_audio_tokenizer(_tiny_config()) + assert isinstance(full, Cosmos3AVAEAudioTokenizer) + assert full.encoder_available is True + assert hasattr(full, "encoder") + assert hasattr(full, "bottleneck") + + decoder_only = create_cosmos3_avae_audio_tokenizer(_tiny_config(encoder_enabled=False)) + assert isinstance(decoder_only, Cosmos3AVAEAudioDecoderOnlyTokenizer) + assert not isinstance(decoder_only, Cosmos3AVAEAudioTokenizer) + assert decoder_only.encoder_available is False + assert not hasattr(decoder_only, "encoder") + assert not hasattr(decoder_only, "bottleneck") + + +def test_select_task_dispatches_on_encoder_enabled(): + """Task selection follows the same encoder-availability signal.""" + assert select_cosmos3_audio_task(_tiny_config()) is Cosmos3AVAEAudioTokenizerTask + assert ( + select_cosmos3_audio_task(_tiny_config(encoder_enabled=False)) + is Cosmos3AVAEAudioDecoderTask + ) + + +def test_tokenizer_task_rejects_decoder_only_config(): + """Building the encoder path on decoder-only weights is a hard error.""" + config = _tiny_config() + module = create_cosmos3_avae_audio_tokenizer(config) + decoder_only_config = _tiny_config(encoder_enabled=False) + with pytest.raises(ValueError, match="encoder_enabled=True"): + Cosmos3AVAEAudioTokenizerTask().build(module, decoder_only_config) + + +def test_tokenizer_task_rejects_module_without_encoder(): + """The component spec catches a decoder-only module on the encoder path.""" + config = _tiny_config() + module = Cosmos3AVAEAudioDecoderOnlyTokenizer(config) + with pytest.raises(TypeError, match="encoder"): + Cosmos3AVAEAudioTokenizerTask().build(module, config) + + +# --------------------------------------------------------------------------- +# Graph construction and I/O contracts +# --------------------------------------------------------------------------- + + +def test_decoder_only_package_has_no_encoder_initializers(): + """Decoder-only builds never emit an initializer no weight can fill.""" + config = _tiny_config(encoder_enabled=False) + module = create_cosmos3_avae_audio_tokenizer(config) + package = Cosmos3AVAEAudioDecoderTask().build(module, config) + + assert set(package) == {"decoder"} + names = list(package["decoder"].graph.initializers) + assert names, "decoder graph must have initializers" + assert all(name.startswith("decoder.") for name in names) + assert not any(name.startswith("encoder.") for name in names) + assert not any("bottleneck" in name for name in names) + + +def test_decoder_graph_io_contract(): + """``latents -> waveform`` shapes, names and dtypes.""" + config = _tiny_config() + module = create_cosmos3_avae_audio_tokenizer(config) + model = Cosmos3AVAEAudioDecoderTask().build(module, config)["decoder"] + + (latents,) = model.graph.inputs + assert latents.name == "latents" + assert latents.dtype == config.dtype + assert latents.shape[1] == config.latent_channels + + (waveform,) = model.graph.outputs + assert waveform.name == "waveform" + + +def test_encoder_graph_io_contract(): + """``audio -> moments/latent_mean/latent_std`` names and shapes.""" + config = _tiny_config() + module = create_cosmos3_avae_audio_tokenizer(config) + package = Cosmos3AVAEAudioTokenizerTask().build(module, config) + + assert set(package) == {"encoder", "decoder"} + encoder = package["encoder"] + + (audio,) = encoder.graph.inputs + assert audio.name == "audio" + assert audio.dtype == config.dtype + assert audio.shape[1] == config.encoder_input_channels + + assert [out.name for out in encoder.graph.outputs] == [ + "moments", + "latent_mean", + "latent_std", + ] + + encoder_names = list(encoder.graph.initializers) + assert all(name.startswith("encoder.") for name in encoder_names) + decoder_names = list(package["decoder"].graph.initializers) + assert all(name.startswith("decoder.") for name in decoder_names) + # The two graphs must not fight over the same weight names. + assert not set(encoder_names) & set(decoder_names) + + +def test_encoder_graph_has_no_random_ops(): + """Posterior sampling stays outside ONNX so the graph is deterministic.""" + config = _tiny_config() + module = create_cosmos3_avae_audio_tokenizer(config) + encoder = Cosmos3AVAEAudioTokenizerTask().build(module, config)["encoder"] + + op_types = {node.op_type for node in encoder.graph} + assert not {"RandomNormal", "RandomNormalLike", "RandomUniform", "Multinomial"} & op_types + # The STFT front-end is part of the exported encoder. + assert "STFT" in op_types + assert "Softplus" in op_types + + +def test_real_config_initializer_names_match_checkpoint_layout(): + """Full-size initializer names reproduce the HF module paths exactly.""" + config = Cosmos3AudioConfig.from_diffusers(_REAL_CONFIG) + module = create_cosmos3_avae_audio_tokenizer(config) + package = Cosmos3AVAEAudioTokenizerTask().build(module, config) + + encoder = package["encoder"].graph.initializers + decoder = package["decoder"].graph.initializers + + # Encoder: nn.Sequential indices, depthwise conv nested at dwconv.1. + assert tuple(encoder["encoder.layers.0.weight"].shape) == (192, 132, 1) + assert "encoder.layers.0.bias" not in encoder + assert tuple(encoder["encoder.layers.1.dwconv.1.weight"].shape) == (192, 1, 7) + assert tuple(encoder["encoder.layers.1.norm.weight"].shape) == (192,) + assert "encoder.layers.1.norm.bias" not in encoder + assert tuple(encoder["encoder.layers.1.act.alpha"].shape) == (1, 768, 1) + assert tuple(encoder["encoder.layers.3.weight"].shape) == (384, 192, 8) + assert tuple(encoder["encoder.layers.6.weight"].shape) == (768, 384, 10) + assert tuple(encoder["encoder.layers.9.weight"].shape) == (768, 768, 12) + assert tuple(encoder["encoder.layers.10.weight"].shape) == (128, 768, 1) + assert "encoder.layers.10.bias" not in encoder + + # Decoder: conv1 / block.{i} / snake1 / conv2, ConvTranspose weight is (in, out, k). + assert tuple(decoder["decoder.conv1.weight"].shape) == (5120, 64, 7) + assert tuple(decoder["decoder.block.0.conv_t1.weight"].shape) == (5120, 2560, 16) + assert tuple(decoder["decoder.block.0.res_unit2.conv1.weight"].shape) == (2560, 2560, 7) + assert tuple(decoder["decoder.block.0.res_unit3.snake2.beta"].shape) == (1, 2560, 1) + assert tuple(decoder["decoder.snake1.alpha"].shape) == (1, 320, 1) + assert tuple(decoder["decoder.conv2.weight"].shape) == (2, 320, 7) + assert "decoder.conv2.bias" not in decoder + + # Exactly the 207 tensors the published checkpoint holds after folding. + assert len(encoder) + len(decoder) == 207 + + +def test_published_config_cannot_reveal_encoder_presence(): + """The shipped config declares ``enc_*`` even for decoder-only checkpoints. + + ``nvidia/Cosmos3-Nano``, ``nvidia/Cosmos3-Super`` (full) and + ``nvidia/Cosmos3-Super-Text2Image`` (decoder-only) ship byte-identical + ``sound_tokenizer/config.json`` files. This pins the reason encoder + presence must come from the checkpoint, not the config. + """ + assert "encoder_enabled" not in _REAL_CONFIG + assert _REAL_CONFIG["enc_type"] == "spec_convnext" + assert _REAL_CONFIG["enc_dim"] == 192 + # Parsing alone therefore optimistically assumes a full checkpoint. + assert Cosmos3AudioConfig.from_diffusers(_REAL_CONFIG).encoder_enabled is True + + +def test_from_diffusers_resolves_encoder_from_weight_names(): + """``weight_names=`` is the single-call, checkpoint-driven entry point.""" + full = Cosmos3AudioConfig.from_diffusers( + _REAL_CONFIG, weight_names=["encoder.layers.0.weight_g", "decoder.conv1.weight_g"] + ) + assert full.encoder_enabled is True + + decoder_only = Cosmos3AudioConfig.from_diffusers( + _REAL_CONFIG, weight_names=["decoder.conv1.weight_g", "decoder.conv2.weight_v"] + ) + assert decoder_only.encoder_enabled is False + assert select_cosmos3_audio_task(decoder_only) is Cosmos3AVAEAudioDecoderTask + + +def test_from_diffusers_rejects_conflicting_encoder_signals(): + """``encoder_enabled`` and ``weight_names`` answer the same question.""" + with pytest.raises(ValueError, match="not both"): + Cosmos3AudioConfig.from_diffusers( + _REAL_CONFIG, encoder_enabled=True, weight_names=["decoder.conv1.weight_g"] + ) + + +def test_real_decoder_only_package_matches_published_layout(): + """A decoder-only build emits exactly the 145 folded ``decoder.*`` tensors. + + The published decoder-only checkpoint holds 182 tensors, 37 of which are + ``weight_g``/``weight_v`` magnitude entries that fold away: 182 - 37 = 145. + """ + config = Cosmos3AudioConfig.from_diffusers( + _REAL_CONFIG, weight_names=["decoder.conv1.weight_g"] + ) + package = Cosmos3AVAEAudioDecoderTask().build( + create_cosmos3_avae_audio_tokenizer(config), config + ) + + assert set(package) == {"decoder"} + names = list(package["decoder"].graph.initializers) + assert len(names) == 145 + assert all(name.startswith("decoder.") for name in names) + + # The full build adds exactly the encoder half: 67 - 5 folded pairs = 62. + full_config = Cosmos3AudioConfig.from_diffusers(_REAL_CONFIG) + full = Cosmos3AVAEAudioTokenizerTask().build( + create_cosmos3_avae_audio_tokenizer(full_config), full_config + ) + assert len(full["encoder"].graph.initializers) == 62 + assert len(full["decoder"].graph.initializers) == 145 + + +def test_graphs_contain_no_runtime_weight_norm_math(): + """``weight_norm`` is folded offline — the graph holds dense kernels only. + + A graph that reconstructed ``g * v / ||v||`` at runtime would need a + norm (``ReduceL2``/``Pow`` + ``ReduceSum`` + ``Sqrt``) and a ``Div`` feeding + each conv. None of those may appear. + """ + weight_norm_ops = {"Div", "Sqrt", "ReduceL2", "ReduceSumSquare", "Pow", "ReduceSum"} + + config = Cosmos3AudioConfig.from_diffusers(_REAL_CONFIG) + package = Cosmos3AVAEAudioTokenizerTask().build( + create_cosmos3_avae_audio_tokenizer(config), config + ) + + decoder_ops = [node.op_type for node in package["decoder"].graph] + assert not weight_norm_ops & set(decoder_ops) + # 37 weight-normed convs in the decoder — matching the checkpoint's 37 pairs. + assert decoder_ops.count("Conv") + decoder_ops.count("ConvTranspose") == 37 + + # The encoder's only Div is the peak-volume normalization; disable it and + # the encoder must likewise be free of every weight-norm op. + plain = Cosmos3AudioConfig.from_diffusers({**_REAL_CONFIG, "normalize_volume": False}) + plain_encoder = Cosmos3AVAEAudioTokenizerTask().build( + create_cosmos3_avae_audio_tokenizer(plain), plain + )["encoder"] + assert not weight_norm_ops & {node.op_type for node in plain_encoder.graph} + + encoder_ops = [node.op_type for node in package["encoder"].graph] + assert encoder_ops.count("Div") == 1 + assert encoder_ops.count("ReduceMax") == 1 + + +def test_default_dtype_is_float32_like_the_checkpoints(): + """Published sound-tokenizer weights are FP32, so that is the build default.""" + config = Cosmos3AudioConfig.from_diffusers(_REAL_CONFIG) + assert config.dtype == ir.DataType.FLOAT + + +def test_config_dtype_propagates_to_graph_and_initializers(): + """``config.dtype`` reaches graph I/O, every initializer, and the FP32 norm.""" + config = _tiny_config(dtype=ir.DataType.FLOAT16) + module = create_cosmos3_avae_audio_tokenizer(config) + package = Cosmos3AVAEAudioTokenizerTask().build(module, config) + + for model in package.values(): + assert all(inp.dtype == ir.DataType.FLOAT16 for inp in model.graph.inputs) + assert all( + init.dtype == ir.DataType.FLOAT16 for init in model.graph.initializers.values() + ), "every parameter must adopt the configured dtype" + + # FP32LayerNorm upcasts around the ConvNeXt norm when the graph is not fp32. + encoder_ops = [node.op_type for node in package["encoder"].graph] + assert "Cast" in encoder_ops + assert "LayerNormalization" in encoder_ops + + # A float32 build needs no cast around the norm. + fp32 = _tiny_config() + fp32_encoder = Cosmos3AVAEAudioTokenizerTask().build( + create_cosmos3_avae_audio_tokenizer(fp32), fp32 + )["encoder"] + assert all( + init.dtype == ir.DataType.FLOAT for init in fp32_encoder.graph.initializers.values() + ) + + +# --------------------------------------------------------------------------- +# weight_norm folding +# --------------------------------------------------------------------------- + + +def test_fold_weight_norm_matches_torch_parametrization(): + """Folding g/v reproduces the effective weight torch computes.""" + torch.manual_seed(0) + conv = torch.nn.utils.parametrizations.weight_norm(tnn.Conv1d(5, 7, 3)) + transposed = torch.nn.utils.parametrizations.weight_norm(tnn.ConvTranspose1d(5, 7, 4)) + with torch.no_grad(): + for module in (conv, transposed): + for param in module.parameters(): + param.copy_(torch.randn(param.shape)) + + state_dict = { + **{f"conv.{k}": v for k, v in conv.state_dict().items()}, + **{f"conv_t.{k}": v for k, v in transposed.state_dict().items()}, + } + folded = fold_weight_norm(state_dict) + + assert set(folded) == {"conv.weight", "conv.bias", "conv_t.weight", "conv_t.bias"} + torch.testing.assert_close(folded["conv.weight"], conv.weight, atol=1e-6, rtol=1e-6) + torch.testing.assert_close( + folded["conv_t.weight"], transposed.weight, atol=1e-6, rtol=1e-6 + ) + + +def test_fold_weight_norm_handles_legacy_key_names(): + """Legacy ``weight_g``/``weight_v`` keys (the shipped format) also fold.""" + direction = torch.randn(4, 3, 5) + magnitude = torch.randn(4, 1, 1) + folded = fold_weight_norm( + { + "layers.0.weight_g": magnitude, + "layers.0.weight_v": direction, + "layers.0.bias": torch.zeros(4), + } + ) + expected = magnitude * direction / direction.pow(2).sum(dim=(1, 2), keepdim=True).sqrt() + assert set(folded) == {"layers.0.weight", "layers.0.bias"} + torch.testing.assert_close(folded["layers.0.weight"], expected, atol=1e-6, rtol=1e-6) + + +def test_fold_weight_norm_requires_matching_direction(): + """A dangling magnitude tensor is an error, not a silent drop.""" + with pytest.raises(ValueError, match="no matching direction"): + fold_weight_norm({"conv.weight_g": torch.ones(2, 1, 1)}) + + +def test_fold_weight_norm_reconstructs_real_checkpoint_kernels(): + """Folding real checkpoint-shaped ``weight_g``/``weight_v`` pairs is exact. + + Uses the published shapes for the first/last weight-normed conv of each + half: ``encoder.layers.0`` (132 packed STFT channels in) and + ``decoder.block.0.conv_t1`` (a ConvTranspose, whose weight-norm axis 0 is + *input* channels, not output). + """ + torch.manual_seed(0) + cases = { + # (out, in, k) for Conv1d; weight_norm dim=0 -> g is (out, 1, 1) + "encoder.layers.0": ((192, 132, 1), (192, 1, 1)), + "encoder.layers.10": ((128, 768, 1), (128, 1, 1)), + # (in, out, k) for ConvTranspose1d; weight_norm dim=0 -> g is (in, 1, 1) + "decoder.block.0.conv_t1": ((5120, 2560, 16), (5120, 1, 1)), + "decoder.conv2": ((2, 320, 7), (2, 1, 1)), + } + state_dict: dict[str, torch.Tensor] = {} + for prefix, (v_shape, g_shape) in cases.items(): + state_dict[f"{prefix}.weight_v"] = torch.randn(v_shape) * 0.05 + state_dict[f"{prefix}.weight_g"] = torch.randn(g_shape).abs() + 0.1 + + folded = fold_weight_norm(state_dict) + assert set(folded) == {f"{prefix}.weight" for prefix in cases} + + for prefix, (v_shape, _) in cases.items(): + direction = state_dict[f"{prefix}.weight_v"] + magnitude = state_dict[f"{prefix}.weight_g"] + expected = torch._weight_norm(direction, magnitude, 0) + assert folded[f"{prefix}.weight"].shape == torch.Size(v_shape) + torch.testing.assert_close(folded[f"{prefix}.weight"], expected, atol=1e-6, rtol=1e-6) + + +def test_fold_weight_norm_leaves_bias_free_convs_bias_free(): + """``decoder.conv2`` and ``encoder.layers.10`` ship without a bias.""" + config = Cosmos3AudioConfig.from_diffusers(_REAL_CONFIG) + package = Cosmos3AVAEAudioTokenizerTask().build( + create_cosmos3_avae_audio_tokenizer(config), config + ) + assert "decoder.conv2.bias" not in package["decoder"].graph.initializers + assert "encoder.layers.10.bias" not in package["encoder"].graph.initializers + assert "encoder.layers.0.bias" not in package["encoder"].graph.initializers + # ...while the strided encoder convs and decoder convs do carry one. + assert "encoder.layers.3.bias" in package["encoder"].graph.initializers + assert "decoder.conv1.bias" in package["decoder"].graph.initializers + + +def test_preprocess_weights_is_exposed_on_both_module_variants(): + """Both module paths expose the builder's ``preprocess_weights`` hook.""" + direction = torch.randn(2, 2, 3) + magnitude = torch.randn(2, 1, 1) + payload = {"decoder.conv1.weight_g": magnitude, "decoder.conv1.weight_v": direction} + for module in ( + create_cosmos3_avae_audio_tokenizer(_tiny_config()), + create_cosmos3_avae_audio_tokenizer(_tiny_config(encoder_enabled=False)), + ): + assert set(module.preprocess_weights(payload)) == {"decoder.conv1.weight"} + + +# --------------------------------------------------------------------------- +# Numerical parity against the PyTorch reference +# --------------------------------------------------------------------------- + + +def test_snake_beta_matches_pytorch(): + """SnakeBeta (log-scaled) matches ``Snake1d`` elementwise.""" + channels = 6 + reference = _RefSnake1d(channels) + state_dict = {f"snake.{k}": v for k, v in _randomize(reference, seed=7).items()} + + model = _wrap_module_graph(_Snake1d(channels), "snake", channels) + session = _session(model, state_dict, None) + + sample = np.random.default_rng(3).standard_normal((2, channels, 9)).astype(np.float32) + got = session.run({"x": sample})["y"] + expected = reference(torch.from_numpy(sample)).detach().numpy() + np.testing.assert_allclose(got, expected, atol=1e-5, rtol=1e-5) + + +def test_snake_beta_without_logscale_matches_pytorch(): + """``snake_logscale=False`` uses the raw alpha/beta parameters.""" + channels = 5 + reference = _RefSnake1d(channels, logscale=False) + _randomize(reference, seed=61) + with torch.no_grad(): + # Keep beta away from -1e-9 so the reciprocal stays well conditioned. + reference.beta.copy_(reference.beta.abs() + 0.5) + state_dict = {f"snake.{k}": v.detach().clone() for k, v in reference.state_dict().items()} + + model = _wrap_module_graph(_Snake1d(channels, logscale=False), "snake", channels) + session = _session(model, state_dict, None) + + sample = np.random.default_rng(67).standard_normal((1, channels, 8)).astype(np.float32) + got = session.run({"x": sample})["y"] + expected = reference(torch.from_numpy(sample)).detach().numpy() + np.testing.assert_allclose(got, expected, atol=1e-5, rtol=1e-5) + + +def test_residual_unit_matches_pytorch(): + """One dilated residual unit matches the reference bit-for-bit (fp32 tol).""" + dim, dilation = 5, 3 + reference = _RefResidualUnit(dim, dilation=dilation) + state_dict = {f"unit.{k}": v for k, v in _randomize(reference, seed=11).items()} + + model = _wrap_module_graph(Cosmos3AudioResidualUnit(dim, dilation=dilation), "unit", dim) + session = _session(model, state_dict, None) + + sample = np.random.default_rng(5).standard_normal((2, dim, 16)).astype(np.float32) + got = session.run({"x": sample})["y"] + expected = reference(torch.from_numpy(sample)).detach().numpy() + + assert got.shape == sample.shape + np.testing.assert_allclose(got, expected, atol=1e-4, rtol=1e-4) + + +def test_convnext_block_matches_pytorch(): + """The SnakeBeta ConvNeXt block (depthwise conv + FP32 LayerNorm) matches.""" + hidden, intermediate = 4, 16 + reference = _RefConvNeXtBlock(hidden, intermediate) + state_dict = {f"block.{k}": v for k, v in _randomize(reference, seed=71).items()} + + model = _wrap_module_graph( + Cosmos3AudioConvNeXtBlock(hidden, intermediate), "block", hidden + ) + session = _session(model, state_dict, None) + + sample = np.random.default_rng(73).standard_normal((2, hidden, 12)).astype(np.float32) + got = session.run({"x": sample})["y"] + expected = reference(torch.from_numpy(sample)).detach().numpy() + np.testing.assert_allclose(got, expected, atol=1e-4, rtol=1e-4) + + +def test_decoder_graph_matches_pytorch(): + """The whole tiny decoder matches the reference, including the clamp.""" + config = _tiny_config() + reference = _RefDecoder(config) + state_dict = {f"decoder.{k}": v for k, v in _randomize(reference, seed=13).items()} + + module = create_cosmos3_avae_audio_tokenizer(config) + model = Cosmos3AVAEAudioDecoderTask().build(module, config)["decoder"] + session = _session(model, state_dict, config) + + frames = 5 + latents = ( + np.random.default_rng(17) + .standard_normal((2, config.latent_channels, frames)) + .astype(np.float32) + ) + got = session.run({"latents": latents})["waveform"] + expected = reference(torch.from_numpy(latents)).clamp(-1.0, 1.0).detach().numpy() + + assert got.shape == (2, config.audio_channels, frames * config.resolved_hop_size) + np.testing.assert_allclose(got, expected, atol=1e-4, rtol=1e-4) + + +def test_encoder_graph_matches_pytorch(): + """The whole tiny encoder — STFT front-end included — matches the reference.""" + config = _tiny_config(normalize_volume=False) + reference = _RefEncoder(config) + state_dict = {f"encoder.{k}": v for k, v in _randomize(reference, seed=19).items()} + + module = create_cosmos3_avae_audio_tokenizer(config) + model = Cosmos3AVAEAudioTokenizerTask().build(module, config)["encoder"] + session = _session(model, state_dict, config) + + num_samples = config.resolved_hop_size * 4 + audio = ( + np.random.default_rng(23) + .standard_normal((2, config.encoder_input_channels, num_samples)) + .astype(np.float32) + ) + outputs = session.run({"audio": audio}) + + # Upstream `_encode` transposes the channels-last encoder output back. + expected_moments = reference(torch.from_numpy(audio)).transpose(1, 2) + expected_mean, expected_scale = expected_moments.chunk(2, dim=1) + expected_std = torch.nn.functional.softplus(expected_scale) + 1e-4 + + frames = num_samples // config.resolved_hop_size + assert outputs["moments"].shape == (2, config.moments_channels, frames) + assert outputs["latent_mean"].shape == (2, config.latent_channels, frames) + assert outputs["latent_std"].shape == (2, config.latent_channels, frames) + + np.testing.assert_allclose( + outputs["moments"], expected_moments.detach().numpy(), atol=1e-4, rtol=1e-4 + ) + np.testing.assert_allclose( + outputs["latent_mean"], expected_mean.detach().numpy(), atol=1e-4, rtol=1e-4 + ) + np.testing.assert_allclose( + outputs["latent_std"], expected_std.detach().numpy(), atol=1e-4, rtol=1e-4 + ) + + +def test_encoder_graph_normalizes_volume_and_pads_to_hop_size(): + """Peak normalization and hop padding are applied inside the graph.""" + config = _tiny_config() + reference = _RefEncoder(config) + state_dict = {f"encoder.{k}": v for k, v in _randomize(reference, seed=29).items()} + + module = create_cosmos3_avae_audio_tokenizer(config) + model = Cosmos3AVAEAudioTokenizerTask().build(module, config)["encoder"] + session = _session(model, state_dict, config) + + hop = config.resolved_hop_size + # Deliberately not a multiple of hop_size — the graph must right-pad. + num_samples = 3 * hop + 5 + audio = ( + np.random.default_rng(31) + .standard_normal((1, config.encoder_input_channels, num_samples)) + .astype(np.float32) + * 4.0 + ) + + torch_audio = torch.from_numpy(audio) + normalized = torch_audio / (torch_audio.abs().max() + 1e-5) * 0.95 + padding = (hop - (num_samples % hop)) % hop + padded = torch.nn.functional.pad(normalized, (0, padding)) + expected = reference(padded).transpose(1, 2) + + outputs = session.run({"audio": audio}) + assert outputs["moments"].shape == (1, config.moments_channels, 4) + np.testing.assert_allclose( + outputs["moments"], expected.detach().numpy(), atol=1e-4, rtol=1e-4 + ) + + +def test_decoder_graph_handles_mono_and_variable_lengths(): + """A mono (non-stereo) config still round-trips through the decoder graph.""" + config = _tiny_config(stereo=False, dec_out_channels=1) + reference = _RefDecoder(config) + state_dict = {f"decoder.{k}": v for k, v in _randomize(reference, seed=37).items()} + + module = create_cosmos3_avae_audio_tokenizer(config) + model = Cosmos3AVAEAudioDecoderTask().build(module, config)["decoder"] + session = _session(model, state_dict, config) + + for frames in (1, 7): + latents = ( + np.random.default_rng(41 + frames) + .standard_normal((1, config.latent_channels, frames)) + .astype(np.float32) + ) + got = session.run({"latents": latents})["waveform"] + expected = reference(torch.from_numpy(latents)).clamp(-1.0, 1.0).detach().numpy() + assert got.shape == (1, 1, frames * config.resolved_hop_size) + np.testing.assert_allclose(got, expected, atol=1e-4, rtol=1e-4) + + +def test_mono_encoder_graph_matches_pytorch(): + """The single-channel STFT path (no channel folding) also matches.""" + config = _tiny_config(stereo=False, dec_out_channels=1, normalize_volume=False) + reference = _RefEncoder(config) + state_dict = {f"encoder.{k}": v for k, v in _randomize(reference, seed=43).items()} + + module = create_cosmos3_avae_audio_tokenizer(config) + model = Cosmos3AVAEAudioTokenizerTask().build(module, config)["encoder"] + session = _session(model, state_dict, config) + + num_samples = config.resolved_hop_size * 3 + audio = np.random.default_rng(47).standard_normal((2, 1, num_samples)).astype(np.float32) + got = session.run({"audio": audio})["moments"] + expected = reference(torch.from_numpy(audio)).transpose(1, 2).detach().numpy() + + assert got.shape == (2, config.moments_channels, 3) + np.testing.assert_allclose(got, expected, atol=1e-4, rtol=1e-4) + + +def test_gelu_convnext_variant_matches_pytorch(): + """``enc_use_snake=False`` swaps SnakeBeta for exact GELU.""" + + class _RefGeluBlock(_RefConvNeXtBlock): + def __init__(self, hidden_dim: int, intermediate_dim: int): + super().__init__(hidden_dim, intermediate_dim) + self.act = tnn.GELU() + + hidden, intermediate = 4, 16 + reference = _RefGeluBlock(hidden, intermediate) + state_dict = {f"block.{k}": v for k, v in _randomize(reference, seed=53).items()} + + model = _wrap_module_graph( + Cosmos3AudioConvNeXtBlock(hidden, intermediate, use_snake=False), "block", hidden + ) + session = _session(model, state_dict, None) + + sample = np.random.default_rng(59).standard_normal((2, hidden, 12)).astype(np.float32) + got = session.run({"x": sample})["y"] + expected = reference(torch.from_numpy(sample)).detach().numpy() + np.testing.assert_allclose(got, expected, atol=1e-4, rtol=1e-4) diff --git a/src/mobius/models/cosmos3_omni_generator.py b/src/mobius/models/cosmos3_omni_generator.py new file mode 100644 index 000000000..403a065ca --- /dev/null +++ b/src/mobius/models/cosmos3_omni_generator.py @@ -0,0 +1,1045 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""NVIDIA Cosmos3-Omni unified MoT transformer (Reasoner + Generator). + +This module builds the **complete neural** ``Cosmos3OmniTransformer`` graph — +the unified Mixture-of-Transformers (MoT) backbone that carries the +understanding ("und" / Reasoner) expert and the rectified-flow diffusion +generation ("gen" / Generator) expert in a *single* stack of layers, plus the +vision, optional Sound and optional Action projection heads. + +It is deliberately **not** the reasoner-only Qwen3-VL wrapper in +:mod:`mobius.models.cosmos3_omni`: that module exports the understanding tower +alone (``input_ids -> logits + KV cache``) and drops every generator weight. +Here the generator is the point — a single denoising step over a packed joint +token sequence. + +Architecture reference +---------------------- +Translated from ``huggingface/diffusers`` +``src/diffusers/models/transformers/transformer_cosmos3.py`` +(Apache License 2.0, Copyright 2025 The NVIDIA Team and The HuggingFace Team) +and the public ``nvidia/Cosmos3-Nano`` ``transformer/config.json``. The +architecture (module structure, parameter names, dual-pathway attention, +interleaved 3-axis mRoPE, timestep handling) is derived from that source; the +ONNX graph construction below is an independent implementation using +``onnxscript.nn``. + +Per-layer dataflow +------------------ +Every layer holds two complete expert weight sets and mixes them only inside +attention:: + + und_norm = input_layernorm(und) gen_norm = input_layernorm_moe_gen(gen) + q,k,v = to_q/to_k/to_v(und_norm) q',k',v' = add_{q,k,v}_proj(gen_norm) + q,k = norm_q/norm_k q',k' = norm_added_q/norm_added_k + k_ufg = k_norm_und_for_gen(k) (optional) + + und_attn = Attention(q, k, v, causal) gen_attn = Attention(q', [k_ufg;k'], [v;v']) + und += to_out(und_attn) gen += to_add_out(gen_attn) + und += mlp(post_attention_layernorm) gen += mlp_moe_gen(post_attention_layernorm_moe_gen) + +The understanding pathway is strictly causal and self-contained; the +generation pathway is non-causal and cross-attends over the concatenated +understanding + generation keys/values. + +Packed ONNX contract +-------------------- +Upstream's ``forward`` takes Python lists of ragged per-item tensors +(``vision_tokens: list[torch.Tensor]``, ``vision_token_shapes: list[tuple]``, +``vision_noisy_frame_indexes: list[torch.Tensor]``, ...) and performs +patchify/unpatchify using host-side shape arithmetic. Python lists and +host-resident shape values cannot appear in an ONNX signature, so this graph +is defined at the **packed-token boundary**: + +* everything upstream does *before* ``proj_in`` / ``audio_proj_in`` / + ``action_proj_in`` (per-item patchify, padding, packing, and the flattening + of ``noisy_frame_indexes`` into row offsets) is host preprocessing; +* everything upstream does *after* ``proj_out`` / ``audio_proj_out`` / + ``action_proj_out`` (unpatchify, per-item scatter back into ``[C, T, H, W]`` + buffers) is host postprocessing; +* the entire neural computation in between is in the graph, with no semantics + dropped: the flattened noisy row offsets that ``_apply_timestep_embeds_to_ + noisy_tokens`` computes from ``token_shapes`` become an explicit + ``*_timestep_token_indexes`` int64 input, and ``und_len`` becomes an + explicit int64 input instead of a Python ``int``. + +See :class:`~mobius.tasks._cosmos3_omni_generator.Cosmos3OmniGeneratorTask` +for the exact input/output names, dtypes and shapes. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius._configs._cosmos3_omni_generator import Cosmos3OmniGeneratorConfig +from mobius.components import ( + FCMLP, + INT64_MAX, + Embedding, + GatedMLP, + Linear, + RMSNorm, + TimestepEmbedding, +) + +if TYPE_CHECKING: + import torch + + +def _const_ints(op: OpBuilder, values: list[int]) -> ir.Value: + """Materialize a 1-D int64 constant.""" + return op.Constant(value_ints=values) + + +def _apply_rope( + op: OpBuilder, x: ir.Value, cos: ir.Value, sin: ir.Value, num_heads: int +) -> ir.Value: + """Apply half-split rotary embedding to a flat multi-head tensor. + + The opset-24 ``RotaryEmbedding`` op with the half-width cos/sin cache and + ``interleaved=0`` computes exactly upstream's + ``x * cos + rotate_half(x) * sin`` (upstream duplicates the frequency table + with ``cat((freqs, freqs), -1)``; the op consumes the un-duplicated half). + + Args: + op: ONNX op builder. + x: ``[1, seq, num_heads * head_dim]``. + cos: ``[1, seq, head_dim // 2]``. + sin: ``[1, seq, head_dim // 2]``. + num_heads: Heads packed into the last axis of *x*. + + Returns: + ``[1, seq, num_heads * head_dim]`` with rotary applied. + """ + return op.RotaryEmbedding(x, cos, sin, num_heads=num_heads, interleaved=0) + + +def _row_indices(op: OpBuilder, indexes: ir.Value) -> ir.Value: + """Turn a 1-D index vector into ScatterND row indices. + + Args: + op: ONNX op builder. + indexes: ``[N]`` int64 row indices. + + Returns: + ``[N, 1]`` int64 indices addressing whole rows of a 2-D tensor. + """ + return op.Unsqueeze(indexes, [-1]) + + +class Cosmos3OmniDomainAwareLinear(nn.Module): + """Per-embodiment-domain linear projection (Action head). + + Replicates upstream ``DomainAwareLinear``: instead of one shared weight, + an ``nn.Embedding`` table stores one flattened ``[in, out]`` weight matrix + and one ``[out]`` bias per embodiment domain, and each *token* selects its + matrix by ``domain_id``. This is a Gather + batched MatMul, **not** a + plain Linear — a plain Linear would silently collapse all domains onto one + weight set. + + Weight layout (matching the published checkpoint exactly):: + + action_proj_in.fc.weight [num_domains, out_features * in_features] + action_proj_in.bias.weight [num_domains, out_features] + + The flattened row is interpreted as ``[in_features, out_features]`` + (upstream does ``self.fc(domain_id).view(D, input_size, output_size)``), + i.e. it is already transposed relative to ``torch.nn.Linear``. + + Args: + in_features: Input width. + out_features: Output width. + num_domains: Number of embodiment domains in the table. + """ + + def __init__(self, in_features: int, out_features: int, num_domains: int): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.num_domains = num_domains + # Named ``fc`` / ``bias`` so parameters land on ``.fc.weight`` + # and ``.bias.weight`` — the published flat names. + self.fc = Embedding(num_domains, out_features * in_features) + self.bias = Embedding(num_domains, out_features) + + def forward(self, op: OpBuilder, x: ir.Value, domain_ids: ir.Value) -> ir.Value: + """Apply the per-token domain weight. + + Args: + op: ONNX op builder. + x: ``[N, in_features]`` packed tokens. + domain_ids: ``[N]`` int64 embodiment domain id per token. + + Returns: + ``[N, out_features]`` projected tokens. + """ + # [N, out*in] -> [N, in, out]; one weight matrix per token. + weight = self.fc(op, domain_ids) + weight = op.Reshape(weight, _const_ints(op, [-1, self.in_features, self.out_features])) + # [N, out] + bias = self.bias(op, domain_ids) + # [N, 1, in] @ [N, in, out] -> [N, 1, out] -> [N, out] + projected = op.MatMul(op.Unsqueeze(x, [1]), weight) + projected = op.Squeeze(projected, [1]) + return op.Add(projected, bias) + + +class Cosmos3OmniTimesteps(nn.Module): + """Sinusoidal timestep projection (diffusers ``Timesteps``). + + Upstream instantiates ``Timesteps(num_channels=256, flip_sin_to_cos=True, + downscale_freq_shift=0)``, i.e. ``concat(cos(t * inv_freq), sin(t * inv_freq))``. + ``inv_freq`` is a derived constant, not a checkpoint weight, and is kept in + float32 (``_keep_float32``) because the timestep path must stay fp32. + + Args: + num_channels: Output width (must be even). + max_period: Sinusoid base period. + downscale_freq_shift: Upstream's exponent shift (``0`` for Cosmos3). + """ + + def __init__( + self, + num_channels: int, + max_period: float = 10_000.0, + downscale_freq_shift: float = 0.0, + ): + super().__init__() + half_dim = num_channels // 2 + exponent = -np.log(max_period) * np.arange(half_dim, dtype=np.float32) + exponent = exponent / (half_dim - downscale_freq_shift) + self.inv_freq = nn.Parameter( + [half_dim], + name="inv_freq", + data=ir.tensor(np.exp(exponent).astype(np.float32)), + ) + self.inv_freq._keep_float32 = True + + def forward(self, op: OpBuilder, timesteps: ir.Value) -> ir.Value: + """Project ``[N]`` float32 timesteps to ``[N, num_channels]`` float32.""" + # [N, 1] * [half_dim] -> [N, half_dim] + freqs = op.Mul(op.Unsqueeze(timesteps, [-1]), self.inv_freq) + # flip_sin_to_cos=True puts the cosine half first. + return op.Concat(op.Cos(freqs), op.Sin(freqs), axis=-1) + + +class Cosmos3OmniRotaryEmbedding(nn.Module): + """Interleaved 3-axis mRoPE over a packed joint sequence. + + Reproduces upstream ``Cosmos3VLTextRotaryEmbedding``. Frequencies for the + three position axes (T, H, W) are computed independently in **float32** + (upstream disables autocast here precisely because bf16 cannot represent + consecutive integers past 256), then merged channel-wise using the + *interleaved* layout:: + + channel c takes the H axis if c in range(1, rope_axes_dim[1] * 3, 3) + channel c takes the W axis if c in range(2, rope_axes_dim[2] * 3, 3) + channel c takes the T axis otherwise + + For ``head_dim=128`` / ``rope_axes_dim=(24, 20, 20)`` this yields the + ``[T H W T H W ... T T]`` pattern with exactly 24/20/20 channels. + + The merged ``[seq, head_dim // 2]`` frequency table is turned into cos/sin + and cast to the model dtype; the ONNX ``RotaryEmbedding`` op consumes the + half-width cache directly (upstream's ``cat((freqs, freqs), -1)`` is the + same rotation in the ``rotate_half`` convention). + + Args: + config: Cosmos3-Omni generator configuration. + """ + + def __init__(self, config: Cosmos3OmniGeneratorConfig): + super().__init__() + self._dtype = config.dtype + rotary_dim = config.rotary_dim + head_dim = config.head_dim + + inv_freq = 1.0 / ( + config.rope_theta ** (np.arange(0, head_dim, 2, dtype=np.float32) / head_dim) + ) + self.inv_freq = nn.Parameter( + [rotary_dim], name="inv_freq", data=ir.tensor(inv_freq.astype(np.float32)) + ) + # Rotary frequencies must stay fp32 even for a bf16 model. + self.inv_freq._keep_float32 = True + + h_mask = np.zeros(rotary_dim, dtype=np.bool_) + w_mask = np.zeros(rotary_dim, dtype=np.bool_) + for channel in range(1, config.rope_axes_dim[1] * 3, 3): + if channel < rotary_dim: + h_mask[channel] = True + for channel in range(2, config.rope_axes_dim[2] * 3, 3): + if channel < rotary_dim: + w_mask[channel] = True + self.h_mask = nn.Parameter([rotary_dim], name="h_mask", data=ir.tensor(h_mask)) + self.w_mask = nn.Parameter([rotary_dim], name="w_mask", data=ir.tensor(w_mask)) + + def forward(self, op: OpBuilder, position_ids: ir.Value) -> tuple[ir.Value, ir.Value]: + """Compute cos/sin for every joint-sequence row. + + Args: + op: ONNX op builder. + position_ids: ``[3, sequence_length]`` int64 (T, H, W) positions. + + Returns: + ``(cos, sin)``, each ``[1, sequence_length, head_dim // 2]`` in the + model dtype — the layout the ONNX ``RotaryEmbedding`` op expects + when no ``position_ids`` input is supplied. + """ + # [3, seq] -> [3, seq, 1] * [rotary_dim] -> [3, seq, rotary_dim], fp32. + positions = op.Cast(position_ids, to=ir.DataType.FLOAT) + freqs = op.Mul(op.Unsqueeze(positions, [-1]), self.inv_freq) + + freqs_t = op.Squeeze(op.Gather(freqs, [0], axis=0), [0]) # (seq, rotary_dim) + freqs_h = op.Squeeze(op.Gather(freqs, [1], axis=0), [0]) + freqs_w = op.Squeeze(op.Gather(freqs, [2], axis=0), [0]) + + # Interleaved channel selection: H and W overwrite disjoint channel + # subsets of the T table, exactly like upstream's in-place slice write. + merged = op.Where(self.h_mask, freqs_h, freqs_t) + merged = op.Where(self.w_mask, freqs_w, merged) + + cos = op.Unsqueeze(op.Cos(merged), [0]) # (1, seq, rotary_dim) + sin = op.Unsqueeze(op.Sin(merged), [0]) + if self._dtype != ir.DataType.FLOAT: + cos = op.Cast(cos, to=self._dtype) + sin = op.Cast(sin, to=self._dtype) + return cos, sin + + +class Cosmos3OmniMoTAttention(nn.Module): + """Dual-pathway packed MoT attention (upstream ``Cosmos3PackedMoTAttention``). + + Holds two complete projection sets. The understanding set + (``to_q``/``to_k``/``to_v``/``to_out`` + ``norm_q``/``norm_k``) runs a + causal self-attention over the understanding tokens only. The generation + set (``add_q_proj``/``add_k_proj``/``add_v_proj``/``to_add_out`` + + ``norm_added_q``/``norm_added_k``) runs a *non-causal* attention whose + keys/values are ``concat(understanding, generation)``. + + Args: + config: Cosmos3-Omni generator configuration. + """ + + def __init__(self, config: Cosmos3OmniGeneratorConfig): + super().__init__() + self.config = config + self._num_heads = config.num_attention_heads + self._num_kv_heads = config.num_key_value_heads + self._head_dim = config.head_dim + self._scale = float(config.head_dim**-0.5) + hidden_size = config.hidden_size + q_size = config.attention_out_size + kv_size = config.key_value_size + bias = config.attention_bias + eps = config.rms_norm_eps + + # Understanding ("und") pathway. + self.to_q = Linear(hidden_size, q_size, bias=bias) + self.to_k = Linear(hidden_size, kv_size, bias=bias) + self.to_v = Linear(hidden_size, kv_size, bias=bias) + self.to_out = Linear(q_size, hidden_size, bias=bias) + # Per-head QK RMSNorm over head_dim only (no reshape needed after it). + self.norm_q = RMSNorm(config.head_dim, eps=eps) if config.qk_norm_for_text else None + self.norm_k = RMSNorm(config.head_dim, eps=eps) if config.qk_norm_for_text else None + # Separate norm for the understanding keys consumed by the generator. + self.k_norm_und_for_gen = ( + RMSNorm(config.head_dim, eps=eps) if config.has_und_k_norm_for_gen else None + ) + + # Generation ("gen") pathway. + self.add_q_proj = Linear(hidden_size, q_size, bias=bias) + self.add_k_proj = Linear(hidden_size, kv_size, bias=bias) + self.add_v_proj = Linear(hidden_size, kv_size, bias=bias) + self.to_add_out = Linear(q_size, hidden_size, bias=bias) + self.norm_added_q = RMSNorm(config.head_dim, eps=eps) + self.norm_added_k = RMSNorm(config.head_dim, eps=eps) + + def _norm_per_head( + self, op: OpBuilder, projected: ir.Value, norm: RMSNorm | None, num_heads: int + ) -> ir.Value: + """RMS-normalize each head slice of a flat ``[N, num_heads * head_dim]``.""" + if norm is None: + return projected + # (N, num_heads * head_dim) -> (N, num_heads, head_dim) + heads = op.Reshape(projected, _const_ints(op, [-1, num_heads, self._head_dim])) + heads = norm(op, heads) + # back to (N, num_heads * head_dim) + return op.Reshape(heads, _const_ints(op, [-1, num_heads * self._head_dim])) + + def forward( + self, + op: OpBuilder, + und_seq: ir.Value, + gen_seq: ir.Value, + rotary_emb: tuple[ir.Value, ir.Value, ir.Value, ir.Value], + ) -> tuple[ir.Value, ir.Value]: + """Run both attention pathways. + + Args: + op: ONNX op builder. + und_seq: ``[und_len, hidden_size]`` pre-normalized understanding tokens. + gen_seq: ``[gen_len, hidden_size]`` pre-normalized generation tokens. + rotary_emb: ``(cos_und, sin_und, cos_gen, sin_gen)``, each + ``[1, *, head_dim // 2]``. + + Returns: + ``(und_out, gen_out)``, each ``[*, hidden_size]``. + """ + cos_und, sin_und, cos_gen, sin_gen = rotary_emb + + # --- projections: (seq, hidden) -> (seq, heads * head_dim) --- + q_und = self._norm_per_head(op, self.to_q(op, und_seq), self.norm_q, self._num_heads) + k_und = self._norm_per_head( + op, self.to_k(op, und_seq), self.norm_k, self._num_kv_heads + ) + v_und = self.to_v(op, und_seq) + q_gen = self._norm_per_head( + op, self.add_q_proj(op, gen_seq), self.norm_added_q, self._num_heads + ) + k_gen = self._norm_per_head( + op, self.add_k_proj(op, gen_seq), self.norm_added_k, self._num_kv_heads + ) + v_gen = self.add_v_proj(op, gen_seq) + + # The generator may consume a separately normalized copy of the + # understanding keys (upstream ``k_norm_und_for_gen``). + k_und_for_gen = self._norm_per_head( + op, k_und, self.k_norm_und_for_gen, self._num_kv_heads + ) + + # --- rotary: the op wants (batch, seq, heads * head_dim) --- + q_und = op.Unsqueeze(q_und, [0]) + k_und_b = op.Unsqueeze(k_und, [0]) + v_und = op.Unsqueeze(v_und, [0]) + q_gen = op.Unsqueeze(q_gen, [0]) + k_gen = op.Unsqueeze(k_gen, [0]) + v_gen = op.Unsqueeze(v_gen, [0]) + + und_pos = (cos_und, sin_und) + gen_pos = (cos_gen, sin_gen) + q_und = _apply_rope(op, q_und, *und_pos, self._num_heads) + k_und_roped = _apply_rope(op, k_und_b, *und_pos, self._num_kv_heads) + if self.k_norm_und_for_gen is None: + # Same tensor upstream — reuse the roped keys instead of recomputing. + k_und_for_gen_roped = k_und_roped + else: + k_und_for_gen_roped = _apply_rope( + op, op.Unsqueeze(k_und_for_gen, [0]), *und_pos, self._num_kv_heads + ) + q_gen = _apply_rope(op, q_gen, *gen_pos, self._num_heads) + k_gen = _apply_rope(op, k_gen, *gen_pos, self._num_kv_heads) + + # --- causal understanding self-attention (GQA) --- + und_attn = op.Attention( + q_und, + k_und_roped, + v_und, + q_num_heads=self._num_heads, + kv_num_heads=self._num_kv_heads, + is_causal=1, + scale=self._scale, + ) + + # --- non-causal generation attention over [und ; gen] K/V (GQA) --- + all_k = op.Concat(k_und_for_gen_roped, k_gen, axis=1) + all_v = op.Concat(v_und, v_gen, axis=1) + gen_attn = op.Attention( + q_gen, + all_k, + all_v, + q_num_heads=self._num_heads, + kv_num_heads=self._num_kv_heads, + is_causal=0, + scale=self._scale, + ) + + # (1, seq, heads * head_dim) -> (seq, hidden_size) + und_out = self.to_out(op, op.Squeeze(und_attn, [0])) + gen_out = self.to_add_out(op, op.Squeeze(gen_attn, [0])) + return und_out, gen_out + + +class Cosmos3OmniMoTDecoderLayer(nn.Module): + """One MoT decoder layer holding both expert weight sets. + + Args: + config: Cosmos3-Omni generator configuration. + """ + + def __init__(self, config: Cosmos3OmniGeneratorConfig): + super().__init__() + eps = config.rms_norm_eps + self.self_attn = Cosmos3OmniMoTAttention(config) + self.mlp = _build_mlp(config) + self.mlp_moe_gen = _build_mlp(config) + self.input_layernorm = RMSNorm(config.hidden_size, eps=eps) + self.input_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=eps) + self.post_attention_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=eps) + + def forward( + self, + op: OpBuilder, + und_seq: ir.Value, + gen_seq: ir.Value, + rotary_emb: tuple[ir.Value, ir.Value, ir.Value, ir.Value], + ) -> tuple[ir.Value, ir.Value]: + """Apply pre-norm attention and per-expert feed-forward with residuals. + + Args: + op: ONNX op builder. + und_seq: ``[und_len, hidden_size]`` understanding hidden states. + gen_seq: ``[gen_len, hidden_size]`` generation hidden states. + rotary_emb: ``(cos_und, sin_und, cos_gen, sin_gen)``. + + Returns: + Updated ``(und_seq, gen_seq)``. + """ + und_attn, gen_attn = self.self_attn( + op, + self.input_layernorm(op, und_seq), + self.input_layernorm_moe_gen(op, gen_seq), + rotary_emb, + ) + und_residual = op.Add(und_seq, und_attn) + gen_residual = op.Add(gen_seq, gen_attn) + + und_mlp = self.mlp(op, self.post_attention_layernorm(op, und_residual)) + gen_mlp = self.mlp_moe_gen(op, self.post_attention_layernorm_moe_gen(op, gen_residual)) + return op.Add(und_residual, und_mlp), op.Add(gen_residual, gen_mlp) + + +def _build_mlp(config: Cosmos3OmniGeneratorConfig) -> nn.Module: + """Create the per-expert feed-forward matching ``hidden_act``. + + ``silu`` -> gated SwiGLU (``gate_proj``/``up_proj``/``down_proj``); + ``relu2`` -> non-gated squared ReLU (``up_proj``/``down_proj``). Both + parameter layouts match the published checkpoint names. + """ + if config.is_gated_mlp: + return GatedMLP( + config.hidden_size, config.intermediate_size, activation="silu", bias=False + ) + return FCMLP(config.hidden_size, config.intermediate_size, activation="relu2", bias=False) + + +# Keys present in the published ``transformer/`` checkpoint that carry no +# neural computation in this graph. ``lm_head`` is constructed by the +# upstream module but never called in ``Cosmos3OmniTransformer.forward`` — the +# understanding logits come from the separately exported Reasoner +# (:mod:`mobius.models.cosmos3_omni`), so it is dead weight here. +_UNUSED_PUBLISHED_KEYS: frozenset[str] = frozenset({"lm_head.weight"}) + +# Non-persistent buffers that some export paths materialize into the state +# dict. They are recomputed as graph constants, so a checkpoint copy is +# redundant rather than unexpected. +_RECOMPUTED_BUFFER_SUFFIXES: tuple[str, ...] = (".inv_freq", ".h_mask", ".w_mask") + +# Reasoner (understanding-tower) vision-encoder prefixes. A *unified* +# Cosmos3-Omni checkpoint carries them alongside the transformer; they belong +# to the separately exported vision tower, never to this graph. +_REASONER_VISION_PREFIXES: tuple[str, ...] = ( + "visual.", + "projector.", + "blocks.", + "patch_embed.", + "merger.", + "deepstack_merger_list.", + "pos_embed", +) + +# Cosmos3-Edge Policy checkpoints carry a tiny framework sidecar with duplicate +# generator-facing key-norm tensors under nested training-framework paths. +_EDGE_FRAMEWORK_K_NORM_PREFIXES: tuple[str, ...] = ( + "net.language_model.model.layers.", + "layers.layers.", +) + +# Accepted container prefixes on a flat published state dict. +_STRIPPABLE_PREFIXES: tuple[str, ...] = ("transformer.", "model.") + + +class Cosmos3OmniGeneratorModel(nn.Module): + """Unified Cosmos3-Omni MoT transformer — one denoising step. + + Builds the whole neural graph: token embedding, per-modality input + projections with timestep conditioning, interleaved 3-axis mRoPE, the + dual-expert layer stack, the two final norms, and the per-modality output + projections. + + Parameter names match the published flat ``transformer/`` checkpoint + one-for-one (``embed_tokens.weight``, ``layers.N.self_attn.to_q.weight``, + ``layers.N.mlp_moe_gen.down_proj.weight``, ``norm_moe_gen.weight``, + ``proj_in.bias``, ``time_embedder.linear_1.weight``, + ``audio_proj_in.weight``, ``action_proj_in.fc.weight``, ...), so + :meth:`preprocess_weights` is a validation + normalization pass rather + than a rename table. + + Args: + config: Cosmos3-Omni generator configuration. + """ + + default_task: str = "cosmos3-omni-generator" + category: str = "Multimodal" + + def __init__(self, config: Cosmos3OmniGeneratorConfig): + super().__init__() + config.validate() + self.config = config + self._dtype = config.dtype + + # Shared backbone. + self.embed_tokens = Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [Cosmos3OmniMoTDecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Cosmos3OmniRotaryEmbedding(config) + + # Vision (diffusion) head + timestep conditioning. + self.proj_in = Linear(config.patch_latent_dim, config.hidden_size, bias=True) + self.proj_out = Linear(config.hidden_size, config.patch_latent_dim, bias=True) + self.time_proj = Cosmos3OmniTimesteps(config.time_proj_channels) + self.time_embedder = TimestepEmbedding(config.time_proj_channels, config.hidden_size) + # Upstream keeps ``time_embedder`` in fp32 (``_keep_in_fp32_modules``); + # the sinusoid -> MLP path must not run in bf16. + for parameter in self.time_embedder.parameters(): + parameter._keep_float32 = True + + # Optional Sound head. + if config.sound_gen: + assert config.sound_dim is not None # guaranteed by config.validate() + self.audio_proj_in = Linear(config.sound_dim, config.hidden_size, bias=True) + self.audio_proj_out = Linear(config.hidden_size, config.sound_dim, bias=True) + self.audio_modality_embed = nn.Parameter( + [config.hidden_size], name="audio_modality_embed" + ) + + # Optional Action head (per-embodiment-domain weights). + if config.action_gen: + assert config.action_dim is not None # guaranteed by config.validate() + self.action_proj_in = Cosmos3OmniDomainAwareLinear( + config.action_dim, config.hidden_size, config.num_embodiment_domains + ) + self.action_proj_out = Cosmos3OmniDomainAwareLinear( + config.hidden_size, config.action_dim, config.num_embodiment_domains + ) + self.action_modality_embed = nn.Parameter( + [config.hidden_size], name="action_modality_embed" + ) + + # ------------------------------------------------------------------ + # Graph construction helpers + # ------------------------------------------------------------------ + + def _timestep_embedding(self, op: OpBuilder, timesteps: ir.Value) -> ir.Value: + """Embed diffusion timesteps, keeping the whole path in float32. + + Args: + op: ONNX op builder. + timesteps: ``[N]`` **float32** timesteps (pre-``timestep_scale``). + + Returns: + ``[N, hidden_size]`` embeddings cast to the model dtype — the cast + happens only at the boundary where the result is added to the + model-dtype token stream. + """ + scaled = op.Mul(timesteps, float(self.config.timestep_scale)) + projected = self.time_proj(op, scaled) + embedded = self.time_embedder(op, projected) + if self._dtype != ir.DataType.FLOAT: + embedded = op.Cast(embedded, to=self._dtype) + return embedded + + def _add_timestep_embeds( + self, + op: OpBuilder, + tokens: ir.Value, + timesteps: ir.Value, + token_indexes: ir.Value, + ) -> ir.Value: + """Scatter-add timestep embeddings onto the noisy rows of a token block. + + Tensorizes upstream ``_apply_timestep_embeds_to_noisy_tokens``: the + host precomputes the flattened noisy row offsets (which upstream + derives from ``noisy_frame_indexes`` + ``token_shapes``) and passes + them as ``token_indexes``. An empty ``token_indexes`` is a no-op, + matching upstream's ``if action_mse_loss_indexes.numel() > 0`` guard. + + Args: + op: ONNX op builder. + tokens: ``[N_tokens, hidden_size]`` projected modality tokens. + timesteps: ``[N_noisy]`` float32 timesteps. + token_indexes: ``[N_noisy]`` int64 rows of *tokens* to add into. + + Returns: + ``[N_tokens, hidden_size]`` tokens with timestep conditioning added. + """ + embeds = self._timestep_embedding(op, timesteps) + return op.ScatterND(tokens, _row_indices(op, token_indexes), embeds, reduction="add") + + def _scatter_into_joint( + self, + op: OpBuilder, + hidden_states: ir.Value, + tokens: ir.Value, + sequence_indexes: ir.Value, + ) -> ir.Value: + """Write a packed modality block into its joint-sequence rows.""" + return op.ScatterND(hidden_states, _row_indices(op, sequence_indexes), tokens) + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def forward( + self, + op: OpBuilder, + *, + input_ids: ir.Value, + text_indexes: ir.Value, + position_ids: ir.Value, + und_len: ir.Value, + vision_tokens: ir.Value, + vision_sequence_indexes: ir.Value, + vision_timesteps: ir.Value, + vision_timestep_token_indexes: ir.Value, + vision_mse_loss_indexes: ir.Value, + sound_tokens: ir.Value | None = None, + sound_sequence_indexes: ir.Value | None = None, + sound_timesteps: ir.Value | None = None, + sound_timestep_token_indexes: ir.Value | None = None, + sound_mse_loss_indexes: ir.Value | None = None, + action_tokens: ir.Value | None = None, + action_domain_ids: ir.Value | None = None, + action_sequence_indexes: ir.Value | None = None, + action_timesteps: ir.Value | None = None, + action_timestep_token_indexes: ir.Value | None = None, + action_mse_loss_indexes: ir.Value | None = None, + action_pred_domain_ids: ir.Value | None = None, + ) -> tuple[ir.Value, ir.Value | None, ir.Value | None]: + """Run one unified denoising step over the packed joint sequence. + + Args: + op: ONNX op builder. + input_ids: ``[num_text_tokens]`` int64 text token ids. + text_indexes: ``[num_text_tokens]`` int64 joint-sequence rows for + the text tokens. + position_ids: ``[3, sequence_length]`` int64 mRoPE (T, H, W) + positions covering the *whole* joint sequence. Its second + dimension defines ``sequence_length``. + und_len: ``[1]`` int64 — number of leading joint-sequence rows + routed through the understanding expert. The remainder goes + through the generation expert. + vision_tokens: ``[num_vision_tokens, patch_latent_dim]`` packed, + host-patchified vision latents in the model dtype. + vision_sequence_indexes: ``[num_vision_tokens]`` int64 joint rows. + vision_timesteps: ``[num_vision_noisy_tokens]`` float32 timesteps. + vision_timestep_token_indexes: ``[num_vision_noisy_tokens]`` int64 + rows of ``vision_tokens`` receiving each timestep embedding. + vision_mse_loss_indexes: ``[num_vision_noisy_tokens]`` int64 joint + rows to decode into velocity predictions. + sound_tokens: ``[num_sound_tokens, sound_dim]`` packed sound latents. + sound_sequence_indexes: ``[num_sound_tokens]`` int64 joint rows. + sound_timesteps: ``[num_sound_noisy_tokens]`` float32 timesteps. + sound_timestep_token_indexes: ``[num_sound_noisy_tokens]`` int64 + rows of ``sound_tokens``. + sound_mse_loss_indexes: ``[num_sound_noisy_tokens]`` int64 joint rows. + action_tokens: ``[num_action_tokens, action_dim]`` packed actions. + action_domain_ids: ``[num_action_tokens]`` int64 embodiment domain + per action token (input projection). + action_sequence_indexes: ``[num_action_tokens]`` int64 joint rows. + action_timesteps: ``[num_action_noisy_tokens]`` float32 timesteps. + action_timestep_token_indexes: ``[num_action_noisy_tokens]`` int64 + rows of ``action_tokens``. + action_mse_loss_indexes: ``[num_action_noisy_tokens]`` int64 joint rows. + action_pred_domain_ids: ``[num_action_noisy_tokens]`` int64 + embodiment domain per predicted action token (output projection). + + Returns: + ``(vision_pred, sound_pred, action_pred)`` where ``vision_pred`` is + ``[num_vision_noisy_tokens, patch_latent_dim]`` and the optional + entries are ``[num_*_noisy_tokens, sound_dim | action_dim]`` or + ``None`` when the corresponding head is disabled in the config. + + Raises: + ValueError: If a configured head is missing any of its inputs. + """ + self._require_head_inputs( + "sound", + self.config.sound_gen, + { + "sound_tokens": sound_tokens, + "sound_sequence_indexes": sound_sequence_indexes, + "sound_timesteps": sound_timesteps, + "sound_timestep_token_indexes": sound_timestep_token_indexes, + "sound_mse_loss_indexes": sound_mse_loss_indexes, + }, + ) + self._require_head_inputs( + "action", + self.config.action_gen, + { + "action_tokens": action_tokens, + "action_domain_ids": action_domain_ids, + "action_sequence_indexes": action_sequence_indexes, + "action_timesteps": action_timesteps, + "action_timestep_token_indexes": action_timestep_token_indexes, + "action_mse_loss_indexes": action_mse_loss_indexes, + "action_pred_domain_ids": action_pred_domain_ids, + }, + ) + + # --- joint hidden-state buffer ------------------------------------ + # sequence_length is read off position_ids so no host int is needed. + sequence_length = op.Shape(position_ids, start=1, end=2) # (1,) + buffer_shape = op.Concat( + sequence_length, _const_ints(op, [self.config.hidden_size]), axis=0 + ) + hidden_states = op.ConstantOfShape( + buffer_shape, value=ir.tensor(np.zeros(1, dtype=np.float32)) + ) + if self._dtype != ir.DataType.FLOAT: + hidden_states = op.Cast(hidden_states, to=self._dtype) + + # --- text tokens --------------------------------------------------- + # (num_text_tokens, hidden_size) written at their joint rows. + text_embeds = self.embed_tokens(op, input_ids) + hidden_states = self._scatter_into_joint(op, hidden_states, text_embeds, text_indexes) + + # --- vision latents ------------------------------------------------ + # (num_vision_tokens, patch_latent_dim) -> (num_vision_tokens, hidden_size) + vision_hidden = self.proj_in(op, vision_tokens) + vision_hidden = self._add_timestep_embeds( + op, vision_hidden, vision_timesteps, vision_timestep_token_indexes + ) + hidden_states = self._scatter_into_joint( + op, hidden_states, vision_hidden, vision_sequence_indexes + ) + + # --- sound latents (optional) -------------------------------------- + if self.config.sound_gen: + # (num_sound_tokens, sound_dim) -> (num_sound_tokens, hidden_size) + sound_hidden = self.audio_proj_in(op, sound_tokens) + sound_hidden = op.Add(sound_hidden, self.audio_modality_embed) + sound_hidden = self._add_timestep_embeds( + op, sound_hidden, sound_timesteps, sound_timestep_token_indexes + ) + hidden_states = self._scatter_into_joint( + op, hidden_states, sound_hidden, sound_sequence_indexes + ) + + # --- action latents (optional) ------------------------------------- + if self.config.action_gen: + # (num_action_tokens, action_dim) -> (num_action_tokens, hidden_size) + # via the per-embodiment-domain weight table. + action_hidden = self.action_proj_in(op, action_tokens, action_domain_ids) + action_hidden = op.Add(action_hidden, self.action_modality_embed) + action_hidden = self._add_timestep_embeds( + op, action_hidden, action_timesteps, action_timestep_token_indexes + ) + hidden_states = self._scatter_into_joint( + op, hidden_states, action_hidden, action_sequence_indexes + ) + + # --- mRoPE, then split the joint sequence into the two experts ----- + cos, sin = self.rotary_emb(op, position_ids) # (1, seq, head_dim // 2) + zero = _const_ints(op, [0]) + end = _const_ints(op, [INT64_MAX]) + cos_und = op.Slice(cos, zero, und_len, _const_ints(op, [1])) + sin_und = op.Slice(sin, zero, und_len, _const_ints(op, [1])) + cos_gen = op.Slice(cos, und_len, end, _const_ints(op, [1])) + sin_gen = op.Slice(sin, und_len, end, _const_ints(op, [1])) + rotary_emb = (cos_und, sin_und, cos_gen, sin_gen) + + und_seq = op.Slice(hidden_states, zero, und_len, zero) # (und_len, hidden) + gen_seq = op.Slice(hidden_states, und_len, end, zero) # (gen_len, hidden) + + # --- MoT layer stack ------------------------------------------------ + for layer in self.layers: + und_seq, gen_seq = layer(op, und_seq, gen_seq, rotary_emb) + und_out = self.norm(op, und_seq) + gen_out = self.norm_moe_gen(op, gen_seq) + # Re-joined so the *_mse_loss_indexes address joint-sequence rows. + last_hidden_state = op.Concat(und_out, gen_out, axis=0) + + # --- per-modality velocity predictions ------------------------------ + vision_pred = self.proj_out( + op, op.Gather(last_hidden_state, vision_mse_loss_indexes, axis=0) + ) + + sound_pred = None + if self.config.sound_gen: + sound_pred = self.audio_proj_out( + op, op.Gather(last_hidden_state, sound_mse_loss_indexes, axis=0) + ) + + action_pred = None + if self.config.action_gen: + action_pred = self.action_proj_out( + op, + op.Gather(last_hidden_state, action_mse_loss_indexes, axis=0), + action_pred_domain_ids, + ) + + return vision_pred, sound_pred, action_pred + + @staticmethod + def _require_head_inputs( + head: str, enabled: bool, inputs: dict[str, ir.Value | None] + ) -> None: + """Fail loudly instead of silently dropping a configured head. + + Args: + head: Head name (for the error message). + enabled: Whether the head is enabled in the config. + inputs: Mapping of input name to the supplied value. + + Raises: + ValueError: If the head is enabled and any input is missing, or if + the head is disabled but inputs were supplied anyway. + """ + missing = sorted(name for name, value in inputs.items() if value is None) + if enabled and missing: + raise ValueError( + f"Cosmos3-Omni {head}_gen=True requires all {head} inputs; missing: {missing}" + ) + if not enabled and len(missing) != len(inputs): + supplied = sorted(name for name, value in inputs.items() if value is not None) + raise ValueError( + f"Cosmos3-Omni {head}_gen=False but {head} inputs were supplied: {supplied}" + ) + + # ------------------------------------------------------------------ + # Weights + # ------------------------------------------------------------------ + + def expected_checkpoint_keys(self) -> set[str]: + """Return the checkpoint keys this graph consumes. + + Derived from the module tree (so it tracks the ``sound_gen`` / + ``action_gen`` / ``hidden_act`` / ``qk_norm_for_text`` gating + automatically), minus the parameters that are recomputed graph + constants rather than checkpoint weights (``rotary_emb.inv_freq``, + ``rotary_emb.h_mask``, ``rotary_emb.w_mask``, ``time_proj.inv_freq``). + """ + return { + name + for name, _ in self.named_parameters() + if not name.endswith(_RECOMPUTED_BUFFER_SUFFIXES) + } + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Normalize and validate a flat published ``transformer/`` state dict. + + The module tree already mirrors the published names, so this performs + no renames beyond two structural normalizations: + + 1. strips an optional ``transformer.`` / ``model.`` container prefix; + 2. collapses ``self_attn.to_out.0.*`` (checkpoints where the output + projection was wrapped in an ``nn.Sequential``) onto + ``self_attn.to_out.*``. + + Keys are then partitioned explicitly — no broad pattern filtering: + + * ``lm_head.weight`` is dropped because upstream constructs it but + never calls it in ``Cosmos3OmniTransformer.forward``; understanding + logits come from the separately exported Reasoner; + * Reasoner vision-tower keys (present only in a *unified* checkpoint) + are dropped because they belong to that separate export; + * recomputed buffers (``*.inv_freq`` / ``*.h_mask`` / ``*.w_mask``) + are dropped because the graph derives them as constants. + + Anything else that does not name a parameter of this graph — and any + graph parameter left without a weight — raises. + + Args: + state_dict: Flat published transformer weights. + + Returns: + A state dict keyed by this graph's initializer names. + + Raises: + ValueError: If the checkpoint carries unexpected keys, or does not + cover every parameter this graph needs (both indicate an + architecture mismatch against the config). + """ + return self._preprocess_weights(state_dict, require_complete=True) + + def preprocess_weight_shard( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Normalize and validate one shard of the unified checkpoint. + + Unlike :meth:`preprocess_weights`, this does not require the shard to + contain every graph parameter. Composite exporters call it for each + shard, accumulate the returned names, then compare that union with + :meth:`expected_checkpoint_keys`. + """ + return self._preprocess_weights(state_dict, require_complete=False) + + def _preprocess_weights( + self, + state_dict: dict[str, torch.Tensor], + *, + require_complete: bool, + ) -> dict[str, torch.Tensor]: + expected = self.expected_checkpoint_keys() + renamed: dict[str, torch.Tensor] = {} + unexpected: list[str] = [] + + for raw_key, value in state_dict.items(): + key = raw_key + for prefix in _STRIPPABLE_PREFIXES: + if key.startswith(prefix): + key = key[len(prefix) :] + break + key = key.replace(".self_attn.to_out.0.", ".self_attn.to_out.") + + if key in _UNUSED_PUBLISHED_KEYS: + continue + if key.endswith(_RECOMPUTED_BUFFER_SUFFIXES): + continue + if key.startswith(_REASONER_VISION_PREFIXES): + continue + if key.startswith(_EDGE_FRAMEWORK_K_NORM_PREFIXES) and key.endswith( + ".self_attn.k_norm_und_for_gen.weight" + ): + continue + if key not in expected: + unexpected.append(raw_key) + continue + renamed[key] = value + + if unexpected: + raise ValueError( + "Unexpected Cosmos3-Omni transformer weights (architecture mismatch " + f"against the config): {sorted(unexpected)[:16]}" + ) + missing = sorted(expected - renamed.keys()) + if require_complete and missing: + raise ValueError( + "Cosmos3-Omni transformer checkpoint is missing weights required by " + f"the configured architecture: {missing[:16]}" + ) + return renamed + + +__all__ = [ + "Cosmos3OmniDomainAwareLinear", + "Cosmos3OmniGeneratorModel", + "Cosmos3OmniMoTAttention", + "Cosmos3OmniMoTDecoderLayer", + "Cosmos3OmniRotaryEmbedding", + "Cosmos3OmniTimesteps", +] diff --git a/src/mobius/models/cosmos3_omni_generator_test.py b/src/mobius/models/cosmos3_omni_generator_test.py new file mode 100644 index 000000000..b2c6716cc --- /dev/null +++ b/src/mobius/models/cosmos3_omni_generator_test.py @@ -0,0 +1,1377 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the unified Cosmos3-Omni MoT transformer (Reasoner + Generator). + +Coverage: + +* config parsing of the public ``nvidia/Cosmos3-Nano`` ``transformer/config.json`` + and every shape-relationship validation; +* exact graph I/O contract for the generator-only and sound+action builds; +* initializer-name coverage for all core branches (both experts, both MLP + variants, both QK-norm variants, the vision/sound/action heads); +* interleaved 3-axis mRoPE channel layout against upstream's index arithmetic; +* ``DomainAwareLinear`` graph semantics and numerics vs a PyTorch reference; +* ``preprocess_weights`` accept/drop/reject behaviour; +* full numerical parity of the packed attention setup (dual-pathway causal + + non-causal attention, mRoPE, timestep conditioning, MoT experts) against a + PyTorch transcription of upstream ``Cosmos3OmniTransformer.forward``. +""" + +from __future__ import annotations + +import dataclasses + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +import torch +import torch.nn.functional as functional +from onnxscript import nn + +from mobius import build_from_module +from mobius._configs._cosmos3_omni_generator import Cosmos3OmniGeneratorConfig +from mobius._model_package import ModelPackage +from mobius.models.cosmos3_omni_generator import ( + Cosmos3OmniDomainAwareLinear, + Cosmos3OmniGeneratorModel, +) +from mobius.tasks._base import ModelTask, _make_graph, _make_model +from mobius.tasks._cosmos3_omni_generator import ( + Cosmos3OmniGeneratorTask, + expected_input_names, + expected_output_names, +) + +# Verbatim ``nvidia/Cosmos3-Nano`` ``transformer/config.json`` (public). +PUBLISHED_CONFIG: dict = { + "_class_name": "Cosmos3OmniTransformer", + "_diffusers_version": "0.37.1", + "action_dim": 64, + "action_gen": True, + "attention_bias": False, + "attention_dropout": 0.0, + "base_fps": 24, + "dtype": "bfloat16", + "enable_fps_modulation": True, + "freeze_und": False, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 12288, + "joint_attn_implementation": "two_way", + "latent_channel": 48, + "latent_patch_size": 2, + "max_action_dim": 64, + "max_position_embeddings": 262144, + "model_type": "qwen3_vl_text", + "num_attention_heads": 32, + "num_embodiment_domains": 32, + "num_hidden_layers": 36, + "num_key_value_heads": 8, + "patch_latent_dim": 192, + "position_embedding_type": "unified_3d_mrope", + "qk_norm": False, + "qk_norm_for_diffusion": True, + "qk_norm_for_text": True, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "mrope_interleaved": True, + "mrope_section": [24, 20, 20], + "rope_type": "default", + }, + "rope_theta": 5000000, + "sound_dim": 64, + "sound_gen": True, + "sound_latent_fps": 25, + "temporal_compression_factor_sound": 1, + "timestep_scale": 0.001, + "unified_3d_mrope_reset_spatial_ids": True, + "unified_3d_mrope_temporal_modality_margin": 15000, + "use_cache": True, + "use_moe": True, + "video_temporal_causal": False, + "vocab_size": 151936, +} + + +def tiny_config(**overrides) -> Cosmos3OmniGeneratorConfig: + """A 2-layer graph-construction config (hidden=32, head_dim=8).""" + fields = { + "vocab_size": 64, + "hidden_size": 32, + "intermediate_size": 48, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "hidden_act": "silu", + "dtype": ir.DataType.FLOAT, + "rope_theta": 5000.0, + # sum == head_dim // 2 == 4 + "rope_axes_dim": (2, 1, 1), + "latent_channel": 3, + "latent_patch_size": 2, + "patch_latent_dim": 12, + "time_proj_channels": 8, + } + fields.update(overrides) + return Cosmos3OmniGeneratorConfig(**fields) + + +def full_config(**overrides) -> Cosmos3OmniGeneratorConfig: + """The tiny config with both optional heads enabled.""" + fields = { + "sound_gen": True, + "sound_dim": 6, + "action_gen": True, + "action_dim": 5, + "num_embodiment_domains": 3, + } + fields.update(overrides) + return tiny_config(**fields) + + +def build_package(config: Cosmos3OmniGeneratorConfig, **kwargs) -> ModelPackage: + """Build the generator package for *config*.""" + return build_from_module( + Cosmos3OmniGeneratorModel(config), + config, + task=Cosmos3OmniGeneratorTask(), + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +class TestConfig: + def test_from_diffusers_parses_published_config(self): + config = Cosmos3OmniGeneratorConfig.from_diffusers(PUBLISHED_CONFIG) + + assert config.hidden_size == 4096 + assert config.intermediate_size == 12288 + assert config.num_hidden_layers == 36 + assert config.num_attention_heads == 32 + assert config.num_key_value_heads == 8 + assert config.head_dim == 128 + assert config.vocab_size == 151936 + assert config.hidden_act == "silu" + assert config.dtype == ir.DataType.BFLOAT16 + assert config.rope_theta == pytest.approx(5_000_000.0) + assert config.rms_norm_eps == pytest.approx(1e-6) + assert config.attention_bias is False + assert config.attn_qkv_bias is False + assert config.attn_o_bias is False + # Vision head shape relationship: 48 * 2**2 == 192. + assert config.latent_channel == 48 + assert config.latent_patch_size == 2 + assert config.patch_latent_dim == 192 + assert config.timestep_scale == pytest.approx(0.001) + # Optional heads. + assert config.sound_gen is True + assert config.sound_dim == 64 + assert config.action_gen is True + assert config.action_dim == 64 + assert config.max_action_dim == 64 + assert config.num_embodiment_domains == 32 + # Host-side fields are parsed, not dropped. + assert config.base_fps == 24 + assert config.enable_fps_modulation is True + assert config.unified_3d_mrope_reset_spatial_ids is True + assert config.unified_3d_mrope_temporal_modality_margin == 15000 + assert config.max_position_embeddings == 262144 + assert config.sound_latent_fps == pytest.approx(25.0) + assert config.temporal_compression_factor_sound == 1 + + def test_from_diffusers_derives_rope_axes_from_mrope_section(self): + config = Cosmos3OmniGeneratorConfig.from_diffusers(PUBLISHED_CONFIG) + assert config.rope_axes_dim == (24, 20, 20) + # mRoPE channel budget must exactly fill the rotary dimension. + assert sum(config.rope_axes_dim) == config.rotary_dim == 64 + + def test_from_diffusers_defaults_rope_axes_without_rope_scaling(self): + raw = dict(PUBLISHED_CONFIG) + raw.pop("rope_scaling") + assert Cosmos3OmniGeneratorConfig.from_diffusers(raw).rope_axes_dim == (24, 20, 20) + + def test_from_diffusers_rejects_unknown_dtype(self): + raw = dict(PUBLISHED_CONFIG, dtype="int8") + with pytest.raises(ValueError, match="Unsupported Cosmos3-Omni dtype"): + Cosmos3OmniGeneratorConfig.from_diffusers(raw) + + def test_derived_properties(self): + config = Cosmos3OmniGeneratorConfig.from_diffusers(PUBLISHED_CONFIG) + assert config.rotary_dim == 64 + assert config.num_key_value_groups == 4 + assert config.attention_out_size == 32 * 128 + assert config.key_value_size == 8 * 128 + assert config.is_gated_mlp is True + # ``use_und_k_norm_for_gen`` is inert while the und pathway has QK norm. + assert config.has_und_k_norm_for_gen is False + + def test_has_und_k_norm_for_gen_requires_qk_norm_off(self): + assert tiny_config(use_und_k_norm_for_gen=True).has_und_k_norm_for_gen is False + assert ( + tiny_config(use_und_k_norm_for_gen=True, qk_norm_for_text=False) + ).has_und_k_norm_for_gen is True + + def test_relu2_backbone_is_not_gated(self): + assert tiny_config(hidden_act="relu2").is_gated_mlp is False + + @pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"hidden_size": 0}, "hidden_size must be a positive integer"), + ({"num_hidden_layers": -1}, "num_hidden_layers must be a positive integer"), + ({"num_attention_heads": 3}, "divisible by num_key_value_heads"), + ({"head_dim": 7, "rope_axes_dim": (1, 1, 1)}, "head_dim must be even"), + ({"hidden_act": "gelu"}, "hidden_act must be one of"), + ({"rms_norm_eps": 0.0}, "rms_norm_eps must be positive"), + ({"dtype": ir.DataType.INT8}, "dtype must be one of"), + ({"attention_dropout": 0.1}, "attention_dropout must be 0.0"), + ({"rope_axes_dim": (2, 2)}, "exactly 3"), + ({"rope_axes_dim": (2, 1, 0)}, "positive integers"), + ({"rope_axes_dim": (2, 2, 2)}, "must equal head_dim // 2"), + ({"rope_theta": 0.0}, "rope_theta must be positive"), + ({"max_position_embeddings": 0}, "max_position_embeddings must be positive"), + ({"latent_channel": 0}, "latent_channel must be positive"), + ({"latent_patch_size": 0}, "latent_patch_size must be positive"), + ({"patch_latent_dim": 13}, "latent_channel \\* latent_patch_size"), + ({"timestep_scale": 0.0}, "timestep_scale must be positive"), + ({"time_proj_channels": 7}, "positive even number"), + ({"sound_gen": True}, "sound_dim must be a positive integer"), + ({"sound_gen": True, "sound_dim": 4, "sound_latent_fps": 0}, "sound_latent_fps"), + ( + {"sound_gen": True, "sound_dim": 4, "temporal_compression_factor_sound": 0}, + "temporal_compression_factor_sound", + ), + ({"action_gen": True}, "action_dim must be a positive integer"), + ( + {"action_gen": True, "action_dim": 4, "num_embodiment_domains": 0}, + "num_embodiment_domains must be positive", + ), + ( + {"action_gen": True, "action_dim": 8, "max_action_dim": 4}, + "max_action_dim", + ), + ({"position_embedding_type": "rope"}, "unified_3d_mrope"), + ({"joint_attn_implementation": "one_way"}, "two_way"), + ({"use_moe": False}, "use_moe=False"), + ({"video_temporal_causal": True}, "video_temporal_causal=True"), + ({"qk_norm_for_diffusion": False}, "qk_norm_for_diffusion=False"), + ], + ) + def test_validate_rejects(self, overrides, match): + config = tiny_config(**overrides) + with pytest.raises(ValueError, match=match): + config.validate() + + def test_validate_accepts_published_and_tiny_configs(self): + Cosmos3OmniGeneratorConfig.from_diffusers(PUBLISHED_CONFIG).validate() + tiny_config().validate() + full_config().validate() + tiny_config(hidden_act="relu2").validate() + + +# --------------------------------------------------------------------------- +# Graph I/O contract +# --------------------------------------------------------------------------- + + +class TestGraphContract: + def test_generator_only_io(self): + config = tiny_config() + model = build_package(config)["model"] + + assert model.graph.name == "cosmos3_omni_generator" + assert [value.name for value in model.graph.inputs] == list( + expected_input_names(config) + ) + assert [value.name for value in model.graph.outputs] == ["vision_pred"] + + shapes = {value.name: list(value.shape) for value in model.graph.inputs} + dtypes = {value.name: value.dtype for value in model.graph.inputs} + assert shapes["input_ids"] == [ir.SymbolicDim("num_text_tokens")] + assert shapes["text_indexes"] == [ir.SymbolicDim("num_text_tokens")] + assert shapes["position_ids"] == [3, ir.SymbolicDim("sequence_length")] + assert shapes["und_len"] == [1] + assert shapes["vision_tokens"] == [ + ir.SymbolicDim("num_vision_tokens"), + config.patch_latent_dim, + ] + assert shapes["vision_mse_loss_indexes"] == [ir.SymbolicDim("num_vision_noisy_tokens")] + assert dtypes["vision_tokens"] == ir.DataType.FLOAT + # Timesteps are float32 regardless of model dtype. + assert dtypes["vision_timesteps"] == ir.DataType.FLOAT + for name in ( + "input_ids", + "text_indexes", + "position_ids", + "und_len", + "vision_sequence_indexes", + "vision_timestep_token_indexes", + "vision_mse_loss_indexes", + ): + assert dtypes[name] == ir.DataType.INT64, name + + output = model.graph.outputs[0] + assert list(output.shape) == [ + ir.SymbolicDim("num_vision_noisy_tokens"), + config.patch_latent_dim, + ] + assert output.dtype == ir.DataType.FLOAT + + def test_sound_and_action_io(self): + config = full_config() + model = build_package(config)["model"] + + assert [value.name for value in model.graph.inputs] == list( + expected_input_names(config) + ) + assert [value.name for value in model.graph.outputs] == list( + expected_output_names(config) + ) + shapes = {value.name: list(value.shape) for value in model.graph.inputs} + assert shapes["sound_tokens"] == [ + ir.SymbolicDim("num_sound_tokens"), + config.sound_dim, + ] + assert shapes["action_tokens"] == [ + ir.SymbolicDim("num_action_tokens"), + config.action_dim, + ] + assert shapes["action_domain_ids"] == [ir.SymbolicDim("num_action_tokens")] + assert shapes["action_pred_domain_ids"] == [ir.SymbolicDim("num_action_noisy_tokens")] + outputs = {value.name: list(value.shape) for value in model.graph.outputs} + assert outputs["sound_pred"] == [ + ir.SymbolicDim("num_sound_noisy_tokens"), + config.sound_dim, + ] + assert outputs["action_pred"] == [ + ir.SymbolicDim("num_action_noisy_tokens"), + config.action_dim, + ] + + @pytest.mark.parametrize( + ("overrides", "absent"), + [ + ({}, ("sound", "action")), + ({"sound_gen": True, "sound_dim": 6}, ("action",)), + ({"action_gen": True, "action_dim": 5, "num_embodiment_domains": 3}, ("sound",)), + ], + ) + def test_disabled_heads_declare_no_inputs_or_outputs(self, overrides, absent): + config = tiny_config(**overrides) + model = build_package(config)["model"] + names = [value.name for value in model.graph.inputs] + names += [value.name for value in model.graph.outputs] + for prefix in absent: + assert not [name for name in names if name.startswith(prefix)] + + def test_bfloat16_keeps_timestep_path_in_float32(self): + config = full_config(dtype=ir.DataType.BFLOAT16) + model = build_package(config)["model"] + initializers = model.graph.initializers + + # Model weights are bf16 ... + assert initializers["proj_in.weight"].dtype == ir.DataType.BFLOAT16 + assert initializers["layers.0.self_attn.to_q.weight"].dtype == ir.DataType.BFLOAT16 + # ... but the timestep MLP and the rotary frequencies stay fp32. + for name in ( + "time_embedder.linear_1.weight", + "time_embedder.linear_1.bias", + "time_embedder.linear_2.weight", + "time_embedder.linear_2.bias", + "time_proj.inv_freq", + "rotary_emb.inv_freq", + ): + assert initializers[name].dtype == ir.DataType.FLOAT, name + + dtypes = {value.name: value.dtype for value in model.graph.inputs} + assert dtypes["vision_tokens"] == ir.DataType.BFLOAT16 + assert dtypes["sound_tokens"] == ir.DataType.BFLOAT16 + assert dtypes["action_tokens"] == ir.DataType.BFLOAT16 + assert dtypes["vision_timesteps"] == ir.DataType.FLOAT + assert dtypes["sound_timesteps"] == ir.DataType.FLOAT + assert dtypes["action_timesteps"] == ir.DataType.FLOAT + assert all(value.dtype == ir.DataType.BFLOAT16 for value in model.graph.outputs) + + def test_task_rejects_wrong_module_output_contract(self): + class BadModule(nn.Module): + def forward(self, op, **kwargs): + return op.Identity(kwargs["vision_tokens"]) + + with pytest.raises(TypeError, match="must return"): + Cosmos3OmniGeneratorTask().build(BadModule(), tiny_config()) + + def test_task_rejects_missing_configured_head_output(self): + class NoSoundModule(nn.Module): + def forward(self, op, **kwargs): + return op.Identity(kwargs["vision_tokens"]), None, None + + config = tiny_config(sound_gen=True, sound_dim=6) + with pytest.raises(TypeError, match="sound_pred=None"): + Cosmos3OmniGeneratorTask().build(NoSoundModule(), config) + + def test_module_requires_all_inputs_of_a_configured_head(self): + config = tiny_config(sound_gen=True, sound_dim=6) + module = Cosmos3OmniGeneratorModel(config) + _, builder = _make_graph() + placeholder = builder.input("x", dtype=ir.DataType.INT64, shape=[1]) + with pytest.raises(ValueError, match="sound_gen=True requires all sound inputs"): + module( + builder.op, + input_ids=placeholder, + text_indexes=placeholder, + position_ids=placeholder, + und_len=placeholder, + vision_tokens=placeholder, + vision_sequence_indexes=placeholder, + vision_timesteps=placeholder, + vision_timestep_token_indexes=placeholder, + vision_mse_loss_indexes=placeholder, + ) + + def test_module_rejects_inputs_for_a_disabled_head(self): + module = Cosmos3OmniGeneratorModel(tiny_config()) + _, builder = _make_graph() + placeholder = builder.input("x", dtype=ir.DataType.INT64, shape=[1]) + with pytest.raises(ValueError, match="action_gen=False but action inputs"): + module( + builder.op, + input_ids=placeholder, + text_indexes=placeholder, + position_ids=placeholder, + und_len=placeholder, + vision_tokens=placeholder, + vision_sequence_indexes=placeholder, + vision_timesteps=placeholder, + vision_timestep_token_indexes=placeholder, + vision_mse_loss_indexes=placeholder, + action_tokens=placeholder, + ) + + +# --------------------------------------------------------------------------- +# Initializer-name coverage +# --------------------------------------------------------------------------- + + +def published_layer_keys(layer: int, *, gated: bool, qk_norm: bool, und_k_norm: bool) -> set: + """Flat published parameter names for one MoT layer.""" + prefix = f"layers.{layer}." + keys = { + prefix + "self_attn.to_q.weight", + prefix + "self_attn.to_k.weight", + prefix + "self_attn.to_v.weight", + prefix + "self_attn.to_out.weight", + prefix + "self_attn.add_q_proj.weight", + prefix + "self_attn.add_k_proj.weight", + prefix + "self_attn.add_v_proj.weight", + prefix + "self_attn.to_add_out.weight", + prefix + "self_attn.norm_added_q.weight", + prefix + "self_attn.norm_added_k.weight", + prefix + "input_layernorm.weight", + prefix + "input_layernorm_moe_gen.weight", + prefix + "post_attention_layernorm.weight", + prefix + "post_attention_layernorm_moe_gen.weight", + prefix + "mlp.up_proj.weight", + prefix + "mlp.down_proj.weight", + prefix + "mlp_moe_gen.up_proj.weight", + prefix + "mlp_moe_gen.down_proj.weight", + } + if gated: + keys |= {prefix + "mlp.gate_proj.weight", prefix + "mlp_moe_gen.gate_proj.weight"} + if qk_norm: + keys |= {prefix + "self_attn.norm_q.weight", prefix + "self_attn.norm_k.weight"} + if und_k_norm: + keys.add(prefix + "self_attn.k_norm_und_for_gen.weight") + return keys + + +class TestInitializerNames: + def test_core_and_head_initializers_match_published_names(self): + config = full_config() + model = build_package(config)["model"] + initializers = set(model.graph.initializers) + + expected = { + "embed_tokens.weight", + "norm.weight", + "norm_moe_gen.weight", + "proj_in.weight", + "proj_in.bias", + "proj_out.weight", + "proj_out.bias", + "time_embedder.linear_1.weight", + "time_embedder.linear_1.bias", + "time_embedder.linear_2.weight", + "time_embedder.linear_2.bias", + "audio_proj_in.weight", + "audio_proj_in.bias", + "audio_proj_out.weight", + "audio_proj_out.bias", + "audio_modality_embed", + "action_proj_in.fc.weight", + "action_proj_in.bias.weight", + "action_proj_out.fc.weight", + "action_proj_out.bias.weight", + "action_modality_embed", + } + for layer in range(config.num_hidden_layers): + expected |= published_layer_keys(layer, gated=True, qk_norm=True, und_k_norm=False) + assert expected <= initializers + # The module's declared checkpoint surface is exactly that set. + assert Cosmos3OmniGeneratorModel(config).expected_checkpoint_keys() == expected + + def test_relu2_backbone_drops_gate_proj(self): + config = tiny_config(hidden_act="relu2") + keys = Cosmos3OmniGeneratorModel(config).expected_checkpoint_keys() + assert "layers.0.mlp.up_proj.weight" in keys + assert "layers.0.mlp_moe_gen.up_proj.weight" in keys + assert not [key for key in keys if "gate_proj" in key] + + def test_qk_norm_off_swaps_in_und_k_norm_for_gen(self): + config = tiny_config(qk_norm_for_text=False, use_und_k_norm_for_gen=True) + keys = Cosmos3OmniGeneratorModel(config).expected_checkpoint_keys() + assert "layers.0.self_attn.k_norm_und_for_gen.weight" in keys + assert "layers.0.self_attn.norm_q.weight" not in keys + assert "layers.0.self_attn.norm_k.weight" not in keys + # The generation pathway always keeps its QK norms. + assert "layers.0.self_attn.norm_added_q.weight" in keys + + def test_disabled_heads_have_no_initializers(self): + keys = Cosmos3OmniGeneratorModel(tiny_config()).expected_checkpoint_keys() + assert not [key for key in keys if key.startswith(("audio_", "action_"))] + + def test_derived_constants_are_not_checkpoint_keys(self): + config = full_config() + model = build_package(config)["model"] + keys = Cosmos3OmniGeneratorModel(config).expected_checkpoint_keys() + for name in ("rotary_emb.inv_freq", "rotary_emb.h_mask", "rotary_emb.w_mask"): + assert name in model.graph.initializers + assert name not in keys + assert "time_proj.inv_freq" in model.graph.initializers + assert "time_proj.inv_freq" not in keys + + +# --------------------------------------------------------------------------- +# Interleaved 3-axis mRoPE layout +# --------------------------------------------------------------------------- + + +def upstream_interleaved_axis_map(rope_axes_dim, rotary_dim: int) -> np.ndarray: + """Axis (0=T, 1=H, 2=W) per rotary channel, using upstream's slice logic. + + Mirrors ``Cosmos3VLTextRotaryEmbedding.apply_interleaved_mrope``: start + from the T frequencies and overwrite ``slice(offset, axis_dim * 3, 3)`` + with the H (offset 1) then W (offset 2) frequencies. + """ + axis_map = np.zeros(rotary_dim, dtype=np.int64) + for dim, offset in enumerate((1, 2), start=1): + length = rope_axes_dim[dim] * 3 + axis_map[slice(offset, min(length, rotary_dim), 3)] = dim + return axis_map + + +class TestRotaryLayout: + @pytest.mark.parametrize( + ("head_dim", "rope_axes_dim"), + [(128, (24, 20, 20)), (8, (2, 1, 1)), (64, (12, 10, 10))], + ) + def test_interleaved_masks_match_upstream(self, head_dim, rope_axes_dim): + config = tiny_config(head_dim=head_dim, rope_axes_dim=rope_axes_dim) + model = build_package(config)["model"] + h_mask = model.graph.initializers["rotary_emb.h_mask"].const_value.numpy() + w_mask = model.graph.initializers["rotary_emb.w_mask"].const_value.numpy() + + axis_map = upstream_interleaved_axis_map(rope_axes_dim, config.rotary_dim) + np.testing.assert_array_equal(h_mask, axis_map == 1) + np.testing.assert_array_equal(w_mask, axis_map == 2) + # Channel budget matches the config's (T, H, W) split exactly. + assert int((axis_map == 0).sum()) == rope_axes_dim[0] + assert int(h_mask.sum()) == rope_axes_dim[1] + assert int(w_mask.sum()) == rope_axes_dim[2] + # Interleaved, not chunked: T/H/W alternate in the low channels. + assert axis_map[:3].tolist() == [0, 1, 2] + + def test_inv_freq_matches_upstream(self): + config = tiny_config(head_dim=128, rope_axes_dim=(24, 20, 20), rope_theta=5e6) + model = build_package(config)["model"] + inv_freq = model.graph.initializers["rotary_emb.inv_freq"].const_value.numpy() + expected = 1.0 / (5e6 ** (np.arange(0, 128, 2, dtype=np.float32) / 128)) + assert inv_freq.shape == (64,) + np.testing.assert_allclose(inv_freq, expected, rtol=1e-6) + + def test_masks_are_disjoint(self): + config = tiny_config(head_dim=128, rope_axes_dim=(24, 20, 20)) + model = build_package(config)["model"] + h_mask = model.graph.initializers["rotary_emb.h_mask"].const_value.numpy() + w_mask = model.graph.initializers["rotary_emb.w_mask"].const_value.numpy() + assert not np.any(h_mask & w_mask) + + +# --------------------------------------------------------------------------- +# DomainAwareLinear +# --------------------------------------------------------------------------- + + +class TorchDomainAwareLinear(torch.nn.Module): + """PyTorch transcription of upstream ``DomainAwareLinear``.""" + + def __init__(self, in_features: int, out_features: int, num_domains: int): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.fc = torch.nn.Embedding(num_domains, out_features * in_features) + self.bias = torch.nn.Embedding(num_domains, out_features) + + def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor: + weight = self.fc(domain_id).view( + domain_id.shape[0], self.in_features, self.out_features + ) + bias = self.bias(domain_id).view(domain_id.shape[0], self.out_features) + return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias + + +class _DomainAwareLinearTask(ModelTask): + """Wraps a bare :class:`Cosmos3OmniDomainAwareLinear` in a 2-input graph.""" + + def build(self, module, config): + graph, builder = _make_graph(name="domain_aware_linear") + tokens = builder.input( + "tokens", dtype=ir.DataType.FLOAT, shape=["n", module.in_features] + ) + domain_ids = builder.input("domain_ids", dtype=ir.DataType.INT64, shape=["n"]) + builder.add_output(module(builder.op, tokens, domain_ids), "projected") + return ModelPackage({"model": _make_model(graph)}, config=config) + + +class TestDomainAwareLinear: + def test_parameter_names_and_shapes(self): + config = full_config() + model = build_package(config)["model"] + initializers = model.graph.initializers + # fc rows are flattened [in, out]; bias rows are [out]. + assert list(initializers["action_proj_in.fc.weight"].shape) == [ + config.num_embodiment_domains, + config.hidden_size * config.action_dim, + ] + assert list(initializers["action_proj_in.bias.weight"].shape) == [ + config.num_embodiment_domains, + config.hidden_size, + ] + assert list(initializers["action_proj_out.fc.weight"].shape) == [ + config.num_embodiment_domains, + config.action_dim * config.hidden_size, + ] + assert list(initializers["action_proj_out.bias.weight"].shape) == [ + config.num_embodiment_domains, + config.action_dim, + ] + + def test_uses_gather_not_a_single_shared_matmul(self): + module = Cosmos3OmniDomainAwareLinear(4, 6, 3) + model = _DomainAwareLinearTask().build(module, tiny_config())["model"] + op_types = [node.op_type for node in model.graph] + # Two Gathers (weight + bias) proves the per-domain table lookup; a + # plain Linear would have neither. + assert op_types.count("Gather") == 2 + assert "MatMul" in op_types + + def test_matches_torch_reference(self, tmp_path): + torch.manual_seed(5) + in_features, out_features, num_domains = 4, 6, 3 + reference = TorchDomainAwareLinear(in_features, out_features, num_domains).eval() + module = Cosmos3OmniDomainAwareLinear(in_features, out_features, num_domains) + package = _DomainAwareLinearTask().build(module, tiny_config()) + package.apply_weights( + { + "fc.weight": reference.fc.weight.data, + "bias.weight": reference.bias.weight.data, + } + ) + package.save(str(tmp_path), progress_bar=False) + + session = ort.InferenceSession( + str(tmp_path / "model.onnx"), providers=["CPUExecutionProvider"] + ) + rng = np.random.default_rng(9) + tokens = rng.standard_normal((5, in_features)).astype(np.float32) + # Repeated and out-of-order ids exercise the per-token gather. + domain_ids = np.array([2, 0, 2, 1, 0], dtype=np.int64) + + actual = session.run(None, {"tokens": tokens, "domain_ids": domain_ids})[0] + with torch.no_grad(): + expected = reference(torch.from_numpy(tokens), torch.from_numpy(domain_ids)) + assert actual.shape == (5, out_features) + np.testing.assert_allclose(actual, expected.numpy(), rtol=1e-5, atol=1e-6) + + def test_distinct_domains_produce_distinct_outputs(self, tmp_path): + torch.manual_seed(6) + module = Cosmos3OmniDomainAwareLinear(3, 3, 2) + package = _DomainAwareLinearTask().build(module, tiny_config()) + package.apply_weights( + { + "fc.weight": torch.tensor( + [[1.0, 0, 0, 0, 1, 0, 0, 0, 1], [0.0, 0, 2, 0, 2, 0, 2, 0, 0]] + ), + "bias.weight": torch.tensor([[0.0, 0, 0], [1.0, 1, 1]]), + } + ) + package.save(str(tmp_path), progress_bar=False) + session = ort.InferenceSession( + str(tmp_path / "model.onnx"), providers=["CPUExecutionProvider"] + ) + tokens = np.array([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]], dtype=np.float32) + actual = session.run( + None, {"tokens": tokens, "domain_ids": np.array([0, 1], dtype=np.int64)} + )[0] + # Domain 0 is the identity; domain 1 permutes/scales and adds a bias. + np.testing.assert_allclose(actual[0], [1.0, 2.0, 3.0], rtol=1e-6) + np.testing.assert_allclose(actual[1], [7.0, 5.0, 3.0], rtol=1e-6) + + +# --------------------------------------------------------------------------- +# preprocess_weights +# --------------------------------------------------------------------------- + + +def fake_checkpoint(module: Cosmos3OmniGeneratorModel) -> dict: + """Build a flat state dict covering exactly the module's expected keys.""" + package = build_package(module.config) + initializers = package["model"].graph.initializers + return { + name: torch.zeros(*[int(dim) for dim in initializers[name].shape]) + for name in sorted(module.expected_checkpoint_keys()) + } + + +class TestPreprocessWeights: + def test_accepts_flat_published_keys_unchanged(self): + module = Cosmos3OmniGeneratorModel(full_config()) + state_dict = fake_checkpoint(module) + assert module.preprocess_weights(state_dict).keys() == state_dict.keys() + + def test_drops_unused_lm_head(self): + module = Cosmos3OmniGeneratorModel(tiny_config()) + state_dict = fake_checkpoint(module) + state_dict["lm_head.weight"] = torch.zeros(4, 4) + assert "lm_head.weight" not in module.preprocess_weights(state_dict) + + def test_drops_reasoner_vision_tower_keys(self): + module = Cosmos3OmniGeneratorModel(tiny_config()) + state_dict = fake_checkpoint(module) + state_dict["blocks.0.attn.qkv.weight"] = torch.zeros(2, 2) + state_dict["patch_embed.proj.weight"] = torch.zeros(2, 2) + state_dict["merger.linear_fc1.weight"] = torch.zeros(2, 2) + state_dict["deepstack_merger_list.0.norm.weight"] = torch.zeros(2) + state_dict["model.projector.linear_fc1.weight"] = torch.zeros(2, 2) + processed = module.preprocess_weights(state_dict) + assert not [ + key + for key in processed + if key.startswith(("blocks.", "patch_embed.", "projector.")) + ] + + def test_drops_edge_framework_key_norm_duplicates(self): + module = Cosmos3OmniGeneratorModel( + tiny_config( + qk_norm_for_text=False, + use_und_k_norm_for_gen=True, + ) + ) + state_dict = fake_checkpoint(module) + state_dict[ + "model.net.language_model.model.layers.0.self_attn.k_norm_und_for_gen.weight" + ] = torch.zeros(module.config.head_dim) + + processed = module.preprocess_weights(state_dict) + + assert not [key for key in processed if key.startswith("net.language_model.")] + + def test_strips_container_prefix(self): + module = Cosmos3OmniGeneratorModel(tiny_config()) + state_dict = {f"transformer.{k}": v for k, v in fake_checkpoint(module).items()} + processed = module.preprocess_weights(state_dict) + assert "embed_tokens.weight" in processed + assert not [key for key in processed if key.startswith("transformer.")] + + def test_normalizes_sequential_to_out(self): + module = Cosmos3OmniGeneratorModel(tiny_config()) + state_dict = fake_checkpoint(module) + value = state_dict.pop("layers.0.self_attn.to_out.weight") + state_dict["layers.0.self_attn.to_out.0.weight"] = value + processed = module.preprocess_weights(state_dict) + assert "layers.0.self_attn.to_out.weight" in processed + assert "layers.0.self_attn.to_out.0.weight" not in processed + + def test_drops_recomputed_buffers(self): + module = Cosmos3OmniGeneratorModel(tiny_config()) + state_dict = fake_checkpoint(module) + state_dict["rotary_emb.inv_freq"] = torch.zeros(4) + assert "rotary_emb.inv_freq" not in module.preprocess_weights(state_dict) + + def test_rejects_unexpected_key(self): + module = Cosmos3OmniGeneratorModel(tiny_config()) + state_dict = fake_checkpoint(module) + state_dict["layers.0.self_attn.mystery_proj.weight"] = torch.zeros(2, 2) + with pytest.raises(ValueError, match="Unexpected Cosmos3-Omni transformer weights"): + module.preprocess_weights(state_dict) + + def test_rejects_missing_key(self): + module = Cosmos3OmniGeneratorModel(tiny_config()) + state_dict = fake_checkpoint(module) + del state_dict["layers.1.mlp_moe_gen.down_proj.weight"] + with pytest.raises(ValueError, match="missing weights"): + module.preprocess_weights(state_dict) + + def test_weight_shard_allows_partial_checkpoint(self): + module = Cosmos3OmniGeneratorModel(tiny_config()) + state_dict = fake_checkpoint(module) + first_name = next(iter(state_dict)) + + processed = module.preprocess_weight_shard({first_name: state_dict[first_name]}) + + assert processed == {first_name: state_dict[first_name]} + + def test_rejects_gated_checkpoint_against_relu2_config(self): + """A gate_proj in the checkpoint contradicts a relu2 backbone.""" + module = Cosmos3OmniGeneratorModel(tiny_config(hidden_act="relu2")) + state_dict = fake_checkpoint(module) + state_dict["layers.0.mlp.gate_proj.weight"] = torch.zeros(48, 32) + with pytest.raises(ValueError, match="Unexpected Cosmos3-Omni transformer weights"): + module.preprocess_weights(state_dict) + + def test_rejects_action_weights_when_head_disabled(self): + module = Cosmos3OmniGeneratorModel(tiny_config()) + state_dict = fake_checkpoint(module) + state_dict["action_proj_in.fc.weight"] = torch.zeros(3, 4) + with pytest.raises(ValueError, match="Unexpected Cosmos3-Omni transformer weights"): + module.preprocess_weights(state_dict) + + def test_reports_missing_action_weights_when_head_enabled(self): + module = Cosmos3OmniGeneratorModel(full_config()) + state_dict = fake_checkpoint(module) + del state_dict["action_proj_out.fc.weight"] + with pytest.raises(ValueError, match=r"action_proj_out\.fc\.weight"): + module.preprocess_weights(state_dict) + + +# --------------------------------------------------------------------------- +# Numerical parity against a PyTorch transcription of upstream forward +# --------------------------------------------------------------------------- + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + """Upstream ``_rotate_half``.""" + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """Upstream ``Cosmos3NemotronRMSNorm`` / diffusers ``RMSNorm``.""" + dtype = x.dtype + x = x.float() + variance = x.pow(2).mean(-1, keepdim=True) + return (weight.float() * (x * torch.rsqrt(variance + eps))).to(dtype) + + +class TorchCosmos3Generator(torch.nn.Module): + """PyTorch transcription of ``Cosmos3OmniTransformer.forward`` (packed). + + Implements exactly the upstream math — interleaved 3-axis mRoPE, the two + QK-normed pathways, causal understanding self-attention, non-causal + generation attention over ``[und ; gen]`` K/V, MoT per-expert MLPs, fp32 + timestep conditioning and the per-domain action heads — but reads the same + packed tensors as the ONNX graph instead of ragged Python lists. + """ + + def __init__(self, config: Cosmos3OmniGeneratorConfig, seed: int = 0): + super().__init__() + self.config = config + generator = torch.Generator().manual_seed(seed) + + def parameter(*shape, scale=0.05): + return torch.nn.Parameter(torch.randn(*shape, generator=generator) * scale) + + hidden, inter = config.hidden_size, config.intermediate_size + head_dim = config.head_dim + q_size = config.attention_out_size + kv_size = config.key_value_size + + self.embed_tokens = parameter(config.vocab_size, hidden) + self.layer_params: list[dict[str, torch.nn.Parameter]] = [] + for index in range(config.num_hidden_layers): + layer = { + "to_q": parameter(q_size, hidden), + "to_k": parameter(kv_size, hidden), + "to_v": parameter(kv_size, hidden), + "to_out": parameter(hidden, q_size), + "add_q_proj": parameter(q_size, hidden), + "add_k_proj": parameter(kv_size, hidden), + "add_v_proj": parameter(kv_size, hidden), + "to_add_out": parameter(hidden, q_size), + "norm_added_q": parameter(head_dim, scale=0.1), + "norm_added_k": parameter(head_dim, scale=0.1), + "input_layernorm": parameter(hidden, scale=0.1), + "input_layernorm_moe_gen": parameter(hidden, scale=0.1), + "post_attention_layernorm": parameter(hidden, scale=0.1), + "post_attention_layernorm_moe_gen": parameter(hidden, scale=0.1), + "mlp.up_proj": parameter(inter, hidden), + "mlp.down_proj": parameter(hidden, inter), + "mlp_moe_gen.up_proj": parameter(inter, hidden), + "mlp_moe_gen.down_proj": parameter(hidden, inter), + } + if config.is_gated_mlp: + layer["mlp.gate_proj"] = parameter(inter, hidden) + layer["mlp_moe_gen.gate_proj"] = parameter(inter, hidden) + if config.qk_norm_for_text: + layer["norm_q"] = parameter(head_dim, scale=0.1) + layer["norm_k"] = parameter(head_dim, scale=0.1) + if config.has_und_k_norm_for_gen: + layer["k_norm_und_for_gen"] = parameter(head_dim, scale=0.1) + for name, value in layer.items(): + self.register_parameter(f"layer{index}_{name.replace('.', '_')}", value) + self.layer_params.append(layer) + + self.norm = parameter(hidden, scale=0.1) + self.norm_moe_gen = parameter(hidden, scale=0.1) + self.proj_in_weight = parameter(hidden, config.patch_latent_dim) + self.proj_in_bias = parameter(hidden) + self.proj_out_weight = parameter(config.patch_latent_dim, hidden) + self.proj_out_bias = parameter(config.patch_latent_dim) + self.time1_weight = parameter(hidden, config.time_proj_channels) + self.time1_bias = parameter(hidden) + self.time2_weight = parameter(hidden, hidden) + self.time2_bias = parameter(hidden) + if config.sound_gen: + self.audio_in_weight = parameter(hidden, config.sound_dim) + self.audio_in_bias = parameter(hidden) + self.audio_out_weight = parameter(config.sound_dim, hidden) + self.audio_out_bias = parameter(config.sound_dim) + self.audio_modality_embed = parameter(hidden) + if config.action_gen: + domains = config.num_embodiment_domains + self.action_in_fc = parameter(domains, hidden * config.action_dim) + self.action_in_bias = parameter(domains, hidden) + self.action_out_fc = parameter(domains, config.action_dim * hidden) + self.action_out_bias = parameter(domains, config.action_dim) + self.action_modality_embed = parameter(hidden) + + def published_state_dict(self) -> dict: + """Return the weights under the published flat checkpoint names.""" + config = self.config + state: dict = {"embed_tokens.weight": self.embed_tokens.data} + attention_names = { + "to_q": "self_attn.to_q", + "to_k": "self_attn.to_k", + "to_v": "self_attn.to_v", + "to_out": "self_attn.to_out", + "norm_q": "self_attn.norm_q", + "norm_k": "self_attn.norm_k", + "add_q_proj": "self_attn.add_q_proj", + "add_k_proj": "self_attn.add_k_proj", + "add_v_proj": "self_attn.add_v_proj", + "to_add_out": "self_attn.to_add_out", + "norm_added_q": "self_attn.norm_added_q", + "norm_added_k": "self_attn.norm_added_k", + "k_norm_und_for_gen": "self_attn.k_norm_und_for_gen", + } + for index, layer in enumerate(self.layer_params): + for key, value in layer.items(): + name = attention_names.get(key, key) + state[f"layers.{index}.{name}.weight"] = value.data + state["norm.weight"] = self.norm.data + state["norm_moe_gen.weight"] = self.norm_moe_gen.data + state["proj_in.weight"] = self.proj_in_weight.data + state["proj_in.bias"] = self.proj_in_bias.data + state["proj_out.weight"] = self.proj_out_weight.data + state["proj_out.bias"] = self.proj_out_bias.data + state["time_embedder.linear_1.weight"] = self.time1_weight.data + state["time_embedder.linear_1.bias"] = self.time1_bias.data + state["time_embedder.linear_2.weight"] = self.time2_weight.data + state["time_embedder.linear_2.bias"] = self.time2_bias.data + if config.sound_gen: + state["audio_proj_in.weight"] = self.audio_in_weight.data + state["audio_proj_in.bias"] = self.audio_in_bias.data + state["audio_proj_out.weight"] = self.audio_out_weight.data + state["audio_proj_out.bias"] = self.audio_out_bias.data + state["audio_modality_embed"] = self.audio_modality_embed.data + if config.action_gen: + state["action_proj_in.fc.weight"] = self.action_in_fc.data + state["action_proj_in.bias.weight"] = self.action_in_bias.data + state["action_proj_out.fc.weight"] = self.action_out_fc.data + state["action_proj_out.bias.weight"] = self.action_out_bias.data + state["action_modality_embed"] = self.action_modality_embed.data + return state + + def rotary(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + config = self.config + head_dim = config.head_dim + inv_freq = 1.0 / ( + config.rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim) + ) + freqs = position_ids.float().unsqueeze(-1) * inv_freq # (3, seq, head_dim // 2) + merged = freqs[0].clone() + for dim, offset in enumerate((1, 2), start=1): + length = config.rope_axes_dim[dim] * 3 + merged[..., slice(offset, length, 3)] = freqs[dim][..., slice(offset, length, 3)] + emb = torch.cat((merged, merged), dim=-1) + return emb.cos(), emb.sin() + + def timestep_embedding(self, timesteps: torch.Tensor) -> torch.Tensor: + config = self.config + scaled = timesteps * config.timestep_scale + half = config.time_proj_channels // 2 + exponent = -np.log(10000.0) * torch.arange(half, dtype=torch.float32) / half + freqs = scaled[:, None] * torch.exp(exponent)[None, :] + emb = torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1) + emb = torch.cat([emb[:, half:], emb[:, :half]], dim=-1) # flip_sin_to_cos + hidden = functional.silu(functional.linear(emb, self.time1_weight, self.time1_bias)) + return functional.linear(hidden, self.time2_weight, self.time2_bias) + + def _feed_forward(self, x: torch.Tensor, layer: dict, prefix: str) -> torch.Tensor: + """Upstream ``Cosmos3VLTextMLP`` for both ``silu`` and ``relu2``.""" + up = functional.linear(x, layer[f"{prefix}.up_proj"]) + if not self.config.is_gated_mlp: + # relu2: down_proj(relu(up_proj(x)) ** 2) + return functional.linear( + functional.relu(up).square(), layer[f"{prefix}.down_proj"] + ) + gate = functional.silu(functional.linear(x, layer[f"{prefix}.gate_proj"])) + return functional.linear(gate * up, layer[f"{prefix}.down_proj"]) + + def _attention(self, query, key, value, is_causal: bool) -> torch.Tensor: + config = self.config + groups = config.num_key_value_groups + query = query.transpose(0, 1).unsqueeze(0) + key = key.transpose(0, 1).unsqueeze(0).repeat_interleave(groups, dim=1) + value = value.transpose(0, 1).unsqueeze(0).repeat_interleave(groups, dim=1) + out = functional.scaled_dot_product_attention(query, key, value, is_causal=is_causal) + return out.squeeze(0).transpose(0, 1).reshape(-1, config.attention_out_size) + + def forward(self, batch: dict) -> dict: + config = self.config + eps = config.rms_norm_eps + head_dim = config.head_dim + heads, kv_heads = config.num_attention_heads, config.num_key_value_heads + + sequence_length = batch["position_ids"].shape[1] + hidden_states = torch.zeros(sequence_length, config.hidden_size) + hidden_states[batch["text_indexes"]] = self.embed_tokens[batch["input_ids"]] + + vision = functional.linear( + batch["vision_tokens"], self.proj_in_weight, self.proj_in_bias + ) + vision = vision.index_add( + 0, + batch["vision_timestep_token_indexes"], + self.timestep_embedding(batch["vision_timesteps"]), + ) + hidden_states[batch["vision_sequence_indexes"]] = vision + + if config.sound_gen: + sound = functional.linear( + batch["sound_tokens"], self.audio_in_weight, self.audio_in_bias + ) + sound = sound + self.audio_modality_embed + sound = sound.index_add( + 0, + batch["sound_timestep_token_indexes"], + self.timestep_embedding(batch["sound_timesteps"]), + ) + hidden_states[batch["sound_sequence_indexes"]] = sound + + if config.action_gen: + domain_ids = batch["action_domain_ids"] + weight = self.action_in_fc[domain_ids].view( + -1, config.action_dim, config.hidden_size + ) + action = torch.bmm(batch["action_tokens"].unsqueeze(1), weight).squeeze(1) + action = action + self.action_in_bias[domain_ids] + self.action_modality_embed + action = action.index_add( + 0, + batch["action_timestep_token_indexes"], + self.timestep_embedding(batch["action_timesteps"]), + ) + hidden_states[batch["action_sequence_indexes"]] = action + + cos, sin = self.rotary(batch["position_ids"]) + und_len = int(batch["und_len"][0]) + und_seq, gen_seq = hidden_states[:und_len], hidden_states[und_len:] + cos_und, sin_und = cos[:und_len].unsqueeze(1), sin[:und_len].unsqueeze(1) + cos_gen, sin_gen = cos[und_len:].unsqueeze(1), sin[und_len:].unsqueeze(1) + + for layer in self.layer_params: + und_norm = rms_norm(und_seq, layer["input_layernorm"], eps) + gen_norm = rms_norm(gen_seq, layer["input_layernorm_moe_gen"], eps) + + q_und = functional.linear(und_norm, layer["to_q"]).view(-1, heads, head_dim) + k_und = functional.linear(und_norm, layer["to_k"]).view(-1, kv_heads, head_dim) + v_und = functional.linear(und_norm, layer["to_v"]).view(-1, kv_heads, head_dim) + q_gen = functional.linear(gen_norm, layer["add_q_proj"]).view(-1, heads, head_dim) + k_gen = functional.linear(gen_norm, layer["add_k_proj"]).view( + -1, kv_heads, head_dim + ) + v_gen = functional.linear(gen_norm, layer["add_v_proj"]).view( + -1, kv_heads, head_dim + ) + + if config.qk_norm_for_text: + q_und = rms_norm(q_und, layer["norm_q"], eps) + k_und = rms_norm(k_und, layer["norm_k"], eps) + # Upstream only builds this norm when the und pathway has none. + k_und_for_gen = ( + rms_norm(k_und, layer["k_norm_und_for_gen"], eps) + if config.has_und_k_norm_for_gen + else k_und + ) + q_gen = rms_norm(q_gen, layer["norm_added_q"], eps) + k_gen = rms_norm(k_gen, layer["norm_added_k"], eps) + + q_und = q_und * cos_und + rotate_half(q_und) * sin_und + k_und = k_und * cos_und + rotate_half(k_und) * sin_und + k_und_for_gen = k_und_for_gen * cos_und + rotate_half(k_und_for_gen) * sin_und + q_gen = q_gen * cos_gen + rotate_half(q_gen) * sin_gen + k_gen = k_gen * cos_gen + rotate_half(k_gen) * sin_gen + + und_attn = self._attention(q_und, k_und, v_und, is_causal=True) + gen_attn = self._attention( + q_gen, + torch.cat([k_und_for_gen, k_gen], dim=0), + torch.cat([v_und, v_gen], dim=0), + is_causal=False, + ) + und_seq = und_seq + functional.linear(und_attn, layer["to_out"]) + gen_seq = gen_seq + functional.linear(gen_attn, layer["to_add_out"]) + + und_post = rms_norm(und_seq, layer["post_attention_layernorm"], eps) + gen_post = rms_norm(gen_seq, layer["post_attention_layernorm_moe_gen"], eps) + und_seq = und_seq + self._feed_forward(und_post, layer, "mlp") + gen_seq = gen_seq + self._feed_forward(gen_post, layer, "mlp_moe_gen") + + last_hidden_state = torch.cat( + [rms_norm(und_seq, self.norm, eps), rms_norm(gen_seq, self.norm_moe_gen, eps)], + dim=0, + ) + + outputs = { + "vision_pred": functional.linear( + last_hidden_state[batch["vision_mse_loss_indexes"]], + self.proj_out_weight, + self.proj_out_bias, + ) + } + if config.sound_gen: + outputs["sound_pred"] = functional.linear( + last_hidden_state[batch["sound_mse_loss_indexes"]], + self.audio_out_weight, + self.audio_out_bias, + ) + if config.action_gen: + domain_ids = batch["action_pred_domain_ids"] + weight = self.action_out_fc[domain_ids].view( + -1, config.hidden_size, config.action_dim + ) + selected = last_hidden_state[batch["action_mse_loss_indexes"]] + outputs["action_pred"] = ( + torch.bmm(selected.unsqueeze(1), weight).squeeze(1) + + self.action_out_bias[domain_ids] + ) + return outputs + + +def packed_batch(config: Cosmos3OmniGeneratorConfig) -> dict: + """Build a packed feed covering text, vision and the enabled heads.""" + rng = np.random.default_rng(3) + und_len, num_vision, num_vision_noisy = 5, 4, 3 + cursor = und_len + num_vision + + batch: dict = { + "input_ids": rng.integers(0, config.vocab_size, und_len).astype(np.int64), + "text_indexes": np.arange(und_len, dtype=np.int64), + "und_len": np.array([und_len], dtype=np.int64), + "vision_tokens": rng.standard_normal((num_vision, config.patch_latent_dim)).astype( + np.float32 + ), + "vision_sequence_indexes": np.arange(und_len, cursor, dtype=np.int64), + "vision_timesteps": rng.uniform(1, 1000, num_vision_noisy).astype(np.float32), + # Non-contiguous noisy rows: only frames 0, 1 and 3 carry noise. + "vision_timestep_token_indexes": np.array([0, 1, 3], dtype=np.int64), + "vision_mse_loss_indexes": np.array( + [und_len, und_len + 1, und_len + 3], dtype=np.int64 + ), + } + if config.sound_gen: + num_sound = 2 + batch.update( + { + "sound_tokens": rng.standard_normal((num_sound, config.sound_dim)).astype( + np.float32 + ), + "sound_sequence_indexes": np.arange( + cursor, cursor + num_sound, dtype=np.int64 + ), + "sound_timesteps": rng.uniform(1, 1000, num_sound).astype(np.float32), + "sound_timestep_token_indexes": np.arange(num_sound, dtype=np.int64), + "sound_mse_loss_indexes": np.arange( + cursor, cursor + num_sound, dtype=np.int64 + ), + } + ) + cursor += num_sound + if config.action_gen: + num_action = 3 + batch.update( + { + "action_tokens": rng.standard_normal((num_action, config.action_dim)).astype( + np.float32 + ), + # Mixed embodiments in one packed batch. + "action_domain_ids": np.array([0, 2, 2], dtype=np.int64), + "action_sequence_indexes": np.arange( + cursor, cursor + num_action, dtype=np.int64 + ), + "action_timesteps": rng.uniform(1, 1000, 2).astype(np.float32), + "action_timestep_token_indexes": np.array([0, 2], dtype=np.int64), + "action_mse_loss_indexes": np.array([cursor, cursor + 2], dtype=np.int64), + "action_pred_domain_ids": np.array([0, 2], dtype=np.int64), + } + ) + cursor += num_action + + sequence_length = cursor + steps = np.arange(sequence_length, dtype=np.int64) + # Distinct T/H/W tracks so a collapsed mRoPE axis would show up. + batch["position_ids"] = np.stack([steps, (steps * 2) % 7, (steps * 3) % 5]) + return batch + + +class TestNumericalParity: + @pytest.mark.parametrize( + "variant", + [ + "generator_only", + "sound_and_action", + "und_k_norm_for_gen", + "relu2_backbone", + ], + ) + def test_matches_torch_reference(self, tmp_path, variant): + config = { + "generator_only": tiny_config(), + "sound_and_action": full_config(), + # qk_norm_for_text=False makes norm_q/norm_k Identity and adds the + # separate k_norm_und_for_gen applied to the raw understanding keys. + "und_k_norm_for_gen": tiny_config( + qk_norm_for_text=False, use_und_k_norm_for_gen=True + ), + # Nemotron backbone: non-gated squared-ReLU feed-forward. + "relu2_backbone": tiny_config(hidden_act="relu2"), + }[variant] + reference = TorchCosmos3Generator(config).eval() + module = Cosmos3OmniGeneratorModel(config) + + package = build_package(config, execution_provider="cpu") + package.apply_weights(module.preprocess_weights(reference.published_state_dict())) + package.save(str(tmp_path), progress_bar=False) + + session = ort.InferenceSession( + str(tmp_path / "model.onnx"), providers=["CPUExecutionProvider"] + ) + batch = packed_batch(config) + assert {value.name for value in package["model"].graph.inputs} == set(batch) + + output_names = [output.name for output in session.get_outputs()] + actual = session.run(output_names, batch) + with torch.no_grad(): + expected = reference( + {name: torch.from_numpy(value) for name, value in batch.items()} + ) + + assert output_names == list(expected_output_names(config)) + for name, value in zip(output_names, actual, strict=True): + np.testing.assert_allclose( + value, expected[name].numpy(), rtol=1e-4, atol=1e-4, err_msg=name + ) + + def test_causal_understanding_pathway_ignores_future_text(self, tmp_path): + """Understanding self-attention must stay causal. + + Perturbing the *last* understanding token cannot change the first + understanding token's contribution, but it must change the generation + predictions (which attend over all understanding keys non-causally). + """ + config = tiny_config() + reference = TorchCosmos3Generator(config).eval() + module = Cosmos3OmniGeneratorModel(config) + package = build_package(config, execution_provider="cpu") + package.apply_weights(module.preprocess_weights(reference.published_state_dict())) + package.save(str(tmp_path), progress_bar=False) + session = ort.InferenceSession( + str(tmp_path / "model.onnx"), providers=["CPUExecutionProvider"] + ) + + batch = packed_batch(config) + baseline = session.run(["vision_pred"], batch)[0] + perturbed_batch = dict(batch) + perturbed_batch["input_ids"] = batch["input_ids"].copy() + perturbed_batch["input_ids"][-1] = (batch["input_ids"][-1] + 7) % config.vocab_size + perturbed = session.run(["vision_pred"], perturbed_batch)[0] + # Generation tokens see every understanding key, so the change lands. + assert not np.allclose(baseline, perturbed, atol=1e-6) + + def test_zero_length_optional_heads_run(self, tmp_path): + """A configured head with no content this step accepts empty tensors.""" + config = full_config() + reference = TorchCosmos3Generator(config).eval() + module = Cosmos3OmniGeneratorModel(config) + package = build_package(config, execution_provider="cpu") + package.apply_weights(module.preprocess_weights(reference.published_state_dict())) + package.save(str(tmp_path), progress_bar=False) + session = ort.InferenceSession( + str(tmp_path / "model.onnx"), providers=["CPUExecutionProvider"] + ) + + batch = packed_batch(config) + empty = dict(batch) + for name in ( + "sound_tokens", + "sound_sequence_indexes", + "sound_timesteps", + "sound_timestep_token_indexes", + "sound_mse_loss_indexes", + "action_tokens", + "action_domain_ids", + "action_sequence_indexes", + "action_timesteps", + "action_timestep_token_indexes", + "action_mse_loss_indexes", + "action_pred_domain_ids", + ): + value = batch[name] + empty[name] = np.zeros((0, *value.shape[1:]), dtype=value.dtype) + + outputs = session.run(None, empty) + shapes = { + output.name: value.shape + for output, value in zip(session.get_outputs(), outputs, strict=True) + } + assert shapes["sound_pred"] == (0, config.sound_dim) + assert shapes["action_pred"] == (0, config.action_dim) + assert shapes["vision_pred"][0] == batch["vision_mse_loss_indexes"].shape[0] + + +def test_config_is_a_dataclass_with_replaceable_fields(): + """``dataclasses.replace`` is how the build pipeline derives variants.""" + config = Cosmos3OmniGeneratorConfig.from_diffusers(PUBLISHED_CONFIG) + replaced = dataclasses.replace(config, dtype=ir.DataType.FLOAT16) + replaced.validate() + assert replaced.dtype == ir.DataType.FLOAT16 + assert replaced.sound_gen is config.sound_gen diff --git a/src/mobius/models/cosmos_test.py b/src/mobius/models/cosmos_test.py index 2dbb29d36..13b707758 100644 --- a/src/mobius/models/cosmos_test.py +++ b/src/mobius/models/cosmos_test.py @@ -1,12 +1,15 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Unit tests for Cosmos3-Edge weight routing and projector configuration.""" +"""Unit tests for Cosmos3-Edge weight routing, vision I/O and token fusion.""" from __future__ import annotations from types import SimpleNamespace +import numpy as np +import onnx_ir as ir +import onnxruntime as ort import pytest import torch @@ -19,13 +22,16 @@ ) from mobius.tasks import Cosmos3EdgeVLTask +IMAGE_TOKEN_ID = 19 +VIDEO_TOKEN_ID = 18 + def _tiny_config( *, tie_word_embeddings: bool = False, spatial_merge_size: int | None = 2, out_hidden_size: int | None = 64, - image_size: int = 28, + num_patches: int | None = 16, patch_size: int = 14, ) -> ArchitectureConfig: return ArchitectureConfig( @@ -40,18 +46,24 @@ def _tiny_config( max_position_embeddings=128, hidden_act="relu2", rms_norm_eps=1e-6, - attn_qk_norm=True, tie_word_embeddings=tie_word_embeddings, + mrope_section=[4, 2, 2], + mrope_interleaved=True, + image_token_id=IMAGE_TOKEN_ID, vision=VisionConfig( hidden_size=32, intermediate_size=64, num_hidden_layers=1, num_attention_heads=2, - image_size=image_size, + image_size=None, patch_size=patch_size, + num_patches=num_patches, spatial_merge_size=spatial_merge_size, + temporal_patch_size=1, out_hidden_size=out_hidden_size, projector_intermediate_size=64, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, ), ) @@ -97,10 +109,57 @@ def test_vl_tied_lm_head_populates_embedding(): assert result["embedding.embed_tokens.weight"] is weight +def test_vl_drops_unified_generator_and_action_weights(): + module = Cosmos3EdgeVLModel(_tiny_config()) + weights = { + "layers.0.self_attn.add_q_proj.weight": torch.zeros(1), + "layers.0.mlp_moe_gen.up_proj.weight": torch.zeros(1), + "time_embedder.linear_1.weight": torch.zeros(1), + "proj_in.weight": torch.zeros(1), + "action_proj_in.fc.weight": torch.zeros(1), + "action_modality_embed": torch.zeros(1), + } + + assert module.preprocess_weights(weights) == {} + + +def test_vision_weights_keep_hf_layout_and_linear_patch_embedding(): + """``model.visual``/``model.projector`` map with only a prefix strip. + + The checkpoint's ``patch_embedding`` is an ``nn.Linear`` over flattened + ``(patch_h, patch_w, channel)`` values, so it must stay 2-D — reshaping it + into a Conv2d ``[out, C, kH, kW]`` kernel would scramble the channel axis. + """ + module = Cosmos3EdgeVLModel(_tiny_config()) + patch_weight = torch.arange(32 * 3 * 14 * 14, dtype=torch.float32).reshape(32, 3 * 14 * 14) + weights = { + "model.visual.embeddings.patch_embedding.weight": patch_weight, + "model.visual.embeddings.patch_embedding.bias": torch.zeros(32), + "model.visual.embeddings.position_embedding.weight": torch.zeros(16, 32), + "model.visual.encoder.layers.0.mlp.fc1.weight": torch.zeros(64, 32), + "model.visual.encoder.layers.0.mlp.fc2.weight": torch.zeros(32, 64), + "model.visual.post_layernorm.weight": torch.zeros(32), + "model.projector.norm.weight": torch.zeros(32), + "model.projector.linear_fc1.weight": torch.zeros(64, 128), + "model.projector.linear_fc2.weight": torch.zeros(64, 64), + } + + result = module.preprocess_weights(weights) + + embed = result["vision_encoder.visual.embeddings.patch_embedding.weight"] + assert embed.shape == (32, 3 * 14 * 14) + assert torch.equal(embed, patch_weight) + assert "vision_encoder.visual.embeddings.position_embedding.weight" in result + assert "vision_encoder.visual.encoder.layers.0.mlp.up_proj.weight" in result + assert "vision_encoder.visual.encoder.layers.0.mlp.down_proj.weight" in result + assert "vision_encoder.visual.post_layernorm.weight" in result + assert "vision_encoder.projector.linear_fc1.weight" in result + + def test_vision_encoder_defaults_missing_spatial_merge_size(): module = _Cosmos3EdgeVisionEncoderModel(_tiny_config(spatial_merge_size=None)) - assert module.multi_modal_projector._ms == 2 + assert module.projector.spatial_merge_size == 2 def test_vision_encoder_rejects_projector_width_mismatch(): @@ -108,32 +167,179 @@ def test_vision_encoder_rejects_projector_width_mismatch(): _Cosmos3EdgeVisionEncoderModel(_tiny_config(out_hidden_size=32)) -def test_vl_task_vision_output_matches_embedding_input(): +def test_vision_encoder_rejects_non_square_reference_grid(): + with pytest.raises(ValueError, match="square reference grid"): + _Cosmos3EdgeVisionEncoderModel(_tiny_config(num_patches=15)) + + +def test_vl_task_declares_packed_vision_and_dual_feature_embedding(): config = _tiny_config() package = Cosmos3EdgeVLTask().build(Cosmos3EdgeVLModel(config), config) - pixel_values = package["vision_encoder"].graph.inputs[0] - image_features = package["vision_encoder"].graph.outputs[0] - embedding_features = next( - value for value in package["embedding"].graph.inputs if value.name == "image_features" - ) + vision = package["vision_encoder"].graph + pixel_values, grid_thw = vision.inputs + assert pixel_values.name == "pixel_values" + # patch_dim = patch * patch * channels * temporal_patch_size + assert pixel_values.shape[1] == 14 * 14 * 3 * 1 + assert grid_thw.name == "grid_thw" + assert grid_thw.dtype == ir.DataType.INT64 + assert list(grid_thw.shape) == [3] - assert pixel_values.shape[0] == 1 + image_features = vision.outputs[0] assert len(image_features.shape) == 2 - assert len(embedding_features.shape) == 2 - assert image_features.shape[-1] == embedding_features.shape[-1] == config.hidden_size + assert image_features.shape[-1] == config.hidden_size + + embedding_inputs = {value.name: value for value in package["embedding"].graph.inputs} + assert set(embedding_inputs) == {"input_ids", "image_features", "video_features"} + for name in ("image_features", "video_features"): + assert len(embedding_inputs[name].shape) == 2 + assert embedding_inputs[name].shape[-1] == config.hidden_size + + # Interleaved MRoPE decoder contract: position_ids is [3, batch, seq]. + decoder_inputs = {value.name: value for value in package["decoder"].graph.inputs} + assert decoder_inputs["position_ids"].shape[0] == 3 -def test_vision_encoder_rejects_non_integral_patch_grid(): - with pytest.raises(ValueError, match=r"image_size .* divisible by patch_size"): - _Cosmos3EdgeVisionEncoderModel(_tiny_config(image_size=30, patch_size=14)) +def _run(model: ir.Model, feeds: dict[str, np.ndarray], tmp_path, name: str): + path = tmp_path / f"{name}.onnx" + ir.save(model, str(path)) + session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + return session.run(None, feeds) -def test_vision_encoder_rejects_unmergeable_patch_grid(): - with pytest.raises(ValueError, match=r"grid_size .* divisible"): - _Cosmos3EdgeVisionEncoderModel( - _tiny_config(image_size=42, patch_size=14, spatial_merge_size=2) - ) +def _random_vision_weights(config: ArchitectureConfig) -> dict[str, torch.Tensor]: + """Small random ``model.visual.*``/``model.projector.*`` checkpoint slice.""" + vc = config.vision + assert vc is not None + generator = torch.Generator().manual_seed(0) + + def randn(*shape: int) -> torch.Tensor: + return torch.randn(*shape, generator=generator) * 0.02 + + patch_dim = vc.patch_size * vc.patch_size * vc.in_channels + weights: dict[str, torch.Tensor] = { + "model.visual.embeddings.patch_embedding.weight": randn(vc.hidden_size, patch_dim), + "model.visual.embeddings.patch_embedding.bias": randn(vc.hidden_size), + "model.visual.embeddings.position_embedding.weight": randn( + vc.num_patches, vc.hidden_size + ), + "model.visual.post_layernorm.weight": torch.ones(vc.hidden_size), + "model.visual.post_layernorm.bias": torch.zeros(vc.hidden_size), + "model.projector.norm.weight": torch.ones(vc.hidden_size), + "model.projector.norm.bias": torch.zeros(vc.hidden_size), + "model.projector.linear_fc1.weight": randn( + vc.projector_intermediate_size, vc.hidden_size * 4 + ), + "model.projector.linear_fc1.bias": randn(vc.projector_intermediate_size), + "model.projector.linear_fc2.weight": randn( + config.hidden_size, vc.projector_intermediate_size + ), + "model.projector.linear_fc2.bias": randn(config.hidden_size), + } + for layer in range(vc.num_hidden_layers): + prefix = f"model.visual.encoder.layers.{layer}" + for norm in ("layer_norm1", "layer_norm2"): + weights[f"{prefix}.{norm}.weight"] = torch.ones(vc.hidden_size) + weights[f"{prefix}.{norm}.bias"] = torch.zeros(vc.hidden_size) + for proj in ("q_proj", "k_proj", "v_proj", "out_proj"): + weights[f"{prefix}.self_attn.{proj}.weight"] = randn( + vc.hidden_size, vc.hidden_size + ) + weights[f"{prefix}.self_attn.{proj}.bias"] = randn(vc.hidden_size) + weights[f"{prefix}.mlp.fc1.weight"] = randn(vc.intermediate_size, vc.hidden_size) + weights[f"{prefix}.mlp.fc1.bias"] = randn(vc.intermediate_size) + weights[f"{prefix}.mlp.fc2.weight"] = randn(vc.hidden_size, vc.intermediate_size) + weights[f"{prefix}.mlp.fc2.bias"] = randn(vc.hidden_size) + + routed = Cosmos3EdgeVLModel(config).preprocess_weights(weights) + return {key: value for key, value in routed.items() if key.startswith("vision_encoder.")} + + +def test_vision_encoder_token_count_scales_with_frames(tmp_path): + """Video frames reuse the image path: tokens = t*h*w / merge**2.""" + config = _tiny_config() + package = Cosmos3EdgeVLTask().build(Cosmos3EdgeVLModel(config), config) + package.apply_weights_partial(_random_vision_weights(config)) + + patch_dim = 14 * 14 * 3 + grid_h, grid_w = 4, 6 + rng = np.random.RandomState(0) + for frames in (1, 3): + total = frames * grid_h * grid_w + features = _run( + package["vision_encoder"], + { + "pixel_values": rng.randn(total, patch_dim).astype(np.float32), + "grid_thw": np.array([frames, grid_h, grid_w], dtype=np.int64), + }, + tmp_path, + f"vision_{frames}", + )[0] + assert features.shape == (total // 4, config.hidden_size) + + +def _set_embedding_table(package, table: np.ndarray) -> None: + package.apply_weights_partial({"embedding.embed_tokens.weight": torch.from_numpy(table)}) + + +def test_embedding_scatters_image_and_video_tokens(tmp_path): + config = _tiny_config() + package = Cosmos3EdgeVLTask().build(Cosmos3EdgeVLModel(config), config) + table = np.arange(config.vocab_size * config.hidden_size, dtype=np.float32).reshape( + config.vocab_size, config.hidden_size + ) + table /= table.max() + _set_embedding_table(package, table) + + input_ids = np.array( + [[7, IMAGE_TOKEN_ID, IMAGE_TOKEN_ID, 8, VIDEO_TOKEN_ID, VIDEO_TOKEN_ID, 9]], + dtype=np.int64, + ) + image_features = np.full((2, config.hidden_size), 3.0, dtype=np.float32) + image_features[1] = 4.0 + video_features = np.full((2, config.hidden_size), -5.0, dtype=np.float32) + video_features[1] = -6.0 + + embeds = _run( + package["embedding"], + { + "input_ids": input_ids, + "image_features": image_features, + "video_features": video_features, + }, + tmp_path, + "embedding", + )[0] + + np.testing.assert_allclose(embeds[0, 1], image_features[0]) + np.testing.assert_allclose(embeds[0, 2], image_features[1]) + np.testing.assert_allclose(embeds[0, 4], video_features[0]) + np.testing.assert_allclose(embeds[0, 5], video_features[1]) + # Text positions keep their token embedding. + np.testing.assert_allclose(embeds[0, 0], table[7]) + np.testing.assert_allclose(embeds[0, 3], table[8]) + np.testing.assert_allclose(embeds[0, 6], table[9]) + + +def test_embedding_tolerates_empty_feature_streams(tmp_path): + config = _tiny_config() + package = Cosmos3EdgeVLTask().build(Cosmos3EdgeVLModel(config), config) + _set_embedding_table( + package, np.zeros((config.vocab_size, config.hidden_size), dtype=np.float32) + ) + + embeds = _run( + package["embedding"], + { + "input_ids": np.array([[1, 2, 3]], dtype=np.int64), + "image_features": np.zeros((0, config.hidden_size), dtype=np.float32), + "video_features": np.zeros((0, config.hidden_size), dtype=np.float32), + }, + tmp_path, + "embedding_empty", + )[0] + + assert embeds.shape == (1, 3, config.hidden_size) def test_vision_config_rejects_non_square_num_patches(): @@ -142,3 +348,37 @@ def test_vision_config_rejects_non_square_num_patches(): with pytest.raises(ValueError, match="num_patches must form a square grid"): _cosmos3_edge_vision(config, None, "cosmos3_edge", {"image_size": None}) + + +def test_vision_config_hook_extracts_edge_fields(): + config = SimpleNamespace( + model_type="cosmos3_edge_text", + vision_config=SimpleNamespace( + num_patches=256, patch_size=16, hidden_act="gelu_pytorch_tanh" + ), + projector_config={ + "merger_intermediate_size": 128, + "out_hidden_size": 64, + "use_postshuffle_norm": False, + "spatial_merge_size": 2, + }, + image_token_id=19, + video_token_id=18, + vision_start_token_id=20, + vision_end_token_id=21, + ) + fields = {"image_size": None, "patch_size": 16, "out_hidden_size": None} + + _cosmos3_edge_vision(config, config, "cosmos3_edge_text", fields) + + assert fields["num_patches"] == 256 + assert fields["image_size"] == 256 + assert fields["projector_intermediate_size"] == 128 + assert fields["use_postshuffle_norm"] is False + assert fields["temporal_patch_size"] == 1 + assert fields["image_token_id"] == 19 + assert fields["video_token_id"] == 18 + assert fields["vision_start_token_id"] == 20 + assert fields["vision_end_token_id"] == 21 + # Cosmos3-Edge M-RoPE assigns axes to interleaved channels, not chunks. + assert fields["mrope_interleaved"] is True diff --git a/src/mobius/models/wan_vae.py b/src/mobius/models/wan_vae.py new file mode 100644 index 000000000..2940a4cf1 --- /dev/null +++ b/src/mobius/models/wan_vae.py @@ -0,0 +1,1167 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Wan 3D causal video VAE (``diffusers.AutoencoderKLWan``). + +Replicates the public HuggingFace diffusers class ``AutoencoderKLWan`` +(``diffusers/models/autoencoders/autoencoder_kl_wan.py``) as a from-scratch +``onnxscript.nn`` graph builder. This is the video VAE used by Wan 2.1 / +Wan 2.2 and, verbatim, by NVIDIA Cosmos3 (``nvidia/Cosmos3-Nano/vae``). + +Pipeline (Cosmos3 / Wan 2.2 shapes, ``patch_size=2``, ``z_dim=48``):: + + video (B, 3, 4k+1, H, W) pixel space, [-1, 1] + patchify fold 2x2 spatial patches into channels + (B, 12, 4k+1, H/2, W/2) + encoder 3 spatial stages, 2 temporal stages + (B, 96, k+1, H/16, W/16) 96 = 2 * z_dim (mean ‖ logvar) + quant_conv 1x1x1 causal conv + (B, 96, k+1, H/16, W/16) + chunk(2, dim=1) DiagonalGaussianDistribution moments + mean/logvar (B, 48, k+1, H/16, W/16) + + latent (B, 48, k+1, H/16, W/16) normalised: (z - latents_mean)/latents_std + denormalise -> post_quant_conv -> decoder + (B, 12, 4k+1, H/2, W/2) + unpatchify + clamp(-1, 1) + (B, 3, 4k+1, H, W) + +**Whole-sequence formulation.** Upstream runs the encoder/decoder in temporal +chunks, threading a two-frame ``feat_cache`` through every ``WanCausalConv3d``. +A static ONNX graph cannot carry that Python-side cache, so every block here is +written in *whole-sequence* form. The two are numerically identical: + +* ``WanCausalConv3d`` (kernel 3, temporal pad ``(2, 0)``) over ``[cache(2), chunk]`` + is exactly the full-sequence causal convolution restricted to that chunk's + output positions. +* ``downsample3d``'s ``time_conv`` (kernel 3, stride 2, no padding) is skipped + for chunk 0 and, from chunk 1 on, is fed ``[last frame of previous chunk, chunk]``. + Concatenated, its windows start at global frames ``0, 2, 4, ...`` — i.e. + ``concat(x[:, :, :1], time_conv(x))`` over the whole sequence. +* ``upsample3d``'s ``time_conv`` is skipped for chunk 0 (no temporal doubling) + and restarts the causal sequence with zero padding at chunk 1 — i.e. + ``concat(x[:, :, :1], interleave(time_conv(x[:, :, 1:])))``. +* ``AvgDown3D`` front-pads the temporal axis to a multiple of ``factor_t``; + chunk 0 (1 frame) pads 1, and a whole ``4k+1``-frame sequence also pads 1, + producing the same groups. +* ``DupUp3D`` with ``first_chunk=True`` drops the first duplicated frame, which + is exactly what chunk 0 contributes when the per-chunk outputs are concatenated. + +The equivalence holds for the supported frame counts: ``T_video = 4k + 1`` +(encoder) and ``T_latent = k + 1`` (decoder), for ``k >= 0``. + +**Single-frame (image) mode.** ``T_video = 1`` / ``T_latent = 1`` is a first +class upstream case — ``_encode`` runs ``1 + (T_video - 1) // 4`` chunks and +``_decode`` one chunk per latent frame, so a lone frame is decoded as chunk 0, +whose ``time_conv`` is skipped entirely (``feat_cache[idx] is None`` -> ``"Rep"`` +for ``upsample3d``, ``feat_cache[idx] = x`` for ``downsample3d``). A static +graph cannot skip a node, and the naive whole-sequence form would hand ``Conv`` +a temporal extent shorter than its kernel (ONNX Runtime rejects that outright), +so both temporal resamplers use a *safe window*: the temporal axis is +zero-padded on the **right** by exactly the number of frames the kernel needs, +and the surplus trailing outputs are sliced off afterwards. Because each +``time_conv`` is causal in time, the retained outputs are bit-identical to +the unpadded convolution, so multi-frame numerics are untouched and the +single-frame case degenerates to "keep frame 0 only", exactly like chunk 0 +upstream. The final ``Slice`` is applied *after* the frame-0 concatenation so +no zero-length tensor is ever materialised. + +Weight names match the HuggingFace ``AutoencoderKLWan`` ``state_dict`` exactly, +so :meth:`AutoencoderKLWanModel.preprocess_weights` performs no renaming. +""" + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import numpy as np +import onnx_ir as ir +from onnxscript import OpBuilder, nn + +from mobius._configs._wan_vae import WanVAEConfig +from mobius.components import INT64_MAX +from mobius.components import Conv2d as _Conv2d +from mobius.components import SiLU as _SiLU + +if TYPE_CHECKING: + import torch + +__all__ = ["AutoencoderKLWanModel"] + +#: ``F.normalize`` clamps the L2 norm at this epsilon before dividing. +_NORMALIZE_EPS = 1e-12 + +#: ``DiagonalGaussianDistribution`` clamps ``logvar`` to this range. +_LOGVAR_MIN = -30.0 +_LOGVAR_MAX = 20.0 + + +def _dim(op: OpBuilder, x: ir.Value, axis: int) -> ir.Value: + """Return ``x.shape[axis]`` as a 1-D int64 tensor of length 1.""" + return op.Shape(x, start=axis, end=axis + 1) + + +def _dims5(op: OpBuilder, x: ir.Value) -> tuple[ir.Value, ...]: + """Return the five dimensions of a ``(B, C, T, H, W)`` tensor.""" + return tuple(_dim(op, x, axis) for axis in range(5)) + + +# --------------------------------------------------------------------------- +# Primitive layers +# --------------------------------------------------------------------------- + + +class _WanCausalConv3d(nn.Module): + """3D convolution with causal (left-only) padding on the temporal axis. + + Mirrors ``WanCausalConv3d``, which subclasses ``nn.Conv3d`` and rewrites the + padding as ``(W, W, H, H, 2 * T, 0)`` in ``F.pad`` order before delegating to + the unpadded convolution. Doubling the temporal padding and placing it all + on the left makes the convolution causal while preserving the frame count. + + Args: + in_channels: Input channel count. + out_channels: Output channel count. + kernel_size: ``(kT, kH, kW)`` or a single int applied to all three axes. + stride: ``(sT, sH, sW)`` or a single int. + padding: ``(pT, pH, pW)`` or a single int, in ``nn.Conv3d`` semantics. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int], + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] = 0, + ) -> None: + super().__init__() + kernel = (kernel_size,) * 3 if isinstance(kernel_size, int) else tuple(kernel_size) + strides = (stride,) * 3 if isinstance(stride, int) else tuple(stride) + pads = (padding,) * 3 if isinstance(padding, int) else tuple(padding) + + self.weight = nn.Parameter((out_channels, in_channels, *kernel)) + self.bias = nn.Parameter((out_channels,)) + self._kernel_shape = list(kernel) + self._strides = list(strides) + # ONNX Pad order for a 5D (B, C, T, H, W) tensor: + # [B_beg, C_beg, T_beg, H_beg, W_beg, B_end, C_end, T_end, H_end, W_end]. + # Temporal padding is 2 * pT on the left and 0 on the right (causal). + self._pads = [0, 0, 2 * pads[0], pads[1], pads[2], 0, 0, 0, pads[1], pads[2]] + self._needs_pad = any(self._pads) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # x: (B, C_in, T, H, W) + if self._needs_pad: + # Pad's constant_value defaults to 0 for any float dtype. + x = op.Pad(x, self._pads) + return op.Conv( + x, + self.weight, + self.bias, + kernel_shape=self._kernel_shape, + strides=self._strides, + pads=[0, 0, 0, 0, 0, 0], + dilations=[1, 1, 1], + group=1, + ) + + +class _WanRMSNorm(nn.Module): + """Channel-wise RMS normalisation (``WanRMS_norm``). + + Computes ``F.normalize(x, dim=1) * dim ** 0.5 * gamma``. ``F.normalize`` + divides by ``max(||x||_2, 1e-12)`` along the channel axis. Upstream forces + the normalisation itself to float32 for half-precision inputs and casts back + before applying ``scale``/``gamma``; the ``Cast``/``CastLike`` pair below + reproduces that and folds away for float32 graphs. + + Upstream's optional ``bias`` term is never enabled by ``AutoencoderKLWan`` + (every call site uses the default ``bias=False``), so no bias is emitted. + + Args: + dim: Channel count. + images: ``True`` broadcasts ``gamma`` over ``(H, W)`` for 4D activations + (the attention block); ``False`` broadcasts over ``(T, H, W)`` for + 5D activations. This matches the upstream parameter shapes + ``(dim, 1, 1)`` and ``(dim, 1, 1, 1)`` respectively. + """ + + def __init__(self, dim: int, images: bool = True) -> None: + super().__init__() + broadcast = (1, 1) if images else (1, 1, 1) + self.gamma = nn.Parameter((dim, *broadcast)) + self._scale = dim**0.5 + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # Normalise in float32, matching WanRMS_norm's fp16/bf16 upcast. + x_f32 = op.Cast(x, to=ir.DataType.FLOAT) + norm = op.ReduceL2(x_f32, [1], keepdims=1) + normalized = op.CastLike(op.Div(x_f32, op.Max(norm, _NORMALIZE_EPS)), x) + scaled = op.Mul(normalized, op.CastLike(self._scale, x)) + return op.Mul(scaled, self.gamma) + + +class _ZeroPad2d(nn.Module): + """``nn.ZeroPad2d((0, 1, 0, 1))`` — one row/column of zeros at bottom/right. + + Emitted as index ``0`` of the ``resample`` ``Sequential`` so the following + ``Conv2d`` keeps the HuggingFace name ``resample.1.{weight,bias}``. + """ + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # x: (N, C, H, W) -> pads [N_beg, C_beg, H_beg, W_beg, N_end, C_end, H_end, W_end] + return op.Pad(x, [0, 0, 0, 0, 0, 0, 1, 1]) + + +class _NearestUpsample2d(nn.Module): + """2x nearest-neighbour spatial upsampling (``WanUpsample``). + + Upstream uses ``mode="nearest-exact"``, which maps output index ``i`` to + ``floor((i + 0.5) / 2)``. For an exact integer scale that is identical to + ``floor(i / 2)``, which ONNX expresses as ``asymmetric`` coordinate + transformation with ``nearest_mode="floor"``. + """ + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # x: (N, C, H, W) -> (N, C, 2H, 2W) + return op.Resize( + x, + None, + [1.0, 1.0, 2.0, 2.0], + mode="nearest", + coordinate_transformation_mode="asymmetric", + nearest_mode="floor", + ) + + +# --------------------------------------------------------------------------- +# Residual shortcuts (parameter-free) +# --------------------------------------------------------------------------- + + +class _AvgDown3D(nn.Module): + """Grouped average-pooling shortcut for Wan 2.2 encoder stages (``AvgDown3D``). + + Folds ``factor_t`` frames and ``factor_s x factor_s`` spatial positions into + the channel axis, then averages consecutive groups of ``group_size`` + channels. The temporal axis is front-padded with zeros so its length is a + multiple of ``factor_t``, exactly as upstream's ``F.pad(x, (0, 0, 0, 0, pad_t, 0))``. + + Args: + in_channels: Shortcut input channels. + out_channels: Shortcut output channels. + factor_t: Temporal reduction factor (1 or 2). + factor_s: Spatial reduction factor (1 or 2). + """ + + def __init__( + self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1 + ) -> None: + super().__init__() + factor = factor_t * factor_s * factor_s + if in_channels * factor % out_channels != 0: + raise ValueError( + f"AvgDown3D requires in_channels * factor ({in_channels} * {factor}) " + f"to be divisible by out_channels ({out_channels})" + ) + self._out_channels = out_channels + self._factor_t = factor_t + self._factor_s = factor_s + self._factor = factor + self._group_size = in_channels * factor // out_channels + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # x: (B, C, T, H, W) + if self._factor == 1 and self._group_size == 1: + # No folding and a single-element group mean: algebraically the identity. + return x + + if self._factor_t > 1: + # pad_t = (factor_t - T % factor_t) % factor_t, applied at the front of T. + t_len = _dim(op, x, 2) + pad_t = op.Mod( + op.Sub(self._factor_t, op.Mod(t_len, self._factor_t)), self._factor_t + ) + pads = op.Concat( + op.Constant(value_ints=[0, 0]), + pad_t, + op.Constant(value_ints=[0, 0, 0, 0, 0, 0, 0]), + axis=0, + ) + x = op.Pad(x, pads) + + batch, channels, t_len, height, width = _dims5(op, x) + t_out = op.Div(t_len, self._factor_t) + h_out = op.Div(height, self._factor_s) + w_out = op.Div(width, self._factor_s) + ft = op.Constant(value_ints=[self._factor_t]) + fs = op.Constant(value_ints=[self._factor_s]) + + # (B, C, T/ft, ft, H/fs, fs, W/fs, fs) + x = op.Reshape(x, op.Concat(batch, channels, t_out, ft, h_out, fs, w_out, fs, axis=0)) + # -> (B, C, ft, fs, fs, T/ft, H/fs, W/fs) + x = op.Transpose(x, perm=[0, 1, 3, 5, 7, 2, 4, 6]) + # Flatten (C, ft, fs, fs) == (out_channels, group_size) and average the group. + x = op.Reshape( + x, + op.Concat( + batch, + op.Constant(value_ints=[self._out_channels, self._group_size]), + t_out, + h_out, + w_out, + axis=0, + ), + ) + # (B, out_channels, T/ft, H/fs, W/fs) + return op.ReduceMean(x, [2], keepdims=0) + + +class _DupUp3D(nn.Module): + """Channel-duplication shortcut for Wan 2.2 decoder stages (``DupUp3D``). + + Repeats each input channel ``repeats`` times, then unfolds those copies into + ``factor_t`` temporal and ``factor_s x factor_s`` spatial positions. + + Args: + in_channels: Shortcut input channels. + out_channels: Shortcut output channels. + factor_t: Temporal expansion factor (1 or 2). + factor_s: Spatial expansion factor (1 or 2). + """ + + def __init__( + self, in_channels: int, out_channels: int, factor_t: int, factor_s: int = 1 + ) -> None: + super().__init__() + factor = factor_t * factor_s * factor_s + if out_channels * factor % in_channels != 0: + raise ValueError( + f"DupUp3D requires out_channels * factor ({out_channels} * {factor}) " + f"to be divisible by in_channels ({in_channels})" + ) + self._out_channels = out_channels + self._factor_t = factor_t + self._factor_s = factor_s + self._repeats = out_channels * factor // in_channels + + def forward(self, op: OpBuilder, x: ir.Value, first_chunk: bool = True) -> ir.Value: + # x: (B, C, T, H, W) + batch, channels, t_len, height, width = _dims5(op, x) + + # repeat_interleave(repeats, dim=1): (B, C, 1, T, H, W) -> (B, C, r, T, H, W) + x = op.Expand( + op.Unsqueeze(x, [2]), + op.Concat( + batch, + channels, + op.Constant(value_ints=[self._repeats]), + t_len, + height, + width, + axis=0, + ), + ) + # -> (B, out_channels, ft, fs, fs, T, H, W) + x = op.Reshape( + x, + op.Concat( + batch, + op.Constant( + value_ints=[ + self._out_channels, + self._factor_t, + self._factor_s, + self._factor_s, + ] + ), + t_len, + height, + width, + axis=0, + ), + ) + # -> (B, out_channels, T, ft, H, fs, W, fs) + x = op.Transpose(x, perm=[0, 1, 5, 2, 6, 3, 7, 4]) + # -> (B, out_channels, T * ft, H * fs, W * fs) + x = op.Reshape( + x, + op.Concat( + batch, + op.Constant(value_ints=[self._out_channels]), + op.Mul(t_len, self._factor_t), + op.Mul(height, self._factor_s), + op.Mul(width, self._factor_s), + axis=0, + ), + ) + if first_chunk and self._factor_t > 1: + # Upstream drops the leading (factor_t - 1) duplicated frames of the + # very first chunk so the shortcut lines up with the main path, whose + # first latent frame is *not* temporally upsampled. + x = op.Slice(x, [self._factor_t - 1], [INT64_MAX], [2]) + return x + + +# --------------------------------------------------------------------------- +# Resampling +# --------------------------------------------------------------------------- + + +class _WanResample(nn.Module): + """Spatial (and optionally temporal) resampling block (``WanResample``). + + ``resample`` is a two-entry ``Sequential`` whose ``Conv2d`` sits at index 1, + matching the HuggingFace parameter names ``resample.1.{weight,bias}``. The + convolution is applied per frame by folding ``T`` into the batch axis. + + Args: + dim: Channel count entering the block. + mode: One of ``upsample2d``, ``upsample3d``, ``downsample2d``, + ``downsample3d``. + upsample_out_dim: Output channels of the upsampling convolution. + ``None`` means ``dim // 2`` (the Wan 2.1 default); the Wan 2.2 + residual decoder passes ``out_dim`` explicitly. + """ + + def __init__(self, dim: int, mode: str, upsample_out_dim: int | None = None) -> None: + super().__init__() + if mode not in ("upsample2d", "upsample3d", "downsample2d", "downsample3d"): + raise ValueError(f"Unsupported WanResample mode {mode!r}") + self._mode = mode + if upsample_out_dim is None: + upsample_out_dim = dim // 2 + + if mode in ("upsample2d", "upsample3d"): + self.resample = nn.Sequential( + _NearestUpsample2d(), + _Conv2d(dim, upsample_out_dim, kernel_size=3, padding=1), + ) + if mode == "upsample3d": + self.time_conv = _WanCausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + else: + self.resample = nn.Sequential( + _ZeroPad2d(), + _Conv2d(dim, dim, kernel_size=3, stride=2, padding=0), + ) + if mode == "downsample3d": + self.time_conv = _WanCausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) + ) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # x: (B, C, T, H, W) + if self._mode == "upsample3d": + x = self._temporal_upsample(op, x) + x = self._spatial_resample(op, x) + if self._mode == "downsample3d": + x = self._temporal_downsample(op, x) + return x + + def _temporal_upsample(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Duplicate every frame after the first via ``time_conv`` + interleave. + + Whole-sequence form of upstream's per-chunk ``upsample3d``: + ``concat(x[:, :, :1], interleave(time_conv(x[:, :, 1:])))``, i.e. frame 0 + passes through untouched (upstream's first decode chunk never runs + ``time_conv``) and the causal sequence restarts, zero-padded, at frame 1. + + ``x[:, :, 1:]`` is empty for a single latent frame, which no ``Conv`` + kernel accepts. Instead the temporal axis is zero-padded by one frame on + the right *before* dropping frame 0, so the convolution always sees + ``T >= 1`` frames; the surplus pair of interleaved output frames is then + sliced off. ``time_conv`` is causal, so the retained frames are exactly + those of the unpadded convolution. + """ + batch, channels, t_len, height, width = _dims5(op, x) + first = op.Slice(x, [0], [1], [2]) + + # (B, C, T, H, W) -> (B, C, T + 1, H, W) -> drop frame 0 -> T frames. + padded = op.Pad(x, [0, 0, 0, 0, 0, 0, 0, 1, 0, 0]) + rest = op.Slice(padded, [1], [INT64_MAX], [2]) + + # (B, C, T, H, W) -> (B, 2C, T, H, W) + doubled = self.time_conv(op, rest) + # Split the 2C channels into the two interleaved half-frames. + even, odd = op.Split(doubled, num_outputs=2, axis=1, _outputs=2) + # Interleave along a new axis: (B, C, T, 2, H, W) -> (B, C, 2T, H, W) + stacked = op.Concat(op.Unsqueeze(even, [3]), op.Unsqueeze(odd, [3]), axis=3) + rest = op.Reshape( + stacked, + op.Concat(batch, channels, op.Mul(t_len, 2), height, width, axis=0), + ) + # (B, C, 1 + 2T, H, W) -> keep 2T - 1 frames (== 1 for a single frame). + out = op.Concat(first, rest, axis=2) + return op.Slice(out, [0], op.Sub(op.Mul(t_len, 2), 1), [2]) + + def _temporal_downsample(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Stride-2 causal-free temporal convolution, keeping frame 0 verbatim. + + Upstream skips ``time_conv`` for chunk 0 and, from chunk 1 on, prepends + the previous chunk's last frame; concatenated, the kernel-3/stride-2 + windows start at global frames ``0, 2, 4, ...`` over the whole sequence, + i.e. ``concat(x[:, :, :1], time_conv(x))``. + + A single-frame clip is shorter than the kernel, so the temporal axis is + zero-padded by two frames on the right and the surplus trailing output + frame is sliced off. The convolution only ever looks backwards in time, + so the retained frames match the unpadded convolution exactly. + """ + t_len = _dim(op, x, 2) + first = op.Slice(x, [0], [1], [2]) + + # (B, C, T, H, W) -> (B, C, T + 2, H, W) -> stride-2 conv -> (T - 1) // 2 + 1. + padded = op.Pad(x, [0, 0, 0, 0, 0, 0, 0, 2, 0, 0]) + strided = self.time_conv(op, padded) + # Keep 1 + (T - 1) // 2 frames, dropping the one window that read padding. + out = op.Concat(first, strided, axis=2) + keep = op.Add(op.Div(op.Sub(t_len, 1), 2), 1) + return op.Slice(out, [0], keep, [2]) + + def _spatial_resample(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Apply the 2D ``resample`` Sequential frame-by-frame.""" + batch, channels, t_len, height, width = _dims5(op, x) + # (B, C, T, H, W) -> (B*T, C, H, W) + frames = op.Reshape( + op.Transpose(x, perm=[0, 2, 1, 3, 4]), + op.Concat(op.Mul(batch, t_len), channels, height, width, axis=0), + ) + frames = self.resample(op, frames) + out_c = _dim(op, frames, 1) + out_h = _dim(op, frames, 2) + out_w = _dim(op, frames, 3) + # (B*T, C', H', W') -> (B, C', T, H', W') + return op.Transpose( + op.Reshape(frames, op.Concat(batch, t_len, out_c, out_h, out_w, axis=0)), + perm=[0, 2, 1, 3, 4], + ) + + +# --------------------------------------------------------------------------- +# Residual / attention blocks +# --------------------------------------------------------------------------- + + +class _WanResidualBlock(nn.Module): + """Pre-norm residual block with two causal 3D convolutions (``WanResidualBlock``). + + ``norm1 -> SiLU -> conv1 -> norm2 -> SiLU -> conv2`` plus a shortcut that is + a 1x1x1 causal convolution when the channel count changes and the identity + otherwise (upstream uses ``nn.Identity``, which contributes no weights). + + Upstream's ``nn.Dropout`` between ``norm2`` and ``conv2`` is an identity in + eval mode and therefore not emitted. + + Args: + in_dim: Input channels. + out_dim: Output channels. + """ + + def __init__(self, in_dim: int, out_dim: int) -> None: + super().__init__() + self.norm1 = _WanRMSNorm(in_dim, images=False) + self.conv1 = _WanCausalConv3d(in_dim, out_dim, 3, padding=1) + self.norm2 = _WanRMSNorm(out_dim, images=False) + self.conv2 = _WanCausalConv3d(out_dim, out_dim, 3, padding=1) + self.conv_shortcut = ( + _WanCausalConv3d(in_dim, out_dim, 1) if in_dim != out_dim else None + ) + self.nonlinearity = _SiLU() + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # x: (B, C_in, T, H, W) + shortcut = x if self.conv_shortcut is None else self.conv_shortcut(op, x) + x = self.conv1(op, self.nonlinearity(op, self.norm1(op, x))) + x = self.conv2(op, self.nonlinearity(op, self.norm2(op, x))) + return op.Add(x, shortcut) + + +class _WanAttentionBlock(nn.Module): + """Single-head spatial self-attention over each frame (``WanAttentionBlock``). + + Frames are folded into the batch axis, so attention is computed + independently per frame over the ``H * W`` spatial positions. + + Args: + dim: Channel count (also the single head's dimension). + """ + + def __init__(self, dim: int) -> None: + super().__init__() + self.norm = _WanRMSNorm(dim, images=True) + self.to_qkv = _Conv2d(dim, dim * 3, kernel_size=1) + self.proj = _Conv2d(dim, dim, kernel_size=1) + self._scale = dim**-0.5 + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # x: (B, C, T, H, W) + identity = x + batch, channels, t_len, height, width = _dims5(op, x) + batch_frames = op.Mul(batch, t_len) + + # (B, C, T, H, W) -> (B*T, C, H, W) + frames = op.Reshape( + op.Transpose(x, perm=[0, 2, 1, 3, 4]), + op.Concat(batch_frames, channels, height, width, axis=0), + ) + frames = self.norm(op, frames) + + # (B*T, 3C, H, W) -> (B*T, 1, H*W, 3C), then split into q, k, v + qkv = self.to_qkv(op, frames) + qkv = op.Reshape( + qkv, + op.Concat( + batch_frames, + op.Constant(value_ints=[1]), + op.Mul(channels, 3), + op.Constant(value_ints=[-1]), + axis=0, + ), + ) + qkv = op.Transpose(qkv, perm=[0, 1, 3, 2]) + query, key, value = op.Split(qkv, num_outputs=3, axis=-1, _outputs=3) + + # Single-head scaled dot-product attention -> (B*T, 1, H*W, C) + attn = op.Attention( + query, key, value, scale=self._scale, q_num_heads=1, kv_num_heads=1 + ) + + # (B*T, 1, H*W, C) -> (B*T, C, H, W) + attn = op.Transpose( + op.Reshape(attn, op.Concat(batch_frames, op.Mul(height, width), channels, axis=0)), + perm=[0, 2, 1], + ) + attn = op.Reshape(attn, op.Concat(batch_frames, channels, height, width, axis=0)) + attn = self.proj(op, attn) + + # (B*T, C, H, W) -> (B, C, T, H, W) + out = op.Transpose( + op.Reshape(attn, op.Concat(batch, t_len, channels, height, width, axis=0)), + perm=[0, 2, 1, 3, 4], + ) + return op.Add(out, identity) + + +class _WanMidBlock(nn.Module): + """Bottleneck block: ``resnet -> (attention -> resnet) * num_layers``. + + Args: + dim: Channel count. + num_layers: Number of attention/resnet pairs after the first resnet. + """ + + def __init__(self, dim: int, num_layers: int = 1) -> None: + super().__init__() + self.attentions = nn.ModuleList([_WanAttentionBlock(dim) for _ in range(num_layers)]) + self.resnets = nn.ModuleList( + [_WanResidualBlock(dim, dim) for _ in range(num_layers + 1)] + ) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + x = self.resnets[0](op, x) + for index, attention in enumerate(self.attentions): + x = attention(op, x) + x = self.resnets[index + 1](op, x) + return x + + +# --------------------------------------------------------------------------- +# Encoder / decoder stages +# --------------------------------------------------------------------------- + + +class _WanResidualDownBlock(nn.Module): + """Wan 2.2 encoder stage: residual blocks + downsampler + ``AvgDown3D`` shortcut. + + Args: + in_dim: Stage input channels. + out_dim: Stage output channels. + num_res_blocks: Residual blocks in the main path. + temporal_downsample: Halve the temporal axis as well as the spatial axes. + down_flag: Whether this stage downsamples at all (the last stage does not). + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + temporal_downsample: bool = False, + down_flag: bool = False, + ) -> None: + super().__init__() + self.avg_shortcut = _AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temporal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + resnets = [] + current_dim = in_dim + for _ in range(num_res_blocks): + resnets.append(_WanResidualBlock(current_dim, out_dim)) + current_dim = out_dim + self.resnets = nn.ModuleList(resnets) + self.downsampler = ( + _WanResample(out_dim, "downsample3d" if temporal_downsample else "downsample2d") + if down_flag + else None + ) + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # x: (B, C_in, T, H, W) + residual = x + for resnet in self.resnets: + x = resnet(op, x) + if self.downsampler is not None: + x = self.downsampler(op, x) + return op.Add(x, self.avg_shortcut(op, residual)) + + +class _WanResidualUpBlock(nn.Module): + """Wan 2.2 decoder stage: residual blocks + upsampler + ``DupUp3D`` shortcut. + + Args: + in_dim: Stage input channels. + out_dim: Stage output channels. + num_res_blocks: ``num_res_blocks + 1`` residual blocks are created, + matching upstream. + temporal_upsample: Double the temporal axis as well as the spatial axes. + up_flag: Whether this stage upsamples at all (the last stage does not). + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + temporal_upsample: bool = False, + up_flag: bool = False, + ) -> None: + super().__init__() + self.avg_shortcut = ( + _DupUp3D(in_dim, out_dim, factor_t=2 if temporal_upsample else 1, factor_s=2) + if up_flag + else None + ) + resnets = [] + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append(_WanResidualBlock(current_dim, out_dim)) + current_dim = out_dim + self.resnets = nn.ModuleList(resnets) + self.upsampler = ( + _WanResample( + out_dim, + "upsample3d" if temporal_upsample else "upsample2d", + upsample_out_dim=out_dim, + ) + if up_flag + else None + ) + + def forward(self, op: OpBuilder, x: ir.Value, first_chunk: bool = True) -> ir.Value: + # x: (B, C_in, T, H, W) + residual = x + for resnet in self.resnets: + x = resnet(op, x) + if self.upsampler is not None: + x = self.upsampler(op, x) + if self.avg_shortcut is not None: + x = op.Add(x, self.avg_shortcut(op, residual, first_chunk=first_chunk)) + return x + + +class _WanUpBlock(nn.Module): + """Wan 2.1 decoder stage: residual blocks + optional upsampler, no shortcut. + + Args: + in_dim: Stage input channels. + out_dim: Stage output channels. + num_res_blocks: ``num_res_blocks + 1`` residual blocks are created. + upsample_mode: ``"upsample2d"``, ``"upsample3d"`` or ``None``. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_res_blocks: int, + upsample_mode: str | None = None, + ) -> None: + super().__init__() + resnets = [] + current_dim = in_dim + for _ in range(num_res_blocks + 1): + resnets.append(_WanResidualBlock(current_dim, out_dim)) + current_dim = out_dim + self.resnets = nn.ModuleList(resnets) + # Upstream stores this as a one-element ModuleList named ``upsamplers``. + self.upsamplers = ( + nn.ModuleList([_WanResample(out_dim, upsample_mode)]) + if upsample_mode is not None + else None + ) + + def forward(self, op: OpBuilder, x: ir.Value, first_chunk: bool = True) -> ir.Value: + # ``first_chunk`` is accepted for signature parity with + # ``_WanResidualUpBlock`` (upstream's ``WanUpBlock.forward`` does the same); + # Wan 2.1 has no ``DupUp3D`` shortcut, so it has no effect here. + del first_chunk + for resnet in self.resnets: + x = resnet(op, x) + if self.upsamplers is not None: + x = self.upsamplers[0](op, x) + return x + + +class _WanEncoder3d(nn.Module): + """3D encoder (``WanEncoder3d``): video -> ``2 * z_dim`` posterior moments. + + Note: + The Wan 2.1 (``is_residual=False``) branch inserts a + :class:`_WanAttentionBlock` after any residual block whose spatial scale + appears in ``attn_scales``. That branch is currently unreachable in + diffusers itself — ``WanEncoder3d.forward`` forwards ``feat_cache`` to + every ``down_blocks`` entry, and ``WanAttentionBlock.forward`` does not + accept it — so it has no runnable upstream reference to compare against. + The layout and weight names implemented here follow the upstream + constructor exactly. + + Args: + config: Parsed :class:`~mobius._configs._wan_vae.WanVAEConfig`. + """ + + def __init__(self, config: WanVAEConfig) -> None: + super().__init__() + dims = config.encoder_dims + last_stage = len(config.dim_mult) - 1 + + self.conv_in = _WanCausalConv3d(config.in_channels, dims[0], 3, padding=1) + + self.down_blocks = nn.ModuleList([]) + scale = 1.0 + for i, (in_dim, out_dim) in enumerate(itertools.pairwise(dims)): + if config.is_residual: + self.down_blocks.append( + _WanResidualDownBlock( + in_dim, + out_dim, + config.num_res_blocks, + temporal_downsample=( + config.temporal_downsample[i] if i != last_stage else False + ), + down_flag=i != last_stage, + ) + ) + continue + # Wan 2.1: flat list of residual (and optional attention) blocks. + stage_in = in_dim + for _ in range(config.num_res_blocks): + self.down_blocks.append(_WanResidualBlock(stage_in, out_dim)) + if scale in config.attn_scales: + self.down_blocks.append(_WanAttentionBlock(out_dim)) + stage_in = out_dim + if i != last_stage: + mode = "downsample3d" if config.temporal_downsample[i] else "downsample2d" + self.down_blocks.append(_WanResample(out_dim, mode=mode)) + scale /= 2.0 + + self.mid_block = _WanMidBlock(dims[-1], num_layers=1) + self.norm_out = _WanRMSNorm(dims[-1], images=False) + # The encoder emits mean ‖ logvar, hence 2 * z_dim output channels. + self.conv_out = _WanCausalConv3d(dims[-1], config.z_dim * 2, 3, padding=1) + self.nonlinearity = _SiLU() + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + # x: (B, in_channels, T, H, W) + x = self.conv_in(op, x) + for block in self.down_blocks: + x = block(op, x) + x = self.mid_block(op, x) + x = self.nonlinearity(op, self.norm_out(op, x)) + # (B, 2 * z_dim, T', H', W') + return self.conv_out(op, x) + + +class _WanDecoder3d(nn.Module): + """3D decoder (``WanDecoder3d``): ``z_dim`` latents -> ``out_channels`` frames. + + Args: + config: Parsed :class:`~mobius._configs._wan_vae.WanVAEConfig`. + """ + + def __init__(self, config: WanVAEConfig) -> None: + super().__init__() + dims = config.decoder_dims + temporal_upsample = config.temporal_upsample + last_stage = len(config.dim_mult) - 1 + + self.conv_in = _WanCausalConv3d(config.z_dim, dims[0], 3, padding=1) + self.mid_block = _WanMidBlock(dims[0], num_layers=1) + + self.up_blocks = nn.ModuleList([]) + for i, (in_dim, out_dim) in enumerate(itertools.pairwise(dims)): + up_flag = i != last_stage + if config.is_residual: + self.up_blocks.append( + _WanResidualUpBlock( + in_dim, + out_dim, + config.num_res_blocks, + temporal_upsample=temporal_upsample[i] if up_flag else False, + up_flag=up_flag, + ) + ) + continue + # Wan 2.1 halves the incoming width from the second stage on, because + # its upsampling convolution emits ``dim // 2`` channels. + stage_in = in_dim // 2 if i > 0 else in_dim + upsample_mode = None + if up_flag: + upsample_mode = "upsample3d" if temporal_upsample[i] else "upsample2d" + self.up_blocks.append( + _WanUpBlock(stage_in, out_dim, config.num_res_blocks, upsample_mode) + ) + + self.norm_out = _WanRMSNorm(dims[-1], images=False) + self.conv_out = _WanCausalConv3d(dims[-1], config.out_channels, 3, padding=1) + self.nonlinearity = _SiLU() + + def forward(self, op: OpBuilder, x: ir.Value, first_chunk: bool = True) -> ir.Value: + # x: (B, z_dim, T, H, W) + x = self.conv_in(op, x) + x = self.mid_block(op, x) + for block in self.up_blocks: + x = block(op, x, first_chunk=first_chunk) + x = self.nonlinearity(op, self.norm_out(op, x)) + # (B, out_channels, T'', H'', W'') + return self.conv_out(op, x) + + +# --------------------------------------------------------------------------- +# Full autoencoder +# --------------------------------------------------------------------------- + + +class AutoencoderKLWanModel(nn.Module): + """Wan 3D causal video VAE (``diffusers.AutoencoderKLWan``). + + Exposes the four HuggingFace sub-modules (``encoder``, ``quant_conv``, + ``post_quant_conv``, ``decoder``) plus the pipeline-level helpers that live + *outside* those sub-modules upstream: :meth:`patchify` / :meth:`unpatchify` + (applied in ``AutoencoderKLWan._encode`` / ``._decode``) and + :meth:`normalize_latents` / :meth:`denormalize_latents` (applied by the + ``WanPipeline``, not by the VAE itself). Keeping them off the sub-modules + is what makes the exported initializer names match the checkpoint exactly. + + Args: + config: Parsed :class:`~mobius._configs._wan_vae.WanVAEConfig`. + """ + + default_task: str = "wan-vae" + category: str = "autoencoder" + + def __init__(self, config: WanVAEConfig) -> None: + super().__init__() + self.config = config + self.encoder = _WanEncoder3d(config) + self.quant_conv = _WanCausalConv3d(config.z_dim * 2, config.z_dim * 2, 1) + self.post_quant_conv = _WanCausalConv3d(config.z_dim, config.z_dim, 1) + self.decoder = _WanDecoder3d(config) + + # ------------------------------------------------------------------ + # Patch folding (``AutoencoderKLWan._encode`` / ``._decode``) + # ------------------------------------------------------------------ + + def patchify(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Fold ``patch_size x patch_size`` spatial patches into the channel axis. + + Args: + op: The ONNX op builder. + x: Video tensor ``(B, C, T, H, W)``. + + Returns: + ``(B, C * p * p, T, H / p, W / p)``, or *x* unchanged when + ``patch_size`` is ``None`` or 1. + """ + patch = self.config.patch_size + if patch is None or patch == 1: + return x + batch, channels, t_len, height, width = _dims5(op, x) + p = op.Constant(value_ints=[patch]) + # (B, C, T, H/p, p, W/p, p) + x = op.Reshape( + x, + op.Concat( + batch, + channels, + t_len, + op.Div(height, patch), + p, + op.Div(width, patch), + p, + axis=0, + ), + ) + # -> (B, C, p_w, p_h, T, H/p, W/p) + x = op.Transpose(x, perm=[0, 1, 6, 4, 2, 3, 5]) + return op.Reshape( + x, + op.Concat( + batch, + op.Mul(channels, patch * patch), + t_len, + op.Div(height, patch), + op.Div(width, patch), + axis=0, + ), + ) + + def unpatchify(self, op: OpBuilder, x: ir.Value) -> ir.Value: + """Unfold channel-packed patches back into spatial resolution. + + Args: + op: The ONNX op builder. + x: Decoder output ``(B, C * p * p, T, H, W)``. + + Returns: + ``(B, C, T, H * p, W * p)``, or *x* unchanged when ``patch_size`` + is ``None`` or 1. + """ + patch = self.config.patch_size + if patch is None or patch == 1: + return x + batch, channels, t_len, height, width = _dims5(op, x) + p = op.Constant(value_ints=[patch]) + # (B, C, p_w, p_h, T, H, W) + x = op.Reshape( + x, + op.Concat( + batch, op.Div(channels, patch * patch), p, p, t_len, height, width, axis=0 + ), + ) + # -> (B, C, T, H, p_h, W, p_w) + x = op.Transpose(x, perm=[0, 1, 4, 5, 3, 6, 2]) + return op.Reshape( + x, + op.Concat( + batch, + op.Div(channels, patch * patch), + t_len, + op.Mul(height, patch), + op.Mul(width, patch), + axis=0, + ), + ) + + # ------------------------------------------------------------------ + # Latent statistics (``WanPipeline``-level, not part of the VAE weights) + # ------------------------------------------------------------------ + + def _latent_stat(self, op: OpBuilder, values: tuple[float, ...], name: str) -> ir.Value: + """Materialise a per-channel latent statistic as a ``(1, z, 1, 1, 1)`` constant.""" + array = np.asarray(values, dtype=np.float32).reshape(1, len(values), 1, 1, 1) + dtype = self.config.dtype + tensor = ir.tensor(array.astype(dtype.numpy()), dtype=dtype, name=name) + return op.initializer(tensor, name) + + def normalize_latents(self, op: OpBuilder, latents: ir.Value) -> ir.Value: + """Apply ``(z - latents_mean) / latents_std``. + + This is the diffusers ``WanPipeline`` convention (which stores the + reciprocal in a variable also called ``latents_std`` and multiplies). + It happens outside ``AutoencoderKLWan`` upstream, so it is emitted at the + graph boundary rather than inside :attr:`encoder`. + + Args: + op: The ONNX op builder. + latents: Raw posterior latents ``(B, z_dim, T, H, W)``. + + Returns: + Normalised latents with the same shape. + """ + mean = self._latent_stat(op, self.config.latents_mean, "latents_mean") + std = self._latent_stat(op, self.config.latents_std, "latents_std") + return op.Div(op.Sub(latents, mean), std) + + def denormalize_latents(self, op: OpBuilder, latents: ir.Value) -> ir.Value: + """Apply ``z * latents_std + latents_mean`` (inverse of :meth:`normalize_latents`). + + Args: + op: The ONNX op builder. + latents: Normalised latents ``(B, z_dim, T, H, W)``. + + Returns: + Raw latents in the VAE's own scale, ready for ``post_quant_conv``. + """ + mean = self._latent_stat(op, self.config.latents_mean, "latents_mean") + std = self._latent_stat(op, self.config.latents_std, "latents_std") + return op.Add(op.Mul(latents, std), mean) + + # ------------------------------------------------------------------ + # Encode / decode + # ------------------------------------------------------------------ + + def encode(self, op: OpBuilder, sample: ir.Value) -> tuple[ir.Value, ir.Value]: + """Encode a video into deterministic posterior moments. + + Mirrors ``AutoencoderKLWan._encode`` followed by + ``DiagonalGaussianDistribution``: patchify, run the encoder, apply + ``quant_conv``, split the ``2 * z_dim`` channels into mean and logvar and + clamp the logvar to ``[-30, 20]``. No sampling happens in the graph. + + Args: + op: The ONNX op builder. + sample: Video tensor ``(B, video_channels, T, H, W)`` in ``[-1, 1]``. + + Returns: + ``(mean, logvar)``, each ``(B, z_dim, T', H', W')``. + """ + hidden = self.patchify(op, sample) + hidden = self.encoder(op, hidden) + moments = self.quant_conv(op, hidden) + mean, logvar = op.Split(moments, num_outputs=2, axis=1, _outputs=2) + logvar = op.Clip( + logvar, op.CastLike(_LOGVAR_MIN, logvar), op.CastLike(_LOGVAR_MAX, logvar) + ) + return mean, logvar + + def decode(self, op: OpBuilder, latents: ir.Value) -> ir.Value: + """Decode raw (un-normalised) latents into a video. + + Mirrors ``AutoencoderKLWan._decode``: ``post_quant_conv``, decoder, + unpatchify and the unconditional clamp to ``[-1, 1]``. + + Args: + op: The ONNX op builder. + latents: Raw latents ``(B, z_dim, T, H, W)``. + + Returns: + Video tensor ``(B, decoded_video_channels, 4 * (T - 1) + 1, H', W')``. + """ + hidden = self.post_quant_conv(op, latents) + hidden = self.decoder(op, hidden) + hidden = self.unpatchify(op, hidden) + # Upstream clamps unconditionally; ``clip_output`` from the public + # config is not a parameter of AutoencoderKLWan and is ignored there. + return op.Clip(hidden, op.CastLike(-1.0, hidden), op.CastLike(1.0, hidden)) + + def forward(self, op: OpBuilder, latents: ir.Value) -> ir.Value: + """Decode normalised latents into a video (the generation-time entry point).""" + return self.decode(op, self.denormalize_latents(op, latents)) + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Return the HuggingFace state dict unchanged. + + Every ``nn.Module`` attribute name here mirrors the corresponding + ``AutoencoderKLWan`` attribute, so the generated initializer names + (``encoder.down_blocks.1.resnets.0.conv1.weight``, + ``decoder.up_blocks.0.upsampler.resample.1.bias``, ...) already match the + checkpoint keys and no renaming is required. + """ + return state_dict diff --git a/src/mobius/models/wan_vae_test.py b/src/mobius/models/wan_vae_test.py new file mode 100644 index 000000000..0a6c2be93 --- /dev/null +++ b/src/mobius/models/wan_vae_test.py @@ -0,0 +1,1048 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Tests for the Wan 3D causal video VAE (``diffusers.AutoencoderKLWan``). + +Layers: + +* Config parsing / validation against the real ``nvidia/Cosmos3-Nano/vae`` + ``config.json`` (embedded verbatim below — no network access). +* L1 graph construction with a tiny config: I/O contract, node construction and + HuggingFace weight-name alignment for both the Wan 2.2 residual architecture + and the Wan 2.1 flat architecture. +* Numerical parity against ``diffusers`` PyTorch for the causal convolution, the + residual block, both temporal resampling modes (whole-sequence vs upstream's + chunked ``feat_cache`` loop) and the full encoder/decoder. +""" + +from __future__ import annotations + +import tempfile +from pathlib import Path + +import numpy as np +import onnx_ir as ir +import pytest + +from mobius._configs._wan_vae import WanVAEConfig +from mobius.models.wan_vae import AutoencoderKLWanModel +from mobius.tasks._base import _make_graph, _make_model +from mobius.tasks._wan_vae import WanVAETask + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +#: Verbatim ``vae/config.json`` of ``nvidia/Cosmos3-Nano`` (== Wan2.2-TI2V-5B). +COSMOS3_VAE_CONFIG: dict = { + "_class_name": "AutoencoderKLWan", + "_diffusers_version": "0.37.1", + "_name_or_path": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", + "attn_scales": [], + "base_dim": 160, + "clip_output": False, + "decoder_base_dim": 256, + "dim_mult": [1, 2, 4, 4], + "dropout": 0.0, + "in_channels": 12, + "is_residual": True, + "latents_mean": [ + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.157, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.123, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.052, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, + ], + "latents_std": [ + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.499, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.06, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, + ], + "num_res_blocks": 2, + "out_channels": 12, + "patch_size": 2, + "scale_factor_spatial": 16, + "scale_factor_temporal": 4, + "temperal_downsample": [False, True, True], + "z_dim": 48, +} + +#: Tiny Wan 2.2 (residual) config: 2 resample stages, 1 of them temporal. +TINY_WAN22: dict = { + "_class_name": "AutoencoderKLWan", + "attn_scales": [], + "base_dim": 8, + "clip_output": False, + "decoder_base_dim": 8, + "dim_mult": [1, 2, 2], + "dropout": 0.0, + "in_channels": 12, + "is_residual": True, + "latents_mean": [0.1, -0.2, 0.3, -0.4], + "latents_std": [1.1, 0.9, 1.2, 0.8], + "num_res_blocks": 1, + "out_channels": 12, + "patch_size": 2, + "scale_factor_spatial": 8, + "scale_factor_temporal": 2, + "temperal_downsample": [False, True], + "z_dim": 4, +} + +#: Tiny Wan 2.1 (flat, non-residual, un-patchified) config. +TINY_WAN21: dict = { + "_class_name": "AutoencoderKLWan", + "attn_scales": [], + "base_dim": 8, + "dim_mult": [1, 2, 4], + "dropout": 0.0, + "in_channels": 3, + "is_residual": False, + "latents_mean": [0.1, -0.2, 0.3, -0.4], + "latents_std": [1.1, 0.9, 1.2, 0.8], + "num_res_blocks": 1, + "out_channels": 3, + "scale_factor_spatial": 4, + "scale_factor_temporal": 2, + "temperal_downsample": [False, True], + "z_dim": 4, +} + +#: Cosmos3-shaped but tiny: 4 stages, **two** temporal resampling stages and +#: ``patch_size = 2``, i.e. the exact topology of ``nvidia/Cosmos3-Nano/vae`` +#: with the channel widths shrunk. Exercises chained temporal upsampling, which +#: the single-stage :data:`TINY_WAN22` cannot. +TINY_COSMOS3: dict = { + "_class_name": "AutoencoderKLWan", + "attn_scales": [], + "base_dim": 8, + "clip_output": False, + "decoder_base_dim": 8, + "dim_mult": [1, 2, 4, 4], + "dropout": 0.0, + "in_channels": 12, + "is_residual": True, + "latents_mean": [0.1, -0.2, 0.3, -0.4], + "latents_std": [1.1, 0.9, 1.2, 0.8], + "num_res_blocks": 1, + "out_channels": 12, + "patch_size": 2, + "scale_factor_spatial": 16, + "scale_factor_temporal": 4, + "temperal_downsample": [False, True, True], + "z_dim": 4, +} + + +def _build(config_dict: dict): + """Parse a config, instantiate the module and build both graphs.""" + config = WanVAEConfig.from_diffusers(config_dict) + module = AutoencoderKLWanModel(config) + return config, module, WanVAETask().build(module, config) + + +def _op_types(model: ir.Model) -> set[str]: + return {node.op_type for node in model.graph} + + +def _run(model: ir.Model, feeds: dict[str, np.ndarray]) -> list[np.ndarray]: + """Serialise *model* and run it through onnxruntime.""" + import onnxruntime as ort + + # Windows keeps the ORT model file mapped; ignore cleanup errors. + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: + path = Path(temp_dir) / "model.onnx" + ir.save(model, path) + session = ort.InferenceSession(path) + return session.run(None, feeds) + + +# --------------------------------------------------------------------------- +# Config parsing and validation +# --------------------------------------------------------------------------- + + +class TestConfig: + """Parsing and validation of the public ``AutoencoderKLWan`` config.""" + + def test_parses_real_cosmos3_config(self): + config = WanVAEConfig.from_diffusers(COSMOS3_VAE_CONFIG) + + assert config.z_dim == 48 + assert config.in_channels == 12 + assert config.out_channels == 12 + assert config.patch_size == 2 + assert config.base_dim == 160 + assert config.decoder_base_dim == 256 + assert config.dim_mult == (1, 2, 4, 4) + assert config.num_res_blocks == 2 + assert config.attn_scales == () + assert config.dropout == pytest.approx(0.0) + assert config.is_residual is True + assert config.temporal_downsample == (False, True, True) + assert config.scale_factor_spatial == 16 + assert config.scale_factor_temporal == 4 + assert len(config.latents_mean) == 48 + assert len(config.latents_std) == 48 + assert config.latents_mean[0] == pytest.approx(-0.2289) + assert config.latents_std[-1] == pytest.approx(0.7744) + assert config.dtype == ir.DataType.FLOAT + + def test_derived_shapes_for_cosmos3(self): + config = WanVAEConfig.from_diffusers(COSMOS3_VAE_CONFIG) + + # in_channels counts patchified channels: 3 RGB * 2 * 2. + assert config.video_channels == 3 + assert config.decoded_video_channels == 3 + assert config.encoder_dims == (160, 160, 320, 640, 640) + assert config.decoder_dims == (1024, 1024, 1024, 512, 256) + assert config.temporal_upsample == (True, True, False) + + def test_misspelled_temperal_downsample_is_accepted_and_aliased(self): + config = WanVAEConfig.from_diffusers(COSMOS3_VAE_CONFIG) + + # Upstream's misspelling is the parsed key; the dataclass exposes the + # corrected name and keeps the misspelled one as a read-only alias. + assert config.temporal_downsample == (False, True, True) + assert config.temperal_downsample == config.temporal_downsample + + def test_corrected_spelling_is_also_accepted(self): + raw = dict(COSMOS3_VAE_CONFIG) + raw.pop("temperal_downsample") + raw["temporal_downsample"] = [False, True, True] + + assert WanVAEConfig.from_diffusers(raw).temporal_downsample == (False, True, True) + + def test_clip_output_is_parsed_but_never_honoured(self): + # ``clip_output`` is not a parameter of AutoencoderKLWan.__init__, so + # diffusers drops it and always clamps. The graph must clamp too. + config = WanVAEConfig.from_diffusers(COSMOS3_VAE_CONFIG) + assert config.clip_output is False + + _, _, package = _build(TINY_WAN22) + assert "Clip" in _op_types(package["decoder"]) + + def test_decoder_base_dim_defaults_to_base_dim(self): + raw = dict(TINY_WAN21) + assert "decoder_base_dim" not in raw + + assert WanVAEConfig.from_diffusers(raw).decoder_base_dim == raw["base_dim"] + + @pytest.mark.parametrize( + ("dtype", "expected"), + [ + (None, ir.DataType.FLOAT), + ("auto", ir.DataType.FLOAT), + ("float32", ir.DataType.FLOAT), + ("float16", ir.DataType.FLOAT16), + ("torch.bfloat16", ir.DataType.BFLOAT16), + ], + ) + def test_dtype_resolution(self, dtype, expected): + raw = dict(TINY_WAN22, dtype=dtype) + + assert WanVAEConfig.from_diffusers(raw).dtype == expected + + def test_rejects_unknown_dtype(self): + with pytest.raises(ValueError, match="Unsupported dtype"): + WanVAEConfig.from_diffusers(dict(TINY_WAN22, dtype="int8")) + + def test_rejects_foreign_class_name(self): + raw = dict(TINY_WAN22, _class_name="AutoencoderKLQwenImage") + + with pytest.raises(ValueError, match="AutoencoderKLWan"): + WanVAEConfig.from_diffusers(raw) + + @pytest.mark.parametrize( + ("override", "message"), + [ + ({"temperal_downsample": [False, True, True]}, "len\\(dim_mult\\) - 1"), + ({"latents_mean": [0.0, 0.0]}, "latents_mean must have z_dim"), + ({"latents_std": [1.0, 1.0]}, "latents_std must have z_dim"), + ({"latents_std": [1.0, 0.0, 1.0, 1.0]}, "non-zero"), + ({"scale_factor_spatial": 4}, "scale_factor_spatial"), + ({"scale_factor_temporal": 4}, "scale_factor_temporal"), + ({"in_channels": 10}, "divisible by patch_size"), + ({"z_dim": 0}, "z_dim must be positive"), + ({"num_res_blocks": 0}, "num_res_blocks must be positive"), + ({"base_dim": 0}, "base_dim must be positive"), + ], + ) + def test_validation_rejects_inconsistent_fields(self, override, message): + with pytest.raises(ValueError, match=message): + WanVAEConfig.from_diffusers(dict(TINY_WAN22, **override)) + + def test_validation_rejects_bad_residual_shortcut_widths(self): + # AvgDown3D groups in_dim * factor channels into out_dim groups; a + # non-divisible pair cannot be expressed as a grouped mean. + raw = dict(TINY_WAN22, dim_mult=[1, 3, 3], scale_factor_spatial=8) + + with pytest.raises(ValueError, match="AvgDown3D"): + WanVAEConfig.from_diffusers(raw) + + +# --------------------------------------------------------------------------- +# L1 graph construction +# --------------------------------------------------------------------------- + + +class TestGraphConstruction: + """Graph I/O contract, node construction and weight-name alignment.""" + + def test_encoder_graph_io(self): + config, _, package = _build(TINY_WAN22) + graph = package["encoder"].graph + + assert [i.name for i in graph.inputs] == ["sample"] + sample = graph.inputs[0] + assert sample.dtype == ir.DataType.FLOAT + # Explicit 5D video: patchification happens inside the graph, so the + # input carries pixel-space (3-channel) video, not the 12 patch channels. + assert [str(d) for d in sample.shape] == [ + "batch", + str(config.video_channels), + "frames", + "height", + "width", + ] + + assert [o.name for o in graph.outputs] == [ + "latent_mean", + "latent_logvar", + "latent", + ] + for output in graph.outputs: + assert output.dtype == ir.DataType.FLOAT + assert [str(d) for d in output.shape] == [ + "batch", + str(config.z_dim), + "latent_frames", + "latent_height", + "latent_width", + ] + + def test_decoder_graph_io(self): + config, _, package = _build(TINY_WAN22) + graph = package["decoder"].graph + + assert [i.name for i in graph.inputs] == ["latent"] + latent = graph.inputs[0] + assert latent.dtype == ir.DataType.FLOAT + assert [str(d) for d in latent.shape] == [ + "batch", + str(config.z_dim), + "latent_frames", + "latent_height", + "latent_width", + ] + + assert [o.name for o in graph.outputs] == ["sample"] + sample = graph.outputs[0] + assert sample.dtype == ir.DataType.FLOAT + assert [str(d) for d in sample.shape] == [ + "batch", + str(config.decoded_video_channels), + "frames", + "height", + "width", + ] + + def test_package_keys_and_roles(self): + _, _, package = _build(TINY_WAN22) + + assert sorted(package.keys()) == ["decoder", "encoder"] + assert WanVAETask.model_roles == {"encoder": "encoder", "decoder": "decoder"} + + def test_graph_metadata_and_opset(self): + from mobius._constants import OPSET_VERSION + + _, _, package = _build(TINY_WAN22) + + for model in package.values(): + assert model.producer_name == "mobius" + assert model.graph.opset_imports[""] == OPSET_VERSION + + def test_encoder_node_construction(self): + _, _, package = _build(TINY_WAN22) + ops = _op_types(package["encoder"]) + + # Causal 3D convolution (Pad + Conv), spatial downsample (ZeroPad2d), + # AvgDown3D shortcut (Pad/Reshape/Transpose/ReduceMean), RMS norm + # (ReduceL2/Max/Div), SiLU (Sigmoid/Mul), single-head attention, and the + # posterior split with the logvar clamp. + assert {"Conv", "Pad", "Reshape", "Transpose", "Concat"} <= ops + assert {"ReduceL2", "Max", "Div", "Sigmoid", "Mul", "Add"} <= ops + assert "ReduceMean" in ops # AvgDown3D grouped mean + assert "Attention" in ops # mid-block single-head attention + assert {"Split", "Clip"} <= ops # mean/logvar split + logvar clamp + # No sampling in the graph: the posterior is returned deterministically. + assert "RandomNormalLike" not in ops + assert "RandomNormal" not in ops + + def test_decoder_node_construction(self): + _, _, package = _build(TINY_WAN22) + ops = _op_types(package["decoder"]) + + assert "Resize" in ops # nearest-exact 2x spatial upsample + assert "Expand" in ops # DupUp3D channel repeat_interleave + assert "Slice" in ops # first-frame passthrough / first_chunk drop + assert "Clip" in ops # unconditional clamp to [-1, 1] + assert "Attention" in ops + + def test_nearest_upsample_uses_floor_semantics(self): + _, _, package = _build(TINY_WAN22) + + resizes = [n for n in package["decoder"].graph if n.op_type == "Resize"] + assert resizes + for node in resizes: + assert node.attributes["mode"].as_string() == "nearest" + assert node.attributes["nearest_mode"].as_string() == "floor" + assert ( + node.attributes["coordinate_transformation_mode"].as_string() == "asymmetric" + ) + + def test_latent_statistics_are_graph_level_constants(self): + config, _, package = _build(TINY_WAN22) + + for key in ("encoder", "decoder"): + initializers = package[key].graph.initializers + assert "latents_mean" in initializers + assert "latents_std" in initializers + mean = initializers["latents_mean"].const_value + assert mean is not None + assert tuple(mean.shape) == (1, config.z_dim, 1, 1, 1) + np.testing.assert_allclose( + np.asarray(mean.numpy()).reshape(-1), config.latents_mean, rtol=1e-6 + ) + + def test_graph_dtype_follows_config(self): + from mobius._builder import _cast_module_dtype + + config = WanVAEConfig.from_diffusers(dict(TINY_WAN22, dtype="float16")) + module = AutoencoderKLWanModel(config) + # ``build_from_module`` casts parameters before building; do the same so + # ONNX type inference sees a uniformly half-precision graph. + _cast_module_dtype(module, config.dtype) + package = WanVAETask().build(module, config) + + for model in package.values(): + graph = model.graph + assert graph.inputs[0].dtype == ir.DataType.FLOAT16 + assert graph.outputs[0].dtype == ir.DataType.FLOAT16 + assert graph.initializers["latents_mean"].dtype == ir.DataType.FLOAT16 + assert graph.initializers["latents_std"].dtype == ir.DataType.FLOAT16 + + def test_weight_names_match_huggingface_layout(self): + _, _, package = _build(TINY_WAN22) + encoder = set(package["encoder"].graph.initializers) + decoder = set(package["decoder"].graph.initializers) + + # Encoder: residual down block with a temporal downsampler. + assert "encoder.conv_in.weight" in encoder + assert "encoder.down_blocks.0.resnets.0.norm1.gamma" in encoder + assert "encoder.down_blocks.0.resnets.0.conv1.weight" in encoder + # nn.Sequential(ZeroPad2d, Conv2d) -> the conv is at index 1. + assert "encoder.down_blocks.0.downsampler.resample.1.weight" in encoder + assert "encoder.down_blocks.1.downsampler.time_conv.weight" in encoder + assert "encoder.mid_block.attentions.0.to_qkv.weight" in encoder + assert "encoder.mid_block.resnets.1.conv2.bias" in encoder + assert "encoder.norm_out.gamma" in encoder + assert "encoder.conv_out.weight" in encoder + assert "quant_conv.weight" in encoder + + # Decoder: residual up block uses the singular ``upsampler`` attribute. + assert "post_quant_conv.weight" in decoder + assert "decoder.conv_in.weight" in decoder + assert "decoder.up_blocks.0.upsampler.resample.1.weight" in decoder + assert "decoder.up_blocks.0.upsampler.time_conv.weight" in decoder + assert "decoder.norm_out.gamma" in decoder + assert "decoder.conv_out.bias" in decoder + + # Parameter-free shortcuts must not create initializers. + assert not [n for n in encoder | decoder if "avg_shortcut" in n] + + def test_shortcut_conv_only_when_channels_change(self): + _, _, package = _build(TINY_WAN22) + encoder = set(package["encoder"].graph.initializers) + + # Stage 0 keeps 8 channels (dim_mult[0] == 1) -> nn.Identity upstream. + assert "encoder.down_blocks.0.resnets.0.conv_shortcut.weight" not in encoder + # Stage 1 widens 8 -> 16 -> a 1x1x1 causal conv shortcut. + assert "encoder.down_blocks.1.resnets.0.conv_shortcut.weight" in encoder + + def test_wan21_non_residual_layout(self): + config, _, package = _build(TINY_WAN21) + encoder = set(package["encoder"].graph.initializers) + decoder = set(package["decoder"].graph.initializers) + + assert config.is_residual is False + assert config.patch_size is None + assert config.video_channels == 3 + # Wan 2.1 flattens residual blocks and resamplers into ``down_blocks``. + assert "encoder.down_blocks.0.conv1.weight" in encoder + assert "encoder.down_blocks.1.resample.1.weight" in encoder + # Wan 2.1 decoder stores upsamplers in a one-element ModuleList. + assert "decoder.up_blocks.0.upsamplers.0.resample.1.weight" in decoder + assert "decoder.up_blocks.0.upsamplers.0.time_conv.weight" in decoder + assert not [n for n in decoder if "up_blocks.0.upsampler." in n] + + def test_attn_scales_insert_encoder_attention_blocks(self): + # Wan 2.1 inserts an attention block after each residual block whose + # spatial scale is listed in ``attn_scales`` (1.0 == full resolution). + _, _, package = _build(dict(TINY_WAN21, attn_scales=[1.0])) + encoder = set(package["encoder"].graph.initializers) + + assert "encoder.down_blocks.1.to_qkv.weight" in encoder + assert "encoder.down_blocks.1.proj.weight" in encoder + + def test_preprocess_weights_is_identity(self): + _, module, _ = _build(TINY_WAN22) + state_dict = {"encoder.conv_in.weight": object()} + + assert module.preprocess_weights(state_dict) is state_dict + + def test_task_rejects_module_without_components(self): + from onnxscript import nn + + config = WanVAEConfig.from_diffusers(TINY_WAN22) + + with pytest.raises(TypeError, match="encoder"): + WanVAETask().build(nn.Module(), config) + + +# --------------------------------------------------------------------------- +# Numerical parity against diffusers +# --------------------------------------------------------------------------- + + +def _wan_module(name: str): + """Import a private class from ``diffusers.models.autoencoders.autoencoder_kl_wan``.""" + module = pytest.importorskip("diffusers.models.autoencoders.autoencoder_kl_wan") + return getattr(module, name) + + +def _single_block_graph(block, torch_input, dtype=ir.DataType.FLOAT): + """Build a one-block ONNX model with ``x`` as its only input/output.""" + graph, builder = _make_graph(name="block") + x = builder.input("x", dtype=dtype, shape=list(torch_input.shape)) + builder.add_output(block(builder.op, x), "y") + return _make_model(graph) + + +def _reference_and_package(config_dict: dict): + """Instantiate the diffusers reference and a weight-loaded mobius package.""" + torch = pytest.importorskip("torch") + from mobius._weight_loading import apply_weights + + autoencoder_kl_wan = _wan_module("AutoencoderKLWan") + torch.manual_seed(0) + kwargs = {k: v for k, v in config_dict.items() if not k.startswith("_")} + kwargs.pop("clip_output", None) + reference = autoencoder_kl_wan(**kwargs).eval() + state_dict = dict(reference.state_dict()) + + config, _, package = _build(config_dict) + for model in package.values(): + initializers = set(model.graph.initializers) + apply_weights(model, {k: v for k, v in state_dict.items() if k in initializers}) + return reference, config, package + + +def _latent_stats(config: WanVAEConfig) -> tuple[np.ndarray, np.ndarray]: + """Return ``(latents_mean, latents_std)`` broadcast to ``(1, z, 1, 1, 1)``.""" + mean = np.asarray(config.latents_mean, dtype=np.float32).reshape(1, -1, 1, 1, 1) + std = np.asarray(config.latents_std, dtype=np.float32).reshape(1, -1, 1, 1, 1) + return mean, std + + +class TestParity: + """Compare individual blocks and the full VAE against diffusers PyTorch.""" + + def test_causal_conv3d_matches_diffusers(self): + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + from mobius._weight_loading import apply_weights + from mobius.models.wan_vae import _WanCausalConv3d + + wan_causal_conv3d = _wan_module("WanCausalConv3d") + torch.manual_seed(0) + reference = wan_causal_conv3d(3, 5, 3, padding=1).eval() + block = _WanCausalConv3d(3, 5, 3, padding=1) + + x = torch.randn(1, 3, 4, 6, 6) + with torch.no_grad(): + expected = reference(x).numpy() + + model = _single_block_graph(block, x) + apply_weights(model, dict(reference.state_dict())) + actual = _run(model, {"x": x.numpy()})[0] + + # Causal padding must keep the frame count and never look ahead in time. + assert actual.shape == expected.shape + assert np.abs(actual - expected).max() < 1e-5 + + def test_residual_block_matches_diffusers(self): + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + from mobius._weight_loading import apply_weights + from mobius.models.wan_vae import _WanResidualBlock + + wan_residual_block = _wan_module("WanResidualBlock") + torch.manual_seed(0) + reference = wan_residual_block(4, 6).eval() + block = _WanResidualBlock(4, 6) + + x = torch.randn(1, 4, 3, 5, 5) + with torch.no_grad(): + expected = reference(x).numpy() + + model = _single_block_graph(block, x) + apply_weights(model, dict(reference.state_dict())) + actual = _run(model, {"x": x.numpy()})[0] + + assert np.abs(actual - expected).max() < 1e-5 + + def test_downsample3d_matches_chunked_diffusers(self): + """Whole-sequence ``downsample3d`` == upstream's cached chunk loop.""" + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + from mobius._weight_loading import apply_weights + from mobius.models.wan_vae import _WanResample + + wan_resample = _wan_module("WanResample") + torch.manual_seed(0) + reference = wan_resample(4, mode="downsample3d").eval() + block = _WanResample(4, mode="downsample3d") + + # The encoder feeds frame 0 alone, then groups of four. + x = torch.randn(1, 4, 9, 8, 8) + chunks = [x[:, :, :1], x[:, :, 1:5], x[:, :, 5:9]] + feat_cache: list = [None] + with torch.no_grad(): + expected = torch.cat( + [reference(c, feat_cache=feat_cache, feat_idx=[0]) for c in chunks], dim=2 + ).numpy() + + model = _single_block_graph(block, x) + apply_weights(model, dict(reference.state_dict())) + actual = _run(model, {"x": x.numpy()})[0] + + assert actual.shape == expected.shape + assert np.abs(actual - expected).max() < 1e-5 + + def test_upsample3d_matches_chunked_diffusers(self): + """Whole-sequence ``upsample3d`` == upstream's per-frame cached loop.""" + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + from mobius._weight_loading import apply_weights + from mobius.models.wan_vae import _WanResample + + wan_resample = _wan_module("WanResample") + torch.manual_seed(0) + reference = wan_resample(4, mode="upsample3d", upsample_out_dim=4).eval() + block = _WanResample(4, mode="upsample3d", upsample_out_dim=4) + + # The decoder feeds one latent frame at a time. + x = torch.randn(1, 4, 3, 4, 4) + feat_cache: list = [None] + with torch.no_grad(): + expected = torch.cat( + [ + reference(x[:, :, i : i + 1], feat_cache=feat_cache, feat_idx=[0]) + for i in range(x.shape[2]) + ], + dim=2, + ).numpy() + + model = _single_block_graph(block, x) + apply_weights(model, dict(reference.state_dict())) + actual = _run(model, {"x": x.numpy()})[0] + + # Frame 0 is not temporally doubled: 3 latent frames -> 2 * 3 - 1 = 5. + assert actual.shape[2] == 2 * x.shape[2] - 1 + assert actual.shape == expected.shape + assert np.abs(actual - expected).max() < 1e-5 + + def test_upsample3d_single_frame_is_a_temporal_no_op(self): + """A lone frame is chunk 0, whose ``time_conv`` upstream never runs.""" + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + from mobius._weight_loading import apply_weights + from mobius.models.wan_vae import _WanResample + + wan_resample = _wan_module("WanResample") + torch.manual_seed(0) + reference = wan_resample(4, mode="upsample3d", upsample_out_dim=4).eval() + block = _WanResample(4, mode="upsample3d", upsample_out_dim=4) + + x = torch.randn(1, 4, 1, 4, 4) + feat_cache: list = [None] + with torch.no_grad(): + expected = reference(x, feat_cache=feat_cache, feat_idx=[0]).numpy() + + model = _single_block_graph(block, x) + apply_weights(model, dict(reference.state_dict())) + actual = _run(model, {"x": x.numpy()})[0] + + # 2 * 1 - 1 == 1: no temporal doubling for the very first chunk. + assert actual.shape[2] == 1 + assert actual.shape == expected.shape + assert np.abs(actual - expected).max() < 1e-5 + + def test_downsample3d_single_frame_is_a_temporal_no_op(self): + """Upstream caches chunk 0 verbatim instead of striding over it.""" + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + from mobius._weight_loading import apply_weights + from mobius.models.wan_vae import _WanResample + + wan_resample = _wan_module("WanResample") + torch.manual_seed(0) + reference = wan_resample(4, mode="downsample3d").eval() + block = _WanResample(4, mode="downsample3d") + + x = torch.randn(1, 4, 1, 8, 8) + feat_cache: list = [None] + with torch.no_grad(): + expected = reference(x, feat_cache=feat_cache, feat_idx=[0]).numpy() + + model = _single_block_graph(block, x) + apply_weights(model, dict(reference.state_dict())) + actual = _run(model, {"x": x.numpy()})[0] + + assert actual.shape[2] == 1 + assert actual.shape == expected.shape + assert np.abs(actual - expected).max() < 1e-5 + + @pytest.mark.parametrize("config_dict", [TINY_WAN22, TINY_WAN21]) + def test_encoder_and_decoder_match_diffusers(self, config_dict): + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + from mobius._weight_loading import apply_weights + + autoencoder_kl_wan = _wan_module("AutoencoderKLWan") + torch.manual_seed(0) + kwargs = {k: v for k, v in config_dict.items() if not k.startswith("_")} + kwargs.pop("clip_output", None) + reference = autoencoder_kl_wan(**kwargs).eval() + state_dict = dict(reference.state_dict()) + + config, _, package = _build(config_dict) + encoder_names = set(package["encoder"].graph.initializers) + decoder_names = set(package["decoder"].graph.initializers) + + # Weight-name alignment: the module-scoped initializers of the two graphs + # are exactly the checkpoint keys — no renaming, nothing missing or extra. + module_prefixes = ("encoder.", "quant_conv.", "decoder.", "post_quant_conv.") + graph_parameters = { + name for name in encoder_names | decoder_names if name.startswith(module_prefixes) + } + assert graph_parameters == set(state_dict) + + apply_weights( + package["encoder"], + {k: v for k, v in state_dict.items() if k in encoder_names}, + ) + apply_weights( + package["decoder"], + {k: v for k, v in state_dict.items() if k in decoder_names}, + ) + + latents_mean = np.asarray(config.latents_mean, dtype=np.float32).reshape( + 1, -1, 1, 1, 1 + ) + latents_std = np.asarray(config.latents_std, dtype=np.float32).reshape(1, -1, 1, 1, 1) + + # 5 video frames == 4 * 1 + 1; two spatial stages plus patchification. + video = torch.randn(1, config.video_channels, 5, 16, 16) + with torch.no_grad(): + posterior = reference.encode(video).latent_dist + mean, logvar, latent = _run(package["encoder"], {"sample": video.numpy()}) + + assert np.abs(mean - posterior.mean.numpy()).max() < 1e-4 + assert np.abs(logvar - posterior.logvar.numpy()).max() < 1e-4 + # The normalisation lives at the graph boundary, not inside ``encoder``. + assert np.abs(latent - (mean - latents_mean) / latents_std).max() < 1e-5 + + latent_frames = mean.shape[2] + z = torch.randn(1, config.z_dim, latent_frames, *mean.shape[3:]) + with torch.no_grad(): + expected = reference.decode(z).sample.numpy() + normalized = ((z.numpy() - latents_mean) / latents_std).astype(np.float32) + actual = _run(package["decoder"], {"latent": normalized})[0] + + assert actual.shape == expected.shape + assert actual.shape[1] == config.decoded_video_channels + # T_video = scale_factor_temporal * (T_latent - 1) + 1 + assert actual.shape[2] == config.scale_factor_temporal * (latent_frames - 1) + 1 + assert np.abs(actual - expected).max() < 1e-4 + + def test_longer_sequence_matches_diffusers(self): + """A 9-frame clip spans three encode chunks, exercising the cache seam.""" + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + from mobius._weight_loading import apply_weights + + autoencoder_kl_wan = _wan_module("AutoencoderKLWan") + torch.manual_seed(0) + kwargs = {k: v for k, v in TINY_WAN22.items() if not k.startswith("_")} + kwargs.pop("clip_output") + reference = autoencoder_kl_wan(**kwargs).eval() + state_dict = dict(reference.state_dict()) + + config, _, package = _build(TINY_WAN22) + encoder_names = set(package["encoder"].graph.initializers) + decoder_names = set(package["decoder"].graph.initializers) + apply_weights( + package["encoder"], + {k: v for k, v in state_dict.items() if k in encoder_names}, + ) + apply_weights( + package["decoder"], + {k: v for k, v in state_dict.items() if k in decoder_names}, + ) + + video = torch.randn(1, config.video_channels, 9, 16, 16) + with torch.no_grad(): + posterior = reference.encode(video).latent_dist + mean, _, latent = _run(package["encoder"], {"sample": video.numpy()}) + + # T_latent = (T_video - 1) / scale_factor_temporal + 1 + assert mean.shape[2] == (9 - 1) // config.scale_factor_temporal + 1 + assert np.abs(mean - posterior.mean.numpy()).max() < 1e-4 + + with torch.no_grad(): + expected = reference.decode(posterior.mean).sample.numpy() + actual = _run(package["decoder"], {"latent": latent})[0] + + assert actual.shape == video.shape + assert actual.shape == expected.shape + assert np.abs(actual - expected).max() < 1e-4 + + +# --------------------------------------------------------------------------- +# Single-frame (image) mode +# --------------------------------------------------------------------------- + + +class TestSingleFrame: + """``T_latent = 1`` / ``T_video = 1``, the Cosmos3 text-to-image path. + + Upstream decodes one latent frame per chunk and skips ``time_conv`` for + chunk 0, so a lone latent frame decodes to a lone video frame. The exported + graph reproduces that with a safe (right-padded, then sliced) temporal + window instead of a Python-level branch. + """ + + @pytest.mark.parametrize("config_dict", [TINY_WAN22, TINY_WAN21, TINY_COSMOS3]) + @pytest.mark.parametrize("latent_frames", [1, 2, 3]) + def test_decoder_matches_diffusers_for_any_latent_frame_count( + self, config_dict, latent_frames + ): + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + + reference, config, package = _reference_and_package(config_dict) + latents_mean, latents_std = _latent_stats(config) + + z = torch.randn(1, config.z_dim, latent_frames, 2, 2) + with torch.no_grad(): + expected = reference.decode(z).sample.numpy() + normalized = ((z.numpy() - latents_mean) / latents_std).astype(np.float32) + actual = _run(package["decoder"], {"latent": normalized})[0] + + # T_video = scale_factor_temporal * (T_latent - 1) + 1, so a single + # latent frame decodes to a single video frame. + assert expected.shape[2] == config.scale_factor_temporal * (latent_frames - 1) + 1 + assert actual.shape == expected.shape + assert actual.shape[1] == config.decoded_video_channels + assert actual.shape[3:] == (2 * config.scale_factor_spatial,) * 2 + assert np.abs(actual - expected).max() < 1e-4 + + @pytest.mark.parametrize("config_dict", [TINY_WAN22, TINY_WAN21, TINY_COSMOS3]) + def test_encoder_matches_diffusers_for_a_single_video_frame(self, config_dict): + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + + reference, config, package = _reference_and_package(config_dict) + latents_mean, latents_std = _latent_stats(config) + + video = torch.randn( + 1, + config.video_channels, + 1, + 2 * config.scale_factor_spatial, + 2 * config.scale_factor_spatial, + ) + with torch.no_grad(): + posterior = reference.encode(video).latent_dist + mean, logvar, latent = _run(package["encoder"], {"sample": video.numpy()}) + + assert mean.shape == (1, config.z_dim, 1, 2, 2) + assert mean.shape == tuple(posterior.mean.shape) + assert np.abs(mean - posterior.mean.numpy()).max() < 1e-4 + assert np.abs(logvar - posterior.logvar.numpy()).max() < 1e-4 + assert np.abs(latent - (mean - latents_mean) / latents_std).max() < 1e-5 + + def test_image_round_trip_through_both_graphs(self): + """Encode one frame, decode the resulting latent, recover one frame.""" + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + + reference, config, package = _reference_and_package(TINY_COSMOS3) + + image = torch.randn(1, config.video_channels, 1, 32, 32) + _, _, latent = _run(package["encoder"], {"sample": image.numpy()}) + assert latent.shape == (1, config.z_dim, 1, 2, 2) + + decoded = _run(package["decoder"], {"latent": latent})[0] + with torch.no_grad(): + posterior = reference.encode(image).latent_dist + expected = reference.decode(posterior.mean).sample.numpy() + + assert decoded.shape == image.shape + assert decoded.shape == expected.shape + assert np.abs(decoded - expected).max() < 1e-4 + + def test_multi_frame_numerics_are_unaffected_by_the_safe_window(self): + """The padded-then-sliced window must not perturb longer sequences. + + ``time_conv`` is causal, so the frames retained after the right-hand + zero padding are the frames the unpadded convolution would produce. A + decode of ``T`` latent frames must therefore agree with independently + decoding its ``T - 1`` frame prefix on the shared video frames. + """ + pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + + _, config, package = _reference_and_package(TINY_COSMOS3) + latents_mean, latents_std = _latent_stats(config) + + z = torch.randn(1, config.z_dim, 3, 2, 2).numpy() + normalized = ((z - latents_mean) / latents_std).astype(np.float32) + full = _run(package["decoder"], {"latent": normalized})[0] + prefix = _run(package["decoder"], {"latent": normalized[:, :, :2]})[0] + + # The decoder is causal in time: the prefix decode is a prefix of the + # full decode, frame for frame. + assert prefix.shape[2] == config.scale_factor_temporal + 1 + assert np.abs(full[:, :, : prefix.shape[2]] - prefix).max() < 1e-6 + + def test_built_package_saves_and_decodes_one_frame_from_disk(self, tmp_path): + """Build -> apply weights -> ``ModelPackage.save`` -> ORT decode at T=1.""" + ort = pytest.importorskip("onnxruntime") + torch = pytest.importorskip("torch") + + reference, config, package = _reference_and_package(TINY_COSMOS3) + package.save(str(tmp_path), progress_bar=False) + + decoder_path = tmp_path / "decoder" / "model.onnx" + assert decoder_path.is_file() + assert (tmp_path / "encoder" / "model.onnx").is_file() + + latents_mean, latents_std = _latent_stats(config) + z = torch.randn(1, config.z_dim, 1, 2, 2) + normalized = ((z.numpy() - latents_mean) / latents_std).astype(np.float32) + + session = ort.InferenceSession(str(decoder_path)) + # The saved graph keeps latent_frames symbolic, so nothing pins it to 1. + latent_input = session.get_inputs()[0] + assert latent_input.name == "latent" + assert latent_input.shape[2] == "latent_frames" + + (sample,) = session.run(None, {"latent": normalized}) + with torch.no_grad(): + expected = reference.decode(z).sample.numpy() + + assert sample.shape == (1, config.decoded_video_channels, 1, 32, 32) + assert sample.shape == expected.shape + assert np.abs(sample - expected).max() < 1e-4 + assert sample.min() >= -1.0 and sample.max() <= 1.0 diff --git a/src/mobius/models/world_model.py b/src/mobius/models/world_model.py index 3b7dd80e3..e17bf3e5c 100644 --- a/src/mobius/models/world_model.py +++ b/src/mobius/models/world_model.py @@ -1,25 +1,25 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Minimal directly declared world-model implementation.""" +"""Minimal directly declared latent-dynamics implementation.""" from __future__ import annotations import torch from onnxscript import OpBuilder, nn -from mobius._configs import WorldModelConfig +from mobius._configs import LatentDynamicsConfig from mobius.components import Linear, get_activation -class MLPWorldModel(nn.Module): - """Deterministic MLP reference model for the world-model task contract.""" +class MLPLatentDynamicsModel(nn.Module): + """Deterministic MLP reference model for the latent-dynamics contract.""" - default_task = "world-model" - config_class = WorldModelConfig - category = "World Model" + default_task = "latent-dynamics" + config_class = LatentDynamicsConfig + category = "World Model / Dynamics" - def __init__(self, config: WorldModelConfig): + def __init__(self, config: LatentDynamicsConfig): super().__init__() config.validate() self.config = config @@ -85,3 +85,8 @@ def preprocess_weights( ) -> dict[str, torch.Tensor]: """Return weights unchanged; provided for parity with other Mobius models.""" return state_dict + + +# Kept for source compatibility. This model is one possible dynamics component, +# not a complete world-model pipeline. +MLPWorldModel = MLPLatentDynamicsModel diff --git a/src/mobius/models/world_model_test.py b/src/mobius/models/world_model_test.py index 02b7e8308..6d4d4515f 100644 --- a/src/mobius/models/world_model_test.py +++ b/src/mobius/models/world_model_test.py @@ -11,12 +11,14 @@ from onnxscript import nn from mobius import ( + LatentDynamicsConfig, + MLPLatentDynamicsModel, MLPWorldModel, WorldModelConfig, WorldModelTask, build_from_module, ) -from mobius.tasks import TASK_REGISTRY, get_task +from mobius.tasks import TASK_REGISTRY, LatentDynamicsTask, get_task class _TorchMLPWorldModel(torch.nn.Module): @@ -100,8 +102,16 @@ def test_invalid_config_raises(self, field, value): class TestWorldModelTask: def test_registered(self): + assert TASK_REGISTRY["latent-dynamics"] is LatentDynamicsTask assert TASK_REGISTRY["world-model"] is WorldModelTask + assert WorldModelTask is LatentDynamicsTask assert isinstance(get_task("world-model"), WorldModelTask) + assert isinstance(get_task("latent-dynamics"), LatentDynamicsTask) + + def test_backward_compatible_aliases(self): + assert WorldModelConfig is LatentDynamicsConfig + assert MLPWorldModel is MLPLatentDynamicsModel + assert MLPLatentDynamicsModel.default_task == "latent-dynamics" def test_graph_contract(self): config = _config() diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index d7f947ffc..3276859b7 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -28,6 +28,9 @@ "CodecTask", "ComponentSpec", "ControlNetTask", + "Cosmos3AVAEAudioDecoderTask", + "Cosmos3AVAEAudioTokenizerTask", + "Cosmos3OmniGeneratorTask", "DeepSeekV4Task", "DFlashDraftTask", "Eagle3DraftTask", @@ -44,6 +47,7 @@ "Cosmos3EdgeVLTask", "HybridQwenVLTask", "ImageClassificationTask", + "LatentDynamicsTask", "ModelTask", "MllamaVisionLanguageTask", "MaskedDiffusionTask", @@ -68,6 +72,7 @@ "VideoDenoisingTask", "VisionLanguageTask", "WorldModelTask", + "WanVAETask", "build_decoder_from_embeds", "build_embedding_from_features", "get_task", @@ -89,6 +94,11 @@ ) from mobius.tasks._codec import CodecTask from mobius.tasks._controlnet import ControlNetTask +from mobius.tasks._cosmos3_audio import ( + Cosmos3AVAEAudioDecoderTask, + Cosmos3AVAEAudioTokenizerTask, +) +from mobius.tasks._cosmos3_omni_generator import Cosmos3OmniGeneratorTask from mobius.tasks._ctc_asr import CTCAsrTask from mobius.tasks._deepseek_v4 import DeepSeekV4Task from mobius.tasks._denoising import DenoisingTask @@ -129,7 +139,8 @@ QwenVLTask, VisionLanguageTask, ) -from mobius.tasks._world_model import WorldModelTask +from mobius.tasks._wan_vae import WanVAETask +from mobius.tasks._world_model import LatentDynamicsTask, WorldModelTask # --------------------------------------------------------------------------- # Task registry @@ -142,6 +153,9 @@ "ctc-asr": CTCAsrTask, "codec": CodecTask, "controlnet": ControlNetTask, + "cosmos3-audio-decoder": Cosmos3AVAEAudioDecoderTask, + "cosmos3-audio-tokenizer": Cosmos3AVAEAudioTokenizerTask, + "cosmos3-omni-generator": Cosmos3OmniGeneratorTask, "denoising": DenoisingTask, "diarization": DiarizationTask, "feature-extraction": FeatureExtractionTask, @@ -158,6 +172,7 @@ "eagle3-draft": Eagle3DraftTask, "qwen35-mtp": Qwen35MtpTask, "vae": VAETask, + "wan-vae": WanVAETask, "qwen-image-vae": QwenImageVAETask, "vision-language": VisionLanguageTask, "cosmos3-edge-vl": Cosmos3EdgeVLTask, @@ -181,6 +196,7 @@ "ssm2-text-generation": SSM2CausalLMTask, "tts": TTSTask, "video-denoising": VideoDenoisingTask, + "latent-dynamics": LatentDynamicsTask, "world-model": WorldModelTask, } diff --git a/src/mobius/tasks/_cosmos3_audio.py b/src/mobius/tasks/_cosmos3_audio.py new file mode 100644 index 000000000..e0c7ba7ba --- /dev/null +++ b/src/mobius/tasks/_cosmos3_audio.py @@ -0,0 +1,224 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Cosmos3 AVAE audio tokenizer tasks (``sound_tokenizer``). + +Two explicit build paths, because AVAE checkpoints are shipped in two shapes: + +* :class:`Cosmos3AVAEAudioDecoderTask` — decoder-only weights. Produces a + single ``"decoder"`` graph. Nothing encoder-related is instantiated, so the + package can never contain an initializer without weight data. +* :class:`Cosmos3AVAEAudioTokenizerTask` — encoder + decoder weights. Produces + ``"encoder"`` and ``"decoder"`` graphs. + +Use :func:`select_cosmos3_audio_task` to pick the right one from a config whose +``encoder_enabled`` flag was resolved from the checkpoint. Resolve it with +``Cosmos3AudioConfig.from_diffusers(config_json, weight_names=state_dict)`` — +the published configs are byte-identical for full and decoder-only releases, so +parsing the config alone always assumes an encoder is present:: + + nvidia/Cosmos3-Nano -> Cosmos3AVAEAudioTokenizerTask (encoder+decoder) + nvidia/Cosmos3-Super -> Cosmos3AVAEAudioTokenizerTask (encoder+decoder) + nvidia/Cosmos3-Super-Text2Image -> Cosmos3AVAEAudioDecoderTask (decoder-only) + +Graph contracts +--------------- + +``decoder`` — latent → waveform:: + + inputs latents [batch, vocoder_input_dim, latent_frames] config.dtype + outputs waveform [batch, dec_out_channels, latent_frames * hop_size] + + latents are consumed as-is: no latent mean/std de-normalization is applied, + matching upstream (which rejects non-null latent_mean/latent_std). The + waveform is clamped to [-1, 1]. + +``encoder`` — waveform → posterior moments:: + + inputs audio [batch, encoder_input_channels, num_samples] config.dtype + outputs moments [batch, 2 * vocoder_input_dim, latent_frames] + latent_mean [batch, vocoder_input_dim, latent_frames] + latent_std [batch, vocoder_input_dim, latent_frames] + + ``latent_frames = ceil(num_samples / hop_size)``. + + Normalization semantics baked into the graph, in order: + + 1. peak volume normalization ``x / (|x|.max() + 1e-5) * 0.95`` when + ``config.normalize_volume`` is set. The maximum is global across the + whole input tensor (batch included), exactly as upstream — use batch + size 1 for per-clip normalization. + 2. right zero-padding to a multiple of ``hop_size``. + 3. VAE bottleneck: ``mean, scale = split(moments, 2, axis=1)`` and + ``std = softplus(scale) + 1e-4``. + + Sampling is intentionally left outside the graph so the ONNX model is + deterministic; draw ``z = mean + std * eps`` in the caller, or use ``mean`` + for the distribution mode. + +Preprocessing boundary +---------------------- + +The STFT front-end is part of ``encoder.forward`` upstream, so it is emitted +inside the ONNX graph (as an ``STFT`` node) and the encoder contract stays +``waveform -> moments``. onnxruntime implements ``STFT`` on CPU only, so under +the CUDA EP that single node falls back to the CPU EP and adds a host/device +copy at the graph entry. Everything after it — the whole ConvNeXt stack — runs +on the accelerator. Deployments that need a pure-GPU encoder should call +:meth:`~mobius.models.cosmos3_audio.Cosmos3AudioSpectrogramConvNeXtEncoder.spectrogram` +out of band and feed ``encoder.layers`` directly. + +Latent normalization (``latent_mean``/``latent_std``) is *not* applied: the +published configs leave both ``null`` and upstream rejects any other value, so +:meth:`Cosmos3AudioConfig.validate` raises rather than invent semantics. + +Weight routing +-------------- + +Both graphs keep the HuggingFace module paths, so initializer names are already +disjoint (``encoder.*`` vs ``decoder.*``). Do **not** set a ``weight_prefix_map`` +on the module: the default "try every weight against every component" routing in +:meth:`ModelPackage.apply_weights` is correct here, and stripping the prefix +would break the match. Run +:meth:`~mobius.models.cosmos3_audio.Cosmos3AVAEAudioDecoderOnlyTokenizer.preprocess_weights` +first to fold ``weight_g``/``weight_v`` into ``weight``. +""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir + +from mobius._configs._cosmos3_audio import Cosmos3AudioConfig +from mobius._model_package import ModelPackage +from mobius.models.cosmos3_audio import ( + Cosmos3AVAEAudioDecoderOnlyTokenizer, + Cosmos3AVAEAudioTokenizer, +) +from mobius.tasks._base import ComponentSpec, ModelTask, _make_graph, _make_model + +__all__ = [ + "Cosmos3AVAEAudioDecoderTask", + "Cosmos3AVAEAudioTokenizerTask", + "select_cosmos3_audio_task", +] + + +class Cosmos3AVAEAudioDecoderTask(ModelTask): + """Build the ``latents -> waveform`` graph of a Cosmos3 AVAE sound tokenizer. + + This is the decoder-only path used for checkpoints that ship without + ``encoder.*`` weights. + """ + + model_roles: ClassVar[dict[str, str]] = {"decoder": "decoder"} + components: ClassVar[ComponentSpec] = ComponentSpec(decoder="decoder") + + def build( + self, + module: Cosmos3AVAEAudioDecoderOnlyTokenizer, + config: Cosmos3AudioConfig, + ) -> ModelPackage: + """Build a package containing only the ``"decoder"`` model.""" + self._validate_components(module) + config.validate() + return ModelPackage({"decoder": self._build_decoder(module, config)}, config=config) + + def _build_decoder( + self, + module: Cosmos3AVAEAudioDecoderOnlyTokenizer, + config: Cosmos3AudioConfig, + ) -> ir.Model: + """Wire ``decode()`` into a graph: latents → clamped waveform.""" + batch = ir.SymbolicDim("batch") + latent_frames = ir.SymbolicDim("latent_frames") + + graph, builder = _make_graph(name="cosmos3_audio_decoder") + latents = builder.input( + "latents", + dtype=config.dtype, + shape=[batch, config.latent_channels, latent_frames], + ) + + # (B, z, T) -> (B, audio_channels, T * hop_size) + waveform = module.decode(builder.op, latents) + + builder.add_output(waveform, "waveform") + return _make_model(graph) + + +class Cosmos3AVAEAudioTokenizerTask(Cosmos3AVAEAudioDecoderTask): + """Build both the encoder and decoder graphs of a Cosmos3 AVAE tokenizer. + + Requires a module with encoder weights available + (:class:`~mobius.models.cosmos3_audio.Cosmos3AVAEAudioTokenizer`). + """ + + model_roles: ClassVar[dict[str, str]] = {"encoder": "encoder", "decoder": "decoder"} + components: ClassVar[ComponentSpec] = ComponentSpec( + encoder="encoder", decoder="decoder", bottleneck="bottleneck" + ) + + def build( + self, + module: Cosmos3AVAEAudioTokenizer, + config: Cosmos3AudioConfig, + ) -> ModelPackage: + """Build a package containing the ``"encoder"`` and ``"decoder"`` models.""" + self._validate_components(module) + config.validate() + if not config.encoder_enabled: + raise ValueError( + "Cosmos3AVAEAudioTokenizerTask requires config.encoder_enabled=True. " + "Decoder-only AVAE checkpoints must use Cosmos3AVAEAudioDecoderTask so no " + "encoder initializer is created without weight data." + ) + return ModelPackage( + { + "encoder": self._build_encoder(module, config), + "decoder": self._build_decoder(module, config), + }, + config=config, + ) + + def _build_encoder( + self, + module: Cosmos3AVAEAudioTokenizer, + config: Cosmos3AudioConfig, + ) -> ir.Model: + """Wire ``encode()`` into a graph: waveform → deterministic moments.""" + batch = ir.SymbolicDim("batch") + num_samples = ir.SymbolicDim("num_samples") + + graph, builder = _make_graph(name="cosmos3_audio_encoder") + audio = builder.input( + "audio", + dtype=config.dtype, + shape=[batch, config.encoder_input_channels, num_samples], + ) + + # (B, C, N) -> moments (B, 2z, T) and the split (mean, std) pair. + moments, mean, std = module.encode(builder.op, audio) + + builder.add_output(moments, "moments") + builder.add_output(mean, "latent_mean") + builder.add_output(std, "latent_std") + return _make_model(graph) + + +def select_cosmos3_audio_task(config: Cosmos3AudioConfig) -> type[Cosmos3AVAEAudioDecoderTask]: + """Return the task class matching ``config.encoder_enabled``. + + Args: + config: A Cosmos3 AVAE audio config whose ``encoder_enabled`` flag has + already been reconciled with the checkpoint (see + :meth:`Cosmos3AudioConfig.with_encoder_from_state_dict`). + + Returns: + :class:`Cosmos3AVAEAudioTokenizerTask` when an encoder is present, + otherwise :class:`Cosmos3AVAEAudioDecoderTask`. + """ + if config.encoder_enabled: + return Cosmos3AVAEAudioTokenizerTask + return Cosmos3AVAEAudioDecoderTask diff --git a/src/mobius/tasks/_cosmos3_omni_generator.py b/src/mobius/tasks/_cosmos3_omni_generator.py new file mode 100644 index 000000000..f8640c69e --- /dev/null +++ b/src/mobius/tasks/_cosmos3_omni_generator.py @@ -0,0 +1,344 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Task for the NVIDIA Cosmos3-Omni unified MoT transformer. + +Builds a **single** ONNX graph performing one rectified-flow denoising step of +:class:`~mobius.models.cosmos3_omni_generator.Cosmos3OmniGeneratorModel`. + +Packed I/O contract +=================== + +Upstream's ``Cosmos3OmniTransformer.forward`` takes Python lists of ragged +per-item tensors and per-item ``(T, H, W)`` shape tuples. Neither can cross an +ONNX boundary, so the graph is cut at the **packed-token boundary**: the host +patchifies/packs latents before ``proj_in`` and unpatchifies/unpacks +predictions after ``proj_out``. Every tensor operation in between — including +the noisy-row offsets that upstream derives from ``token_shapes`` and the +``und_len`` split — is an explicit tensor input, so no semantics are lost. + +Symbolic dimensions +------------------- + +=========================== ========================================== +``sequence_length`` rows of the packed joint sequence +``num_text_tokens`` text tokens written into the joint sequence +``num_vision_tokens`` packed vision patch tokens +``num_vision_noisy_tokens`` vision patch tokens carrying noise +``num_sound_tokens`` packed sound frames (``sound_gen``) +``num_sound_noisy_tokens`` noisy sound frames (``sound_gen``) +``num_action_tokens`` packed action steps (``action_gen``) +``num_action_noisy_tokens`` noisy action steps (``action_gen``) +=========================== ========================================== + +Inputs (``D`` = ``config.dtype``, the model compute dtype) +---------------------------------------------------------- + +=============================== ======= =============================================== +``input_ids`` int64 ``[num_text_tokens]`` +``text_indexes`` int64 ``[num_text_tokens]`` joint rows for the text +``position_ids`` int64 ``[3, sequence_length]`` mRoPE (T, H, W) +``und_len`` int64 ``[1]`` leading rows on the understanding expert +``vision_tokens`` D ``[num_vision_tokens, patch_latent_dim]`` +``vision_sequence_indexes`` int64 ``[num_vision_tokens]`` joint rows +``vision_timesteps`` float32 ``[num_vision_noisy_tokens]`` +``vision_timestep_token_indexes``int64 ``[num_vision_noisy_tokens]`` rows of ``vision_tokens`` +``vision_mse_loss_indexes`` int64 ``[num_vision_noisy_tokens]`` joint rows to decode +=============================== ======= =============================================== + +Present only when ``config.sound_gen``: + +=============================== ======= =============================================== +``sound_tokens`` D ``[num_sound_tokens, sound_dim]`` +``sound_sequence_indexes`` int64 ``[num_sound_tokens]`` +``sound_timesteps`` float32 ``[num_sound_noisy_tokens]`` +``sound_timestep_token_indexes`` int64 ``[num_sound_noisy_tokens]`` +``sound_mse_loss_indexes`` int64 ``[num_sound_noisy_tokens]`` +=============================== ======= =============================================== + +Present only when ``config.action_gen``: + +================================ ======= ============================================== +``action_tokens`` D ``[num_action_tokens, action_dim]`` +``action_domain_ids`` int64 ``[num_action_tokens]`` embodiment per token +``action_sequence_indexes`` int64 ``[num_action_tokens]`` +``action_timesteps`` float32 ``[num_action_noisy_tokens]`` +``action_timestep_token_indexes`` int64 ``[num_action_noisy_tokens]`` +``action_mse_loss_indexes`` int64 ``[num_action_noisy_tokens]`` +``action_pred_domain_ids`` int64 ``[num_action_noisy_tokens]`` embodiment per pred +================================ ======= ============================================== + +Outputs +------- + +=============== ======= =================================================== +``vision_pred`` D ``[num_vision_noisy_tokens, patch_latent_dim]`` +``sound_pred`` D ``[num_sound_noisy_tokens, sound_dim]`` (gated) +``action_pred`` D ``[num_action_noisy_tokens, action_dim]`` (gated) +=============== ======= =================================================== + +Notes: +----- +* ``*_timesteps`` are **float32 regardless of the model dtype**: upstream keeps + ``time_embedder`` in fp32 (``_keep_in_fp32_modules``) and the sinusoidal + projection would lose integer resolution in bf16. The embedding is cast to + the model dtype only where it is added to the token stream. +* Optional heads are gated at *graph construction* time by the config, never + at run time. A ``sound_gen`` model that has no sound content in a given + step passes **zero-length** sound tensors (``num_sound_tokens = 0``); the + scatters and gathers then degenerate to no-ops, matching upstream's + ``numel() > 0`` guard. A configured head is never silently dropped. +* The packed vision predictions are returned **before** unpatchify: the host + reshapes ``[num_vision_noisy_tokens, patch_latent_dim]`` back into + ``[C, T, H, W]`` using the same per-item ``(T, H, W)`` shapes it used to + patchify. Expressing that reshape in the graph would require host-resident + per-item shape lists, so it is a documented, precise boundary rather than a + faked graph op. +""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir +from onnxscript import GraphBuilder, nn + +from mobius._configs._cosmos3_omni_generator import Cosmos3OmniGeneratorConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ModelTask, _make_graph, _make_model + +#: Inputs that every Cosmos3-Omni generator graph carries. +_CORE_INPUT_NAMES: tuple[str, ...] = ( + "input_ids", + "text_indexes", + "position_ids", + "und_len", + "vision_tokens", + "vision_sequence_indexes", + "vision_timesteps", + "vision_timestep_token_indexes", + "vision_mse_loss_indexes", +) + +#: Additional inputs when ``config.sound_gen``. +_SOUND_INPUT_NAMES: tuple[str, ...] = ( + "sound_tokens", + "sound_sequence_indexes", + "sound_timesteps", + "sound_timestep_token_indexes", + "sound_mse_loss_indexes", +) + +#: Additional inputs when ``config.action_gen``. +_ACTION_INPUT_NAMES: tuple[str, ...] = ( + "action_tokens", + "action_domain_ids", + "action_sequence_indexes", + "action_timesteps", + "action_timestep_token_indexes", + "action_mse_loss_indexes", + "action_pred_domain_ids", +) + + +def expected_input_names(config: Cosmos3OmniGeneratorConfig) -> tuple[str, ...]: + """Return the graph input names for *config*, in declaration order.""" + names = list(_CORE_INPUT_NAMES) + if config.sound_gen: + names.extend(_SOUND_INPUT_NAMES) + if config.action_gen: + names.extend(_ACTION_INPUT_NAMES) + return tuple(names) + + +def expected_output_names(config: Cosmos3OmniGeneratorConfig) -> tuple[str, ...]: + """Return the graph output names for *config*, in declaration order.""" + names = ["vision_pred"] + if config.sound_gen: + names.append("sound_pred") + if config.action_gen: + names.append("action_pred") + return tuple(names) + + +class Cosmos3OmniGeneratorTask(ModelTask): + """Build the unified Cosmos3-Omni MoT transformer graph. + + Produces a single ``"model"`` entry. The optimization role is + ``"encoder"``: the graph has no KV cache, so decoder-oriented fusions + (GroupQueryAttention with past/present) must not run on it. + """ + + model_roles: ClassVar[dict[str, str]] = {"model": "encoder"} + + def build( + self, + module: nn.Module, + config: Cosmos3OmniGeneratorConfig, + ) -> ModelPackage: + """Wire *module* into the packed denoising-step graph. + + Args: + module: A :class:`~mobius.models.cosmos3_omni_generator.Cosmos3OmniGeneratorModel` + (or any module with a compatible ``forward``). + config: Cosmos3-Omni generator configuration. ``sound_gen`` and + ``action_gen`` gate the optional inputs/outputs. + + Returns: + A :class:`ModelPackage` with a single ``"model"`` entry. + + Raises: + TypeError: If *module* does not return the + ``(vision_pred, sound_pred, action_pred)`` contract. + """ + config.validate() + + graph, builder = _make_graph(name="cosmos3_omni_generator") + kwargs = self._declare_core_inputs(builder, config) + if config.sound_gen: + kwargs.update(self._declare_sound_inputs(builder, config)) + if config.action_gen: + kwargs.update(self._declare_action_inputs(builder, config)) + + outputs = module(builder.op, **kwargs) + if not isinstance(outputs, tuple) or len(outputs) != 3: + raise TypeError( + f"{type(module).__name__} must return (vision_pred, sound_pred, action_pred)" + ) + vision_pred, sound_pred, action_pred = outputs + + builder.add_output(vision_pred, "vision_pred") + if config.sound_gen: + if sound_pred is None: + raise TypeError( + f"{type(module).__name__} returned sound_pred=None although " + "config.sound_gen is True" + ) + builder.add_output(sound_pred, "sound_pred") + if config.action_gen: + if action_pred is None: + raise TypeError( + f"{type(module).__name__} returned action_pred=None although " + "config.action_gen is True" + ) + builder.add_output(action_pred, "action_pred") + + return ModelPackage({"model": _make_model(graph)}, config=config) + + @staticmethod + def _declare_core_inputs( + builder: GraphBuilder, config: Cosmos3OmniGeneratorConfig + ) -> dict[str, ir.Value]: + """Declare the always-present inputs (text + joint layout + vision).""" + sequence_length = ir.SymbolicDim("sequence_length") + num_text_tokens = ir.SymbolicDim("num_text_tokens") + num_vision_tokens = ir.SymbolicDim("num_vision_tokens") + num_vision_noisy = ir.SymbolicDim("num_vision_noisy_tokens") + + return { + "input_ids": builder.input( + "input_ids", dtype=ir.DataType.INT64, shape=[num_text_tokens] + ), + "text_indexes": builder.input( + "text_indexes", dtype=ir.DataType.INT64, shape=[num_text_tokens] + ), + # 3-axis mRoPE positions; dim 1 also defines sequence_length. + "position_ids": builder.input( + "position_ids", dtype=ir.DataType.INT64, shape=[3, sequence_length] + ), + # Boundary between the causal understanding expert and the + # non-causal generation expert within the joint sequence. + "und_len": builder.input("und_len", dtype=ir.DataType.INT64, shape=[1]), + "vision_tokens": builder.input( + "vision_tokens", + dtype=config.dtype, + shape=[num_vision_tokens, config.patch_latent_dim], + ), + "vision_sequence_indexes": builder.input( + "vision_sequence_indexes", dtype=ir.DataType.INT64, shape=[num_vision_tokens] + ), + # Timesteps stay fp32 — see the module docstring. + "vision_timesteps": builder.input( + "vision_timesteps", dtype=ir.DataType.FLOAT, shape=[num_vision_noisy] + ), + "vision_timestep_token_indexes": builder.input( + "vision_timestep_token_indexes", + dtype=ir.DataType.INT64, + shape=[num_vision_noisy], + ), + "vision_mse_loss_indexes": builder.input( + "vision_mse_loss_indexes", dtype=ir.DataType.INT64, shape=[num_vision_noisy] + ), + } + + @staticmethod + def _declare_sound_inputs( + builder: GraphBuilder, config: Cosmos3OmniGeneratorConfig + ) -> dict[str, ir.Value]: + """Declare the Sound-head inputs (only when ``config.sound_gen``).""" + num_sound_tokens = ir.SymbolicDim("num_sound_tokens") + num_sound_noisy = ir.SymbolicDim("num_sound_noisy_tokens") + return { + "sound_tokens": builder.input( + "sound_tokens", + dtype=config.dtype, + shape=[num_sound_tokens, config.sound_dim], + ), + "sound_sequence_indexes": builder.input( + "sound_sequence_indexes", dtype=ir.DataType.INT64, shape=[num_sound_tokens] + ), + "sound_timesteps": builder.input( + "sound_timesteps", dtype=ir.DataType.FLOAT, shape=[num_sound_noisy] + ), + "sound_timestep_token_indexes": builder.input( + "sound_timestep_token_indexes", + dtype=ir.DataType.INT64, + shape=[num_sound_noisy], + ), + "sound_mse_loss_indexes": builder.input( + "sound_mse_loss_indexes", dtype=ir.DataType.INT64, shape=[num_sound_noisy] + ), + } + + @staticmethod + def _declare_action_inputs( + builder: GraphBuilder, config: Cosmos3OmniGeneratorConfig + ) -> dict[str, ir.Value]: + """Declare the Action-head inputs (only when ``config.action_gen``).""" + num_action_tokens = ir.SymbolicDim("num_action_tokens") + num_action_noisy = ir.SymbolicDim("num_action_noisy_tokens") + return { + "action_tokens": builder.input( + "action_tokens", + dtype=config.dtype, + shape=[num_action_tokens, config.action_dim], + ), + # Embodiment domain selects the DomainAwareLinear weight per token. + "action_domain_ids": builder.input( + "action_domain_ids", dtype=ir.DataType.INT64, shape=[num_action_tokens] + ), + "action_sequence_indexes": builder.input( + "action_sequence_indexes", dtype=ir.DataType.INT64, shape=[num_action_tokens] + ), + "action_timesteps": builder.input( + "action_timesteps", dtype=ir.DataType.FLOAT, shape=[num_action_noisy] + ), + "action_timestep_token_indexes": builder.input( + "action_timestep_token_indexes", + dtype=ir.DataType.INT64, + shape=[num_action_noisy], + ), + "action_mse_loss_indexes": builder.input( + "action_mse_loss_indexes", dtype=ir.DataType.INT64, shape=[num_action_noisy] + ), + "action_pred_domain_ids": builder.input( + "action_pred_domain_ids", dtype=ir.DataType.INT64, shape=[num_action_noisy] + ), + } + + +__all__ = [ + "Cosmos3OmniGeneratorTask", + "expected_input_names", + "expected_output_names", +] diff --git a/src/mobius/tasks/_vision_language_3model.py b/src/mobius/tasks/_vision_language_3model.py index 42d804ed9..c83dd4094 100644 --- a/src/mobius/tasks/_vision_language_3model.py +++ b/src/mobius/tasks/_vision_language_3model.py @@ -98,11 +98,21 @@ def _build_vision( class Cosmos3EdgeVLTask(VisionLanguageTask): """NVIDIA Cosmos3-Edge VL 3-model split. - Builds the text decoder with 3D multimodal RoPE + The decoder uses interleaved 3D multimodal RoPE (``mrope_section=[24, 20, 20]``), so ``position_ids`` has shape - ``[3, batch, seq]``. The vision runtime processes one image at a time and - removes the image batch dimension so its output matches the embedding - model's rank-2 ``[num_image_tokens, hidden]`` input contract. + ``[3, batch, seq]``. + + The vision encoder mirrors the published packed SigLIP2 contract: it takes + pre-patchified ``pixel_values [total_patches, patch_dim]`` plus a + ``grid_thw [3]`` triple describing one visual item, and emits rank-2 + ``image_features [total_patches / merge**2, text_hidden]``. Because + ``Cosmos3EdgeModel.get_video_features`` simply delegates to + ``get_image_features``, this single graph serves both images + (``grid_t == 1``) and videos (``grid_t == num_frames``). + + The embedding model accepts two feature streams so the runtime can route + each modality to its own placeholder id: ``image_features`` (token 19) and + ``video_features`` (token 18). Either may have zero rows. """ def build( @@ -114,12 +124,7 @@ def build( models: dict[str, ir.Model] = {} models["decoder"] = build_decoder_from_embeds(module.decoder, config, mrope=True) models["vision_encoder"] = self._build_vision(module.vision_encoder, config) - models["embedding"] = build_embedding_from_features( - module.embedding, - config, - feature_name="image_features", - feature_dim=config.hidden_size, - ) + models["embedding"] = self._build_embedding(module.embedding, config) return ModelPackage(models, config=config) def _build_vision( @@ -127,21 +132,62 @@ def _build_vision( vision: nn.Module, config: ArchitectureConfig, ) -> ir.Model: - """Build the single-image vision encoder with rank-2 feature output.""" - image_size = (config.vision.image_size if config.vision else None) or 224 + """Build the packed, variable-resolution vision encoder.""" + total_patches = ir.SymbolicDim("total_patches") graph, builder = _make_graph(name="vision_encoder") pixel_values = builder.input( "pixel_values", dtype=config.dtype, - shape=[1, 3, image_size, image_size], + shape=[total_patches, vision.patch_dim], + ) + # (grid_t, grid_h, grid_w) for the single packed image or video. + grid_thw = builder.input( + "grid_thw", + dtype=ir.DataType.INT64, + shape=[3], ) - image_features = vision(builder.op, pixel_values=pixel_values) - image_features = builder.op.Squeeze(image_features, [0]) + image_features = vision(builder.op, pixel_values=pixel_values, grid_thw=grid_thw) builder.add_output(image_features, "image_features") return _make_model(graph) + def _build_embedding( + self, + embedding: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build ``input_ids + image_features + video_features → inputs_embeds``.""" + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + num_image_tokens = ir.SymbolicDim("num_image_tokens") + num_video_tokens = ir.SymbolicDim("num_video_tokens") + + graph, builder = _make_graph(name="embedding") + input_ids = builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], + ) + image_features = builder.input( + "image_features", + dtype=config.dtype, + shape=[num_image_tokens, config.hidden_size], + ) + video_features = builder.input( + "video_features", + dtype=config.dtype, + shape=[num_video_tokens, config.hidden_size], + ) + inputs_embeds = embedding( + builder.op, + input_ids=input_ids, + image_features=image_features, + video_features=video_features, + ) + builder.add_output(inputs_embeds, "inputs_embeds") + return _make_model(graph) + class QwenVLTask(VisionLanguageTask): """Qwen-family VL 3-model split with packed-attention vision and MRoPE. diff --git a/src/mobius/tasks/_wan_vae.py b/src/mobius/tasks/_wan_vae.py new file mode 100644 index 000000000..f502d21c7 --- /dev/null +++ b/src/mobius/tasks/_wan_vae.py @@ -0,0 +1,136 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Wan video VAE task: 5D encoder and decoder ONNX graphs. + +Builds a :class:`~mobius._model_package.ModelPackage` with two graphs: + +``encoder`` + ``sample`` ``(B, video_channels, frames, height, width)`` + -> ``latent_mean`` / ``latent_logvar`` ``(B, z_dim, T', H', W')`` + -> ``latent`` — the deterministic, pipeline-normalised latent + ``(latent_mean - latents_mean) / latents_std``. + +``decoder`` + ``latent`` ``(B, z_dim, T', H', W')`` (pipeline-normalised) + -> ``sample`` ``(B, video_channels, frames, height, width)``. + +Both graphs are fully dynamic in batch, temporal and spatial extent. The +temporal extent must satisfy ``frames = scale_factor_temporal * k + 1`` with +``k >= 0`` (encoder) and ``latent_frames = k + 1 >= 1`` (decoder); ``k == 0`` +is the single-frame image mode used by text-to-image pipelines such as Cosmos3. +See the "Single-frame (image) mode" note in :mod:`mobius.models.wan_vae`. +``height``/``width`` must be divisible by ``scale_factor_spatial``. + +Why a dedicated task instead of reusing :class:`~mobius.tasks._vae.VAETask` or +:class:`~mobius.tasks._qwen_image_vae.QwenImageVAETask`: + +* ``VAETask`` declares 4D image I/O (``batch, C, height, width``) and reads + ``config.latent_channels``; the Wan VAE is inherently 5D video. +* Both existing tasks emit a single fused ``latent_dist`` tensor and perform no + posterior split, so they cannot expose deterministic mean/logvar moments. +* Neither models the ``WanPipeline`` latent mean/std normalisation, which has to + straddle the graph boundary (see below). +* Neither applies ``patchify``/``unpatchify``, which upstream performs inside + ``AutoencoderKLWan._encode``/``._decode`` rather than inside the encoder or + decoder sub-module. + +**Where the latent statistics live.** Upstream keeps ``latents_mean`` / +``latents_std`` out of ``AutoencoderKLWan`` entirely — the ``WanPipeline`` +normalises after ``vae.encode`` and denormalises before ``vae.decode``. Mobius +mirrors that split: the ``encoder``/``decoder`` ``nn.Module``s stay byte-identical +to the checkpoint, and the normalisation is emitted here, at the graph boundary, +via :meth:`~mobius.models.wan_vae.AutoencoderKLWanModel.normalize_latents` and +:meth:`~mobius.models.wan_vae.AutoencoderKLWanModel.denormalize_latents`. The +raw ``latent_mean``/``latent_logvar`` moments are exported alongside the +normalised latent so callers that need to sample the posterior (or reuse the +un-normalised scale) can do so outside the graph. +""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir + +from mobius._configs._wan_vae import WanVAEConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ComponentSpec, ModelTask, _make_graph, _make_model + + +class WanVAETask(ModelTask): + """Build the Wan 3D causal video VAE encoder and decoder graphs.""" + + model_roles: ClassVar[dict[str, str]] = {"encoder": "encoder", "decoder": "decoder"} + components: ClassVar[ComponentSpec] = ComponentSpec(encoder="encoder", decoder="decoder") + + def build(self, module, config: WanVAEConfig) -> ModelPackage: + """Build both graphs. + + Args: + module: An :class:`~mobius.models.wan_vae.AutoencoderKLWanModel`. + config: The parsed Wan VAE configuration. + + Returns: + A package with ``"encoder"`` and ``"decoder"`` models. + """ + self._validate_components(module) + return ModelPackage( + { + "encoder": self._build_encoder_graph(module, config), + "decoder": self._build_decoder_graph(module, config), + }, + config=config, + ) + + def _build_encoder_graph(self, module, config: WanVAEConfig) -> ir.Model: + graph, builder = _make_graph(name="wan_vae_encoder") + op = builder.op + + # Pixel-space video in [-1, 1]. ``frames`` must be + # ``scale_factor_temporal * k + 1`` (k >= 0) so the temporal downsampling + # stages line up with upstream's chunked encode; k == 0 is a single image. + sample = builder.input( + "sample", + dtype=config.dtype, + shape=["batch", config.video_channels, "frames", "height", "width"], + ) + + latent_mean, latent_logvar = module.encode(op, sample) + latent = module.normalize_latents(op, latent_mean) + + # Annotate the latent grid symbolically: shape inference cannot see + # through the dynamic Reshape/Pad chain of the resampling blocks. + latent_shape = ir.Shape( + ["batch", config.z_dim, "latent_frames", "latent_height", "latent_width"] + ) + for value, name in ( + (latent_mean, "latent_mean"), + (latent_logvar, "latent_logvar"), + (latent, "latent"), + ): + value.shape = latent_shape + builder.add_output(value, name) + + return _make_model(graph) + + def _build_decoder_graph(self, module, config: WanVAEConfig) -> ir.Model: + graph, builder = _make_graph(name="wan_vae_decoder") + op = builder.op + + # Pipeline-normalised latents, i.e. the ``latent`` output of the encoder + # graph or the denoiser's sample. ``latent_frames`` may be 1 (image mode). + latent = builder.input( + "latent", + dtype=config.dtype, + shape=["batch", config.z_dim, "latent_frames", "latent_height", "latent_width"], + ) + + sample = module.decode(op, module.denormalize_latents(op, latent)) + + sample.shape = ir.Shape( + ["batch", config.decoded_video_channels, "frames", "height", "width"] + ) + builder.add_output(sample, "sample") + + return _make_model(graph) diff --git a/src/mobius/tasks/_world_model.py b/src/mobius/tasks/_world_model.py index d6b6a3863..52bcfcd29 100644 --- a/src/mobius/tasks/_world_model.py +++ b/src/mobius/tasks/_world_model.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Single-step world-model task.""" +"""Single-step latent-dynamics task.""" from __future__ import annotations @@ -10,13 +10,13 @@ import onnx_ir as ir from onnxscript import nn -from mobius._configs import WorldModelConfig +from mobius._configs import LatentDynamicsConfig from mobius._model_package import ModelPackage from mobius.tasks._base import ModelTask, _make_graph, _make_model -class WorldModelTask(ModelTask): - """Build a stateful one-step world-model graph. +class LatentDynamicsTask(ModelTask): + """Build a stateful one-step latent-dynamics graph. Inputs: - observation: ``[batch, *observation_shape]`` @@ -42,7 +42,7 @@ class WorldModelTask(ModelTask): def build( self, module: nn.Module, - config: WorldModelConfig, + config: LatentDynamicsConfig, ) -> ModelPackage: config.validate() batch = ir.SymbolicDim("batch") @@ -82,3 +82,8 @@ def build( builder.add_output(value, name) return ModelPackage({"model": _make_model(graph)}, config=config) + + +# Backward-compatible alias for the task name introduced by the initial +# single-step implementation. +WorldModelTask = LatentDynamicsTask diff --git a/testdata/cases/vision-language/cosmos3-omni-nano.yaml b/testdata/cases/vision-language/cosmos3-omni-nano.yaml index f492c46a7..0f25cdbee 100644 --- a/testdata/cases/vision-language/cosmos3-omni-nano.yaml +++ b/testdata/cases/vision-language/cosmos3-omni-nano.yaml @@ -17,8 +17,8 @@ generation: do_sample: false notes: > - Cosmos3-Omni understanding tower ("Reasoner"). Architecturally identical to - Qwen3-VL (interleaved 3D M-RoPE, QK-norm, DeepStack layers [8, 16, 24]); the - reasoner sub-tower of the unified Cosmos3 checkpoint is exported via the - Qwen3-VL 3-model split with DeepStack preserved. Generator / Sound / Action - towers are dropped. + This VLM case validates the Cosmos3-Omni understanding tower ("Reasoner") + independently. Architecturally it is Qwen3-VL (interleaved 3D M-RoPE, + QK-norm, DeepStack layers [8, 16, 24]). Complete world-model export, + including Generator, Wan VAE, Sound, and Action, uses + build_world_model() / --features world-model. diff --git a/tests/_cosmos3_edge_reference.py b/tests/_cosmos3_edge_reference.py new file mode 100644 index 000000000..76e62c569 --- /dev/null +++ b/tests/_cosmos3_edge_reference.py @@ -0,0 +1,519 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Faithful PyTorch reference for the public NVIDIA Cosmos3-Edge Reasoner. + +This module re-implements the *published* Cosmos3-Edge vision tower, merger +projector, interleaved M-RoPE and text decoder so mobius' ONNX graphs can be +compared numerically against a trusted implementation. Every routine below is +a direct transcription of authoritative upstream sources: + +- ``huggingface/transformers`` — + ``src/transformers/models/cosmos3_edge/modular_cosmos3_edge.py`` + (``Cosmos3EdgeVisionEmbeddings.resize_positional_embeddings``, + ``Cosmos3EdgeVisionAttention``, ``Cosmos3EdgePatchMerger``, + ``Cosmos3EdgeTextRotaryEmbedding``, ``Cosmos3EdgeImageProcessor.patchify``, + ``Cosmos3EdgeVideoProcessor.patchify``). +- ``vllm-project/vllm`` — ``vllm/model_executor/models/cosmos3_edge.py`` + (independent confirmation of the projector shape contract). + +The reference is intentionally dependency-light: it consumes plain tensors and +a small config dataclass so it can run against the real checkpoint without +requiring a ``transformers`` release that ships ``cosmos3_edge``. + +.. note:: + Only the *understanding* (Reasoner) tower is reproduced. The Cosmos3-Edge + Generator/Action towers that share the checkpoint are proprietary + rectified-flow components and are out of scope here. +""" + +from __future__ import annotations + +import dataclasses +import math + +import torch +import torch.nn.functional as functional + +__all__ = [ + "EdgeRefConfig", + "merge_patches", + "patchify_images", + "patchify_videos", + "ref_interleaved_mrope_cos_sin", + "ref_projector", + "ref_text_decoder_logits", + "ref_vision_features", + "ref_vision_tower", + "resize_positional_embeddings", + "smart_resize", +] + + +@dataclasses.dataclass +class EdgeRefConfig: + """The subset of ``config.json`` the reference needs.""" + + # Vision tower (``vision_config``). + vision_hidden_size: int = 1152 + vision_intermediate_size: int = 4304 + vision_num_layers: int = 27 + vision_num_heads: int = 16 + patch_size: int = 16 + num_channels: int = 3 + num_patches: int = 256 + layer_norm_eps: float = 1e-6 + spatial_merge_size: int = 2 + # Projector (``projector_config``). + projector_hidden_size: int = 11520 + use_postshuffle_norm: bool = False + # Text tower (``text_config``). + hidden_size: int = 2048 + intermediate_size: int = 9216 + num_hidden_layers: int = 28 + num_attention_heads: int = 16 + num_key_value_heads: int = 8 + head_dim: int = 128 + rms_norm_eps: float = 1e-5 + rope_theta: float = 100_000_000.0 + mrope_section: tuple[int, int, int] = (24, 20, 20) + vocab_size: int = 131072 + + @classmethod + def from_hf_config(cls, config: dict) -> EdgeRefConfig: + """Build from a raw ``config.json`` mapping.""" + vision = config["vision_config"] + projector = config["projector_config"] + text = config["text_config"] + rope = text["rope_parameters"] + return cls( + vision_hidden_size=vision["hidden_size"], + vision_intermediate_size=vision["intermediate_size"], + vision_num_layers=vision["num_hidden_layers"], + vision_num_heads=vision["num_attention_heads"], + patch_size=vision["patch_size"], + num_channels=vision["num_channels"], + num_patches=vision["num_patches"], + layer_norm_eps=vision["layer_norm_eps"], + spatial_merge_size=vision["spatial_merge_size"], + projector_hidden_size=projector["merger_intermediate_size"], + use_postshuffle_norm=projector.get("use_postshuffle_norm", False), + hidden_size=text["hidden_size"], + intermediate_size=text["intermediate_size"], + num_hidden_layers=text["num_hidden_layers"], + num_attention_heads=text["num_attention_heads"], + num_key_value_heads=text["num_key_value_heads"], + head_dim=text["head_dim"], + rms_norm_eps=text["rms_norm_eps"], + rope_theta=rope["rope_theta"], + mrope_section=tuple(rope["mrope_section"]), + vocab_size=text["vocab_size"], + ) + + +# ────────────────────────────────────────────────────────────────────────── +# Preprocessing (Cosmos3EdgeImageProcessor / Cosmos3EdgeVideoProcessor) +# ────────────────────────────────────────────────────────────────────────── + + +def smart_resize( + height: int, + width: int, + *, + num_frames: int = 1, + temporal_factor: int = 1, + factor: int = 32, + min_pixels: int = 256 * 256, + max_pixels: int = 4096 * 4096, +) -> tuple[int, int]: + """Exact transcription of ``Cosmos3EdgeImageProcessor.smart_resize``. + + The output side lengths are multiples of ``factor`` + (``patch_size * merge_size`` = 32 for Cosmos3-Edge) and the total + ``num_frames * height * width`` volume is clamped to the processor's + ``size.shortest_edge`` / ``size.longest_edge``, which are *areas* + (``256*256`` / ``4096*4096`` for images, ``64*64`` / ``24*1024*1024`` for + videos), not edge lengths. + + The processor resamples with **bicubic** interpolation, then rescales by + ``1/255`` and normalises with ``mean = std = 0.5``. + """ + if num_frames < temporal_factor: + raise ValueError( + f"t:{num_frames} must be larger than temporal_factor:{temporal_factor}" + ) + if height < factor or width < factor: + scale = max(factor / height, factor / width) + height = int(height * scale) + width = int(width * scale) + + if max(height, width) / min(height, width) > 200: + raise ValueError( + "absolute aspect ratio must be smaller than 200, got " + f"{max(height, width) / min(height, width)}" + ) + h_bar = round(height / factor) * factor + w_bar = round(width / factor) * factor + t_bar = round(num_frames / temporal_factor) * temporal_factor + + if t_bar * h_bar * w_bar > max_pixels: + beta = math.sqrt((num_frames * height * width) / max_pixels) + h_bar = max(factor, math.floor(height / beta / factor) * factor) + w_bar = max(factor, math.floor(width / beta / factor) * factor) + elif t_bar * h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (num_frames * height * width)) + h_bar = math.ceil(height * beta / factor) * factor + w_bar = math.ceil(width * beta / factor) * factor + + return h_bar, w_bar + + +def patchify_images( + images: torch.Tensor, + *, + patch_size: int, + merge_size: int, + temporal_patch_size: int = 1, +) -> tuple[torch.Tensor, int, int]: + """``Cosmos3EdgeImageProcessor.patchify``: block-major, HWC-in-patch. + + ``images`` is ``[B, C, H, W]``. Returns ``([B, grid_h*grid_w, patch_dim], + grid_h, grid_w)`` where ``patch_dim = patch*patch*C*temporal_patch_size`` + and the values inside a patch are ordered ``(patch_h, patch_w, channel)`` + (channel-**last**), not the usual ``(channel, patch_h, patch_w)``. + """ + batch_size, channel, height, width = images.shape + grid_h, grid_w = height // patch_size, width // patch_size + patches = images.reshape( + batch_size, + channel, + grid_h // merge_size, + merge_size, + patch_size, + grid_w // merge_size, + merge_size, + patch_size, + ) + # (B, gh/m, gw/m, m, m, ph, pw, C) — 2x2 blocks are contiguous. + patches = patches.permute(0, 2, 5, 3, 6, 4, 7, 1) + flatten_patches = ( + patches.unsqueeze(-1) + .expand(-1, -1, -1, -1, -1, -1, -1, -1, temporal_patch_size) + .reshape( + batch_size, + grid_h * grid_w, + patch_size * patch_size * channel * temporal_patch_size, + ) + ) + return flatten_patches, grid_h, grid_w + + +def patchify_videos( + videos: torch.Tensor, + *, + patch_size: int, + merge_size: int, + temporal_patch_size: int = 1, +) -> tuple[torch.Tensor, int, int, int]: + """``Cosmos3EdgeVideoProcessor.patchify``: ``[B, T, C, H, W]`` → packed. + + Cosmos3-Edge only supports ``temporal_patch_size=1``, so every frame + contributes its own ``grid_h*grid_w`` block-major patch run. + """ + batch_size, num_frames, channel, height, width = videos.shape + grid_t = num_frames // temporal_patch_size + grid_h, grid_w = height // patch_size, width // patch_size + patches = videos.view( + batch_size, + grid_t, + temporal_patch_size, + channel, + grid_h // merge_size, + merge_size, + patch_size, + grid_w // merge_size, + merge_size, + patch_size, + ) + patches = patches.permute(0, 1, 4, 7, 5, 8, 6, 9, 3, 2) + flatten_patches = patches.reshape( + batch_size, + grid_t * grid_h * grid_w, + patch_size * patch_size * channel * temporal_patch_size, + ) + return flatten_patches, grid_t, grid_h, grid_w + + +# ────────────────────────────────────────────────────────────────────────── +# Vision tower +# ────────────────────────────────────────────────────────────────────────── + + +def resize_positional_embeddings( + positional_embeddings: torch.Tensor, + grid_thw: torch.Tensor, + spatial_merge_size: int, +) -> torch.Tensor: + """``Cosmos3EdgeVisionEmbeddings.resize_positional_embeddings``. + + ``positional_embeddings`` is the learned ``(g, g, D)`` reference grid. It + is bilinearly resampled (antialiased, ``align_corners=False``) to each + item's ``(H, W)`` grid and then reordered into the processor's block-major + 2x2 layout before being repeated ``T`` times. + """ + positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0) + source_dtype = positional_embeddings.dtype + if positional_embeddings.device.type == "cpu": + positional_embeddings = positional_embeddings.float() + + position_chunks = [] + for temporal, height, width in grid_thw.tolist(): + resized = functional.interpolate( + positional_embeddings, + size=(height, width), + mode="bilinear", + align_corners=False, + antialias=True, + ) + resized = resized.squeeze(0).permute(1, 2, 0).to(source_dtype) + resized = resized.reshape( + height // spatial_merge_size, + spatial_merge_size, + width // spatial_merge_size, + spatial_merge_size, + -1, + ) + resized = resized.transpose(1, 2).reshape(height * width, -1) + position_chunks.append(resized.repeat(temporal, 1)) + + return torch.cat(position_chunks, dim=0) + + +def _vision_attention( + hidden: torch.Tensor, + weights: dict[str, torch.Tensor], + prefix: str, + num_heads: int, +) -> torch.Tensor: + """Packed SigLIP2 attention; ``hidden`` is ``[frames, seq, D]``.""" + frames, seq, dim = hidden.shape + head_dim = dim // num_heads + scale = head_dim**-0.5 + + def linear(name: str, x: torch.Tensor) -> torch.Tensor: + return functional.linear( + x, weights[f"{prefix}.{name}.weight"], weights[f"{prefix}.{name}.bias"] + ) + + query = linear("q_proj", hidden).view(frames, seq, num_heads, head_dim).transpose(1, 2) + key = linear("k_proj", hidden).view(frames, seq, num_heads, head_dim).transpose(1, 2) + value = linear("v_proj", hidden).view(frames, seq, num_heads, head_dim).transpose(1, 2) + attn = functional.scaled_dot_product_attention(query, key, value, scale=scale) + attn = attn.transpose(1, 2).reshape(frames, seq, dim) + return linear("out_proj", attn) + + +def _layer_norm( + hidden: torch.Tensor, weights: dict[str, torch.Tensor], prefix: str, eps: float +) -> torch.Tensor: + return functional.layer_norm( + hidden, + (hidden.shape[-1],), + weights[f"{prefix}.weight"], + weights[f"{prefix}.bias"], + eps, + ) + + +def ref_vision_tower( + pixel_values: torch.Tensor, + grid_thw: torch.Tensor, + weights: dict[str, torch.Tensor], + config: EdgeRefConfig, +) -> torch.Tensor: + """SigLIP2 tower: packed ``[N, patch_dim]`` → ``[N, vision_hidden]``. + + ``grid_thw`` must describe a single item (``[1, 3]``); attention runs + independently per frame, exactly as ``Cosmos3EdgeEncoder`` does with + per-frame ``cu_seqlens``. + """ + assert grid_thw.shape[0] == 1, "reference handles one packed item at a time" + temporal, height, width = (int(v) for v in grid_thw[0].tolist()) + + hidden = functional.linear( + pixel_values, + weights["visual.embeddings.patch_embedding.weight"], + weights["visual.embeddings.patch_embedding.bias"], + ) + grid = math.isqrt(config.num_patches) + position_table = weights["visual.embeddings.position_embedding.weight"].reshape( + grid, grid, -1 + ) + hidden = hidden + resize_positional_embeddings( + position_table, grid_thw, config.spatial_merge_size + ) + + # Per-frame attention == batched attention over equal-length sequences. + hidden = hidden.reshape(temporal, height * width, -1) + for index in range(config.vision_num_layers): + prefix = f"visual.encoder.layers.{index}" + residual = hidden + normed = _layer_norm(hidden, weights, f"{prefix}.layer_norm1", config.layer_norm_eps) + hidden = residual + _vision_attention( + normed, weights, f"{prefix}.self_attn", config.vision_num_heads + ) + residual = hidden + normed = _layer_norm(hidden, weights, f"{prefix}.layer_norm2", config.layer_norm_eps) + normed = functional.linear( + normed, + weights[f"{prefix}.mlp.fc1.weight"], + weights[f"{prefix}.mlp.fc1.bias"], + ) + normed = functional.gelu(normed, approximate="tanh") + normed = functional.linear( + normed, + weights[f"{prefix}.mlp.fc2.weight"], + weights[f"{prefix}.mlp.fc2.bias"], + ) + hidden = residual + normed + + hidden = _layer_norm(hidden, weights, "visual.post_layernorm", config.layer_norm_eps) + return hidden.reshape(temporal * height * width, -1) + + +def merge_patches(features: torch.Tensor, spatial_merge_size: int) -> torch.Tensor: + """Group ``spatial_merge_size**2`` consecutive (block-major) patches.""" + return features.reshape(-1, spatial_merge_size**2, features.shape[-1]) + + +def ref_projector( + features: torch.Tensor, + weights: dict[str, torch.Tensor], + config: EdgeRefConfig, +) -> torch.Tensor: + """``Cosmos3EdgePatchMerger``: LayerNorm → shuffle → fc1 → GELU → fc2.""" + merged_dim = config.vision_hidden_size * config.spatial_merge_size**2 + hidden = merge_patches(features, config.spatial_merge_size) + if config.use_postshuffle_norm: + hidden = _layer_norm(hidden.reshape(-1, merged_dim), weights, "projector.norm", 1e-6) + else: + hidden = _layer_norm(hidden, weights, "projector.norm", 1e-6).reshape(-1, merged_dim) + hidden = functional.linear( + hidden, weights["projector.linear_fc1.weight"], weights["projector.linear_fc1.bias"] + ) + hidden = functional.gelu(hidden) + return functional.linear( + hidden, weights["projector.linear_fc2.weight"], weights["projector.linear_fc2.bias"] + ) + + +def ref_vision_features( + pixel_values: torch.Tensor, + grid_thw: torch.Tensor, + weights: dict[str, torch.Tensor], + config: EdgeRefConfig, +) -> torch.Tensor: + """Full ``pixel_values`` → projected ``[num_vision_tokens, text_hidden]``.""" + return ref_projector( + ref_vision_tower(pixel_values, grid_thw, weights, config), weights, config + ) + + +# ────────────────────────────────────────────────────────────────────────── +# Text tower +# ────────────────────────────────────────────────────────────────────────── + + +def ref_interleaved_mrope_cos_sin( + position_ids: torch.Tensor, config: EdgeRefConfig +) -> tuple[torch.Tensor, torch.Tensor]: + """``Cosmos3EdgeTextRotaryEmbedding``: axis-interleaved 3D M-RoPE. + + ``position_ids`` is ``[3, batch, seq]``. Frequency channel ``i`` is driven + by the height axis when ``i % 3 == 1 and i < 3*mrope_section[1]``, by the + width axis when ``i % 3 == 2 and i < 3*mrope_section[2]``, and by the + temporal axis otherwise. + """ + dim = config.head_dim + inv_freq = 1.0 / ( + config.rope_theta ** (torch.arange(0, dim, 2, dtype=torch.float64) / dim) + ) + indices = torch.arange(inv_freq.shape[0]) + height_mask = (indices % 3 == 1) & (indices < config.mrope_section[1] * 3) + width_mask = (indices % 3 == 2) & (indices < config.mrope_section[2] * 3) + temporal_mask = ~(height_mask | width_mask) + inv_freq = torch.stack( + (inv_freq * temporal_mask, inv_freq * height_mask, inv_freq * width_mask) + ) + freqs = position_ids.permute(1, 2, 0).double() @ inv_freq + emb = torch.cat((freqs, freqs), dim=-1) + return emb.cos().float(), emb.sin().float() + + +def _rms_norm(hidden: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + variance = hidden.float().pow(2).mean(-1, keepdim=True) + return weight * (hidden.float() * torch.rsqrt(variance + eps)).to(hidden.dtype) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + return torch.cat((-x[..., half:], x[..., :half]), dim=-1) + + +def ref_text_decoder_logits( + inputs_embeds: torch.Tensor, + position_ids: torch.Tensor, + weights: dict[str, torch.Tensor], + config: EdgeRefConfig, +) -> torch.Tensor: + """Causal GQA decoder with squared-ReLU FFN and interleaved M-RoPE. + + Mirrors ``Cosmos3EdgeTextModel`` (a ``LlamaModel`` with + ``Cosmos3EdgeTextMLP`` = ``fc1 → relu² → fc2`` and no QK-norm). + ``inputs_embeds`` is ``[batch, seq, hidden]`` and ``position_ids`` is + ``[3, batch, seq]``. + """ + batch, seq, _ = inputs_embeds.shape + heads = config.num_attention_heads + kv_heads = config.num_key_value_heads + head_dim = config.head_dim + + cos, sin = ref_interleaved_mrope_cos_sin(position_ids, config) + cos = cos.unsqueeze(1).to(inputs_embeds.dtype) + sin = sin.unsqueeze(1).to(inputs_embeds.dtype) + + hidden = inputs_embeds + for index in range(config.num_hidden_layers): + prefix = f"layers.{index}" + residual = hidden + normed = _rms_norm( + hidden, weights[f"{prefix}.input_layernorm.weight"], config.rms_norm_eps + ) + query = functional.linear(normed, weights[f"{prefix}.self_attn.to_q.weight"]) + key = functional.linear(normed, weights[f"{prefix}.self_attn.to_k.weight"]) + value = functional.linear(normed, weights[f"{prefix}.self_attn.to_v.weight"]) + query = query.view(batch, seq, heads, head_dim).transpose(1, 2) + key = key.view(batch, seq, kv_heads, head_dim).transpose(1, 2) + value = value.view(batch, seq, kv_heads, head_dim).transpose(1, 2) + query = query * cos + _rotate_half(query) * sin + key = key * cos + _rotate_half(key) * sin + attn = functional.scaled_dot_product_attention( + query, key, value, is_causal=True, enable_gqa=True + ) + attn = attn.transpose(1, 2).reshape(batch, seq, heads * head_dim) + hidden = residual + functional.linear( + attn, weights[f"{prefix}.self_attn.to_out.weight"] + ) + + residual = hidden + normed = _rms_norm( + hidden, weights[f"{prefix}.post_attention_layernorm.weight"], config.rms_norm_eps + ) + normed = functional.linear(normed, weights[f"{prefix}.mlp.up_proj.weight"]) + normed = torch.square(functional.relu(normed)) + hidden = residual + functional.linear( + normed, weights[f"{prefix}.mlp.down_proj.weight"] + ) + + hidden = _rms_norm(hidden, weights["norm.weight"], config.rms_norm_eps) + return functional.linear(hidden, weights["lm_head.weight"]) diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 19ea16689..06605a0fe 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -2095,12 +2095,20 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: intermediate_size=64, num_hidden_layers=1, num_attention_heads=2, - image_size=28, - patch_size=14, + image_size=64, + patch_size=16, + # Learned position-embedding reference grid (4x4), resampled per image. + num_patches=16, norm_eps=1e-6, spatial_merge_size=2, + temporal_patch_size=1, out_hidden_size=64, projector_intermediate_size=64, + use_postshuffle_norm=False, + image_token_id=19, + video_token_id=18, + vision_start_token_id=20, + vision_end_token_id=21, ) @@ -2121,6 +2129,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: "hidden_act": "relu2", "mlp_bias": False, "mrope_section": [24, 20, 20], + "mrope_interleaved": True, }, True, ), diff --git a/tests/cli_test.py b/tests/cli_test.py index 7a3e7cf6e..35c7a3953 100644 --- a/tests/cli_test.py +++ b/tests/cli_test.py @@ -143,6 +143,68 @@ def test_text_only_skips_diffusers_autodetect(self): mock_build.assert_called_once() assert mock_build.call_args.kwargs.get("text_only") is True + def test_world_model_feature_dispatches_complete_pipeline(self): + package = mock.MagicMock() + with ( + tempfile.TemporaryDirectory() as tmpdir, + mock.patch( + "mobius._diffusers_builder._load_diffusers_pipeline_index" + ) as mock_diffusers, + mock.patch( + "mobius._world_model_builder.build_world_model", + return_value=package, + ) as mock_build_world_model, + mock.patch("mobius.__main__._save_package") as mock_save, + ): + main( + [ + "build", + "--model", + "nvidia/Cosmos3-Nano", + tmpdir, + "--features", + "world-model", + "--no-weights", + "--dtype", + "bf16", + "--ep", + "cuda", + ] + ) + + mock_diffusers.assert_not_called() + mock_build_world_model.assert_called_once() + kwargs = mock_build_world_model.call_args.kwargs + assert kwargs["load_weights"] is False + assert kwargs["execution_provider"] == "cuda" + mock_save.assert_called_once_with( + package, tmpdir, mock.ANY, mock.ANY, component_filter=None + ) + + @pytest.mark.parametrize( + ("extra_args", "message"), + [ + (["--task", "qwen-vl"], "--task"), + (["--component", "generator"], "--component"), + (["--runtime", "onnx-genai"], "--runtime"), + (["--features", "text-only"], "decoder-specific"), + ], + ) + def test_world_model_rejects_incompatible_options(self, extra_args, message): + with tempfile.TemporaryDirectory() as tmpdir, pytest.raises(SystemExit, match=message): + main( + [ + "build", + "--model", + "nvidia/Cosmos3-Nano", + tmpdir, + "--features", + "world-model", + "--no-weights", + *extra_args, + ] + ) + def test_build_static_cache(self): with tempfile.TemporaryDirectory() as tmpdir: main( diff --git a/tests/cosmos3_edge_integration_test.py b/tests/cosmos3_edge_integration_test.py new file mode 100644 index 000000000..0e983710c --- /dev/null +++ b/tests/cosmos3_edge_integration_test.py @@ -0,0 +1,411 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Integration tests: NVIDIA Cosmos3-Edge Reasoner against real weights. + +Verifies the exported Reasoner graphs stage by stage using the real +``nvidia/Cosmos3-Edge`` checkpoint and the PyTorch transcription of the +published ``cosmos3_edge`` modeling code in +:mod:`tests._cosmos3_edge_reference`:: + + pytest tests/cosmos3_edge_integration_test.py -m integration -sv + +Stages checked: SigLIP2 vision tower + merger projector (image and video), +image/video token fusion in the embedding graph, decoder logits under +interleaved 3D M-RoPE, and a cached greedy-decode generation smoke test. + +.. note:: + Only the *understanding* (Reasoner) tower is verified. The Cosmos3-Edge + Generator, Action head and Sound tower share the same checkpoint but are + proprietary rectified-flow components with no published reference + implementation, so their numerics remain unverifiable here. +""" + +from __future__ import annotations + +import json +import os + +import numpy as np +import pytest +import torch + +from tests._cosmos3_edge_reference import ( + EdgeRefConfig, + patchify_images, + patchify_videos, + ref_text_decoder_logits, + ref_vision_features, + smart_resize, +) + +MODEL_ID = "nvidia/Cosmos3-Edge" +PATCH_SIZE = 16 +MERGE_SIZE = 2 +IMAGE_TOKEN_ID = 19 +VIDEO_TOKEN_ID = 18 + +pytestmark = [pytest.mark.integration, pytest.mark.integration_slow] + + +def _checkpoint_dir() -> str: + from huggingface_hub import snapshot_download + + return snapshot_download(MODEL_ID, allow_patterns=["*.json", "*/*.safetensors"]) + + +def _normalized(image: torch.Tensor) -> torch.Tensor: + """Apply the checkpoint's ``image_mean``/``image_std`` of 0.5.""" + return (image - 0.5) / 0.5 + + +def _split_image(height: int, width: int) -> torch.Tensor: + """Left half red, right half green — a deterministic, describable image.""" + image = torch.zeros(1, 3, height, width, dtype=torch.float32) + image[:, 0, :, : width // 2] = 1.0 + image[:, 1, :, width // 2 :] = 1.0 + return _normalized(image) + + +@pytest.fixture(scope="module") +def edge_package(tmp_path_factory): + """Build the Reasoner graphs with real fp32 weights and open ORT sessions.""" + import onnx_ir as ir + import onnxruntime as ort + + import mobius + from mobius._weight_loading import iter_weight_shards + from mobius.models.cosmos import Cosmos3EdgeVLModel + + snapshot = _checkpoint_dir() + with open(os.path.join(snapshot, "config.json"), encoding="utf-8") as handle: + ref_config = EdgeRefConfig.from_hf_config(json.load(handle)) + + package = mobius.build(MODEL_ID, task="cosmos3-edge-vl", dtype="f32", load_weights=False) + module = Cosmos3EdgeVLModel(package.config) + reference: dict[str, torch.Tensor] = {} + for shard in iter_weight_shards(MODEL_ID): + package.apply_weights_partial(module.preprocess_weights(shard)) + for key, value in shard.items(): + if "k_norm_und_for_gen" in key or "moe_gen" in key: + continue + if key.startswith(("model.visual.", "model.projector.")): + reference[key.removeprefix("model.")] = value.float() + elif key.startswith("layers.") or key in ( + "embed_tokens.weight", + "norm.weight", + "lm_head.weight", + ): + reference[key] = value.float() + package.finalize_weights() + package.validate_weights() + + directory = tmp_path_factory.mktemp("cosmos3_edge_real") + sessions = {} + for name in ("vision_encoder", "embedding", "decoder"): + path = directory / f"{name}.onnx" + ir.save(package[name], str(path), external_data=f"{name}.onnx.data") + sessions[name] = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + return package.config, ref_config, sessions, reference + + +def test_vision_encoder_input_contract(edge_package): + """The exported vision graph takes packed patches, not a single image.""" + _, _, sessions, _ = edge_package + inputs = {value.name: value for value in sessions["vision_encoder"].get_inputs()} + assert set(inputs) == {"pixel_values", "grid_thw"} + assert inputs["pixel_values"].shape[1] == PATCH_SIZE * PATCH_SIZE * 3 + assert inputs["grid_thw"].shape == [3] + + +@pytest.mark.parametrize(("height", "width"), [(256, 256), (128, 512), (320, 192)]) +def test_vision_features_match_reference(edge_package, height, width): + _, ref_config, sessions, reference = edge_package + image = _split_image(height, width) + packed, grid_h, grid_w = patchify_images( + image, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + packed = packed[0] + + got = sessions["vision_encoder"].run( + None, + { + "pixel_values": packed.numpy(), + "grid_thw": np.array([1, grid_h, grid_w], dtype=np.int64), + }, + )[0] + expected = ref_vision_features( + packed, torch.tensor([[1, grid_h, grid_w]]), reference, ref_config + ).numpy() + + assert got.shape == (grid_h * grid_w // MERGE_SIZE**2, ref_config.hidden_size) + np.testing.assert_allclose(got, expected, atol=1e-3, rtol=1e-3) + correlation = np.corrcoef(got.reshape(-1), expected.reshape(-1))[0, 1] + assert correlation > 0.9999 + + +def test_smart_resized_photo_resolution_matches_reference(edge_package): + """Drive the processor's own ``smart_resize`` policy, not a hand-picked size. + + A natural 1000x750 photo is resized to a multiple of ``patch*merge`` (32) + inside the checkpoint's pixel-area bounds, then patchified and encoded. + """ + _, ref_config, sessions, reference = edge_package + height, width = smart_resize( + 750, + 1000, + factor=PATCH_SIZE * MERGE_SIZE, + min_pixels=256 * 256, + max_pixels=4096 * 4096, + ) + assert height % (PATCH_SIZE * MERGE_SIZE) == 0 + assert width % (PATCH_SIZE * MERGE_SIZE) == 0 + assert 256 * 256 <= height * width <= 4096 * 4096 + + image = _split_image(height, width) + packed, grid_h, grid_w = patchify_images( + image, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + got = sessions["vision_encoder"].run( + None, + { + "pixel_values": packed[0].numpy(), + "grid_thw": np.array([1, grid_h, grid_w], dtype=np.int64), + }, + )[0] + expected = ref_vision_features( + packed[0], torch.tensor([[1, grid_h, grid_w]]), reference, ref_config + ).numpy() + + assert got.shape == (grid_h * grid_w // MERGE_SIZE**2, ref_config.hidden_size) + np.testing.assert_allclose(got, expected, atol=1e-3, rtol=1e-3) + + +def test_video_features_match_reference_and_token_count(edge_package): + _, ref_config, sessions, reference = edge_package + frames = 4 + video = torch.stack([_split_image(128, 160)[0] for _ in range(frames)]).unsqueeze(0) + packed, grid_t, grid_h, grid_w = patchify_videos( + video, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + packed = packed[0] + + got = sessions["vision_encoder"].run( + None, + { + "pixel_values": packed.numpy(), + "grid_thw": np.array([grid_t, grid_h, grid_w], dtype=np.int64), + }, + )[0] + expected = ref_vision_features( + packed, torch.tensor([[grid_t, grid_h, grid_w]]), reference, ref_config + ).numpy() + + tokens_per_frame = grid_h * grid_w // MERGE_SIZE**2 + assert got.shape == (frames * tokens_per_frame, ref_config.hidden_size) + np.testing.assert_allclose(got, expected, atol=1e-3, rtol=1e-3) + + +def test_image_and_video_fusion_and_decoder_logits(edge_package): + config, ref_config, sessions, reference = edge_package + + image = _split_image(256, 256) + image_packed, img_h, img_w = patchify_images( + image, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + image_features = sessions["vision_encoder"].run( + None, + { + "pixel_values": image_packed[0].numpy(), + "grid_thw": np.array([1, img_h, img_w], dtype=np.int64), + }, + )[0] + + frames = 2 + video = torch.stack([_split_image(64, 96)[0] for _ in range(frames)]).unsqueeze(0) + video_packed, vid_t, vid_h, vid_w = patchify_videos( + video, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + video_features = sessions["vision_encoder"].run( + None, + { + "pixel_values": video_packed[0].numpy(), + "grid_thw": np.array([vid_t, vid_h, vid_w], dtype=np.int64), + }, + )[0] + + ids = [101, 102] + image_start = len(ids) + ids += [IMAGE_TOKEN_ID] * image_features.shape[0] + video_start = len(ids) + ids += [VIDEO_TOKEN_ID] * video_features.shape[0] + ids += [201] + input_ids = np.array([ids], dtype=np.int64) + + inputs_embeds = sessions["embedding"].run( + None, + { + "input_ids": input_ids, + "image_features": image_features, + "video_features": video_features, + }, + )[0] + + expected_embeds = torch.nn.functional.embedding( + torch.from_numpy(input_ids), reference["embed_tokens.weight"] + ) + expected_embeds = expected_embeds.masked_scatter( + torch.from_numpy(input_ids == IMAGE_TOKEN_ID).unsqueeze(-1), + torch.from_numpy(image_features), + ) + expected_embeds = expected_embeds.masked_scatter( + torch.from_numpy(input_ids == VIDEO_TOKEN_ID).unsqueeze(-1), + torch.from_numpy(video_features), + ) + np.testing.assert_allclose(inputs_embeds, expected_embeds.numpy(), atol=1e-3) + + length = len(ids) + positions = np.zeros((3, 1, length), dtype=np.int64) + for index in range(image_start): + positions[:, 0, index] = index + base = image_start + merged_w = img_w // MERGE_SIZE + for token in range(image_features.shape[0]): + positions[0, 0, image_start + token] = base + positions[1, 0, image_start + token] = base + token // merged_w + positions[2, 0, image_start + token] = base + token % merged_w + base += max(img_h, img_w) // MERGE_SIZE + tokens_per_frame = video_features.shape[0] // frames + merged_vw = vid_w // MERGE_SIZE + for token in range(video_features.shape[0]): + frame, spatial = divmod(token, tokens_per_frame) + positions[0, 0, video_start + token] = base + frame + positions[1, 0, video_start + token] = base + spatial // merged_vw + positions[2, 0, video_start + token] = base + spatial % merged_vw + base += max(frames, vid_h // MERGE_SIZE, merged_vw) + positions[:, 0, length - 1] = base + + feeds = { + "inputs_embeds": inputs_embeds, + "attention_mask": np.ones((1, length), dtype=np.int64), + "position_ids": positions, + } + empty = np.zeros((1, config.num_key_value_heads, 0, config.head_dim), dtype=np.float32) + for layer in range(config.num_hidden_layers): + feeds[f"past_key_values.{layer}.key"] = empty + feeds[f"past_key_values.{layer}.value"] = empty + logits = sessions["decoder"].run(["logits"], feeds)[0] + + expected_logits = ref_text_decoder_logits( + expected_embeds, torch.from_numpy(positions), reference, ref_config + ).numpy() + np.testing.assert_allclose(logits, expected_logits, atol=5e-3, rtol=5e-3) + assert (logits.argmax(-1) == expected_logits.argmax(-1)).all() + + +def test_text_only_decoder_matches_reference(edge_package): + """Regression guard: the text path must stay correct after the vision fix.""" + config, ref_config, sessions, reference = edge_package + input_ids = np.array([[5, 77, 900, 12, 34, 56, 78, 90]], dtype=np.int64) + length = input_ids.shape[1] + empty_features = np.zeros((0, config.hidden_size), dtype=np.float32) + + inputs_embeds = sessions["embedding"].run( + None, + { + "input_ids": input_ids, + "image_features": empty_features, + "video_features": empty_features, + }, + )[0] + expected_embeds = torch.nn.functional.embedding( + torch.from_numpy(input_ids), reference["embed_tokens.weight"] + ) + np.testing.assert_allclose(inputs_embeds, expected_embeds.numpy(), atol=0) + + positions = np.tile(np.arange(length, dtype=np.int64), (3, 1, 1)) + feeds = { + "inputs_embeds": inputs_embeds, + "attention_mask": np.ones((1, length), dtype=np.int64), + "position_ids": positions, + } + empty = np.zeros((1, config.num_key_value_heads, 0, config.head_dim), dtype=np.float32) + for layer in range(config.num_hidden_layers): + feeds[f"past_key_values.{layer}.key"] = empty + feeds[f"past_key_values.{layer}.value"] = empty + logits = sessions["decoder"].run(["logits"], feeds)[0] + + expected = ref_text_decoder_logits( + expected_embeds, torch.from_numpy(positions), reference, ref_config + ).numpy() + np.testing.assert_allclose(logits, expected, atol=5e-3, rtol=5e-3) + assert (logits.argmax(-1) == expected.argmax(-1)).all() + + +def test_greedy_generation_smoke_with_image(edge_package): + """Prefill with an image, then decode a few tokens through the KV cache.""" + config, _, sessions, _ = edge_package + image = _split_image(256, 256) + packed, grid_h, grid_w = patchify_images( + image, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + image_features = sessions["vision_encoder"].run( + None, + { + "pixel_values": packed[0].numpy(), + "grid_thw": np.array([1, grid_h, grid_w], dtype=np.int64), + }, + )[0] + empty_features = np.zeros((0, config.hidden_size), dtype=np.float32) + + ids = [101, 20] + [IMAGE_TOKEN_ID] * image_features.shape[0] + [21, 102] + input_ids = np.array([ids], dtype=np.int64) + length = input_ids.shape[1] + + inputs_embeds = sessions["embedding"].run( + None, + { + "input_ids": input_ids, + "image_features": image_features, + "video_features": empty_features, + }, + )[0] + feeds = { + "inputs_embeds": inputs_embeds, + "attention_mask": np.ones((1, length), dtype=np.int64), + "position_ids": np.tile(np.arange(length, dtype=np.int64), (3, 1, 1)), + } + empty = np.zeros((1, config.num_key_value_heads, 0, config.head_dim), dtype=np.float32) + for layer in range(config.num_hidden_layers): + feeds[f"past_key_values.{layer}.key"] = empty + feeds[f"past_key_values.{layer}.value"] = empty + + names = [value.name for value in sessions["decoder"].get_outputs()] + generated: list[int] = [] + position = length + for _ in range(4): + outputs = dict(zip(names, sessions["decoder"].run(None, feeds), strict=True)) + logits = outputs["logits"] + assert np.isfinite(logits).all() + token = int(logits[0, -1].argmax()) + assert 0 <= token < config.vocab_size + generated.append(token) + step_embeds = sessions["embedding"].run( + None, + { + "input_ids": np.array([[token]], dtype=np.int64), + "image_features": empty_features, + "video_features": empty_features, + }, + )[0] + feeds = { + "inputs_embeds": step_embeds, + "attention_mask": np.ones((1, position + 1), dtype=np.int64), + "position_ids": np.full((3, 1, 1), position, dtype=np.int64), + } + for layer in range(config.num_hidden_layers): + feeds[f"past_key_values.{layer}.key"] = outputs[f"present.{layer}.key"] + feeds[f"past_key_values.{layer}.value"] = outputs[f"present.{layer}.value"] + position += 1 + + assert len(generated) == 4 diff --git a/tests/cosmos3_edge_vision_test.py b/tests/cosmos3_edge_vision_test.py new file mode 100644 index 000000000..8a91d6f81 --- /dev/null +++ b/tests/cosmos3_edge_vision_test.py @@ -0,0 +1,571 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Numerical parity tests for the Cosmos3-Edge Reasoner (tiny configs). + +These are L2 tests: they build the real ONNX graphs from a tiny architecture +config, fill them with random weights named exactly as the HuggingFace +checkpoint names them, run them under onnxruntime and compare against +:mod:`tests._cosmos3_edge_reference` — a transcription of the published +``transformers``/vLLM ``cosmos3_edge`` modeling code. + +They cover the pieces that silently corrupt image understanding when wrong: + +- the ``nn.Linear`` patch embedding over channel-last, block-major patches, +- the antialiased position-embedding resample for non-square patch grids, +- the pixel-shuffle merger ordering, +- per-frame (video) attention and token counts, +- interleaved 3D M-RoPE in the text decoder. +""" + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import onnxruntime as ort +import pytest +import torch + +from mobius._configs import ArchitectureConfig, VisionConfig +from mobius.models.cosmos import Cosmos3EdgeVLModel +from mobius.tasks import Cosmos3EdgeVLTask +from tests._cosmos3_edge_reference import ( + EdgeRefConfig, + patchify_images, + patchify_videos, + ref_text_decoder_logits, + ref_vision_features, + smart_resize, +) + +IMAGE_TOKEN_ID = 19 +VIDEO_TOKEN_ID = 18 +PATCH_SIZE = 8 +MERGE_SIZE = 2 +NUM_PATCHES = 16 # 4x4 learned reference grid +VISION_HIDDEN = 32 +TEXT_HIDDEN = 64 +HEAD_DIM = 24 # mrope_section sums to head_dim // 2 +MROPE_SECTION = [4, 4, 4] + + +def _config() -> ArchitectureConfig: + return ArchitectureConfig( + model_type="cosmos3_edge", + hidden_size=TEXT_HIDDEN, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=HEAD_DIM, + vocab_size=64, + max_position_embeddings=256, + hidden_act="relu2", + rms_norm_eps=1e-5, + rope_theta=1_000_000.0, + mrope_section=MROPE_SECTION, + mrope_interleaved=True, + image_token_id=IMAGE_TOKEN_ID, + dtype=ir.DataType.FLOAT, + vision=VisionConfig( + hidden_size=VISION_HIDDEN, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + image_size=None, + patch_size=PATCH_SIZE, + num_patches=NUM_PATCHES, + norm_eps=1e-6, + spatial_merge_size=MERGE_SIZE, + temporal_patch_size=1, + out_hidden_size=TEXT_HIDDEN, + projector_intermediate_size=48, + use_postshuffle_norm=False, + image_token_id=IMAGE_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + ), + ) + + +def _ref_config(config: ArchitectureConfig) -> EdgeRefConfig: + vision = config.vision + assert vision is not None + return EdgeRefConfig( + vision_hidden_size=vision.hidden_size, + vision_intermediate_size=vision.intermediate_size, + vision_num_layers=vision.num_hidden_layers, + vision_num_heads=vision.num_attention_heads, + patch_size=vision.patch_size, + num_channels=vision.in_channels, + num_patches=vision.num_patches, + layer_norm_eps=vision.norm_eps, + spatial_merge_size=vision.spatial_merge_size, + projector_hidden_size=vision.projector_intermediate_size, + use_postshuffle_norm=vision.use_postshuffle_norm, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + num_hidden_layers=config.num_hidden_layers, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + head_dim=config.head_dim, + rms_norm_eps=config.rms_norm_eps, + rope_theta=config.rope_theta, + mrope_section=tuple(MROPE_SECTION), + vocab_size=config.vocab_size, + ) + + +def _hf_weights(config: ArchitectureConfig) -> dict[str, torch.Tensor]: + """Random weights keyed exactly as the Cosmos3-Edge checkpoint keys them.""" + vision = config.vision + assert vision is not None + generator = torch.Generator().manual_seed(1234) + + def randn(*shape: int) -> torch.Tensor: + return torch.randn(*shape, generator=generator) * 0.05 + + patch_dim = vision.patch_size * vision.patch_size * vision.in_channels + weights: dict[str, torch.Tensor] = { + "model.visual.embeddings.patch_embedding.weight": randn(vision.hidden_size, patch_dim), + "model.visual.embeddings.patch_embedding.bias": randn(vision.hidden_size), + "model.visual.embeddings.position_embedding.weight": randn( + vision.num_patches, vision.hidden_size + ), + "model.visual.post_layernorm.weight": 1.0 + randn(vision.hidden_size), + "model.visual.post_layernorm.bias": randn(vision.hidden_size), + "model.projector.norm.weight": 1.0 + randn(vision.hidden_size), + "model.projector.norm.bias": randn(vision.hidden_size), + "model.projector.linear_fc1.weight": randn( + vision.projector_intermediate_size, vision.hidden_size * MERGE_SIZE**2 + ), + "model.projector.linear_fc1.bias": randn(vision.projector_intermediate_size), + "model.projector.linear_fc2.weight": randn( + config.hidden_size, vision.projector_intermediate_size + ), + "model.projector.linear_fc2.bias": randn(config.hidden_size), + "embed_tokens.weight": randn(config.vocab_size, config.hidden_size), + "norm.weight": 1.0 + randn(config.hidden_size), + "lm_head.weight": randn(config.vocab_size, config.hidden_size), + } + for layer in range(vision.num_hidden_layers): + prefix = f"model.visual.encoder.layers.{layer}" + for norm in ("layer_norm1", "layer_norm2"): + weights[f"{prefix}.{norm}.weight"] = 1.0 + randn(vision.hidden_size) + weights[f"{prefix}.{norm}.bias"] = randn(vision.hidden_size) + for proj in ("q_proj", "k_proj", "v_proj", "out_proj"): + weights[f"{prefix}.self_attn.{proj}.weight"] = randn( + vision.hidden_size, vision.hidden_size + ) + weights[f"{prefix}.self_attn.{proj}.bias"] = randn(vision.hidden_size) + weights[f"{prefix}.mlp.fc1.weight"] = randn( + vision.intermediate_size, vision.hidden_size + ) + weights[f"{prefix}.mlp.fc1.bias"] = randn(vision.intermediate_size) + weights[f"{prefix}.mlp.fc2.weight"] = randn( + vision.hidden_size, vision.intermediate_size + ) + weights[f"{prefix}.mlp.fc2.bias"] = randn(vision.hidden_size) + + heads = config.num_attention_heads * config.head_dim + kv_heads = config.num_key_value_heads * config.head_dim + for layer in range(config.num_hidden_layers): + prefix = f"layers.{layer}" + weights[f"{prefix}.input_layernorm.weight"] = 1.0 + randn(config.hidden_size) + weights[f"{prefix}.post_attention_layernorm.weight"] = 1.0 + randn(config.hidden_size) + weights[f"{prefix}.self_attn.to_q.weight"] = randn(heads, config.hidden_size) + weights[f"{prefix}.self_attn.to_k.weight"] = randn(kv_heads, config.hidden_size) + weights[f"{prefix}.self_attn.to_v.weight"] = randn(kv_heads, config.hidden_size) + weights[f"{prefix}.self_attn.to_out.weight"] = randn(config.hidden_size, heads) + # Generator-tower artifact that the Reasoner must drop. + weights[f"{prefix}.self_attn.k_norm_und_for_gen.weight"] = randn(config.head_dim) + weights[f"{prefix}.mlp.up_proj.weight"] = randn( + config.intermediate_size, config.hidden_size + ) + weights[f"{prefix}.mlp.down_proj.weight"] = randn( + config.hidden_size, config.intermediate_size + ) + return weights + + +@pytest.fixture(scope="module") +def edge_fixture(tmp_path_factory): + """Built + weighted ONNX sessions plus the matching reference weights.""" + config = _config() + weights = _hf_weights(config) + package = Cosmos3EdgeVLTask().build(Cosmos3EdgeVLModel(config), config) + package.apply_weights(Cosmos3EdgeVLModel(config).preprocess_weights(weights)) + package.validate_weights() + + directory = tmp_path_factory.mktemp("cosmos3_edge") + sessions = {} + for name in ("vision_encoder", "embedding", "decoder"): + path = directory / f"{name}.onnx" + ir.save(package[name], str(path)) + sessions[name] = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + + reference = { + key.removeprefix("model."): value.float() + for key, value in weights.items() + if "k_norm_und_for_gen" not in key + } + return config, _ref_config(config), sessions, reference + + +def _random_image(height: int, width: int, seed: int) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.randn(1, 3, height, width, generator=generator) + + +def test_patchify_is_block_major_with_channel_last_values(): + """Pin the processor's packed layout — the tower's weights depend on it. + + ``Cosmos3EdgeImageProcessor.patchify`` emits ``merge x merge`` blocks + contiguously (block-major) and orders the values inside a patch as + ``(patch_h, patch_w, channel)``. Reading a patch as ``(C, ph, pw)`` — the + Conv2d kernel layout — permutes colour into space and scrambles the image. + """ + patch, merge, channels = 2, 2, 3 + height = width = patch * merge * 2 # 2x2 blocks of 2x2 patches + # Encode each pixel as a unique value so ordering is fully observable. + image = torch.arange(channels * height * width, dtype=torch.float32).reshape( + 1, channels, height, width + ) + packed, grid_h, grid_w = patchify_images(image, patch_size=patch, merge_size=merge) + packed = packed[0] + + assert (grid_h, grid_w) == (height // patch, width // patch) + assert packed.shape == (grid_h * grid_w, patch * patch * channels) + + def pixel(channel: int, row: int, col: int) -> float: + return float(image[0, channel, row, col]) + + for block_row in range(grid_h // merge): + for block_col in range(grid_w // merge): + for merge_row in range(merge): + for merge_col in range(merge): + # Block-major sequence index. + index = ( + ((block_row * (grid_w // merge)) + block_col) * merge + merge_row + ) * merge + merge_col + patch_row = (block_row * merge + merge_row) * patch + patch_col = (block_col * merge + merge_col) * patch + for inner_row in range(patch): + for inner_col in range(patch): + for channel in range(channels): + # Channel-last inside the patch. + offset = (inner_row * patch + inner_col) * channels + channel + assert packed[index, offset] == pixel( + channel, patch_row + inner_row, patch_col + inner_col + ) + + +@pytest.mark.parametrize( + ("height", "width", "expected"), + [ + # Already aligned and inside the area bounds -> unchanged. + (256, 256, (256, 256)), + # Rounded to the nearest multiple of patch*merge = 32. + (750, 1000, (736, 992)), + # Below min_pixels (256*256) -> scaled up. + (64, 64, (256, 256)), + ], +) +def test_smart_resize_matches_processor_policy(height, width, expected): + """Sides are multiples of 32 and the area stays inside the processor bounds. + + Uses the real checkpoint geometry (``patch_size=16``, ``merge_size=2``), + not this module's tiny-config values, because ``smart_resize`` describes + ``Cosmos3EdgeImageProcessor``'s policy. + """ + factor = 16 * 2 + resized = smart_resize( + height, + width, + factor=factor, + min_pixels=256 * 256, + max_pixels=4096 * 4096, + ) + assert resized == expected + assert resized[0] % factor == 0 + assert resized[1] % factor == 0 + + +@pytest.mark.parametrize( + ("height", "width"), + [ + (32, 32), # grid 4x4 — identical to the learned reference grid + (64, 64), # grid 8x8 — upsampled position embeddings + (16, 96), # grid 2x12 — height downsampled (antialias path) + (80, 32), # grid 10x4 — non-square, both axes resampled + ], +) +def test_vision_encoder_matches_reference_for_image_grids(edge_fixture, height, width): + _, ref_config, sessions, reference = edge_fixture + image = _random_image(height, width, seed=height * 1000 + width) + packed, grid_h, grid_w = patchify_images( + image, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + packed = packed[0] + + got = sessions["vision_encoder"].run( + None, + { + "pixel_values": packed.numpy(), + "grid_thw": np.array([1, grid_h, grid_w], dtype=np.int64), + }, + )[0] + expected = ref_vision_features( + packed, torch.tensor([[1, grid_h, grid_w]]), reference, ref_config + ).numpy() + + assert got.shape == (grid_h * grid_w // MERGE_SIZE**2, TEXT_HIDDEN) + np.testing.assert_allclose(got, expected, atol=1e-4, rtol=1e-4) + + +def test_vision_encoder_matches_reference_for_video(edge_fixture): + _, ref_config, sessions, reference = edge_fixture + frames = 4 + video = torch.stack( + [_random_image(32, 64, seed=7 + index)[0] for index in range(frames)] + ).unsqueeze(0) + packed, grid_t, grid_h, grid_w = patchify_videos( + video, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + packed = packed[0] + + got = sessions["vision_encoder"].run( + None, + { + "pixel_values": packed.numpy(), + "grid_thw": np.array([grid_t, grid_h, grid_w], dtype=np.int64), + }, + )[0] + expected = ref_vision_features( + packed, torch.tensor([[grid_t, grid_h, grid_w]]), reference, ref_config + ).numpy() + + # One token per merged 2x2 block, per frame. + assert got.shape == (grid_t * grid_h * grid_w // MERGE_SIZE**2, TEXT_HIDDEN) + assert got.shape[0] == frames * (grid_h * grid_w // MERGE_SIZE**2) + np.testing.assert_allclose(got, expected, atol=1e-4, rtol=1e-4) + + +def test_video_frames_are_encoded_independently(edge_fixture): + """Attention must not leak across frames (per-frame ``cu_seqlens``).""" + _, _, sessions, _ = edge_fixture + frame_a = _random_image(16, 16, seed=3) + frame_b = _random_image(16, 16, seed=4) + packed_a, grid_h, grid_w = patchify_images( + frame_a, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + packed_b, _, _ = patchify_images(frame_b, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE) + + single = sessions["vision_encoder"].run( + None, + { + "pixel_values": packed_a[0].numpy(), + "grid_thw": np.array([1, grid_h, grid_w], dtype=np.int64), + }, + )[0] + both = sessions["vision_encoder"].run( + None, + { + "pixel_values": torch.cat([packed_a[0], packed_b[0]]).numpy(), + "grid_thw": np.array([2, grid_h, grid_w], dtype=np.int64), + }, + )[0] + + np.testing.assert_allclose(both[: single.shape[0]], single, atol=1e-5, rtol=1e-5) + + +def test_embedding_and_decoder_match_reference_with_image_and_video(edge_fixture): + config, ref_config, sessions, reference = edge_fixture + + image = _random_image(32, 32, seed=11) + image_packed, img_h, img_w = patchify_images( + image, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + image_features = sessions["vision_encoder"].run( + None, + { + "pixel_values": image_packed[0].numpy(), + "grid_thw": np.array([1, img_h, img_w], dtype=np.int64), + }, + )[0] + + frames = 2 + video = torch.stack( + [_random_image(16, 16, seed=21 + index)[0] for index in range(frames)] + ).unsqueeze(0) + video_packed, vid_t, vid_h, vid_w = patchify_videos( + video, patch_size=PATCH_SIZE, merge_size=MERGE_SIZE + ) + video_features = sessions["vision_encoder"].run( + None, + { + "pixel_values": video_packed[0].numpy(), + "grid_thw": np.array([vid_t, vid_h, vid_w], dtype=np.int64), + }, + )[0] + + ids = [5, 6] + image_start = len(ids) + ids += [IMAGE_TOKEN_ID] * image_features.shape[0] + video_start = len(ids) + ids += [VIDEO_TOKEN_ID] * video_features.shape[0] + ids += [7] + input_ids = np.array([ids], dtype=np.int64) + + inputs_embeds = sessions["embedding"].run( + None, + { + "input_ids": input_ids, + "image_features": image_features, + "video_features": video_features, + }, + )[0] + + expected_embeds = torch.nn.functional.embedding( + torch.from_numpy(input_ids), reference["embed_tokens.weight"] + ) + expected_embeds = expected_embeds.masked_scatter( + torch.from_numpy(input_ids == IMAGE_TOKEN_ID).unsqueeze(-1), + torch.from_numpy(image_features), + ) + expected_embeds = expected_embeds.masked_scatter( + torch.from_numpy(input_ids == VIDEO_TOKEN_ID).unsqueeze(-1), + torch.from_numpy(video_features), + ) + np.testing.assert_allclose(inputs_embeds, expected_embeds.numpy(), atol=1e-5) + + # 3D M-RoPE position ids: visual spans get (t, h, w) axes, text is diagonal. + length = len(ids) + positions = np.zeros((3, 1, length), dtype=np.int64) + for index in range(image_start): + positions[:, 0, index] = index + base = image_start + merged_w = img_w // MERGE_SIZE + for token in range(image_features.shape[0]): + positions[0, 0, image_start + token] = base + positions[1, 0, image_start + token] = base + token // merged_w + positions[2, 0, image_start + token] = base + token % merged_w + base = base + max(img_h, img_w) // MERGE_SIZE + tokens_per_frame = video_features.shape[0] // frames + merged_vw = vid_w // MERGE_SIZE + for token in range(video_features.shape[0]): + frame, spatial = divmod(token, tokens_per_frame) + positions[0, 0, video_start + token] = base + frame + positions[1, 0, video_start + token] = base + spatial // merged_vw + positions[2, 0, video_start + token] = base + spatial % merged_vw + base += max(frames, vid_h // MERGE_SIZE, merged_vw) + positions[:, 0, length - 1] = base + + feeds = { + "inputs_embeds": inputs_embeds, + "attention_mask": np.ones((1, length), dtype=np.int64), + "position_ids": positions, + } + empty = np.zeros((1, config.num_key_value_heads, 0, config.head_dim), dtype=np.float32) + for layer in range(config.num_hidden_layers): + feeds[f"past_key_values.{layer}.key"] = empty + feeds[f"past_key_values.{layer}.value"] = empty + logits = sessions["decoder"].run(["logits"], feeds)[0] + + expected_logits = ref_text_decoder_logits( + expected_embeds, torch.from_numpy(positions), reference, ref_config + ).numpy() + np.testing.assert_allclose(logits, expected_logits, atol=2e-4, rtol=2e-4) + + +def test_text_only_decoder_matches_reference(edge_fixture): + """Text-only inference must stay bit-comparable to the reference.""" + config, ref_config, sessions, reference = edge_fixture + input_ids = np.array([[3, 9, 12, 40, 2, 61]], dtype=np.int64) + length = input_ids.shape[1] + + inputs_embeds = sessions["embedding"].run( + None, + { + "input_ids": input_ids, + "image_features": np.zeros((0, TEXT_HIDDEN), dtype=np.float32), + "video_features": np.zeros((0, TEXT_HIDDEN), dtype=np.float32), + }, + )[0] + expected_embeds = torch.nn.functional.embedding( + torch.from_numpy(input_ids), reference["embed_tokens.weight"] + ) + np.testing.assert_allclose(inputs_embeds, expected_embeds.numpy(), atol=0) + + positions = np.tile(np.arange(length, dtype=np.int64), (3, 1, 1)) + feeds = { + "inputs_embeds": inputs_embeds, + "attention_mask": np.ones((1, length), dtype=np.int64), + "position_ids": positions, + } + empty = np.zeros((1, config.num_key_value_heads, 0, config.head_dim), dtype=np.float32) + for layer in range(config.num_hidden_layers): + feeds[f"past_key_values.{layer}.key"] = empty + feeds[f"past_key_values.{layer}.value"] = empty + logits = sessions["decoder"].run(["logits"], feeds)[0] + + expected = ref_text_decoder_logits( + expected_embeds, torch.from_numpy(positions), reference, ref_config + ).numpy() + np.testing.assert_allclose(logits, expected, atol=2e-4, rtol=2e-4) + + +def test_decoder_kv_cache_decode_step_matches_full_prefill(edge_fixture): + """ORT generation smoke: one cached decode step equals a full re-prefill.""" + config, _, sessions, reference = edge_fixture + input_ids = np.array([[3, 9, 12, 40]], dtype=np.int64) + empty_features = np.zeros((0, TEXT_HIDDEN), dtype=np.float32) + + def embed(ids: np.ndarray) -> np.ndarray: + return sessions["embedding"].run( + None, + { + "input_ids": ids, + "image_features": empty_features, + "video_features": empty_features, + }, + )[0] + + length = input_ids.shape[1] + feeds = { + "inputs_embeds": embed(input_ids), + "attention_mask": np.ones((1, length), dtype=np.int64), + "position_ids": np.tile(np.arange(length, dtype=np.int64), (3, 1, 1)), + } + empty = np.zeros((1, config.num_key_value_heads, 0, config.head_dim), dtype=np.float32) + for layer in range(config.num_hidden_layers): + feeds[f"past_key_values.{layer}.key"] = empty + feeds[f"past_key_values.{layer}.value"] = empty + outputs = sessions["decoder"].run(None, feeds) + names = [value.name for value in sessions["decoder"].get_outputs()] + present = dict(zip(names, outputs, strict=True)) + next_token = int(present["logits"][0, -1].argmax()) + + step_ids = np.array([[next_token]], dtype=np.int64) + feeds = { + "inputs_embeds": embed(step_ids), + "attention_mask": np.ones((1, length + 1), dtype=np.int64), + "position_ids": np.full((3, 1, 1), length, dtype=np.int64), + } + for layer in range(config.num_hidden_layers): + feeds[f"past_key_values.{layer}.key"] = present[f"present.{layer}.key"] + feeds[f"past_key_values.{layer}.value"] = present[f"present.{layer}.value"] + cached_logits = sessions["decoder"].run(["logits"], feeds)[0] + + full_ids = np.concatenate([input_ids, step_ids], axis=1) + feeds = { + "inputs_embeds": embed(full_ids), + "attention_mask": np.ones((1, length + 1), dtype=np.int64), + "position_ids": np.tile(np.arange(length + 1, dtype=np.int64), (3, 1, 1)), + } + for layer in range(config.num_hidden_layers): + feeds[f"past_key_values.{layer}.key"] = empty + feeds[f"past_key_values.{layer}.value"] = empty + full_logits = sessions["decoder"].run(["logits"], feeds)[0] + + np.testing.assert_allclose(cached_logits[0, 0], full_logits[0, -1], atol=2e-4, rtol=2e-4) + assert reference # reference weights were used to build the graphs