Duplicate from nvidia/Cosmos3-Edge

Co-authored-by: Yue Tang <tangyue0820@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,54 @@
+*.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
+assets/example_reasoning_input.png filter=lfs diff=lfs merge=lfs -text
+images/benchmark-reasoning.png filter=lfs diff=lfs merge=lfs -text
+assets/example_action_fd_agibotworld_4chunk_output.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/example_action_fd_agibotworld_first_frame.png filter=lfs diff=lfs merge=lfs -text
+assets/example_action_id_av_0_input.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/example_action_id_av_1_input.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/example_i2v_input.jpg filter=lfs diff=lfs merge=lfs -text
+assets/example_i2v_output.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/example_action_id_av_0_output.png filter=lfs diff=lfs merge=lfs -text
+assets/example_action_id_av_1_output.png filter=lfs diff=lfs merge=lfs -text
+images/benchmark-overall.png filter=lfs diff=lfs merge=lfs -text
+assets/benchmark-overall.png filter=lfs diff=lfs merge=lfs -text
+assets/edge_i2v_output.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/edge_action_fd_umi_2chunk_output.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/edge_action_id_av_0_output.png filter=lfs diff=lfs merge=lfs -text
+assets/edge_action_id_av_1_output.png filter=lfs diff=lfs merge=lfs -text
+assets/diffusers_outputs/edge_i2v_diffusers.mp4 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,1093 @@
+---
+license: other
+license_name: openmdw1.1-license
+license_link: >-
+  https://openmdw.ai/license/1-1/
+library_name: cosmos
+tags:
+  - nvidia
+  - cosmos
+  - cosmos3
+---
+
+# **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
+
+## 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.
+
+> **Update — August 25, 2026:** The Cosmos3-Edge generator checkpoint, runtime defaults, usage examples, and benchmark results have been updated. Users pulling from `main` should refresh their local snapshot. See the [update announcement](https://huggingface.co/nvidia/Cosmos3-Edge/discussions/62) for compatibility and reproducibility details.
+
+**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, Action Trajectory
+  - **Input Format(s)**:
+    - Text: String
+    - Image: jpg, png, jpeg, webp
+    - Action: JSON — 2D array shaped (T, D), where T is the number of frames and D is the embodiment-specific dimensionality
+  - **Input Parameters**:
+    - Text: One-dimensional (1D)
+    - Image: Two-dimensional (2D)
+    - Action trajectory: Two-dimensional (2D)
+  - **Other Properties Related to Input**:
+    - Image input is 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 compatible embodiments, including general camera motion (9D), autonomous vehicle (9D), egocentric motion (57D), single Franka Panda arm with RobotiQ gripper (10D), dual Franka Panda arm with RobotiQ gripper (20D), Agibot (29D), UR (10D), Google robot (10D), WidowX 250 (10D), UMI (10D).
+  - **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)
+    - **Action:** 16 – 400 sequence length
+- **Generator Output**
+  - **Output Type(s)**: Image, video, action, text
+  - **Output Format(s)**:
+    - Image: JPG
+    - Video: MP4
+    - Action: JSON
+    - Text: string
+  - **Output Parameters**:
+    - Image: Two-dimensional (2D)
+    - Video: Three-dimensional (3D)
+    - Action: Two-dimensional (2D)
+    - 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 256p and 480p resolution, 12–30 fps, and 50–150 frames. These are set per request via the `size`, `fps`, and `num_frames` fields.
+    - The generated action is only supported for compatible embodiments, including general camera motion (9D), autonomous vehicle (9D), egocentric motion (57D), single Franka Panda arm with RobotiQ gripper (10D), dual Franka Panda arm with RobotiQ gripper (20D), Agibot (29D), UR (10D), Google robot (10D), WidowX 250 (10D), UMI (10D).
+    - 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):**
+
+- [vLLM-Omni](https://github.com/vllm-project/vllm-omni)
+- [vLLM](https://github.com/vllm-project/vllm)
+- [PyTorch](https://github.com/nvidia/cosmos3)
+- [Hugging Face Diffusers](https://github.com/huggingface/diffusers)
+
+**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, 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, 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.
+
+- For more information about the datasets used to train this model, please see the [Public Summary of Training Content](https://docs.nvidia.com/cosmos/latest/_downloads/e482b7114ce8dbfbb07d2d4b42cafe4e/training-content-Cosmos-3.pdf).
+
+**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                 |
+| 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, 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, 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
+
+For detailed evaluations of the base model, see our [technical paper](https://research.nvidia.com/labs/cosmos-lab/cosmos3/technical-report.pdf).
+
+### Overall
+
+The table below summarizes Cosmos3-Edge across reasoning and generation. Each reasoning column (General, Robotics, Smart Infrastructure, Driving) reports the average score over that capability's benchmarks. For generation, **Image2Video** is the PAIBench overall score and **Policy: Robot** is the RoboLab success rate. In each column, the best result is in **bold** and the second-best is <u>underlined</u>.
+\* denotes post-trained Cosmos3 variants: [Cosmos3-Nano-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Nano-Policy-DROID) and [Cosmos3-Edge-Policy-DROID](https://huggingface.co/nvidia/Cosmos3-Edge-Policy-DROID#robolab).
+
+![Overall benchmark results](assets/benchmark-overall.png)
+
+### Reasoning Benchmarks
+
+![Reasoning benchmarks](images/benchmark-reasoning.png)
+
+### Generation Benchmarks
+
+#### Image-to-Video Generation
+
+All models are evaluated on image-to-video generation at **480p, 24 fps**. Throughput is the number of frames generated per second, measured in eager mode on a single **NVIDIA H100** GPU. Cosmos3-Edge delivers the highest generation throughput while achieving competitive quality across PAIBench, RBench, and PhysicsIQ.
+
+![Image-to-Video benchmark results](assets/benchmark-image2video.png)
+
+
+#### Action
+
+The Edge model is a strong initialization for downstream action tasks. For example, post-training it on the DROID dataset produces a policy whose RoboLab success rate is reported in the [Cosmos3-Edge-Policy-DROID model card](https://huggingface.co/nvidia/Cosmos3-Edge-Policy-DROID#robolab).
+
+### PBR (Performance Benchmark Reporting)
+
+The following tables report single-GPU or single-platform inference performance for the Cosmos3-Edge **Generator** and **Reasoner** towers.
+
+Generator results are measured using end-to-end or generation latency in seconds; lower is better. Reasoner results include serving and token-generation metrics, such as time to first token, request latency, and throughput.
+
+All results were measured using a single GPU and a batch size of 1.
+
+#### Generator
+
+Unless otherwise noted, visual-generation benchmarks use **480p resolution**. Image-to-video benchmarks generate **189 frames**.
+
+##### vLLM-Omni
+
+| GPU or Platform | Image-to-Video | Forward Dynamics | Inverse Dynamics | Policy DROID |
+| --- | ---: | ---: | ---: | ---: |
+| B200 SXM 192 GB | — | 2.44 s | 3.98 s | 0.99 s |
+| H100 SXM 80 GB | 27.64 s | 3.91 s | 5.60 s | 1.41 s |
+| H100 NVL 96 GB | 35.60 s | 4.73 s | 6.39 s | 1.37 s |
+| H20 SXM 96 GB | 108.16 s | 12.77 s | 15.49 s | 3.41 s |
+| RTX PRO 6000 Blackwell Server Edition | 36.29 s | 5.65 s | 7.46 s | 1.87 s |
+| DGX Station | 12.17 s | 4.33 s | 6.34 s | 8.11 s |
+| DGX Spark | 165.96 s | 26.43 s | 30.86 s | 7.66 s |
+| Jetson AGX Thor T5000, 128 GB, MAXN | 137.50 s | 6.05 s | 7.19 s | 6.32 s |
+| Jetson T3000, 32 GB, 1100 MHz | 194.76 s | 8.67 s | 10.25 s | 8.63 s |
+| Jetson T2000, 16 GB, 702 MHz, THOR_NANO | 101.20 s | — | — | — |
+
+##### PyTorch
+
+| GPU or Platform | Image-to-Video | Forward Dynamics | Inverse Dynamics | Policy DROID |
+| --- | ---: | ---: | ---: | ---: |
+| H100 SXM 80 GB | 23.92 s | 3.69 s | 3.56 s | 1.25 s |
+| H100 NVL 96 GB | 32.24 s | 4.64 s | 4.52 s | 1.28 s |
+| H20 SXM 96 GB | 97.51 s | 12.78 s | 12.64 s | 2.92 s |
+| RTX PRO 6000 Blackwell Server Edition | 38.98 s | 5.26 s | 5.66 s | 1.32 s |
+| DGX Station | 10.57 s | 2.16 s | 2.26 s | 1.30 s |
+| DGX Spark | 179.80 s | 24.59 s | 26.76 s | 5.44 s |
+| Jetson AGX Thor T5000, 128 GB, MAXN | 153.00 s | — | — | — |
+| Jetson T3000, 32 GB, 1100 MHz | 227.80 s | — | — | — |
+
+#### Reasoner
+
+The following tables report Cosmos3-Edge Reasoner performance. Reasoner workloads produce autoregressively generated text and therefore use different metrics from the Generator workloads:
+
+- **Time To First Token (TTFT):** Time from request submission until the first output token is produced. Lower is better.
+- **Request Latency:** End-to-end latency for the complete request. Lower is better.
+- **Request Throughput:** Completed requests per second. Higher is better.
+- **Output Token Throughput:** Generated output tokens per second. Higher is better.
+
+##### vLLM serving benchmarks
+
+These measurements use the `nvidia/Cosmos3-Edge` checkpoint with one GPU. Metrics were collected at client-side concurrency levels of 1, 64, 128, and 256.
+
+The workload notation is **input sequence length / output sequence length / video FPS**.
+
+###### RTX PRO 4500 Blackwell Server Edition
+
+| Input / Output / Video FPS | Metric | Concurrency 1 | Concurrency 64 | Concurrency 128 | Concurrency 256 |
+| --- | --- | ---: | ---: | ---: | ---: |
+| 50 / 1 / 1 | Time To First Token (ms) ↓ | 165.79 | 8817.33 | 14702.20 | 29482.39 |
+|  | Request Latency (ms) ↓ | 165.79 | 8817.33 | 14702.20 | 29482.39 |
+|  | Request Count (requests) | 50 | 320 | 256 | 512 |
+|  | Request Throughput (req/s) ↑ | 6.00 | 6.55 | 6.55 | 6.52 |
+|  | Output Token Throughput (tok/s) ↑ | 6.00 | 6.55 | 6.55 | 6.52 |
+| 50 / 1 / 2 | Time To First Token (ms) ↓ | 371.67 | 20375.98 | 33812.45 | 68201.55 |
+|  | Request Latency (ms) ↓ | 371.67 | 20375.98 | 33812.45 | 68201.55 |
+|  | Request Count (requests) | 50 | 313 | 249 | 492 |
+|  | Request Throughput (req/s) ↑ | 2.68 | 2.77 | 2.76 | 2.71 |
+|  | Output Token Throughput (tok/s) ↑ | 2.68 | 2.77 | 2.76 | 2.71 |
+| 50 / 100 / 1 | Time To First Token (ms) ↓ | 166.86 | 6900.90 | 19625.83 | 45729.55 |
+|  | Request Latency (ms) ↓ | 764.15 | 16667.01 | 29196.84 | 55749.62 |
+|  | Request Count (requests) | 50 | 320 | 256 | 512 |
+|  | Request Throughput (req/s) ↑ | 1.31 | 3.73 | 3.74 | 3.70 |
+|  | Output Token Throughput (tok/s) ↑ | 130.63 | 372.40 | 373.98 | 369.87 |
+| 50 / 100 / 2 | Time To First Token (ms) ↓ | 374.93 | 23526.65 | 47550.99 | 101553.31 |
+|  | Request Latency (ms) ↓ | 1041.29 | 33712.54 | 57641.53 | 111895.20 |
+|  | Request Count (requests) | 50 | 320 | 256 | 512 |
+|  | Request Throughput (req/s) ↑ | 0.96 | 1.79 | 1.79 | 1.78 |
+|  | Output Token Throughput (tok/s) ↑ | 95.74 | 178.73 | 178.89 | 178.15 |
+
+###### RTX PRO 6000 Blackwell Server Edition
+
+| Input / Output / Video FPS | Metric | Concurrency 1 | Concurrency 64 | Concurrency 128 | Concurrency 256 |
+| --- | --- | ---: | ---: | ---: | ---: |
+| 50 / 1 / 1 | Time To First Token (ms) ↓ | 141.99 | 3213.91 | 5384.51 | 10792.72 |
+|  | Request Latency (ms) ↓ | 141.99 | 3213.91 | 5384.51 | 10792.72 |
+|  | Request Count (requests) | 50 | 320 | 254 | 512 |
+|  | Request Throughput (req/s) ↑ | 6.96 | 18.00 | 17.95 | 17.89 |
+|  | Output Token Throughput (tok/s) ↑ | 6.96 | 18.00 | 17.95 | 17.89 |
+| 50 / 1 / 2 | Time To First Token (ms) ↓ | 239.86 | 7483.22 | 12552.69 | 25259.11 |
+|  | Request Latency (ms) ↓ | 239.86 | 7483.22 | 12552.69 | 25259.11 |
+|  | Request Count (requests) | 49 | 303 | 249 | 491 |
+|  | Request Throughput (req/s) ↑ | 4.06 | 7.28 | 7.49 | 7.34 |
+|  | Output Token Throughput (tok/s) ↑ | 4.06 | 7.28 | 7.49 | 7.34 |
+| 50 / 100 / 1 | Time To First Token (ms) ↓ | 138.74 | 943.46 | 2680.17 | 11599.63 |
+|  | Request Latency (ms) ↓ | 503.44 | 6188.90 | 13022.07 | 26388.89 |
+|  | Request Count (requests) | 50 | 320 | 256 | 512 |
+|  | Request Throughput (req/s) ↑ | 1.98 | 10.27 | 9.57 | 8.95 |
+|  | Output Token Throughput (tok/s) ↑ | 197.75 | 1026.14 | 956.47 | 893.91 |
+| 50 / 100 / 2 | Time To First Token (ms) ↓ | 239.24 | 1798.96 | 11644.84 | 33293.32 |
+|  | Request Latency (ms) ↓ | 638.71 | 13599.89 | 26299.90 | 49165.91 |
+|  | Request Count (requests) | 50 | 320 | 256 | 512 |
+|  | Request Throughput (req/s) ↑ | 1.56 | 4.66 | 4.50 | 4.45 |
+|  | Output Token Throughput (tok/s) ↑ | 155.93 | 465.28 | 449.57 | 444.17 |
+
+##### Embedded-platform eager Transformers benchmarks
+
+These preliminary measurements use raw Hugging Face Transformers in eager mode rather than vLLM. They are presented separately because their runtime, workload, and metric definitions differ from the vLLM serving benchmarks above.
+
+| Board | Specification | Input | Prompt Tokens | Prefill Throughput | Prefill Latency | Decode Throughput | E2E Latency |
+| --- | --- | --- | ---: | ---: | ---: | ---: | ---: |
+| Jetson AGX Thor T5000 | 128 GB / MAXN | Text | 1705 | 8717 tok/s | 0.20 s | 37.3 tok/s | 3.60 s |
+| Jetson AGX Thor T5000 | 128 GB / MAXN | Image | 911 | 4845 tok/s | 0.19 s | 42.6 tok/s | 3.17 s |
+| Jetson AGX Thor T5000 | 128 GB / MAXN | Video | 1263 | 6032 tok/s | 0.21 s | 41.8 tok/s | 3.25 s |
+| Jetson AGX Thor T4000 | 64 GB / MAXN, 1530 MHz | Text | 1705 | 6519 tok/s | 0.26 s | 34.1 tok/s | 3.99 s |
+| Jetson AGX Thor T4000 | 64 GB / MAXN, 1530 MHz | Image | 911 | 3471 tok/s | 0.26 s | 40.3 tok/s | 3.41 s |
+| Jetson AGX Thor T4000 | 64 GB / MAXN, 1530 MHz | Video | 1263 | 4164 tok/s | 0.30 s | 38.1 tok/s | 3.64 s |
+| Jetson Thor T3000 | 32 GB / 1100 MHz | Text | 1705 | 5230 tok/s | 0.33 s | 29.7 tok/s | 4.61 s |
+| Jetson Thor T3000 | 32 GB / 1100 MHz | Image | 911 | 2710 tok/s | 0.34 s | 36.3 tok/s | 3.83 s |
+| Jetson Thor T3000 | 32 GB / 1100 MHz | Video | 1263 | 3388 tok/s | 0.37 s | 33.7 tok/s | 4.14 s |
+| Jetson Thor T2000 | 16 GB / 702 MHz, THOR_NANO | Text | 1705 | 2355 tok/s | 0.72 s | 15.7 tok/s | 8.80 s |
+| Jetson Thor T2000 | 16 GB / 702 MHz, THOR_NANO | Image | 911 | 1233 tok/s | 0.74 s | 19.6 tok/s | 7.21 s |
+| Jetson Thor T2000 | 16 GB / 702 MHz, THOR_NANO | Video | 1263 | 1543 tok/s | 0.82 s | 18.0 tok/s | 7.87 s |
+| Jetson AGX Orin | 64 GB | Text | 1705 | 3260 tok/s | 0.52 s | 12.3 tok/s | 10.83 s |
+| Jetson AGX Orin | 64 GB | Image | 911 | 1840 tok/s | 0.50 s | 12.3 tok/s | 10.81 s |
+| Jetson AGX Orin | 64 GB | Video | 1263 | 2103 tok/s | 0.60 s | 12.2 tok/s | 10.97 s |
+
+**Benchmark notes**
+
+- All Generator measurements use one GPU or one integrated computing platform.
+- Generator values are average end-to-end or generation latency in seconds; lower is better.
+- Unless otherwise specified, visual-generation measurements use **480p resolution**.
+- Image-to-video measurements generate **189 output frames**.
+- Jetson AGX Thor T5000 and Jetson T3000 visual-generation measurements use **832 × 480** resolution.
+- Jetson T2000 visual-generation measurements use **448 × 256** resolution and therefore should not be compared directly with the 480p results. Its image-to-video values are warm-run measurements generating 189 frames.
+- PyTorch Generator values report average generation latency rather than diffusion-only latency.
+- Datacenter and enterprise forward- and inverse-dynamics results use the autonomous-driving (`AV`) configuration.
+- Jetson AGX Thor T5000 and Jetson T3000 forward-dynamics, inverse-dynamics, and policy measurements use the DROID configuration with action chunk `[16, 8]`.
+
+## Usage
+
+- See [Cosmos](https://github.com/nvidia/cosmos) for details.
+
+### Prompt upsampling
+
+For optimal quality, prompts should be upsampled into a specific JSON structure. Description and code can be found [here](https://github.com/nvidia/cosmos-framework/blob/main/docs/prompt_upsampling.md).
+
+For example, for image-to-video upsampling using Opus-4.6:
+
+```bash
+git clone https://github.com/NVIDIA/cosmos-framework.git packages/cosmos-framework
+pip install -e packages/cosmos-framework
+
+export PROMPT_UPSAMPLER_ENDPOINT_URL="https://api.anthropic.com/v1/"
+export PROMPT_UPSAMPLER_MODEL_NAME="claude-opus-4-6"
+export PROMPT_UPSAMPLER_API_TOKEN="<your_token>"
+
+python -m cosmos_framework.inference.prompt_upsampling \
+    --input inputs/prompt_upsampler/prompts_i2v.txt \
+    --image-list inputs/prompt_upsampler/images.txt \
+    --output outputs/prompt_upsampler/upsampled_i2v_prompts_opus \
+    --mode image2video \
+    --endpoint-url "${PROMPT_UPSAMPLER_ENDPOINT_URL}" \
+    --model "${PROMPT_UPSAMPLER_MODEL_NAME}" \
+    --api-token "${PROMPT_UPSAMPLER_API_TOKEN}" \
+    --resolution 480 \
+    --aspect-ratio "16,9" \
+    --duration "5s" \
+    --fps 24
+```
+
+For image-to-video, provide either one shared image via `--image-url` or one image per prompt via `--image-list` (the image-list file must have the same number of non-empty lines as the prompt file). Accepted image formats: local paths, HTTP(S) URLs, and `data:` URLs.
+
+### vLLM-Omni
+
+#### Container
+
+```
+docker pull vllm/vllm-omni:cosmos3
+```
+
+#### General Invocation
+
+You can use the release-tested `vllm-omni` package for deploying an OpenAI-compatible API inference endpoint.
+The recommended vLLM-Omni serving configuration for nvidia/Cosmos3-Edge on a single GPU is:
+
+```bash
+vllm serve nvidia/Cosmos3-Edge \
+  --omni \
+  --host 0.0.0.0 \
+  --port 8000 \
+  --init-timeout 1800
+```
+
+To speed up inference with additional GPUs, enable context parallelism with `--ulysses-degree` or switch to tensor parallelism with `--tensor-parallel-size`. Setting `--enable-layerwise-offload` can help reduce memory usage on GPUs with less available memory.
+
+#### Examples
+
+##### Download example prompts
+
+The example inputs (`assets/`) live in this model repo. Download just this folder with the Hugging Face CLI:
+
+```bash
+pip install -U "huggingface_hub[cli]"
+hf download nvidia/Cosmos3-Edge assets/ --local-dir Cosmos3-Edge
+cd Cosmos3-Edge
+```
+
+Run all commands below from the downloaded repo root.
+
+---
+
+##### Image to video generation
+
+```python
+import json
+import mimetypes
+from pathlib import Path
+
+import requests
+
+# 1. Read JSON-upsampled prompt and negative prompt
+json_prompt = json.load(open("assets/example_i2v_prompt.json"))
+negative_prompt = json.load(open("assets/negative_prompt.json"))
+
+# 2. Build and send the multipart API request
+url = "http://localhost:8000/v1/videos/sync"
+image_path = Path("assets/example_i2v_input.jpg")
+mime_type = mimetypes.guess_type(image_path)[0] or "image/png"
+data = {
+    "prompt": json.dumps(json_prompt),
+    "negative_prompt": json.dumps(negative_prompt),
+    "size": "832x480",
+    "num_frames": "121",
+    "fps": "24",
+    "num_inference_steps": "20",
+    "guidance_scale": "6.0",
+    "max_sequence_length": "4096",
+    "flow_shift": "12.0",
+    "extra_params": json.dumps(
+        {
+          "use_resolution_template": False,
+          "use_duration_template": False,
+          "guardrails": True,
+        }
+    ),
+    "seed": "0",
+}
+
+with image_path.open("rb") as image_file:
+    files = {
+        "input_reference": (image_path.name, image_file, mime_type),
+    }
+    print("Sending request to server...")
+    response = requests.post(
+        url,
+        data=data,
+        files=files,
+        headers={"Accept": "video/mp4"},
+    )
+    response.raise_for_status()
+
+# 3. Save the generated video
+output_path = Path("/tmp/cosmos3_edge_i2v.mp4")
+output_path.write_bytes(response.content)
+print(f"Saved video to {output_path}")
+```
+
+Example output:
+
+<video controls width="832" height="480" src="https://huggingface.co/nvidia/Cosmos3-Edge/resolve/main/assets/edge_i2v_output.mp4"></video>
+
+---
+
+##### Action generation
+
+The forward-dynamics example uses UMI robotics action trajectories, and the inverse-dynamics examples use autonomous-vehicle (AV) action trajectories. Source files:
+
+- Forward dynamics first frame: `assets/example_action_fd_umi_first_frame.png`
+- Forward dynamics action chunks: `assets/example_action_fd_umi_action_chunks.json`
+- Forward dynamics output video: `assets/edge_action_fd_umi_2chunk_output.mp4`
+- Inverse dynamics source videos: `assets/example_action_id_av_0_input.mp4`, `assets/example_action_id_av_1_input.mp4`
+- Inverse dynamics predicted actions: `assets/edge_action_id_av_0_output.json`, `assets/edge_action_id_av_1_output.json`
+
+###### Action forward dynamics
+
+The example below performs a 2-chunk UMI robotics rollout with the vLLM-Omni `/v1/videos/sync` inference endpoint. Each request sends one conditioning frame through `input_reference` and one 16-step normalized action chunk through `extra_params["action"]`. The request also sets the top-level `size` field to the input image resolution, so vLLM-Omni returns each chunk at the same resolution as the conditioning image without reflection padding. The stitched output drops each chunk's conditioning frame, producing 32 generated frames. The script extracts the last generated frame from each chunk and uses it as the next chunk's conditioning frame.
+
+```python
+import json
+import mimetypes
+from pathlib import Path
+
+import imageio.v3 as iio
+import numpy as np
+import requests
+from PIL import Image
+
+url = "http://localhost:8000/v1/videos/sync"
+first_frame_path = Path("assets/example_action_fd_umi_first_frame.png")
+action_spec = json.loads(Path("assets/example_action_fd_umi_action_chunks.json").read_text())
+action_chunks = action_spec["action_chunks"]
+
+prompt = action_spec.get("prompt", "mouse arrangement")
+fps = int(action_spec.get("fps", 20))
+action_chunk_size = int(action_spec.get("action_chunk_size", 16))
+current_frame_path = first_frame_path
+input_width, input_height = Image.open(first_frame_path).size
+chunk_video_paths = []
+stitch_frames = []
+
+for chunk_idx, action_chunk in enumerate(action_chunks):
+    mime_type = mimetypes.guess_type(current_frame_path)[0] or "image/png"
+    extra_params = {
+        "action_mode": "forward_dynamics",
+        "domain_name": action_spec.get("domain_name", "umi"),
+        "action_chunk_size": action_chunk_size,
+        "image_size": action_spec.get("image_size", 256),
+        "view_point": action_spec.get("view_point", "ego_view"),
+        "action": action_chunk,
+        "guardrails": True,
+    }
+    data = {
+        "prompt": prompt,
+        "num_frames": str(action_chunk_size + 1),  # conditioning frame + generated frames
+        "fps": str(fps),
+        "size": f"{input_width}x{input_height}",  # return chunks at input resolution
+        "num_inference_steps": "30",
+        "guidance_scale": "1.0",
+        "flow_shift": "10.0",
+        "seed": str(chunk_idx),
+        "extra_params": json.dumps(extra_params),
+    }
+
+    with current_frame_path.open("rb") as image_file:
+        files = {"input_reference": (current_frame_path.name, image_file, mime_type)}
+        print(f"Sending action FD chunk {chunk_idx} to vLLM-Omni...")
+        response = requests.post(
+            url,
+            data=data,
+            files=files,
+            headers={"Accept": "video/mp4"},
+            timeout=600,
+        )
+        response.raise_for_status()
+
+    chunk_video_path = Path(f"/tmp/cosmos3_edge_action_fd_chunk_{chunk_idx:02d}.mp4")
+    chunk_video_path.write_bytes(response.content)
+    chunk_video_paths.append(chunk_video_path)
+
+    # The returned chunk contains the conditioning frame followed by generated frames.
+    # Drop the conditioning frame when stitching the generated-only rollout.
+    frames = iio.imread(chunk_video_path)
+    stitch_frames.extend(frames[1:])
+
+    # Autoregressive conditioning: use the final generated frame from this chunk
+    # as the input image for the next vLLM-Omni request.
+    if chunk_idx + 1 < len(action_chunks):
+        current_frame_path = Path(f"/tmp/cosmos3_edge_action_fd_ar_frame_{chunk_idx + 1:02d}.png")
+        iio.imwrite(current_frame_path, frames[-1])
+
+stitched_path = Path("/tmp/cosmos3_edge_action_fd_umi_2chunk.mp4")
+iio.imwrite(stitched_path, np.asarray(stitch_frames), fps=fps)
+print("Generated chunk videos:", chunk_video_paths)
+print("Saved stitched rollout:", stitched_path)
+print("stitched resolution:", f"{input_width}x{input_height}")
+```
+
+
+Example output:
+
+<video width="512" controls src="https://huggingface.co/nvidia/Cosmos3-Edge/resolve/main/assets/edge_action_fd_umi_2chunk_output.mp4"></video>
+
+###### Action inverse dynamics
+
+```python
+import json
+import time
+from pathlib import Path
+
+import requests
+
+base_url = "http://localhost:8000"
+input_videos = {
+    "av_inverse_0": Path("assets/example_action_id_av_0_input.mp4"),
+    "av_inverse_1": Path("assets/example_action_id_av_1_input.mp4"),
+}
+
+for name, video_path in input_videos.items():
+    extra_params = {
+        "action_mode": "inverse_dynamics",
+        "domain_name": "av",
+        "action_chunk_size": 60,
+        "image_size": 480,
+        "view_point": "ego_view",
+        "raw_action_dim": 9,
+        "guardrails": True,
+    }
+    data = {
+        "prompt": "You are an autonomous vehicle planning system.",
+        "num_frames": "61",
+        "fps": "10",
+        "num_inference_steps": "30",
+        "guidance_scale": "1.0",
+        "flow_shift": "10.0",
+        "seed": "0",
+        "extra_params": json.dumps(extra_params),
+    }
+
+    with video_path.open("rb") as video_file:
+        files = {
+            "input_reference": (video_path.name, video_file, "video/mp4"),
+        }
+        print(f"Submitting {name} request to server...")
+        response = requests.post(f"{base_url}/v1/videos", data=data, files=files)
+        response.raise_for_status()
+    initial = response.json()
+
+    while True:
+        response = requests.get(f"{base_url}/v1/videos/{initial['id']}", timeout=30)
+        response.raise_for_status()
+        final = response.json()
+        print(initial["id"], final.get("status"), f"{final.get('progress', 0)}%")
+        if final.get("status") == "completed":
+            break
+        if final.get("status") in {"failed", "cancelled"}:
+            raise RuntimeError(json.dumps(final, indent=2))
+        time.sleep(2)
+
+    action = final.get("action")
+    if not action or "data" not in action:
+        raise RuntimeError(f"Response did not include action data: {json.dumps(final, indent=2)}")
+
+    output_path = Path(f"/tmp/cosmos3_edge_action_id_{name}.json")
+    output_path.write_text(json.dumps(action, indent=2))
+    print(f"Saved predicted action to {output_path}")
+    print("action shape:", action.get("shape"), "dtype:", action.get("dtype"))
+```
+
+Example outputs:
+
+- [av_inverse_0 predicted action JSON](https://huggingface.co/nvidia/Cosmos3-Edge/blob/main/assets/edge_action_id_av_0_output.json)
+- [av_inverse_1 predicted action JSON](https://huggingface.co/nvidia/Cosmos3-Edge/blob/main/assets/edge_action_id_av_1_output.json)
+
+<img width="1280" src="https://huggingface.co/nvidia/Cosmos3-Edge/resolve/main/assets/edge_action_id_av_0_output.png">
+
+<img width="1280" src="https://huggingface.co/nvidia/Cosmos3-Edge/resolve/main/assets/edge_action_id_av_1_output.png">
+
+### Diffusers
+
+Cosmos3-Edge is supported in the Hugging Face Diffusers library through [`Cosmos3OmniPipeline`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/cosmos3). The pipeline supports image-to-video generation and action-conditioned forward and inverse dynamics.
+
+#### Container
+
+To install Diffusers with `Cosmos3OmniPipeline`:
+
+```bash
+uv venv --python 3.12 --seed --managed-python
+source .venv/bin/activate
+uv pip install \
+  "diffusers @ git+https://github.com/huggingface/diffusers.git" \
+  accelerate \
+  av \
+  cosmos_guardrail \
+  huggingface_hub \
+  imageio \
+  imageio-ffmpeg \
+  torch \
+  torchvision \
+  transformers
+```
+
+#### Examples
+
+The examples use the inputs in this repository's `assets/` directory. Run them from the model repository root.
+
+##### Image to video generation
+
+```python
+import json
+from pathlib import Path
+
+import torch
+from diffusers import Cosmos3OmniPipeline
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+from diffusers.utils import export_to_video, load_image
+
+# Read the JSON-upsampled positive and negative prompts.
+json_prompt = json.loads(Path("assets/example_i2v_prompt.json").read_text())
+negative_prompt = json.loads(Path("assets/negative_prompt.json").read_text())
+
+pipe = Cosmos3OmniPipeline.from_pretrained(
+    "nvidia/Cosmos3-Edge",
+    torch_dtype=torch.bfloat16,
+    enable_safety_checker=True,
+)
+pipe.to("cuda")
+pipe.scheduler = UniPCMultistepScheduler.from_config(
+    pipe.scheduler.config, flow_shift=12.0, use_karras_sigmas=False
+)
+
+result = pipe(
+    prompt=json.dumps(json_prompt),
+    negative_prompt=json.dumps(negative_prompt),
+    image=load_image("assets/example_i2v_input.jpg"),
+    num_frames=121,
+    height=480,
+    width=832,
+    fps=24.0,
+    num_inference_steps=20,
+    guidance_scale=6.0,
+    generator=torch.Generator(device="cuda").manual_seed(0),
+    add_resolution_template=False,
+    add_duration_template=False,
+)
+
+output_path = Path("assets/diffusers_outputs/edge_i2v_diffusers.mp4")
+output_path.parent.mkdir(parents=True, exist_ok=True)
+# macro_block_size=1 preserves the requested 832x480 resolution.
+export_to_video(result.video, str(output_path), fps=24, macro_block_size=1)
+print(f"Saved video to {output_path}")
+```
+
+Example output:
+
+<video controls width="832" height="480" src="https://huggingface.co/nvidia/Cosmos3-Edge/resolve/main/assets/diffusers_outputs/edge_i2v_diffusers.mp4"></video>
+
+---
+
+##### Action generation
+
+The forward-dynamics example uses UMI robotics action trajectories, and the inverse-dynamics examples use autonomous-vehicle (AV) action trajectories. The pipeline's `CosmosActionCondition` groups the action task, embodiment, action chunk, and conditioning image or video. Unlike the image-to-video example, action prompts are plain task descriptions and should not be prompt-upsampled.
+
+###### Action forward dynamics
+
+This two-chunk UMI rollout conditions the first chunk on `assets/example_action_fd_umi_first_frame.png`, then conditions each subsequent chunk on the final generated frame of the previous chunk. Dropping each chunk's conditioning frame produces a stitched 32-frame rollout.
+
+```python
+import json
+from pathlib import Path
+
+import torch
+from diffusers import Cosmos3OmniPipeline, CosmosActionCondition
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+from diffusers.utils import export_to_video, load_image
+
+action_spec = json.loads(Path("assets/example_action_fd_umi_action_chunks.json").read_text())
+output_dir = Path("assets/diffusers_outputs")
+output_dir.mkdir(parents=True, exist_ok=True)
+
+pipe = Cosmos3OmniPipeline.from_pretrained(
+    "nvidia/Cosmos3-Edge",
+    torch_dtype=torch.bfloat16,
+    enable_safety_checker=True,
+)
+pipe.to("cuda")
+pipe.scheduler = UniPCMultistepScheduler.from_config(
+    pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
+)
+
+prompt = action_spec.get("prompt", "mouse arrangement")
+fps = int(action_spec.get("fps", 20))
+chunk_size = int(action_spec.get("action_chunk_size", 16))
+current_frame = load_image("assets/example_action_fd_umi_first_frame.png")
+stitched_frames = []
+chunk_paths = []
+
+for chunk_idx, action_chunk in enumerate(action_spec["action_chunks"]):
+    result = pipe(
+        prompt=prompt,
+        action=CosmosActionCondition(
+            mode="forward_dynamics",
+            chunk_size=chunk_size,
+            domain_name=action_spec.get("domain_name", "umi"),
+            resolution_tier=int(action_spec.get("image_size", 256)),
+            raw_actions=torch.tensor(action_chunk, dtype=torch.float32),
+            image=current_frame,
+            view_point=action_spec.get("view_point", "ego_view"),
+        ),
+        fps=fps,
+        num_inference_steps=30,
+        guidance_scale=1.0,
+        generator=torch.Generator(device="cuda").manual_seed(chunk_idx),
+        use_system_prompt=False,
+    )
+
+    chunk_path = output_dir / f"edge_action_fd_diffusers_chunk_{chunk_idx:02d}.mp4"
+    export_to_video(result.video, str(chunk_path), fps=fps, macro_block_size=1)
+    chunk_paths.append(chunk_path)
+    stitched_frames.extend(result.video[1:])
+    current_frame = result.video[-1]
+
+stitched_path = output_dir / "edge_action_fd_umi_2chunk_diffusers.mp4"
+export_to_video(stitched_frames, str(stitched_path), fps=fps, macro_block_size=1)
+print("Generated chunk videos:", chunk_paths)
+print("Saved stitched rollout:", stitched_path)
+```
+
+Example output:
+
+<video controls width="256" height="256" src="https://huggingface.co/nvidia/Cosmos3-Edge/resolve/main/assets/diffusers_outputs/edge_action_fd_umi_2chunk_diffusers.mp4"></video>
+
+###### Action inverse dynamics
+
+This example predicts the 60-step, 9-D AV action sequence that connects each 61-frame conditioning video. It writes one JSON output per input video.
+
+```python
+import json
+from pathlib import Path
+
+import torch
+from diffusers import Cosmos3OmniPipeline, CosmosActionCondition
+from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
+from diffusers.utils import load_video
+
+input_videos = {
+    "av_inverse_0": Path("assets/example_action_id_av_0_input.mp4"),
+    "av_inverse_1": Path("assets/example_action_id_av_1_input.mp4"),
+}
+output_dir = Path("assets/diffusers_outputs")
+output_dir.mkdir(parents=True, exist_ok=True)
+
+pipe = Cosmos3OmniPipeline.from_pretrained(
+    "nvidia/Cosmos3-Edge",
+    torch_dtype=torch.bfloat16,
+    enable_safety_checker=True,
+)
+pipe.to("cuda")
+pipe.scheduler = UniPCMultistepScheduler.from_config(
+    pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
+)
+
+for name, video_path in input_videos.items():
+    result = pipe(
+        prompt="You are an autonomous vehicle planning system.",
+        action=CosmosActionCondition(
+            mode="inverse_dynamics",
+            chunk_size=60,
+            domain_name="av",
+            resolution_tier=480,
+            video=load_video(str(video_path)),
+            view_point="ego_view",
+        ),
+        fps=10,
+        num_inference_steps=30,
+        guidance_scale=1.0,
+        generator=torch.Generator(device="cuda").manual_seed(0),
+        use_system_prompt=False,
+    )
+    if result.action is None:
+        raise RuntimeError(f"{name} did not return an action tensor")
+
+    action = result.action[0].cpu()
+    output_path = output_dir / f"edge_action_id_{name}_diffusers.json"
+    output_path.write_text(
+        json.dumps(
+            {
+                "data": action.tolist(),
+                "shape": list(action.shape),
+                "dtype": str(action.dtype),
+                "raw_action_dim": 9,
+                "action_mode": "inverse_dynamics",
+                "domain_id": 1,
+            },
+            indent=2,
+        )
+        + "\n"
+    )
+    print(f"Saved predicted action to {output_path}; shape={tuple(action.shape)}")
+```
+
+Example outputs:
+
+- [av_inverse_0 predicted action JSON](https://huggingface.co/nvidia/Cosmos3-Edge/blob/main/assets/diffusers_outputs/edge_action_id_av_inverse_0_diffusers.json)
+- [av_inverse_1 predicted action JSON](https://huggingface.co/nvidia/Cosmos3-Edge/blob/main/assets/diffusers_outputs/edge_action_id_av_inverse_1_diffusers.json)
+
+### vLLM
+
+#### Container
+
+```bash
+docker pull vllm/vllm-openai:cosmos3
+```
+
+#### General Invocation
+
+You can use the `vllm` package to deploy the Cosmos3-Edge reasoner as an OpenAI-compatible API endpoint. The recommended vLLM serving configuration for `nvidia/Cosmos3-Edge` on a single GPU is:
+
+```bash
+vllm serve nvidia/Cosmos3-Edge \
+  --host 0.0.0.0 \
+  --port 8000 \
+  --max-model-len 131072 \
+  --allowed-local-media-path / \
+  --mm-processor-kwargs '{"do_resize": true, "min_pixels": 4096, "max_pixels": 16777216}' \
+  --media-io-kwargs '{"video": {"num_frames": 256}}'
+```
+
+#### Examples
+
+##### Reasoning
+
+Image input:
+
+<img src="assets/example_reasoning_input.png" alt="Reasoning example input showing a tabletop robot-manipulation scene" width="640">
+
+User prompt:
+
+```text
+The task is to put flower into the red bottle. Generate a plan consisting of subtasks for accomplish the task.
+```
+
+```python
+import base64
+import json
+from pathlib import Path
+
+import openai
+
+# 1. Read the image reasoning prompt
+example = json.load(open("assets/example_reasoning_prompt.json"))
+image_path = Path("assets/example_reasoning_input.png").resolve()
+image_url = (
+    "data:image/png;base64," + base64.b64encode(image_path.read_bytes()).decode()
+)
+
+# 2. Query the OpenAI-compatible vLLM server
+client = openai.OpenAI(
+    api_key="EMPTY",
+    base_url="http://localhost:8000/v1",
+)
+
+response = client.chat.completions.create(
+    model=client.models.list().data[0].id,
+    messages=[
+        {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": image_url}},
+                {"type": "text", "text": example["prompt"]},
+            ],
+        },
+    ],
+    max_tokens=example["max_tokens"],
+)
+
+# 3. Print the default response (thinking is enabled)
+print(response.choices[0].message.content)
+```
+
+Example output:
+
+```text
+Got it, let's try to figure out how to put the flower into the red bottle. First, I need to identify the objects: the flower is the green one with the red center, and the red bottle is the red container. So the steps would involve moving the flower to the red bottle.
+
+First, the robot arm needs to locate the flower. Then, grasp the flower. After grasping, lift it, then move it towards the red bottle. Then, position the flower over the red bottle's opening, and finally release it into the bottle. Wait, but maybe I need to check if the red bottle has an opening that can accept the flower. Assuming the red bottle is a container with an opening, so the steps would be: move to flower, grasp, lift, move to red bottle, position over opening, release. Let me make sure each step is clear.
+
+First subtask: Position the robot arm above the flower. Then, lower the gripper to grasp the flower. Then, lift the flower from the table. Then, move the robot arm with the flower towards the red bottle. Then, position the flower over the red bottle's opening. Then, release the flower into the red bottle. That should cover the task.</think>
+To accomplish the task of putting the flower into the red bottle, the plan should follow these subtasks:
+
+1. **Position the robot arm above the flower** – Move the robot arm to a position directly over the green flower with the red center.
+2. **Lower the gripper to grasp the flower** – Lower the gripper to make contact with the flower and close it around the stem.
+3. **Lift the flower from the table** – Once the flower is securely grasped, lift it upward to clear the table surface.
+4. **Move the robot arm with the flower toward the red bottle** – Transport the flower (still held by the gripper) from its current position to the location of the red bottle.
+5. **Position the flower over the red bottle’s opening** – Align the flower so its stem is directly above the opening of the red bottle.
+6. **Release the flower into the red bottle** – Open the gripper to let the flower fall into the red bottle, completing the task.
+
+Each step ensures the flower is picked up, transported, and placed into the red bottle in a controlled manner.
+```
+
+Thinking is enabled by default. To disable it, add the following argument to the `client.chat.completions.create(...)` call:
+
+```python
+extra_body={"chat_template_kwargs": {"enable_thinking": False}}
+```
+
+## Limitations
+
+Cosmos3 may produce imperfect outputs in challenging scenarios. Generation artifacts include temporal inconsistency, unstable camera or object motion, imprecise physical interactions, 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:** [vLLM-Omni](https://github.com/vllm-project/vllm-omni), [vLLM](https://github.com/vllm-project/vllm), [PyTorch](https://pytorch.org/)
+
+**Test Hardware:** B200, H100, H20, RTX PRO 6000, DGX Station, DGX Spark, Jetson Thor, Jetson AGX Orin
+
+## 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. |

assets/diffusers_outputs/edge_action_id_av_inverse_0_diffusers.json

@@ -0,0 +1,672 @@
+{
+  "data": [
+    [
+      -0.013148754835128784,
+      -0.007273353636264801,
+      0.008258908987045288,
+      0.9254555702209473,
+      -0.018639683723449707,
+      0.004321083426475525,
+      -0.00694487988948822,
+      0.9603041410446167,
+      0.007375245913863182
+    ],
+    [
+      0.0006928369402885437,
+      0.00114353746175766,
+      0.04833635687828064,
+      0.9550113677978516,
+      0.007736995816230774,
+      -0.0007977494969964027,
+      -0.001905590295791626,
+      0.9611708521842957,
+      -0.0071649253368377686
+    ],
+    [
+      0.00024224817752838135,
+      -0.00632748007774353,
+      0.01743471622467041,
+      0.9503518342971802,
+      0.006431609392166138,
+      -0.0021766857244074345,
+      0.004798382520675659,
+      0.9540749192237854,
+      -0.002829238772392273
+    ],
+    [
+      -0.003921394236385822,
+      0.008617669343948364,
+      -0.012175090610980988,
+      0.9608586430549622,
+      -0.008618205785751343,
+      -0.008708074688911438,
+      -0.013742625713348389,
+      0.9400895833969116,
+      0.006977975368499756
+    ],
+    [
+      0.005376532673835754,
+      0.003240346908569336,
+      0.018488049507141113,
+      0.9523720145225525,
+      0.01081898808479309,
+      -0.00035040080547332764,
+      0.00804641842842102,
+      0.9556390643119812,
+      0.012530684471130371
+    ],
+    [
+      0.005463764071464539,
+      0.0012700557708740234,
+      -0.012111283838748932,
+      0.9668004512786865,
+      0.009593963623046875,
+      0.0014052540063858032,
+      0.00772630050778389,
+      0.935010552406311,
+      0.011620238423347473
+    ],
+    [
+      -0.003778284415602684,
+      0.012686282396316528,
+      0.012819170951843262,
+      0.9748083353042603,
+      -0.015680819749832153,
+      -0.0014512911438941956,
+      -0.004034742712974548,
+      0.9654931426048279,
+      0.0029178932309150696
+    ],
+    [
+      0.009581685066223145,
+      -0.0018721558153629303,
+      0.015696480870246887,
+      0.9498556852340698,
+      -0.007028155028820038,
+      -0.003986746072769165,
+      -0.0010101348161697388,
+      0.9363378286361694,
+      0.0034978240728378296
+    ],
+    [
+      0.011020243167877197,
+      0.015359997749328613,
+      0.0042057037353515625,
+      0.9537041783332825,
+      -0.008601069450378418,
+      -0.0073888301849365234,
+      0.0013249143958091736,
+      0.94188392162323,
+      0.0001823045313358307
+    ],
+    [
+      0.008615091443061829,
+      0.005651026964187622,
+      -0.009945794939994812,
+      0.9452113509178162,
+      0.019210517406463623,
+      0.004503607749938965,
+      0.015948235988616943,
+      0.9488579630851746,
+      0.008347108960151672
+    ],
+    [
+      0.001062244176864624,
+      -0.01891157031059265,
+      -0.0280563086271286,
+      0.9730226993560791,
+      0.011333107948303223,
+      -0.000612400472164154,
+      -0.004964202642440796,
+      0.9605361819267273,
+      0.00998196005821228
+    ],
+    [
+      0.0016751736402511597,
+      -0.024208784103393555,
+      -0.024219349026679993,
+      0.9559570550918579,
+      -0.009602159261703491,
+      -0.004984036087989807,
+      0.0061578452587127686,
+      0.950946033000946,
+      0.008079171180725098
+    ],
+    [
+      -0.011940449476242065,
+      -0.0034200847148895264,
+      0.0017094612121582031,
+      0.9588748812675476,
+      0.006712079048156738,
+      -0.004783693701028824,
+      -0.00014981627464294434,
+      0.9792852997779846,
+      -0.00029080547392368317
+    ],
+    [
+      0.0016713226214051247,
+      -0.0047683920711278915,
+      0.0011281678453087807,
+      0.9475633502006531,
+      0.007233217358589172,
+      -9.954720735549927e-05,
+      -0.0002157837152481079,
+      0.9581102132797241,
+      -0.001299828290939331
+    ],
+    [
+      0.006934747099876404,
+      0.008136779069900513,
+      -0.039625704288482666,
+      0.9607610106468201,
+      0.010957300662994385,
+      0.0001928955316543579,
+      0.004511512815952301,
+      0.9533595442771912,
+      0.011611387133598328
+    ],
+    [
+      0.007694117724895477,
+      -0.01480981707572937,
+      0.01341133564710617,
+      0.9478503465652466,
+      0.0021869465708732605,
+      -0.0011091921478509903,
+      0.003402979113161564,
+      0.9389904737472534,
+      -0.0016643255949020386
+    ],
+    [
+      -0.006875067949295044,
+      0.0049460530281066895,
+      0.0421476811170578,
+      0.9672802090644836,
+      -0.01006162166595459,
+      0.0008940547704696655,
+      0.007537230849266052,
+      0.9556049704551697,
+      0.0027959682047367096
+    ],
+    [
+      -0.0024900436401367188,
+      -0.005135815590620041,
+      0.05198696255683899,
+      0.9647979140281677,
+      0.01649463176727295,
+      0.003368660807609558,
+      -0.019760966300964355,
+      0.9663780331611633,
+      -0.004725992679595947
+    ],
+    [
+      -0.0026149116456508636,
+      -0.007602095603942871,
+      0.03872739523649216,
+      0.9427680969238281,
+      -0.0020749270915985107,
+      -0.004558362066745758,
+      -0.011337906122207642,
+      0.9802517890930176,
+      0.0020113736391067505
+    ],
+    [
+      -0.0024757981300354004,
+      -0.007645174860954285,
+      0.07559360563755035,
+      0.9630593061447144,
+      -0.0025666803121566772,
+      -0.00292031466960907,
+      -0.012108653783798218,
+      0.9655929803848267,
+      0.0013676062226295471
+    ],
+    [
+      0.0007936935871839523,
+      0.0024557560682296753,
+      0.036987900733947754,
+      0.9483445882797241,
+      0.006975963711738586,
+      0.0030634403228759766,
+      -0.0004689544439315796,
+      0.9642810225486755,
+      -0.001054992899298668
+    ],
+    [
+      -0.003186628222465515,
+      0.00013031065464019775,
+      0.06438089907169342,
+      0.967340350151062,
+      -0.008410289883613586,
+      -0.007106497883796692,
+      0.010625198483467102,
+      0.9285515546798706,
+      -0.0012674331665039062
+    ],
+    [
+      -0.02860349416732788,
+      -8.131563663482666e-05,
+      0.08170737326145172,
+      0.9743189811706543,
+      0.004199020564556122,
+      0.006051763892173767,
+      0.004216104745864868,
+      0.9849252700805664,
+      -0.01118806004524231
+    ],
+    [
+      0.00029268115758895874,
+      -0.0008304938673973083,
+      0.08716102689504623,
+      0.9572217464447021,
+      0.00250987708568573,
+      -0.0032331272959709167,
+      -0.004110082983970642,
+      0.9600695371627808,
+      -0.0021736174821853638
+    ],
+    [
+      0.016955643892288208,
+      0.0030432939529418945,
+      0.04992994666099548,
+      0.9569373726844788,
+      -0.011746704578399658,
+      -0.0003188829869031906,
+      -0.007667466998100281,
+      0.9866983890533447,
+      0.004928670823574066
+    ],
+    [
+      -0.0039969682693481445,
+      0.00038054585456848145,
+      0.12631714344024658,
+      0.9718866348266602,
+      0.018440932035446167,
+      -0.005601778626441956,
+      -0.009917289018630981,
+      0.9650411605834961,
+      0.005232140421867371
+    ],
+    [
+      -0.0012138709425926208,
+      -0.011934876441955566,
+      0.11297653615474701,
+      0.9682045578956604,
+      0.0012573739513754845,
+      0.0006018783897161484,
+      0.009243190288543701,
+      0.9470562934875488,
+      0.012865453958511353
+    ],
+    [
+      -0.009654894471168518,
+      0.004492349922657013,
+      0.14582459628582,
+      0.941580057144165,
+      -0.027970850467681885,
+      -0.01605433225631714,
+      -0.003179602324962616,
+      0.9500089883804321,
+      -0.005389504134654999
+    ],
+    [
+      0.014612823724746704,
+      -0.004633047617971897,
+      0.14664553105831146,
+      0.9637702703475952,
+      -0.0011444943957030773,
+      -0.0005822107195854187,
+      0.005188122391700745,
+      0.9636419415473938,
+      -0.007603704929351807
+    ],
+    [
+      -0.014512777328491211,
+      -0.00830838829278946,
+      0.14723633229732513,
+      0.9575201272964478,
+      0.009052962064743042,
+      0.002744898200035095,
+      0.00654207170009613,
+      0.9790866374969482,
+      -0.006706550717353821
+    ],
+    [
+      -0.01381877064704895,
+      0.013246327638626099,
+      0.1438756138086319,
+      0.9541648626327515,
+      0.0031017810106277466,
+      -0.010317087173461914,
+      0.004141256213188171,
+      0.9406745433807373,
+      -0.0032720565795898438
+    ],
+    [
+      0.0049401819705963135,
+      0.016335099935531616,
+      0.1370936781167984,
+      0.9569792747497559,
+      0.009543389081954956,
+      -0.0038724213372915983,
+      -0.011453330516815186,
+      0.9742096066474915,
+      -0.013819366693496704
+    ],
+    [
+      0.016965091228485107,
+      -0.014787629246711731,
+      0.18282562494277954,
+      0.952957034111023,
+      0.00769689679145813,
+      -0.006391674280166626,
+      0.00217418372631073,
+      0.9660205841064453,
+      0.006829023361206055
+    ],
+    [
+      0.00011231936514377594,
+      0.022837340831756592,
+      0.21823179721832275,
+      0.9779958128929138,
+      0.0016702637076377869,
+      -0.01641136407852173,
+      0.008365646004676819,
+      0.9648821949958801,
+      -0.015010237693786621
+    ],
+    [
+      -0.003881126642227173,
+      0.008091509342193604,
+      0.16629643738269806,
+      0.9440470933914185,
+      0.017186373472213745,
+      -0.0032056421041488647,
+      0.015325725078582764,
+      0.9715288877487183,
+      0.005628049373626709
+    ],
+    [
+      0.015618979930877686,
+      -0.016067102551460266,
+      0.1358950436115265,
+      0.9455870985984802,
+      0.015390992164611816,
+      0.0032478421926498413,
+      0.007700107991695404,
+      0.949781060218811,
+      0.00139538012444973
+    ],
+    [
+      0.003707513213157654,
+      0.0012879809364676476,
+      0.18519513309001923,
+      0.9606454968452454,
+      -0.016481131315231323,
+      -0.0009821802377700806,
+      -0.0053266435861587524,
+      0.9794431924819946,
+      0.002626672387123108
+    ],
+    [
+      0.016481012105941772,
+      -0.017973721027374268,
+      0.14953568577766418,
+      0.9501946568489075,
+      0.013054758310317993,
+      -0.0001233518123626709,
+      0.002798639237880707,
+      0.963482141494751,
+      0.0029630735516548157
+    ],
+    [
+      0.027081698179244995,
+      -0.008018411695957184,
+      0.19762511551380157,
+      0.9853548407554626,
+      0.007277920842170715,
+      0.0018384754657745361,
+      -0.0006946176290512085,
+      0.9530956149101257,
+      -0.00409291684627533
+    ],
+    [
+      0.009928137063980103,
+      -0.01654413342475891,
+      0.18209686875343323,
+      0.9711376428604126,
+      0.01371067762374878,
+      -0.01334737241268158,
+      -0.011193245649337769,
+      0.9716130495071411,
+      0.005922876298427582
+    ],
+    [
+      0.012314438819885254,
+      -0.0027115456759929657,
+      0.18243494629859924,
+      0.9741570949554443,
+      0.008049450814723969,
+      -0.004693128168582916,
+      -0.003039032220840454,
+      0.9581626057624817,
+      0.012616187334060669
+    ],
+    [
+      0.0076538920402526855,
+      -0.003307037055492401,
+      0.20810538530349731,
+      0.948331892490387,
+      -0.015447914600372314,
+      -0.01176406443119049,
+      -0.0022220220416784286,
+      0.9713965058326721,
+      -0.0014124363660812378
+    ],
+    [
+      0.022586889564990997,
+      -0.008036494255065918,
+      0.21718934178352356,
+      0.9554579257965088,
+      0.028451621532440186,
+      -0.008404174819588661,
+      -0.0012748166918754578,
+      0.9658167362213135,
+      -0.0026244521141052246
+    ],
+    [
+      0.015891581773757935,
+      0.0012519285082817078,
+      0.2355501353740692,
+      0.9612299799919128,
+      -0.005295991897583008,
+      0.00039255619049072266,
+      -0.00497812032699585,
+      0.9328453540802002,
+      -0.002513475716114044
+    ],
+    [
+      0.024546697735786438,
+      0.004739120602607727,
+      0.24198676645755768,
+      0.9556739926338196,
+      0.001108948141336441,
+      -0.010602131485939026,
+      -0.0015903003513813019,
+      0.9675569534301758,
+      0.004108220338821411
+    ],
+    [
+      0.02232934534549713,
+      0.018866419792175293,
+      0.24329781532287598,
+      0.9607975482940674,
+      -0.008478760719299316,
+      -0.00401577353477478,
+      0.005954533815383911,
+      0.961711585521698,
+      -0.008179426193237305
+    ],
+    [
+      0.014459583908319473,
+      0.013193011283874512,
+      0.2807546854019165,
+      0.9628647565841675,
+      0.02138429880142212,
+      -0.009663410484790802,
+      0.012736231088638306,
+      0.9604622721672058,
+      0.004635050892829895
+    ],
+    [
+      0.03173169493675232,
+      -0.008499197661876678,
+      0.2548338770866394,
+      0.9557382464408875,
+      0.0049033742398023605,
+      0.006540834903717041,
+      0.013527542352676392,
+      0.9523492455482483,
+      0.0032977014780044556
+    ],
+    [
+      0.03362193703651428,
+      0.005373731255531311,
+      0.28878676891326904,
+      0.9541832804679871,
+      0.00421798974275589,
+      -0.022053956985473633,
+      -0.007141679525375366,
+      0.9740029573440552,
+      0.002051055431365967
+    ],
+    [
+      0.0156448632478714,
+      0.007966414093971252,
+      0.26515382528305054,
+      0.9523801803588867,
+      0.0037920400500297546,
+      -0.011845819652080536,
+      -0.005416005849838257,
+      0.9764019846916199,
+      0.0025727003812789917
+    ],
+    [
+      0.025211334228515625,
+      -0.0038858801126480103,
+      0.3200596272945404,
+      0.9874192476272583,
+      0.01006472110748291,
+      -0.005867555737495422,
+      0.0030220672488212585,
+      0.950593113899231,
+      -0.002622779458761215
+    ],
+    [
+      0.016667742282152176,
+      0.016612201929092407,
+      0.3063710629940033,
+      0.9813263416290283,
+      0.010637477040290833,
+      -0.0004864633083343506,
+      0.009502381086349487,
+      0.9399935007095337,
+      -0.0036408603191375732
+    ],
+    [
+      0.03205466270446777,
+      -0.005740225315093994,
+      0.3219253718852997,
+      0.9442741870880127,
+      -0.0003297049552202225,
+      -0.007991593331098557,
+      -0.012274682521820068,
+      0.9457107782363892,
+      0.001197025179862976
+    ],
+    [
+      0.019879944622516632,
+      -0.00796326994895935,
+      0.34765785932540894,
+      0.9492185711860657,
+      0.008849918842315674,
+      -0.003306932747364044,
+      0.005508631467819214,
+      0.9591333866119385,
+      0.008243441581726074
+    ],
+    [
+      0.015053778886795044,
+      -0.00721488893032074,
+      0.3408358097076416,
+      0.9282618761062622,
+      0.002898307517170906,
+      -0.00362338125705719,
+      0.027731001377105713,
+      0.9604060053825378,
+      -0.00824989378452301
+    ],
+    [
+      0.015081822872161865,
+      -0.006049111485481262,
+      0.35627633333206177,
+      0.9522294998168945,
+      0.014767050743103027,
+      -0.013559997081756592,
+      -0.0005444735288619995,
+      0.9336357116699219,
+      -0.008866548538208008
+    ],
+    [
+      0.012881480157375336,
+      -0.008715510368347168,
+      0.35815778374671936,
+      0.9767504334449768,
+      0.011995196342468262,
+      -0.0015369206666946411,
+      -0.00036385655403137207,
+      0.9670126438140869,
+      0.006960272789001465
+    ],
+    [
+      0.022443406283855438,
+      -0.002777688205242157,
+      0.3705684542655945,
+      0.9601130485534668,
+      -0.010128200054168701,
+      -0.00759703665971756,
+      0.0009619146585464478,
+      0.968835711479187,
+      -0.006335556507110596
+    ],
+    [
+      0.0191623717546463,
+      -0.003963552415370941,
+      0.3500891327857971,
+      0.9584890007972717,
+      -0.0029276683926582336,
+      -0.015166878700256348,
+      -0.002148374915122986,
+      0.9575344920158386,
+      0.010722845792770386
+    ],
+    [
+      0.006769150495529175,
+      -0.0050390660762786865,
+      0.4090595245361328,
+      0.9641549587249756,
+      -0.002060607075691223,
+      -0.015374422073364258,
+      0.005776405334472656,
+      0.9640921354293823,
+      0.0018393322825431824
+    ]
+  ],
+  "shape": [
+    60,
+    9
+  ],
+  "dtype": "torch.float32",
+  "raw_action_dim": 9,
+  "action_mode": "inverse_dynamics",
+  "domain_id": 1
+}

assets/diffusers_outputs/edge_action_id_av_inverse_1_diffusers.json

@@ -0,0 +1,672 @@
+{
+  "data": [
+    [
+      0.059141844511032104,
+      -0.003840334713459015,
+      1.0373594760894775,
+      0.9279366731643677,
+      -0.01645451784133911,
+      -0.020578712224960327,
+      -0.010958120226860046,
+      0.9640744924545288,
+      0.0024261921644210815
+    ],
+    [
+      0.07163303345441818,
+      0.002209894359111786,
+      1.0704994201660156,
+      0.9624326229095459,
+      0.011429950594902039,
+      -0.02556397207081318,
+      -0.005203992128372192,
+      0.968371570110321,
+      -0.005390077829360962
+    ],
+    [
+      0.06948350369930267,
+      -0.014859065413475037,
+      1.0249040126800537,
+      0.9573811292648315,
+      0.007525742053985596,
+      -0.027830475941300392,
+      0.006637901067733765,
+      0.9581300616264343,
+      0.00023680925369262695
+    ],
+    [
+      0.07753375917673111,
+      0.004523158073425293,
+      0.9877839684486389,
+      0.9649670124053955,
+      -0.006604701280593872,
+      -0.037096068263053894,
+      -0.021954774856567383,
+      0.9435235261917114,
+      0.013597071170806885
+    ],
+    [
+      0.09455021470785141,
+      -0.0016047954559326172,
+      1.0100252628326416,
+      0.9584341645240784,
+      0.014761030673980713,
+      -0.029446259140968323,
+      0.00870533287525177,
+      0.968270480632782,
+      0.0247458815574646
+    ],
+    [
+      0.10178711265325546,
+      -0.0029810070991516113,
+      0.9616026878356934,
+      0.9771547913551331,
+      0.01504051685333252,
+      -0.031167365610599518,
+      0.0053897276520729065,
+      0.9344217777252197,
+      0.012324824929237366
+    ],
+    [
+      0.10360411554574966,
+      0.01793801784515381,
+      0.9839838743209839,
+      0.986229658126831,
+      -0.011833727359771729,
+      -0.038624923676252365,
+      -0.007526323199272156,
+      0.9758810997009277,
+      0.005369395017623901
+    ],
+    [
+      0.12239452451467514,
+      0.010958259925246239,
+      0.9665457010269165,
+      0.9559676647186279,
+      -0.0019251927733421326,
+      -0.039464741945266724,
+      -0.005040377378463745,
+      0.9393422603607178,
+      0.00036060065031051636
+    ],
+    [
+      0.13781818747520447,
+      0.04306331276893616,
+      0.9548782110214233,
+      0.9634062051773071,
+      -0.0013715922832489014,
+      -0.045522540807724,
+      -0.0034839287400245667,
+      0.9445981979370117,
+      -0.0009270161390304565
+    ],
+    [
+      0.13575460016727448,
+      0.026487261056900024,
+      0.9226974844932556,
+      0.9517682790756226,
+      0.024003714323043823,
+      -0.03737533092498779,
+      0.013349950313568115,
+      0.9547375440597534,
+      0.00824420154094696
+    ],
+    [
+      0.12451967597007751,
+      -0.0071739256381988525,
+      0.9000017642974854,
+      0.9793621301651001,
+      0.011195510625839233,
+      -0.04407672584056854,
+      -0.00877419114112854,
+      0.9616803526878357,
+      0.011973768472671509
+    ],
+    [
+      0.12706036865711212,
+      -0.014932453632354736,
+      0.9002858400344849,
+      0.9622393250465393,
+      -0.00951644778251648,
+      -0.04734916239976883,
+      0.003892853856086731,
+      0.9567041397094727,
+      0.013047799468040466
+    ],
+    [
+      0.11344760656356812,
+      0.01170305535197258,
+      0.9190186858177185,
+      0.9637442231178284,
+      0.011624082922935486,
+      -0.04697572439908981,
+      -0.0024839788675308228,
+      0.9862824082374573,
+      0.0046536242589354515
+    ],
+    [
+      0.1317308247089386,
+      0.01151534914970398,
+      0.9060279726982117,
+      0.950132429599762,
+      0.013150960206985474,
+      -0.042732544243335724,
+      -0.006613433361053467,
+      0.9610092043876648,
+      0.0009927898645401
+    ],
+    [
+      0.13878925144672394,
+      0.027730464935302734,
+      0.8691807985305786,
+      0.9631220698356628,
+      0.016683101654052734,
+      -0.038141220808029175,
+      -0.002742782235145569,
+      0.9560869932174683,
+      0.009326145052909851
+    ],
+    [
+      0.13446137309074402,
+      0.0007270127534866333,
+      0.9065344333648682,
+      0.9507459402084351,
+      0.007336311042308807,
+      -0.04232313856482506,
+      -0.001765731256455183,
+      0.9377214908599854,
+      -0.004297852516174316
+    ],
+    [
+      0.11476296186447144,
+      0.02131199836730957,
+      0.9385601878166199,
+      0.9792501926422119,
+      -0.0066051185131073,
+      -0.04161176085472107,
+      0.0019875168800354004,
+      0.9670274257659912,
+      0.005515553057193756
+    ],
+    [
+      0.12210288643836975,
+      0.014036169275641441,
+      0.9407962560653687,
+      0.9701698422431946,
+      0.020105689764022827,
+      -0.040228381752967834,
+      -0.020383119583129883,
+      0.9672138094902039,
+      0.0005619972944259644
+    ],
+    [
+      0.12542369961738586,
+      0.012022823095321655,
+      0.9342747926712036,
+      0.9483041763305664,
+      0.0013997424393892288,
+      -0.046562936156988144,
+      -0.02045336365699768,
+      0.9865255355834961,
+      0.00695318728685379
+    ],
+    [
+      0.11753255128860474,
+      0.016480877995491028,
+      0.977397084236145,
+      0.9714816212654114,
+      0.0014084279537200928,
+      -0.04178629815578461,
+      -0.01719522476196289,
+      0.9729042649269104,
+      0.004064872860908508
+    ],
+    [
+      0.12327518314123154,
+      0.026617400348186493,
+      0.9441711902618408,
+      0.9547913670539856,
+      0.009748861193656921,
+      -0.03178030252456665,
+      -0.005975648760795593,
+      0.9724594950675964,
+      0.0015931297093629837
+    ],
+    [
+      0.11693878471851349,
+      0.021891258656978607,
+      0.9779582023620605,
+      0.9773543477058411,
+      -0.0021515488624572754,
+      -0.043881744146347046,
+      0.005668759346008301,
+      0.9312025308609009,
+      -0.0018804669380187988
+    ],
+    [
+      0.09713619947433472,
+      0.022913970053195953,
+      0.9865879416465759,
+      0.984277606010437,
+      0.008582152426242828,
+      -0.035678133368492126,
+      -0.003005385398864746,
+      0.9906625747680664,
+      -0.010433465242385864
+    ],
+    [
+      0.11487491428852081,
+      0.02062314748764038,
+      1.004561185836792,
+      0.9658318161964417,
+      0.00811738520860672,
+      -0.0393141433596611,
+      -0.008064255118370056,
+      0.9643077850341797,
+      -0.00047591328620910645
+    ],
+    [
+      0.12689977884292603,
+      0.02446255087852478,
+      0.9813741445541382,
+      0.9599205255508423,
+      -0.007350906729698181,
+      -0.03434314206242561,
+      -0.0102815181016922,
+      0.99027019739151,
+      0.002850934863090515
+    ],
+    [
+      0.10348488390445709,
+      0.01788690686225891,
+      1.0414836406707764,
+      0.9772853851318359,
+      0.019330114126205444,
+      -0.041947878897190094,
+      -0.013509660959243774,
+      0.9699065089225769,
+      0.005952954292297363
+    ],
+    [
+      0.10493399202823639,
+      0.010198116302490234,
+      1.0407824516296387,
+      0.9719204902648926,
+      0.0058035897091031075,
+      -0.03458989039063454,
+      0.003740936517715454,
+      0.9522305130958557,
+      0.012766420841217041
+    ],
+    [
+      0.09770180284976959,
+      0.024997010827064514,
+      1.0794150829315186,
+      0.9405001401901245,
+      -0.02436143159866333,
+      -0.04618346691131592,
+      -0.008885281160473824,
+      0.9424073696136475,
+      -0.0033924728631973267
+    ],
+    [
+      0.11438962817192078,
+      0.01623646914958954,
+      1.0927131175994873,
+      0.9766831398010254,
+      0.0005866186693310738,
+      -0.030141858384013176,
+      0.0021623969078063965,
+      0.9683176875114441,
+      -0.0047433823347091675
+    ],
+    [
+      0.07521477341651917,
+      0.008214041590690613,
+      1.098933219909668,
+      0.9660043120384216,
+      0.007547542452812195,
+      -0.025479033589363098,
+      0.004106804728507996,
+      0.9826686978340149,
+      -0.001127287745475769
+    ],
+    [
+      0.07278415560722351,
+      0.0321023166179657,
+      1.1175589561462402,
+      0.9674540758132935,
+      0.005659177899360657,
+      -0.038438498973846436,
+      -0.0007940381765365601,
+      0.9435781240463257,
+      -0.0032588541507720947
+    ],
+    [
+      0.08530683070421219,
+      0.02760574221611023,
+      1.123960018157959,
+      0.9651504158973694,
+      0.010333657264709473,
+      -0.02847376838326454,
+      -0.014152109622955322,
+      0.9765793681144714,
+      -0.007708162069320679
+    ],
+    [
+      0.08623477816581726,
+      -2.8833746910095215e-05,
+      1.1860765218734741,
+      0.9632298350334167,
+      0.009518980979919434,
+      -0.028353288769721985,
+      0.0003351382911205292,
+      0.9702121019363403,
+      0.011311382055282593
+    ],
+    [
+      0.06764379888772964,
+      0.028459995985031128,
+      1.228102207183838,
+      0.9854918718338013,
+      0.004853567108511925,
+      -0.030850589275360107,
+      0.005490615963935852,
+      0.9692810773849487,
+      -0.004521816968917847
+    ],
+    [
+      0.05355161428451538,
+      0.020747363567352295,
+      1.2010810375213623,
+      0.951269805431366,
+      0.016863375902175903,
+      -0.017977416515350342,
+      0.01373264193534851,
+      0.9801374077796936,
+      0.014254122972488403
+    ],
+    [
+      0.07317647337913513,
+      0.0030185580253601074,
+      1.18898344039917,
+      0.9577667713165283,
+      0.01679345965385437,
+      -0.010508865118026733,
+      0.006877075880765915,
+      0.9590687155723572,
+      0.005443556234240532
+    ],
+    [
+      0.04951860010623932,
+      0.01835135743021965,
+      1.2524373531341553,
+      0.9677692651748657,
+      -0.012610375881195068,
+      -0.01456555724143982,
+      -0.005562201142311096,
+      0.9804211854934692,
+      0.005360782146453857
+    ],
+    [
+      0.057762667536735535,
+      -0.0012004673480987549,
+      1.2260611057281494,
+      0.9578285217285156,
+      0.0164794921875,
+      -0.00892224907875061,
+      -0.002040460705757141,
+      0.9750438332557678,
+      0.010551542043685913
+    ],
+    [
+      0.0635719895362854,
+      0.009304247796535492,
+      1.2985970973968506,
+      0.997957170009613,
+      0.005984649062156677,
+      -0.006568968296051025,
+      -0.005957365036010742,
+      0.9624383449554443,
+      0.0034569650888442993
+    ],
+    [
+      0.030685484409332275,
+      -0.0031789839267730713,
+      1.2745897769927979,
+      0.9833477735519409,
+      0.015771299600601196,
+      -0.02075822651386261,
+      -0.018386155366897583,
+      0.9756336212158203,
+      0.011255372315645218
+    ],
+    [
+      0.023339152336120605,
+      0.010662931948900223,
+      1.300084114074707,
+      0.989835262298584,
+      0.006620235741138458,
+      -0.0073957666754722595,
+      -0.006816834211349487,
+      0.9638521075248718,
+      0.01855221390724182
+    ],
+    [
+      0.012639999389648438,
+      0.007987678050994873,
+      1.3306738138198853,
+      0.9491333961486816,
+      -0.020535647869110107,
+      -0.01028469204902649,
+      -0.002563953399658203,
+      0.9858134388923645,
+      0.0020048022270202637
+    ],
+    [
+      0.02348978817462921,
+      -0.0019411742687225342,
+      1.3654818534851074,
+      0.969693124294281,
+      0.027309060096740723,
+      -0.004834961146116257,
+      -0.0036605969071388245,
+      0.9743373990058899,
+      0.003153502941131592
+    ],
+    [
+      0.013904020190238953,
+      0.013629376888275146,
+      1.394413948059082,
+      0.972601592540741,
+      -0.00408954918384552,
+      0.004027083516120911,
+      -0.008378058671951294,
+      0.9446616172790527,
+      0.0034879446029663086
+    ],
+    [
+      0.02238820493221283,
+      0.01910078525543213,
+      1.4235327243804932,
+      0.968029797077179,
+      0.004553822800517082,
+      -0.0062204450368881226,
+      -0.003272419795393944,
+      0.977702260017395,
+      0.00832560658454895
+    ],
+    [
+      0.02112443745136261,
+      0.04720121622085571,
+      1.4362165927886963,
+      0.9765746593475342,
+      -0.013532817363739014,
+      0.001711629331111908,
+      0.007163316011428833,
+      0.9686256647109985,
+      -0.010016411542892456
+    ],
+    [
+      0.007742851972579956,
+      0.03149166703224182,
+      1.4930720329284668,
+      0.9732664823532104,
+      0.01966017484664917,
+      -0.0022897571325302124,
+      0.011702001094818115,
+      0.96928870677948,
+      0.006121553480625153
+    ],
+    [
+      0.022700637578964233,
+      0.010335221886634827,
+      1.4739594459533691,
+      0.9702088832855225,
+      0.0024269153364002705,
+      0.01349705457687378,
+      0.013742953538894653,
+      0.9670001864433289,
+      0.006703421473503113
+    ],
+    [
+      0.02724406123161316,
+      0.02931670844554901,
+      1.5207319259643555,
+      0.9660757184028625,
+      0.004963122308254242,
+      -0.017152994871139526,
+      -0.009395450353622437,
+      0.9865334630012512,
+      0.009139366447925568
+    ],
+    [
+      0.008380383253097534,
+      0.024726346135139465,
+      1.4987715482711792,
+      0.9654471278190613,
+      0.0028583072125911713,
+      -0.005337700247764587,
+      -0.0050354450941085815,
+      0.9871041774749756,
+      0.006967529654502869
+    ],
+    [
+      0.01739254593849182,
+      0.012995205819606781,
+      1.5692707300186157,
+      1.0089263916015625,
+      0.013926923274993896,
+      0.0009343177080154419,
+      0.004500538110733032,
+      0.9502257108688354,
+      0.006055116653442383
+    ],
+    [
+      0.009858133271336555,
+      0.03461769223213196,
+      1.5677274465560913,
+      0.9968569278717041,
+      0.011664733290672302,
+      0.004488259553909302,
+      0.00662684440612793,
+      0.9470456838607788,
+      0.0005231946706771851
+    ],
+    [
+      0.02299785614013672,
+      0.009481832385063171,
+      1.5700998306274414,
+      0.9577478766441345,
+      0.00069449283182621,
+      -0.0011343881487846375,
+      -0.016278564929962158,
+      0.958492636680603,
+      0.006463430821895599
+    ],
+    [
+      0.005318776238709688,
+      0.008516699075698853,
+      1.615574836730957,
+      0.9617098569869995,
+      0.007522448897361755,
+      0.0025394558906555176,
+      0.005760625004768372,
+      0.9720036387443542,
+      0.015061944723129272
+    ],
+    [
+      0.007067561149597168,
+      0.0065644532442092896,
+      1.6096922159194946,
+      0.9413715600967407,
+      -0.0019346103072166443,
+      0.002095624804496765,
+      0.02600991725921631,
+      0.974121630191803,
+      -0.006266683340072632
+    ],
+    [
+      0.007169589400291443,
+      0.009554922580718994,
+      1.6439979076385498,
+      0.964790403842926,
+      0.014039307832717896,
+      -0.006732136011123657,
+      0.00041218847036361694,
+      0.9473226070404053,
+      -0.01044127345085144
+    ],
+    [
+      0.007147058844566345,
+      0.00785742700099945,
+      1.644051194190979,
+      0.985714852809906,
+      0.0075702965259552,
+      0.0006187558174133301,
+      -0.000648103654384613,
+      0.9734295010566711,
+      0.009049922227859497
+    ],
+    [
+      0.015374734997749329,
+      0.014419212937355042,
+      1.6574026346206665,
+      0.9712061285972595,
+      -0.01914668083190918,
+      -0.0018936395645141602,
+      -0.0007403641939163208,
+      0.9860180616378784,
+      -0.006495922803878784
+    ],
+    [
+      0.013697490096092224,
+      0.013768978416919708,
+      1.642235517501831,
+      0.9747515916824341,
+      -0.002270542085170746,
+      -0.011506795883178711,
+      -0.006279811263084412,
+      0.9668055772781372,
+      0.013116374611854553
+    ],
+    [
+      -0.0011235475540161133,
+      0.013980269432067871,
+      1.7224175930023193,
+      0.9761761426925659,
+      -0.0001882612705230713,
+      -0.008678078651428223,
+      0.002798914909362793,
+      0.9733132719993591,
+      0.00561833381652832
+    ]
+  ],
+  "shape": [
+    60,
+    9
+  ],
+  "dtype": "torch.float32",
+  "raw_action_dim": 9,
+  "action_mode": "inverse_dynamics",
+  "domain_id": 1
+}

assets/diffusers_outputs/edge_i2v_diffusers.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2fe93251039d8640625f660f42514d8fddd0812ae6918a4ea3cf5fc2363e4a16
+size 1278049

assets/edge_action_fd_umi_2chunk_output.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10037bf0c2a20c46f24e616d7b822932babb45f72113833f320e2bb524c8c57b
+size 32755

assets/edge_action_id_av_0_output.json

@@ -0,0 +1,672 @@
+{
+  "data": [
+    [
+      0.0107421875,
+      -0.00048828125,
+      0.0859375,
+      0.95703125,
+      0.0025634765625,
+      0.00390625,
+      0.0087890625,
+      0.9609375,
+      0.00390625
+    ],
+    [
+      -0.004638671875,
+      0.005859375,
+      -0.0146484375,
+      0.96875,
+      0.01171875,
+      -0.0009765625,
+      -0.0029296875,
+      0.9609375,
+      0.0068359375
+    ],
+    [
+      0.0078125,
+      0.0009765625,
+      0.0048828125,
+      0.921875,
+      0.015625,
+      -0.013671875,
+      0.0234375,
+      0.9765625,
+      -0.00341796875
+    ],
+    [
+      -0.003173828125,
+      0.005859375,
+      0.0078125,
+      0.9609375,
+      -0.0029296875,
+      -0.000244140625,
+      0.009765625,
+      0.9609375,
+      0.0029296875
+    ],
+    [
+      -0.005615234375,
+      -0.01953125,
+      -0.00927734375,
+      0.9609375,
+      0.005859375,
+      -0.009765625,
+      -0.0087890625,
+      0.9609375,
+      -0.001953125
+    ],
+    [
+      -0.013671875,
+      -0.01953125,
+      -0.012939453125,
+      0.9453125,
+      -0.001953125,
+      0.001953125,
+      -0.0029296875,
+      0.92578125,
+      -0.001953125
+    ],
+    [
+      0.00341796875,
+      -0.0086669921875,
+      0.0078125,
+      0.9609375,
+      0.00244140625,
+      0.0048828125,
+      0.009765625,
+      0.9453125,
+      0.00146484375
+    ],
+    [
+      -0.015625,
+      0.00390625,
+      -0.02734375,
+      0.94140625,
+      -0.0048828125,
+      -0.0029296875,
+      -0.0078125,
+      0.96484375,
+      0.0028076171875
+    ],
+    [
+      0.013671875,
+      0.0009765625,
+      -0.0390625,
+      0.9453125,
+      0.0078125,
+      0.005859375,
+      -0.01953125,
+      0.921875,
+      -0.009765625
+    ],
+    [
+      0.005859375,
+      0.005859375,
+      -0.006591796875,
+      0.9765625,
+      -0.0029296875,
+      -0.00244140625,
+      -0.0009765625,
+      0.96875,
+      0.015625
+    ],
+    [
+      -0.013671875,
+      -0.013671875,
+      -0.0078125,
+      0.96875,
+      -0.0009765625,
+      0.002197265625,
+      -0.015625,
+      0.95703125,
+      -0.01171875
+    ],
+    [
+      0.001708984375,
+      0.01953125,
+      -0.0078125,
+      0.9453125,
+      -0.0087890625,
+      0.005859375,
+      0.0126953125,
+      0.9375,
+      0.013671875
+    ],
+    [
+      0.005859375,
+      -0.0009765625,
+      -0.0111083984375,
+      0.94921875,
+      -0.013671875,
+      -0.001708984375,
+      0.009765625,
+      0.953125,
+      0.002197265625
+    ],
+    [
+      -0.0009765625,
+      -0.021484375,
+      0.021484375,
+      0.96484375,
+      -0.0087890625,
+      -0.0078125,
+      0.00341796875,
+      0.953125,
+      -0.005859375
+    ],
+    [
+      0.02734375,
+      -0.01953125,
+      0.0068359375,
+      0.95703125,
+      0.009765625,
+      0.00244140625,
+      -0.0078125,
+      0.953125,
+      0.013671875
+    ],
+    [
+      0.017578125,
+      -0.00927734375,
+      0.046875,
+      0.95703125,
+      0.0087890625,
+      0.001953125,
+      -0.009765625,
+      0.9453125,
+      0.01171875
+    ],
+    [
+      0.01171875,
+      0.0087890625,
+      0.01611328125,
+      0.9453125,
+      -0.004638671875,
+      -0.0087890625,
+      0.00537109375,
+      0.953125,
+      -0.01171875
+    ],
+    [
+      0.001220703125,
+      0.00390625,
+      0.0458984375,
+      0.9609375,
+      0.0078125,
+      0.0029296875,
+      0.0087890625,
+      0.94921875,
+      0.0048828125
+    ],
+    [
+      -0.005859375,
+      0.001953125,
+      0.02392578125,
+      0.98046875,
+      0.01953125,
+      0.005859375,
+      0.00634765625,
+      0.95703125,
+      0.005859375
+    ],
+    [
+      -0.0068359375,
+      -0.003662109375,
+      0.0634765625,
+      0.9296875,
+      -0.0068359375,
+      0.0001220703125,
+      0.015625,
+      0.94921875,
+      0.000244140625
+    ],
+    [
+      0.0068359375,
+      0.005859375,
+      0.06591796875,
+      0.9296875,
+      -0.01171875,
+      0.0029296875,
+      0.013671875,
+      0.953125,
+      -0.00390625
+    ],
+    [
+      0.01171875,
+      -0.02734375,
+      0.07421875,
+      0.9609375,
+      0.0078125,
+      0.0078125,
+      0.013671875,
+      0.9765625,
+      0.00390625
+    ],
+    [
+      -0.0018310546875,
+      0.013671875,
+      0.0888671875,
+      0.95703125,
+      -0.015625,
+      0.001953125,
+      0.009765625,
+      0.96875,
+      0.0078125
+    ],
+    [
+      0.015625,
+      -0.025390625,
+      0.08251953125,
+      0.94921875,
+      0.0068359375,
+      -0.013671875,
+      -0.002197265625,
+      0.96875,
+      0.001953125
+    ],
+    [
+      0.017578125,
+      0.00146484375,
+      0.09619140625,
+      0.94140625,
+      0.0029296875,
+      0.01171875,
+      0.0067138671875,
+      0.9296875,
+      0.0029296875
+    ],
+    [
+      -0.0001220703125,
+      -0.025390625,
+      0.111328125,
+      0.96875,
+      0.00390625,
+      -0.001220703125,
+      0.0006103515625,
+      0.94140625,
+      0.00048828125
+    ],
+    [
+      0.00146484375,
+      0.0078125,
+      0.08935546875,
+      0.94921875,
+      -0.00238037109375,
+      -0.0087890625,
+      -0.005859375,
+      0.95703125,
+      0.0029296875
+    ],
+    [
+      -0.00390625,
+      0.0078125,
+      0.10302734375,
+      0.9609375,
+      0.005859375,
+      0.0078125,
+      0.0068359375,
+      0.9453125,
+      -0.0009765625
+    ],
+    [
+      -0.0087890625,
+      0.0029296875,
+      0.10986328125,
+      0.9765625,
+      0.00048828125,
+      -0.0078125,
+      0.017578125,
+      0.95703125,
+      -0.0087890625
+    ],
+    [
+      -0.00665283203125,
+      0.0078125,
+      0.1064453125,
+      0.953125,
+      0.00244140625,
+      -0.015625,
+      0.000244140625,
+      0.96875,
+      0.00146484375
+    ],
+    [
+      -0.00146484375,
+      -0.00384521484375,
+      0.15625,
+      0.9609375,
+      -0.01171875,
+      0.01171875,
+      0.013671875,
+      0.95703125,
+      -0.0029296875
+    ],
+    [
+      0.0048828125,
+      0.013671875,
+      0.1318359375,
+      0.94921875,
+      0.00439453125,
+      -0.005859375,
+      0.003662109375,
+      0.96875,
+      -0.00244140625
+    ],
+    [
+      -0.01171875,
+      0.01171875,
+      0.138671875,
+      0.953125,
+      -0.0078125,
+      0.0009765625,
+      0.009765625,
+      0.953125,
+      0.0078125
+    ],
+    [
+      -0.00390625,
+      -0.006103515625,
+      0.138671875,
+      0.97265625,
+      0.002197265625,
+      0.0009765625,
+      0.0048828125,
+      0.9453125,
+      0.00341796875
+    ],
+    [
+      0.0009765625,
+      -0.0028076171875,
+      0.17578125,
+      0.953125,
+      0.00048828125,
+      0.0,
+      0.005859375,
+      0.96875,
+      0.0009765625
+    ],
+    [
+      0.0078125,
+      -0.00439453125,
+      0.150390625,
+      0.953125,
+      -0.0068359375,
+      0.0,
+      -0.0078125,
+      0.95703125,
+      0.0029296875
+    ],
+    [
+      -0.0009765625,
+      0.0048828125,
+      0.15234375,
+      0.9453125,
+      -0.0068359375,
+      -0.005859375,
+      0.00390625,
+      0.92578125,
+      0.0078125
+    ],
+    [
+      0.00341796875,
+      0.001953125,
+      0.14453125,
+      0.96484375,
+      0.0009765625,
+      -0.017578125,
+      -0.001953125,
+      0.96875,
+      0.00146484375
+    ],
+    [
+      0.009765625,
+      0.0,
+      0.197265625,
+      0.9609375,
+      0.00341796875,
+      -0.00390625,
+      0.005859375,
+      0.96484375,
+      0.00390625
+    ],
+    [
+      0.0087890625,
+      -0.0015869140625,
+      0.1826171875,
+      0.97265625,
+      0.0054931640625,
+      -0.00390625,
+      0.00048828125,
+      0.9765625,
+      0.009765625
+    ],
+    [
+      0.00390625,
+      0.0001220703125,
+      0.19921875,
+      0.9609375,
+      0.015625,
+      -0.0126953125,
+      -0.0078125,
+      0.984375,
+      0.015625
+    ],
+    [
+      0.0205078125,
+      0.005859375,
+      0.21875,
+      0.9609375,
+      0.01953125,
+      -0.01025390625,
+      -0.0029296875,
+      0.96875,
+      0.00390625
+    ],
+    [
+      0.029296875,
+      -0.00927734375,
+      0.22265625,
+      0.9609375,
+      0.00341796875,
+      -0.001953125,
+      0.00390625,
+      0.96484375,
+      -0.00390625
+    ],
+    [
+      0.021484375,
+      -0.004638671875,
+      0.1875,
+      0.9765625,
+      0.004638671875,
+      -0.0126953125,
+      -0.003173828125,
+      0.97265625,
+      0.001708984375
+    ],
+    [
+      0.01239013671875,
+      0.0068359375,
+      0.236328125,
+      0.96875,
+      -0.009765625,
+      -0.013671875,
+      0.01953125,
+      0.953125,
+      0.0068359375
+    ],
+    [
+      0.025390625,
+      -0.017578125,
+      0.244140625,
+      0.95703125,
+      -0.01953125,
+      -0.0087890625,
+      0.01171875,
+      0.95703125,
+      0.01953125
+    ],
+    [
+      0.0194091796875,
+      0.0078125,
+      0.251953125,
+      0.953125,
+      -0.00048828125,
+      0.001953125,
+      0.00152587890625,
+      0.95703125,
+      -0.0078125
+    ],
+    [
+      0.0166015625,
+      -0.0068359375,
+      0.2265625,
+      0.9453125,
+      0.000732421875,
+      -0.01953125,
+      0.00390625,
+      0.9453125,
+      -0.00341796875
+    ],
+    [
+      0.016845703125,
+      0.00341796875,
+      0.26171875,
+      0.9453125,
+      -0.001953125,
+      -0.005859375,
+      0.00390625,
+      0.9375,
+      0.00390625
+    ],
+    [
+      0.014892578125,
+      0.0087890625,
+      0.298828125,
+      0.953125,
+      0.005859375,
+      -0.009033203125,
+      0.0009765625,
+      0.953125,
+      -0.0029296875
+    ],
+    [
+      0.0224609375,
+      -0.0234375,
+      0.30078125,
+      0.9609375,
+      0.005859375,
+      -0.0048828125,
+      -0.0003662109375,
+      0.9609375,
+      0.009765625
+    ],
+    [
+      0.017822265625,
+      0.005859375,
+      0.322265625,
+      0.9609375,
+      -0.0029296875,
+      -0.00634765625,
+      -0.001220703125,
+      0.96875,
+      0.0029296875
+    ],
+    [
+      0.02392578125,
+      -0.00390625,
+      0.30078125,
+      0.9375,
+      0.0006103515625,
+      0.00390625,
+      -0.005859375,
+      0.93359375,
+      0.003662109375
+    ],
+    [
+      0.0185546875,
+      -0.0020751953125,
+      0.333984375,
+      0.94921875,
+      0.001220703125,
+      -0.015625,
+      -0.00048828125,
+      0.96875,
+      0.0048828125
+    ],
+    [
+      0.0185546875,
+      0.01171875,
+      0.326171875,
+      0.9453125,
+      0.0023193359375,
+      -0.008056640625,
+      0.00146484375,
+      0.95703125,
+      -0.0068359375
+    ],
+    [
+      0.0162353515625,
+      0.015625,
+      0.35546875,
+      0.953125,
+      0.009765625,
+      -0.0068359375,
+      0.01171875,
+      0.96875,
+      0.0
+    ],
+    [
+      0.02099609375,
+      0.000244140625,
+      0.375,
+      0.95703125,
+      -0.0009765625,
+      -0.0029296875,
+      -0.0009765625,
+      0.96484375,
+      0.00390625
+    ],
+    [
+      0.0145263671875,
+      -0.015625,
+      0.369140625,
+      0.96484375,
+      0.00164794921875,
+      -0.01171875,
+      -0.00537109375,
+      0.953125,
+      -0.0009765625
+    ],
+    [
+      0.021484375,
+      0.00390625,
+      0.37109375,
+      0.953125,
+      0.00048828125,
+      -0.013671875,
+      0.0048828125,
+      0.95703125,
+      0.001953125
+    ],
+    [
+      0.018798828125,
+      -0.0078125,
+      0.41015625,
+      0.95703125,
+      -0.01171875,
+      -0.00390625,
+      -0.0068359375,
+      0.9609375,
+      0.005859375
+    ]
+  ],
+  "shape": [
+    60,
+    9
+  ],
+  "dtype": "torch.bfloat16",
+  "raw_action_dim": 9,
+  "action_mode": "inverse_dynamics",
+  "domain_id": 1
+}

assets/edge_action_id_av_1_output.json

@@ -0,0 +1,672 @@
+{
+  "data": [
+    [
+      0.076171875,
+      0.0,
+      1.0859375,
+      0.96875,
+      0.004150390625,
+      -0.017578125,
+      0.005859375,
+      0.9765625,
+      0.0078125
+    ],
+    [
+      0.0654296875,
+      0.00390625,
+      0.9921875,
+      0.9765625,
+      0.0107421875,
+      -0.02490234375,
+      -0.0048828125,
+      0.96875,
+      0.0107421875
+    ],
+    [
+      0.087890625,
+      -0.0068359375,
+      1.0,
+      0.9375,
+      0.013671875,
+      -0.037109375,
+      0.01953125,
+      0.984375,
+      0.00146484375
+    ],
+    [
+      0.08203125,
+      -0.001953125,
+      0.984375,
+      0.97265625,
+      -0.0029296875,
+      -0.027099609375,
+      0.009765625,
+      0.97265625,
+      0.0126953125
+    ],
+    [
+      0.0849609375,
+      -0.0322265625,
+      0.95703125,
+      0.97265625,
+      0.009765625,
+      -0.04296875,
+      -0.0126953125,
+      0.96875,
+      0.00537109375
+    ],
+    [
+      0.080078125,
+      -0.021484375,
+      0.9375,
+      0.953125,
+      0.000732421875,
+      -0.03271484375,
+      -0.0068359375,
+      0.93359375,
+      -0.0009765625
+    ],
+    [
+      0.1015625,
+      -0.0042724609375,
+      0.9453125,
+      0.97265625,
+      0.00244140625,
+      -0.03173828125,
+      0.0078125,
+      0.953125,
+      0.00244140625
+    ],
+    [
+      0.09375,
+      0.0224609375,
+      0.91015625,
+      0.9453125,
+      -0.001953125,
+      -0.0361328125,
+      -0.015625,
+      0.96875,
+      0.0018310546875
+    ],
+    [
+      0.1318359375,
+      0.017578125,
+      0.90625,
+      0.95703125,
+      0.0107421875,
+      -0.0322265625,
+      -0.0234375,
+      0.93359375,
+      -0.017578125
+    ],
+    [
+      0.130859375,
+      0.025390625,
+      0.91796875,
+      0.984375,
+      0.0009765625,
+      -0.04345703125,
+      -0.0078125,
+      0.9765625,
+      0.017578125
+    ],
+    [
+      0.107421875,
+      -0.001953125,
+      0.90234375,
+      0.9765625,
+      0.002685546875,
+      -0.04443359375,
+      -0.017578125,
+      0.95703125,
+      -0.009765625
+    ],
+    [
+      0.1220703125,
+      0.03125,
+      0.90234375,
+      0.94921875,
+      -0.0078125,
+      -0.037109375,
+      0.0087890625,
+      0.9453125,
+      0.017578125
+    ],
+    [
+      0.130859375,
+      0.009765625,
+      0.88671875,
+      0.953125,
+      -0.009765625,
+      -0.040283203125,
+      0.0029296875,
+      0.953125,
+      0.005126953125
+    ],
+    [
+      0.12353515625,
+      -0.0078125,
+      0.91015625,
+      0.96875,
+      0.0029296875,
+      -0.0478515625,
+      -0.00537109375,
+      0.95703125,
+      -0.0029296875
+    ],
+    [
+      0.15625,
+      -0.001953125,
+      0.8828125,
+      0.96484375,
+      0.0126953125,
+      -0.038330078125,
+      -0.01171875,
+      0.953125,
+      0.01171875
+    ],
+    [
+      0.1328125,
+      0.005859375,
+      0.9140625,
+      0.96875,
+      0.0087890625,
+      -0.03955078125,
+      -0.013671875,
+      0.953125,
+      0.013671875
+    ],
+    [
+      0.1318359375,
+      0.0205078125,
+      0.8984375,
+      0.953125,
+      0.0,
+      -0.0458984375,
+      0.0009765625,
+      0.95703125,
+      -0.0078125
+    ],
+    [
+      0.119140625,
+      0.013671875,
+      0.921875,
+      0.96875,
+      0.0087890625,
+      -0.037109375,
+      0.001953125,
+      0.95703125,
+      0.0068359375
+    ],
+    [
+      0.115234375,
+      0.015625,
+      0.91015625,
+      0.984375,
+      0.01953125,
+      -0.0361328125,
+      -0.000732421875,
+      0.9609375,
+      0.0078125
+    ],
+    [
+      0.1103515625,
+      0.01312255859375,
+      0.9453125,
+      0.9375,
+      -0.00244140625,
+      -0.038330078125,
+      0.005859375,
+      0.953125,
+      -0.001068115234375
+    ],
+    [
+      0.1220703125,
+      0.0244140625,
+      0.953125,
+      0.9375,
+      -0.005859375,
+      -0.033203125,
+      0.005859375,
+      0.95703125,
+      -0.00390625
+    ],
+    [
+      0.1279296875,
+      -0.00390625,
+      0.96875,
+      0.96875,
+      0.0087890625,
+      -0.029296875,
+      0.00390625,
+      0.98046875,
+      0.0048828125
+    ],
+    [
+      0.11572265625,
+      0.033203125,
+      0.984375,
+      0.9609375,
+      -0.0078125,
+      -0.0341796875,
+      0.0029296875,
+      0.97265625,
+      0.009765625
+    ],
+    [
+      0.126953125,
+      -0.00390625,
+      0.9765625,
+      0.9609375,
+      0.0087890625,
+      -0.044921875,
+      -0.00390625,
+      0.9765625,
+      0.00244140625
+    ],
+    [
+      0.12109375,
+      0.01953125,
+      0.9921875,
+      0.953125,
+      0.0048828125,
+      -0.0234375,
+      0.0010986328125,
+      0.9375,
+      0.0029296875
+    ],
+    [
+      0.10546875,
+      -0.0078125,
+      1.015625,
+      0.9765625,
+      0.0078125,
+      -0.03466796875,
+      -0.0035400390625,
+      0.94921875,
+      0.0029296875
+    ],
+    [
+      0.1064453125,
+      0.0263671875,
+      1.0078125,
+      0.94921875,
+      0.00042724609375,
+      -0.0400390625,
+      -0.013671875,
+      0.9609375,
+      0.005859375
+    ],
+    [
+      0.095703125,
+      0.0244140625,
+      1.03125,
+      0.96875,
+      0.0068359375,
+      -0.0234375,
+      0.00048828125,
+      0.953125,
+      0.002197265625
+    ],
+    [
+      0.08837890625,
+      0.02099609375,
+      1.0390625,
+      0.984375,
+      0.0052490234375,
+      -0.03515625,
+      0.009765625,
+      0.96484375,
+      -0.0048828125
+    ],
+    [
+      0.08154296875,
+      0.029296875,
+      1.0625,
+      0.953125,
+      0.0048828125,
+      -0.04296875,
+      -0.005615234375,
+      0.96875,
+      0.00634765625
+    ],
+    [
+      0.0849609375,
+      0.00823974609375,
+      1.109375,
+      0.96875,
+      -0.0078125,
+      -0.01953125,
+      0.00390625,
+      0.96484375,
+      0.00146484375
+    ],
+    [
+      0.08203125,
+      0.029296875,
+      1.1171875,
+      0.95703125,
+      0.0048828125,
+      -0.02734375,
+      0.0,
+      0.97265625,
+      0.004852294921875
+    ],
+    [
+      0.056640625,
+      0.021484375,
+      1.1328125,
+      0.9609375,
+      -0.005859375,
+      -0.02001953125,
+      0.0078125,
+      0.9609375,
+      0.015625
+    ],
+    [
+      0.060546875,
+      0.001220703125,
+      1.140625,
+      0.984375,
+      0.0029296875,
+      -0.017578125,
+      0.001953125,
+      0.94921875,
+      0.009765625
+    ],
+    [
+      0.060546875,
+      0.007568359375,
+      1.1875,
+      0.96484375,
+      -0.000732421875,
+      -0.0166015625,
+      0.0029296875,
+      0.96875,
+      0.005615234375
+    ],
+    [
+      0.060302734375,
+      0.01123046875,
+      1.1875,
+      0.9609375,
+      -0.0087890625,
+      -0.0166015625,
+      -0.01171875,
+      0.9609375,
+      0.0078125
+    ],
+    [
+      0.046875,
+      0.01953125,
+      1.2109375,
+      0.953125,
+      -0.0048828125,
+      -0.0184326171875,
+      0.0029296875,
+      0.9296875,
+      0.009765625
+    ],
+    [
+      0.04248046875,
+      0.01953125,
+      1.2109375,
+      0.97265625,
+      0.005859375,
+      -0.029296875,
+      -0.0068359375,
+      0.9765625,
+      0.0078125
+    ],
+    [
+      0.041015625,
+      0.01611328125,
+      1.2734375,
+      0.96875,
+      0.0048828125,
+      -0.00927734375,
+      0.0009765625,
+      0.97265625,
+      0.0087890625
+    ],
+    [
+      0.0302734375,
+      0.009521484375,
+      1.265625,
+      0.984375,
+      0.004638671875,
+      -0.005859375,
+      -0.00537109375,
+      0.98046875,
+      0.013671875
+    ],
+    [
+      0.015625,
+      0.0086669921875,
+      1.296875,
+      0.96875,
+      0.013671875,
+      -0.013671875,
+      -0.015625,
+      0.984375,
+      0.0234375
+    ],
+    [
+      0.0283203125,
+      0.0146484375,
+      1.3359375,
+      0.96875,
+      0.01953125,
+      -0.0068359375,
+      -0.0078125,
+      0.9765625,
+      0.01318359375
+    ],
+    [
+      0.033203125,
+      0.0,
+      1.34375,
+      0.97265625,
+      0.00146484375,
+      0.00390625,
+      0.001953125,
+      0.97265625,
+      0.005615234375
+    ],
+    [
+      0.0234375,
+      0.0048828125,
+      1.328125,
+      0.9921875,
+      0.00439453125,
+      -0.0078125,
+      -0.003662109375,
+      0.98046875,
+      0.010009765625
+    ],
+    [
+      0.009033203125,
+      0.0166015625,
+      1.3984375,
+      0.984375,
+      -0.01171875,
+      -0.009765625,
+      0.0234375,
+      0.96484375,
+      0.01171875
+    ],
+    [
+      0.021484375,
+      -0.01171875,
+      1.40625,
+      0.96875,
+      -0.01953125,
+      -0.0048828125,
+      0.01171875,
+      0.96484375,
+      0.0234375
+    ],
+    [
+      0.0155029296875,
+      0.021484375,
+      1.4375,
+      0.9609375,
+      -0.00341796875,
+      0.009765625,
+      0.0006103515625,
+      0.96484375,
+      -0.00390625
+    ],
+    [
+      0.0108642578125,
+      0.005859375,
+      1.421875,
+      0.953125,
+      -0.002685546875,
+      -0.013671875,
+      0.0068359375,
+      0.953125,
+      -0.00048828125
+    ],
+    [
+      0.01068115234375,
+      0.01171875,
+      1.46875,
+      0.953125,
+      -0.00341796875,
+      -0.00048828125,
+      0.005859375,
+      0.9453125,
+      0.009765625
+    ],
+    [
+      0.00830078125,
+      0.0224609375,
+      1.515625,
+      0.96875,
+      0.005859375,
+      0.00048828125,
+      0.0009765625,
+      0.96484375,
+      0.001953125
+    ],
+    [
+      0.01171875,
+      -0.015625,
+      1.515625,
+      0.9765625,
+      0.0029296875,
+      0.001953125,
+      -0.001220703125,
+      0.96875,
+      0.01171875
+    ],
+    [
+      0.010986328125,
+      0.0185546875,
+      1.546875,
+      0.97265625,
+      -0.00390625,
+      0.00244140625,
+      -0.00390625,
+      0.984375,
+      0.0078125
+    ],
+    [
+      0.01318359375,
+      0.0087890625,
+      1.5390625,
+      0.94921875,
+      0.0,
+      0.015625,
+      -0.0078125,
+      0.953125,
+      0.005859375
+    ],
+    [
+      0.00830078125,
+      0.01165771484375,
+      1.578125,
+      0.953125,
+      0.0003662109375,
+      -0.009765625,
+      0.0,
+      0.9765625,
+      0.0087890625
+    ],
+    [
+      0.011962890625,
+      0.025390625,
+      1.5703125,
+      0.953125,
+      -0.00164794921875,
+      -0.000244140625,
+      0.001953125,
+      0.96875,
+      -0.0068359375
+    ],
+    [
+      0.00970458984375,
+      0.025390625,
+      1.6171875,
+      0.9609375,
+      0.0068359375,
+      0.0009765625,
+      0.0078125,
+      0.984375,
+      0.0048828125
+    ],
+    [
+      0.01220703125,
+      0.014892578125,
+      1.640625,
+      0.96875,
+      -0.00390625,
+      0.0,
+      -0.0048828125,
+      0.97265625,
+      0.0078125
+    ],
+    [
+      0.006103515625,
+      -0.001953125,
+      1.640625,
+      0.9765625,
+      0.002471923828125,
+      -0.001953125,
+      -0.00537109375,
+      0.96875,
+      0.00341796875
+    ],
+    [
+      0.0166015625,
+      0.017578125,
+      1.640625,
+      0.96875,
+      0.00146484375,
+      -0.009765625,
+      0.0048828125,
+      0.96484375,
+      0.0087890625
+    ],
+    [
+      0.005126953125,
+      0.0087890625,
+      1.6796875,
+      0.9609375,
+      -0.005859375,
+      -0.00146484375,
+      -0.005859375,
+      0.96875,
+      0.0068359375
+    ]
+  ],
+  "shape": [
+    60,
+    9
+  ],
+  "dtype": "torch.bfloat16",
+  "raw_action_dim": 9,
+  "action_mode": "inverse_dynamics",
+  "domain_id": 1
+}

assets/edge_i2v_output.mp4

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

assets/example_action_fd_umi_action_chunks.json

@@ -0,0 +1,404 @@
+{
+  "prompt": "mouse arrangement",
+  "domain_name": "umi",
+  "view_point": "ego_view",
+  "fps": 20,
+  "image_size": 256,
+  "action_chunk_size": 16,
+  "num_chunks": 2,
+  "action_shape_per_chunk": [
+    16,
+    10
+  ],
+  "description": "Two normalized 10-D UMI action chunks for a single-view ego_view rollout. Use the first frame as the initial conditioning image. For chunk 1, condition on the final generated frame from chunk 0.",
+  "action_chunks": [
+    [
+      [
+        0.053812503814697266,
+        -0.004779123235493898,
+        0.002008301904425025,
+        0.9671046733856201,
+        0.03711317479610443,
+        -0.17747420072555542,
+        -0.036594998091459274,
+        0.9977647662162781,
+        0.02953636832535267,
+        -0.13823197782039642
+      ],
+      [
+        0.07370813935995102,
+        -0.00015990791143849492,
+        -0.00526643730700016,
+        0.966162919998169,
+        0.05513376370072365,
+        -0.17547594010829926,
+        -0.05433010682463646,
+        0.994891881942749,
+        0.04623955488204956,
+        -0.13932132720947266
+      ],
+      [
+        0.1894039809703827,
+        -0.0023578035179525614,
+        0.012279288843274117,
+        0.9756994843482971,
+        0.09475686401128769,
+        -0.12368101626634598,
+        -0.09487718343734741,
+        0.9908984303474426,
+        -0.009143882431089878,
+        -0.14122092723846436
+      ],
+      [
+        0.1646755337715149,
+        0.014868665486574173,
+        -0.02827119641005993,
+        0.9662106037139893,
+        0.11741045117378235,
+        -0.14138658344745636,
+        -0.11718737334012985,
+        0.9859751462936401,
+        0.01658531278371811,
+        -0.14161521196365356
+      ],
+      [
+        0.20245812833309174,
+        0.039910294115543365,
+        -0.025805803015828133,
+        0.8995737433433533,
+        0.16153031587600708,
+        -0.27254194021224976,
+        -0.1620442271232605,
+        0.9734464883804321,
+        -0.016664521768689156,
+        -0.14406844973564148
+      ],
+      [
+        0.18459202349185944,
+        0.03604526445269585,
+        -0.02222990058362484,
+        0.941689670085907,
+        0.08134357631206512,
+        -0.22731782495975494,
+        -0.08176705986261368,
+        0.9929845333099365,
+        -0.01771094650030136,
+        -0.14325366914272308
+      ],
+      [
+        0.19590386748313904,
+        0.02676231600344181,
+        -0.02729889750480652,
+        0.8968558311462402,
+        0.16832612454891205,
+        -0.2734377384185791,
+        -0.16814984381198883,
+        0.97162264585495,
+        0.008744634687900543,
+        -0.1433658003807068
+      ],
+      [
+        0.2117065191268921,
+        0.019320685416460037,
+        -0.013418853282928467,
+        0.8548710346221924,
+        0.22257667779922485,
+        -0.30904778838157654,
+        -0.22233201563358307,
+        0.9504156708717346,
+        0.011350061744451523,
+        -0.1437632143497467
+      ],
+      [
+        0.22299973666667938,
+        0.0327562540769577,
+        -0.009134029038250446,
+        0.8706302046775818,
+        0.21322639286518097,
+        -0.28956693410873413,
+        -0.21412642300128937,
+        0.953336238861084,
+        -0.02801540493965149,
+        -0.14188432693481445
+      ],
+      [
+        0.2140849381685257,
+        0.02660888060927391,
+        -0.01170671172440052,
+        0.8908835649490356,
+        0.20415504276752472,
+        -0.2596054971218109,
+        -0.20452779531478882,
+        0.9580091238021851,
+        -0.01170840673148632,
+        -0.14164277911186218
+      ],
+      [
+        0.22304928302764893,
+        0.028626402840018272,
+        -0.006994225084781647,
+        0.8915749192237854,
+        0.12412955611944199,
+        -0.30490365624427795,
+        -0.12411271035671234,
+        0.9845566153526306,
+        0.0024451441131532192,
+        -0.14154405891895294
+      ],
+      [
+        0.2222731113433838,
+        0.025864848867058754,
+        -0.03064606711268425,
+        0.8916106820106506,
+        0.11292119324207306,
+        -0.3092045783996582,
+        -0.11282884329557419,
+        0.9872149229049683,
+        0.004730388056486845,
+        -0.14219173789024353
+      ],
+      [
+        0.2097492665052414,
+        0.05057718604803085,
+        -0.007311991881579161,
+        0.9075368046760559,
+        0.1548784226179123,
+        -0.26160797476768494,
+        -0.15552769601345062,
+        0.9752822518348694,
+        -0.022801857441663742,
+        -0.1428474634885788
+      ],
+      [
+        0.2123514711856842,
+        0.042986635118722916,
+        0.006450705695897341,
+        0.9269675612449646,
+        0.16701975464820862,
+        -0.2123788744211197,
+        -0.16792836785316467,
+        0.9700967669487,
+        -0.041070736944675446,
+        -0.14396123588085175
+      ],
+      [
+        0.2005469650030136,
+        0.04303627461194992,
+        0.0021174526773393154,
+        0.9145819544792175,
+        0.15997713804244995,
+        -0.24451406300067902,
+        -0.16089586913585663,
+        0.9728385210037231,
+        -0.0356520377099514,
+        -0.143580824136734
+      ],
+      [
+        0.1828203797340393,
+        0.06162899360060692,
+        0.015765784308314323,
+        0.9442406892776489,
+        0.12680575251579285,
+        -0.199128195643425,
+        -0.12898683547973633,
+        0.9715272784233093,
+        -0.10863088071346283,
+        -0.14364691078662872
+      ]
+    ],
+    [
+      [
+        0.1800091415643692,
+        0.05415700003504753,
+        0.0220723208039999,
+        0.9427744746208191,
+        0.15702207386493683,
+        -0.18039608001708984,
+        -0.1585177183151245,
+        0.9681656360626221,
+        -0.08177943527698517,
+        -0.14226296544075012
+      ],
+      [
+        0.1782533973455429,
+        0.046168215572834015,
+        0.03214065358042717,
+        0.9491162896156311,
+        0.15013249218463898,
+        -0.16828377544879913,
+        -0.15176372230052948,
+        0.9677245616912842,
+        -0.09608176350593567,
+        -0.1437339335680008
+      ],
+      [
+        0.17733417451381683,
+        0.05023511126637459,
+        0.042507752776145935,
+        0.9599641561508179,
+        0.11787547171115875,
+        -0.16160599887371063,
+        -0.11985509097576141,
+        0.9706928133964539,
+        -0.12208050489425659,
+        -0.14515265822410583
+      ],
+      [
+        0.17919816076755524,
+        0.07686550915241241,
+        0.06860226392745972,
+        0.9644821286201477,
+        0.10539975762367249,
+        -0.15618394315242767,
+        -0.10708234459161758,
+        0.9769988656044006,
+        -0.10728096216917038,
+        -0.1453562080860138
+      ],
+      [
+        0.1575465053319931,
+        0.029737284407019615,
+        0.05468235909938812,
+        0.9541349411010742,
+        0.15530435740947723,
+        -0.14736345410346985,
+        -0.1561610847711563,
+        0.9723378419876099,
+        -0.057166263461112976,
+        -0.14554671943187714
+      ],
+      [
+        0.17658387124538422,
+        -0.0008871768950484693,
+        0.09531591832637787,
+        0.983400285243988,
+        0.09690801054239273,
+        -0.08479148149490356,
+        -0.09749053418636322,
+        0.9857844114303589,
+        -0.06855253875255585,
+        -0.1436685174703598
+      ],
+      [
+        0.1308804750442505,
+        0.016017800197005272,
+        0.04266123101115227,
+        0.9813141822814941,
+        0.10064378380775452,
+        -0.09238820523023605,
+        -0.10080896317958832,
+        0.9895036816596985,
+        -0.017434898763895035,
+        -0.14607854187488556
+      ],
+      [
+        0.12789754569530487,
+        0.0020461133681237698,
+        0.06620839238166809,
+        0.99201899766922,
+        0.03137550875544548,
+        -0.08352884650230408,
+        -0.03163702413439751,
+        0.9980270266532898,
+        -0.031180214136838913,
+        -0.14278197288513184
+      ],
+      [
+        0.11177432537078857,
+        -0.02513059414923191,
+        0.08431962132453918,
+        0.9990403056144714,
+        -0.010432509705424309,
+        -0.02879752404987812,
+        0.010669702664017677,
+        0.9930679798126221,
+        0.08243173360824585,
+        -0.14342190325260162
+      ],
+      [
+        0.0631430447101593,
+        -0.010469917207956314,
+        0.08317670226097107,
+        0.999850869178772,
+        -0.011552712880074978,
+        0.0031132197473198175,
+        0.011538979597389698,
+        0.9979555010795593,
+        0.04372014105319977,
+        -0.1447361707687378
+      ],
+      [
+        0.09046280384063721,
+        -0.012558498419821262,
+        0.09904492646455765,
+        0.998515784740448,
+        0.032513268291950226,
+        0.02011110819876194,
+        -0.0324835367500782,
+        0.9986945986747742,
+        -0.01484532468020916,
+        -0.14469973742961884
+      ],
+      [
+        0.07805871963500977,
+        -0.019602082669734955,
+        0.11864962428808212,
+        0.9934136867523193,
+        0.054121777415275574,
+        0.06036890670657158,
+        -0.05413716286420822,
+        0.9970375895500183,
+        0.002376429969444871,
+        -0.1444096863269806
+      ],
+      [
+        0.08657136559486389,
+        -0.0348740816116333,
+        0.12534210085868835,
+        0.9954044818878174,
+        0.005703791975975037,
+        0.067320816218853,
+        -0.00547851761803031,
+        0.9988376498222351,
+        -0.033448826521635056,
+        -0.14501787722110748
+      ],
+      [
+        0.07565143704414368,
+        -0.044480662792921066,
+        0.1283297836780548,
+        0.9972044825553894,
+        0.022084418684244156,
+        0.04793533310294151,
+        -0.02211857959628105,
+        0.999433696269989,
+        0.007050265092402697,
+        -0.14408022165298462
+      ],
+      [
+        0.06433551013469696,
+        -0.04466487839818001,
+        0.13651199638843536,
+        0.9936521053314209,
+        -0.04951975867152214,
+        0.0623551569879055,
+        0.049662359058856964,
+        0.9970018267631531,
+        -0.022745873779058456,
+        -0.14406241476535797
+      ],
+      [
+        0.05908937379717827,
+        -0.05935889855027199,
+        0.1395753026008606,
+        0.9970375895500183,
+        -0.024985427036881447,
+        0.04805263504385948,
+        0.02504965290427208,
+        0.9991833567619324,
+        -0.013274282217025757,
+        -0.14403760433197021
+      ]
+    ]
+  ]
+}

assets/example_action_id_av_0_input.mp4

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

assets/example_action_id_av_1_input.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:169e65cee76ef7c6366987a8e4c8e4ec6f803f5e624ba2ef39ba1f0c672925aa
+size 1639407

assets/example_i2v_prompt.json

@@ -0,0 +1,124 @@
+{
+    "subjects": [
+      {
+        "description": "A car (viewed from the driver's perspective/dashcam) traveling along a coastal mountain road. The vehicle's grey dashboard and windshield frame are visible at the bottom of the frame.",
+        "appearance_details": "Grey/dark interior dashboard, windshield wipers visible, clean windshield providing clear view of the road ahead",
+        "relationship": "Primary subject and point-of-view vehicle navigating the mountain road",
+        "location": "Bottom foreground, providing the POV perspective",
+        "relative_size": "Large within frame (dashboard occupies bottom portion)",
+        "orientation": "Facing forward along the road direction",
+        "pose": "",
+        "action": "Driving fast along the curved coastal road, then making a sudden emergency stop",
+        "state_changes": "Transitions from fast forward motion to abrupt deceleration and complete stop",
+        "clothing": "",
+        "expression": "",
+        "gender": "",
+        "age": "",
+        "skin_tone_and_texture": "",
+        "facial_features": "",
+        "number_of_subjects": 1,
+        "number_of_arms": 0,
+        "number_of_legs": 0
+      },
+      {
+        "description": "Falling rocks and debris from the steep rocky mountain cliff face, creating a landslide that tumbles down toward the road surface",
+        "appearance_details": "Grey and brown angular rocks of various sizes, accompanied by dust and smaller debris, mixed with displaced green vegetation",
+        "relationship": "Obstacle that forces the car to make an emergency stop",
+        "location": "Center and center-right of frame, falling from the mountain cliff face toward the road",
+        "relative_size": "Medium within frame, growing larger as rocks approach the road",
+        "orientation": "Falling downward from the cliff face",
+        "pose": "",
+        "action": "Tumbling and cascading down the steep rocky mountainside toward the road",
+        "state_changes": "Begins as small movement on the cliff face, escalates into a significant rockfall with dust cloud",
+        "clothing": "",
+        "expression": "",
+        "gender": "",
+        "age": "",
+        "skin_tone_and_texture": "",
+        "facial_features": "",
+        "number_of_subjects": 0,
+        "number_of_arms": 0,
+        "number_of_legs": 0
+      }
+    ],
+    "background_setting": "A dramatic coastal mountain highway carved along steep rocky cliffs. The road is a well-maintained two-lane highway with white dashed center lines, solid white edge lines, and raised reflective road markers. A metal guardrail runs along the right side protecting from a drop toward the ocean. The mountain rises steeply on the left with exposed grey and brown rock faces partially covered with lush green native vegetation including ferns and shrubs. The ocean is visible in the distance to the right under a partly cloudy blue sky. The road curves sharply to the left around the mountain headland.",
+    "lighting": {
+      "conditions": "Bright natural daylight with partly cloudy skies, strong ambient illumination",
+      "direction": "Top-lit and slightly front-lit from the sun positioned high in the sky",
+      "shadows": "Moderate shadows cast by the mountain cliff creating areas of shade on the road and vegetation, soft shadows from clouds",
+      "illumination_effect": "Bright, clear visibility with good contrast between the sunlit road surface and shaded mountain areas, natural coastal light"
+    },
+    "aesthetics": {
+      "composition": "Driver's POV/dashcam perspective with the road curving ahead as the central leading line, mountain dominating the upper-left portion, ocean visible to the right, dashboard framing the bottom",
+      "color_scheme": "Dominant greens from vegetation, grey-brown rocky cliff faces, dark grey asphalt road with white markings, blue sky and ocean, yellow warning sign providing accent color",
+      "mood_atmosphere": "Initially scenic and adventurous, transitioning to tense and dangerous as the landslide occurs",
+      "patterns": "Repeating metal guardrail posts along the right side, dashed center line markings, layered rock strata on the cliff face"
+    },
+    "cinematography": {
+      "camera_motion": "Forward-moving dashcam perspective with slight vibration from vehicle speed, sudden forward lurch during emergency braking",
+      "framing": "Wide shot from driver's POV through windshield",
+      "camera_angle": "Eye-level from driver's seated position",
+      "depth_of_field": "Deep",
+      "focus": "Road ahead and mountain cliff face in sharp focus",
+      "lens_focal_length": "Wide-angle (approximately 28-35mm equivalent)"
+    },
+    "style_medium": "Live-action video",
+    "artistic_style": "Realistic dashcam footage, dramatic documentary style",
+    "context": "Dashcam footage capturing a dangerous landslide event on a coastal mountain highway, requiring emergency driving maneuver",
+    "actions": [
+      {
+        "time": "0:00-0:02",
+        "description": "The car drives fast along the winding coastal mountain road, approaching the left curve with the road rushing beneath"
+      },
+      {
+        "time": "0:02-0:04",
+        "description": "Rocks begin to dislodge from the steep cliff face ahead, tumbling down with increasing intensity as a landslide develops, dust rising from impacts"
+      },
+      {
+        "time": "0:04-0:05",
+        "description": "The car makes a sudden emergency stop, the camera lurching forward from the deceleration as rocks and debris continue falling onto the road ahead"
+      }
+    ],
+    "text_and_signage_elements": [
+      {
+        "text": "Left curve arrow",
+        "category": "scene_sign",
+        "appearance": "Black arrow on yellow diamond-shaped warning sign, mounted on a metal pole",
+        "spatial_temporal": "Left side of road, visible from 0:00 and passing by as car moves forward",
+        "context": "Road warning sign indicating a sharp left curve ahead, standard highway signage"
+      }
+    ],
+    "segments": [
+      {
+        "segment_index": 0,
+        "time_range": "0:00-0:02",
+        "description": "The car speeds along the scenic coastal mountain highway, navigating the curve. The road surface rushes beneath, white lines blur slightly from speed, and the mountain looms ahead.",
+        "key_changes": "Road perspective shifts as car navigates the curve, passing the warning sign, guardrail moves in peripheral view",
+        "camera": "Forward-moving dashcam with slight road vibration, steady forward momentum"
+      },
+      {
+        "segment_index": 1,
+        "time_range": "0:02-0:04",
+        "description": "Rocks begin breaking loose from the cliff face ahead. Small stones first, then larger boulders cascade down the mountainside creating a growing dust cloud and debris field on and near the road.",
+        "key_changes": "Mountain transitions from static backdrop to active hazard, dust cloud forms, debris appears on road surface ahead",
+        "camera": "Still moving forward but beginning to slow, slight camera shake increases"
+      },
+      {
+        "segment_index": 2,
+        "time_range": "0:04-0:05",
+        "description": "The car performs an emergency stop. The camera pitches forward from sudden braking force. The vehicle comes to a complete halt with the landslide debris visible on the road ahead, dust still settling.",
+        "key_changes": "Rapid deceleration, camera lurches forward then settles, motion stops completely, dust and small rocks still visible falling",
+        "camera": "Abrupt forward pitch from braking, then stabilizes to static as vehicle stops"
+      }
+    ],
+    "transitions": [],
+    "temporal_caption": "The video begins with a dashcam view of a car driving fast along a scenic coastal mountain road, the asphalt rushing beneath and a sharp left curve sign visible on the left. The car navigates the curve at speed with the steep rocky cliff towering above and the ocean visible to the right. Around the 2-second mark, small rocks begin to dislodge from the cliff face ahead, quickly escalating into a significant rockfall with larger boulders tumbling down the mountainside and a dust cloud forming. By second 4, the driver reacts with emergency braking - the camera lurches forward dramatically from the sudden deceleration. The car comes to a complete stop between seconds 4 and 5, with the road ahead partially blocked by fallen rocks and debris, dust still settling in the air around the landslide zone.",
+    "audio_description": "Initially the sound of a car engine at moderate-high RPM and tire noise on asphalt with wind passing over the vehicle. Around second 2, the rumbling and cracking sounds of rocks breaking loose begin, growing louder with impacts of stones hitting the road surface and each other. Loud thuds and crashes as larger rocks land. At second 4, the sharp screech of tires braking hard on asphalt, followed by the settling sounds of smaller pebbles and dust, with distant rumbling of the remaining rockfall subsiding.",
+    "resolution": {
+      "W": 1280,
+      "H": 720
+    },
+    "aspect_ratio": "16,9",
+    "duration": "5s",
+    "fps": 24
+  }

assets/example_reasoning_prompt.json

@@ -0,0 +1,4 @@
+{
+  "prompt": "The task is to put flower into the red bottle. Generate a plan consisting of subtasks for accomplish the task.",
+  "max_tokens": 4096
+}

assets/negative_prompt.json

@@ -0,0 +1,131 @@
+{
+  "subjects": [
+    {
+      "description": "A subject suffering from extreme temporal instability — flickering outlines, morphing silhouette, and visible frame-to-frame discontinuities that make it look like a corrupted video file.",
+      "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.",
+      "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.",
+      "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.",
+      "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.",
+      "orientation": "Unnatural orientations that defy physics and spatial logic.",
+      "pose": "Poses snap between positions with no interpolation, creating a strobe-like effect.",
+      "action": "Motion is completely incoherent — limbs teleport between positions, movement appears as discrete jumps rather than continuous flow, and speed varies wildly between frames.",
+      "state_changes": "The subject flickers between completely different visual states every few frames, as if multiple renders were randomly interleaved.",
+      "clothing": "Garments phase in and out of visibility, textures swim across surfaces untethered to the geometry.",
+      "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.",
+      "gender": "",
+      "age": "",
+      "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.",
+      "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.",
+      "number_of_subjects": 0,
+      "number_of_arms": 0,
+      "number_of_legs": 0
+    },
+    {
+      "description": "A subject with catastrophically broken surface rendering — visible polygon edges, texture seams, UV mapping errors, and normal map inversions.",
+      "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.",
+      "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.",
+      "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.",
+      "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.",
+      "orientation": "Unnatural orientations that defy physics and spatial logic.",
+      "pose": "Static poses reveal texture swimming, moire patterns, and aliased edges on every contour.",
+      "action": "Movement reveals progressively worse texture breakdown — stretching, tearing, and resolution collapse in deformed regions.",
+      "state_changes": "Surface materials randomly swap between frames — metal becomes cloth, glass becomes stone, with no transition.",
+      "clothing": "Fabric renders as a flat solid color with no folds, wrinkles, or material response to form.",
+      "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.",
+      "gender": "",
+      "age": "",
+      "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.",
+      "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.",
+      "number_of_subjects": 0,
+      "number_of_arms": 0,
+      "number_of_legs": 0
+    },
+    {
+      "description": "A subject that appears composited into the scene rather than belonging to it — mismatched lighting direction, wrong color temperature, and edges that bleed into the background.",
+      "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.",
+      "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.",
+      "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.",
+      "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.",
+      "orientation": "Unnatural orientations that defy physics and spatial logic.",
+      "pose": "Proportions shift between frames — limbs lengthen and shorten, heads resize, creating an unsettling instability.",
+      "action": "Movement creates visible matting artifacts — green screen fringe, edge chattering, and translucent boundaries.",
+      "state_changes": "The subject randomly changes size, opacity, and z-depth relative to other scene elements.",
+      "clothing": "Clothing appears to be from a different render pass — different noise level, different motion blur kernel, different color space.",
+      "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.",
+      "gender": "",
+      "age": "",
+      "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.",
+      "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.",
+      "number_of_subjects": 0,
+      "number_of_arms": 0,
+      "number_of_legs": 0
+    }
+  ],
+  "background_setting": "A poorly rendered, flat background with visible seams, repeated textures, and inconsistent depth cues. The environment lacks volumetric depth and appears as a painted backdrop rather than a three-dimensional space. Vegetation looks like flat cutouts with no volumetric depth. The background appears to have been composited from multiple source materials at different resolutions, creating visible seams and edge artifacts where elements meet. Textures swim and shift across surfaces in a way that breaks the illusion of solidity — patterns drift laterally rather than staying anchored to the geometry they belong to. Background elements flicker in and out of existence between frames, particularly at the edges of the field of view. The rendering resolution is visibly lower for distant elements, creating a jarring transition between near and far objects. Cloud textures repeat obviously in the sky with visible tiling. Water surfaces lack proper reflection and refraction, appearing as flat animated textures. Fog and atmospheric effects pop in and out rather than smoothly transitioning. Trees and vegetation exhibit obvious LOD (level-of-detail) switching. Building facades have inconsistent window spacing and pattern repetition. The overall scene feels like a poorly assembled collage of individually rendered elements rather than a coherent whole.",
+  "lighting": {
+    "conditions": "Harsh, flat lighting with no natural variation. The scene appears uniformly lit as if by a single overhead fluorescent light, removing all sense of depth and atmosphere.",
+    "direction": "Inconsistent light sources — shadows point in multiple contradictory directions, breaking physical plausibility.",
+    "shadows": "Hard-edged, unrealistic shadows that pop in and out of existence between frames. Some objects cast no shadows while others have impossibly dark ones that don't animate smoothly with the object's motion. Shadow edges exhibit visible staircase aliasing artifacts. Shadow maps appear to have been rendered at extremely low resolution, creating blocky patterns. Self-shadowing on characters shows visible peter-panning artifacts where shadows detach from their source. Contact shadows between objects and the ground appear and disappear as objects move slightly. Shadow color is pure black with no ambient contribution, creating an unnaturally harsh contrast that flattens the image. Multiple shadow cascades have visible boundaries where resolution changes. The shadow rendering appears to be temporally unstable — even static objects have shadows that shimmer and crawl frame to frame, breaking the illusion of a stable light source.",
+    "illumination_effect": "No bounce light, no ambient occlusion, no subtle color interactions between surfaces. The scene looks like a poorly lit 3D render from the early 2000s."
+  },
+  "aesthetics": {
+    "composition": "Cluttered, poorly framed composition with no clear focal point. Important elements are cut off by the frame edges. The rule of thirds is completely ignored, leading to an unbalanced and visually unpleasant arrangement.",
+    "color_scheme": "Oversaturated, garish colors that clash violently. Color banding is visible in gradient areas. The overall palette feels artificial and digitally processed rather than natural.",
+    "mood_atmosphere": "Unsettling, uncanny atmosphere that fails to evoke any intended emotional response. The scene feels lifeless and sterile despite attempting to portray dynamic action.",
+    "patterns": "Visible tiling artifacts in textures, moiré patterns, and aliasing on edges."
+  },
+  "cinematography": {
+    "camera_motion": "Extremely shaky, unstable camera with visible rolling shutter artifacts. The motion is jerky and discontinuous, causing motion sickness and making the scene impossible to follow. The camera also executes an over-aggressive, unrelenting push-in: it dollies and optically zooms toward the subject far too fast — magnifying the scene roughly 2.5x to 6x within a few seconds — combined with a large, uncontrolled lateral drift, so the framing never settles. Instead of holding a stable composition, the viewpoint continuously accelerates forward and slides sideways, blowing past the intended framing into an unwanted extreme close-up by the end of the clip. This push-in is a distinct failure mode from the shake/jitter described above — it has little to no rotational wobble — a smooth but excessive, monotonic magnitude of forward dolly, zoom, and translation that never eases off, never holds, and never returns to a steady frame. On scenes with strong perspective lines (conveyor belts, quarries, aerial industrial sites, long corridors) the runaway zoom-and-translate makes the perspective rush inward, reading as erratic, claustrophobic, nauseating camerawork.",
+    "framing": "Poorly framed shots that cut off important elements and include unnecessary empty space. The shot also relentlessly tightens into an unintended extreme close-up as the camera barrels forward, cropping out the intended subject and context because the push-in never stops.",
+    "camera_angle": "Awkward, disorienting camera angles that provide no useful spatial information about the scene. The camera path exhibits visible mathematical artifacts suggesting simple interpolation between keyframes rather than natural camera operation. Camera motion is completely disconnected from the scene content — panning away from action, dollying during dialogue, and shaking during still moments. The camera appears to pass through solid objects occasionally. Zoom is applied digitally rather than optically, revealing progressively worse resolution. Camera motion exhibits non-physical acceleration profiles — instant starts and stops rather than smooth ease-in/ease-out. Rolling shutter simulation is applied inconsistently, present in some frames but not others. The camera occasionally exhibits impossible motion like teleporting between positions. Virtual camera stabilization creates an uncanny floating sensation disconnected from any physical camera rig.",
+    "depth_of_field": "Uniform focus throughout, creating a flat, documentary-like appearance with no cinematic depth separation.",
+    "focus": "Soft, out-of-focus imagery with visible chromatic aberration and lens distortion that was not corrected in post-processing.",
+    "lens_focal_length": "Inappropriate focal length causing barrel distortion and unnatural perspective compression."
+  },
+  "style_medium": "Low quality compressed digital video with visible encoding artifacts",
+  "artistic_style": "Amateur, unpolished with inconsistent visual style",
+  "context": "A poorly produced video with numerous technical and artistic flaws that detract from any intended narrative or visual impact.",
+  "actions": [
+    {
+      "time": "0:00-0:02",
+      "description": "As the camera starts moving, rigid foreground objects (cars, signs) warp and change shape; straight painted road lines bend and smear."
+    },
+    {
+      "time": "0:02-0:04",
+      "description": "A clearly-present background object vanishes without occlusion while a new object pops into the frame edge; feet/wheels slide over the ground without contact."
+    },
+    {
+      "time": "0:04-0:06",
+      "description": "A hand manipulating an object fuses with it into a morphing blob; poured/chopped material changes mass, portions disappearing and multiplying."
+    },
+    {
+      "time": "0:06-0:08",
+      "description": "Thin structures (wires, lane lines, foliage) flicker and shimmer; object count and world geometry no longer match the opening frame."
+    },
+    {
+      "time": "0:00-0:06",
+      "description": "Throughout, the camera pushes in far too aggressively — a fast continuous forward dolly and zoom (about 2.5x-6x) with a large lateral drift and no rotational shake — never holding a steady framing and ending in an unwanted extreme close-up."
+    }
+  ],
+  "text_and_signage_elements": [],
+  "segments": [
+    {
+      "segment_index": 0,
+      "time_range": "0:00-0:04",
+      "description": "Camera-motion segment where rigid-body stability collapses: vehicles, signage, and road markings morph, warp, and slide as the viewpoint translates.",
+      "key_changes": "World geometry un-locks — ground plane and lane lines drift; a background object vanishes and another appears without cause.",
+      "camera": "Forward-moving camera that drives rigid-object deformation; simultaneously an over-fast forward dolly-and-zoom with lateral drift that magnifies the scene far too quickly and never holds a stable frame."
+    },
+    {
+      "segment_index": 1,
+      "time_range": "0:04-0:08",
+      "description": "Manipulation and material-transfer segment where hand-object contact fails.",
+      "key_changes": "Hand and handled object merge into a morphing blob; transferred food/liquid violates mass conservation, portions vanishing and multiplying; thin high-frequency structures flicker.",
+      "camera": "Unstable framing with incorrect motion blur; the excessive push-in continues, sliding sideways and blowing past the composition into an extreme close-up."
+    }
+  ],
+  "transitions": [],
+  "temporal_caption": "The clip opens with objects appearing sharp and stable, but as the camera moves forward every rigid object begins to betray its geometry: parked cars warp and bulge, painted lane markings smear and slide across the asphalt, stop signs and poles bend and reshape frame to frame. Midway, a vehicle that was clearly present silently vanishes while a new object pops into existence at the frame edge without cause. Hands manipulating objects fuse with them into a morphing blob; poured or chopped food changes mass and volume, portions disappearing and multiplying. Thin structures — wires, foliage, railings, lane lines — flicker and shimmer throughout. Feet, wheels, and paws glide over the ground without firm contact. By the end, world geometry, object count, and material mass are all inconsistent with how the clip began. Over the whole clip the camera also barrels forward far too fast: a smooth, unrelenting push-in — fast forward dolly plus optical zoom, roughly 2.5x to 6x magnification in a few seconds — paired with a large lateral drift, with essentially no rotational shake. The framing never stabilizes; the viewpoint continuously accelerates past the intended composition and collapses into an uncomfortable extreme close-up by the end, which on industrial and aerial scenes with strong perspective lines reads as erratic, disorienting camerawork.",
+  "audio_description": "",
+  "physical_realism": "No adherence to physical laws. Objects defy gravity, pass through solid surfaces, and change mass and momentum without cause. Fluid dynamics, cloth simulation, and rigid body physics are all fundamentally broken. Furthermore, conservation of energy is violated as objects gain or lose kinetic energy spontaneously. Elastic collisions produce inelastic results and vice versa. Surface friction is inconsistent — objects slide on rough surfaces while sticking to smooth ones. Air resistance appears to affect only some objects while others move through the atmosphere unimpeded. Rigid bodies morph under camera motion; objects slide without ground contact; mass is not conserved in transfer; objects enter/vanish breaking permanence; thin structures flicker."
+}

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 %}

config.json

@@ -0,0 +1,67 @@
+{
+  "allow_patterns_overrides": [
+    "*/*.safetensors"
+  ],
+  "architectures": [
+    "Cosmos3EdgeForConditionalGeneration"
+  ],
+  "image_token_id": 19,
+  "model_type": "cosmos3_edge",
+  "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": "4.57.1",
+  "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
+}

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": "4.57.1"
+}

model.safetensors.index.json

@@ -0,0 +1,705 @@
+{
+  "metadata": {
+    "total_size": 4871240160
+  },
+  "weight_map": {
+    "embed_tokens.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.0.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.1.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.10.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.11.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.12.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.13.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.14.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.15.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.16.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.17.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.18.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.19.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.2.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.20.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.21.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.input_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.22.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.22.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.23.input_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.23.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.input_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.24.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.input_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.25.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.input_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.26.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.input_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.27.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "layers.3.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.3.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.4.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.5.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.6.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.7.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.8.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.input_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.mlp.down_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.mlp.up_proj.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.post_attention_layernorm.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.k_norm_und_for_gen.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.to_k.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.to_out.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.to_q.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "layers.9.self_attn.to_v.weight": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors",
+    "lm_head.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors",
+    "model.projector.linear_fc1.bias": "vision_encoder/model.safetensors",
+    "model.projector.linear_fc1.weight": "vision_encoder/model.safetensors",
+    "model.projector.linear_fc2.bias": "vision_encoder/model.safetensors",
+    "model.projector.linear_fc2.weight": "vision_encoder/model.safetensors",
+    "model.projector.norm.bias": "vision_encoder/model.safetensors",
+    "model.projector.norm.weight": "vision_encoder/model.safetensors",
+    "model.visual.embeddings.patch_embedding.bias": "vision_encoder/model.safetensors",
+    "model.visual.embeddings.patch_embedding.weight": "vision_encoder/model.safetensors",
+    "model.visual.embeddings.position_embedding.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.0.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.1.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.10.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.11.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.12.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.13.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.14.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.15.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.16.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.17.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.18.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.19.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.2.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.20.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.21.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.22.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.23.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.24.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.25.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.26.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.3.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.4.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.5.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.6.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.7.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.8.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.layer_norm1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.layer_norm1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.layer_norm2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.layer_norm2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.mlp.fc1.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.mlp.fc1.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.mlp.fc2.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.mlp.fc2.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.self_attn.k_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.self_attn.k_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.self_attn.out_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.self_attn.out_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.self_attn.q_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.self_attn.q_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.self_attn.v_proj.bias": "vision_encoder/model.safetensors",
+    "model.visual.encoder.layers.9.self_attn.v_proj.weight": "vision_encoder/model.safetensors",
+    "model.visual.post_layernorm.bias": "vision_encoder/model.safetensors",
+    "model.visual.post_layernorm.weight": "vision_encoder/model.safetensors",
+    "norm.weight": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors"
+  }
+}

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,60 @@
+{
+  "_blocks_class_name": "Cosmos3OmniBlocks",
+  "_class_name": "Cosmos3OmniModularPipeline",
+  "_diffusers_version": "0.40.0.dev0",
+  "is_distilled": false,
+  "default_use_system_prompt": false,
+  "use_native_flow_schedule": true,
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler",
+    {
+      "pretrained_model_name_or_path": "nvidia/Cosmos3-Edge",
+      "subfolder": "scheduler",
+      "type_hint": [
+        "diffusers",
+        "UniPCMultistepScheduler"
+      ],
+      "variant": null
+    }
+  ],
+  "text_tokenizer": [
+    "transformers",
+    "PreTrainedTokenizerFast",
+    {
+      "pretrained_model_name_or_path": "nvidia/Cosmos3-Edge",
+      "subfolder": "text_tokenizer",
+      "type_hint": [
+        "transformers",
+        "PreTrainedTokenizerFast"
+      ],
+      "variant": null
+    }
+  ],
+  "transformer": [
+    "diffusers",
+    "Cosmos3OmniTransformer",
+    {
+      "pretrained_model_name_or_path": "nvidia/Cosmos3-Edge",
+      "subfolder": "transformer",
+      "type_hint": [
+        "diffusers",
+        "Cosmos3OmniTransformer"
+      ],
+      "variant": null
+    }
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan",
+    {
+      "pretrained_model_name_or_path": "nvidia/Cosmos3-Edge",
+      "subfolder": "vae",
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLWan"
+      ],
+      "variant": null
+    }
+  ]
+}

preprocessor_config.json

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

processor_config.json

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

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": 1.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": 200.0,
+  "sigma_min": 0.147,
+  "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": true
+}

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

text_tokenizer/tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "is_local": true,
+  "local_files_only": false,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 262144,
+  "pad_token": "<|im_end|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>"
+}

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,46 @@
+{
+  "_class_name": "Cosmos3OmniTransformer",
+  "_diffusers_version": "0.40.0.dev0",
+  "action_dim": 64,
+  "action_gen": true,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "backbone_type": "cosmos3_edge_nemotron_dense",
+  "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,
+  "sound_dim": null,
+  "sound_gen": false,
+  "sound_latent_fps": 25,
+  "temporal_compression_factor": 4,
+  "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/diffusion_pytorch_model-00001-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:116134e5492c514df79a2f07214e577c397aa961387dbe1c5dda488e7994ff6b
+size 5000040072

transformer/diffusion_pytorch_model-00002-of-00002.safetensors

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

transformer/diffusion_pytorch_model.safetensors.index.json

@@ -0,0 +1,556 @@
+{
+  "metadata": {
+    "total_size": 6739314048
+  },
+  "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-00002-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-00002-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-00002-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-00002-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-00002-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"
+  }
+}