Add composable world-model pipelines and Cosmos3 export - #466
Open
xiaoyu-work wants to merge 11 commits into
Open
Add composable world-model pipelines and Cosmos3 export#466xiaoyu-work wants to merge 11 commits into
xiaoyu-work wants to merge 11 commits into
Conversation
Introduce executable pipeline schema 1.1 with typed generated-input programs, recurrent state lifecycle, scheduler and autoregressive controls, parameterized transforms, runtime assets, and dtype/EP hints. Implement complete Cosmos3 Omni and Edge exports with the shared MoT generator, Wan video VAE, optional AVAE audio tokenizer, domain-aware action heads, real checkpoint weight routing, CLI support, documentation, and numerical tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 997ab62c-69cb-4802-a5c6-fddd51bfd583 Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
|
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Focus the world-model guides on export commands, package contents, supported checkpoints, runtime usage, and practical options. Remove internal architecture and manifest implementation details from the user guide. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 997ab62c-69cb-4802-a5c6-fddd51bfd583 Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Add model-agnostic build, generation, and pipeline configuration for world-model exporters, plus safe shared helpers for diffusers-style checkpoint inspection and asset resolution. Keep LLM ArchitectureConfig and every model/component implementation unchanged; Cosmos3 Omni and Edge only consume the new outer pipeline configuration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 997ab62c-69cb-4802-a5c6-fddd51bfd583 Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Match upstream chunk-zero temporal resampling so Cosmos3 can decode a single latent frame without running Conv on an empty temporal extent. Preserve bit-identical multi-frame behavior and add ORT/PyTorch parity coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 997ab62c-69cb-4802-a5c6-fddd51bfd583 Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Replace the fixed SigLIP1-style vision graph with the published packed variable-resolution SigLIP2 contract, correct block-major projection ordering, and use interleaved multimodal RoPE. Add image/video feature routing, processor metadata, real-checkpoint numerical parity tests, and image/video understanding coverage while preserving text-only behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 997ab62c-69cb-4802-a5c6-fddd51bfd583 Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Describe conditioned video encoding, classifier-free guidance, prompt packing, and official Cosmos3 Edge generation recipes in pipeline.json. Ship the recommended negative prompt when available and preserve checkpoint-specific action behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 997ab62c-69cb-4802-a5c6-fddd51bfd583 Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Apply the repository Ruff formatter to resolve CI formatting findings. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 997ab62c-69cb-4802-a5c6-fddd51bfd583 Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> Signed-off-by: Xiaoyu <85524621+xiaoyu-work@users.noreply.github.com>
Fix the Ruff unused-import failure after the latest world-model cleanup. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 997ab62c-69cb-4802-a5c6-fddd51bfd583 Signed-off-by: Xiaoyu Zhang <xiaoyuzhang@microsoft.com>
Member
Author
|
@copilot resolve the merge conflicts in this pull request |
Resolve conflicts between the world-model/Cosmos3 export work and main: - --features: keep world-model alongside main's prune-lm-head -> prune-prefill-prefix rename (CLI, README, CHANGELOG, cli_reference). - _diffusers_builder: keep both the component allowlist and pipeline component task overrides from main and the EP/trace-optimization plumbing from this branch, plus the corresponding tests. - components/tasks __init__: export main's new components/tasks together with the Cosmos3 Edge vision and Cosmos3 audio/generator additions; drop the projector removed by this branch. - _weight_loading: thread main's revision pin through the new iter_weight_shards generator. - _sub_configs: drop duplicated Cosmos3-Edge vision token fields that main added independently. - wan_vae_test: expect the fused Swish op introduced by main. Signed-off-by: GitHub <noreply@github.com> Co-authored-by: xiaoyu-work <85524621+xiaoyu-work@users.noreply.github.com>
Contributor
Merged
Also fixed two silent (non-conflicting) merge breakages:
|
Member
Author
|
@copilot resolve the merge conflicts in this pull request |
…orld-model Signed-off-by: GitHub <noreply@github.com> # Conflicts: # src/mobius/_model_package_test.py Co-authored-by: xiaoyu-work <85524621+xiaoyu-work@users.noreply.github.com>
Contributor
Merged Verified after the merge: |
xiaoyu-work
marked this pull request as ready for review
August 17, 2026 23:57
Contributor
There was a problem hiding this comment.
Pull request overview
This PR introduces a compositional “world-model” export path in Mobius (multi-ONNX-component packages plus a pipeline manifest) and implements full Cosmos3 exports (Cosmos3-Omni and Cosmos3-Edge), including new tasks/components for the shared MoT generator, Wan video VAE, and optional Cosmos3 audio tokenizer. It also upgrades Cosmos3-Edge’s VLM export to match the published packed SigLIP2 + interleaved 3D M-RoPE contracts, and adds extensive tests + documentation for the new pipeline workflow.
Changes:
- Add a world-model pipeline builder/registry (
build_world_model,--features world-model) plus frozen configs for build + generation defaults, and diffusers-checkpoint inspection utilities. - Implement/compose Cosmos3 components (Cosmos3-Omni generator task, Wan VAE task, Cosmos3 audio tasks) and update Cosmos3-Edge vision-language export to the packed, variable-resolution SigLIP2 contract with dual (image/video) feature fusion.
- Improve weight-loading for composite exports (streaming shard iteration + partial weight application + deferred constant folding) and add numerical/integration tests + docs for world-model packages.
Reviewed changes
Copilot reviewed 60 out of 61 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/cosmos3_edge_integration_test.py | New integration test validating Cosmos3-Edge Reasoner (vision→fusion→decoder) numerics against a PyTorch reference with real weights. |
| tests/cli_test.py | Add CLI tests for --features world-model dispatch and option incompatibility checks. |
| tests/_test_configs.py | Update tiny Cosmos3-Edge-related vision settings (patch size, tokens, interleaved M-RoPE flag) used by test configs. |
| testdata/cases/vision-language/cosmos3-omni-nano.yaml | Update case notes to distinguish VLM Reasoner validation from full world-model export flow. |
| src/mobius/tasks/_world_model.py | Rename the single-step “world model” task to LatentDynamicsTask with backward-compatible alias. |
| src/mobius/tasks/_wan_vae.py | New task exporting Wan VAE encoder/decoder as 5D video graphs with explicit latent normalization boundary. |
| src/mobius/tasks/_vision_language_3model.py | Update Cosmos3EdgeVLTask to packed vision input (pixel_values, grid_thw) and dual-stream embedding (image_features, video_features). |
| src/mobius/tasks/_cosmos3_omni_generator.py | New task building the packed denoising-step graph for the Cosmos3-Omni unified MoT transformer with gated optional heads. |
| src/mobius/tasks/_cosmos3_audio.py | New tasks for Cosmos3 AVAE audio tokenizer: decoder-only vs full encoder+decoder builds, plus selector. |
| src/mobius/tasks/init.py | Export/register new tasks (Cosmos3 audio, Cosmos3 Omni generator, Wan VAE, latent-dynamics) and keep world-model alias. |
| src/mobius/models/world_model.py | Rename reference module to MLPLatentDynamicsModel, keep MLPWorldModel alias. |
| src/mobius/models/world_model_test.py | Update tests to cover new task/config/model names and backward-compatible aliases. |
| src/mobius/models/cosmos.py | Major Cosmos3-Edge update: packed SigLIP2 vision tower + patch merger, interleaved M-RoPE, dual image/video feature scatter, and unified-checkpoint key dropping. |
| src/mobius/models/cosmos_test.py | Expand unit tests for Edge weight routing, packed vision I/O, token fusion behavior, and config hook extraction. |
| src/mobius/models/init.py | Export new Cosmos3/Wan/latent-dynamics model symbols. |
| src/mobius/components/_multimodal.py | Remove the old fixed-grid Cosmos3-Edge projector component (superseded by new packed SigLIP2 path). |
| src/mobius/components/_cosmos3_edge_vision.py | New packed, variable-resolution SigLIP2 embeddings/tower + pixel-shuffle merger projector with exact resampling. |
| src/mobius/components/init.py | Export new Cosmos3-Edge vision components. |
| src/mobius/_world_model_config.py | New frozen build/generation/pipeline config dataclasses for composed world-model exporters. |
| src/mobius/_world_model_config_test.py | Tests for world-model config immutability, dtype resolution, defaults, and manifest blocks. |
| src/mobius/_world_model_builder.py | New builder registry and build_world_model() entry point (loads model_type from config/model_index). |
| src/mobius/_world_model_builder_test.py | Tests for registry semantics and model_type dispatch behavior (local + pure diffusers pipeline). |
| src/mobius/_weight_loading.py | Add iter_weight_shards() for shard-by-shard streaming; _download_weights() now merges shards then dequantizes FP8. |
| src/mobius/_weight_loading_test.py | Tests for shard iteration, local shard ordering, and FP8 scale application after merge. |
| src/mobius/_model_package.py | Add apply_weights_partial(), finalize_weights(), and validate_weights() to support streaming + deferred folding. |
| src/mobius/_model_package_test.py | Tests for deferred folding semantics and missing-initializer validation. |
| src/mobius/_diffusers_checkpoint.py | New safe, read-only diffusers checkpoint inspection (component class, shard paths, asset resolution, metadata-only weight names). |
| src/mobius/_diffusers_checkpoint_test.py | Tests for traversal protection, metadata reading, shard ordering, and asset resolution behavior. |
| src/mobius/_diffusers_builder.py | Extend diffusers builder for Cosmos3 Omni transformer + Wan VAE, and forward trace_optimization into component builds. |
| src/mobius/_diffusers_builder_test.py | Update/extend tests for new diffusers-supported components and forwarded build kwargs. |
| src/mobius/_cosmos3_edge_world_model.py | New full Cosmos3-Edge world-model exporter composing Reasoner + shared Generator + Wan VAE (+ policy metadata + vision contract metadata). |
| src/mobius/_cosmos3_edge_world_model_test.py | Tests for Edge world-model dispatch, metadata shaping, and checkpoint validation rules. |
| src/mobius/_configs/per_model/_cosmos3_edge_vision.py | Update Cosmos3-Edge vision hook: num_patches semantics, projector flags, token ids, temporal_patch_size=1, and interleaved M-RoPE flag. |
| src/mobius/_configs/_world_model.py | Rename WorldModelConfig → LatentDynamicsConfig with backward-compatible alias. |
| src/mobius/_configs/_wan_vae.py | New Wan VAE config parser/validator for diffusers AutoencoderKLWan. |
| src/mobius/_configs/_sub_configs.py | Extend VisionConfig with Edge-specific fields (use_postshuffle_norm, num_patches, mrope_interleaved). |
| src/mobius/_configs/_extractors.py | Include mrope_interleaved in shared vision-config extraction. |
| src/mobius/_configs/init.py | Export new config types (Cosmos3*, WanVAEConfig, LatentDynamicsConfig alias). |
| src/mobius/main.py | Add --features world-model path with incompatibility checks and dispatch to build_world_model(). |
| src/mobius/init.py | Export new world-model APIs/types and Cosmos3/Wan components from the public API surface. |
| README.md | Document world-model as a build feature and provide example commands. |
| docs/world-models.md | New documentation page describing world-model export output layout and how to run packages. |
| docs/index.md | Add world-model docs page to the documentation TOC. |
| docs/cli_reference.md | Document --features world-model and usage examples. |
| docs/api/index.md | Add build_world_model() to API docs index. |
| docs/api/build_world_model.md | New API doc page for build_world_model() options and supported model_types. |
| CHANGELOG.md | Changelog entries for pipeline schema, world-model exporters, Cosmos3-Edge fixes, and renamed latent-dynamics API. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+423
to
+424
| scale_factor_temporal=int(raw.get("scale_factor_temporal") or 4), | ||
| scale_factor_spatial=int(raw.get("scale_factor_spatial") or 8), |
Comment on lines
+285
to
+286
| "temperal_downsample must have len(dim_mult) - 1 = " | ||
| f"{len(self.dim_mult) - 1} entries, got {len(self.temporal_downsample)}" |
rui-ren
self-requested a review
August 18, 2026 00:20
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Batch feature routing, component dtype overrides, VAE validation, and pipeline persistence currently contain correctness or safety defects.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 60/61 changed files
- Comments generated: 6
- Review effort level: Balanced
Comment on lines
+2883
to
+2889
| 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: |
Comment on lines
+374
to
+375
| 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)) |
Comment on lines
+303
to
+304
| if any(std == 0 for std in self.latents_std): | ||
| raise ValueError("latents_std entries must be non-zero (used as a divisor)") |
Comment on lines
+268
to
+269
| vae_config = WanVAEConfig.from_diffusers(vae_config_dict) | ||
| vae_module = AutoencoderKLWanModel(vae_config) |
Comment on lines
+291
to
+295
| audio_config = Cosmos3AudioConfig.from_diffusers( | ||
| audio_config_dict, | ||
| weight_names=audio_weight_names, | ||
| ) | ||
| audio_module = create_cosmos3_avae_audio_tokenizer(audio_config) |
| with os.fdopen(handle, "w", encoding="utf-8") as file: | ||
| json.dump(document, file, indent=2) | ||
| file.write("\n") | ||
| os.replace(staged, marker) |
Contributor
Does this runtime support |
Member
|
Iterating on the design. Right now it's not updated yet. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Introduce executable pipeline schema 1.1 with typed generated-input programs, recurrent state lifecycle, scheduler and autoregressive controls, parameterized transforms, runtime assets, and dtype/EP hints.
Implement complete Cosmos3 Omni and Edge exports with the shared MoT generator, Wan video VAE, optional AVAE audio tokenizer, domain-aware action heads, real checkpoint weight routing, CLI support, documentation, and numerical tests.