Initial release

Co-authored-by: fsergiy <fsergiy@users.noreply.huggingface.co>
Co-authored-by: yaoxuhf <yaoxuhf@users.noreply.huggingface.co>
Co-authored-by: rickznv <rickznv@users.noreply.huggingface.co>
Co-authored-by: atharvajoshi10 <atharvajoshi10@users.noreply.huggingface.co>
Co-authored-by: ychao-nvidia <ychao-nvidia@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,38 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+text_tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

BIAS.md

@@ -0,0 +1,11 @@
+## Bias
+
+| Field | Response |
+| :---- | :---- |
+| Participation considerations from adversely impacted groups [protected classes](https://www.senate.ca.gov/content/protected-classes) in model design and testing | None. |
+| Measures taken to mitigate against unwanted bias | Training, evaluation, and testing data are curated before release to filter restricted content, including content relating to protected classes. Model behavior is evaluated across Physical AI domains — robotics, autonomous vehicles, human-centric scenes, common scenes, industry, miscellaneous, and physics-oriented benchmarks — with attention to coverage across diverse demographic and contextual characteristics that affect protected-class outcomes. |
+| Which characteristic (feature) show(s) the greatest difference in performance?: | Greatest performance differences are observed in tasks requiring long-horizon temporal consistency, fine-grained physical interactions, and embodiment-specific action generation. Performance is generally stronger on common visual reasoning and world-generation tasks than on complex multi-agent, robotics-control, or tightly synchronized multimodal generation scenarios. |
+| Which feature(s) have the worst performance overall? | Performance is generally weakest in tasks requiring long-horizon temporal consistency, precise physical interactions, embodiment-specific action control, and strict audio-visual synchronization. |
+| If using internal data, description of methods implemented in data acquisition or processing, if any, to address the prevalence of identifiable biases in the training, testing, and validation data: | Bias-specific methods applied during data processing include person-presence screening, demographic-taxonomy classification (age, gender, ethnicity), embedding-based diversity analysis, and dataset balancing across sources. Internal analysis surfaced: non-person scenes are more prevalent than person-centric content; demographic-taxonomy outputs on person-present samples are most frequently "uncertain" across age, gender, and ethnicity dimensions; and source-type variation, with people-centric image and video datasets showing higher demographic signal than document-, object-, robotics-, or scene-focused datasets. *(Quantitative details in the row below.)* Downstream deployments should add bias audits, fairness evaluation, red-teaming, demographically balanced fine-tuning, or counterfactual augmentation as mitigations. |
+| Tools used to assess statistical imbalances and highlight patterns that may introduce bias into AI models: | Dataset analytics pipelines, metadata distribution analysis, heuristic quality checks, embedding-based clustering, model-assisted filtering systems, and benchmark evaluation suites are used to assess statistical imbalances and identify patterns that may introduce bias into model behavior. |
+| Tools used to assess statistical imbalances and highlight patterns that may introduce bias into AI models: | These datasets, such as OpenImages-derived detection-to-NLP datasets, visual grounding and VQA datasets, document/image understanding datasets, video/action understanding datasets, and NVIDIA-created or curated visual datasets, do not collectively or exhaustively represent all demographic groups (and proportionally therein). For instance, automated person-presence screening did not identify a person in approximately 58% of visual samples analyzed across approximately 400 datasets, while person-present signals were identified in approximately 42% of analyzed samples. In the subset where person-present signals were identified, these datasets contain uneven representation splits across the measured visual taxonomies: age outputs were most frequently uncertain, followed by child and adult; gender outputs were most frequently uncertain, followed by male and female; and ethnicity outputs were most frequently uncertain, followed by Hispanic and White as the most frequent identified categories. Dataset-level results vary by source type, with people-centric image and video datasets containing higher person-present and demographic-taxonomy signals than document-, object-, robotics-, or scene-focused datasets. To mitigate these imbalances, we recommend considering evaluation techniques such as bias audits, task-specific fairness evaluation, and red-teaming, along with fine-tuning with demographically balanced datasets and counterfactual data augmentation to align with the desired model behavior. This evaluation used a baseline of 200 samples across all datasets, with larger subsets of up to 3,000 samples utilized for certain in-depth analyses, identified as optimal thresholds for maximizing embedder accuracy. |

EXPLAINABILITY.md

@@ -0,0 +1,16 @@
+## Explainability
+
+| Field | Response |
+| :---- | :---- |
+| Intended Application & Domain | World reasoning and generation for Physical AI. |
+| Model Type | Mixture-of-Transformers architecture with two towers. One is an autoregressive model for Physical AI reasoning; the other is a diffusion model for Physical AI generation. |
+| Intended Users | Physical AI developers, researchers, and practitioners building or evaluating autonomous vehicle, robotics, and world-generation workflows. |
+| Output | Images, videos, audio, and action commands. |
+| Tools used to evaluate datasets to identify synthetic data and ensure data authenticity. | Dataset provenance analysis, metadata validation, watermark and artifact detection, embedding-based clustering, heuristic quality checks, and model-assisted data validation pipelines are used to identify synthetic content patterns, assess dataset authenticity, and improve data quality during dataset curation. |
+| Describe how the model works | Cosmos3 is an Omni world foundation model that generates texts, images, videos, audio, and action commands from combinations of text, images, videos, and action trajectory inputs. Input tokens from multiple modalities are packed into a shared sequence and processed by our mixture-of-transformer backbone with modality-specific output heads. |
+| Name the adversely impacted groups this has been tested to deliver comparable outcomes regardless of: | None. |
+| Technical Limitations | The model may not follow text, image, video, audio, or action trajectory inputs accurately in challenging cases, especially where the input contains complex scene composition, unusual camera motion, multiple interacting agents, low lighting, high motion blur, or fine-grained physical interactions. Generated outputs may contain temporal inconsistency, object morphing, inaccurate 3D structure, or implausible physical dynamics. Generated audio may not accurately render intelligible speech, or maintain strict temporal and semantic alignment with the visual context. |
+| Verified to have met prescribed NVIDIA quality standards | Yes. |
+| Performance Metrics | Video generation is measured using PAIBench-G, RBench, PhysicsIQ, and Artifical Analysis Image2Video benchmark. Image generation uses UniGenBench and Artifical Analysis Text2Image benchmark. For transfer evaluation, we use PAIBench-C and AVBench-C. Audio generation uses internal benchmarks. Action prediction uses metrics such as action MSE, Absolute Translation Error, Relative Translation Error, Relative Rotation Error, PSNR, and robotic task completion success rate. |
+| Potential Known Risks | This model can generate synthetic media and may produce content that is offensive, unsafe, misleading, indecent, or unsuitable for a target deployment. Users should implement robust safety guardrails — including content filtering, abuse monitoring, and access controls — to reduce the risk of harmful outputs. Users are responsible for ensuring that their use of the model complies with all applicable laws and regulations, and for regularly reviewing and updating their guardrails as risks evolve. |
+| Licensing | [OpenMDW1.1](https://openmdw.ai/)  |

PRIVACY.md

@@ -0,0 +1,6 @@
+## Privacy
+| Privacy Information |
+|---|
+| The model was trained on large-scale publicly available data that may contain images, audio-video, and text relating to people. NVIDIA collected and used this data in compliance with applicable data protection and privacy laws. This model was not designed to derive insights or otherwise learn from any personal data contained in the datasets. |
+| NVIDIA uses a combination of filters, data minimization techniques, and other guardrails to help prevent personal data from being recited by our models. We employ automated tools and data processing techniques during pre-training or training to identify and filter certain categories of personal data. For example, for text-bearing source and document components, our automated tools identified potential personal data such as person names, locations, and possible business or public-facing contact information such as email addresses and phone numbers.  We reviewed and removed any verified instances of personal data through a combination of automated filtering and human-in-the-loop validation. |
+| Please review NVIDIA's [Privacy Policy](https://www.nvidia.com/en-us/about-nvidia/privacy-policy/) for more information. |

README.md

@@ -0,0 +1,475 @@
+---
+license: other
+license_name: openmdw1.1-license
+license_link: >-
+  https://openmdw.ai/license/1-1/
+library_name: cosmos
+tags:
+  - nvidia
+  - cosmos
+  - cosmos3
+  - world action model
+  - policy model
+countDownloads:
+  - checkpoint.json
+  - config.json
+  - generation_config.json
+  - model.safetensors.index.json
+  - model_index.json
+  - tokenizer.json
+  - tokenizer_config.json
+  - sound_tokenizer/config.json
+  - sound_tokenizer/diffusion_pytorch_model.safetensors
+  - text_tokenizer/tokenizer.json
+  - text_tokenizer/tokenizer_config.json
+  - transformer/config.json
+  - transformer/diffusion_pytorch_model-00001-of-00007.safetensors
+  - transformer/diffusion_pytorch_model-00002-of-00007.safetensors
+  - transformer/diffusion_pytorch_model-00003-of-00007.safetensors
+  - transformer/diffusion_pytorch_model-00004-of-00007.safetensors
+  - transformer/diffusion_pytorch_model-00005-of-00007.safetensors
+  - transformer/diffusion_pytorch_model-00006-of-00007.safetensors
+  - transformer/diffusion_pytorch_model-00007-of-00007.safetensors
+  - transformer/diffusion_pytorch_model.safetensors.index.json
+  - vae/config.json
+  - vae/diffusion_pytorch_model.safetensors
+  - vision_encoder/config.json
+  - vision_encoder/model.safetensors  
+---
+
+# **Cosmos 3: Omnimodal World Models for Physical AI**
+**[Model Collection](https://huggingface.co/collections/nvidia/cosmos3)** | **[Code](https://github.com/nvidia/cosmos)** | **[White Paper](https://research.nvidia.com/labs/cosmos-lab/cosmos3/technical-report.pdf)** | **[Website](https://research.nvidia.com/labs/cosmos-lab/cosmos3/)**
+
+[NVIDIA Cosmos™](https://github.com/nvidia/cosmos) is a world foundation model platform designed to accelerate the development of Physical AI by enabling machines to understand, simulate, and interact with the physical world across robotics, autonomous driving, and smart space environments, including industrial and factory-scale applications.
+
+# Model Overview: Cosmos3-Edge-Policy-DROID
+
+## Description
+
+Cosmos3 is a collection of Omnimodal world models capable of generating dynamic, high-quality video, image, audio, and action commands from combinations of text, image, video, and action trajectory inputs. It serves as a foundational building block for a broad range of Physical AI applications and research spanning world understanding, world generation, simulation, and embodied policy learning.
+
+This model is ready for commercial and non-commercial use.
+
+**Model Developer:** NVIDIA
+
+### Model Versions
+
+**Released on: 07/20/2026**
+- Cosmos3-Edge:
+  - Given multimodal inputs including text, images, video, and action trajectories, generate coherent text, images, video, and action outputs for multimodal understanding, world simulation, future prediction, action reasoning, and Physical AI applications.
+
+- Cosmos3-Edge-Policy-DROID:
+  - Given language instructions and visual observations from the DROID robot platform, generate robot action trajectories for manipulation and control tasks.
+
+- Cosmos3-Super-Image2Video-4Step:
+  - Given one or more input images and optional text instructions, generate temporally coherent video sequences that are consistent with the provided visual content. Distilled from Cosmos3-Super-Image2Video using [Improved Distribution Matching Distillation (DMD2)](https://arxiv.org/abs/2405.14867), enabling high-quality generation in 4 steps.
+
+- Cosmos3-Super-Text2Image-4Step:
+  - Given text input, generate high-fidelity images that are consistent with the provided description. Distilled from Cosmos3-Super-Text2Image using [Improved Distribution Matching Distillation (DMD2)](https://arxiv.org/abs/2405.14867), enabling high-quality generation in 4 steps.
+
+**Released on: 05/31/2026**
+- Cosmos3-Nano:
+  - Given multimodal inputs including text, images, video, audio, and action trajectories, generate coherent text, images, video, audio, and action outputs for multimodal understanding, world simulation, future prediction, action reasoning, and Physical AI applications.
+
+- Cosmos3-Super:
+  - Given multimodal inputs including text, images, video, audio, and action trajectories, generate coherent text, images, video, audio, and action outputs for multimodal understanding, world simulation, future prediction, action reasoning, and Physical AI applications.
+
+- Cosmos3-Nano-Policy-DROID:
+  - Given language instructions and visual observations from the DROID robot platform, generate robot action trajectories for manipulation and control tasks.
+
+- Cosmos3-Super-Image2Video:
+  - Given one or more input images and optional text instructions, generate temporally coherent video sequences that are consistent with the provided visual content.
+
+- Cosmos3-Super-Text2Image:
+  - Given text input, generate high-fidelity images that are consistent with the provided description.
+
+### License
+
+This model is released under the [OpenMDW1.1](https://openmdw.ai/license/1-1/)
+
+### Deployment Geography
+
+Global
+
+### Use Case
+
+Physical AI: Encompassing robotics, autonomous vehicles (AV), and smart space environments, including industrial and factory-scale applications.
+
+### Release Date
+
+Hugging Face 07/20/2026 via [https://huggingface.co/collections/nvidia/cosmos3](https://huggingface.co/collections/nvidia/cosmos3)
+GitHub 07/20/2026 via [https://github.com/nvidia/cosmos](https://github.com/nvidia/cosmos)
+
+## Model Architecture
+
+**Architecture Type:** Transformer
+
+**Network Architecture:** Mixture-of-Transformers (MoT)
+
+Cosmos3 is an Omni-modal foundation model built on a Mixture-of-Transformers (MoT) architecture consisting of two complementary transformer towers: an autoregressive transformer for discrete token generation and a diffusion transformer for continuous multimodal generation. During inference, text is generated through standard next-token autoregressive decoding, while non-text modalities, such as images, video, audio, and actions, are synthesized through iterative denoising. This unified architecture enables Cosmos3 to model heterogeneous modalities within a single framework while preserving generation mechanisms best suited to each modality.
+
+**This model was developed based on:**  [Cosmos Framework](https://github.com/nvidia/cosmos-framework)
+
+**Number of trainable model parameters:**
+
+**Released on: 07/20/2026**
+- Cosmos3-Edge: 4B
+- Cosmos3-Edge-Policy-DROID: 4B
+- Cosmos3-Super-Image2Video-4Step: 64B
+- Cosmos3-Super-Text2Image-4Step: 64B
+
+**Released on: 05/31/2026**
+- Cosmos3-Nano: 16B
+- Cosmos3-Super: 64B
+- Cosmos3-Nano-Policy-DROID: 16B
+- Cosmos3-Super-Image2Video: 64B
+- Cosmos3-Super-Text2Image: 64B
+
+## Input/Output Specifications
+
+- **Generator Input**
+  - **Input Type(s)**: Text, Image, Video, Action Trajectory
+  - **Input Format(s)**:
+    - Text: String
+    - Image: jpg, png, jpeg, webp
+    - Video: mp4
+    - Action: json (1D list)
+  - **Input Parameters**:
+    - Text: One-dimensional (1D)
+    - Image: Two-dimensional (2D)
+    - Video: Three-dimensional (3D)
+    - Action trajectory: One-dimensional (1D)
+  - **Other Properties Related to Input**:
+    - For video inputs, we accept various resolutions, including 480p and 256p.
+    - Image and video inputs are RGB color (8 bits per channel, sRGB color space); grayscale inputs are not supported.
+    - Action input is a per-frame sequence of robot/agent state or control values (e.g., joint positions, gripper state, camera pose). The full input is a 2D array shaped (T, D), where T is the number of frames and D is the embodiment-specific dimensionality listed below.
+    - Input action is only supported for the DROID robot platform (8D).
+  - **Input Size and Length limits:**
+    - **Text:** 4096 tokens
+    - **Image:** 256p and 480p resolution at one of these aspect ratios (16:9, 4:3, 1:1, 3:4, 9:16)
+    - **Video:** 256p and 480p resolution at one of these aspect ratios (16:9, 4:3, 1:1, 3:4, 9:16). Max number of frames = 5.
+    - **Action:** 16 – 400 video frames
+- **Generator Output**
+  - **Output Type(s)**: Image, video, action, text
+  - **Output Format(s)**:
+    - Image: JPG
+    - Video: MP4
+    - Action: 1D list (.json)
+    - Text: string
+  - **Output Parameters**:
+    - Image: Two-dimensional (2D)
+    - Video: Three-dimensional (3D)
+    - Action: One-dimensional (1D)
+    - Text: One-dimensional (1D)
+  - **Other Properties Related to Output**:
+    - The generated video is an MP4 file, with the resolution, frame rate, and duration specified in the input.
+    - Video generation supports durations from 5 to 400 frames, with 189 frames as the default generation duration.
+    - The generated action is only supported for the DROID robot platform (8D).
+    - Video: mp4 at the FPS specified in input
+    - Image: JPEG
+- **Reasoner Input**
+  - **Input Type(s)**: Text, Text+Image, Text+Video
+  - **Input Format(s)**:
+    - Text: String
+    - Image: jpg, png, jpeg, webp
+    - Video: mp4
+  - **Input Parameters**:
+    - Text: One-dimensional (1D)
+    - Image: Two-dimensional (2D)
+    - Video: Three-dimensional (3D)
+  - **Other Properties Related to Input**:
+    - Video inputs are recommended at a frame rate of 4 fps.
+    - Long-context inputs supported up to 256K tokens.
+  - **Input Size and Length limits:**
+    - **Text:** Up to 256K tokens (context window).
+    - **Image:** Standard input image formats; passed as file or URL.
+    - **Video:** mp4 at the recommended 4 fps.
+- **Reasoner Output**
+  - **Output Type(s)**: Text
+  - **Output Format(s)**:
+    - Text: string
+  - **Output Parameters**:
+    - Text: One-dimensional (1D)
+  - **Other Properties Related to Output**:
+    - Default `max_tokens=4096+` is recommended for reasoning outputs; longer outputs may be requested.
+    - Reasoning outputs may include structured chain-of-thought, 2D/3D point localization, and bounding-box coordinates for vision-based tasks.
+
+The video content visualizes the input text description as a short animated scene, capturing key elements within the specified time constraints.
+
+Our AI models are designed and/or optimized to run on NVIDIA GPU-accelerated systems. By leveraging NVIDIA's hardware (e.g., GPU cores) and software frameworks (e.g., CUDA libraries), the model achieves faster training and inference times compared to CPU-only solutions.
+
+## Software Integration
+
+**Runtime Engine(s):**
+
+- [PyTorch](https://github.com/nvidia/cosmos3)
+
+**Supported Hardware Microarchitecture Compatibility:**
+
+- NVIDIA Ampere
+- NVIDIA Blackwell
+- NVIDIA Hopper
+
+**Operating System(s):**
+
+- Linux (We have not tested on other operating systems.)
+
+**Note:** Only BF16 precision is tested. Other precisions like FP4, FP8, and FP16 are not officially supported.
+
+The integration of foundation and fine-tuned models into AI systems requires additional testing using use-case-specific data to ensure safe and effective deployment. Following the V-model methodology, iterative testing and validation at both unit and system levels are essential to mitigate risks, meet technical and functional requirements, and ensure compliance with safety and ethical standards before deployment.
+
+## Training, Testing, and Evaluation Datasets
+
+### Dataset Overview
+
+- **Total Size:** 1.3B data points
+- **Total Number of Datasets:** 393 dataset entries
+- **Dataset partition:** Training [100%], Testing [N/A — evaluation benchmarks used separately], Validation [N/A — evaluation benchmarks used separately]
+- **Time period for training data collection:** 2024–2026
+- **Time period for testing data collection:** N/A (standard public benchmarks)
+- **Time period for validation data collection:** N/A (standard public benchmarks)
+
+Raw data from internal and external sources is transformed into training-ready data through multiple stages of curation, filtering, and quality review. Data acquisition spans diverse multimodal sources — robotics, autonomous driving, industrial environments, indoor and outdoor scenes, varied lighting and weather conditions, camera viewpoints, object categories, and human activities — to broaden coverage across Physical AI operating environments. Automated filtering pipelines remove corrupted, duplicate, low-quality, and restricted content. Metadata analysis, heuristic rules, and model-assisted classifiers are applied during preprocessing to flag anomalous distributions and low-diversity subsets. Human review supplements automated filtering for selected datasets, benchmark construction, and targeted quality analysis. Datasets are balanced across modalities and task categories — visual reasoning, text-to-image, text-to-video, image-to-video, audio generation, video transfer, action-conditioned generation, and action command generation — to reduce overrepresentation of narrow domains. Synthetic and simulation-based augmentation supplements coverage of rare physical interactions and edge-case scenarios. Deduplication and provenance tracking are applied across the corpus. The resulting processed data is converted into model-ready tokenized or encoded representations through modality-specific preprocessors before training begins.
+
+Training datasets passed through multiple layers of automated and manual safeguards designed to reduce the presence of harmful or policy-violating content across categories including weapons and weapons-related instructional content, criminal planning, child sexual abuse material (CSAM), non-consensual intimate imagery (NCII), sexual content involving minors, harassment, hate speech, profanity, threats and incitement to violence, self-harm or suicide-related content, and graphic violence. Data sources are reviewed for licensing compatibility, provenance, and alignment with internal data governance and safety policies before admission into training corpora. Automated filtering pipelines combine multiple detection strategies: hash-matching against known CSAM and NCII reference databases; classifier-based moderation models trained for explicit sexual content, hate speech, violence, weapons imagery, and other restricted categories; keyword and regex-based screening for criminal-planning, threats, and self-harm phrases in text data; metadata and provenance heuristics for source-level risk signals; and embedding-based anomaly detection to surface samples that fall outside expected distributions. Human review and targeted audits supplement automated filtering for selected datasets, benchmark construction, and safety-sensitive evaluation. For multimodal Physical AI data (robotics, autonomous driving, industrial scenes), additional filtering targets invalid action trajectories, physically implausible interactions, and unsafe control sequences. Synthetic and simulation-generated data are evaluated through internal validation before inclusion. Benchmark evaluations and red-team testing are applied post-training to surface remaining safety gaps across world generation, reasoning, audio, and action tasks. No large-scale data-filtering process can guarantee complete removal of all harmful content; residual risks may remain, particularly in rare edge cases or open-world deployment settings. Ongoing monitoring and dataset review continue post-release.
+
+**Data Modality and Training Data Size**
+
+| Modality | Reasoning Data Sample Count | Generation Data Sample Count |
+| -------- | ------------------- | -------------------- |
+| Text     | 22M                 | Not Applicable       |
+| Image    | 19M                 | 767M                 |
+| Video    | 1M                  | 348M                 |
+| Audio    | Not Applicable      | 139M                 |
+| Action   | Not Applicable      | 7M                   |
+
+**Data Collection Method by dataset**
+
+- Hybrid: Automatic/Sensors, Synthetic, Automated
+
+**Labeling Method by dataset**
+
+- Hybrid: Human, Automated
+
+**Properties:** The training, testing, and evaluation datasets consist of diverse multimodal video, image, audio, action, synthetic, and sensor-conditioned data sourced from NVIDIA-owned data and publicly available, commercially permissive datasets. These datasets are curated to exclude known restricted content and to support building an Omni model that learns to generate and reason about dynamic physical environments across world reasoning and generation tasks.
+
+### Public Datasets
+
+| Dataset                   | Samples     |
+|---------------------------|-------------|
+| OpenImage                 | 1.2M        |
+| Coyo700M                  | 100M        |
+| YouTube Video             | 340M        |
+| UMI                       | 4.5M        |
+
+### Private Datasets
+
+| Dataset                   | Samples     |
+|---------------------------|-------------|
+| Egocentric                | 7M          |
+| Nexar                     | 0.6M        |
+| AgiBot                    | 0.2M        |
+| HOI                       | 0.3M        |
+
+### Synthetic Datasets
+
+| Dataset                                 | Samples     |
+|-----------------------------------------|-------------|
+| synthetic images generated using HiDream-I1       | 15M         |
+| synthetic images generated using Qwen-Image-2512  | 14M         |
+| synthetic captions generated using Qwen3-VL       | 1115M       |
+
+## Evaluation Datasets
+
+**Data Collection Method by dataset**
+
+- Hybrid: Automatic/Sensors, Synthetic, Automated
+
+**Labeling Method by dataset**
+
+- Hybrid: Human, Automated
+
+**Properties:** The training, testing, and evaluation datasets consist of diverse multimodal video, image, audio, action, synthetic, and sensor-conditioned data sourced from NVIDIA-owned data and publicly available, commercially permissive datasets. These datasets are curated to exclude known restricted content and to support building an Omni model that learns to generate and reason about dynamic physical environments across world reasoning and generation tasks.
+
+## Benchmarks
+
+### Action
+
+#### RoboLab
+
+![RoboLab — task success rates across language specificity and difficulty levels](images/benchmark-robolab.png)
+
+### PBR (Performance Benchmark Reporting)
+
+
+- Single-GPU end-to-end inference latency (seconds) for Cosmos3-Edge-Policy-DROID.
+- Policy DROID generates an action chunk with shape `[16, 8]`: 16 future action timesteps with 8 action values per timestep.
+- The canonical policy request uses one input image, `num_frames=17`, `fps=5`, and `action_chunk_size=16`.
+- The 17 generated frames represent the rollout associated with the action chunk; they are not a 17-frame input video.
+- An empty cell means not measured, not unsupported.
+
+| Platform | Memory | Configuration | vLLM-Omni E2E Latency (s) | PyTorch E2E Latency (s) |
+|---|---:|---|---:|---:|
+| NVIDIA B200 SXM | 192 GB | Default | 0.99 | — |
+| NVIDIA H100 SXM | 80 GB | Default | 1.41 | 1.25 |
+| NVIDIA H100 NVL | 96 GB | Default | 1.37 | 1.28 |
+| NVIDIA H20 SXM | 96 GB | Default | 3.41 | 2.92 |
+| NVIDIA RTX PRO 6000 Blackwell Server Edition | 96 GB | Default | 1.87 | 1.65 |
+| NVIDIA DGX Station | 128 GB unified memory | Default | 8.11 | 1.77 |
+| NVIDIA DGX Spark | 128 GB unified memory | Default | 7.66 | 6.56 |
+| NVIDIA Jetson AGX Thor T5000 | 128 GB | MAXN | 6.32 | — |
+| NVIDIA Jetson T3000 | 32 GB | 1100 MHz | 8.63 | — |
+| NVIDIA Jetson T2000 | 16 GB | 702 MHz, `THOR_NANO` | — | — |
+
+#### Real-time vLLM-Omni Policy DROID with `[32, 8]` action chunks
+
+These measurements use 320 × 192 observations, 33 generated rollout frames, 30 denoising steps, guidance 1.0, and a 5 Hz control rate. For a 32-action chunk, the real-time budget at 5 Hz is 6.4 seconds. End-to-end (E2E) latency includes client-side HTTP overhead, while inference latency is reported by the server. Each value follows one discarded warmup and two timed runs.
+
+| Platform | E2E Latency | Inference Latency | RTF at 5 Hz | Real-time at 5 Hz |
+|---|---:|---:|---:|---|
+| NVIDIA Jetson AGX Thor T5000, 128 GB, MAXN | 2.59 s | 2.43 s | 2.47 | Yes |
+| NVIDIA Jetson AGX Thor T4000, 64 GB, MAXN, 1530 MHz | 3.24 s | 3.14 s | 1.98 | Yes |
+| NVIDIA Jetson Thor T3000, 32 GB, 1100 MHz (emulated) | 4.05 s | 3.84 s | 1.58 | Yes |
+| NVIDIA Jetson Thor T2000, 16 GB, 765 MHz, `THOR_NANO` | 7.63 s | 7.38 s | 0.84 | No |
+
+#### Real-time PyTorch Policy DROID with `[32, 8]` action chunks
+
+These measurements use an action chunk of `[32, 8]`: 32 future action timesteps with 8 action values per timestep. They are reported separately from the `[16, 8]` Policy DROID measurements above because the action chunk, input processing, and inference configuration differ.
+
+Real-time factor (RTF) is the available time per action chunk divided by measured end-to-end latency; values greater than or equal to 1.0 meet the corresponding real-time budget. For a 32-action chunk, the budgets are 6.4 seconds at 5 Hz and 2.133 seconds at 15 Hz.
+
+These loopback measurements use the PyTorch Policy DROID implementation with one request at a time over a persistent connection, a 640 × 540 observation mapped to a 544 × 736 processing bucket, four UniPC denoising steps, guidance 3.0, and `conditioning_fps=15`. The table evaluates each measured latency against both the 5 Hz and 15 Hz budgets; these are not separate timing runs at the two control rates.
+
+| Platform | Median E2E | Observed Range | Timed Requests | RTF at 5 Hz | RTF at 15 Hz | Real-time at 5 / 15 Hz |
+|---|---:|---:|---:|---:|---:|---|
+| NVIDIA Jetson AGX Thor T5000, 128 GB, MAXN, 1575 MHz | 1.528 s | 1.503–1.673 s | 16 | 4.19 | 1.40 | Yes / Yes |
+| NVIDIA Jetson AGX Thor T4000, 64 GB-class, MAXN, 1530 MHz | 2.208 s | 2.148–2.425 s | 9 | 2.90 | 0.97 | Yes / No |
+| NVIDIA Jetson Thor T3000, 32 GB, 1100 MHz (emulated) | 2.632 s | 2.595–2.687 s | 4 | 2.43 | 0.81 | Yes / No |
+| NVIDIA Jetson Thor T2000, 16 GB, 765 MHz, `THOR_NANO` | 5.195 s | 5.176–5.224 s | 9 | 1.23 | 0.41 | Yes / No |
+
+The T4000 median is approximately 3.5% above the 15 Hz budget. Its fastest observed request was 2.148 seconds, approximately 0.7% above the budget.
+
+The real-time vLLM-Omni and PyTorch measurements should not be compared directly because they use different input processing, denoising configurations, runtime stacks, and measurement protocols. The T3000 results were measured on an emulated configuration with a 32 GB memory limit, 12 SMs, a 1100 MHz GPU clock cap, and CPU caps. DRAM bandwidth was not capped, so a production T3000 may be slower. The T2000 PyTorch measurement used expandable-segment memory allocation and 2.6 GB of swap, with other Cosmos containers stopped; no swap-out was observed during the audited inference calls.
+
+## Usage
+
+- See [Cosmos](https://github.com/nvidia/cosmos) for details.
+
+### Quickstart
+
+Cosmos3-Edge-Policy-DROID is served by a policy **Server** that streams actions to a **Client** driving a simulated or real robot. This example uses [`RoboLab`](https://github.com/NVlabs/RoboLab), a simulation benchmark for task-generalist policies, as the client. Start the server first, then connect the client.
+
+#### Server
+
+First, clone [`cosmos-framework`](https://github.com/NVIDIA/cosmos-framework):
+
+```bash
+git clone https://github.com/NVIDIA/cosmos-framework.git
+cd cosmos-framework
+```
+
+Build the Docker image:
+
+```bash
+docker build \
+  -t cosmos-framework:latest \
+  .
+```
+
+Set your Hugging Face token and launch the container, which installs the dependencies:
+
+```bash
+# Set your Hugging Face token (https://huggingface.co/settings/tokens):
+export HF_TOKEN=<your_hf_token>
+docker run \
+  -it \
+  -e HF_HOME=/workspace/.cache/huggingface \
+  -e HF_TOKEN=$HF_TOKEN \
+  --net host \
+  --rm \
+  --runtime nvidia \
+  -v .:/workspace \
+  -v /workspace/.venv \
+  -v $HOME/.cache/huggingface:/root/.cache/huggingface \
+  cosmos-framework:latest \
+  bash -c '\
+    uv sync \
+      --all-extras \
+      --group=cu130-train \
+      --group=policy-server && \
+    exec bash; \
+  '
+```
+
+Inside the container, start the policy server:
+
+```
+python -m cosmos_framework.scripts.action_policy_server_robolab \
+  --checkpoint_path nvidia/Cosmos3-Edge-Policy-DROID \
+  --port 8000 \
+  --format-prompt-as-json True
+```
+
+#### Client
+
+Clone [`RoboLab`](https://github.com/NVlabs/RoboLab):
+
+```bash
+git clone https://github.com/NVlabs/RoboLab.git
+cd RoboLab
+```
+
+Build the Docker image:
+
+```bash
+./docker/build_docker.sh latest
+```
+
+Launch the container:
+
+```bash
+./docker/run_docker.sh latest
+```
+
+Run a task against the policy server. This opens a viewer window for real-time visualization of the simulation:
+
+```bash
+python policies/cosmos3/run.py \
+  --task BananaInBowlTask
+```
+
+To evaluate across multiple sub-environments in parallel in headless mode:
+
+```bash
+python policies/cosmos3/run.py \
+  --task BananaInBowlTask \
+  --num-envs 10 \
+  --headless
+```
+
+Example output:
+
+<video controls width="864" height="480" src="
+https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID/resolve/main/assets/Pick_up_the_banana_and_place_it_in_the_bowl_0_env0_viewport.mp4"></video>
+
+## Limitations
+
+Cosmos3 may produce imperfect outputs in challenging scenarios. Generation artifacts include temporal inconsistency, unstable camera or object motion, imprecise physical interactions, inaccurate audio-video synchronization, and action-state drift — especially in long-horizon or high-resolution outputs. Reasoning may also be incorrect: object states, causal relationships, spatial geometry, temporal ordering, agent intent, and future outcomes can be misinferred, and complex or long-context inputs may yield hallucinated entities, inconsistent interpretations, or implausible predictions. Because the model lacks an explicit physics simulator, 3D geometry, 4D space-time evolution, object permanence, contact dynamics, and physical laws are only approximated — producing artifacts such as disappearing or morphing objects, unrealistic collisions, and physically implausible motions. Quality further degrades in out-of-distribution environments, safety-critical edge cases, and domains underrepresented in training.
+
+Cosmos3 outputs should not be treated as physically accurate simulation, reliable ground-truth reasoning, or safety-certified decision making. Applications involving robotics control, autonomous systems, scientific simulation, or safety-critical planning require additional validation, external constraints, system-level safety analysis, and domain-specific guardrails before deployment.
+
+## Inference
+
+**Acceleration Engine:** [PyTorch](https://pytorch.org/)
+
+**Test Hardware:** B200, H100, H20, RTX PRO 6000, DGX Station, DGX Spark, Jetson Thor
+
+## Ethical Considerations
+
+NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications.  Developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.
+
+Please make sure you have proper rights and permissions for all input image and video content; if image or video includes people, personal health information, or intellectual property, the image or video generated will not blur or maintain proportions of image subjects included.
+
+Users are responsible for model inputs and outputs. Users are responsible for ensuring safe integration of this model, including implementing guardrails as well as other safety mechanisms, prior to deployment.
+
+For more detailed information on ethical considerations for this model, please see the Model Card++ [Explainability](EXPLAINABILITY.md), [Bias](BIAS.md), [Safety & Security](SAFETY.md), and [Privacy](PRIVACY.md) subcards. Please report model quality, risk, security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).

SAFETY.md

@@ -0,0 +1,11 @@
+## Safety & Security
+
+| Field | Response |
+| :---- | :---- |
+| Model Application(s) | World reasoning and generation for Physical AI. |
+| Describe the life critical impact: | This model is not a safety-certified component and must not be used as the sole basis for life-critical decisions or control without additional system-level validation, safety analysis, and safeguards. The model is not designed or tested by NVIDIA for use in any system or application where the use of or failure of such system or application developed with the model could result in injury, death, or catastrophic damage. NVIDIA is not liable to any party, in whole or in part, for any claims or damages arising from those uses. Any system or application developed with the model must include sufficient safety and redundancy features and comply with applicable legal and regulatory standards and requirements. |
+| Description of methods implemented in data acquisition or processing, if any, to address other types of potentially harmful data in the training, testing, and validation data:  | Training, evaluation, and validation datasets pass through multi-stage automated and manual filtering to reduce harmful, unsafe, restricted, or policy-violating content. Pipelines include source-licensing review, deduplication, metadata-based and classifier-based moderation, embedding-based anomaly detection, and human audits on selected datasets. For Physical AI data (robotics, autonomous driving, industrial scenes), filtering also targets invalid action trajectories, physically implausible interactions, and unsafe control sequences. Synthetic and simulation-generated data are evaluated through internal validation before inclusion. Benchmark and red-team testing surface remaining safety gaps across world generation, reasoning, audio, and action tasks. No data-filtering process can guarantee complete removal; developers are responsible for application-specific safeguards and validation before deployment. |
+| Description of any methods implemented in data acquisition or processing, if any, to address illegal or harmful content in the training data, including, but not limited to, child sexual abuse material (CSAM) and non-consensual intimate imagery (NCII)  | In addition to the general unsafe-content filtering described above, training data acquisition and preprocessing apply CSAM- and NCII-specific safeguards: hash-matching systems against known CSAM databases, classifier-based moderation models trained specifically for explicit content and NCII detection, and provenance and licensing review for sources containing human imagery. Identified content is removed at ingest, with human review and targeted audits supplementing automated filtering for selected datasets. Despite these safeguards, no large-scale data-filtering system can guarantee complete detection. Ongoing monitoring and dataset review continue post-release. |
+| Use Case Restrictions | Use is governed by the [OpenMDW1.1](https://openmdw.ai/) |
+| Model and dataset restrictions | The Principle of least privilege (PoLP) is applied limiting access for dataset generation and model development. Restrictions enforce dataset access during training, and dataset license constraints adhered to. |
+| Responsible Data Handling | This AI model was developed based on our policies to ensure responsible data handling and risk mitigation. The datasets used for training have been scanned for harmful content and illegal content, consistent with our policies including scanning for Child Sexual Abuse Material (CSAM). Ongoing review and monitoring mechanisms are in place based on our policies and to maintain data integrity. |

chat_template.jinja

@@ -0,0 +1,324 @@
+
+{% macro render_extra_keys(json_dict, handled_keys) %}
+    {%- if json_dict is mapping %}
+        {%- for json_key in json_dict if json_key not in handled_keys %}
+            {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+                {{- '
+<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
+            {%- else %}
+                {{-'
+<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{% endmacro %}
+{%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
+{%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
+
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- set loop_messages = messages %}
+{%- for m in loop_messages %}
+  {%- if m["role"] == "user" %}
+    {%- set ns.last_user_idx = loop.index0 %}
+  {%- endif %}
+{%- endfor %}
+
+{%- if messages[0]["role"] == "system" %}
+    {%- if messages[0]["content"] is string %}
+        {%- set system_message = messages[0]["content"] %}
+    {%- else %}
+        {%- set ns = namespace(system_text="") %}
+        {%- for part in messages[0]["content"] %}
+            {%- if 'text' in part %}
+                {%- set ns.system_text = ns.system_text + part.text %}
+            {%- endif %}
+        {%- endfor %}
+        {%- set system_message = ns.system_text %}
+    {%- endif %}
+
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set system_message = "" %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = [] %}
+{%- endif %}
+{# Recompute last_user_idx relative to loop_messages after handling system #}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- for m in loop_messages %}
+  {%- if m["role"] == "user" %}
+    {%- set ns.last_user_idx = loop.index0 %}
+  {%- endif %}
+{%- endfor %}
+{%- if system_message is defined %}
+    {{- "<|im_start|>system
+" + system_message }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- "<|im_start|>system
+" }}
+    {%- endif %}
+{%- endif %}
+{%- if tools is iterable and tools | length > 0 %}
+    {%- if system_message is defined and system_message | length > 0 %}
+        {{- "
+
+" }}
+    {%- endif %}
+    {{- "# Tools
+
+You have access to the following functions:
+
+" }}
+    {{- "<tools>" }}
+    {%- for tool in tools %}
+        {%- if tool.function is defined %}
+            {%- set tool = tool.function %}
+        {%- endif %}
+        {{- "
+<function>
+<name>" ~ tool.name ~ "</name>" }}
+        {%- if tool.description is defined %}
+            {{- '
+<description>' ~ (tool.description | trim) ~ '</description>' }}
+        {%- endif %}
+        {{- '
+<parameters>' }}
+        {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+            {%- for param_name, param_fields in tool.parameters.properties|items %}
+                {{- '
+<parameter>' }}
+                {{- '
+<name>' ~ param_name ~ '</name>' }}
+                {%- if param_fields.type is defined %}
+                    {{- '
+<type>' ~ (param_fields.type | string) ~ '</type>' }}
+                {%- endif %}
+                {%- if param_fields.description is defined %}
+                    {{- '
+<description>' ~ (param_fields.description | trim) ~ '</description>' }}
+                {%- endif %}
+                {%- if param_fields.enum is defined %}
+                    {{- '
+<enum>' ~ (param_fields.enum | tojson | safe) ~ '</enum>' }}
+                {%- endif %}
+                {%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
+                {{- render_extra_keys(param_fields, handled_keys) }}
+                {{- '
+</parameter>' }}
+            {%- endfor %}
+        {%- endif %}
+        {% set handled_keys = ['type', 'properties', 'required'] %}
+        {{- render_extra_keys(tool.parameters, handled_keys) }}
+        {%- if tool.parameters is defined and tool.parameters.required is defined %}
+            {{- '
+<required>' ~ (tool.parameters.required | tojson | safe) ~ '</required>' }}
+        {%- endif %}
+        {{- '
+</parameters>' }}
+        {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+        {{- render_extra_keys(tool, handled_keys) }}
+        {{- '
+</function>' }}
+    {%- endfor %}
+    {{- "
+</tools>" }}
+
+    {{- '
+
+If you choose to call a function ONLY reply in the following format with NO suffix:
+
+<tool_call>
+<function=example_function_name>
+<parameter=example_parameter_1>
+value_1
+</parameter>
+<parameter=example_parameter_2>
+This is the value for the second parameter
+that can span
+multiple lines
+</parameter>
+</function>
+</tool_call>
+
+<IMPORTANT>
+Reminder:
+- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags
+- Required parameters MUST be specified
+- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after
+- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls
+</IMPORTANT>' }}
+{%- endif %}
+
+
+{%- if system_message is defined %}
+    {{- '<|im_end|>
+' }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- '<|im_end|>
+' }}
+    {%- endif %}
+{%- endif %}
+
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- for message in loop_messages %}
+    {%- if message.role == "assistant" %}
+        {# Add reasoning content in to content field for unified processing below. #}
+        {%- set text_content = namespace(text_content="") -%}
+        {%- if message.content is string %}
+            {%- set text_content.text_content = message.content %}
+        {%- else %}
+            {%- for content in message.content %}
+                {%- if 'text' in content %}
+                    {%- set text_content.text_content = content.text %}
+                    {%- break %}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
+            {%- set content = "<think>
+" ~ message.reasoning_content ~ "
+</think>
+" ~ (text_content.text_content | default('', true)) %}
+        {%- else %}
+            {%- set content = text_content.text_content | default('', true) %}
+            {# Allow downstream logic to to take care of broken thought, only handle coherent reasoning here. #}
+            {%- if '<think>' not in content and '</think>' not in content %}
+                {%- set content = "<think></think>" ~ content %}
+            {%- else %}
+                {%- set content = content %}
+            {%- endif -%}
+        {%- endif %}
+        {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
+            {# Assistant message has tool calls. #}
+            {{- '<|im_start|>assistant
+' }}
+                {%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {%- if content is string and content | trim | length > 0 %}
+                    {%- if include_content %}
+                        {{- (content | trim) ~ '
+' -}}
+                    {%- else %}
+                        {%- set c = (content | string) %}
+                        {%- if '</think>' in c %}
+                            {# Keep only content after the last closing think. Also generation prompt causes this. #}
+                            {%- set c = c.split('</think>')[-1] %}
+                        {%- elif '<think>' in c %}
+                            {# If <think> was opened but never closed, drop the trailing think segment #}
+                            {%- set c = c.split('<think>')[0] %}
+                        {%- endif %}
+                        {%- set c = "<think></think>" ~ c | trim %}
+                        {%- if c | length > 0 %}
+                            {{- c ~ '
+' -}}
+                        {%- endif %}
+                    {%- endif %}
+                {%- else %}
+                    {{- "<think></think>" -}}
+                {%- endif %}
+                {%- for tool_call in message.tool_calls %}
+                    {%- if tool_call.function is defined %}
+                        {%- set tool_call = tool_call.function %}
+                    {%- endif %}
+                    {{- '<tool_call>
+<function=' ~ tool_call.name ~ '>
+' -}}
+                        {%- if tool_call.arguments is defined %}
+                            {%- for args_name, args_value in tool_call.arguments|items %}
+                                {{- '<parameter=' ~ args_name ~ '>
+' -}}
+                                    {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                                {{- args_value ~ '
+</parameter>
+' -}}
+                            {%- endfor %}
+                        {%- endif %}
+                    {{- '</function>
+</tool_call>
+' -}}
+                {%- endfor %}
+                {{- '<|im_end|>
+' }}
+        {%- else %}
+            {# Assistant message doesn't have tool calls. #}
+            {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {{- '<|im_start|>assistant
+' ~ (content | default('', true) | string | trim) ~ '<|im_end|>
+' }}
+            {%- else %}
+                {%- set c = (content | default('', true) | string) %}
+                {%- if '<think>' in c and '</think>' in c %}
+                    {%- set c = "<think></think>" ~ c.split('</think>')[-1] %}
+                {%- endif %}
+                {%- set c = c | trim %}
+                {%- if c | length > 0 %}
+                    {{- '<|im_start|>assistant
+' ~ c ~ '<|im_end|>
+' }}
+                {%- else %}
+                    {{- '<|im_start|>assistant
+<|im_end|>
+' }}
+                {%- endif %}
+            {%- endif %}
+        {%- endif %}
+    {%- elif message.role == "user" or message.role == "system" %}
+        {{- '<|im_start|>' + message.role + '
+' }}
+        {%- if message.content is string %}
+            {{- message.content }}
+        {%- else %}
+            {%- for content in message.content %}
+                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}
+                    <|vision_start|><|image_pad|><|vision_end|>
+                {%- elif content.type == 'video' or 'video' in content %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}
+                    <|vision_start|><|video_pad|><|vision_end|>
+                {%- elif 'text' in content %}
+                    {{- content.text }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>
+' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user
+' }}
+        {%- endif %}
+        {{- '<tool_response>
+' }}
+        {{- message.content }}
+        {{- '
+</tool_response>
+' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>
+' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>
+' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>' + message.role + '
+' + message.content + '<|im_end|>
+' }}
+    {%- endif %}
+{%- endfor %}
+
+{%- if add_generation_prompt %}
+    {%- if enable_thinking %}
+        {{- '<|im_start|>assistant
+<think>
+' }}
+    {%- else %}
+        {{- '<|im_start|>assistant
+<think></think>' }}
+    {%- endif %}
+{%- endif %}

checkpoint.json

@@ -0,0 +1,10 @@
+{
+  "config_file": null,
+  "experiment": null,
+  "experiment_overrides": null,
+  "policy": {
+    "action_chunk_size": 32,
+    "conditioning_fps": 15.0,
+    "domain_name": "droid_lerobot"
+  }
+}

config.json

@@ -0,0 +1,244 @@
+{
+  "allow_patterns_overrides": [
+    "*/*.safetensors"
+  ],
+  "architectures": [
+    "Cosmos3ForConditionalGeneration"
+  ],
+  "image_token_id": 19,
+  "model": {
+    "_recursive_": false,
+    "_target": "omni_mot_model",
+    "config": {
+      "_type": "omni_mot_model_config",
+      "action_gen": true,
+      "activation_checkpointing": {
+        "_type": "activation_checkpointing_config",
+        "determinism_check": "default",
+        "mode": "full",
+        "preserve_rng_state": true,
+        "save_ops_regex": [
+          "fmha"
+        ]
+      },
+      "causal_training_strategy": "none",
+      "compile": {
+        "_type": "compile_config",
+        "compile_dynamic": true,
+        "compiled_region": "language",
+        "coordinate_descent_tuning": false,
+        "enabled": false,
+        "max_autotune_pointwise": false,
+        "use_cuda_graphs": false
+      },
+      "diffusion_expert_config": {
+        "_type": "diffusion_expert_config",
+        "base_fps": 24,
+        "enable_fps_modulation": true,
+        "load_weights_from_pretrained": false,
+        "max_vae_latent_side_after_patchify": 20,
+        "patch_spatial": 2,
+        "sound_base_temporal_compression_factor": null,
+        "timestep_range": 1.0,
+        "unified_3d_mrope_reset_spatial_ids": true,
+        "unified_3d_mrope_temporal_modality_margin": 15000,
+        "vision_temporal_position_mode": "latent_index"
+      },
+      "ema": {
+        "_type": "ema_config",
+        "enabled": false,
+        "iteration_shift": 0,
+        "rate": 0.1
+      },
+      "enable_input_bias": true,
+      "exclude_reasoner_weights_from_checkpoint": false,
+      "fixed_step_sampler_config": null,
+      "input_caption_key": "ai_caption",
+      "input_image_key": "images",
+      "input_video_key": "video",
+      "joint_attn_implementation": "two_way",
+      "latent_downsample_factor": 16,
+      "lbl": {
+        "_type": "lbl_config",
+        "coeff_gen": null,
+        "coeff_und": null,
+        "method": "local"
+      },
+      "log_enc_time_every_n": 100,
+      "lora_alpha": 32,
+      "lora_enabled": false,
+      "lora_rank": 16,
+      "lora_target_modules": "q_proj_moe_gen,k_proj_moe_gen,v_proj_moe_gen,o_proj_moe_gen",
+      "max_action_dim": 64,
+      "max_num_tokens_after_packing": -1,
+      "natten_parameter_list": null,
+      "net": null,
+      "num_embodiment_domains": 32,
+      "parallelism": {
+        "_type": "parallelism_config",
+        "attention_io_layout": "sequence_sharded",
+        "cfg_parallel_shard_degree": 1,
+        "context_parallel_shard_degree": 1,
+        "data_parallel_replicate_degree": 1,
+        "data_parallel_shard_degree": 8,
+        "enable_inference_mode": false,
+        "fsdp_master_dtype": "float32"
+      },
+      "precision": "bfloat16",
+      "rectified_flow_inference_config": {
+        "_type": "rectified_flow_inference_config",
+        "num_train_timesteps": 1000,
+        "scheduler_type": "unipc",
+        "shift": 1,
+        "use_dynamic_shifting": false
+      },
+      "rectified_flow_training_config": {
+        "_type": "rectified_flow_training_config",
+        "action_loss_weight": 10.0,
+        "image_loss_scale": null,
+        "independent_action_schedule": false,
+        "independent_sound_schedule": false,
+        "loss_scale": 10.0,
+        "normalize_loss_by_active": false,
+        "shift": {
+          "256": 3,
+          "480": 5
+        },
+        "shift_action": null,
+        "shift_sound": null,
+        "sound_loss_scale": null,
+        "train_time_action_distribution": "logitnormal",
+        "train_time_image_distribution": "logitnormal",
+        "train_time_sound_distribution": "logitnormal",
+        "train_time_video_distribution": "waver",
+        "train_time_weight": "uniform",
+        "use_discrete_rf": false,
+        "use_dynamic_shift": false
+      },
+      "resolution": "480",
+      "sound_dim": null,
+      "sound_gen": false,
+      "sound_latent_fps": 25,
+      "sound_tokenizer": null,
+      "state_ch": 48,
+      "state_t": 300,
+      "tokenizer": {
+        "_target": "wan2pt2_vae_interface",
+        "bucket_name": "",
+        "causal": true,
+        "chunk_duration": 93,
+        "encode_chunk_frames": {
+          "256": 68,
+          "480": 24,
+          "720": 12
+        },
+        "encode_exact_durations": [
+          33
+        ],
+        "keep_decoder_cache": false,
+        "object_store_credential_path_pretrained": "",
+        "spatial_compression_factor": 16,
+        "temporal_compression_factor": 4,
+        "use_streaming_encode": false,
+        "vae_path": "pretrained/tokenizers/video/wan2pt2/Wan2.2_VAE.pth"
+      },
+      "video_temporal_causal": false,
+      "vision_gen": true,
+      "vlm_config": {
+        "_type": "vlm_config",
+        "layer_module": null,
+        "model_instance": {
+          "_target": "nemotron3_dense_vl_text_for_causal_lm",
+          "config": {
+            "_target": "create_vlm_config",
+            "base_config": {
+              "_target": "nemotron3_dense_vl_mot_config_from_json_file",
+              "json_file": "cosmos3://vfm/models/vlm/nemotron_3_dense_vl/configs/Nemotron-2B-Dense-VL.json"
+            },
+            "qk_norm_for_text": false,
+            "use_und_k_norm_for_gen": true
+          }
+        },
+        "model_name": "nvidia/Cosmos3-Edge-Policy-DROID",
+        "pretrained_weights": {
+          "_type": "pretrained_weights_config",
+          "backbone_path": "",
+          "checkpoint_format": "nemotron_3_dense_vl",
+          "credentials_path": "",
+          "enable_gcs_patch_in_boto3": false,
+          "enabled": false
+        },
+        "qk_norm": false,
+        "safetensors_path": "",
+        "tie_word_embeddings": false,
+        "tokenizer": {
+          "_target": "build_processor_lazy",
+          "config_variant": "hf",
+          "repository": null,
+          "revision": null,
+          "subdir": "",
+          "tokenizer_type": "nvidia/Cosmos3-Edge-Policy-DROID"
+        },
+        "use_system_prompt": false
+      }
+    }
+  },
+  "model_type": "cosmos3_omni",
+  "projector_config": {
+    "input_hidden_size": 1152,
+    "merger_intermediate_size": 11520,
+    "model_type": "cosmos3_edge_projector",
+    "out_hidden_size": 2048,
+    "spatial_merge_size": 2,
+    "use_postshuffle_norm": false
+  },
+  "projector_hidden_size": 11520,
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "dtype": "bfloat16",
+    "eos_token_id": 11,
+    "head_dim": 128,
+    "hidden_act": "relu2",
+    "hidden_size": 2048,
+    "initializer_range": 0.02,
+    "intermediate_size": 9216,
+    "max_position_embeddings": 131072,
+    "mlp_bias": false,
+    "model_type": "cosmos3_edge_text",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 8,
+    "rms_norm_eps": 1e-05,
+    "rope_parameters": {
+      "mrope_section": [
+        24,
+        20,
+        20
+      ],
+      "rope_theta": 100000000,
+      "rope_type": "default"
+    },
+    "use_cache": true,
+    "vocab_size": 131072
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "5.14.0.dev0",
+  "video_token_id": 18,
+  "vision_config": {
+    "attention_dropout": 0.0,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "intermediate_size": 4304,
+    "layer_norm_eps": 1e-06,
+    "model_type": "cosmos3_edge_vision",
+    "num_attention_heads": 16,
+    "num_channels": 3,
+    "num_hidden_layers": 27,
+    "num_patches": 256,
+    "patch_size": 16,
+    "spatial_merge_size": 2
+  },
+  "vision_end_token_id": 21,
+  "vision_start_token_id": 20
+}

convert_cosmos3_to_diffusers.py

@@ -0,0 +1,1286 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Convert a Cosmos3 DCP checkpoint to diffusers format.
+
+Example:
+CUDA_VISIBLE_DEVICES=0 python scripts/convert_cosmos3_to_diffusers.py \
+    --checkpoint-path Cosmos3-Nano \
+    --output converted/cosmos3-nano-pipeline \
+    --save-pipeline
+
+A raw Cosmos3 Edge DCP checkpoint is detected automatically. The Edge path
+loads its EMA denoiser weights directly from DCP, packages the matching
+Wan 2.2 VAE, and writes a shared-weight Omni checkpoint: the Diffusers
+transformer is reused by the reasoner and only the missing vision tower is
+stored separately.
+
+Edge checkpoints trained with action generation are also detected
+automatically: the `action2llm`/`llm2action` domain-aware projections and the
+action modality embedding are exported onto the transformer's
+`action_proj_in`/`action_proj_out`/`action_modality_embed` modules, with
+`action_dim` and `num_embodiment_domains` derived from the DCP tensor shapes.
+"""
+
+import argparse
+import contextlib
+import json
+import pathlib
+import re
+import shutil
+
+import torch
+
+from diffusers.models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer
+
+
+DEFAULT_SOUND_TOKENIZER_CONFIG = {
+    "sampling_rate": 48000,
+    "vocoder_input_dim": 64,
+    "dec_dim": 320,
+    "dec_c_mults": [1, 2, 4, 8, 16],
+    "dec_strides": [2, 4, 5, 6, 8],
+    "dec_out_channels": 2,
+}
+
+
+COSMOS3_EDGE_REASONER = "nvidia/Cosmos3-Edge-Reasoner"
+COSMOS3_EDGE_REASONER_REVISION = "590c1c0f1cd7146162d478a2180556055c1a252b"
+COSMOS3_EDGE_VAE = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+
+# Keep the Transformers metadata, but avoid copying its two source shards:
+# Edge's language tower is already byte-identical to the base part of the
+# converted Diffusers transformer. Only the visual tower and projector have
+# to be extracted from source weights.
+COSMOS3_EDGE_REASONER_METADATA_FILES = (
+    "chat_template.jinja",
+    "config.json",
+    "generation_config.json",
+    "preprocessor_config.json",
+    "special_tokens_map.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "video_preprocessor_config.json",
+)
+COSMOS3_EDGE_REASONER_INDEX_FILE = "model.safetensors.index.json"
+COSMOS3_EDGE_REASONER_VISION_SHARD = "00000.safetensors"
+COSMOS3_EDGE_REASONER_FILES = (
+    *COSMOS3_EDGE_REASONER_METADATA_FILES,
+    COSMOS3_EDGE_REASONER_INDEX_FILE,
+    COSMOS3_EDGE_REASONER_VISION_SHARD,
+)
+
+
+_EDGE_ATTN_KEY_REMAP = (
+    (".q_proj_moe_gen.", ".add_q_proj."),
+    (".k_proj_moe_gen.", ".add_k_proj."),
+    (".v_proj_moe_gen.", ".add_v_proj."),
+    (".o_proj_moe_gen.", ".to_add_out."),
+    (".q_norm_moe_gen.", ".norm_added_q."),
+    (".k_norm_moe_gen.", ".norm_added_k."),
+    (".q_proj.", ".to_q."),
+    (".k_proj.", ".to_k."),
+    (".v_proj.", ".to_v."),
+    (".o_proj.", ".to_out."),
+)
+
+_TIME_EMBEDDER_REMAP = {
+    "mlp.0.weight": "linear_1.weight",
+    "mlp.0.bias": "linear_1.bias",
+    "mlp.2.weight": "linear_2.weight",
+    "mlp.2.bias": "linear_2.bias",
+}
+
+
+def _get_config_value(*configs, name, default=None):
+    for config in configs:
+        if config is None:
+            continue
+        if hasattr(config, name):
+            value = getattr(config, name)
+            if value is not None:
+                return value
+        if isinstance(config, dict) and config.get(name) is not None:
+            return config[name]
+    return default
+
+
+def _find_edge_dcp_model_dir(checkpoint_path: pathlib.Path) -> pathlib.Path | None:
+    candidates = [checkpoint_path, checkpoint_path / "model"]
+    if checkpoint_path.is_dir():
+        candidates.extend(metadata_path.parent for metadata_path in checkpoint_path.glob("iter_*/model/.metadata"))
+        candidates.extend(metadata_path.parent for metadata_path in checkpoint_path.glob("*/iter_*/model/.metadata"))
+
+    for candidate in candidates:
+        if not (candidate / ".metadata").is_file():
+            continue
+
+        from torch.distributed.checkpoint.filesystem import FileSystemReader
+
+        state_dict_metadata = FileSystemReader(str(candidate)).read_metadata().state_dict_metadata
+        has_edge_mlp = "net.language_model.model.layers.0.mlp.up_proj.weight" in state_dict_metadata
+        has_qwen_mlp = "net.language_model.model.layers.0.mlp.gate_proj.weight" in state_dict_metadata
+        if has_edge_mlp and not has_qwen_mlp:
+            return candidate
+
+    return None
+
+
+def _remap_edge_dcp_key(key: str) -> str:
+    if key.startswith("language_model."):
+        key = key.removeprefix("language_model.").removeprefix("model.")
+        for old, new in _EDGE_ATTN_KEY_REMAP:
+            if old in key:
+                return key.replace(old, new)
+        return key
+    if key.startswith("vae2llm."):
+        return f"proj_in.{key.removeprefix('vae2llm.')}"
+    if key.startswith("llm2vae."):
+        return f"proj_out.{key.removeprefix('llm2vae.')}"
+    if key.startswith("time_embedder."):
+        time_embedder_key = key.removeprefix("time_embedder.")
+        return f"time_embedder.{_TIME_EMBEDDER_REMAP[time_embedder_key]}"
+    if key.startswith("action2llm."):
+        return f"action_proj_in.{key.removeprefix('action2llm.')}"
+    if key.startswith("llm2action."):
+        return f"action_proj_out.{key.removeprefix('llm2action.')}"
+    if key == "action_modality_embed":
+        return key
+    raise ValueError(f"Unsupported Cosmos3 Edge DCP key: {key!r}.")
+
+
+def _detect_edge_action_config(checkpoint_path: pathlib.Path, use_ema: bool) -> dict | None:
+    from torch.distributed.checkpoint.filesystem import FileSystemReader
+
+    prefix = "net_ema." if use_ema else "net."
+    metadata = FileSystemReader(str(checkpoint_path)).read_metadata().state_dict_metadata
+    action_suffixes = (
+        "action2llm.fc.weight",
+        "action2llm.bias.weight",
+        "llm2action.fc.weight",
+        "llm2action.bias.weight",
+        "action_modality_embed",
+    )
+    action_metadata = {suffix: metadata.get(f"{prefix}{suffix}") for suffix in action_suffixes}
+    present_suffixes = [suffix for suffix, value in action_metadata.items() if value is not None]
+    if not present_suffixes:
+        return None
+
+    # `action2llm`/`llm2action` are DomainAwareLinear modules: `fc` is an
+    # Embedding(num_domains, output_size * input_size) and `bias` an
+    # Embedding(num_domains, output_size), so the shapes fix every config value.
+    missing_suffixes = [suffix for suffix, value in action_metadata.items() if value is None]
+    if missing_suffixes:
+        raise ValueError(
+            f"Cosmos3 Edge {'EMA' if use_ema else 'regular'} action weights are incomplete: "
+            f"present={present_suffixes}, missing={missing_suffixes}."
+        )
+
+    action_shapes = {suffix: tuple(value.size) for suffix, value in action_metadata.items()}
+    expected_ranks = {
+        "action2llm.fc.weight": 2,
+        "action2llm.bias.weight": 2,
+        "llm2action.fc.weight": 2,
+        "llm2action.bias.weight": 2,
+        "action_modality_embed": 1,
+    }
+    invalid_shapes = {
+        suffix: shape for suffix, shape in action_shapes.items() if len(shape) != expected_ranks[suffix]
+    }
+    if invalid_shapes:
+        raise ValueError(f"Cosmos3 Edge action tensors have invalid ranks: {invalid_shapes}.")
+
+    num_embodiment_domains, proj_in_flat = action_shapes["action2llm.fc.weight"]
+    proj_in_domains, hidden_size = action_shapes["action2llm.bias.weight"]
+    proj_out_domains, proj_out_flat = action_shapes["llm2action.fc.weight"]
+    proj_out_bias_domains, action_dim = action_shapes["llm2action.bias.weight"]
+    (action_modality_embed_size,) = action_shapes["action_modality_embed"]
+    domain_counts = {
+        "action2llm.fc.weight": num_embodiment_domains,
+        "action2llm.bias.weight": proj_in_domains,
+        "llm2action.fc.weight": proj_out_domains,
+        "llm2action.bias.weight": proj_out_bias_domains,
+    }
+    if len(set(domain_counts.values())) != 1:
+        raise ValueError(f"Cosmos3 Edge action tensors disagree on embodiment domains: {domain_counts}.")
+    if min(num_embodiment_domains, hidden_size, action_dim) < 1:
+        raise ValueError(f"Cosmos3 Edge action tensor dimensions must be positive: {action_shapes}.")
+
+    expected_projection_size = hidden_size * action_dim
+    if proj_in_flat != expected_projection_size or proj_out_flat != expected_projection_size:
+        raise ValueError(
+            "Cosmos3 Edge action projection shapes are inconsistent: "
+            f"action2llm.fc.weight={action_shapes['action2llm.fc.weight']}, "
+            f"llm2action.fc.weight={action_shapes['llm2action.fc.weight']}; both must factor into "
+            f"hidden_size={hidden_size} x action_dim={action_dim}."
+        )
+    if action_modality_embed_size != hidden_size:
+        raise ValueError(
+            "Cosmos3 Edge action modality embedding has the wrong width: "
+            f"action_modality_embed={action_shapes['action_modality_embed']}, expected ({hidden_size},)."
+        )
+
+    return {
+        "action_gen": True,
+        "action_dim": action_dim,
+        "num_embodiment_domains": num_embodiment_domains,
+    }
+
+
+def _build_edge_transformer(dtype: torch.dtype, action_config: dict | None = None):
+    import inspect
+
+    from accelerate import init_empty_weights
+
+    from diffusers.models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
+
+    action_kwargs = dict(action_config or {})
+    init_params = inspect.signature(Cosmos3OmniTransformer.__init__).parameters
+    edge_required_params = {
+        "hidden_act",
+        "qk_norm_for_text",
+        "backbone_type",
+        "temporal_compression_factor",
+    }
+    action_required_params = set(action_kwargs)
+    missing_edge_params = sorted(edge_required_params - set(init_params))
+    missing_action_params = sorted(action_required_params - set(init_params))
+    if missing_edge_params or missing_action_params:
+        missing_descriptions = []
+        if missing_edge_params:
+            missing_descriptions.append(f"Edge backbone parameters {missing_edge_params}")
+        if missing_action_params:
+            missing_descriptions.append(f"action-generation parameters {missing_action_params}")
+        raise RuntimeError(
+            "The installed diffusers build cannot construct this Cosmos3 Edge checkpoint; it is missing "
+            f"{', '.join(missing_descriptions)}. Use a build that combines the Edge Nemotron-dense backbone "
+            "with upstream action generation support (commit bcc20e452)."
+        )
+
+    with init_empty_weights():
+        transformer = Cosmos3OmniTransformer(
+            **action_kwargs,
+            attention_bias=False,
+            attention_dropout=0.0,
+            base_fps=24,
+            enable_fps_modulation=True,
+            head_dim=128,
+            hidden_act="relu2",
+            hidden_size=2048,
+            intermediate_size=9216,
+            latent_channel=48,
+            latent_patch_size=2,
+            num_attention_heads=16,
+            num_hidden_layers=28,
+            num_key_value_heads=8,
+            patch_latent_dim=192,
+            qk_norm_for_text=False,
+            rms_norm_eps=1e-5,
+            rope_scaling={"mrope_section": [24, 20, 20]},
+            rope_theta=100_000_000.0,
+            timestep_scale=0.001,
+            unified_3d_mrope_reset_spatial_ids=True,
+            unified_3d_mrope_temporal_modality_margin=15000,
+            vocab_size=131072,
+            backbone_type="cosmos3_edge_nemotron_dense",
+            temporal_compression_factor=4,
+        )
+        transformer = transformer.to(dtype=dtype)
+        transformer.time_embedder.to(dtype=torch.float32)
+    return transformer.to_empty(device="cpu")
+
+
+def _validate_edge_action_pipeline_support() -> None:
+    import inspect
+
+    from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline
+
+    pipeline_params = inspect.signature(Cosmos3OmniPipeline.__call__).parameters
+    required_params = {"action", "action_latents"}
+    missing_params = sorted(required_params - set(pipeline_params))
+    if missing_params:
+        raise RuntimeError(
+            "The checkpoint has action generation weights and --save-pipeline was requested, but this diffusers "
+            f"build's Cosmos3OmniPipeline does not accept {missing_params}. Use a build that also includes the "
+            "Cosmos3 action pipeline support from commit bcc20e452."
+        )
+
+
+def _load_edge_dcp_weights(transformer, checkpoint_path: pathlib.Path, use_ema: bool) -> None:
+    import torch.distributed.checkpoint as dcp
+    from torch.distributed.checkpoint.filesystem import FileSystemReader
+
+    prefix = "net_ema." if use_ema else "net."
+    state_dict_metadata = FileSystemReader(str(checkpoint_path)).read_metadata().state_dict_metadata
+    target_state_dict = transformer.state_dict()
+    dcp_state_dict: dict[str, torch.Tensor] = {}
+    for key in state_dict_metadata:
+        if not key.startswith(prefix):
+            continue
+        target_key = _remap_edge_dcp_key(key.removeprefix(prefix))
+        if target_key not in target_state_dict:
+            raise RuntimeError(f"Cosmos3 Edge DCP key {key!r} maps to unknown transformer key {target_key!r}.")
+        dcp_state_dict[key] = target_state_dict[target_key]
+
+    mapped_keys = {_remap_edge_dcp_key(key.removeprefix(prefix)) for key in dcp_state_dict}
+    missing_keys = sorted(set(target_state_dict) - mapped_keys)
+    if missing_keys:
+        raise RuntimeError(f"Cosmos3 Edge DCP is missing transformer weights: {missing_keys}.")
+
+    print(f"Loading {'EMA' if use_ema else 'regular'} Cosmos3 Edge weights from {checkpoint_path} …")
+    dcp.load(state_dict=dcp_state_dict, storage_reader=FileSystemReader(str(checkpoint_path)), no_dist=True)
+    del dcp_state_dict, target_state_dict
+
+
+def _resolve_edge_reasoner_path(args) -> pathlib.Path:
+    if args.reasoner_path is not None:
+        reasoner_path = pathlib.Path(args.reasoner_path).expanduser().absolute()
+    else:
+        from huggingface_hub import snapshot_download
+
+        print(
+            "Downloading the pinned Cosmos3 Edge reasoner snapshot "
+            f"({args.reasoner_repo_id}@{args.reasoner_revision}) …"
+        )
+        reasoner_path = pathlib.Path(
+            snapshot_download(
+                repo_id=args.reasoner_repo_id,
+                revision=args.reasoner_revision,
+                allow_patterns=list(COSMOS3_EDGE_REASONER_FILES),
+            )
+        )
+
+    if not reasoner_path.is_dir():
+        raise FileNotFoundError(f"Cosmos3 Edge reasoner directory not found: {reasoner_path}")
+
+    missing_files = [filename for filename in COSMOS3_EDGE_REASONER_FILES if not (reasoner_path / filename).is_file()]
+    if missing_files:
+        raise FileNotFoundError(
+            f"Cosmos3 Edge reasoner at {reasoner_path} is missing required files: {missing_files}"
+        )
+    return reasoner_path
+
+
+def _load_json(path: pathlib.Path) -> dict:
+    with open(path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _save_json(payload: dict, path: pathlib.Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(payload, f, indent=2, sort_keys=True)
+        f.write("\n")
+
+
+def _edge_reasoner_language_to_generator_key_map() -> dict[str, str]:
+    """Map source Nemotron reasoner keys to their shared Diffusers tensor keys."""
+    mapping = {
+        "lm_head.weight": "lm_head.weight",
+        "model.language_model.embeddings.weight": "embed_tokens.weight",
+        "model.language_model.norm_f.weight": "norm.weight",
+    }
+    attention_key_map = {
+        "q_proj": "to_q",
+        "k_proj": "to_k",
+        "v_proj": "to_v",
+        "o_proj": "to_out",
+    }
+    for layer_idx in range(28):
+        attention_layer_idx = 2 * layer_idx
+        mlp_layer_idx = attention_layer_idx + 1
+        reasoner_attention_prefix = f"model.language_model.layers.{attention_layer_idx}"
+        generator_attention_prefix = f"layers.{layer_idx}"
+        mapping[f"{reasoner_attention_prefix}.norm.weight"] = f"{generator_attention_prefix}.input_layernorm.weight"
+        for source_name, generator_name in attention_key_map.items():
+            mapping[f"{reasoner_attention_prefix}.mixer.{source_name}.weight"] = (
+                f"{generator_attention_prefix}.self_attn.{generator_name}.weight"
+            )
+
+        reasoner_mlp_prefix = f"model.language_model.layers.{mlp_layer_idx}"
+        generator_mlp_prefix = f"layers.{layer_idx}"
+        mapping[f"{reasoner_mlp_prefix}.norm.weight"] = f"{generator_mlp_prefix}.post_attention_layernorm.weight"
+        for source_name in ("up_proj", "down_proj"):
+            mapping[f"{reasoner_mlp_prefix}.mixer.{source_name}.weight"] = (
+                f"{generator_mlp_prefix}.mlp.{source_name}.weight"
+            )
+    return mapping
+
+
+def _native_edge_text_config(source_config: dict) -> dict:
+    """Strip inactive hybrid fields and express the dense reasoner layout explicitly."""
+    layers_block_type = source_config.get("layers_block_type")
+    if layers_block_type is None:
+        legacy_pattern = source_config.get("hybrid_override_pattern")
+        if legacy_pattern is None:
+            num_hidden_layers = source_config.get("num_hidden_layers", 56)
+            if num_hidden_layers % 2:
+                raise ValueError("Cosmos3 Edge requires an even number of alternating attention/MLP layers.")
+            layers_block_type = ["full_attention", "mlp"] * (num_hidden_layers // 2)
+        else:
+            layer_type_map = {"*": "full_attention", "-": "mlp"}
+            try:
+                layers_block_type = [layer_type_map[layer] for layer in legacy_pattern]
+            except KeyError as exc:
+                raise ValueError(f"Unsupported Cosmos3 Edge legacy layer type: {exc.args[0]!r}.") from exc
+    else:
+        layers_block_type = list(layers_block_type)
+
+    if set(layers_block_type) - {"full_attention", "mlp"}:
+        raise ValueError(f"Cosmos3 Edge only supports dense attention/MLP layers, got {layers_block_type!r}.")
+
+    native_config = {"model_type": "cosmos3_edge_text", "layers_block_type": layers_block_type}
+    for key in (
+        "attention_bias",
+        "attention_dropout",
+        "bos_token_id",
+        "dtype",
+        "eos_token_id",
+        "head_dim",
+        "hidden_size",
+        "initializer_range",
+        "intermediate_size",
+        "layer_norm_epsilon",
+        "max_position_embeddings",
+        "mlp_bias",
+        "mlp_hidden_act",
+        "num_attention_heads",
+        "num_key_value_heads",
+        "num_logits_to_keep",
+        "pad_token_id",
+        "rope_theta",
+        "use_cache",
+        "vocab_size",
+    ):
+        if key in source_config:
+            native_config[key] = source_config[key]
+    native_config["num_hidden_layers"] = len(layers_block_type)
+    native_config["rope_parameters"] = {
+        "rope_type": "default",
+        "rope_theta": native_config.get("rope_theta", 100_000_000.0),
+        "mrope_section": source_config.get("mrope_section", [24, 20, 20]),
+    }
+    return native_config
+
+
+def _native_edge_projector_config(source_config: dict) -> dict:
+    merger_intermediate_size = source_config.get("merger_intermediate_size")
+    if merger_intermediate_size is None:
+        merger_intermediate_size = source_config["merger_intermedia"]
+    native_config = {
+        "model_type": "cosmos3_edge_projector",
+        "input_hidden_size": source_config["input_hidden_size"],
+        "merger_intermediate_size": merger_intermediate_size,
+        "out_hidden_size": source_config["out_hidden_size"],
+        "spatial_merge_size": source_config["spatial_merge_size"],
+        "use_postshuffle_norm": source_config["use_postshuffle_norm"],
+    }
+    return native_config
+
+
+def _native_edge_vision_config(source_config: dict) -> dict:
+    native_config = {"model_type": "cosmos3_edge_vision"}
+    for key in (
+        "attention_dropout",
+        "hidden_act",
+        "hidden_size",
+        "intermediate_size",
+        "layer_norm_eps",
+        "num_attention_heads",
+        "num_channels",
+        "num_hidden_layers",
+        "num_patches",
+        "patch_size",
+        "spatial_merge_size",
+    ):
+        if key in source_config:
+            native_config[key] = source_config[key]
+    return native_config
+
+
+def _native_edge_config(source_config: dict) -> dict:
+    native_config = {
+        key: value
+        for key, value in source_config.items()
+        if key not in {"architectures", "auto_map", "model_type", "projector_config", "text_config", "vision_config"}
+    }
+    native_config["architectures"] = ["Cosmos3EdgeForConditionalGeneration"]
+    native_config["model_type"] = "cosmos3_edge"
+    native_config["allow_patterns_overrides"] = ["*/*.safetensors"]
+    native_config["text_config"] = _native_edge_text_config(source_config["text_config"])
+    native_config["vision_config"] = _native_edge_vision_config(source_config["vision_config"])
+    native_config["projector_config"] = _native_edge_projector_config(source_config["projector_config"])
+    return native_config
+
+
+def _native_edge_image_processor_config(source_config: dict) -> dict:
+    native_config = dict(source_config)
+    native_config.pop("auto_map", None)
+    native_config["processor_class"] = "Cosmos3EdgeProcessor"
+    native_config["image_processor_type"] = "Cosmos3EdgeImageProcessor"
+    return native_config
+
+
+def _native_edge_video_processor_config(source_config: dict) -> dict:
+    native_config = dict(source_config)
+    native_config.pop("auto_map", None)
+    native_config["processor_class"] = "Cosmos3EdgeProcessor"
+    native_config["video_processor_type"] = "Cosmos3EdgeVideoProcessor"
+    return native_config
+
+
+def _validate_edge_reasoner_mapping(source_weight_map: dict[str, str], generator_weight_map: dict[str, str]) -> dict[str, str]:
+    language_keys = {
+        key for key in source_weight_map if key == "lm_head.weight" or key.startswith("model.language_model.")
+    }
+    vision_keys = {
+        key for key in source_weight_map if key.startswith(("model.visual.", "model.projector."))
+    }
+    unexpected_keys = set(source_weight_map) - language_keys - vision_keys
+    if unexpected_keys:
+        raise RuntimeError(f"Unexpected Cosmos3 Edge reasoner keys: {sorted(unexpected_keys)}")
+
+    mapping = _edge_reasoner_language_to_generator_key_map()
+    if set(mapping) != language_keys:
+        missing = sorted(language_keys - set(mapping))
+        extra = sorted(set(mapping) - language_keys)
+        raise RuntimeError(f"Incomplete Edge language mapping: missing={missing}, extra={extra}")
+    missing_generator_keys = sorted(set(mapping.values()) - set(generator_weight_map))
+    if missing_generator_keys:
+        raise RuntimeError(f"Diffusers transformer is missing shared reasoner tensors: {missing_generator_keys}")
+    if len(vision_keys) != 443:
+        raise RuntimeError(f"Expected 443 Edge vision/projector tensors, found {len(vision_keys)}.")
+    return mapping
+
+
+def _write_edge_vision_encoder(
+    reasoner_path: pathlib.Path,
+    output_dir: pathlib.Path,
+    source_weight_map: dict[str, str],
+) -> set[str]:
+    try:
+        from safetensors import safe_open
+        from safetensors.torch import save_file
+    except ImportError as exc:
+        raise ImportError("Saving the Cosmos3 Edge vision encoder requires safetensors.") from exc
+
+    vision_keys = {
+        key for key in source_weight_map if key.startswith(("model.visual.", "model.projector."))
+    }
+    non_vision_shards = {source_weight_map[key] for key in vision_keys} - {COSMOS3_EDGE_REASONER_VISION_SHARD}
+    if non_vision_shards:
+        raise RuntimeError(f"Edge visual tensors unexpectedly span shards: {sorted(non_vision_shards)}")
+
+    vision_dir = output_dir / "vision_encoder"
+    vision_dir.mkdir(parents=True, exist_ok=True)
+    vision_path = vision_dir / "model.safetensors"
+    print(f"Extracting {len(vision_keys)} Cosmos3 Edge vision/projector tensors to {vision_path} …")
+    with safe_open(reasoner_path / COSMOS3_EDGE_REASONER_VISION_SHARD, framework="pt", device="cpu") as source_file:
+        state_dict = {key: source_file.get_tensor(key).contiguous() for key in sorted(vision_keys)}
+    save_file(state_dict, str(vision_path), metadata={"format": "pt"})
+    del state_dict
+    # This is an auxiliary shard consumed through the root reasoner manifest;
+    # it is not a standalone Transformers model. Do not leave a stale legacy
+    # config beside it when converting into an existing output directory.
+    (vision_dir / "config.json").unlink(missing_ok=True)
+    return vision_keys
+
+
+def _write_edge_reasoner_manifest(reasoner_path: pathlib.Path, output_dir: pathlib.Path) -> None:
+    source_index = _load_json(reasoner_path / COSMOS3_EDGE_REASONER_INDEX_FILE)
+    source_weight_map = source_index["weight_map"]
+    source_config = _load_json(reasoner_path / "config.json")
+    generator_index = _load_json(output_dir / "transformer" / "diffusion_pytorch_model.safetensors.index.json")
+    generator_weight_map = generator_index["weight_map"]
+    language_mapping = _validate_edge_reasoner_mapping(source_weight_map, generator_weight_map)
+    vision_keys = _write_edge_vision_encoder(reasoner_path, output_dir, source_weight_map)
+
+    weight_map = {
+        generator_key: f"transformer/{generator_weight_map[generator_key]}"
+        for generator_key in language_mapping.values()
+    }
+    weight_map.update({key: "vision_encoder/model.safetensors" for key in vision_keys})
+    if len(weight_map) != len(source_weight_map):
+        raise RuntimeError(
+            f"Edge unified weight index has {len(weight_map)} entries, expected {len(source_weight_map)}."
+        )
+    _save_json(
+        {"metadata": source_index["metadata"], "weight_map": weight_map},
+        output_dir / COSMOS3_EDGE_REASONER_INDEX_FILE,
+    )
+
+    _save_json(_native_edge_config(source_config), output_dir / "config.json")
+
+
+def _copy_edge_reasoner_metadata(reasoner_path: pathlib.Path, output_dir: pathlib.Path) -> None:
+    print(f"Writing the shared Cosmos3 Edge reasoner into {output_dir} …")
+    for filename in COSMOS3_EDGE_REASONER_METADATA_FILES:
+        if filename == "config.json":
+            continue
+        if filename == "preprocessor_config.json":
+            _save_json(
+                _native_edge_image_processor_config(_load_json(reasoner_path / filename)),
+                output_dir / filename,
+            )
+            continue
+        if filename == "video_preprocessor_config.json":
+            _save_json(
+                _native_edge_video_processor_config(_load_json(reasoner_path / filename)),
+                output_dir / filename,
+            )
+            continue
+        shutil.copy2(reasoner_path / filename, output_dir / filename)
+
+    for filename in (
+        "configuration_nemotron_siglip2_h.py",
+        "modeling_cosmos3_edge_omni.py",
+        "modeling_nemotron_siglip2_h.py",
+        "processing.py",
+    ):
+        (output_dir / filename).unlink(missing_ok=True)
+    _write_edge_reasoner_manifest(reasoner_path, output_dir)
+    for filename in ("00000.safetensors", "00001.safetensors"):
+        (output_dir / filename).unlink(missing_ok=True)
+
+
+def _copy_edge_conversion_script(output_dir: pathlib.Path) -> None:
+    source = pathlib.Path(__file__).resolve()
+    destination = output_dir / source.name
+    if source != destination:
+        shutil.copy2(source, destination)
+
+
+def _add_edge_reasoner_to_pipeline(args) -> None:
+    output_dir = pathlib.Path(args.output).expanduser().absolute()
+    expected_paths = ("model_index.json", "scheduler", "text_tokenizer", "transformer", "vae")
+    missing_paths = [str(output_dir / path) for path in expected_paths if not (output_dir / path).exists()]
+    if missing_paths:
+        raise FileNotFoundError(
+            "Expected an existing Cosmos3 Edge Diffusers pipeline before adding its reasoner; "
+            f"missing paths: {missing_paths}"
+        )
+
+    reasoner_path = _resolve_edge_reasoner_path(args)
+    _copy_edge_reasoner_metadata(reasoner_path, output_dir)
+    _copy_edge_conversion_script(output_dir)
+    print("Done.")
+
+
+def _convert_edge_dcp(args, checkpoint_path: pathlib.Path, dtype: torch.dtype) -> None:
+    from transformers import PreTrainedTokenizerFast
+
+    from diffusers import AutoencoderKLWan, UniPCMultistepScheduler
+    from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline
+
+    if args.include_sound_tokenizer or args.sound_tokenizer_path is not None:
+        raise ValueError("The supplied Cosmos3 Edge checkpoint is video-only and cannot include a sound tokenizer.")
+    if args.include_reasoner and not args.save_pipeline:
+        raise ValueError(
+            "A Cosmos3 Edge reasoner can only be included with --save-pipeline because its root Transformers "
+            "config.json would conflict with a transformer-only Diffusers save. Use --no-include-reasoner instead."
+        )
+
+    action_config = _detect_edge_action_config(checkpoint_path, args.use_ema)
+    if action_config is not None:
+        print(
+            "Detected Cosmos3 Edge action generation weights "
+            f"(action_dim={action_config['action_dim']}, "
+            f"num_embodiment_domains={action_config['num_embodiment_domains']})."
+        )
+    transformer = _build_edge_transformer(dtype, action_config)
+    if action_config is not None and args.save_pipeline:
+        _validate_edge_action_pipeline_support()
+    _load_edge_dcp_weights(transformer, checkpoint_path, args.use_ema)
+
+    reasoner_path = _resolve_edge_reasoner_path(args) if args.include_reasoner else None
+
+    output_dir = pathlib.Path(args.output).expanduser().absolute()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    if not args.save_pipeline:
+        print(f"Saving Cosmos3 Edge transformer to {output_dir} …")
+        transformer.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB")
+        print("Done.")
+        return
+
+    tokenizer_source = str(reasoner_path) if reasoner_path is not None else args.reasoner_repo_id
+    tokenizer_kwargs = {}
+    if reasoner_path is None:
+        tokenizer_kwargs["revision"] = args.reasoner_revision
+    text_tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_source, **tokenizer_kwargs)
+    for token in ("<|vision_start|>", "<|vision_end|>"):
+        token_id = text_tokenizer.convert_tokens_to_ids(token)
+        if token_id is None or token_id < 0 or token_id >= transformer.config.vocab_size:
+            raise ValueError(
+                f"Cosmos3 Edge tokenizer token {token!r} has invalid ID {token_id!r} for "
+                f"vocab_size={transformer.config.vocab_size}."
+            )
+
+    diffusers_vae = AutoencoderKLWan.from_pretrained(COSMOS3_EDGE_VAE, subfolder="vae", torch_dtype=torch.bfloat16)
+    scheduler = UniPCMultistepScheduler(
+        num_train_timesteps=1000,
+        solver_order=2,
+        predict_x0=True,
+        solver_type="bh2",
+        lower_order_final=True,
+        final_sigmas_type="zero",
+        use_flow_sigmas=True,
+        prediction_type="flow_prediction",
+        flow_shift=5.0,
+    )
+    pipeline = Cosmos3OmniPipeline(
+        transformer=transformer,
+        text_tokenizer=text_tokenizer,
+        vae=diffusers_vae,
+        scheduler=scheduler,
+        enable_safety_checker=False,
+        default_use_system_prompt=False,
+        use_native_flow_schedule=True,
+    )
+    print(f"Saving Cosmos3 Edge pipeline to {output_dir} …")
+    pipeline.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB")
+    if reasoner_path is not None:
+        _copy_edge_reasoner_metadata(reasoner_path, output_dir)
+    _copy_edge_conversion_script(output_dir)
+    print("Done.")
+
+
+def _load_sound_tokenizer_state_dict(checkpoint_path: pathlib.Path) -> dict[str, torch.Tensor]:
+    if checkpoint_path.suffix == ".safetensors":
+        try:
+            from safetensors.torch import load_file
+        except ImportError as exc:
+            raise ImportError("Loading AVAE .safetensors checkpoints requires safetensors.") from exc
+        checkpoint = load_file(str(checkpoint_path), device="cpu")
+    else:
+        checkpoint = torch.load(checkpoint_path, map_location="cpu")
+
+    if not isinstance(checkpoint, dict):
+        raise TypeError(f"AVAE checkpoint must be a dict, got {type(checkpoint)!r}.")
+
+    for key in ("generator", "state_dict", "model"):
+        value = checkpoint.get(key)
+        if isinstance(value, dict):
+            checkpoint = value
+            break
+
+    state_dict = {
+        key: value.detach().cpu().contiguous() for key, value in checkpoint.items() if isinstance(value, torch.Tensor)
+    }
+    if not state_dict:
+        raise RuntimeError(f"No tensor state dict found in AVAE checkpoint keys: {list(checkpoint.keys())[:16]}")
+    return state_dict
+
+
+def _load_sound_tokenizer_config(config_path: pathlib.Path | None, fallback_config_path: pathlib.Path) -> dict:
+    selected_config_path = config_path
+    if selected_config_path is None and fallback_config_path.exists():
+        selected_config_path = fallback_config_path
+    if selected_config_path is None:
+        return dict(DEFAULT_SOUND_TOKENIZER_CONFIG)
+    with open(selected_config_path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+_SOUND_TOKENIZER_PER_KEY_PREFIXES = ("module.", "generator.", "model.", "state_dict.")
+_SOUND_TOKENIZER_RES_UNIT_INNER_NAMES = {0: "snake1", 1: "conv1", 2: "snake2", 3: "conv2"}
+
+
+def _sound_tokenizer_strip_per_key_prefixes(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    out = dict(state_dict)
+    changed = True
+    while changed:
+        changed = False
+        for prefix in _SOUND_TOKENIZER_PER_KEY_PREFIXES:
+            if any(key.startswith(prefix) for key in out):
+                out = {(key[len(prefix) :] if key.startswith(prefix) else key): value for key, value in out.items()}
+                changed = True
+                break
+        if any(key.startswith(("decoder.", "encoder.", "bottleneck.")) for key in out):
+            break
+    return out
+
+
+def _sound_tokenizer_filter_decoder(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    return {key: value for key, value in state_dict.items() if key.startswith("decoder.")}
+
+
+def _sound_tokenizer_infer_num_blocks(state_dict: dict[str, torch.Tensor]) -> int:
+    block_indices: set[int] = set()
+    for key in state_dict:
+        match = re.match(r"decoder\.layers\.(\d+)\.layers\.\d+\.", key)
+        if match:
+            block_indices.add(int(match.group(1)))
+    return len(block_indices)
+
+
+def _sound_tokenizer_remap_flat_layout(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Convert legacy AVAE `decoder.layers.*` keys to OobleckDecoder attribute keys."""
+    if not any(re.match(r"decoder\.layers\.\d+\.", key) for key in state_dict):
+        return state_dict
+
+    num_blocks = _sound_tokenizer_infer_num_blocks(state_dict)
+    if num_blocks == 0:
+        raise RuntimeError("Detected flat `decoder.layers.*` layout but no decoder blocks were found; cannot remap.")
+    snake1_idx = num_blocks + 1
+    conv2_idx = num_blocks + 2
+
+    def _remap(key: str) -> str:
+        match = re.fullmatch(r"decoder\.layers\.(\d+)\.layers\.(\d+)\.layers\.(\d+)\.(.+)", key)
+        if match:
+            block_n, res_n, inner_n, rest = (
+                int(match.group(1)),
+                int(match.group(2)),
+                int(match.group(3)),
+                match.group(4),
+            )
+            if res_n not in (2, 3, 4):
+                raise RuntimeError(f"Unexpected residual position {res_n} in {key!r}.")
+            inner_name = _SOUND_TOKENIZER_RES_UNIT_INNER_NAMES.get(inner_n)
+            if inner_name is None:
+                raise RuntimeError(f"Unexpected residual inner index {inner_n} in {key!r}.")
+            return f"decoder.block.{block_n - 1}.res_unit{res_n - 1}.{inner_name}.{rest}"
+
+        match = re.fullmatch(r"decoder\.layers\.(\d+)\.layers\.(\d+)\.(.+)", key)
+        if match:
+            block_n, sub_n, rest = int(match.group(1)), int(match.group(2)), match.group(3)
+            block_idx = block_n - 1
+            if sub_n == 0:
+                return f"decoder.block.{block_idx}.snake1.{rest}"
+            if sub_n == 1:
+                return f"decoder.block.{block_idx}.conv_t1.{rest}"
+            raise RuntimeError(f"Unexpected decoder block sub-index {sub_n} in {key!r}.")
+
+        match = re.fullmatch(r"decoder\.layers\.(\d+)\.(.+)", key)
+        if match:
+            layer_n, rest = int(match.group(1)), match.group(2)
+            if layer_n == 0:
+                return f"decoder.conv1.{rest}"
+            if layer_n == snake1_idx:
+                return f"decoder.snake1.{rest}"
+            if layer_n == conv2_idx:
+                return f"decoder.conv2.{rest}"
+            raise RuntimeError(
+                f"Unexpected decoder leaf layer index {layer_n} (expected 0, {snake1_idx}, or {conv2_idx}) in {key!r}."
+            )
+
+        return key
+
+    return {_remap(key): value for key, value in state_dict.items()}
+
+
+def _sound_tokenizer_reshape_snake_params(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    out: dict[str, torch.Tensor] = {}
+    for key, value in state_dict.items():
+        if (key.endswith(".alpha") or key.endswith(".beta")) and value.ndim == 1:
+            value = value.unsqueeze(0).unsqueeze(-1).contiguous()
+        out[key] = value
+    return out
+
+
+def _sound_tokenizer_reapply_weight_norm(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Reconstruct weight-norm parameters if the source checkpoint has folded conv weights."""
+    out = dict(state_dict)
+    candidate_keys = [
+        key
+        for key in state_dict
+        if key.endswith(".weight") and any(f".{layer}." in key for layer in ("conv1", "conv2", "conv_t1"))
+    ]
+    for key in candidate_keys:
+        stem = key[: -len(".weight")]
+        weight_g_key = f"{stem}.weight_g"
+        weight_v_key = f"{stem}.weight_v"
+        if weight_g_key in state_dict or weight_v_key in state_dict:
+            continue
+        weight = state_dict[key]
+        norm_dims = tuple(range(1, weight.ndim))
+        out.pop(key)
+        out[weight_g_key] = weight.norm(p=2, dim=norm_dims, keepdim=True).contiguous()
+        out[weight_v_key] = weight.contiguous()
+    return out
+
+
+def _remap_avae_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Convert a legacy AVAE state dict into the Cosmos3AVAEAudioTokenizer state dict."""
+    state_dict = _sound_tokenizer_strip_per_key_prefixes(state_dict)
+    state_dict = _sound_tokenizer_filter_decoder(state_dict)
+    if not state_dict:
+        raise RuntimeError("Sound tokenizer state dict has no `decoder.*` keys after prefix stripping.")
+    state_dict = _sound_tokenizer_remap_flat_layout(state_dict)
+    state_dict = _sound_tokenizer_reshape_snake_params(state_dict)
+    state_dict = _sound_tokenizer_reapply_weight_norm(state_dict)
+    if any(re.match(r"decoder\.layers\.\d+", key) for key in state_dict):
+        raise RuntimeError("Flat `decoder.layers.*` keys remain after remap; conversion is incomplete.")
+    return state_dict
+
+
+def _build_sound_tokenizer(
+    checkpoint_path: pathlib.Path,
+    config_path: pathlib.Path | None,
+) -> Cosmos3AVAEAudioTokenizer:
+    config = _load_sound_tokenizer_config(config_path, fallback_config_path=pathlib.Path())
+    print(f"Loading AVAE sound tokenizer weights from {checkpoint_path} …")
+    raw_state_dict = _load_sound_tokenizer_state_dict(checkpoint_path)
+    state_dict = _remap_avae_state_dict(raw_state_dict)
+    print(f"  Remapped {len(raw_state_dict)} → {len(state_dict)} decoder keys.")
+
+    sound_tokenizer = Cosmos3AVAEAudioTokenizer(
+        sampling_rate=config.get("sampling_rate", DEFAULT_SOUND_TOKENIZER_CONFIG["sampling_rate"]),
+        vocoder_input_dim=config.get("vocoder_input_dim", DEFAULT_SOUND_TOKENIZER_CONFIG["vocoder_input_dim"]),
+        dec_dim=config.get("dec_dim", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_dim"]),
+        dec_c_mults=tuple(config.get("dec_c_mults", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_c_mults"])),
+        dec_strides=tuple(config.get("dec_strides", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_strides"])),
+        dec_out_channels=config.get("dec_out_channels", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_out_channels"]),
+    )
+    load_result = sound_tokenizer.load_state_dict(state_dict, strict=True)
+    if load_result.missing_keys or load_result.unexpected_keys:
+        raise RuntimeError(
+            "Cosmos3 AVAE sound tokenizer load did not match strictly: "
+            f"missing={load_result.missing_keys}, unexpected={load_result.unexpected_keys}."
+        )
+    return sound_tokenizer
+
+
+@contextlib.contextmanager
+def _skip_source_sound_tokenizer_load(omni_mot_model_cls):
+    original_set_up_tokenizers = omni_mot_model_cls.set_up_tokenizers
+
+    def set_up_tokenizers_without_sound(self):
+        if not getattr(self.config, "sound_gen", False):
+            return original_set_up_tokenizers(self)
+
+        sound_gen = self.config.sound_gen
+        self.config.sound_gen = False
+        try:
+            return original_set_up_tokenizers(self)
+        finally:
+            self.config.sound_gen = sound_gen
+
+    omni_mot_model_cls.set_up_tokenizers = set_up_tokenizers_without_sound
+    try:
+        yield
+    finally:
+        omni_mot_model_cls.set_up_tokenizers = original_set_up_tokenizers
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--checkpoint-path",
+        default="Cosmos3-Nano",
+        help="Named checkpoint (e.g. 'Cosmos3-Nano') or path to a DCP checkpoint directory.",
+    )
+    parser.add_argument("--output", required=True, help="Directory to save the converted diffusers model.")
+    parser.add_argument(
+        "--save-pipeline",
+        action="store_true",
+        help="Save the full pipeline (transformer + VAE + tokenizer + scheduler).",
+    )
+    parser.add_argument(
+        "--dtype", default="bf16", choices=["fp32", "fp16", "bf16"], help="Dtype to save the transformer in."
+    )
+    parser.add_argument(
+        "--sound-tokenizer-path", help="Optional AVAE sound tokenizer checkpoint to save under sound_tokenizer/."
+    )
+    parser.add_argument(
+        "--sound-tokenizer-config-path", help="Optional AVAE config JSON to save under sound_tokenizer/config.json."
+    )
+    parser.add_argument(
+        "--include-sound-tokenizer",
+        action="store_true",
+        help="Require saving sound_tokenizer/ even if the source transformer is video-only.",
+    )
+    parser.add_argument(
+        "--use-ema",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Use `net_ema` weights when converting a raw Cosmos3 Edge DCP checkpoint.",
+    )
+    parser.add_argument(
+        "--include-reasoner",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help=(
+            "Add the pinned Cosmos3 Edge reasoner in shared-weight form so the pipeline output is usable by "
+            "Transformers as well as Diffusers."
+        ),
+    )
+    parser.add_argument(
+        "--reasoner-repo-id",
+        default=COSMOS3_EDGE_REASONER,
+        help="Hugging Face repository containing the Cosmos3 Edge reasoner checkpoint.",
+    )
+    parser.add_argument(
+        "--reasoner-revision",
+        default=COSMOS3_EDGE_REASONER_REVISION,
+        help="Pinned revision of the Cosmos3 Edge reasoner checkpoint.",
+    )
+    parser.add_argument(
+        "--reasoner-path",
+        help="Optional local Cosmos3 Edge reasoner snapshot, used instead of downloading --reasoner-repo-id.",
+    )
+    parser.add_argument(
+        "--copy-edge-reasoner",
+        action="store_true",
+        help=(
+            "Add the pinned shared-weight reasoner and this converter to an existing Cosmos3 Edge Diffusers "
+            "pipeline at --output."
+        ),
+    )
+    args = parser.parse_args()
+
+    if args.copy_edge_reasoner:
+        _add_edge_reasoner_to_pipeline(args)
+        return
+
+    dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[args.dtype]
+    raw_checkpoint_path = pathlib.Path(args.checkpoint_path).expanduser()
+    edge_dcp_model_dir = _find_edge_dcp_model_dir(raw_checkpoint_path) if raw_checkpoint_path.exists() else None
+    if edge_dcp_model_dir is not None:
+        _convert_edge_dcp(args, edge_dcp_model_dir, dtype)
+        return
+
+    from cosmos3.common.init import init_script
+
+    init_script()
+
+    from accelerate import init_empty_weights
+    from cosmos3.args import _CHECKPOINTS
+    from cosmos3.model import Cosmos3OmniModel
+    from projects.cosmos3.vfm.models.omni_mot_model import OmniMoTModel
+    from transformers import AutoTokenizer
+
+    from diffusers import AutoencoderKLWan, UniPCMultistepScheduler
+    from diffusers.models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
+    from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline
+
+    sound_tokenizer_path = (
+        pathlib.Path(args.sound_tokenizer_path).expanduser().absolute() if args.sound_tokenizer_path else None
+    )
+    sound_tokenizer_config_path = (
+        pathlib.Path(args.sound_tokenizer_config_path).expanduser().absolute()
+        if args.sound_tokenizer_config_path
+        else None
+    )
+    if args.include_sound_tokenizer and sound_tokenizer_path is None:
+        raise ValueError("Sound tokenizer output was requested, but --sound-tokenizer-path was not provided.")
+    if sound_tokenizer_path is not None and not sound_tokenizer_path.exists():
+        raise FileNotFoundError(f"Sound tokenizer checkpoint not found: {sound_tokenizer_path}")
+    if sound_tokenizer_config_path is not None and not sound_tokenizer_config_path.exists():
+        raise FileNotFoundError(f"Sound tokenizer config not found: {sound_tokenizer_config_path}")
+
+    checkpoint_name = args.checkpoint_path
+    if checkpoint_name in _CHECKPOINTS:
+        checkpoint_path = pathlib.Path(_CHECKPOINTS[checkpoint_name].download())
+    else:
+        checkpoint_path = pathlib.Path(checkpoint_name).expanduser().absolute()
+    print(f"Resolved checkpoint path: {checkpoint_path}")
+
+    print("Instantiating model and loading weights from DCP checkpoint …")
+    print("Skipping source AVAE tokenizer instantiation during converter-only model load …")
+    with _skip_source_sound_tokenizer_load(OmniMoTModel):
+        _tmp = Cosmos3OmniModel.from_pretrained_dcp(checkpoint_path).model
+
+    # Extract network components and architecture config from DCP model
+    language_model = _tmp.net.language_model
+    vae2llm = _tmp.net.vae2llm
+    llm2vae = _tmp.net.llm2vae
+    time_embedder = _tmp.net.time_embedder
+    lm_cfg = _tmp.net.language_model.config
+    net_cfg = _tmp.net.config
+    model_cfg = _tmp.config
+    patch_latent_dim = _tmp.net.patch_latent_dim
+    hidden_size = _tmp.net.hidden_size
+    num_attention_heads = _tmp.net.num_heads
+    num_key_value_heads = _tmp.net.num_kv_heads
+    head_dim = _tmp.net.head_dim
+    num_hidden_layers = _tmp.net.num_hidden_layers
+    latent_patch_size = _tmp.net.latent_patch_size
+    latent_channel = _tmp.net.latent_channel
+    timestep_scale = _tmp.net.timestep_scale
+    base_fps = int(net_cfg.base_fps)
+    enable_fps_modulation = net_cfg.enable_fps_modulation
+    unified_3d_mrope_reset_spatial_ids = _tmp.config.diffusion_expert_config.unified_3d_mrope_reset_spatial_ids
+    unified_3d_mrope_temporal_modality_margin = (
+        _tmp.config.diffusion_expert_config.unified_3d_mrope_temporal_modality_margin
+    )
+    sound2llm = getattr(_tmp.net, "sound2llm", None)
+    llm2sound = getattr(_tmp.net, "llm2sound", None)
+    sound_modality_embed = getattr(_tmp.net, "sound_modality_embed", None)
+    has_sound_projection_weights = any(module is not None for module in (sound2llm, llm2sound, sound_modality_embed))
+    sound_gen = bool(
+        _get_config_value(net_cfg, model_cfg, name="sound_gen", default=False) or has_sound_projection_weights
+    )
+    sound_dim = _get_config_value(net_cfg, model_cfg, name="sound_dim", default=None)
+    if sound_dim is None and sound2llm is not None:
+        sound_dim = sound2llm.in_features
+    sound_latent_fps = _get_config_value(net_cfg, model_cfg, name="sound_latent_fps", default=25.0)
+    if sound_gen:
+        missing_sound_modules = [
+            name
+            for name, module in (
+                ("sound2llm", sound2llm),
+                ("llm2sound", llm2sound),
+                ("sound_modality_embed", sound_modality_embed),
+            )
+            if module is None
+        ]
+        if missing_sound_modules:
+            raise RuntimeError(
+                "Source checkpoint is configured for sound generation but is missing "
+                f"sound projection weights: {missing_sound_modules}."
+            )
+        if sound_dim is None:
+            raise RuntimeError("Source checkpoint is configured for sound generation but sound_dim is missing.")
+    del _tmp
+    torch.cuda.empty_cache()
+
+    # Init diffusers Cosmos3OmniTransformer with full architecture config from DCP
+    with init_empty_weights():
+        transformer = Cosmos3OmniTransformer(
+            attention_bias=lm_cfg.attention_bias,
+            attention_dropout=lm_cfg.attention_dropout,
+            base_fps=base_fps,
+            enable_fps_modulation=enable_fps_modulation,
+            head_dim=head_dim,
+            hidden_size=hidden_size,
+            intermediate_size=lm_cfg.intermediate_size,
+            latent_channel=latent_channel,
+            latent_patch_size=latent_patch_size,
+            num_attention_heads=num_attention_heads,
+            num_hidden_layers=num_hidden_layers,
+            num_key_value_heads=num_key_value_heads,
+            patch_latent_dim=patch_latent_dim,
+            rms_norm_eps=lm_cfg.rms_norm_eps,
+            rope_scaling=lm_cfg.rope_scaling,
+            rope_theta=lm_cfg.rope_theta,
+            sound_dim=sound_dim,
+            sound_gen=sound_gen,
+            sound_latent_fps=sound_latent_fps,
+            timestep_scale=timestep_scale,
+            unified_3d_mrope_reset_spatial_ids=unified_3d_mrope_reset_spatial_ids,
+            unified_3d_mrope_temporal_modality_margin=unified_3d_mrope_temporal_modality_margin,
+            vocab_size=lm_cfg.vocab_size,
+            backbone_type="cosmos3_edge_nemotron_dense",
+
+        )
+    # The source language_model nests its transformer stack under a `model.` attribute
+    # (HF Qwen-style). Diffusers Cosmos3OmniTransformer holds those layers flat, so
+    # strip the leading `model.` prefix from the language-model state-dict keys.
+    state_dict = {
+        (k[len("model.") :] if k.startswith("model.") else k): v for k, v in language_model.state_dict().items()
+    }
+    # Remap PackedAttentionMoT attribute names from the source (Qwen-style q_proj/k_proj/...
+    # plus cosmos-specific *_moe_gen) to the diffusers AttentionModuleMixin canonical names.
+    # Order matters: the *_moe_gen substrings must be substituted before the plain ones.
+    _ATTN_KEY_REMAP = [
+        (".q_proj_moe_gen.", ".add_q_proj."),
+        (".k_proj_moe_gen.", ".add_k_proj."),
+        (".v_proj_moe_gen.", ".add_v_proj."),
+        (".o_proj_moe_gen.", ".to_add_out."),
+        (".q_norm_moe_gen.", ".norm_added_q."),
+        (".k_norm_moe_gen.", ".norm_added_k."),
+        (".q_proj.", ".to_q."),
+        (".k_proj.", ".to_k."),
+        (".v_proj.", ".to_v."),
+        (".o_proj.", ".to_out."),
+        (".q_norm.", ".norm_q."),
+        (".k_norm.", ".norm_k."),
+    ]
+    remapped_state_dict: dict[str, torch.Tensor] = {}
+    for k, v in state_dict.items():
+        for old, new in _ATTN_KEY_REMAP:
+            if old in k:
+                k = k.replace(old, new)
+                break
+        remapped_state_dict[k] = v
+    state_dict = remapped_state_dict
+    for k, v in vae2llm.state_dict().items():
+        state_dict[f"proj_in.{k}"] = v
+    for k, v in llm2vae.state_dict().items():
+        state_dict[f"proj_out.{k}"] = v
+    _TIME_EMBEDDER_REMAP = {
+        "mlp.0.weight": "linear_1.weight",
+        "mlp.0.bias": "linear_1.bias",
+        "mlp.2.weight": "linear_2.weight",
+        "mlp.2.bias": "linear_2.bias",
+    }
+    for k, v in time_embedder.state_dict().items():
+        state_dict[f"time_embedder.{_TIME_EMBEDDER_REMAP[k]}"] = v
+    if sound_gen:
+        for k, v in sound2llm.state_dict().items():
+            state_dict[f"audio_proj_in.{k}"] = v
+        for k, v in llm2sound.state_dict().items():
+            state_dict[f"audio_proj_out.{k}"] = v
+        state_dict["audio_modality_embed"] = sound_modality_embed
+    transformer.load_state_dict(state_dict, strict=True, assign=True)
+    del (
+        language_model,
+        vae2llm,
+        llm2vae,
+        time_embedder,
+        sound2llm,
+        llm2sound,
+        sound_modality_embed,
+        state_dict,
+    )
+    torch.cuda.empty_cache()
+
+    transformer = transformer.to(dtype=dtype)
+
+    output_dir = pathlib.Path(args.output)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    include_sound_tokenizer = (
+        args.include_sound_tokenizer or sound_tokenizer_path is not None or (sound_gen and args.save_pipeline)
+    )
+    if include_sound_tokenizer and sound_tokenizer_path is None:
+        raise ValueError(
+            "The source checkpoint is configured for sound generation, so --sound-tokenizer-path "
+            "is required when saving a full pipeline."
+        )
+
+    if args.save_pipeline:
+        text_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-VL-8B-Instruct")
+
+        diffusers_vae = AutoencoderKLWan.from_pretrained(
+            "Wan-AI/Wan2.2-TI2V-5B-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16
+        )
+        sound_tokenizer = None
+        if include_sound_tokenizer:
+            assert sound_tokenizer_path is not None
+            sound_tokenizer = _build_sound_tokenizer(sound_tokenizer_path, sound_tokenizer_config_path)
+
+        # Karras schedule approximating FlowUniPCMultistepScheduler with shift=5, 35 steps.
+        # Measured from that schedule: first flow-sigma=0.9998, last flow-sigma=0.1281.
+        # EDM sigma = flow_sigma / (1 - flow_sigma), so:
+        #   sigma_max = 0.9998 / 0.0002 = 4999  (but capped at 200 to avoid duplicate
+        #               integer timesteps from Karras clustering near the top)
+        #   sigma_min = 0.1281 / (1 - 0.1281)  = 0.1281 / 0.8719 ≈ 0.147
+        scheduler = UniPCMultistepScheduler(
+            use_karras_sigmas=True,
+            use_flow_sigmas=True,
+            prediction_type="flow_prediction",
+            sigma_max=200.0,
+            sigma_min=0.147,
+        )
+
+        pipeline = Cosmos3OmniPipeline(
+            transformer=transformer,
+            text_tokenizer=text_tokenizer,
+            vae=diffusers_vae,
+            scheduler=scheduler,
+            sound_tokenizer=sound_tokenizer,
+        )
+        print(f"Saving full pipeline to {output_dir} …")
+        pipeline.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB")
+    else:
+        print(f"Saving transformer to {output_dir} …")
+        transformer.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB")
+        if include_sound_tokenizer:
+            print("Skipping sound_tokenizer/ save because --save-pipeline was not set.")
+
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "do_sample": true,
+  "eos_token_id": 11,
+  "pad_token_id": 0,
+  "transformers_version": "5.14.0.dev0"
+}

model_index.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "Cosmos3OmniPipeline",
+  "_diffusers_version": "0.40.0.dev0",
+  "default_use_system_prompt": false,
+  "enable_safety_checker": false,
+  "safety_checker": [
+    null,
+    null
+  ],
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler"
+  ],
+  "sound_tokenizer": [
+    null,
+    null
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "PreTrainedTokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "Cosmos3OmniTransformer"
+  ],
+  "use_native_flow_schedule": true,
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan"
+  ]
+}

modular_model_index.json

@@ -0,0 +1,58 @@
+{
+  "_blocks_class_name": "Cosmos3OmniBlocks",
+  "_class_name": "Cosmos3OmniModularPipeline",
+  "_diffusers_version": "0.40.0.dev0",
+  "is_distilled": false,
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler",
+    {
+      "pretrained_model_name_or_path": "nvidia/Cosmos3-Edge-Policy-DROID",
+      "subfolder": "scheduler",
+      "type_hint": [
+        "diffusers",
+        "UniPCMultistepScheduler"
+      ],
+      "variant": null
+    }
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "PreTrainedTokenizerFast",
+    {
+      "pretrained_model_name_or_path": "nvidia/Cosmos3-Edge-Policy-DROID",
+      "subfolder": "text_tokenizer",
+      "type_hint": [
+        "transformers",
+        "PreTrainedTokenizerFast"
+      ],
+      "variant": null
+    }
+  ],
+  "transformer": [
+    "diffusers",
+    "Cosmos3OmniTransformer",
+    {
+      "pretrained_model_name_or_path": "nvidia/Cosmos3-Edge-Policy-DROID",
+      "subfolder": "transformer",
+      "type_hint": [
+        "diffusers",
+        "Cosmos3OmniTransformer"
+      ],
+      "variant": null
+    }
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan",
+    {
+      "pretrained_model_name_or_path": "nvidia/Cosmos3-Edge-Policy-DROID",
+      "subfolder": "vae",
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLWan"
+      ],
+      "variant": null
+    }
+  ]
+}

preprocessor_config.json

@@ -0,0 +1,21 @@
+{
+  "size": {
+    "longest_edge": 16777216,
+    "shortest_edge": 65536
+  },
+  "do_resize": true,
+  "patch_size": 16,
+  "merge_size": 2,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "processor_class": "Cosmos3EdgeProcessor",
+  "image_processor_type": "Cosmos3EdgeImageProcessor"
+}

processor_config.json

@@ -0,0 +1,44 @@
+{
+  "processor_class": "Cosmos3EdgeProcessor",
+  "image_processor": {
+    "size": {
+      "longest_edge": 16777216,
+      "shortest_edge": 65536
+    },
+    "do_resize": true,
+    "patch_size": 16,
+    "merge_size": 2,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_processor_type": "Cosmos3EdgeImageProcessor"
+  },
+  "video_processor": {
+    "size": {
+      "longest_edge": 25165824,
+      "shortest_edge": 4096
+    },
+    "temporal_patch_size": 1,
+    "patch_size": 16,
+    "merge_size": 2,
+    "image_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "video_processor_type": "Cosmos3EdgeVideoProcessor",
+    "return_metadata": true
+  }
+}

scheduler/scheduler_config.json

@@ -0,0 +1,33 @@
+{
+  "_class_name": "UniPCMultistepScheduler",
+  "_diffusers_version": "0.40.0.dev0",
+  "beta_end": 0.02,
+  "beta_schedule": "linear",
+  "beta_start": 0.0001,
+  "disable_corrector": [],
+  "dynamic_thresholding_ratio": 0.995,
+  "final_sigmas_type": "zero",
+  "flow_shift": 5.0,
+  "lower_order_final": true,
+  "num_train_timesteps": 1000,
+  "predict_x0": true,
+  "prediction_type": "flow_prediction",
+  "rescale_betas_zero_snr": false,
+  "sample_max_value": 1.0,
+  "shift_terminal": null,
+  "sigma_max": null,
+  "sigma_min": null,
+  "solver_order": 2,
+  "solver_p": null,
+  "solver_type": "bh2",
+  "steps_offset": 0,
+  "thresholding": false,
+  "time_shift_type": "exponential",
+  "timestep_spacing": "linspace",
+  "trained_betas": null,
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": false,
+  "use_exponential_sigmas": false,
+  "use_flow_sigmas": true,
+  "use_karras_sigmas": false
+}

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

text_tokenizer/chat_template.jinja

@@ -0,0 +1,324 @@
+
+{% macro render_extra_keys(json_dict, handled_keys) %}
+    {%- if json_dict is mapping %}
+        {%- for json_key in json_dict if json_key not in handled_keys %}
+            {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+                {{- '
+<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
+            {%- else %}
+                {{-'
+<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{% endmacro %}
+{%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
+{%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
+
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- set loop_messages = messages %}
+{%- for m in loop_messages %}
+  {%- if m["role"] == "user" %}
+    {%- set ns.last_user_idx = loop.index0 %}
+  {%- endif %}
+{%- endfor %}
+
+{%- if messages[0]["role"] == "system" %}
+    {%- if messages[0]["content"] is string %}
+        {%- set system_message = messages[0]["content"] %}
+    {%- else %}
+        {%- set ns = namespace(system_text="") %}
+        {%- for part in messages[0]["content"] %}
+            {%- if 'text' in part %}
+                {%- set ns.system_text = ns.system_text + part.text %}
+            {%- endif %}
+        {%- endfor %}
+        {%- set system_message = ns.system_text %}
+    {%- endif %}
+
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set system_message = "" %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = [] %}
+{%- endif %}
+{# Recompute last_user_idx relative to loop_messages after handling system #}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- for m in loop_messages %}
+  {%- if m["role"] == "user" %}
+    {%- set ns.last_user_idx = loop.index0 %}
+  {%- endif %}
+{%- endfor %}
+{%- if system_message is defined %}
+    {{- "<|im_start|>system
+" + system_message }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- "<|im_start|>system
+" }}
+    {%- endif %}
+{%- endif %}
+{%- if tools is iterable and tools | length > 0 %}
+    {%- if system_message is defined and system_message | length > 0 %}
+        {{- "
+
+" }}
+    {%- endif %}
+    {{- "# Tools
+
+You have access to the following functions:
+
+" }}
+    {{- "<tools>" }}
+    {%- for tool in tools %}
+        {%- if tool.function is defined %}
+            {%- set tool = tool.function %}
+        {%- endif %}
+        {{- "
+<function>
+<name>" ~ tool.name ~ "</name>" }}
+        {%- if tool.description is defined %}
+            {{- '
+<description>' ~ (tool.description | trim) ~ '</description>' }}
+        {%- endif %}
+        {{- '
+<parameters>' }}
+        {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+            {%- for param_name, param_fields in tool.parameters.properties|items %}
+                {{- '
+<parameter>' }}
+                {{- '
+<name>' ~ param_name ~ '</name>' }}
+                {%- if param_fields.type is defined %}
+                    {{- '
+<type>' ~ (param_fields.type | string) ~ '</type>' }}
+                {%- endif %}
+                {%- if param_fields.description is defined %}
+                    {{- '
+<description>' ~ (param_fields.description | trim) ~ '</description>' }}
+                {%- endif %}
+                {%- if param_fields.enum is defined %}
+                    {{- '
+<enum>' ~ (param_fields.enum | tojson | safe) ~ '</enum>' }}
+                {%- endif %}
+                {%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
+                {{- render_extra_keys(param_fields, handled_keys) }}
+                {{- '
+</parameter>' }}
+            {%- endfor %}
+        {%- endif %}
+        {% set handled_keys = ['type', 'properties', 'required'] %}
+        {{- render_extra_keys(tool.parameters, handled_keys) }}
+        {%- if tool.parameters is defined and tool.parameters.required is defined %}
+            {{- '
+<required>' ~ (tool.parameters.required | tojson | safe) ~ '</required>' }}
+        {%- endif %}
+        {{- '
+</parameters>' }}
+        {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+        {{- render_extra_keys(tool, handled_keys) }}
+        {{- '
+</function>' }}
+    {%- endfor %}
+    {{- "
+</tools>" }}
+
+    {{- '
+
+If you choose to call a function ONLY reply in the following format with NO suffix:
+
+<tool_call>
+<function=example_function_name>
+<parameter=example_parameter_1>
+value_1
+</parameter>
+<parameter=example_parameter_2>
+This is the value for the second parameter
+that can span
+multiple lines
+</parameter>
+</function>
+</tool_call>
+
+<IMPORTANT>
+Reminder:
+- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags
+- Required parameters MUST be specified
+- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after
+- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls
+</IMPORTANT>' }}
+{%- endif %}
+
+
+{%- if system_message is defined %}
+    {{- '<|im_end|>
+' }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- '<|im_end|>
+' }}
+    {%- endif %}
+{%- endif %}
+
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- for message in loop_messages %}
+    {%- if message.role == "assistant" %}
+        {# Add reasoning content in to content field for unified processing below. #}
+        {%- set text_content = namespace(text_content="") -%}
+        {%- if message.content is string %}
+            {%- set text_content.text_content = message.content %}
+        {%- else %}
+            {%- for content in message.content %}
+                {%- if 'text' in content %}
+                    {%- set text_content.text_content = content.text %}
+                    {%- break %}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
+            {%- set content = "<think>
+" ~ message.reasoning_content ~ "
+</think>
+" ~ (text_content.text_content | default('', true)) %}
+        {%- else %}
+            {%- set content = text_content.text_content | default('', true) %}
+            {# Allow downstream logic to to take care of broken thought, only handle coherent reasoning here. #}
+            {%- if '<think>' not in content and '</think>' not in content %}
+                {%- set content = "<think></think>" ~ content %}
+            {%- else %}
+                {%- set content = content %}
+            {%- endif -%}
+        {%- endif %}
+        {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
+            {# Assistant message has tool calls. #}
+            {{- '<|im_start|>assistant
+' }}
+                {%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {%- if content is string and content | trim | length > 0 %}
+                    {%- if include_content %}
+                        {{- (content | trim) ~ '
+' -}}
+                    {%- else %}
+                        {%- set c = (content | string) %}
+                        {%- if '</think>' in c %}
+                            {# Keep only content after the last closing think. Also generation prompt causes this. #}
+                            {%- set c = c.split('</think>')[-1] %}
+                        {%- elif '<think>' in c %}
+                            {# If <think> was opened but never closed, drop the trailing think segment #}
+                            {%- set c = c.split('<think>')[0] %}
+                        {%- endif %}
+                        {%- set c = "<think></think>" ~ c | trim %}
+                        {%- if c | length > 0 %}
+                            {{- c ~ '
+' -}}
+                        {%- endif %}
+                    {%- endif %}
+                {%- else %}
+                    {{- "<think></think>" -}}
+                {%- endif %}
+                {%- for tool_call in message.tool_calls %}
+                    {%- if tool_call.function is defined %}
+                        {%- set tool_call = tool_call.function %}
+                    {%- endif %}
+                    {{- '<tool_call>
+<function=' ~ tool_call.name ~ '>
+' -}}
+                        {%- if tool_call.arguments is defined %}
+                            {%- for args_name, args_value in tool_call.arguments|items %}
+                                {{- '<parameter=' ~ args_name ~ '>
+' -}}
+                                    {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                                {{- args_value ~ '
+</parameter>
+' -}}
+                            {%- endfor %}
+                        {%- endif %}
+                    {{- '</function>
+</tool_call>
+' -}}
+                {%- endfor %}
+                {{- '<|im_end|>
+' }}
+        {%- else %}
+            {# Assistant message doesn't have tool calls. #}
+            {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {{- '<|im_start|>assistant
+' ~ (content | default('', true) | string | trim) ~ '<|im_end|>
+' }}
+            {%- else %}
+                {%- set c = (content | default('', true) | string) %}
+                {%- if '<think>' in c and '</think>' in c %}
+                    {%- set c = "<think></think>" ~ c.split('</think>')[-1] %}
+                {%- endif %}
+                {%- set c = c | trim %}
+                {%- if c | length > 0 %}
+                    {{- '<|im_start|>assistant
+' ~ c ~ '<|im_end|>
+' }}
+                {%- else %}
+                    {{- '<|im_start|>assistant
+<|im_end|>
+' }}
+                {%- endif %}
+            {%- endif %}
+        {%- endif %}
+    {%- elif message.role == "user" or message.role == "system" %}
+        {{- '<|im_start|>' + message.role + '
+' }}
+        {%- if message.content is string %}
+            {{- message.content }}
+        {%- else %}
+            {%- for content in message.content %}
+                {%- if content.type == 'image' or 'image' in content or 'image_url' in content %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                    {%- if add_vision_id %}Picture {{ image_count.value }}: {% endif -%}
+                    <|vision_start|><|image_pad|><|vision_end|>
+                {%- elif content.type == 'video' or 'video' in content %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                    {%- if add_vision_id %}Video {{ video_count.value }}: {% endif -%}
+                    <|vision_start|><|video_pad|><|vision_end|>
+                {%- elif 'text' in content %}
+                    {{- content.text }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>
+' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user
+' }}
+        {%- endif %}
+        {{- '<tool_response>
+' }}
+        {{- message.content }}
+        {{- '
+</tool_response>
+' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>
+' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>
+' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>' + message.role + '
+' + message.content + '<|im_end|>
+' }}
+    {%- endif %}
+{%- endfor %}
+
+{%- if add_generation_prompt %}
+    {%- if enable_thinking %}
+        {{- '<|im_start|>assistant
+<think>
+' }}
+    {%- else %}
+        {{- '<|im_start|>assistant
+<think></think>' }}
+    {%- endif %}
+{%- endif %}

text_tokenizer/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

text_tokenizer/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4dc692a99dca6d92a44e9529ffb53541eb79bb260eec7281bc51268c05d7311c
+size 17077500

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4dc692a99dca6d92a44e9529ffb53541eb79bb260eec7281bc51268c05d7311c
+size 17077500

transformer/config.json

@@ -0,0 +1,44 @@
+{
+  "_class_name": "Cosmos3OmniTransformer",
+  "_diffusers_version": "0.40.0.dev0",
+  "action_dim": 64,
+  "action_gen": true,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "base_fps": 24,
+  "enable_fps_modulation": true,
+  "head_dim": 128,
+  "hidden_act": "relu2",
+  "hidden_size": 2048,
+  "intermediate_size": 9216,
+  "latent_channel": 48,
+  "latent_patch_size": 2,
+  "num_attention_heads": 16,
+  "num_embodiment_domains": 32,
+  "num_hidden_layers": 28,
+  "num_key_value_heads": 8,
+  "patch_latent_dim": 192,
+  "qk_norm_for_text": false,
+  "rms_norm_eps": 1e-05,
+  "rope_axes_dim": [
+    24,
+    20,
+    20
+  ],
+  "rope_scaling": {
+    "mrope_section": [
+      24,
+      20,
+      20
+    ]
+  },
+  "rope_theta": 100000000.0,
+  "sound_dim": null,
+  "sound_gen": false,
+  "sound_latent_fps": 25.0,
+  "timestep_scale": 0.001,
+  "unified_3d_mrope_reset_spatial_ids": true,
+  "unified_3d_mrope_temporal_modality_margin": 15000,
+  "use_und_k_norm_for_gen": true,
+  "vocab_size": 131072
+}

transformer/cosmos_framework_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1815af3a075108820c1f3c26d50b999299a8d831526cf561a20249bf1bb6fcd1
+size 11032

transformer/diffusion_pytorch_model-00001-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:99b0ff1f9dbdcb497fac16cbc0ebe33dee19653b4c182fa0d474433dcc845206
+size 5000041944

transformer/diffusion_pytorch_model-00002-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1842fefdacd0dcbb8ab543b59848a0b696e0c9634676bc76db9e0c971fe328d5
+size 1748779008

transformer/diffusion_pytorch_model.safetensors.index.json

@@ -0,0 +1,556 @@
+{
+  "metadata": {
+    "total_size": 6748759424
+  },
+  "weight_map": {
+    "action_modality_embed": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "action_proj_in.bias.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "action_proj_in.fc.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "action_proj_out.bias.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "action_proj_out.fc.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "embed_tokens.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.23.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.23.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.24.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.25.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.26.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.input_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.mlp.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.mlp.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.post_attention_layernorm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.27.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.to_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.3.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.input_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.input_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.mlp.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.mlp.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.mlp_moe_gen.down_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.mlp_moe_gen.up_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.post_attention_layernorm.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.post_attention_layernorm_moe_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.k_norm_und_for_gen.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.to_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors",
+    "lm_head.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "norm.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "norm_moe_gen.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "proj_in.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "proj_in.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "proj_out.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "time_embedder.linear_1.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "time_embedder.linear_1.weight": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "time_embedder.linear_2.bias": "diffusion_pytorch_model-00002-of-00002.safetensors",
+    "time_embedder.linear_2.weight": "diffusion_pytorch_model-00002-of-00002.safetensors"
+  }
+}

vae/config.json

@@ -0,0 +1,129 @@
+{
+  "_class_name": "AutoencoderKLWan",
+  "_diffusers_version": "0.40.0.dev0",
+  "_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
+}

vae/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:230496cb59ff85bc9c040487737c4062480cb61c71e697b197b4c30142f2a0da
+size 1409400600

video_preprocessor_config.json

@@ -0,0 +1,22 @@
+{
+  "size": {
+    "longest_edge": 25165824,
+    "shortest_edge": 4096
+  },
+  "temporal_patch_size": 1,
+  "patch_size": 16,
+  "merge_size": 2,
+  "image_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "processor_class": "Cosmos3EdgeProcessor",
+  "video_processor_type": "Cosmos3EdgeVideoProcessor",
+  "return_metadata": true
+}

vision_encoder/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2180ad739ecc96b5c1e9386892d3c5c08bfa42b9cdab9aabc53b028671db89b3
+size 978739880