copy from personal repo

LICENSE

@@ -0,0 +1,201 @@
+                                Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,115 @@
+## How to Use
+### Example outputs can be found in assets folder
+
+```python
+
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "Nvidia-CMU25/ARVideo2WorldGeneration", 
+    cache_dir="./cache", 
+    trust_remote_code=True,
+    
+    input_type = "text_and_image",
+    num_input_frames = 1,
+    prompt = "A video recorded from a moving vehicle's perspective, capturing roads, buildings, landscapes, and changing weather and lighting conditions." ,
+    input_image_or_video_path = "AutoregressiveVideo2WorldGeneration/cosmos1/models/autoregressive/assets/v1p0/input.jpg",
+    video_save_name = "diffusion_decoder_image_output",
+    ar_model_dir = "Cosmos-1.0-Autoregressive-5B-Video2World",
+    
+    # input_type = "text_and_video",
+    # num_input_frames = 9,
+    # prompt = "A video recorded from a moving vehicle's perspective, capturing roads, buildings, landscapes, and changing weather and lighting conditions." ,
+    # input_image_or_video_path = "AutoregressiveVideo2WorldGeneration/cosmos1/models/autoregressive/assets/v1p0/input.mp4",
+    # video_save_name = "diffusion_decoder_video_output",
+    
+    # turn on offloading on a low GPU memory machine:
+    disable_diffusion_decoder=False,
+    offload_guardrail_models=True,
+    offload_diffusion_decoder=True,
+    offload_network=True,
+    offload_tokenizer=True,
+    offload_text_encoder_model=True,
+)
+
+model()
+```
+
+
+![Cosmos Logo](assets/cosmos-logo.png)
+
+--------------------------------------------------------------------------------
+### [Website](https://www.nvidia.com/en-us/ai/cosmos/) | [HuggingFace](https://huggingface.co/collections/nvidia/cosmos-6751e884dc10e013a0a0d8e6) | [GPU-free Preview](https://build.nvidia.com/explore/discover) | [Paper](https://arxiv.org/abs/2501.03575) | [Paper Website](https://research.nvidia.com/labs/dir/cosmos1/)
+
+[NVIDIA Cosmos](https://www.nvidia.com/cosmos/) is a developer-first world foundation model platform designed to help Physical AI developers build their Physical AI systems better and faster. Cosmos contains
+
+1. pre-trained models, available via [Hugging Face](https://huggingface.co/collections/nvidia/cosmos-6751e884dc10e013a0a0d8e6) under the [NVIDIA Open Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/) that allows commercial use of the models for free
+2. training scripts under the [Apache 2 License](https://www.apache.org/licenses/LICENSE-2.0), offered through [NVIDIA Nemo Framework](https://github.com/NVIDIA/NeMo) for post-training the models for various downstream Physical AI applications
+
+Details of the platform is described in the [Cosmos paper](https://research.nvidia.com/publication/2025-01_cosmos-world-foundation-model-platform-physical-ai). Preview access is avaiable at [build.nvidia.com](https://build.nvidia.com).
+
+## Key Features
+
+- [Pre-trained Diffusion-based world foundation models](cosmos1/models/diffusion/README.md) for Text2World and Video2World generation where a user can generate visual simulation based on text prompts and video prompts.
+- [Pre-trained Autoregressive-based world foundation models](cosmos1/models/autoregressive/README.md) for Video2World generation where a user can generate visual simulation based on video prompts and optional text prompts.
+- [Video tokenizers](https://github.com/NVIDIA/Cosmos-Tokenizer) for tokenizing videos into continuous tokens (latent vectors) and discrete tokens (integers) efficiently and effectively.
+- Video curation pipeline for building your own video dataset. [Coming soon]
+- [Post-training scripts](cosmos1/models/POST_TRAINING.md) via NeMo Framework to post-train the pre-trained world foundation models for various Physical AI setup.
+- Pre-training scripts via NeMo Framework for building your own world foundation model. [[Diffusion](https://github.com/NVIDIA/NeMo/tree/main/nemo/collections/diffusion)] [[Autoregressive](https://github.com/NVIDIA/NeMo/tree/main/nemo/collections/multimodal_autoregressive)] [[Tokenizer](https://github.com/NVIDIA/NeMo/tree/main/nemo/collections/diffusion/vae)].
+
+## Model Family
+
+| Model name | Description | Try it out |
+|------------|----------|----------|
+| [Cosmos-1.0-Diffusion-7B-Text2World](https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Text2World) | Text to visual world generation  | [Inference](cosmos1/models/diffusion/README.md)   |
+| [Cosmos-1.0-Diffusion-14B-Text2World](https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-14B-Text2World) | Text to visual world generation  | [Inference](cosmos1/models/diffusion/README.md)   |
+| [Cosmos-1.0-Diffusion-7B-Video2World](https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Video2World) | Video + Text based future visual world generation  | [Inference](cosmos1/models/diffusion/README.md)   |
+| [Cosmos-1.0-Diffusion-14B-Video2World](https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-14B-Video2World) | Video + Text based future visual world generation  | [Inference](cosmos1/models/diffusion/README.md)   |
+| [Cosmos-1.0-Autoregressive-4B](https://huggingface.co/nvidia/Cosmos-1.0-Autoregressive-4B) | Future visual world generation  | [Inference](cosmos1/models/autoregressive/README.md)   |
+| [Cosmos-1.0-Autoregressive-12B](https://huggingface.co/nvidia/Cosmos-1.0-Autoregressive-12B) | Future visual world generation  | [Inference](cosmos1/models/autoregressive/README.md)   |
+| [Cosmos-1.0-Autoregressive-5B-Video2World](https://huggingface.co/nvidia/Cosmos-1.0-Autoregressive-5B-Video2World) | Video + Text based future visual world generation | [Inference](cosmos1/models/autoregressive/README.md)   |
+| [Cosmos-1.0-Autoregressive-13B-Video2World](https://huggingface.co/nvidia/Cosmos-1.0-Autoregressive-13B-Video2World) | Video + Text based future visual world generation | [Inference](cosmos1/models/autoregressive/README.md)   |
+| [Cosmos-1.0-Guardrail](https://huggingface.co/nvidia/Cosmos-1.0-Guardrail) | Guardrail contains pre-Guard and post-Guard for safe use | Embedded in model inference scripts |
+
+## Example Usage
+
+### Inference
+
+Follow the [Cosmos Installation Guide](INSTALL.md) to setup the docker. For inference with the pretrained models, please refer to [Cosmos Diffusion Inference](cosmos1/models/diffusion/README.md) and [Cosmos Autoregressive Inference](cosmos1/models/autoregressive/README.md).
+
+The code snippet below provides a gist of the inference usage.
+
+```bash
+PROMPT="A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves. \
+The robot's metallic body gleams under the bright, even lighting, highlighting its futuristic design and intricate joints. \
+A glowing blue light emanates from its chest, adding a touch of advanced technology. The background is dominated by rows of boxes, \
+suggesting a highly organized storage system. The floor is lined with wooden pallets, enhancing the industrial setting. \
+The camera remains static, capturing the robot's poised stance amidst the orderly environment, with a shallow depth of \
+field that keeps the focus on the robot while subtly blurring the background for a cinematic effect."
+
+# Example using 7B model
+PYTHONPATH=$(pwd) python cosmos1/models/diffusion/inference/text2world.py \
+    --checkpoint_dir checkpoints \
+    --diffusion_transformer_dir Cosmos-1.0-Diffusion-7B-Text2World \
+    --prompt "$PROMPT" \
+    --offload_prompt_upsampler \
+    --video_save_name Cosmos-1.0-Diffusion-7B-Text2World
+```
+
+<video src="https://github.com/user-attachments/assets/db7bebfe-5314-40a6-b045-4f6ce0a87f2a">
+  Your browser does not support the video tag.
+</video>
+
+We also offer [multi-GPU inference](cosmos1/models/diffusion/nemo/inference/README.md) support for Diffusion Text2World WFM models through NeMo Framework.
+
+### Post-training
+
+NeMo Framework provides GPU accelerated post-training with general post-training for both [diffusion](cosmos1/models/diffusion/nemo/post_training/README.md) and [autoregressive](cosmos1/models/autoregressive/nemo/post_training/README.md) models, with other types of post-training coming soon.
+
+## License and Contact
+
+This project will download and install additional third-party open source software projects. Review the license terms of these open source projects before use.
+
+NVIDIA Cosmos source code is released under the [Apache 2 License](https://www.apache.org/licenses/LICENSE-2.0).
+
+NVIDIA Cosmos models are released under the [NVIDIA Open Model License](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). For a custom license, please contact [[email protected]](mailto:[email protected]).

RELEASE.md

@@ -0,0 +1,7 @@
+# Release Cadence
+
+
+| Version | Description | Date |
+|------------|----------|----------|
+| [v1.0](release_notes/v0p1.md) | Initial diffusion and autoregressive WFMs release | 2025-01-06 |
+| [v0.1](release_notes/v0p1.md) | Initial tokenizer release | 2024-11-06 |

ar_config_base_model.py

@@ -0,0 +1,118 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import attrs
+
+from .ar_config_base_tokenizer import TokenizerConfig
+
+
+@attrs.define
+class ModelConfig:
+    """
+    A class to hold model configuration arguments.
+
+    Args:
+        dim (int): The dimensionality of the input and output of each transformer block.
+        n_layers (int): Number of layers in the transformer.
+        n_heads (int): Number of attention heads.
+        n_kv_heads (Optional[int]): Number of key-value heads. If None, defaults to n_heads. Note: this is equivalent to
+            `num_gqa_groups` in TransformerEngine, where GQA means Grouped Query Attention.
+        head_dim (Optional[int]): Dimensionality of each head. If None, defaults to dim // n_heads.
+        vocab_size (int): Vocabulary size.
+        ffn_hidden_size (int): Hidden size for feedforward network.
+        norm_eps (float): Epsilon value for normalization.
+        rope_theta (float): Theta value for rotary positional embeddings.
+        apply_abs_pos_emb (bool): Whether to apply absolute position embeddings.
+        max_batch_size (int): Maximum batch size for inference.
+        max_seq_len (int): Maximum sequence length for input text.
+        fuse_qkv (bool): Whether to fuse QKV in attention. Defaults to True.
+        causal_mask (bool): Whether to use causal mask. Defaults to True.
+        norm_type (str): Type of normalization layer. Choices: "rmsnorm", "fused_rmsnorm", "layernorm", "np_layernorm".
+        precision (str): Data type for the model.
+        use_qk_normalization (bool): Whether to enable QK normalization.
+        ckpt_dir (str): Checkpoint directory.
+        ckpt_path (str): Checkpoint path.
+        apply_yarn (Optional[bool]): Whether to apply YaRN (long-context extension).
+        yarn_scale (Optional[float]): Scale factor for YaRN.
+        yarn_beta_fast (Optional[int]): Beta fast variable for YaRN (i.e., low_freq_factor in Llama 3.1 RoPE scaling code)
+        yarn_beta_slow (Optional[int]): Beta slow variable for YaRN (i.e., high_freq_factor in Llama 3.1 RoPE scaling code)
+        original_seq_len (Optional[int]): Original sequence length.
+        vision_encoder (Optional[str]): Vision encoder name.
+        mm_projector (Optional[str]): Multi-modal projector name.
+        vision_encoder_in_channels (Optional[int]): Number of channels in the input image for the vision encoder. Default is 3, you can specify to int larger than 3. E.g. if you have 4-channel images with the last channel as the alpha channel, set this to 4.
+        rope_dim (Optional[str]): Dimensionality of the RoPE. Choices: "1D", "3D".
+        pytorch_rope_version (Optional[str]): Version of the PyTorch RoPE implementation. Choices: "v1", "v2".
+        original_latent_shape (Optional[list]): Original shape of the latent tensor needed for rope extension.
+        pad_to_multiple_of (Optional[int]): Pad the position embedding to a multiple of this value.
+        vision_encoder_in_channels (Optional[int]): Number of channels in the input image for the vision encoder. Default is 3.
+        insert_cross_attn (bool): Whether to insert the cross-attention layers after each multi-head self-attention (MSA) layer.
+        insert_cross_attn_every_k_layers (int): Insert cross-attention layers every k TransformerLayers.
+        context_dim (Optional[int]): The dimensionality of cross-attention embedding, e.g., T5 embed feature dim.
+        num_video_frames (Optional[int]): Number of video frames.
+        video_height (Optional[int]): Raw video pixel height dimension.
+        video_width (Optional[int]): Raw video pixel width dimension.
+        video_latent_shape (Optional[list]): Video tokenizer output dimension, in (T,H,W).
+    """
+
+    dim: int = attrs.field(default=4096)
+    n_layers: int = attrs.field(default=32)
+    n_heads: int = attrs.field(default=32)
+    n_kv_heads: Optional[int] = attrs.field(default=8)
+    head_dim: Optional[int] = attrs.field(default=None)
+    vocab_size: int = attrs.field(default=128256)
+    ffn_hidden_size: int = attrs.field(default=14336)
+    norm_eps: float = attrs.field(default=1e-5)
+    rope_theta: float = attrs.field(default=500000)
+    apply_abs_pos_emb: bool = attrs.field(default=False)
+    max_batch_size: int = attrs.field(default=1)
+    max_seq_len: int = attrs.field(default=8192)
+    fuse_qkv: bool = attrs.field(default=False)
+    causal_mask: bool = attrs.field(default=True)
+    norm_type: str = attrs.field(default="rmsnorm")
+    precision: str = attrs.field(default="bfloat16")
+    use_qk_normalization: bool = False
+    tokenizer: Optional[TokenizerConfig] = None
+    ckpt_dir: Optional[str] = attrs.field(default=None)
+    ckpt_path: Optional[str] = attrs.field(
+        default=None
+    )  # If not None, load the model from this path instead of ckpt_dir
+    apply_yarn: Optional[bool] = attrs.field(default=False)
+    yarn_scale: Optional[float] = attrs.field(default=None)
+    yarn_beta_fast: Optional[int] = attrs.field(default=None)
+    yarn_beta_slow: Optional[int] = attrs.field(default=None)
+    original_seq_len: Optional[int] = attrs.field(default=None)
+    vision_encoder: Optional[str] = attrs.field(default=None)
+    vision_encoder_in_channels: Optional[int] = attrs.field(default=3)
+    mm_projector: Optional[str] = attrs.field(default=None)
+    rope_dim: Optional[str] = attrs.field(default="1D")
+    pytorch_rope_version: Optional[str] = attrs.field(default="v2")
+    original_latent_shape: Optional[list] = None
+    pad_to_multiple_of: Optional[int] = None
+    vision_encoder_in_channels: Optional[int] = attrs.field(default=3)
+    insert_cross_attn: bool = False
+    insert_cross_attn_every_k_layers: int = 1
+    context_dim: Optional[int] = attrs.field(default=1024)
+    # For video training
+    num_video_frames: Optional[int] = None
+    # Raw video pixel dimension
+    video_height: Optional[int] = None
+    video_width: Optional[int] = None
+    # Video tokenizer output dimension, in (T,H,W), it's computed by num_video_frames/temporal_compress_factor, video_height/spatial_compression_fact, video_width/spatial_compression_fact
+    video_latent_shape: Optional[list] = None
+
+    def __getitem__(self, item):
+        return getattr(self, item)

ar_config_base_model_config.py

@@ -0,0 +1,421 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from typing import Callable, List, Optional
+
+from .ar_config_base_model import ModelConfig
+from .ar_config_base_tokenizer import (
+    TextTokenizerConfig,
+    TokenizerConfig,
+    VideoTokenizerConfig,
+    create_discrete_video_fsq_tokenizer_state_dict_config,
+)
+from .ar_tokenizer_image_text_tokenizer import ImageTextTokenizer
+from .ar_tokenizer_text_tokenizer import TextTokenizer
+from .log import log
+from .lazy_config_init import LazyCall as L
+
+# Common architecture specifications
+BASE_CONFIG = {"n_kv_heads": 8, "norm_type": "rmsnorm", "norm_eps": 1e-5, "ffn_hidden_size": 14336}
+COSMOS_ARCHITECTURES = {
+    "4b": {
+        "n_layers": 16,
+        "dim": 4096,
+        "n_heads": 32,
+    },
+    "12b": {
+        "n_layers": 40,
+        "dim": 5120,
+        "n_heads": 32,
+        "head_dim": 128,
+    },
+}
+
+COSMOS_YARN_CONFIG = {
+    "original_latent_shape": [3, 40, 64],
+    "apply_yarn": True,
+    "yarn_beta_fast": 4,
+    "yarn_beta_slow": 1,
+    "yarn_scale": 2,
+}
+
+# Llama3 architecture specifications for different model sizes
+LLAMA3_ARCHITECTURES = {
+    "8b": {
+        "n_layers": 32,
+        "dim": 4096,
+        "n_heads": 32,
+        "ffn_hidden_size": 14336,
+    },
+}
+# Llama3.1 uses YaRN for long context support (context of 128k tokens)
+LLAMA_YARN_CONFIG = {
+    "apply_yarn": True,
+    "yarn_scale": 8,
+    "yarn_beta_fast": 4,
+    "yarn_beta_slow": 1,
+}
+
+# Mistral architecture specifications for different model sizes
+MISTRAL_ARCHITECTURES = {
+    "12b": {
+        "n_layers": 40,
+        "dim": 5120,
+        "n_heads": 32,
+        "ffn_hidden_size": 14336,
+        "head_dim": 128,
+    },
+}
+
+PIXTRAL_VISION_ARCHITECTURES = {
+    "12b": {"vision_encoder": "pixtral-12b-vit", "mm_projector": "mlp"},
+}
+
+
+def get_model_arch_specs(model_size: str, model_family: str = "mistral", pretrained: bool = False) -> dict:
+    """
+    Get the model architecture specifications for the given model size, model family and pretrained status.
+
+    Args:
+        model_size (str): Model size. Choices: "1b", "3b", "4b", "7b", etc.
+        model_family (str): Model family. Choices: "llama", "llama3", "llama3.1", "mistral"
+        pretrained (bool): Whether to load pretrained weights.
+
+    Returns:
+        dict: A dictionary containing the model architecture specifications.
+    """
+    arch_specs = copy.deepcopy(BASE_CONFIG)
+    model_size = model_size.lower()
+    if model_family.startswith("cosmos"):
+        arch_specs.update(COSMOS_ARCHITECTURES[model_size])
+    elif model_family.startswith("llama"):
+        arch_specs.update(LLAMA3_ARCHITECTURES[model_size])
+    elif model_family in ["mistral", "pixtral"]:
+        arch_specs.update(MISTRAL_ARCHITECTURES[model_size])
+        if model_family == "pixtral":
+            arch_specs.update(PIXTRAL_VISION_ARCHITECTURES[model_size])
+    else:
+        raise ValueError(f"Model family {model_family} is not supported.")
+
+    if pretrained:
+        if model_family == "cosmos":
+            if model_size == "12b":
+                arch_specs.update(COSMOS_YARN_CONFIG)
+                log.debug(f"Using YaRN for RoPE extension with config: {COSMOS_YARN_CONFIG}")
+            else:
+                pass
+        elif model_family in ["llama", "llama3"]:
+            pretrained_specs = {
+                "rope_theta": 500000,
+                "max_seq_len": 8192,
+                "vocab_size": 128256,
+            }
+            arch_specs.update(pretrained_specs)
+        elif model_family == "llama3.1":
+            pretrained_specs = {
+                "rope_theta": 500000,
+                "max_seq_len": 131072,
+                "original_seq_len": 8192,
+                "vocab_size": 128256,
+                **LLAMA_YARN_CONFIG,
+            }
+            arch_specs.update(pretrained_specs)
+        elif model_family == "mistral":
+            assert model_size == "12b", "We only support Mistral-Nemo-12B model."
+            pretrained_specs = {
+                "rope_theta": 1000000,
+                "max_seq_len": 128000,
+                "vocab_size": 131072,
+            }
+            arch_specs.update(pretrained_specs)
+        elif model_family == "pixtral":
+            assert model_size == "12b", "We only support Pixtral 12B model."
+            pretrained_specs = {"rope_theta": 1000000000, "max_seq_len": 128000, "vocab_size": 131072}
+            arch_specs.update(pretrained_specs)
+        else:
+            raise ValueError(f"Model family {model_family} doesn't have a pretrained config.")
+
+    return arch_specs
+
+
+def create_text_model_config(
+    model_ckpt_path: str,
+    tokenizer_path: str,
+    model_family: str = "mistral",
+    model_size: str = "12b",
+    is_instruct_model: bool = True,
+    max_seq_len: int = None,
+    max_batch_size: int = 1,
+    rope_dim: str = "1D",
+    add_special_tokens: bool = True,
+    pytorch_rope_version: str = None,
+) -> dict:
+    """Create a text model for training or inference.
+    Args:
+        model_ckpt_path (str): Path to the model checkpoint.
+        tokenizer_path (str): Path to the tokenizer folder.
+        model_family (str): Model family. Choices: "llama", "llama3", "llama3.1", "mistral".
+        model_size (str): Model size. Choices: "1b", "3b", "4b", "7b", "8b", "72b", etc.
+        is_instruct_model (bool): Whether the model is an instruct model.
+        inference (bool): Whether to create the model for inference.
+        max_seq_len (int): Maximum sequence length.
+        max_batch_size (int): Maximum batch size.
+        rope_dim (str): RoPE dimension. Choices: "1D", "3D".
+        add_special_tokens (bool): Whether to add special tokens.
+    Returns:
+        dict: A dictionary containing the model configuration, which can be used to instantiate the model object.
+    """
+    # Model size specific parameters
+    model_arch_specs = get_model_arch_specs(model_family=model_family, model_size=model_size, pretrained=True)
+    if max_seq_len is not None:
+        # Override the max_seq_len if provided
+        model_arch_specs["max_seq_len"] = max_seq_len
+    if pytorch_rope_version is not None:
+        model_arch_specs["pytorch_rope_version"] = pytorch_rope_version
+    model_config = ModelConfig(
+        max_batch_size=max_batch_size,
+        precision="bfloat16",
+        ckpt_path=model_ckpt_path,
+        use_qk_normalization=False,
+        rope_dim=rope_dim,
+        **model_arch_specs,
+    )
+
+    tokenizer_config = TokenizerConfig(
+        text_tokenizer=TextTokenizerConfig(
+            config=L(TextTokenizer)(
+                model_family=model_family,
+                is_instruct_model=is_instruct_model,
+                local_path=tokenizer_path,
+            ),
+            data_key="text",
+            tokenizer_offset=model_config.vocab_size,
+            tokenize_here=False,
+            vocab_size=model_config.vocab_size,
+        ),
+        seq_len=model_config.max_seq_len,
+        training_type="text_only",
+        add_special_tokens=add_special_tokens,
+    )
+    return model_config, tokenizer_config
+
+
+def create_vision_language_model_config(
+    model_ckpt_path: str,
+    tokenizer_ckpt_path: str,
+    model_family: str = "pixtral",
+    model_size: str = "12b",
+    is_instruct_model: bool = True,
+    max_batch_size: int = 1,
+    rope_dim: str = "1D",
+    add_special_tokens: bool = True,
+    max_seq_len: int = None,
+    vision_encoder_in_channels: int = 3,
+    fuse_qkv: bool = False,
+    pytorch_rope_version: str = None,
+) -> dict:
+    """Create a vision-language model for training or inference.
+    Args:
+        model_ckpt_path (str): Path to the model checkpoint.
+        tokenizer_ckpt_path (str): Path to the tokenizer checkpoint.
+        model_family (str): Model family. Choices: "pixtral".
+        model_size (str): Model size. Choices: "12b".
+        is_instruct_model (bool): Whether the model is an instruct model.
+        rope_dim (str): RoPE dimension. Choices: "1D".
+        add_special_tokens (bool): Whether to add special tokens.
+        max_seq_len (int): Maximum sequence length.
+        vision_encoder_in_channels (int): Number of channels in the input image for the vision encoder. Default is 3, you can specify to int larger than 3. E.g. if you have 4 channel images where last channel is binary mask, set this to 4.
+        fuse_qkv (bool): Whether to fuse the QKV linear layers.
+    Returns:
+        dict: A dictionary containing the model configuration, which can be used to instantiate the model object.
+    """
+    # Model size specific parameters
+    model_arch_specs = get_model_arch_specs(model_family=model_family, model_size=model_size, pretrained=True)
+    if max_seq_len is not None:
+        # Override the max_seq_len if provided
+        model_arch_specs["max_seq_len"] = max_seq_len
+    if pytorch_rope_version is not None:
+        model_arch_specs["pytorch_rope_version"] = pytorch_rope_version
+
+    model_config = ModelConfig(
+        max_batch_size=max_batch_size,
+        precision="bfloat16",
+        ckpt_path=model_ckpt_path,
+        use_qk_normalization=False,
+        rope_dim=rope_dim,
+        vision_encoder_in_channels=vision_encoder_in_channels,
+        fuse_qkv=fuse_qkv,
+        **model_arch_specs,
+    )
+    # Vision-language tokenizer
+    tokenizer_config = TokenizerConfig(
+        text_tokenizer=TextTokenizerConfig(
+            config=L(ImageTextTokenizer)(
+                model_family=model_family,
+                is_instruct_model=is_instruct_model,
+                image_processor_path=tokenizer_ckpt_path,
+                tokenizer_path=tokenizer_ckpt_path,
+            ),
+            data_key="image_text_interleaved",
+            tokenizer_offset=model_config.vocab_size,
+            tokenize_here=False,
+            vocab_size=model_config.vocab_size,
+        ),
+        seq_len=model_config.max_seq_len,
+        training_type="image_text_interleaved",
+        add_special_tokens=add_special_tokens,
+    )
+    return model_config, tokenizer_config
+
+
+def create_video2world_model_config(
+    model_ckpt_path: str,
+    tokenizer_ckpt_path: str,
+    model_family: str = "cosmos",
+    model_size: str = "4b",
+    pixel_chunk_duration: int = 9,
+    num_video_frames: int = 36,
+    compression_ratio: List[int] = [8, 16, 16],
+    original_seq_len: int = 8192,
+    num_condition_latents_t: int = 1,
+    num_tokens_to_ignore: int = -1,
+    batch_size: int = 2,
+    video_tokenizer_config_creator: Callable = create_discrete_video_fsq_tokenizer_state_dict_config,
+    rope_dim: str = "3D",
+    add_special_tokens: bool = True,
+    video_height: int = 384,
+    video_width: int = 640,
+    use_qk_normalization: bool = True,
+    insert_cross_attn: bool = False,
+    insert_cross_attn_every_k_layers: int = 1,
+    context_dim: int = 1024,
+    training_type: str = "video_to_video",
+    pad_to_multiple_of: Optional[int] = 64,
+    vocab_size: int = 64000,
+    apply_abs_pos_emb: bool = False,
+) -> dict:
+    """Create a video-to-world model config.
+    Args:
+        model_family (str): Model family. Choices: "llama", "llama3", "llama3.1", "mistral".
+        model_size (str): Model size. Choices: "1b", "8b", "3b".
+        pixel_chunk_duration (int): Number of frames in each chunk.
+        num_video_frames (int): Number of video frames.
+        compression_ratio (List[int]): Compression ratio for the video frames. Choices: [8, 16, 16] or [4, 8, 8].
+        original_seq_len (int): Original sequence length.
+        apply_yarn (bool): Whether to apply YaRN for long context scaling.
+        yarn_beta_fast (Optional[int]): Fast beta for YaRN.
+        yarn_beta_slow (Optional[int]): Slow beta for YaRN.
+        yarn_scale (Optional[int]): Scale factor for ctx extension.
+        use_qk_normalization (bool): Whether to use Query-Key normalization.
+        training_type (str): Type of training task.
+        batch_size (int): Batch size.
+        video_tokenizer_config_creator (Callable): Method that takes "pixel_chunk_duration: int" and "version: str" as arguments and returns video tokenizer config
+        video_tokenizer_version (str): Version of the video tokenizer.
+        num_condition_latents_t (int): Number of conditioning latent channels
+        num_tokens_to_ignore (int) = Number of tokens to ignore. This takes the precedence
+        video_height (int): Height of the video frame. Defaults to 384.
+        video_width (int): Width of the video frame. Defaults to 640.
+        rope_dim (str): RoPE dimension. Choices: "1D", "3D".
+        add_special_tokens (bool): Whether to add special tokens, use False for 2D/3D RoPE.
+        pad_to_multiple_of (int): Pad the token sequence length to the nearest multiple of this number. Defaults to 64.
+        vocab_size (int): Vocabulary size.
+        apply_abs_pos_emb (bool): Whether to apply absolute positional embeddings.
+    Returns:
+        dict: A dictionary containing the model configuration representing the model object, can be instantiated.
+    """
+    assert (
+        pixel_chunk_duration % compression_ratio[0] == 1
+    ), f"pixel_chunk_duration({pixel_chunk_duration}) should be k*n + 1 (k={compression_ratio[0]})"
+    latent_chunk_duration = (pixel_chunk_duration - 1) // compression_ratio[0] + 1
+    latent_height = video_height // compression_ratio[1]
+    latent_width = video_width // compression_ratio[2]
+    # Do some math to compute the video latent shape and sequence length
+    assert (
+        num_video_frames % pixel_chunk_duration == 0
+    ), f"num_video_frames {num_video_frames} should be divisible by pixel_chunk_duration {pixel_chunk_duration}"
+    video_latent_shape = [
+        num_video_frames // pixel_chunk_duration * latent_chunk_duration,
+        latent_height,
+        latent_width,
+    ]
+    # product of video_latent_shape
+    num_token_video_latent = video_latent_shape[0] * video_latent_shape[1] * video_latent_shape[2]
+    if add_special_tokens:
+        seq_len = num_token_video_latent + 3  # Sequence length per batch, max_seq_len + 3
+        seq_len = (seq_len + 63) // 64 * 64  # Round up to multiple of 64
+    # for text to video, we need to add <bov> token to indicate the start of the video
+    elif training_type == "text_to_video":
+        seq_len = num_token_video_latent + 1
+    else:
+        seq_len = num_token_video_latent
+
+    if seq_len % pad_to_multiple_of != 0:
+        # Round up to the nearest multiple of pad_to_multiple_of
+        seq_len = ((seq_len + pad_to_multiple_of - 1) // pad_to_multiple_of) * pad_to_multiple_of
+
+    # Model size specific parameters
+    model_arch_specs = get_model_arch_specs(model_family=model_family, model_size=model_size, pretrained=True)
+
+    # Whether skip the loss for first chunk or not, note the first token is already skipped when computing the loss
+    # If num_tokens_to_ignore is specified, use it.
+    # Else compute it from num_condition_latents_t
+    if num_tokens_to_ignore < 0:
+        num_tokens_to_ignore = latent_height * latent_width * num_condition_latents_t
+        if not add_special_tokens and num_condition_latents_t > 0:
+            # If there are no special tokens (bov), do a -1 so that you can compute the loss
+            # from the first token of the next chunk
+            num_tokens_to_ignore -= 1
+
+    model_config = ModelConfig(
+        video_height=video_height,
+        video_width=video_width,
+        max_seq_len=seq_len,
+        max_batch_size=batch_size,
+        precision="bfloat16",
+        ckpt_path=model_ckpt_path,
+        use_qk_normalization=use_qk_normalization,
+        vocab_size=64000,
+        original_seq_len=original_seq_len,
+        video_latent_shape=video_latent_shape,
+        num_video_frames=num_video_frames,
+        rope_dim=rope_dim,
+        pad_to_multiple_of=pad_to_multiple_of,
+        insert_cross_attn=insert_cross_attn,
+        insert_cross_attn_every_k_layers=insert_cross_attn_every_k_layers,
+        context_dim=context_dim,
+        apply_abs_pos_emb=apply_abs_pos_emb,
+        **model_arch_specs,
+    )
+
+    video_tokenizer_config = video_tokenizer_config_creator(
+        tokenizer_ckpt_path, pixel_chunk_duration, compression_ratio
+    )
+    tokenizer_config = TokenizerConfig(
+        text_tokenizer=None,
+        video_tokenizer=VideoTokenizerConfig(
+            config=video_tokenizer_config,
+            data_key="video",
+            tokenizer_offset=0,  # Since there is no text embeddings in the model. Note this only apply when the model is trained from scratch. If we use text pretrained model, the offset will be vocab_size of text token.
+            tokenize_here=True,
+            max_seq_len=num_token_video_latent,
+            vocab_size=vocab_size,
+        ),
+        seq_len=seq_len,
+        training_type=training_type,
+        add_special_tokens=add_special_tokens,
+        pad_to_multiple_of=pad_to_multiple_of,
+    )
+    return model_config, tokenizer_config

ar_config_base_tokenizer.py

@@ -0,0 +1,137 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import attrs
+
+from .ar_tokenizer_discrete_video import DiscreteVideoFSQStateDictTokenizer
+from .ar_tokenizer_networks import CausalDiscreteVideoTokenizer
+from .lazy_config_init import LazyCall as L
+from .lazy_config_init import LazyDict
+
+
+def create_discrete_video_fsq_tokenizer_state_dict_config(
+    ckpt_path, pixel_chunk_duration=33, compression_ratio=[8, 16, 16]
+) -> LazyDict:
+    CausalDiscreteFactorizedVideoTokenizerConfig: LazyDict = L(CausalDiscreteVideoTokenizer)(
+        # The new causal discrete tokenizer, that is at least 2x more efficient in memory and runtime.
+        # - It relies on fully 3D discrete wavelet transform
+        # - Uses a layer norm instead of a group norm
+        # - Factorizes full convolutions into spatial and temporal convolutions
+        # - Factorizes full attention into spatial and temporal attention
+        # - Strictly causal, with flexible temporal length at inference.
+        attn_resolutions=[32],
+        channels=128,
+        channels_mult=[2, 4, 4],
+        dropout=0.0,
+        in_channels=3,
+        num_res_blocks=2,
+        out_channels=3,
+        resolution=1024,
+        patch_size=4,
+        patch_method="haar",
+        z_channels=16,
+        z_factor=1,
+        num_groups=1,
+        legacy_mode=False,
+        spatial_compression=16,
+        temporal_compression=8,
+        embedding_dim=6,
+        levels=[8, 8, 8, 5, 5, 5],
+        name="CausalDiscreteFactorizedVideoTokenizer",
+    )
+
+    return L(DiscreteVideoFSQStateDictTokenizer)(
+        enc_fp=ckpt_path.replace("ema.jit", "encoder.jit"),
+        dec_fp=ckpt_path.replace("ema.jit", "decoder.jit"),
+        tokenizer_module=CausalDiscreteFactorizedVideoTokenizerConfig,
+        name="discrete_video_fsq",
+        latent_ch=6,
+        is_bf16=True,
+        pixel_chunk_duration=pixel_chunk_duration,
+        latent_chunk_duration=1 + (pixel_chunk_duration - 1) // compression_ratio[0],
+        max_enc_batch_size=8,
+        max_dec_batch_size=4,
+        levels=[8, 8, 8, 5, 5, 5],
+        compression_ratio=compression_ratio,
+    )
+
+
+@attrs.define(slots=False)
+class TextTokenizerConfig:
+    """
+    Text tokenizer config
+
+    Args:
+        config: Config file to define the text tokenizer class.
+        data_key (str): The input key from data_dict that will be passed to the text tokenizer.
+        tokenize_here (bool): Whether to use the tokenizer to perform online tokenization.
+        tokenizer_offset (int): Offset that is added to the tokens.
+        vocab_size (int): Vocabulary size of the tokenizer.
+    """
+
+    config: LazyDict
+    data_key: str = ""
+    tokenize_here: bool = False
+    tokenizer_offset: int = 0
+    vocab_size: int = 0
+
+
+@attrs.define(slots=False)
+class VideoTokenizerConfig:
+    """
+    Video tokenizer config
+
+    Args:
+        config: Config file to define the video tokenizer class.
+        data_key (str): The input key from data_dict that will be passed to the video tokenizer.
+        tokenize_here (bool): Whether to use the tokenizer to perform online tokenization.
+        tokenizer_offset (int): Offset that is added to the tokens. In case of joint text-video tokenizers, we
+            add an offset to make sure that video tokens and text tokens don't overlap.
+        vocab_size (int): Vocabulary size of the tokenizer.
+        max_seq_len (int): Maximum token length for an input video.
+    """
+
+    config: LazyDict
+    data_key: str = ""
+    tokenize_here: bool = True
+    tokenizer_offset: int = 0
+    vocab_size: int = 0
+    max_seq_len: int = -1
+
+
+@attrs.define(slots=False)
+class TokenizerConfig:
+    """
+    Joint tokenizer config
+
+    Args:
+        text_tokenizer (TextTokenizerConfig): Text tokenizer config file
+        class_tokenizer (ClassTokenizerConfig): Class tokenizer config file
+        video_tokenizer (VideoTokenizerConfig): Video tokenizer config file
+        image_tokenizer (ImageTokenizerConfig): Image tokenizer config file
+        seq_len (int): Final token sequence length
+        training_type (str): Type of training we use. Supports ["text_only", "text_to_video", "class_to_image", "image_text_interleaved"]
+        add_special_tokens (bool): Whether to add special tokens to the output tokens
+        pad_to_multiple_of (int): Pad the token sequence length to the nearest multiple of this number. Defaults to 64.
+    """
+
+    text_tokenizer: Optional[TextTokenizerConfig] = None
+    video_tokenizer: Optional[VideoTokenizerConfig] = None
+    seq_len: int = 4096
+    training_type: str = None
+    add_special_tokens: bool = True
+    pad_to_multiple_of: Optional[int] = 64

ar_config_inference_inference_config.py

@@ -0,0 +1,102 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, List, Union
+
+import attrs
+
+from .ar_config_base_model import ModelConfig, TokenizerConfig
+
+
+@attrs.define(slots=False)
+class DataShapeConfig:
+    latent_shape: list = []
+    num_video_frames: Union[None, int] = None
+    height: Union[None, int] = None
+    width: Union[None, int] = None
+
+
+@attrs.define(slots=False)
+class SamplingConfig:
+    """
+    Sampling config
+    Args:
+        temperature (float): Temperature value for controlling randomness in sampling. Defaults to 0.6.
+        top_p (float): Top-p probability threshold for nucleus sampling. Defaults to 0.9.
+        logprobs (bool): Flag indicating whether to compute token log probabilities. Defaults to False.
+        echo (bool): Flag indicating whether to include prompt tokens in the generated output. Defaults to False.
+
+    """
+
+    temperature: float = 0.6
+    top_k: int = None
+    top_p: float = 0.9
+    compile_prefill: bool = False
+    compile_sampling: bool = True
+    logprobs: bool = False
+    echo: bool = False
+
+
+@attrs.define(slots=False)
+class DiffusionDecoderSamplingConfig:
+    """
+    Diffusion decoder sampling config
+    Args:
+        guidance (float): Guidance scale for the diffusion process. Controls how much the model follows the conditioning. Defaults to 0.8.
+        sigma_min (float): Minimum noise level for the diffusion process. Defaults to 0.02.
+        sigma (float): Initial noise level for the diffusion process. Defaults to 8.
+        num_steps (int): Number of denoising steps to perform. Defaults to 35.
+        overlap (int): Number of overlapping frames between video chunks during processing. Defaults to 2.
+        continuous_tokenizer_channel (int): Number of channels in the continuous tokenizer of diffusion decoder. Defaults to 16.
+        continuous_tokenizer_spatial_compression_ratio (int): Spatial compression ratio for the continuous tokenizer of diffusion decoder. Defaults to 8.
+        dd_train_num_video_frames (int): Number of video frames used during training for diffusion decoder. Defaults to 57.
+    """
+
+    guidance: float = 1.8
+    sigma_min: float = 0.02
+    sigma: float = 8
+    num_steps: int = 15
+    overlap: int = 2
+    continuous_tokenizer_channel = 16
+    continuous_tokenizer_spatial_compression_ratio = 8
+    dd_train_num_video_frames: int = 57
+    max_iter: int = 99
+    fps: int = 24
+
+
+@attrs.define(slots=False)
+class InferenceConfig:
+    """
+    Inference config
+    Args:
+        model_config (ModelConfig): Model config
+        tokenizer_config (TokenizerConfig): Tokenizer config
+        ckpt_path (str): Path to the checkpoint
+        latent_shape (list): Shape of the latent
+    """
+
+    model_config: ModelConfig = None
+    tokenizer_config: TokenizerConfig = None
+    ckpt_path: str = ""
+    data_shape_config: DataShapeConfig = None
+
+    defaults: List[Any] = attrs.field(
+        factory=lambda: [
+            "_self_",
+            {"data_val": None},
+            {"data_shape_config": "video_shape_as_model_config"},
+            {"eval_job": None},
+        ]
+    )

ar_diffusion_decoder_config_base_conditioner.py

@@ -0,0 +1,61 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+import torch
+
+from .df_conditioner import BaseVideoCondition, GeneralConditioner
+from .df_config_base_conditioner import (
+    FPSConfig,
+    ImageSizeConfig,
+    LatentConditionConfig,
+    LatentConditionSigmaConfig,
+    NumFramesConfig,
+    PaddingMaskConfig,
+    TextConfig,
+)
+from .lazy_config_init import LazyCall as L
+from .lazy_config_init import LazyDict
+
+
+@dataclass
+class VideoLatentDiffusionDecoderCondition(BaseVideoCondition):
+    # latent_condition will concat to the input of network, along channel dim;
+    # cfg will make latent_condition all zero padding.
+    latent_condition: Optional[torch.Tensor] = None
+    latent_condition_sigma: Optional[torch.Tensor] = None
+
+
+class VideoDiffusionDecoderConditioner(GeneralConditioner):
+    def forward(
+        self,
+        batch: Dict,
+        override_dropout_rate: Optional[Dict[str, float]] = None,
+    ) -> VideoLatentDiffusionDecoderCondition:
+        output = super()._forward(batch, override_dropout_rate)
+        return VideoLatentDiffusionDecoderCondition(**output)
+
+
+VideoLatentDiffusionDecoderConditionerConfig: LazyDict = L(VideoDiffusionDecoderConditioner)(
+    text=TextConfig(),
+    fps=FPSConfig(),
+    num_frames=NumFramesConfig(),
+    image_size=ImageSizeConfig(),
+    padding_mask=PaddingMaskConfig(),
+    latent_condition=LatentConditionConfig(),
+    latent_condition_sigma=LatentConditionSigmaConfig(),
+)

ar_diffusion_decoder_config_config_latent_diffusion_decoder.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, List
+
+import attrs
+
+from .ar_diffusion_decoder_config_registry import register_configs as register_dd_configs
+from .df_config_base_model import LatentDiffusionDecoderModelConfig
+from .df_config_registry import register_configs
+from .config import Config as ori_Config
+from .config_helper import import_all_modules_from_package
+
+from .ar_diffusion_decoder_config_inference_cosmos_diffusiondecoder_7b import LazyDict
+
+@attrs.define(slots=False)
+class Config(ori_Config):
+    # default config groups that will be used unless overwritten
+    # see config groups in registry.py
+    defaults: List[Any] = attrs.field(
+        factory=lambda: [
+            "_self_",
+            {"net": None},
+            {"conditioner": "basic"},
+            {"tokenizer": "tokenizer"},
+            {"tokenizer_corruptor": None},
+            {"latent_corruptor": None},
+            {"pixel_corruptor": None},
+            {"experiment": None},
+        ]
+    )
+
+
+def make_config():
+    c = Config(model=LatentDiffusionDecoderModelConfig())
+
+    # Specifying values through instances of attrs
+    c.job.project = "cosmos_video4"
+    c.job.group = "debug"
+    c.job.name = "delete_${now:%Y-%m-%d}_${now:%H-%M-%S}"
+
+    # # Call this function to register config groups for advanced overriding.
+    register_configs()
+    register_dd_configs()
+
+    # # experiment config are defined in the experiment folder
+    # # call import_all_modules_from_package to register them
+    # import_all_modules_from_package("cosmos1.models.diffusion.config.inference", reload=True)
+    # import_all_modules_from_package("cosmos1.models.autoregressive.diffusion_decoder.config.inference", reload=True)
+    return c

ar_diffusion_decoder_config_inference_cosmos_diffusiondecoder_7b.py

@@ -0,0 +1,85 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from hydra.core.config_store import ConfigStore
+
+from .ar_diffusion_decoder_network import DiffusionDecoderGeneralDIT
+from .lazy_config_init import LazyCall as L
+from .lazy_config_init import LazyDict
+
+num_frames = 57
+Cosmos_DiffusionDecoder_7B_INFERENCE_ONLY: LazyDict = LazyDict(
+    dict(
+        defaults=[
+            {"override /net": "faditv2_7b"},
+            {"override /tokenizer": "cosmos_video_tokenizer_res720_comp8x8x8_t121_ver092624"},
+            {"override /conditioner": "video_latent_diffusion_decoder_cond"},
+            {"override /tokenizer_corruptor": "cosmos_video_discrete_tokenizer_res720_comp8x16x16_t49_ver110224"},
+            "_self_",
+        ],
+        job=dict(
+            group="diffusion_deocder_FT_7Bv1_001",
+            name="DD_FT_7Bv1_003_002_tokenizer888_spatch2_discrete_cond_on_token",
+        ),
+        model=dict(
+            diffusion_decoder_cond_sigma_low=0.0,
+            diffusion_decoder_cond_sigma_high=0.0,
+            diffusion_decoder_corrupt_prob=0.0,
+            condition_on_tokenizer_corruptor_token=True,
+            latent_shape=[
+                16,
+                num_frames,
+                88,
+                160,
+            ],
+            tokenizer_corruptor=dict(
+                pixel_chunk_duration=num_frames,
+                latent_chunk_duration=1 + (num_frames - 1) // 8,
+            ),
+            net=L(DiffusionDecoderGeneralDIT)(
+                diffusion_decoder_condition_on_sigma=False,
+                max_img_h=240,
+                max_img_w=240,
+                rope_h_extrapolation_ratio=1.5,
+                rope_w_extrapolation_ratio=1.5,
+                rope_t_extrapolation_ratio=1,
+                block_x_format="THWBD",
+                is_diffusion_decoder=True,
+                patch_spatial=2,
+                diffusion_decoder_condition_on_token=True,
+                diffusion_decoder_token_condition_voc_size=64000,
+                diffusion_decoder_token_condition_dim=32,
+            ),
+            tokenizer=dict(
+                video_vae=dict(
+                    pixel_chunk_duration=num_frames,
+                )
+            ),
+            conditioner=dict(
+                latent_condition=dict(
+                    dropout_rate=0.2,
+                )
+            ),
+        ),
+    )
+)
+
+cs = ConfigStore.instance()
+cs.store(
+    group="experiment",
+    package="_global_",
+    name=Cosmos_DiffusionDecoder_7B_INFERENCE_ONLY["job"]["name"],
+    node=Cosmos_DiffusionDecoder_7B_INFERENCE_ONLY,
+)

ar_diffusion_decoder_config_registry.py

@@ -0,0 +1,118 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from hydra.core.config_store import ConfigStore
+
+from .ar_diffusion_decoder_config_base_conditioner import (
+    VideoLatentDiffusionDecoderConditionerConfig,
+)
+from .ar_tokenizer_discrete_video import DiscreteVideoFSQJITTokenizer
+from .df_module_pretrained_vae import JITVAE, JointImageVideoSharedJITTokenizer, VideoJITTokenizer
+from .lazy_config_init import LazyCall as L
+
+
+def get_cosmos_video_discrete_tokenizer_comp8x16x16(
+    resolution: str,
+    chunk_duration: int,
+    checkpoint_path: str,
+):
+    assert resolution in ["720"]
+
+    pixel_chunk_duration = chunk_duration
+    temporal_compression_factor = 8
+    spatial_compression_factor = 16
+
+    return L(DiscreteVideoFSQJITTokenizer)(
+        enc_fp=checkpoint_path.replace(".jit", "encoder.jit"),
+        dec_fp=checkpoint_path.replace(".jit", "decoder.jit"),
+        name="discrete_video_fsq",
+        latent_ch=6,
+        is_bf16=True,
+        pixel_chunk_duration=pixel_chunk_duration,
+        latent_chunk_duration=1 + (pixel_chunk_duration - 1) // temporal_compression_factor,
+        max_enc_batch_size=8,
+        max_dec_batch_size=4,
+        levels=[8, 8, 8, 5, 5, 5],
+        compression_ratio=[temporal_compression_factor, spatial_compression_factor, spatial_compression_factor],
+    )
+
+
+def get_cosmos_video_tokenizer_comp8x8x8(resolution: str, chunk_duration: int, checkpoint_path=None):
+    pixel_chunk_duration = chunk_duration
+    temporal_compression_factor = 8
+    spatial_compression_factor = 8
+
+    return L(JointImageVideoSharedJITTokenizer)(
+        video_vae=L(VideoJITTokenizer)(
+            name="cosmos_1_0_diffusion_tokenizer",
+            latent_ch=16,
+            is_bf16=True,
+            pixel_chunk_duration=pixel_chunk_duration,
+            temporal_compression_factor=temporal_compression_factor,
+            spatial_compression_factor=spatial_compression_factor,
+            spatial_resolution=resolution,
+        ),
+        image_vae=L(JITVAE)(
+            name="cosmos_1_0_diffusion_tokenizer",
+            latent_ch=16,
+            is_image=False,
+            is_bf16=True,
+        ),
+        name="cosmos_diffusion_tokenizer_res720_comp8x8x8_t121_ver092624",
+        latent_ch=16,
+    )
+
+
+def register_tokenizer(cs):
+    cs.store(
+        group="tokenizer",
+        package="model.tokenizer",
+        name="cosmos_video_tokenizer_res720_comp8x8x8_t121_ver092624",
+        node=get_cosmos_video_tokenizer_comp8x8x8(
+            resolution="720",
+            chunk_duration=121,
+            checkpoint_path="checkpoints/Cosmos-1.0-Tokenizer-CV8x8x8/.jit",
+        ),
+    )
+
+
+def register_corruptor(cs):
+    cs.store(
+        group="tokenizer_corruptor",
+        package="model.tokenizer_corruptor",
+        name="cosmos_video_discrete_tokenizer_res720_comp8x16x16_t49_ver110224",
+        node=get_cosmos_video_discrete_tokenizer_comp8x16x16(
+            resolution="720",
+            chunk_duration=49,
+            checkpoint_path="checkpoints/Cosmos-1.0-Tokenizer-DV8x16x16/.jit",
+        ),
+    )
+
+
+def register_conditioner(cs):
+    cs.store(
+        group="conditioner",
+        package="model.conditioner",
+        name="video_latent_diffusion_decoder_cond",
+        node=VideoLatentDiffusionDecoderConditionerConfig,
+    )
+
+
+def register_configs():
+    cs = ConfigStore.instance()
+
+    register_conditioner(cs)
+    register_corruptor(cs)
+    register_tokenizer(cs)

ar_diffusion_decoder_inference.py

@@ -0,0 +1,120 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+import gc
+from typing import List
+
+import torch
+
+from .ar_config_inference_inference_config import DiffusionDecoderSamplingConfig
+from .ar_diffusion_decoder_model import LatentDiffusionDecoderModel
+from .ar_diffusion_decoder_utils import linear_blend_video_list, split_with_overlap
+from .log import log
+
+
+def diffusion_decoder_process_tokens(
+    model: LatentDiffusionDecoderModel,
+    indices_tensor: List[torch.Tensor],
+    dd_sampling_config: DiffusionDecoderSamplingConfig = None,
+    original_video_example: torch.Tensor = None,
+    t5_emb_batch: List[torch.Tensor] = None,
+):
+    _, T, H, W = original_video_example.shape
+    if dd_sampling_config is None:
+        dd_sampling_config = DiffusionDecoderSamplingConfig()
+    # indices_tensor is assumed to be a list of tensors with shape 1LHW
+    data_batch_list = []
+    for sample_num, token_CTHW in enumerate(indices_tensor):
+        token_BCTHW = token_CTHW.unsqueeze(0).unsqueeze(1)
+        token_BCTHW = split_with_overlap(
+            token_BCTHW,
+            (dd_sampling_config.dd_train_num_video_frames - 1) // 8 + 1,
+            overlap=dd_sampling_config.overlap,
+            tobf16=False,
+        )
+        data_batch_list.append(
+            {
+                "token_chunks": token_BCTHW,
+                "t5_text_embeddings": t5_emb_batch[sample_num].to(torch.bfloat16),
+                "t5_text_mask": torch.ones(1, 512, dtype=torch.bfloat16).cuda(),
+                # other conditions
+                "image_size": torch.tensor([[H, W, H, W]] * 1, dtype=torch.bfloat16).cuda(),
+                "fps": torch.tensor([dd_sampling_config.fps] * 1, dtype=torch.bfloat16).cuda(),
+                "num_frames": torch.tensor(
+                    [dd_sampling_config.dd_train_num_video_frames] * 1, dtype=torch.bfloat16
+                ).cuda(),
+                "padding_mask": torch.zeros((1, 1, H, W), dtype=torch.bfloat16).cuda(),
+            }
+        )
+
+    out_videos_batch = []
+
+    for idx, data_batch_template in enumerate(data_batch_list):
+        full_length_sample = []
+        iterations = min(len(data_batch_template["token_chunks"]), dd_sampling_config.max_iter)
+        for iter in range(iterations):
+            gc.collect()
+            torch.cuda.empty_cache()
+
+            data_batch = copy.deepcopy(data_batch_template)
+            data_batch["video"] = data_batch_template["token_chunks"][iter].cuda().to("cuda")
+
+            log.debug(f"Run iter {iter} for video # {idx} at length {data_batch['video'].shape[2]}")
+            # org_video,
+            with torch.no_grad():
+                samples_latent = model.generate_samples_from_batch(
+                    data_batch,
+                    guidance=dd_sampling_config.guidance,
+                    sigma_min=dd_sampling_config.sigma_min,
+                    state_shape=[
+                        dd_sampling_config.continuous_tokenizer_channel,
+                        dd_sampling_config.continuous_tokenizer_spatial_compression_ratio,
+                        H // 8,
+                        W // 8,
+                    ],
+                    apply_corruptor=False,
+                    return_recon_x=False,
+                    # corrupt_sigma=dd_sampling_config.sigma,
+                    preencode_condition=True,  # We are using discrete model, so the input is already pre-encoded
+                    num_steps=dd_sampling_config.num_steps,
+                )
+                log.debug(f"Current sample shape {samples_latent.shape} for video # {idx} ")
+            full_length_sample.append(samples_latent.detach())
+
+            # Turn off because we remove CP
+            # distributed.barrier()
+            del data_batch
+
+            torch.cuda.empty_cache()
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        # Decode full-length samples and free GPU memory
+        full_length_sample_pixs = [model.decode(item).clamp(-1, 1).cpu() for item in full_length_sample]
+        torch.cuda.empty_cache()
+
+        # Blend pixel samples
+        if len(full_length_sample_pixs) > 1:
+            full_length_sample_pixel_blend = linear_blend_video_list(
+                full_length_sample_pixs, dd_sampling_config.overlap
+            )[:, :, :T]
+        else:
+            full_length_sample_pixel_blend = full_length_sample_pixs[0][:, :, :T]
+
+        # Batch size of full_length_sample_pixel_blend is always 1
+        out_videos_batch.append((1 + full_length_sample_pixel_blend[0].cpu()) / 2)
+    return out_videos_batch

ar_diffusion_decoder_model.py

@@ -0,0 +1,231 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from .df_conditioner import BaseVideoCondition
+from .df_df_functional_batch_ops import batch_mul
+from .df_df_module_res_sampler import COMMON_SOLVER_OPTIONS
+from .df_model_model_t2w import DiffusionT2WModel as VideoDiffusionModel
+from .lazy_config_init import instantiate as lazy_instantiate
+
+
+@dataclass
+class VideoLatentDiffusionDecoderCondition(BaseVideoCondition):
+    # latent_condition will concat to the input of network, along channel dim;
+    # cfg will make latent_condition all zero padding.
+    latent_condition: Optional[torch.Tensor] = None
+    latent_condition_sigma: Optional[torch.Tensor] = None
+
+
+class LatentDiffusionDecoderModel(VideoDiffusionModel):
+    def __init__(self, config):
+        super().__init__(config)
+        """
+        latent_corruptor: the corruption module is used to corrupt the latents. It add gaussian noise to the latents.
+        pixel_corruptor: the corruption module is used to corrupt the pixels. It apply gaussian blur kernel to pixels in a temporal consistent way.
+        tokenizer_corruptor: the corruption module is used to simulate tokenizer reconstruction errors.
+
+        diffusion decoder noise augmentation pipeline for continuous token condition model:
+        condition: GT_video [T, H, W]
+                        -> tokenizer_corruptor~(8x8x8) encode -> latent_corruptor -> tokenizer_corruptor~(8x8x8) decode
+                        -> pixel corruptor
+                        -> tokenizer~(1x8x8) encode -> condition [T, H/8, W/8]
+        GT: GT_video [T, H, W] -> tokenizer~(1x8x8) -> x_t [T, H/8, W/8].
+
+        diffusion decoder noise augmentation pipeline for discrete token condition model:
+        condition: GT_video [T, H, W]
+                -> pixel corruptor
+                -> discrete tokenizer encode -> condition [T, T/8, H/16, W/16]
+        GT: GT_video [T, H, W] -> tokenizer~(8x8x8) -> x_t [T, T/8, H/8, W/8].
+
+        """
+        self.latent_corruptor = lazy_instantiate(config.latent_corruptor)
+        self.pixel_corruptor = lazy_instantiate(config.pixel_corruptor)
+        self.tokenizer_corruptor = lazy_instantiate(config.tokenizer_corruptor)
+
+        if self.latent_corruptor:
+            self.latent_corruptor.to(**self.tensor_kwargs)
+        if self.pixel_corruptor:
+            self.pixel_corruptor.to(**self.tensor_kwargs)
+
+        if self.tokenizer_corruptor:
+            if hasattr(self.tokenizer_corruptor, "reset_dtype"):
+                self.tokenizer_corruptor.reset_dtype()
+        else:
+            assert self.pixel_corruptor is not None
+
+        self.diffusion_decoder_cond_sigma_low = config.diffusion_decoder_cond_sigma_low
+        self.diffusion_decoder_cond_sigma_high = config.diffusion_decoder_cond_sigma_high
+        self.diffusion_decoder_corrupt_prob = config.diffusion_decoder_corrupt_prob
+        if hasattr(config, "condition_on_tokenizer_corruptor_token"):
+            self.condition_on_tokenizer_corruptor_token = config.condition_on_tokenizer_corruptor_token
+        else:
+            self.condition_on_tokenizer_corruptor_token = False
+
+    def is_image_batch(self, data_batch: dict[str, Tensor]) -> bool:
+        """We hanlde two types of data_batch. One comes from a joint_dataloader where "dataset_name" can be used to differenciate image_batch and video_batch.
+        Another comes from a dataloader which we by default assumes as video_data for video model training.
+        """
+        is_image = self.input_image_key in data_batch
+        is_video = self.input_data_key in data_batch
+        assert (
+            is_image != is_video
+        ), "Only one of the input_image_key or input_data_key should be present in the data_batch."
+        return is_image
+
+    def get_x0_fn_from_batch(
+        self,
+        data_batch: Dict,
+        guidance: float = 1.5,
+        is_negative_prompt: bool = False,
+        apply_corruptor: bool = True,
+        corrupt_sigma: float = 1.5,
+        preencode_condition: bool = False,
+    ) -> Callable:
+        """
+        Generates a callable function `x0_fn` based on the provided data batch and guidance factor.
+
+        This function first processes the input data batch through a conditioning workflow (`conditioner`) to obtain conditioned and unconditioned states. It then defines a nested function `x0_fn` which applies a denoising operation on an input `noise_x` at a given noise level `sigma` using both the conditioned and unconditioned states.
+
+        Args:
+        - data_batch (Dict): A batch of data used for conditioning. The format and content of this dictionary should align with the expectations of the `self.conditioner`
+        - guidance (float, optional): A scalar value that modulates the influence of the conditioned state relative to the unconditioned state in the output. Defaults to 1.5.
+        - is_negative_prompt (bool): use negative prompt t5 in uncondition if true
+
+        Returns:
+        - Callable: A function `x0_fn(noise_x, sigma)` that takes two arguments, `noise_x` and `sigma`, and return x0 predictoin
+
+        The returned function is suitable for use in scenarios where a denoised state is required based on both conditioned and unconditioned inputs, with an adjustable level of guidance influence.
+        """
+        input_key = self.input_data_key  # by default it is video key
+        # Latent state
+        raw_state = data_batch[input_key]
+
+        if self.condition_on_tokenizer_corruptor_token:
+            if preencode_condition:
+                latent_condition = raw_state.to(torch.int32).contiguous()
+                corrupted_pixel = self.tokenizer_corruptor.decode(latent_condition[:, 0])
+            else:
+                corrupted_pixel = (
+                    self.pixel_corruptor(raw_state) if apply_corruptor and self.pixel_corruptor else raw_state
+                )
+                latent_condition = self.tokenizer_corruptor.encode(corrupted_pixel)
+                latent_condition = latent_condition[1] if isinstance(latent_condition, tuple) else latent_condition
+                corrupted_pixel = self.tokenizer_corruptor.decode(latent_condition)
+                latent_condition = latent_condition.unsqueeze(1)
+        else:
+            if preencode_condition:
+                latent_condition = raw_state
+                corrupted_pixel = self.decode(latent_condition)
+            else:
+                corrupted_pixel = (
+                    self.pixel_corruptor(raw_state) if apply_corruptor and self.pixel_corruptor else raw_state
+                )
+                latent_condition = self.encode(corrupted_pixel).contiguous()
+
+        sigma = (
+            torch.rand((latent_condition.shape[0],)).to(**self.tensor_kwargs) * corrupt_sigma
+        )  # small value to indicate clean video
+        _, _, _, c_noise_cond = self.scaling(sigma=sigma)
+        if corrupt_sigma != self.diffusion_decoder_cond_sigma_low and self.diffusion_decoder_corrupt_prob > 0:
+            noise = batch_mul(sigma, torch.randn_like(latent_condition))
+            latent_condition = latent_condition + noise
+        data_batch["latent_condition_sigma"] = batch_mul(torch.ones_like(latent_condition[:, 0:1, ::]), c_noise_cond)
+        data_batch["latent_condition"] = latent_condition
+        if is_negative_prompt:
+            condition, uncondition = self.conditioner.get_condition_with_negative_prompt(data_batch)
+        else:
+            condition, uncondition = self.conditioner.get_condition_uncondition(data_batch)
+
+        def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
+            cond_x0 = self.denoise(noise_x, sigma, condition).x0
+            uncond_x0 = self.denoise(noise_x, sigma, uncondition).x0
+            return cond_x0 + guidance * (cond_x0 - uncond_x0)
+
+        return x0_fn, corrupted_pixel
+
+    def generate_samples_from_batch(
+        self,
+        data_batch: Dict,
+        guidance: float = 1.5,
+        seed: int = 1,
+        state_shape: Tuple | None = None,
+        n_sample: int | None = None,
+        is_negative_prompt: bool = False,
+        num_steps: int = 35,
+        solver_option: COMMON_SOLVER_OPTIONS = "2ab",
+        sigma_min: float = 0.02,
+        apply_corruptor: bool = False,
+        return_recon_x: bool = False,
+        corrupt_sigma: float = 0.01,
+        preencode_condition: bool = False,
+    ) -> Tensor:
+        """
+        Generate samples from the batch. Based on given batch, it will automatically determine whether to generate image or video samples.
+        Args:
+            data_batch (dict): raw data batch draw from the training data loader.
+            iteration (int): Current iteration number.
+            guidance (float): guidance weights
+            seed (int): random seed
+            state_shape (tuple): shape of the state, default to self.state_shape if not provided
+            n_sample (int): number of samples to generate
+            is_negative_prompt (bool): use negative prompt t5 in uncondition if true
+            num_steps (int): number of steps for the diffusion process
+            solver_option (str): differential equation solver option, default to "2ab"~(mulitstep solver)
+            preencode_condition (bool): use pre-computed condition if true, save tokenizer's inference time memory/
+        """
+        if not preencode_condition:
+            self._normalize_video_databatch_inplace(data_batch)
+            self._augment_image_dim_inplace(data_batch)
+        is_image_batch = False
+        if n_sample is None:
+            input_key = self.input_image_key if is_image_batch else self.input_data_key
+            n_sample = data_batch[input_key].shape[0]
+        if state_shape is None:
+            if is_image_batch:
+                state_shape = (self.state_shape[0], 1, *self.state_shape[2:])  # C,T,H,W
+
+        x0_fn, recon_x = self.get_x0_fn_from_batch(
+            data_batch,
+            guidance,
+            is_negative_prompt=is_negative_prompt,
+            apply_corruptor=apply_corruptor,
+            corrupt_sigma=corrupt_sigma,
+            preencode_condition=preencode_condition,
+        )
+        generator = torch.Generator(device=self.tensor_kwargs["device"])
+        generator.manual_seed(seed)
+        x_sigma_max = (
+            torch.randn(n_sample, *state_shape, **self.tensor_kwargs, generator=generator) * self.sde.sigma_max
+        )
+
+        samples = self.sampler(
+            x0_fn,
+            x_sigma_max,
+            num_steps=num_steps,
+            sigma_min=sigma_min,
+            sigma_max=self.sde.sigma_max,
+            solver_option=solver_option,
+        )
+
+        if return_recon_x:
+            return samples, recon_x
+        else:
+            return samples

ar_diffusion_decoder_network.py

@@ -0,0 +1,163 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Tuple
+
+import torch
+from einops import rearrange
+from torch import nn
+from torchvision import transforms
+
+from .df_module_blocks import PatchEmbed
+from .df_network_general_dit import GeneralDIT
+
+
+class DiffusionDecoderGeneralDIT(GeneralDIT):
+    def __init__(
+        self,
+        *args,
+        is_diffusion_decoder: bool = True,
+        diffusion_decoder_condition_on_sigma: bool = False,
+        diffusion_decoder_condition_on_token: bool = False,
+        diffusion_decoder_token_condition_voc_size: int = 64000,
+        diffusion_decoder_token_condition_dim: int = 32,
+        **kwargs,
+    ):
+        # diffusion decoder setting
+        self.is_diffusion_decoder = is_diffusion_decoder
+        self.diffusion_decoder_condition_on_sigma = diffusion_decoder_condition_on_sigma
+        self.diffusion_decoder_condition_on_token = diffusion_decoder_condition_on_token
+        self.diffusion_decoder_token_condition_voc_size = diffusion_decoder_token_condition_voc_size
+        self.diffusion_decoder_token_condition_dim = diffusion_decoder_token_condition_dim
+        super().__init__(*args, **kwargs)
+
+    def initialize_weights(self):
+        # Initialize transformer layers:
+        super().initialize_weights()
+        if self.diffusion_decoder_condition_on_token:
+            nn.init.constant_(self.token_embedder.weight, 0)
+
+    def build_patch_embed(self):
+        (
+            concat_padding_mask,
+            in_channels,
+            patch_spatial,
+            patch_temporal,
+            model_channels,
+            is_diffusion_decoder,
+            diffusion_decoder_token_condition_dim,
+            diffusion_decoder_condition_on_sigma,
+        ) = (
+            self.concat_padding_mask,
+            self.in_channels,
+            self.patch_spatial,
+            self.patch_temporal,
+            self.model_channels,
+            self.is_diffusion_decoder,
+            self.diffusion_decoder_token_condition_dim,
+            self.diffusion_decoder_condition_on_sigma,
+        )
+        in_channels = (
+            in_channels + in_channels
+            if (is_diffusion_decoder and not self.diffusion_decoder_condition_on_token)
+            else in_channels
+        )
+        in_channels = in_channels + 1 if diffusion_decoder_condition_on_sigma else in_channels
+        in_channels = (
+            in_channels + self.diffusion_decoder_token_condition_dim
+            if self.diffusion_decoder_condition_on_token
+            else in_channels
+        )
+        in_channels = in_channels + 1 if concat_padding_mask else in_channels
+
+        self.x_embedder = PatchEmbed(
+            spatial_patch_size=patch_spatial,
+            temporal_patch_size=patch_temporal,
+            in_channels=in_channels,
+            out_channels=model_channels,
+            bias=False,
+        )
+
+        if self.diffusion_decoder_condition_on_token:
+            self.token_embedder = nn.Embedding(
+                self.diffusion_decoder_token_condition_voc_size, self.diffusion_decoder_token_condition_dim
+            )
+
+    def prepare_embedded_sequence(
+        self,
+        x_B_C_T_H_W: torch.Tensor,
+        fps: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+        latent_condition: Optional[torch.Tensor] = None,
+        latent_condition_sigma: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Prepares an embedded sequence tensor by applying positional embeddings and handling padding masks.
+
+        Args:
+            x_B_C_T_H_W (torch.Tensor): video
+            fps (Optional[torch.Tensor]): Frames per second tensor to be used for positional embedding when required.
+                                    If None, a default value (`self.base_fps`) will be used.
+            padding_mask (Optional[torch.Tensor]): current it is not used
+
+        Returns:
+            Tuple[torch.Tensor, Optional[torch.Tensor]]:
+                - A tensor of shape (B, T, H, W, D) with the embedded sequence.
+                - An optional positional embedding tensor, returned only if the positional embedding class
+                (`self.pos_emb_cls`) includes 'rope'. Otherwise, None.
+
+        Notes:
+            - If `self.concat_padding_mask` is True, a padding mask channel is concatenated to the input tensor.
+            - The method of applying positional embeddings depends on the value of `self.pos_emb_cls`.
+            - If 'rope' is in `self.pos_emb_cls` (case insensitive), the positional embeddings are generated using
+                the `self.pos_embedder` with the shape [T, H, W].
+            - If "fps_aware" is in `self.pos_emb_cls`, the positional embeddings are generated using the `self.pos_embedder`
+                with the fps tensor.
+            - Otherwise, the positional embeddings are generated without considering fps.
+        """
+        if self.diffusion_decoder_condition_on_token:
+            latent_condition = self.token_embedder(latent_condition)
+            B, _, T, H, W, _ = latent_condition.shape
+            latent_condition = rearrange(latent_condition, "B 1 T H W D -> (B T) (1 D) H W")
+
+            latent_condition = transforms.functional.resize(
+                latent_condition, list(x_B_C_T_H_W.shape[-2:]), interpolation=transforms.InterpolationMode.BILINEAR
+            )
+            latent_condition = rearrange(latent_condition, "(B T) D H W -> B D T H W ", B=B, T=T)
+        x_B_C_T_H_W = torch.cat([x_B_C_T_H_W, latent_condition], dim=1)
+        if self.diffusion_decoder_condition_on_sigma:
+            x_B_C_T_H_W = torch.cat([x_B_C_T_H_W, latent_condition_sigma], dim=1)
+        if self.concat_padding_mask:
+            padding_mask = transforms.functional.resize(
+                padding_mask, list(x_B_C_T_H_W.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
+            )
+            x_B_C_T_H_W = torch.cat(
+                [x_B_C_T_H_W, padding_mask.unsqueeze(1).repeat(1, 1, x_B_C_T_H_W.shape[2], 1, 1)], dim=1
+            )
+        x_B_T_H_W_D = self.x_embedder(x_B_C_T_H_W)
+
+        if self.extra_per_block_abs_pos_emb:
+            extra_pos_emb = self.extra_pos_embedder(x_B_T_H_W_D, fps=fps)
+        else:
+            extra_pos_emb = None
+
+        if "rope" in self.pos_emb_cls.lower():
+            return x_B_T_H_W_D, self.pos_embedder(x_B_T_H_W_D, fps=fps), extra_pos_emb
+
+        if "fps_aware" in self.pos_emb_cls:
+            x_B_T_H_W_D = x_B_T_H_W_D + self.pos_embedder(x_B_T_H_W_D, fps=fps)  # [B, T, H, W, D]
+        else:
+            x_B_T_H_W_D = x_B_T_H_W_D + self.pos_embedder(x_B_T_H_W_D)  # [B, T, H, W, D]
+        return x_B_T_H_W_D, None, extra_pos_emb

ar_diffusion_decoder_utils.py

@@ -0,0 +1,119 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn.functional as F
+
+
+def split_with_overlap(video_BCTHW, num_video_frames, overlap=2, tobf16=True):
+    """
+    Splits the video tensor into chunks of num_video_frames with a specified overlap.
+
+    Args:
+    - video_BCTHW (torch.Tensor): Input tensor with shape [Batch, Channels, Time, Height, Width].
+    - num_video_frames (int): Number of frames per chunk.
+    - overlap (int): Number of overlapping frames between chunks.
+
+    Returns:
+    - List of torch.Tensors: List of video chunks with overlap.
+    """
+    # Get the dimensions of the input tensor
+    B, C, T, H, W = video_BCTHW.shape
+
+    # Ensure overlap is less than num_video_frames
+    assert overlap < num_video_frames, "Overlap should be less than num_video_frames."
+
+    # List to store the chunks
+    chunks = []
+
+    # Step size for the sliding window
+    step = num_video_frames - overlap
+
+    # Loop through the time dimension (T) with the sliding window
+    for start in range(0, T - overlap, step):
+        end = start + num_video_frames
+        # Handle the case when the last chunk might go out of bounds
+        if end > T:
+            # Get the last available frame
+            num_padding_frames = end - T
+            chunk = F.pad(video_BCTHW[:, :, start:T, :, :], (0, 0, 0, 0, 0, num_padding_frames), mode="reflect")
+        else:
+            # Regular case: no padding needed
+            chunk = video_BCTHW[:, :, start:end, :, :]
+        if tobf16:
+            chunks.append(chunk.to(torch.bfloat16))
+        else:
+            chunks.append(chunk)
+    return chunks
+
+
+def linear_blend_video_list(videos, D):
+    """
+    Linearly blends a list of videos along the time dimension with overlap length D.
+
+    Parameters:
+    - videos: list of video tensors, each of shape [b, c, t, h, w]
+    - D: int, overlap length
+
+    Returns:
+    - output_video: blended video tensor of shape [b, c, L, h, w]
+    """
+    assert len(videos) >= 2, "At least two videos are required."
+    b, c, t, h, w = videos[0].shape
+    N = len(videos)
+
+    # Ensure all videos have the same shape
+    for video in videos:
+        assert video.shape == (b, c, t, h, w), "All videos must have the same shape."
+
+    # Calculate total output length
+    L = N * t - D * (N - 1)
+    output_video = torch.zeros((b, c, L, h, w), device=videos[0].device)
+
+    output_index = 0  # Current index in the output video
+
+    for i in range(N):
+        if i == 0:
+            # Copy frames from the first video up to t - D
+            output_video[:, :, output_index : output_index + t - D, :, :] = videos[i][:, :, : t - D, :, :]
+            output_index += t - D
+        else:
+            # Blend overlapping frames between videos[i-1] and videos[i]
+            blend_weights = torch.linspace(0, 1, steps=D, device=videos[0].device)
+
+            for j in range(D):
+                w1 = 1 - blend_weights[j]
+                w2 = blend_weights[j]
+                frame_from_prev = videos[i - 1][:, :, t - D + j, :, :]
+                frame_from_curr = videos[i][:, :, j, :, :]
+                output_frame = w1 * frame_from_prev + w2 * frame_from_curr
+                output_video[:, :, output_index, :, :] = output_frame
+                output_index += 1
+
+            if i < N - 1:
+                # Copy non-overlapping frames from current video up to t - D
+                frames_to_copy = t - 2 * D
+                if frames_to_copy > 0:
+                    output_video[:, :, output_index : output_index + frames_to_copy, :, :] = videos[i][
+                        :, :, D : t - D, :, :
+                    ]
+                    output_index += frames_to_copy
+            else:
+                # For the last video, copy frames from D to t
+                frames_to_copy = t - D
+                output_video[:, :, output_index : output_index + frames_to_copy, :, :] = videos[i][:, :, D:, :, :]
+                output_index += frames_to_copy
+
+    return output_video

ar_model.py

@@ -0,0 +1,596 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Set
+
+from .misc import misc, Color, timer
+import torch
+from safetensors.torch import load_file
+from torch.nn.modules.module import _IncompatibleKeys
+
+from .ar_config_base_model import ModelConfig
+from .ar_config_base_tokenizer import TokenizerConfig
+from .ar_module_mm_projector import MultimodalProjector
+from .ar_network_transformer import Transformer
+from .ar_network_vit import VisionTransformer, get_vit_config
+from .ar_tokenizer_tokenizer import DiscreteMultimodalTokenizer, update_vocab_size
+from .ar_utils_checkpoint import (
+    get_partial_state_dict,
+    process_state_dict,
+    substrings_to_ignore,
+)
+from .ar_utils_sampling import decode_n_tokens, decode_one_token, prefill
+from .log import log
+
+
+class AutoRegressiveModel(torch.nn.Module):
+    """
+    A class to build and use a AutoRegressiveModel model for text generation.
+
+    Methods:
+        build: Build a AutoRegressiveModel instance by initializing and loading a model checkpoint.
+        generate: Generate text sequences based on provided prompts using the language generation model.
+    """
+
+    def __init__(
+        self,
+        model: Transformer = None,
+        tokenizer: DiscreteMultimodalTokenizer = None,
+        config: ModelConfig = None,
+        vision_encoder: VisionTransformer = None,
+        mm_projector: MultimodalProjector = None,
+    ):
+        """
+        Initialize the AutoRegressiveModel instance with a model and tokenizer.
+
+        Args:
+            model (Transformer): The Transformer model for text generation.
+            tokenizer (Tokenizer): The tokenizer for encoding and decoding text.
+            config (Config): The configuration for the AutoRegressiveModel model.
+            vision_encoder (VisionTransformer): The vision encoder for the AutoRegressiveModel model.
+            mm_projector (MultimodalProjector): The multi-modal projector for the AutoRegressiveModel model.
+        """
+        super().__init__()
+        self.model = model
+        self.tokenizer = tokenizer
+        self.config = config
+
+        self.vision_encoder = vision_encoder
+        self.mm_projector = mm_projector
+
+    @property
+    def precision(self):
+        return self.model.precision
+
+    def get_num_params(
+        self,
+    ) -> int:
+        """
+        Return the number of parameters in the model.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        return n_params
+
+    def load_ar_model(
+        self,
+        tokenizer_config,
+    ):
+        """
+        Load the AR model.
+        """
+        model_config = self.config
+        ckpt_path = model_config.ckpt_path
+        with timer(f"loading checkpoint from {ckpt_path}"):
+            if ckpt_path.endswith("safetensors"):
+                # Load with safetensors API
+                checkpoint = load_file(ckpt_path, device="cpu")
+            else:
+                # The pytorch version
+                checkpoint = torch.load(
+                    ckpt_path,
+                    map_location="cpu",
+                    mmap=True,  # load the checkpoint in memory-mapped mode
+                    weights_only=True,
+                )
+        llm_checkpoint = checkpoint["model"] if "model" in checkpoint else checkpoint
+        orig_precision = torch.get_default_dtype()
+        precision = getattr(torch, model_config.precision)
+        torch.set_default_dtype(precision)
+        log.debug(f"Setting torch default dtype to {precision}")
+
+        model = Transformer(
+            params=model_config,
+            tokenizer_config=tokenizer_config,
+        )
+        log.debug(
+            f"tokenizer tokenizer_config.video_tokenizer.vocab_size {tokenizer_config.video_tokenizer.vocab_size}"
+        )
+        vocab_size = update_vocab_size(
+            existing_vocab_size=0,
+            to_be_added_vocab_size=tokenizer_config.video_tokenizer.vocab_size,
+            training_type=tokenizer_config.training_type,
+            add_special_tokens=False,
+        )
+        log.debug(
+            f"tokenizer tokenizer_config.video_tokenizer.vocab_size {tokenizer_config.video_tokenizer.vocab_size}  vocab_size {vocab_size}"
+        )
+        # Perform vocab expansion
+        if vocab_size > model.vocab_size:
+            log.debug(f"Expanding vocab size to {vocab_size}")
+            # For text-to-video training, we only expand the embedding layer but not the output (unembedding) layer,
+            expand_output_layer = not (tokenizer_config.training_type == "text_to_video")
+            model.expand_vocab(
+                vocab_size,
+                init_method="gaussian",
+                expand_output_layer=expand_output_layer,
+            )
+        # Remove the "model." prefix in the state_dict
+        llm_checkpoint = process_state_dict(llm_checkpoint, prefix_to_remove="model.")
+        with timer("loading state_dict into model"):
+            missing_keys, _ = model.load_state_dict(llm_checkpoint, strict=True)
+        # Remove keys with "_extra_state" suffix in missing_keys (defined by TransformerEngine for FP8 usage)
+        missing_keys = [k for k in missing_keys if not k.endswith("_extra_state")]
+        assert len(missing_keys) == 0, f"Missing keys: {missing_keys}"
+
+        self.model = model.to(precision).to("cuda")
+        torch.set_default_dtype(orig_precision)  # Reset the default dtype to the original value
+
+    def load_tokenizer(self, tokenizer_config):
+        """
+        Load the tokenizer.
+        """
+        self.tokenizer = DiscreteMultimodalTokenizer(tokenizer_config)
+
+    @staticmethod
+    def build(
+        model_config: ModelConfig = ModelConfig(),
+        tokenizer_config: TokenizerConfig = None,
+    ) -> "AutoRegressiveModel":
+        """
+        Build a AutoRegressiveModel instance by initializing and loading a model checkpoint.
+
+        Args:
+            model_config (ModelConfig, optional): The model configuration for the AutoRegressiveModel instance. Defaults to ModelConfig().
+            tokenizer_config (TokenizerConfig, optional): The tokenizer configuration for the AutoRegressiveModel instance. Defaults to None.
+            download_rank_sync (bool, optional): Whether to download the checkpoint in a rank-synchronized manner. Defaults to True.
+        Returns:
+            AutoRegressiveModel: An instance of the AutoRegressiveModel class with the loaded model and tokenizer.
+
+        Raises:
+            AssertionError: If there are no checkpoint files in the specified directory.
+
+        Note:
+            This method sets the device to CUDA and loads the pre-trained model and tokenizer.
+        """
+        # Initialize model configuration parameters
+        config_params = {}
+
+        # Load checkpoint and model parameters
+
+        if model_config.ckpt_path is None:
+            # If ckpt_path is not provided, we assume the model checkpoint is saved in the ckpt_dir
+            ckpt_dir = model_config.ckpt_dir
+
+            # We prioritize safetensors version over the pytorch version, since the former is
+            # much faster for checkpoint loading.
+            checkpoints = sorted(Path(ckpt_dir).glob("*.safetensors"))
+            if len(checkpoints) == 0:
+                checkpoints = sorted(Path(ckpt_dir).glob("*.pth"))
+
+            assert len(checkpoints) > 0, f"no checkpoint files found in {ckpt_dir}"
+            assert (
+                len(checkpoints) == 1
+            ), f"multiple checkpoint files found in {ckpt_dir} (currently only one is supported)"
+            ckpt_path = str(checkpoints[0])  # Assuming single checkpoint for non-parallel case
+
+            if os.path.exists(Path(ckpt_dir) / "config.json"):
+                with open(Path(ckpt_dir) / "config.json", "r") as f:
+                    config_params = json.loads(f.read())
+            else:
+                log.info(
+                    f"No params.json found in the checkpoint directory ({ckpt_dir}). " f"Using default model config."
+                )
+
+        else:
+            # If ckpt_path is provided, we load the model from the specified path,
+            # and use the default model configuration
+            ckpt_path = model_config.ckpt_path
+
+        for key, value in config_params.items():
+            if hasattr(model_config, key):
+                # Override the default model configuration with the parameters from the checkpoint
+                setattr(model_config, key, value)
+
+        with timer(f"loading checkpoint from {ckpt_path}"):
+            if ckpt_path.endswith("safetensors"):
+                # Load with safetensors API
+                checkpoint = load_file(ckpt_path, device="cpu")
+            else:
+                # The pytorch version
+                checkpoint = torch.load(
+                    ckpt_path,
+                    map_location="cpu",
+                    mmap=True,  # load the checkpoint in memory-mapped mode
+                    weights_only=True,
+                )
+        llm_checkpoint = checkpoint["model"] if "model" in checkpoint else checkpoint
+
+        if model_config.vision_encoder is not None:
+            # Take the LLM weights (starting with "model.") from the VLM checkpoint
+            llm_checkpoint = get_partial_state_dict(llm_checkpoint, prefix="model.")
+        if model_config.vision_encoder is not None:
+            # For vanilla VLM ckpt before fine-tuning, `checkpoint['model']` only contains LLM weights, and `checkpoint['vision_encoder']`
+            #   and `checkpoint['mm_projector']` are both for those weights
+            # For fine-tuned VLM ckpt, `checkpoint['model']` contains all LLM, mm_projector and vision_encoder weights
+            if "vision_encoder" in checkpoint:
+                log.debug("Using pretrained vision_encoder")
+                vit_checkpoint = checkpoint["vision_encoder"]
+            else:
+                log.debug("Using fine-tuned vision_encoder")
+                vit_checkpoint = get_partial_state_dict(llm_checkpoint, prefix="vision_encoder.")
+                vit_checkpoint = process_state_dict(vit_checkpoint, prefix_to_remove="vision_encoder.")
+            if "mm_projector" in checkpoint:
+                log.debug("Using pretrained mm_projector")
+                projector_checkpoint = checkpoint["mm_projector"]
+            else:
+                log.debug("Using fine-tuned mm_projector")
+                projector_checkpoint = get_partial_state_dict(llm_checkpoint, prefix="mm_projector.")
+                projector_checkpoint = process_state_dict(projector_checkpoint, prefix_to_remove="mm_projector.")
+            assert (
+                len(vit_checkpoint) > 0 and len(projector_checkpoint) > 0
+            ), "vit_checkpoint and projector_checkpoint cannot be empty. We do not support random initialization for vision_encoder and mm_projector."
+
+        tokenizer = DiscreteMultimodalTokenizer(tokenizer_config)
+        orig_precision = torch.get_default_dtype()
+        precision = getattr(torch, model_config.precision)
+        torch.set_default_dtype(precision)
+        log.debug(f"Setting torch default dtype to {precision}")
+
+        model = Transformer(
+            params=model_config,
+            tokenizer_config=tokenizer_config,
+        )
+        model_kwargs = {}
+
+        if model_config.vision_encoder is not None:
+            assert model_config.mm_projector is not None, "mm_projector must be provided if vision_encoder is provided."
+            vit_config = get_vit_config(model_config.vision_encoder)
+            vision_encoder = VisionTransformer.build(
+                vit_config,
+            )
+
+            mm_projector = MultimodalProjector(
+                mm_projector_type=model_config.mm_projector, in_dim=vit_config["dim"], out_dim=model_config["dim"]
+            )
+            model_kwargs.update({"vision_encoder": vision_encoder, "mm_projector": mm_projector})
+
+        # Perform vocab expansion
+        if tokenizer.vocab_size > model.vocab_size:
+            log.debug(f"Expanding vocab size to {tokenizer.vocab_size}")
+            # For text-to-video training, we only expand the embedding layer but not the output (unembedding) layer,
+            expand_output_layer = not (tokenizer.training_type == "text_to_video")
+            model.expand_vocab(
+                tokenizer.vocab_size,
+                init_method="gaussian",
+                expand_output_layer=expand_output_layer,
+            )
+
+        # Remove the "model." prefix in the state_dict
+        llm_checkpoint = process_state_dict(llm_checkpoint, prefix_to_remove="model.")
+        with timer("loading state_dict into model"):
+            missing_keys, unexpected_keys = model.load_state_dict(llm_checkpoint, strict=True)
+        # Remove keys with "_extra_state" suffix in missing_keys (defined by TransformerEngine for FP8 usage)
+        missing_keys = [k for k in missing_keys if not k.endswith("_extra_state")]
+        assert len(missing_keys) == 0, f"Missing keys: {missing_keys}"
+
+        if model_config.vision_encoder is not None:
+            vision_encoder.load_state_dict(vit_checkpoint)
+            mm_projector.load_state_dict(projector_checkpoint)
+            if model_config.vision_encoder_in_channels != 3:
+                vision_encoder.expand_in_channels(model_config.vision_encoder_in_channels)
+
+        model = model.to(precision)  # ensure model parameters are in the correct precision
+        log.debug(f"Model config: {model_config}")
+
+        model_class = AutoRegressiveModel
+
+        torch.set_default_dtype(orig_precision)  # Reset the default dtype to the original value
+
+        return model_class(model, tokenizer, model_config, **model_kwargs)
+
+    @torch.no_grad()
+    def generate(
+        self,
+        prompt_tokens: List[List[int]] | torch.Tensor,
+        max_gen_len: int,
+        temperature: float = 1.0,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        num_gen_seq: int = 1,
+        logprobs: bool = False,
+        echo: bool = False,
+        seed: int = None,
+        context: Optional[torch.Tensor] = None,
+        context_mask: Optional[torch.Tensor] = None,
+        compile_sampling: bool = True,
+        compile_prefill: bool = False,
+        verbose: bool = True,
+        stop_tokens: Optional[Set[int]] = None,
+        images: Optional[torch.Tensor] = None,
+    ):
+        """
+        Autoregressive generation built upon the gpt-fast implementation (https://github.com/pytorch-labs/gpt-fast).
+
+        Args:
+            prompt_tokens (List[List[int]] | torch.Tensor): A single prompt of shape (1, seq_len).
+            max_gen_len (int): Maximum length of the generated text sequence.
+            temperature (float, optional): Temperature value for controlling randomness in sampling. Defaults to 0.6.
+            top_k (int, optional): Top-k value for top-k sampling. Defaults to None.
+            top_p (float, optional): Top-p probability threshold for nucleus sampling. Defaults to None.
+            num_gen_seq (int, optional): Number of outputs to generate given the same prompt. Defaults to 1. When temperature == 0, num_gen_seq must be 1 because the generation is deterministic.
+            echo (bool, optional): Flag indicating whether to include prompt tokens in the generated output. Defaults to False.
+            logit_clipping_range (list, optional): Range of logits to clip. Defaults to [].
+            seed (int, optional): Random seed for reproducibility. Defaults to None.
+            compile_sampling (bool, optional): Flag indicating whether to compile the decoding function. Defaults to True.
+            compile_prefill (bool, optional): Flag indicating whether to compile the prefill function. Defaults to False.
+            verbose (bool, optional): Flag indicating whether to print the the time. Defaults to False.
+        """
+        assert top_k is None or top_p is None, f"Only one of top_k ({top_k} or top_p ({top_p} should be specified."
+        if temperature == 0:
+            top_p, top_k = None, None
+            log.debug("Setting top_p and top_k to None because temperature is 0")
+        if top_p is not None:
+            log.debug(f"Using top-p sampling with p={top_p} and temperature={temperature}")
+        elif top_k is not None:
+            log.debug(f"Using top-k sampling with k={top_k} and temperature={temperature}")
+        else:
+            log.debug("Not applying top-k or top-p sampling. Will use top-k sampling with k=None")
+
+        orig_precision = torch.get_default_dtype()
+        torch.set_default_dtype(self.precision)
+
+        torch._inductor.config.coordinate_descent_tuning = True
+        torch._inductor.config.triton.unique_kernel_names = True
+        # Experimental features to reduce compilation times, will be on by default in future
+        torch._inductor.config.fx_graph_cache = True
+
+        if seed is not None:
+            misc.set_random_seed(seed)
+
+        assert not logprobs, "logprobs are not supported for fast_generate yet"
+        # Examine if the function prefil and decode_one_token functions are compiled yet. If not, compile them based on the flags
+        if compile_sampling and not getattr(self, "inference_decode_compiled", False):
+            self.decode_one_token = torch.compile(decode_one_token, mode="reduce-overhead", fullgraph=True)
+            self.inference_decode_compiled = True
+            log.info("Compiled AR sampling function. Note: the first run will be slower due to compilation")
+        if compile_prefill and not getattr(self, "inference_prefill_compiled", False):
+            self.prefill = torch.compile(prefill, fullgraph=True, dynamic=True)
+            self.inference_prefill_compiled = True
+            log.info("Compiled prefill function. Note: the first run will be slower due to compilation")
+
+        if not hasattr(self, "decode_one_token"):
+            self.decode_one_token = decode_one_token
+        if not hasattr(self, "prefill"):
+            self.prefill = prefill
+
+        # Initialization and Assertions
+        if isinstance(self.model.params, list):
+            # During training, model.params is a list
+            log.debug(
+                f"Find self.model.params is a list, use self.config instead. Get max_batch_size={self.config.max_batch_size}, max_seq_len={self.config.max_seq_len}"
+            )
+            params = self.config
+        else:
+            params = self.model.params
+        if isinstance(prompt_tokens, list):
+            prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.long, device="cuda")
+        if prompt_tokens.ndim == 1:
+            prompt_tokens = prompt_tokens.view(1, -1)
+        else:
+            assert prompt_tokens.ndim == 2, f"prompt_tokens has shape {prompt_tokens.shape}"
+        batch_size, prompt_len = prompt_tokens.shape
+        total_len = min(params.max_seq_len, max_gen_len + prompt_len)
+        if max_gen_len + prompt_len > params.max_seq_len:
+            log.warning(
+                f"max_gen_len + prompt_len={max_gen_len + prompt_len} exceeds max_seq_len={params.max_seq_len}, truncate max_gen_len to {params.max_seq_len - prompt_len}"
+            )
+            max_gen_len = params.max_seq_len - prompt_len
+
+        if context_mask is not None:
+            context_mask = context_mask.to(dtype=torch.bool)
+            if context_mask.ndim == 2:
+                assert (
+                    context_mask.shape[0] == batch_size
+                ), f"batch_size mismatch: {context_mask.shape[0]} != {batch_size}"
+                # Unsqueeze it to make it of shape [batch_size, 1, 1, context_seq_len]
+                context_mask = context_mask.view(batch_size, 1, 1, -1)
+
+        if num_gen_seq > 1:
+            assert (
+                batch_size == 1
+            ), f"num_gen_seq > 1 is only supported for a single prompt, got {len(prompt_tokens)} prompts"
+            log.debug(f"Generating {num_gen_seq} sequences with the same prompt")
+            assert (
+                num_gen_seq <= params.max_batch_size
+            ), f"num_gen_seq={num_gen_seq} exceeds max_batch_size={params.max_batch_size}"
+            # repeat the prompt tokens for num_gen_seq times
+            prompt_tokens = prompt_tokens.repeat(num_gen_seq, 1)
+            assert prompt_tokens.shape == (
+                num_gen_seq,
+                prompt_len,
+            ), f"prompt_tokens must be of shape (num_gen_seq, seq_len), got {prompt_tokens.shape}"
+            batch_size = len(prompt_tokens)
+
+        # create an empty tensor of the expected final shape and fill in the current tokens
+        empty = torch.empty(batch_size, total_len, dtype=prompt_tokens.dtype, device=prompt_tokens.device)
+        empty[:, :prompt_len] = prompt_tokens
+        seq = empty
+        input_pos = torch.arange(0, prompt_len, device="cuda")
+
+        if verbose:
+            prefill_start = time.time()
+
+        if images is not None:
+            images = images.to(device=prompt_tokens.device, dtype=torch.bfloat16)
+            prompt_token_embeddings = self.embed_vision_language_features(prompt_tokens, images)
+        else:
+            prompt_token_embeddings = None
+
+        if context is not None:
+            context = context.to(device=prompt_tokens.device, dtype=self.precision)
+
+        # Prefill stage
+        next_token = self.prefill(
+            self.model,
+            input_pos=input_pos,
+            tokens=prompt_tokens if prompt_token_embeddings is None else None,
+            token_embeddings=prompt_token_embeddings,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            context=context,
+            context_mask=context_mask,
+        )
+        if verbose:
+            prefill_time = time.time() - prefill_start
+
+        seq[:, [prompt_len]] = next_token.to(dtype=seq.dtype)
+        input_pos = torch.tensor([prompt_len], dtype=torch.long, device="cuda")
+        stop_tokens = self.tokenizer.stop_tokens if stop_tokens is None else stop_tokens
+        stop_tokens = torch.tensor(list(stop_tokens), dtype=torch.long, device="cuda")
+
+        if verbose:
+            decode_start = time.time()
+        # Decode stage
+        generated_tokens = decode_n_tokens(
+            self.model,
+            next_token.view(batch_size, -1),
+            input_pos,
+            max_gen_len - 1,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            stop_tokens=stop_tokens,
+            decode_one_token_function=self.decode_one_token,
+            context=context,
+            context_mask=context_mask,
+        )
+        gen_len = len(generated_tokens)
+        if verbose:
+            decode_time = time.time() - decode_start
+            prefill_throughput = prompt_len / prefill_time
+            decode_throughput = gen_len / decode_time
+            log.debug(f"[Prefill] Time: {prefill_time:.2f}s; Throughput: {prefill_throughput:.2f} tokens/s")
+            log.debug(f"[Decode] Time: {decode_time:.2f}s; Throughput: {decode_throughput:.2f} tokens/s")
+
+        generated_tokens = torch.cat(generated_tokens, dim=1)
+
+        log.debug(f"generated_tokens: {generated_tokens.shape}")
+        seq = seq[:, : prompt_len + 1 + gen_len]
+        seq[:, prompt_len + 1 :] = generated_tokens
+        if not echo:
+            seq = seq[:, prompt_len:]
+
+        torch.set_default_dtype(orig_precision)  # Reset the default dtype to the original value
+
+        return seq, None
+
+    def embed_vision_language_features(self, input_ids: torch.Tensor, images: torch.tensor) -> torch.Tensor:
+        """
+        Embed vision and language features into a combined representation.
+
+        Args:
+            input_ids (torch.Tensor): Input token IDs.
+            images (torch.tensor): Input images.
+
+        Returns:
+            torch.Tensor: Combined vision-language features.
+
+        Raises:
+            AssertionError: If vision encoder or mm projector is not initialized,
+                            or if dimensions mismatch.
+        """
+        # Ensure vision encoder and mm projector are initialized
+        assert self.vision_encoder is not None
+        assert self.mm_projector is not None
+
+        # Get image token ID and validate it
+        image_token_id = self.vision_encoder.image_token_id
+        assert isinstance(image_token_id, int) and image_token_id >= 0, f"Invalid image_token_id: {image_token_id}"
+
+        # Identify text and image locations in the input
+        text_locations = input_ids != image_token_id
+        image_locations = input_ids == image_token_id
+
+        # Process text features
+        text_features = self.model.tok_embeddings(input_ids[text_locations])
+
+        # Process image features
+        images = images.to(device=text_features.device, dtype=text_features.dtype)
+        vit_outputs = self.vision_encoder(images)
+        image_features = self.mm_projector(vit_outputs)
+
+        # Get dimensions
+        B, seq_len = input_ids.shape
+        N_total = B * seq_len
+        N_txt, D_txt = text_features.shape
+        N_img, N_patch, D_img = image_features.shape
+
+        # Reshape image features
+        image_features = image_features.reshape(N_img * N_patch, D_img)
+
+        # Validate dimensions
+        assert D_txt == D_img, f"Text features dim {D_txt} should be equal to image features dim {D_img}"
+        assert (
+            N_total == N_txt + N_img * N_patch
+        ), f"seq_len {seq_len} should be equal to N_txt + N_img*N_Patch {(N_txt, N_img * N_patch, image_locations.sum().item())}"
+
+        # Combine text and image features
+        combined_features = torch.empty(
+            (B, seq_len, D_txt),
+            dtype=text_features.dtype,
+            device=text_features.device,
+        )
+        combined_features[text_locations, :] = text_features
+        combined_features[image_locations, :] = image_features
+
+        return combined_features
+
+    def state_dict(self, *args, **kwargs):
+        """
+        Process the state dict (e.g., remove "_extra_state" keys imposed by TransformerEngine for FP8).
+        """
+        state_dict = super().state_dict(*args, **kwargs)
+        return process_state_dict(state_dict)
+
+    def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True, assign: bool = False):
+        """
+        Ignore the missing keys with substrings matching `substring_to_ignore` (e.g., "_extra_state" keys imposed by
+        TransformerEngine for FP8).
+        """
+        state_dict = process_state_dict(state_dict)
+        missing_keys, unexpected_keys = super().load_state_dict(state_dict, strict=False, assign=assign)
+        actual_missing_keys = []
+        for key in missing_keys:
+            if not any(substring in key for substring in substrings_to_ignore):
+                actual_missing_keys.append(key)
+        if strict:
+            if len(actual_missing_keys) > 0 or len(unexpected_keys) > 0:
+                raise ValueError(f"Missing keys: {actual_missing_keys}\n\nUnexpected keys: {unexpected_keys}")
+        return _IncompatibleKeys(actual_missing_keys, unexpected_keys)

ar_module_attention.py

@@ -0,0 +1,262 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import Optional, Union
+
+import torch
+from torch import nn
+
+from .ar_module_embedding import RotaryPositionEmbedding
+from .ar_module_normalization import create_norm
+
+
+class Attention(nn.Module):
+    """
+    Attenion layer with KV cache.
+    """
+
+    def __init__(
+        self,
+        n_heads: int,
+        n_kv_heads: Union[int, None],
+        dim: int,
+        max_batch_size: int,
+        max_seq_len: int,
+        context_dim: Optional[int] = None,
+        use_qk_normalization: bool = False,
+        norm_type: str = "rmsnorm",
+        norm_eps: float = 1e-5,
+        causal_mask: Optional[bool] = True,
+        head_dim: Optional[int] = None,
+        fuse_qkv: bool = False,
+        precision: str = "bfloat16",
+        attn_type: str = "self",
+    ):
+        """
+        Initializes the GQA module.
+
+        Args:
+            n_heads (int): The number of attention heads.
+            n_kv_heads (int, optional): The number of key-value attention heads. None defaults to n_heads.
+            dim (int): The dimensionality of the input and output.
+            max_batch_size (int): The maximum batch size.
+            max_seq_len (int): The maximum sequence length.
+            context_dim (int, optional): The dimensionality of the context for cross-attn. Defaults to None.
+            use_qk_normalization (bool, optional): Whether to apply QK normalization. Defaults to False.
+            norm_type (str, optional): The type of normalization layer. Defaults to "rmsnorm".
+            norm_eps (float, optional): The epsilon value for normalization. Defaults to 1e-5.
+            causal_mask (bool, optional): Whether to use causal mask. Defaults to True.
+            head_dim (int, optional): The dimensionality of each attention head. If None, defaults to dim // n_heads.
+            fuse_qkv (bool, optional): Whether to fuse QKV. Defaults to False.
+            precision (str, optional): The precision of the module. Defaults to "bfloat16".
+            attn_type (str, optional): The type of attention. Defaults to "self".
+        """
+        super().__init__()
+        assert attn_type in ["self", "cross", "full"], f"Invalid attention type: {attn_type}"
+        self.attn_type = attn_type
+        context_dim = dim if context_dim is None else context_dim
+
+        self.dim = dim
+        self.context_dim = context_dim
+        self.n_kv_heads = n_heads if n_kv_heads is None else n_kv_heads
+        self.n_local_kv_heads = self.n_kv_heads
+        self.n_local_heads = n_heads
+        self.n_rep = self.n_local_heads // self.n_local_kv_heads
+        self.head_dim = dim // n_heads if head_dim is None else head_dim
+        self.causal_mask = causal_mask
+        self.fuse_qkv = fuse_qkv
+        self.precision = precision
+
+        if fuse_qkv:
+            assert context_dim == dim, f"Fuse QKV requires context_dim ({context_dim}) to be equal to dim ({dim})"
+            self.total_local_head_dim = (self.n_local_heads + 2 * self.n_local_kv_heads) * self.head_dim
+            self.wqkv = nn.Linear(dim, self.total_local_head_dim, bias=False)
+            # Register hook to load fused QKV weights
+            self._register_load_state_dict_pre_hook(self.load_hook)
+        else:
+            self.wq = nn.Linear(dim, self.n_local_heads * self.head_dim, bias=False)
+            self.wk = nn.Linear(context_dim, self.n_local_kv_heads * self.head_dim, bias=False)
+            self.wv = nn.Linear(context_dim, self.n_local_kv_heads * self.head_dim, bias=False)
+        self.wo = nn.Linear(self.n_local_heads * self.head_dim, dim, bias=False)
+
+        self.max_batch_size = max_batch_size
+        self.max_seq_len = max_seq_len
+
+        if self.attn_type == "self":
+            # Cache for key and value tensors
+            self.init_kv_cache()
+
+        # QK normalization layers
+        if use_qk_normalization:
+            self.q_norm = create_norm(norm_type, dim=self.head_dim, eps=norm_eps)
+            self.k_norm = create_norm(norm_type, dim=self.head_dim, eps=norm_eps)
+
+        self.use_qk_normalization = use_qk_normalization
+
+        self.to(dtype=getattr(torch, self.precision))
+
+    def load_hook(self, state_dict, prefix, *args):
+        if prefix + "wq.weight" in state_dict:
+            wq = state_dict.pop(prefix + "wq.weight")
+            wk = state_dict.pop(prefix + "wk.weight")
+            wv = state_dict.pop(prefix + "wv.weight")
+            state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv])
+
+    def init_kv_cache(self, dtype=None):
+        cache_shape = (self.max_batch_size, self.n_local_kv_heads, self.max_seq_len, self.head_dim)
+        if dtype is None:
+            dtype = getattr(torch, self.precision)
+        if self.attn_type == "self":
+            self.cache_k = torch.zeros(cache_shape, dtype=dtype).cuda()
+            self.cache_v = torch.zeros(cache_shape, dtype=dtype).cuda()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: RotaryPositionEmbedding,
+        input_pos: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        context: Optional[torch.Tensor] = None,
+    ):
+        """
+        Forward pass of GQA.
+
+        Args:
+            x: The input tensor of shape (batch_size, seq_len, dim).
+            rope: The rotary positional embedding module.
+            input_pos: The starting position of the current sequence.
+            mask: The attention mask tensor.
+            context: The context tensor of shape (batch_size, context_len, dim).
+
+        Returns:
+            The output tensor after applying GQA.
+        """
+        bsz, seqlen, _ = x.shape
+
+        # Use one single module to handle both self-attn and cross-attn
+        context = x if context is None else context
+        context_len = seqlen if context is None else context.shape[1]
+
+        if self.fuse_qkv:
+            q_size = self.n_local_heads * self.head_dim
+            kv_size = self.n_local_kv_heads * self.head_dim
+            xq, xk, xv = self.wqkv(x).split([q_size, kv_size, kv_size], dim=-1)
+        else:
+            # Compute query, key, and value projections
+            xq, xk, xv = self.wq(x), self.wk(context), self.wv(context)
+
+        # Reshape projections
+        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+        xk = xk.view(bsz, context_len, self.n_local_kv_heads, self.head_dim)
+        xv = xv.view(bsz, context_len, self.n_local_kv_heads, self.head_dim)
+
+        # QK normalization
+        if self.use_qk_normalization:
+            xq = self.q_norm(xq)
+            xk = self.k_norm(xk)
+
+        # Apply rotary positional embeddings to queries and keys
+        # Only apply RoPE to self-attention!
+        if self.attn_type in ["self", "full"]:
+            xq, xk = rope(xq, xk, input_pos, seqlen)
+
+        xq, xk, xv = map(lambda x: x.transpose(1, 2), (xq, xk, xv))
+        # xq: (bs, n_local_heads, seqlen, head_dim)
+        # xk: (bs, n_kv_heads, cache_len + context_len, head_dim)
+        # xv: (bs, n_kv_heads, cache_len + context_len, head_dim)
+        if self.attn_type == "self":
+            # Update cache with current key and value tensors
+            assert input_pos is not None
+            self.cache_k[:bsz, :, input_pos] = xk
+            self.cache_v[:bsz, :, input_pos] = xv
+            keys, values = (
+                self.cache_k[:bsz, :, :],
+                self.cache_v[:bsz, :, :],
+            )
+        else:
+            keys, values = xk, xv
+
+        # Repeat keys and values if necessary
+        keys = keys.repeat_interleave(self.n_rep, dim=1)  # (bs, n_local_heads, cache_len + context_len, head_dim)
+        values = values.repeat_interleave(self.n_rep, dim=1)  # (bs, n_local_heads, cache_len + context_len, head_dim)
+
+        # For self-attention, `is_causal` should be set to False when KV cache is pre-computed and used,
+        # since the masking is handled outside this attention module.
+        # For cross-attention, it's always full-attn without causal mask
+        is_causal = False
+        output = scaled_dot_product_attention(
+            xq,
+            keys,
+            values,
+            head_dim=self.head_dim,
+            mask=mask,
+            is_causal=is_causal,
+            dropout_p=0.0,
+        )
+        output = output.view(bsz, seqlen, -1)
+        output = self.wo(output)
+        return output
+
+
+def scaled_dot_product_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    head_dim: int,
+    mask: Optional[torch.Tensor] = None,
+    is_causal: Optional[bool] = None,
+    dropout_p: float = 0.0,
+) -> torch.Tensor:
+    """
+    PyTorch's native implementation of Flash Attention 2.
+
+    If `is_causal` is given, then the causal attention mask is applied accordingly:
+    - If `is_causal` is True, the standard upper-left causal attention masking is applied.
+    - If `is_causal` is False, no attention mask is applied, unless an explicit mask tensor is
+      provided (i.e., `mask is not None`).
+
+    If `is_causal` is not given (i.e., `is_causal is None`), then the attention mask is applied
+    based on the provided mask tensor:
+    - If no explicit attention mask is given (i.e., `mask is None`), `is_causal` is set to True,
+    leading to the standard upper-left causal attention masking.
+    - If an attention mask is given (i.e., `mask is not None`), the provided mask is used,
+    and `is_causal` is set to False.
+
+    Args:
+        q (torch.Tensor): Query tensor
+        k (torch.Tensor): Key tensor
+        v (torch.Tensor): Value tensor
+        head_dim (int): Dimension of each attention head
+        mask (Optional[torch.Tensor], optional): Attention mask. Defaults to None.
+        is_causal (Optional[bool], optional): Whether to apply causal attention mask. Defaults to None.
+        dropout_p (float, optional): Dropout rate. Defaults to 0.0.
+
+    Returns:
+        torch.Tensor: Output tensor after applying scaled dot-product attention
+    """
+    scale = 1.0 / math.sqrt(head_dim)
+    if is_causal is None:
+        is_causal = mask is None
+    y = torch.nn.functional.scaled_dot_product_attention(
+        q,
+        k,
+        v,
+        attn_mask=mask,
+        dropout_p=dropout_p,
+        scale=scale,
+        is_causal=is_causal,
+    )
+    return y.transpose(1, 2).contiguous()

ar_module_embedding.py

@@ -0,0 +1,491 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import List, Optional, Tuple
+
+import numpy as np
+import torch
+from einops import rearrange, repeat
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum("m,d->md", pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out)  # (M, D/2)
+    emb_cos = np.cos(out)  # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+def _rotate_half_te(x: torch.Tensor) -> torch.Tensor:
+    """
+    change sign so the last dimension becomes [-odd, +even].
+    Adopted from TransformerEngine.
+    Source: https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py
+    """
+    x = x.view(x.shape[:-1] + torch.Size((2, x.shape[-1] // 2)))
+    x1, x2 = x.unbind(dim=-2)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_pos_emb_te(
+    t: torch.Tensor,
+    cos_freqs: torch.Tensor,
+    sin_freqs: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Apply rotary positional embedding tensor to the input tensor.
+    Adopted from TransformerEngine.
+    Source: https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py
+
+    Parameters
+    ----------
+    t: torch.Tensor
+        Input tensor of shape `[b, s, h, d]`, on which
+        rotary positional embedding will be applied.
+    cos_freqs: torch.Tensor
+        Cosine component of rotary positional embedding tensor of shape `[s, 1, 1, d]` and dtype 'float',
+    sin_freqs: torch.Tensor
+        Sine component of rotary positional embedding tensor of shape `[s, 1, 1, d]` and dtype 'float',
+    """
+    rot_dim = cos_freqs.shape[-1]
+    # ideally t_pass is empty so rotary pos embedding is applied to all tensor t
+    t, t_pass = t[..., :rot_dim], t[..., rot_dim:]
+    # first part is cosine component
+    # second part is sine component, need to change signs with _rotate_half method
+    t = (t * cos_freqs) + (_rotate_half_te(t) * sin_freqs)
+    output = torch.cat((t, t_pass), dim=-1)
+    return output
+
+
+class RotaryPositionEmbedding(torch.nn.Module):
+    """
+    Rotary Position Embedding module as described in the paper:
+    https://arxiv.org/abs/2104.09864
+
+    This module implements rotary positional embeddings, which are used to
+    enhance the performance of transformer models.
+
+    Args:
+        dim (int): Dimensionality of the input tensor.
+        max_position_embeddings (Optional[int]): Maximum position embeddings.
+        original_max_position_embeddings (Optional[int]): Original maximum position embeddings.
+        rope_theta (Optional[float]): Base for the frequency calculation.
+        apply_yarn (Optional[bool]): Whether to apply YaRN (Yet another Rotary).
+        scale (Optional[int]): Scaling factor for the frequency calculation.
+        extrapolation_factor (Optional[int]): Extrapolation factor for the frequency extension.
+        attn_factor (Optional[int]): Attention factor for the frequency calculation.
+        beta_fast (Optional[int]): Fast beta value for the YaRN frequency calculation.
+        beta_slow (Optional[int]): Slow beta value for the YaRN frequency calculation.
+        rope_dim (Optional[str]): Dimensionality of the RoPE. Choices: "1D", "2D", "3D".
+        latent_shape (Optional[List[int]]): Shape of the latent tensor for video or image inputs.
+        original_latent_shape (Optional[List[int]]): Original shape of the latent tensor for video or image inputs.
+        pad_to_multiple_of (Optional[int]): Pad the position embedding to a multiple of this value.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        max_position_embeddings: Optional[int] = None,
+        original_max_position_embeddings: Optional[int] = None,
+        rope_theta: Optional[float] = 10000.0,
+        apply_yarn: Optional[bool] = False,
+        scale: Optional[int] = None,
+        extrapolation_factor: Optional[int] = 1,
+        attn_factor: Optional[int] = 1,
+        beta_fast: Optional[int] = 32,
+        beta_slow: Optional[int] = 1,
+        rope_dim: Optional[str] = "1D",
+        latent_shape: Optional[List[int]] = None,
+        original_latent_shape: Optional[List[int]] = None,
+        pad_to_multiple_of: Optional[int] = None,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.original_max_position_embeddings = original_max_position_embeddings
+        self.rope_theta = rope_theta
+        self.apply_yarn = apply_yarn
+        self.scale = scale
+        self.extrapolation_factor = extrapolation_factor
+        self.attn_factor = attn_factor
+        self.beta_fast = beta_fast
+        self.beta_slow = beta_slow
+        self.mscale = 1.0
+        self.rope_dim = rope_dim
+        self.latent_shape = latent_shape
+        self.original_latent_shape = original_latent_shape
+        self.pad_to_multiple_of = pad_to_multiple_of
+        self.get_inv_freq(torch.cuda.current_device())
+
+    def get_mscale(self, scale: float = 1.0) -> float:
+        """Get the magnitude scaling factor for YaRN."""
+        if scale <= 1:
+            return 1.0
+        return 0.1 * math.log(scale) + 1.0
+
+    def forward(self, seq_len: Optional[int] = None) -> torch.Tensor:
+        """
+        Forward pass for the rotary position embedding.
+
+        Args:
+            seq_len (Optional[int]): Length of the sequence.
+
+        Returns:
+            torch.Tensor: The computed frequencies for positional embedding.
+        """
+
+        if self.apply_yarn and seq_len > self.max_seq_len_cached:
+            self.max_seq_len_cached = seq_len
+        self.freqs = self.compute_freqs()
+
+        return self.freqs
+
+    def compute_freqs(
+        self,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Compute the spatial frequencies for the latent tensor."""
+        self.seq = torch.arange(self.max_seq_len_cached, dtype=torch.float).cuda()
+        if self.rope_dim == "1D":
+            emb = torch.einsum("i,j->ij", self.seq, self.inv_freq)
+
+        elif self.rope_dim == "2D":
+            H, W = self.latent_shape
+            half_emb_h = torch.outer(self.seq[:H], self.spatial_inv_freq)
+            half_emb_w = torch.outer(self.seq[:W], self.spatial_inv_freq)
+            emb = torch.cat(
+                [
+                    repeat(half_emb_h, "h d -> h w d", w=W),
+                    repeat(half_emb_w, "w d -> h w d", h=H),
+                ]
+                * 2,
+                dim=-1,
+            )
+            emb = rearrange(emb, "h w d -> (h w) 1 1 d").float()
+
+        elif self.rope_dim == "3D":
+            T, H, W = self.latent_shape
+            half_emb_t = torch.outer(self.seq[:T], self.temporal_inv_freq)
+            half_emb_h = torch.outer(self.seq[:H], self.spatial_inv_freq)
+            half_emb_w = torch.outer(self.seq[:W], self.spatial_inv_freq)
+            emb = torch.cat(
+                [
+                    repeat(half_emb_t, "t d -> t h w d", h=H, w=W),
+                    repeat(half_emb_h, "h d -> t h w d", t=T, w=W),
+                    repeat(half_emb_w, "w d -> t h w d", t=T, h=H),
+                ]
+                * 2,
+                dim=-1,
+            )
+            emb = rearrange(emb, "t h w d -> (t h w) 1 1 d").float()
+        else:
+            raise ValueError(f"Invalid RoPE dimensionality: {self.rope_dim}")
+        return emb
+
+    def get_scale_factors(self, inv_freq: torch.Tensor, original_seq_len: int) -> torch.Tensor:
+        """Get the scale factors for YaRN."""
+        # Calculate the high and low frequency cutoffs for YaRN. Note: `beta_fast` and `beta_slow` are called
+        # `high_freq_factor` and `low_freq_factor` in the Llama 3.1 RoPE scaling code.
+        high_freq_cutoff = 2 * math.pi * self.beta_fast / original_seq_len
+        low_freq_cutoff = 2 * math.pi * self.beta_slow / original_seq_len
+        # Obtain a smooth mask that has a value of 0 for low frequencies and 1 for high frequencies, with linear
+        # interpolation in between.
+        smooth_mask = torch.clamp((inv_freq - low_freq_cutoff) / (high_freq_cutoff - low_freq_cutoff), min=0, max=1)
+        # For low frequencies, we scale the frequency by 1/self.scale. For high frequencies, we keep the frequency.
+        scale_factors = (1 - smooth_mask) / self.scale + smooth_mask
+        return scale_factors
+
+    def get_inv_freq(self, device: torch.device) -> None:
+        """Get the inverse frequency."""
+        if self.rope_dim == "1D":
+            assert self.max_position_embeddings is not None, "Max position embeddings required."
+            inv_freq = 1.0 / (
+                self.rope_theta ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
+            )
+            if self.apply_yarn:
+                assert self.original_max_position_embeddings is not None, "Original max position embeddings required."
+                assert self.beta_slow is not None, "Beta slow value required."
+                assert self.beta_fast is not None, "Beta fast value required."
+
+                scale_factors = self.get_scale_factors(inv_freq, self.original_max_position_embeddings)
+                # Apply the scaling factors to inv_freq.
+                inv_freq = inv_freq * scale_factors
+                # Set the magnitude scaling factor.
+                self.mscale = float(self.get_mscale(self.scale) * self.attn_factor)
+            self.max_seq_len_cached = self.max_position_embeddings
+            self.inv_freq = inv_freq
+
+        elif self.rope_dim == "2D":
+            assert self.latent_shape is not None, "Latent shape required."
+            dim_h = self.dim // 2
+            spatial_inv_freq = 1.0 / (
+                self.rope_theta ** torch.arange(0, dim_h, 2, dtype=torch.float32, device=device) / dim_h
+            )
+            if self.apply_yarn:
+                assert self.original_latent_shape is not None, "Original latent shape required."
+                assert self.beta_slow is not None, "Beta slow value required."
+                assert self.beta_fast is not None, "Beta fast value required."
+
+                scale_factors = self.get_scale_factors(spatial_inv_freq, self.original_latent_shape[0])
+                spatial_inv_freq = spatial_inv_freq * scale_factors
+                self.mscale = float(self.get_mscale(self.scale) * self.attn_factor)
+            self.spatial_inv_freq = spatial_inv_freq
+            self.max_seq_len_cached = max(self.latent_shape)
+
+        elif self.rope_dim == "3D":
+            assert self.latent_shape is not None, "Latent shape required."
+            dim_h = self.dim // 6 * 2
+            dim_t = self.dim - 2 * dim_h
+            self.dim_spatial_range = torch.arange(0, dim_h, 2)[: (dim_h // 2)].float().to(device) / dim_h
+            spatial_inv_freq = 1.0 / (self.rope_theta**self.dim_spatial_range)
+            self.dim_temporal_range = torch.arange(0, dim_t, 2)[: (dim_t // 2)].float().to(device) / dim_t
+            temporal_inv_freq = 1.0 / (self.rope_theta**self.dim_temporal_range)
+            if self.apply_yarn:
+                assert self.original_latent_shape is not None, "Original latent shape required."
+                assert self.beta_slow is not None, "Beta slow value required."
+                assert self.beta_fast is not None, "Beta fast value required."
+                scale_factors_spatial = self.get_scale_factors(spatial_inv_freq, self.original_latent_shape[1])
+                spatial_inv_freq = spatial_inv_freq * scale_factors_spatial
+                scale_factors_temporal = self.get_scale_factors(temporal_inv_freq, self.original_latent_shape[0])
+                temporal_inv_freq = temporal_inv_freq * scale_factors_temporal
+                self.mscale = float(self.get_mscale(self.scale) * self.attn_factor)
+            self.spatial_inv_freq = spatial_inv_freq
+            self.temporal_inv_freq = temporal_inv_freq
+            self.max_seq_len_cached = max(self.latent_shape)
+        else:
+            raise ValueError(f"Invalid RoPE dimensionality: {self.rope_dim}")
+
+        self.freqs = self.compute_freqs()
+
+
+class RotaryPositionEmbeddingPytorchV2(RotaryPositionEmbedding):
+    """
+    Rotary Position Embedding that works in the same way as the TransformerEngine RoPE
+    (https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py)
+
+    """
+
+    def __init__(
+        self,
+        seq_len: int,
+        training_type: str = None,
+        **kwargs,
+    ):
+        super().__init__(
+            **kwargs,
+        )
+        emb = self.create_rope_freqs(seq_len=seq_len, training_type=training_type)
+        emb = emb.transpose(0, 1).contiguous()  # [seq, 1, 1, dim] -> [1, seq, 1, dim]
+        assert emb.shape[0] == 1 and emb.shape[2] == 1, f"emb shape: {emb.shape}"
+        # cos/sin first then dtype conversion for better precision
+        self.register_buffer("cos_cached", torch.cos(emb), persistent=False)
+        self.register_buffer("sin_cached", torch.sin(emb), persistent=False)
+
+    def create_rope_freqs(self, seq_len: int, training_type: str = None) -> torch.Tensor:
+        """
+        Create rotary position embedding frequencies.
+
+        Args:
+            seq_len (int): Sequence length of a sample.
+
+        Returns:
+            torch.Tensor: The computed positional embeddings.
+        """
+        if self.rope_dim == "1D":
+            freqs = super().forward(seq_len=seq_len)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            emb = emb.reshape(emb.size(0), 1, 1, emb.size(1))
+
+        elif self.rope_dim in ["2D", "3D"]:
+            emb = super().forward(seq_len=seq_len)
+            if training_type == "text_to_video":
+                # since we added <bov> token at the beginning of the video for text2world, we also extend the position embedding by one token in the beginning
+                bov_pe = torch.zeros((1, *emb.shape[1:]), device=emb.device)
+                emb = torch.cat((bov_pe, emb), dim=0)
+        else:
+            raise ValueError(f"Invalid RoPE dimensionality: {self.rope_dim}")
+        if self.pad_to_multiple_of is not None and emb.shape[0] % self.pad_to_multiple_of != 0:
+            # Round up to the nearest multiple of pad_to_multiple_of
+            pad_len = self.pad_to_multiple_of - emb.shape[0] % self.pad_to_multiple_of
+            emb = torch.cat((emb, torch.zeros((pad_len, *emb.shape[1:]), device=emb.device)), dim=0)
+
+        return emb
+
+    def forward(
+        self, q: torch.Tensor, k: torch.Tensor, input_pos: Optional[torch.Tensor] = None, seq_len: Optional[int] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if q.dtype != self.cos_cached.dtype:
+            self.cos_cached = self.cos_cached.to(q.dtype)
+            self.sin_cached = self.sin_cached.to(q.dtype)
+
+        cos_emb = self.cos_cached
+        sin_emb = self.sin_cached
+        if input_pos is not None:
+            cos_emb = cos_emb[:, input_pos, :, :]
+            sin_emb = sin_emb[:, input_pos, :, :]
+        elif seq_len is not None:
+            cos_emb = cos_emb[:, :seq_len, :, :]
+            sin_emb = sin_emb[:, :seq_len, :, :]
+        q = _apply_rotary_pos_emb_te(q, cos_emb, sin_emb)
+        k = _apply_rotary_pos_emb_te(k, cos_emb, sin_emb)
+        return q, k
+
+
+class RotaryPositionEmbeddingPytorchV1(RotaryPositionEmbedding):
+    """
+    Rotary Position Embedding that works in the same way as
+    mistral_inference (https://github.com/mistralai/mistral-inference/blob/main/src/mistral_inference/rope.py)
+    or llama3 (https://github.com/meta-llama/llama3/blob/main/llama/model.py)
+
+    """
+
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        super().__init__(
+            **kwargs,
+        )
+        if self.rope_dim == "1D":
+            emb = torch.stack((self.freqs, self.freqs), dim=-1).reshape(*self.freqs.shape[:-1], -1)
+        elif self.rope_dim in ["2D", "3D"]:
+            emb = rearrange(self.freqs, "s 1 1 d -> s d").float()
+        self.register_buffer("cos_cached", (emb.cos() * self.mscale)[None, :, None, :], persistent=False)
+        self.register_buffer("sin_cached", (emb.sin() * self.mscale)[None, :, None, :], persistent=False)
+
+    def rotate_half(self, x: torch.Tensor) -> torch.Tensor:
+        """Rotate half the hidden dimensions of the input tensor."""
+        x_reshaped = x.reshape(*x.shape[:-1], -1, 2)
+        x1 = x_reshaped[..., 0]
+        x2 = x_reshaped[..., 1]
+        output = torch.stack((-x2, x1), dim=-1).reshape(*x.shape)
+        return output
+
+    def forward(
+        self, q: torch.Tensor, k: torch.Tensor, input_pos: Optional[torch.Tensor] = None, seq_len: Optional[int] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward pass for the rotary position embedding.
+
+        Args:
+            q (torch.Tensor): Query tensor.
+            k (torch.Tensor): Key tensor.
+            input_pos (Optional[torch.Tensor]): Starting position for the sequence.
+            seq_len (Optional[int]): Length of the sequence.
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor]: Rotated query and key tensors.
+        """
+        if self.apply_yarn and seq_len > self.max_seq_len_cached:
+            freqs = super().forward(seq_len)
+            if self.rope_dim == "1D":
+                emb = torch.stack((freqs, freqs), dim=-1).reshape(*freqs.shape[:-1], -1)
+            elif self.rope_dim in ["2D", "3D"]:
+                emb = rearrange(freqs, "s 1 1 d -> s d").float()
+            else:
+                raise ValueError(f"Invalid RoPE dimensionality: {self.rope_dim}")
+            self.register_buffer(
+                "cos_cached", (emb.cos() * self.mscale)[None, :, None, :].to(q.dtype), persistent=False
+            )
+            self.register_buffer(
+                "sin_cached", (emb.sin() * self.mscale)[None, :, None, :].to(q.dtype), persistent=False
+            )
+
+        if input_pos is not None:
+            cos_cached = self.cos_cached[:, input_pos]
+            sin_cached = self.sin_cached[:, input_pos]
+        else:
+            assert (
+                self.cos_cached.shape[1] >= seq_len
+            ), f"Invalid sequence length; cos_cached.shape {self.cos_cached.shape}, seq_len {seq_len}."
+            cos_cached = self.cos_cached[:, :seq_len, ...]
+            sin_cached = self.sin_cached[:, :seq_len, ...]
+        xq = q * cos_cached + self.rotate_half(q) * sin_cached
+        xk = k * cos_cached + self.rotate_half(k) * sin_cached
+
+        return xq.type_as(q), xk.type_as(k)
+
+
+class SinCosPosEmbAxisTE(torch.nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        latent_shape: Optional[List[int]] = None,
+        pad_to_multiple_of: Optional[int] = None,
+        dtype: torch.dtype = torch.bfloat16,
+        **kwargs,
+    ):
+        """
+        Args:
+            dim (int): Dimensionality of the input tensor.
+            latent_shape (Optional[List[int]]): Shape of the latent tensor for video or image inputs.
+            pad_to_multiple_of (Optional[int]): Pad the position embedding to a multiple of this value.
+            dtype (torch.dtype): Data type of the position embedding tensor.
+        """
+        super().__init__()
+        dim_h = dim // 6 * 2
+        dim_w = dim_h
+        dim_t = dim - 2 * dim_h
+        assert dim == dim_h + dim_w + dim_t, f"bad dim: {dim} != {dim_h} + {dim_w} + {dim_t}"
+        self.latent_shape = latent_shape
+        T, H, W = latent_shape
+        emb_h = get_1d_sincos_pos_embed_from_grid(dim_h, pos=np.arange(H))
+        emb_w = get_1d_sincos_pos_embed_from_grid(dim_w, pos=np.arange(W))
+        emb_t = get_1d_sincos_pos_embed_from_grid(dim_t, pos=np.arange(T))
+
+        self.register_buffer("pos_emb_h", torch.from_numpy(emb_h).to(dtype=dtype, device="cuda"), persistent=False)
+        self.register_buffer("pos_emb_w", torch.from_numpy(emb_w).to(dtype=dtype, device="cuda"), persistent=False)
+        self.register_buffer("pos_emb_t", torch.from_numpy(emb_t).to(dtype=dtype, device="cuda"), persistent=False)
+        self.pad_to_multiple_of = pad_to_multiple_of
+
+    def forward(
+        self,
+        training_type: str = None,
+    ) -> torch.Tensor:
+        T, H, W = self.latent_shape
+        emb = torch.cat(
+            [
+                repeat(self.pos_emb_t, "t d-> t h w d", h=H, w=W),
+                repeat(self.pos_emb_h, "h d-> t h w d", t=T, w=W),
+                repeat(self.pos_emb_w, "w d-> t h w d", t=T, h=H),
+            ],
+            dim=-1,
+        )
+        # Flatten the T,H,W dimensions
+        emb = rearrange(emb, "t h w d -> (t h w) d")
+
+        if training_type == "text_to_video":
+            bov_pe = torch.zeros((1, *emb.shape[1:]), device=emb.device, dtype=emb.dtype)
+            emb = torch.cat((bov_pe, emb), dim=0)
+        if self.pad_to_multiple_of is not None and emb.shape[0] % self.pad_to_multiple_of != 0:
+            pad_len = self.pad_to_multiple_of - emb.shape[0] % self.pad_to_multiple_of
+            emb = torch.cat((emb, torch.zeros((pad_len, *emb.shape[1:]), device=emb.device, dtype=emb.dtype)), dim=0)
+        seq_len, dim = emb.shape
+        emb = emb.reshape(1, seq_len, dim)
+        return emb

ar_module_mlp.py

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+    ):
+        """
+        Initializes the multilayer perceptron (MLP) module.
+
+        Args:
+            dim: The input and output dimensionality.
+            hidden_dim: The dimensionality of the hidden layer.
+        """
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Performs the forward pass of the MLP module.
+
+        Args:
+            x: The input tensor of shape (batch_size, dim).
+
+        Returns:
+            The output tensor of shape (batch_size, dim).
+        """
+        output = self.w2(F.silu(self.w1(x)) * self.w3(x))
+        return output

ar_module_mm_projector.py

@@ -0,0 +1,109 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Multimodal projector to connect vision encoder / tokenizer with the LLM."""
+
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+
+
+class DownSampleBlock(nn.Module):
+    """Downsample block."""
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        """
+        Performs the forward pass of the downsample block.
+
+        Args:
+            x (torch.Tensor): The input tensor from ViT's output of a sequence of embeddings.
+            Shape: (b, seq_len, c).
+
+        Returns:
+            torch.Tensor: The output tensor. Shape: (b, seq_len/4, c*4).
+        """
+        vit_embeds = x
+        # Get h and w as the sqrt of seq length. This assumes that the input is square-shaped.
+        h = w = int(vit_embeds.shape[1] ** 0.5)
+        b = vit_embeds.shape[0]
+        vit_embeds = vit_embeds.reshape(b, h, w, -1)
+        vit_embeds = self.flat_square(vit_embeds)
+        vit_embeds = vit_embeds.reshape(b, -1, vit_embeds.shape[-1])
+        return vit_embeds
+
+    def flat_square(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Performs spatial downsampling while increasing the number of channels.
+
+        Args:
+            x (torch.Tensor): The input tensor reshaped to a 2D grid.
+            Shape: (b, h, w, c)
+
+        Returns:
+            torch.Tensor: The output tensor after the spatial downsampling.
+            Shape: (b, h/2, w/2, c*4)
+        """
+        b, h, w, c = x.size()
+        # If w or h is odd, pad a column or a row of zeros.
+        if h % 2 == 1:
+            x = torch.concat([x, torch.zeros((b, 1, w, c), dtype=x.dtype).to(x.device)], dim=1).contiguous()
+            b, h, w, c = x.size()
+        if w % 2 == 1:
+            x = torch.concat([x, torch.zeros((b, h, 1, c), dtype=x.dtype).to(x.device)], dim=2).contiguous()
+            b, h, w, c = x.size()
+        # 2x spatial downsampling, 4x channel increasing.
+        x = x.view(b, h, int(w / 2), int(c * 2))
+        x = x.permute(0, 2, 1, 3).contiguous()
+        x = x.view(b, int(h / 2), int(w / 2), int(c * 4))
+        x = x.permute(0, 2, 1, 3).contiguous()
+        return x
+
+
+class MultimodalProjector(nn.Module):
+    """Multimodal projector."""
+
+    def __init__(
+        self,
+        mm_projector_type: str,
+        in_dim: int,
+        out_dim: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        super().__init__()
+        if out_dim is None:
+            out_dim = in_dim
+        if mm_projector_type == "identity":
+            self.projector = nn.Identity()
+        elif mm_projector_type == "linear":
+            self.projector = nn.Linear(in_dim, out_dim)
+        elif mm_projector_type == "mlp":
+            self.projector = nn.Sequential(nn.Linear(in_dim, out_dim), nn.GELU(), nn.Linear(out_dim, out_dim))
+        elif mm_projector_type == "mlp_downsample":
+            self.projector = nn.Sequential(
+                DownSampleBlock(),
+                nn.LayerNorm(in_dim * 4),
+                nn.Linear(in_dim * 4, out_dim),
+                nn.GELU(),
+                nn.Linear(out_dim, out_dim),
+            )
+        else:
+            raise ValueError(f"Unknown projector type: {mm_projector_type}")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.projector(x)

ar_module_normalization.py

@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+
+
+def create_norm(norm_type: str, dim: int, eps: float = 1e-6):
+    """
+    Creates the specified normalization layer based on the norm_type.
+    Adopted from TorchTriton: https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/norms.py
+
+    Args:
+        norm_type (str): The type of normalization layer to create.
+            Supported types: 1. rmsnorm 2. fused_rmsnorm 3. layernorm 4. np_layernorm
+        dim (int): The dimension of the normalization layer.
+        eps (float, optional): The epsilon value for numerical stability. Defaults to 1e-6.
+
+    Returns:
+        The created normalization layer.
+
+    Raises:
+        NotImplementedError: If an unknown norm_type is provided.
+    """
+    norm_type = norm_type.lower()  # Normalize to lowercase
+
+    if norm_type == "layernorm":
+        return nn.LayerNorm(dim, eps=eps, bias=False)
+    elif norm_type == "np_layernorm":
+        return nn.LayerNorm(dim, eps=eps, elementwise_affine=False, bias=False)
+    elif norm_type == "rmsnorm":
+        return RMSNorm(dim, eps=eps, compile=False)
+    elif norm_type == "compiled_rmsnorm":
+        return RMSNorm(dim, eps=eps, compile=True)
+    elif norm_type == "fused_rmsnorm":
+        raise NotImplementedError("Fused RMSNorm is not supported yet.")
+    else:
+        raise NotImplementedError(f"Unknown norm_type: '{norm_type}'")
+
+
+class RMSNorm(nn.Module):
+    """
+    Initialize the RMSNorm normalization layer.
+    Reference implementation: https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/norms.py
+
+    Args:
+        dim (int): The dimension of the input tensor.
+        eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+        compile (bool, optional): Whether to compile the forward function. Default is False.
+
+    Attributes:
+        eps (float): A small value added to the denominator for numerical stability.
+        weight (nn.Parameter): Learnable scaling parameter.
+
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6, compile: bool = False):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.rmsnorm_fn = torch.compile(self.compute_rmsnorm, fullgraph=True) if compile else self.compute_rmsnorm
+
+    @staticmethod
+    def compute_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float):
+        def _norm(x, eps):
+            # Computes the root-mean-square norm of the input tensor.
+            return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+
+        output = _norm(x.float(), eps).type_as(x)
+        return output * weight
+
+    def forward(self, x: torch.Tensor):
+        return self.rmsnorm_fn(x, self.weight, self.eps)
+
+    def reset_parameters(self):
+        torch.nn.init.ones_(self.weight)

ar_network_transformer.py

@@ -0,0 +1,461 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, Optional
+
+import torch
+import torch.nn as nn
+from torch.nn.modules.module import _IncompatibleKeys
+
+from .ar_module_attention import Attention
+from .ar_module_embedding import (
+    RotaryPositionEmbeddingPytorchV1,
+    RotaryPositionEmbeddingPytorchV2,
+    SinCosPosEmbAxisTE,
+)
+from .ar_module_mlp import MLP
+from .ar_module_normalization import create_norm
+from .ar_utils_checkpoint import process_state_dict, substrings_to_ignore
+from .ar_utils_misc import maybe_convert_to_namespace
+from .log import log
+
+
+class TransformerBlock(nn.Module):
+    """
+    A single transformer block consisting of an attention layer and a feed-forward layer.
+    """
+
+    def __init__(self, layer_id: int, args=None):
+        """
+        Initializes the TransformerBlock module.
+
+        Args:
+            layer_id: The ID of the transformer block.
+            args: The model arguments containing hyperparameters.
+        """
+        super().__init__()
+        args = maybe_convert_to_namespace(args)
+        attention_args = {
+            "n_heads": args["n_heads"],
+            "n_kv_heads": args["n_kv_heads"],
+            "dim": args["dim"],
+            "context_dim": None,
+            "max_batch_size": args["max_batch_size"],
+            "max_seq_len": args["max_seq_len"],
+            "use_qk_normalization": args["use_qk_normalization"],
+            "causal_mask": args["causal_mask"],
+            "head_dim": args["head_dim"],
+            "fuse_qkv": getattr(args, "fuse_qkv", False),
+            "precision": getattr(args, "precision", "bfloat16"),
+            "attn_type": getattr(args, "attn_type", "self"),
+        }
+        self.attention = Attention(**attention_args)
+
+        self.has_cross_attention = False
+        self.cross_attention, self.cross_attention_norm = None, None
+
+        if args["insert_cross_attn"] and layer_id % args["insert_cross_attn_every_k_layers"] == 0:
+            self.has_cross_attention = True
+            cross_attention_args = attention_args.copy()
+            cross_attention_args.update({"context_dim": args["context_dim"], "fuse_qkv": False, "attn_type": "cross"})
+            self.cross_attention = Attention(**cross_attention_args)
+            self.cross_attention_norm = create_norm(args["norm_type"], dim=args["dim"], eps=args["norm_eps"])
+
+        self.feed_forward = MLP(
+            dim=args["dim"],
+            hidden_dim=args["ffn_hidden_size"],
+        )
+        self.layer_id = layer_id
+        self.attention_norm = create_norm(args["norm_type"], dim=args["dim"], eps=args["norm_eps"])
+        self.ffn_norm = create_norm(args["norm_type"], dim=args["dim"], eps=args["norm_eps"])
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: RotaryPositionEmbeddingPytorchV2,
+        input_pos: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        context: Optional[torch.Tensor] = None,
+        context_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Performs the forward pass of the TransformerBlock module.
+
+        Args:
+            x: The input tensor.
+            input_pos: The position of the current sequence. Used in inference (with KV cache) only.
+            freqs_cis: The precomputed frequency values for rotary position embeddings.
+            mask: The attention mask tensor.
+            context (Optional[torch.Tensor]): The context tensor added via cross-attn.
+            context_mask (Optional[torch.Tensor]): The context cross-attn mask tensor.
+
+        Returns:
+            The output tensor after applying the transformer block.
+        """
+        # Apply attention and residual connection
+        h = x + self.attention(self.attention_norm(x), rope=rope, input_pos=input_pos, mask=mask)
+
+        # If insert cross-attention, apply CA and residual connection
+        if self.has_cross_attention:
+            h = h + self.cross_attention(
+                self.cross_attention_norm(h), rope=rope, input_pos=input_pos, mask=context_mask, context=context
+            )
+
+        # Apply feed-forward network and residual connection
+        out = h + self.feed_forward(self.ffn_norm(h))
+        return out
+
+    def init_weights(self):
+        """
+        Initializes the weights of the transformer block.
+        """
+        for norm in (self.attention_norm, self.ffn_norm):
+            norm.reset_parameters()
+        self.attention.init_weights(self.weight_init_std)
+        self.feed_forward.init_weights(self.weight_init_std)
+
+        if self.has_cross_attention:
+            self.cross_attention_norm.reset_parameters()
+            self.cross_attention.init_weights(self.weight_init_std)
+            # zero-init the final output layer of cross-attention
+            # nn.init.zeros_(self.cross_attention.wo.weight)
+
+
+class Transformer(nn.Module):
+    """
+    The Transformer network consisting of transformer blocks.
+    """
+
+    def __init__(self, params, tokenizer_config=None, init_weights: bool = True):
+        """
+        Initializes the Transformer module.
+
+        Args:
+            params: The model parameters containing hyperparameters.
+            tokenizer_config: The model tokenizer configuration.
+            init_weights (bool): Whether to initialize the weights of the transformer following
+                TorchTitan's Llama3 initialization scheme.
+        """
+        super().__init__()
+        # Check if self.params is an OmegaConf DictConfig instance
+        self.params = maybe_convert_to_namespace(params)
+        self.vocab_size = params["vocab_size"]
+        self.n_layers = params["n_layers"]
+        self.precision = getattr(torch, params["precision"])
+        self.tokenizer_config = tokenizer_config
+        self.num_video_frames = params["num_video_frames"]
+
+        # Token embeddings
+        self.tok_embeddings = self._create_token_embeddings()
+        self.rope_config = self._create_rope_config()
+
+        # Transformer layers
+        self.layers = nn.ModuleList(
+            [TransformerBlock(layer_id, self.params).to(self.precision) for layer_id in range(self.n_layers)]
+        )
+
+        # Final layer normalization
+        self.norm = create_norm(self.params["norm_type"], dim=self.params["dim"], eps=self.params["norm_eps"]).to(
+            self.precision
+        )
+        if self.params["pytorch_rope_version"] == "v1":
+            self.rope = RotaryPositionEmbeddingPytorchV1(**self.rope_config)
+        elif self.params["pytorch_rope_version"] == "v2":
+            # Rotary position embeddings
+            training_type = self.tokenizer_config.training_type if self.tokenizer_config is not None else None
+            self.rope = RotaryPositionEmbeddingPytorchV2(
+                seq_len=self.params["max_seq_len"], training_type=training_type, **self.rope_config
+            )
+        else:
+            raise ValueError(f"Invalid PyTorch RoPE version: {self.params['pytorch_rope_version']}")
+        # Causal mask
+        self.causal_mask = torch.tril(
+            torch.ones(self.params["max_seq_len"], self.params["max_seq_len"], dtype=torch.bool)
+        ).cuda()
+
+        # Output projection
+        self.output = self._create_output_projection()
+
+        # Freeze network parameters for finetuning w/ cross-attention
+        self.has_cross_attention = getattr(params, "insert_cross_attn", False)
+
+        # Absolute position embeddings
+        if self.params["apply_abs_pos_emb"]:
+            self.pos_emb_config = self._create_abs_pos_emb_config()
+            self.pos_emb, self.abs_pos_emb = self._initialize_abs_pos_emb()
+
+    def _create_rope_config(self) -> Dict:
+        shape_map = {
+            "3D": self.params["video_latent_shape"],
+            "1D": None,
+        }
+        latent_shape = shape_map.get(self.params["rope_dim"], None)
+        head_dim = self.params["head_dim"]
+        if head_dim is None:
+            head_dim = self.params["dim"] // self.params["n_heads"]
+        return {
+            "dim": head_dim,
+            "max_position_embeddings": self.params["max_seq_len"],
+            "original_max_position_embeddings": self.params["original_seq_len"],
+            "rope_theta": self.params["rope_theta"],
+            "apply_yarn": self.params["apply_yarn"],
+            "scale": self.params["yarn_scale"],
+            "beta_fast": self.params["yarn_beta_fast"],
+            "beta_slow": self.params["yarn_beta_slow"],
+            "rope_dim": self.params["rope_dim"],
+            "latent_shape": latent_shape,
+            "original_latent_shape": self.params["original_latent_shape"],
+            "pad_to_multiple_of": self.params["pad_to_multiple_of"],
+        }
+
+    def _create_abs_pos_emb_config(self):
+        shape_map = {
+            "3D": self.params["video_latent_shape"],
+            "1D": None,
+        }
+        latent_shape = shape_map.get(self.params["rope_dim"], None)
+        return {
+            "dim": self.params["dim"],
+            "latent_shape": latent_shape,
+            "pad_to_multiple_of": self.params["pad_to_multiple_of"],
+        }
+
+    def _create_token_embeddings(self, vocab_size: int = None):
+        """
+        Create token embeddings.
+
+        Returns:
+            nn.Module: Token embeddings module.
+        """
+        if vocab_size is None:
+            vocab_size = self.params["vocab_size"]
+        return nn.Embedding(vocab_size, self.params["dim"]).to(self.precision)
+
+    def _create_output_projection(self, vocab_size: int = None):
+        """
+        Create the output projection layer.
+
+        Args:
+            vocab_size (int): Vocabulary size (to override the default vocab size).
+        Returns:
+            LinearTE: Output projection layer.
+        """
+        if vocab_size is None:
+            vocab_size = self.params["vocab_size"]
+        return nn.Linear(self.params["dim"], vocab_size, bias=False).to(self.precision)
+
+    def _initialize_abs_pos_emb(self):
+        pos_emb = SinCosPosEmbAxisTE(**self.pos_emb_config)
+        training_type = self.tokenizer_config.training_type if self.tokenizer_config is not None else None
+        abs_pos_emb = pos_emb.forward(training_type=training_type)
+        return pos_emb, abs_pos_emb
+
+    def forward(
+        self,
+        tokens: Optional[torch.Tensor] = None,
+        input_pos: Optional[torch.Tensor] = None,
+        token_embeddings: Optional[torch.Tensor] = None,
+        context: Optional[torch.Tensor] = None,
+        context_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Performs the forward pass of the Transformer module.
+
+        Args:
+            tokens (torch.Tensor, optional): The input tensor of token IDs.
+            input_pos (Optional[torch.Tensor]): The position of the current sequence. Used in inference with KV cache.
+            token_embeddings (torch.Tensor, optional): Precomputed token embeddings. If provided, tokens should be None.
+            context (Optional[torch.Tensor]): The context tensor added via cross-attn.
+            context_mask (Optional[torch.Tensor]): The context cross-attn mask tensor.
+        Returns:
+            The output tensor after applying the transformer layers.
+        """
+        # Token embeddings
+        assert (
+            tokens is None or token_embeddings is None
+        ), "Either tokens or token_embeddings should be provided, not both."
+
+        if token_embeddings is None:
+            seq_len = tokens.shape[1]
+            h = self.tok_embeddings(tokens)
+        else:
+            seq_len = token_embeddings.shape[1]
+            h = token_embeddings
+
+        # Create attention mask
+        mask = self._create_attention_mask(input_pos=input_pos)
+
+        # Prepare layer arguments
+        layer_kwargs = self._prepare_layer_kwargs(
+            input_pos=input_pos,
+            mask=mask,
+            context=context,
+            context_mask=context_mask,
+        )
+
+        # Apply transformer layers
+        for layer in self.layers:
+            if self.params["apply_abs_pos_emb"]:
+                h = self.apply_abs_pos_emb(h, input_pos=input_pos)
+            h = layer(h, **layer_kwargs)
+
+        # Apply final layer normalization
+        h = self.norm(h)
+
+        # Output linear projection
+        output = self.output(h)
+        return output
+
+    def _create_attention_mask(self, input_pos: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+        """
+        Creates an attention mask for the transformer layers.
+
+        Args:
+            input_pos[torch.Tensor]: The position of input sequence (used for inference only).
+
+        Returns:
+            Optional[torch.Tensor]: The attention mask, or None for causal mask.
+        """
+
+        assert input_pos is not None, "input_pos must be provided for inference"
+        mask = self.causal_mask[input_pos]
+        return mask
+
+    def _prepare_layer_kwargs(
+        self,
+        input_pos: Optional[torch.Tensor],
+        mask: Optional[torch.Tensor],
+        context: Optional[torch.Tensor],
+        context_mask: Optional[torch.Tensor],
+    ) -> Dict[str, Any]:
+        """
+        Prepares the keyword arguments for transformer layers.
+
+        Args:
+            input_pos (Optional[torch.Tensor]): The position of the current sequence.
+            mask (Optional[torch.Tensor]): The attention mask.
+            context (Optional[torch.Tensor]): The context tensor added via cross-attn.
+            context_mask (Optional[torch.Tensor]): The context cross-attn mask tensor.
+
+        Returns:
+            Dict[str, Any]: A dictionary of keyword arguments for the transformer layers.
+        """
+        if context is not None:
+            context = context.to(self.precision)
+
+        if isinstance(mask, torch.Tensor) and mask.ndim == 2:
+            mask = mask[None, None, :, :]
+        if isinstance(context_mask, torch.Tensor) and context_mask.ndim == 2:
+            context_mask = context_mask[None, None, :, :]
+
+        layer_kwargs = {
+            "mask": mask,
+            "context": context,
+            "context_mask": context_mask,
+        }
+
+        layer_kwargs["input_pos"] = input_pos
+        layer_kwargs["rope"] = self.rope
+
+        return layer_kwargs
+
+    def apply_abs_pos_emb(self, x: torch.Tensor, input_pos: int = None) -> torch.Tensor:
+        """
+        Applies the absolute position embeddings to the input tensor.
+        """
+        abs_pos_emb = self.abs_pos_emb
+        abs_pos_emb = abs_pos_emb[:, input_pos, :] if input_pos is not None else abs_pos_emb
+        return x + abs_pos_emb
+
+    @torch.no_grad()
+    def expand_vocab(
+        self, new_vocab_size: int, init_method: str = "gaussian", multiple_of=64, expand_output_layer=True
+    ):
+        """
+        Expands the vocabulary of the model to the new size.
+
+        Args:
+            new_vocab_size (int): The new vocabulary size.
+            init_method (str): The initialization method for new embeddings.
+                               Can be "zero" or "gaussian". Default is "gaussian".
+            multiple_of (int): The new vocabulary size must be a multiple of this value. Defaults to 64 to fully
+                leverage the power of NVIDIA TensorCore (source 1: https://x.com/karpathy/status/1621578354024677377,
+                source 2: https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc)
+            expand_output_layer (bool): Whether to also expand the output layer. Defaults to True.
+
+        Returns:
+            None
+        """
+        if new_vocab_size <= self.vocab_size:
+            raise ValueError(
+                f"New vocabulary size ({new_vocab_size}) must be " f"larger than current size ({self.vocab_size})"
+            )
+        if new_vocab_size % multiple_of != 0:
+            log.debug(f"New vocabulary size must be a multiple of {multiple_of}. Obtained {new_vocab_size}.")
+            new_vocab_size = (new_vocab_size // multiple_of + 1) * multiple_of
+            log.debug(f"Rounded vocabulary size to {new_vocab_size}.")
+        # Resize token embeddings
+        old_embeddings = self.tok_embeddings
+        tensor_kwargs = {"device": old_embeddings.weight.device, "dtype": old_embeddings.weight.dtype}
+        self.tok_embeddings = self._create_token_embeddings(vocab_size=new_vocab_size).to(**tensor_kwargs)
+        # Initialize new embeddings
+        if init_method not in ["zero", "gaussian"]:
+            raise ValueError(f"Unknown initialization method: {init_method}")
+        # The default initialization of nn.Embedding is Gaussian, so we don't need to do anything
+        # if init_method == "gaussian". Only if init_method == "zero", we need to zero out the new embeddings.
+        if init_method == "zero":
+            self.tok_embeddings.weight.data[self.vocab_size :].zero_()
+
+        # Copy old embeddings
+        log.debug(
+            f"old_embeddings: {old_embeddings.weight.data.shape}, new_embeddings: {self.tok_embeddings.weight.data.shape}, vocab_size: {self.vocab_size}"
+        )
+        self.tok_embeddings.weight.data[: self.vocab_size] = old_embeddings.weight.data
+        # Resize output layer
+        old_output = self.output
+        self.output = self._create_output_projection(vocab_size=new_vocab_size if expand_output_layer else None)
+
+        # Initialize new output weights
+        self.output.weight.data[self.vocab_size :].zero_()
+        # Copy old output weights
+        self.output.weight.data[: self.vocab_size] = old_output.weight.data
+
+        # Update vocab size
+        self.vocab_size = new_vocab_size
+        log.debug(f"Expanded vocabulary size to {new_vocab_size}")
+
+    def state_dict(self, *args, **kwargs):
+        """
+        Process the state dict (e.g., remove "_extra_state" keys imposed by TransformerEngine for FP8).
+        """
+        state_dict = super().state_dict(*args, **kwargs)
+        return process_state_dict(state_dict)
+
+    def load_state_dict(self, state_dict: Dict[str, Any], strict: bool = True, assign: bool = False):
+        """
+        Ignore the missing keys with substrings matching `substring_to_ignore` (e.g., "_extra_state" keys imposed by
+        TransformerEngine for FP8).
+        """
+        state_dict = process_state_dict(state_dict)
+        missing_keys, unexpected_keys = super().load_state_dict(state_dict, strict=False, assign=assign)
+        if strict:
+            actual_missing_keys = []
+            for key in missing_keys:
+                if not any(substring in key for substring in substrings_to_ignore):
+                    actual_missing_keys.append(key)
+            if len(actual_missing_keys) > 0 or len(unexpected_keys) > 0:
+                raise ValueError(f"Missing keys: {actual_missing_keys}\n\nUnexpected keys: {unexpected_keys}")
+            missing_keys = actual_missing_keys
+        return _IncompatibleKeys(missing_keys, unexpected_keys)

ar_network_vit.py

@@ -0,0 +1,410 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+This module implements a Vision Transformer (ViT) with 2D Rotary Position Embeddings,
+designed for processing image inputs in vision-language models.
+
+This module follows Mistral's vision encoder implementation (for their Pistral-12B VLM):
+https://github.com/mistralai/mistral-inference/blob/main/src/mistral_inference/vision_encoder.py
+"""
+from functools import partial
+from typing import Any, Callable, Mapping, Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from .ar_module_normalization import create_norm
+from .ar_network_transformer import TransformerBlock
+from .log import log
+
+
+def get_vit_config(model_name: str) -> Mapping[str, Any]:
+    """
+    Get the ViT configuration for a given model name.
+    """
+    if model_name == "pixtral-12b-vit":
+        # The 400M ViT of Pixtral 12B VLM
+        return dict(
+            dim=1024,
+            num_channels=3,
+            image_size=1024,
+            patch_size=16,
+            rope_theta=10000,
+            ffn_hidden_size=4096,
+            n_layers=24,
+            n_heads=16,
+            n_kv_heads=16,
+            norm_type="rmsnorm",
+            norm_eps=1e-5,
+            image_token_id=10,
+        )
+    else:
+        raise ValueError(f"Unknown model name: {model_name}")
+
+
+def precompute_freqs_cis_2d(
+    dim: int,
+    height: int,
+    width: int,
+    theta: float,
+) -> torch.Tensor:
+    """
+    Precompute 2D complex tensor for rotary position embedding.
+
+    This function generates a 2D complex tensor used for rotary position embeddings,
+    which helps the model understand spatial relationships in the input image.
+
+    Args:
+        dim (int): Dimension of the model (typically the hidden size divided by number of heads).
+        height (int): Height of the image in patches.
+        width (int): Width of the image in patches.
+        theta (float): Base value for the angle calculation, controls the frequency range.
+
+    Returns:
+        torch.Tensor: 2D complex tensor of shape (height, width, dim // 2).
+    """
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+
+    h = torch.arange(height, device=freqs.device)
+    w = torch.arange(width, device=freqs.device)
+
+    freqs_h = torch.outer(h, freqs[::2]).float()
+    freqs_w = torch.outer(w, freqs[1::2]).float()
+    freqs_2d = torch.cat(
+        [
+            freqs_h[:, None, :].repeat(1, width, 1),
+            freqs_w[None, :, :].repeat(height, 1, 1),
+        ],
+        dim=-1,
+    )
+    return torch.polar(torch.ones_like(freqs_2d), freqs_2d)
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    """
+    Reshape frequency tensor for broadcasting with input tensor.
+
+    This function ensures that the frequency tensor can be properly broadcast
+    with the input tensor during the rotary embedding process.
+
+    Args:
+        freqs_cis (torch.Tensor): Frequency tensor from precompute_freqs_cis_2d.
+        x (torch.Tensor): Input tensor to be embedded.
+
+    Returns:
+        torch.Tensor: Reshaped frequency tensor ready for broadcasting.
+    """
+    ndim = x.ndim
+    assert 0 <= 1 < ndim, f"ndim is {ndim} but index is {1}"
+    assert freqs_cis.shape == (
+        x.shape[1],
+        x.shape[-1],
+    ), f"freqs_cis shape is {freqs_cis.shape} but x shape is {x.shape}"
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    *args,
+    freqs_cis: torch.Tensor,
+    **kwargs,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary positional embeddings to input tensors.
+
+    This function applies the rotary positional embeddings to the query and key tensors,
+    which helps the model understand spatial relationships in the input.
+
+    Args:
+        xq (torch.Tensor): Query tensor.
+        xk (torch.Tensor): Key tensor.
+        freqs_cis (torch.Tensor): Precomputed frequencies from precompute_freqs_cis_2d.
+        *args: Variable length argument list (unused).
+        **kwargs: Arbitrary keyword arguments (unused).
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Rotated query and key tensors.
+    """
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class VisionTransformer(nn.Module):
+    """
+    Vision Transformer model for image processing.
+
+    This class implements a Vision Transformer that processes images using a patch-based approach
+    and applies transformer layers with rotary position embeddings.
+
+    Args:
+        dim (int): Dimension of the model (hidden size).
+        num_channels (int): Number of input image channels (e.g., 3 for RGB).
+        patch_size (int): Size of each image patch (e.g., 16x16 pixels).
+        n_layers (int): Number of transformer layers.
+        n_heads (int): Number of attention heads.
+        ffn_hidden_size (int): Hidden size of the feed-forward network in transformer blocks.
+        norm_type (str): Type of normalization to use (e.g., "rmsnorm").
+        norm_eps (float): Epsilon value for normalization layers.
+        image_size (int): Size of the input image (assumed square).
+        rope_theta (float): Base value for rotary position embedding calculation.
+        attention_dropout (float): Dropout rate for attention layers.
+        hidden_dropout (float): Dropout rate for hidden layers.
+        image_token_id (int): Token ID for the image token (if present).
+    """
+
+    def __init__(
+        self,
+        dim: int = 1024,
+        num_channels: int = 3,
+        patch_size: int = 16,
+        n_layers: int = 24,
+        n_heads: int = 16,
+        n_kv_heads: int = None,
+        ffn_hidden_size: int = 4096,
+        norm_type: str = "rmsnorm",
+        norm_eps: float = 1e-5,
+        image_size: int = 1024,
+        rope_theta: float = 1000000.0,
+        image_token_id: int = None,
+    ):
+        super().__init__()
+        self.patch_conv = nn.Conv2d(
+            in_channels=num_channels,
+            out_channels=dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=False,
+        )
+        self.ln_pre = create_norm(norm_type=norm_type, dim=dim, eps=norm_eps)
+        if n_kv_heads is None:
+            n_kv_heads = n_heads
+        layer_args = dict(
+            n_layers=n_layers,
+            n_heads=n_heads,
+            n_kv_heads=n_kv_heads,
+            dim=dim,
+            use_qk_normalization=False,
+            max_seq_len=None,
+            max_batch_size=None,
+            ffn_hidden_size=ffn_hidden_size,
+            norm_type=norm_type,
+            norm_eps=norm_eps,
+            causal_mask=False,  # Full attention in ViT
+            head_dim=None,
+            insert_cross_attn=False,
+            attn_type="full",
+        )
+
+        self.transformer = VisionTransformerBlocks(n_layers=n_layers, args=layer_args)
+
+        head_dim = dim // n_heads
+        assert head_dim % 2 == 0, "ROPE requires even head_dim"
+
+        self.dim = dim
+        self.n_heads = n_heads
+        self.max_patches_per_side = image_size // patch_size
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.rope_theta = rope_theta
+        self._freqs_cis: Optional[torch.Tensor] = None
+        self.image_token_id = image_token_id
+
+        num_params = self.get_num_params()
+        log.debug(f"Number  of model parameters: {round(num_params / 1e6, 3)}M")
+
+    @classmethod
+    def build(
+        cls,
+        config: Mapping[str, Any],
+    ) -> "VisionTransformer":
+        """
+        Create a Vision Transformer from a configuration dictionary.
+
+        This class method creates a Vision Transformer from a configuration dictionary,
+        which is typically loaded from a JSON file or other configuration source.
+
+        Args:
+            config (Mapping[str, Any]): Configuration dictionary for the Vision Transformer.
+
+        Returns:
+            VisionTransformer: Vision Transformer model instance.
+        """
+        necessary_keys = ["dim", "num_channels", "patch_size", "n_layers", "n_heads", "ffn_hidden_size", "rope_theta"]
+        missing_keys = [k for k in necessary_keys if k not in config]
+        assert len(missing_keys) == 0, f"Missing keys in config: {missing_keys}"
+        return cls(
+            **config,
+        )
+
+    def expand_in_channels(self, new_in_channels: int):
+        """
+        Expand the input channels of the patch convolution layer.
+        This is useful when the input is non-standard, e.g. a 4-channel image with the last channel as the alpha channel.
+        Note that you should only call this method after the weight is loaded.
+        """
+        assert (
+            new_in_channels > self.patch_conv.in_channels
+        ), "Cannot expand the input channels of the patch convolution layer to be less than the original number of channels."
+        log.debug(
+            f"Vision encoder in_channels is {self.patch_conv.in_channels}. But you have specified to be {new_in_channels}. We will change it to {new_in_channels} channels with {new_in_channels - self.patch_conv.in_channels} channels of 0s."
+        )
+        new_conv = nn.Conv2d(
+            in_channels=new_in_channels,
+            out_channels=self.patch_conv.out_channels,
+            kernel_size=self.patch_conv.kernel_size,
+            stride=self.patch_conv.stride,
+            bias=False,
+        )
+        new_conv.weight.data[:, : self.patch_conv.in_channels].copy_(self.patch_conv.weight.data)
+        new_conv.weight.data[
+            :, self.patch_conv.in_channels :
+        ].zero_()  # zeroize, such that initially it has no effect to output
+        self.patch_conv = new_conv
+
+    @property
+    def device(self) -> torch.device:
+        """Get the device of the model."""
+        return next(self.parameters()).device
+
+    @property
+    def freqs_cis(self) -> torch.Tensor:
+        """
+        Get or compute the frequency tensor for rotary position embedding.
+
+        This property lazily initializes and caches the frequency tensor used for
+        rotary position embeddings, ensuring it's on the correct device.
+
+        Returns:
+            torch.Tensor: The frequency tensor for rotary position embeddings.
+        """
+        if self._freqs_cis is None:
+            self._freqs_cis = precompute_freqs_cis_2d(
+                dim=self.dim // self.n_heads,
+                height=self.max_patches_per_side,
+                width=self.max_patches_per_side,
+                theta=self.rope_theta,
+            )
+
+        if self._freqs_cis.device != self.device:
+            self._freqs_cis = self._freqs_cis.to(device=self.device)
+
+        return self._freqs_cis
+
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Forward pass of the Vision Transformer.
+
+        This method processes the input image through the Vision Transformer,
+        including patch embedding, position embedding, and transformer layers.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (B, C, H, W), where B is batch size,
+                              C is number of channels, and H, W are height and width.
+
+        Returns:
+            torch.Tensor: Output features of shape (B, N, D), where N is the number of patches
+                          and D is the embedding dimension.
+        """
+
+        patch_embeds = self.patch_conv(x)  # (B, D, Hp, Wp)
+        _, _, Hp, Wp = patch_embeds.shape  # Patch embeds dim
+        patch_embeds = patch_embeds.flatten(2)  # (B, D, Hp*Wp)
+        patch_embeds = patch_embeds.transpose(1, 2)  # (B, Hp*Wp, D)
+        patch_embeds = self.ln_pre(patch_embeds)  # (B, Hp*Wp, D)
+        positions = torch.stack(
+            torch.meshgrid(
+                torch.arange(Hp),
+                torch.arange(Wp),
+                indexing="ij",
+            ),
+            dim=-1,
+        ).reshape(-1, 2)
+
+        freqs_cis = self.freqs_cis[positions[:, 0], positions[:, 1]]
+        rope = partial(apply_rotary_emb, freqs_cis=freqs_cis)
+        out = self.transformer(patch_embeds, rope=rope)
+
+        return out
+
+    def get_num_params(
+        self,
+    ) -> int:
+        """
+        Return the number of parameters in the model.
+        """
+        n_params = sum(p.numel() for p in self.parameters())
+        return n_params
+
+
+class VisionTransformerBlocks(nn.Module):
+    """
+    Vision Transformer Blocks.
+
+    This class implements a stack of Transformer blocks used in the Vision Transformer.
+
+    Args:
+        n_layers (int): Number of transformer layers.
+        args (Mapping[str, Any]): Arguments for each transformer block, including dimensions,
+    """
+
+    def __init__(
+        self,
+        n_layers: int,
+        args: Mapping[str, Any],
+    ):
+        super().__init__()
+        self.layers = torch.nn.ModuleList()
+
+        for layer_id in range(n_layers):
+            self.layers.append(
+                TransformerBlock(
+                    layer_id=layer_id,
+                    args=args,
+                )
+            )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rope: Callable,
+    ) -> torch.Tensor:
+        """
+        Forward pass through the Vision Transformer Blocks.
+
+        This method applies a series of Transformer blocks to the input tensor,
+        using the provided rotary position embedding function.
+
+        Args:
+            x (torch.Tensor): Input tensor of shape (B, N, D), where B is batch size,
+                              N is the number of patches, and D is the embedding dimension.
+            rope (Callable): Rotary position embedding function to be applied in each layer.
+
+        Returns:
+            torch.Tensor: Output tensor after passing through all transformer layers,
+                          with the same shape as the input.
+        """
+        for layer in self.layers:
+            x = layer(x, input_pos=None, mask=None, rope=rope)
+        return x

ar_tokenizer_discrete_video.py

@@ -0,0 +1,360 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+import torch
+from einops import rearrange
+
+from .ar_tokenizer_quantizers import FSQuantizer
+
+# Make sure jit model output consistenly during consecutive calls
+# Check here: https://github.com/pytorch/pytorch/issues/74534
+torch._C._jit_set_texpr_fuser_enabled(False)
+
+
+def load_jit_model(jit_filepath: str = None, device: str = "cuda") -> torch.jit.ScriptModule:
+    """Loads a torch.jit.ScriptModule from a filepath.
+
+    Args:
+        jit_filepath: The filepath to the JIT-compiled model.
+        device: The device to load the model onto, default=cuda.
+    Returns:
+        The JIT compiled model loaded to device and on eval mode.
+    """
+    # Make sure jit model output consistenly during consecutive calls
+    # Check here: https://github.com/pytorch/pytorch/issues/74534
+    torch._C._jit_set_texpr_fuser_enabled(False)
+
+    model = torch.jit.load(jit_filepath)
+    return model.eval().to(device)
+
+
+class BaseDiscreteVideoFSQTokenizer(torch.nn.Module):
+    """
+    A base class for  Discrete Video FSQ Tokenizer that handles data type conversions, and normalization
+    using provided mean and standard deviation values for latent space representation.
+    Derived classes should load pre-trained encoder and decoder components into a encoder and decoder attributes.
+
+    Attributes:
+        encoder (Module | Callable): Encoder loaded from storage.
+        decoder (Module | Callable): Decoder loaded from storage.
+        dtype (dtype): Data type for model tensors, determined by whether bf16 is enabled.
+
+    Args:
+        name (str): Name of the model, used for differentiating cache file paths.
+        latent_ch (int, optional): Number of latent channels (default is 6).
+        is_bf16 (bool, optional): Flag to use Brain Floating Point 16-bit data type (default is True).
+        pixel_chunk_duration (int): The duration (in number of frames) of each chunk of video data at the pixel level.
+        latent_chunk_duration (int): The duration (in number of frames) of each chunk at the latent representation level.
+        max_enc_batch_size (int): The maximum batch size to process in one go to avoid memory overflow.
+        level (list[int]): The level defined in FSQ quantizer.
+        compression_ratio (list[int]): The compression factor for (T, H, W).
+    """
+
+    def __init__(
+        self,
+        name: str,
+        latent_ch: int = 6,
+        is_bf16: bool = True,
+        pixel_chunk_duration: int = 25,
+        latent_chunk_duration: int = 4,
+        max_enc_batch_size: int = 8,
+        max_dec_batch_size: int = 4,
+        levels: list[int] = [8, 8, 8, 5, 5, 5],
+        compression_ratio: list[int] = [8, 16, 16],
+    ):
+        super().__init__()
+        self.channel = latent_ch
+        self.name = name
+        dtype = torch.bfloat16 if is_bf16 else torch.float32
+        self.dtype = dtype
+        self.pixel_chunk_duration = pixel_chunk_duration
+        self.latent_chunk_duration = latent_chunk_duration
+        self.max_enc_batch_size = max_enc_batch_size
+        self.max_dec_batch_size = max_dec_batch_size
+        self.levels = levels
+        self.compress_ratio = compression_ratio
+        self.fsq_quantizer = FSQuantizer(levels)
+
+    @property
+    def latent_ch(self) -> int:
+        """
+        Returns the number of latent channels in the tokenizer.
+        """
+        return self.channel
+
+    @torch.no_grad()
+    def encode(self, state: torch.Tensor, pixel_chunk_duration: Optional[int] = None) -> torch.Tensor:
+        B, C, T, H, W = state.shape
+        if pixel_chunk_duration is None:
+            # Use the default pixel chunk duration and latent chunk duration
+            pixel_chunk_duration = self.pixel_chunk_duration
+            latent_chunk_duration = self.latent_chunk_duration
+        else:
+            # Update the latent chunk duration based on the given pixel chunk duration
+            latent_chunk_duration = 1 + (pixel_chunk_duration - 1) // self.compress_ratio[0]
+
+        assert (
+            T % pixel_chunk_duration == 0
+        ), f"Temporal dimension {T} is not divisible by chunk_length {pixel_chunk_duration}"
+        state = rearrange(state, "b c (n t) h w -> (b n) c t h w", t=pixel_chunk_duration)
+
+        # use max_enc_batch_size to avoid OOM
+        if state.shape[0] > self.max_enc_batch_size:
+            quantized_out_list = []
+            indices_list = []
+            for i in range(0, state.shape[0], self.max_enc_batch_size):
+                indices, quantized_out, _ = self.encoder(state[i : i + self.max_enc_batch_size].to(self.dtype))
+                quantized_out_list.append(quantized_out)
+                indices_list.append(indices)
+            quantized_out = torch.cat(quantized_out_list, dim=0)
+            indices = torch.cat(indices_list, dim=0)
+        else:
+            indices, quantized_out, _ = self.encoder(state.to(self.dtype))
+        assert quantized_out.shape[2] == latent_chunk_duration
+        return rearrange(quantized_out, "(b n) c t h w -> b c (n t) h w", b=B), rearrange(
+            indices, "(b n) t h w -> b (n t) h w", b=B
+        )
+
+    @torch.no_grad()
+    def decode(self, indices: torch.Tensor, pixel_chunk_duration: Optional[int] = None) -> torch.Tensor:
+        B, T, _, _ = indices.shape
+        if pixel_chunk_duration is None:
+            pixel_chunk_duration = self.pixel_chunk_duration
+            latent_chunk_duration = self.latent_chunk_duration
+        else:
+            latent_chunk_duration = 1 + (pixel_chunk_duration - 1) // self.compress_ratio[0]
+        assert (
+            T % latent_chunk_duration == 0
+        ), f"Temporal dimension {T} is not divisible by chunk_length {latent_chunk_duration}"
+        indices = rearrange(indices, "b (n t) h w -> (b n) t h w", t=latent_chunk_duration)
+
+        # use max_dec_batch_size to avoid OOM
+        if indices.shape[0] > self.max_dec_batch_size:
+            state = []
+            for i in range(0, indices.shape[0], self.max_dec_batch_size):
+                state.append(self.decoder(indices[i : i + self.max_dec_batch_size]))
+            state = torch.cat(state, dim=0)
+        else:
+            state = self.decoder(indices)
+
+        assert state.shape[2] == pixel_chunk_duration
+        return rearrange(state, "(b n) c t h w -> b c (n t) h w", b=B)
+
+    def reset_dtype(self, *args, **kwargs):
+        """
+        Resets the data type of the encoder and decoder to the model's default data type.
+
+        Args:
+            *args, **kwargs: Unused, present to allow flexibility in method calls.
+        """
+        del args, kwargs
+        self.decoder.to(self.dtype)
+        self.encoder.to(self.dtype)
+
+
+class DiscreteVideoFSQJITTokenizer(BaseDiscreteVideoFSQTokenizer):
+    """
+    A JIT compiled Discrete Video FSQ Tokenizer that loads pre-trained encoder
+    and decoder components from a remote store, handles data type conversions, and normalization
+    using provided mean and standard deviation values for latent space representation.
+
+    Attributes:
+        encoder (Module): The JIT compiled encoder loaded from storage.
+        decoder (Module): The JIT compiled decoder loaded from storage.
+        dtype (dtype): Data type for model tensors, determined by whether bf16 is enabled.
+
+    Args:
+        enc_fp (str): File path to the encoder's JIT file on the remote store.
+        dec_fp (str): File path to the decoder's JIT file on the remote store.
+        name (str): Name of the model, used for differentiating cache file paths.
+        latent_ch (int, optional): Number of latent channels (default is 6).
+        is_bf16 (bool, optional): Flag to use Brain Floating Point 16-bit data type (default is True).
+        pixel_chunk_duration (int): The duration (in number of frames) of each chunk of video data at the pixel level.
+        latent_chunk_duration (int): The duration (in number of frames) of each chunk at the latent representation level.
+        max_enc_batch_size (int): The maximum batch size to process in one go to avoid memory overflow.
+        level (list[int]): The level defined in FSQ quantizer.
+        compression_ratio (list[int]): The compression factor for (T, H, W).
+    """
+
+    def __init__(
+        self,
+        enc_fp: str,
+        dec_fp: str,
+        name: str,
+        latent_ch: int = 6,
+        is_bf16: bool = True,
+        pixel_chunk_duration: int = 25,
+        latent_chunk_duration: int = 4,
+        max_enc_batch_size: int = 8,
+        max_dec_batch_size: int = 4,
+        levels: list[int] = [8, 8, 8, 5, 5, 5],
+        compression_ratio: list[int] = [8, 16, 16],
+    ):
+        super().__init__(
+            name,
+            latent_ch,
+            is_bf16,
+            pixel_chunk_duration,
+            latent_chunk_duration,
+            max_enc_batch_size,
+            max_dec_batch_size,
+            levels,
+            compression_ratio,
+        )
+
+        self.load_encoder(enc_fp)
+        self.load_decoder(dec_fp)
+
+    def load_encoder(self, enc_fp: str) -> None:
+        """
+        Load the encoder from the remote store.
+
+        Args:
+        - enc_fp (str): File path to the encoder's JIT file on the remote store.
+        """
+        self.encoder = load_jit_model(enc_fp, device="cuda")
+        self.encoder.eval()
+        for param in self.encoder.parameters():
+            param.requires_grad = False
+        self.encoder.to(self.dtype)
+
+    def load_decoder(self, dec_fp: str) -> None:
+        """
+        Load the decoder from the remote store.
+
+        Args:
+        - dec_fp (str): File path to the decoder's JIT file on the remote store.
+        """
+        self.decoder = load_jit_model(dec_fp, device="cuda")
+        self.decoder.eval()
+        for param in self.decoder.parameters():
+            param.requires_grad = False
+        self.decoder.to(self.dtype)
+
+
+class DiscreteVideoFSQStateDictTokenizer(BaseDiscreteVideoFSQTokenizer):
+    """
+    A Discrete Video FSQ Tokenizer that loads weights from pre-trained JITed encoder
+    into as nn.Module so that encoder can be "torch.compile()" and JITed decoder, so it can be torch.compiled,
+    handles data type conversions, and normalization using provided mean and standard deviation values for latent
+    space representation.
+
+    Attributes:
+        tokenizer_module (Module): Tokenizer module with weights loaded from JIT checkpoints
+        encoder (Callable): tokenizer_module's encode method
+        decoder (Callable): tokenizer_module's decode method
+        dtype (dtype): Data type for model tensors, determined by whether bf16 is enabled.
+
+    Args:
+        enc_fp (str): File path to the encoder's JIT file on the remote store.
+        dec_fp (str): File path to the decoder's JIT file on the remote store.
+        tokenizer_module (Module): Tokenizer module that will have it's weights loaded
+        name (str): Name of the model, used for differentiating cache file paths.
+        latent_ch (int, optional): Number of latent channels (default is 6).
+        is_bf16 (bool, optional): Flag to use Brain Floating Point 16-bit data type (default is True).
+        pixel_chunk_duration (int): The duration (in number of frames) of each chunk of video data at the pixel level.
+        latent_chunk_duration (int): The duration (in number of frames) of each chunk at the latent representation level.
+        max_enc_batch_size (int): The maximum batch size to process in one go to avoid memory overflow.
+        level (list[int]): The level defined in FSQ quantizer.
+        compression_ratio (list[int]): The compression factor for (T, H, W).
+    """
+
+    def __init__(
+        self,
+        enc_fp: str,
+        dec_fp: str,
+        tokenizer_module: torch.nn.Module,
+        name: str,
+        latent_ch: int = 6,
+        is_bf16: bool = True,
+        pixel_chunk_duration: int = 25,
+        latent_chunk_duration: int = 4,
+        max_enc_batch_size: int = 8,
+        max_dec_batch_size: int = 4,
+        levels: list[int] = [8, 8, 8, 5, 5, 5],
+        compression_ratio: list[int] = [8, 16, 16],
+    ):
+        super().__init__(
+            name,
+            latent_ch,
+            is_bf16,
+            pixel_chunk_duration,
+            latent_chunk_duration,
+            max_enc_batch_size,
+            max_dec_batch_size,
+            levels,
+            compression_ratio,
+        )
+
+        self.load_encoder_and_decoder(enc_fp, dec_fp, tokenizer_module)
+
+    def load_encoder_and_decoder(self, enc_fp: str, dec_fp: str, tokenizer_module: torch.nn.Module) -> None:
+        """
+        Load the encoder from the remote store.
+
+        Args:
+        - enc_fp (str): File path to the encoder's JIT file on the remote store.
+        - def_fp (str): File path to the decoder's JIT file on the remote store.
+        - tokenizer_module (Module): Tokenizer module that was used to create JIT checkpoints
+        """
+        self.decoder = load_jit_model(dec_fp)
+
+        self.decoder.eval()
+        for param in self.decoder.parameters():
+            param.requires_grad = False
+        self.decoder.to(self.dtype)
+
+        encoder_sd = load_jit_model(enc_fp).state_dict()
+
+        del tokenizer_module.post_quant_conv
+        del tokenizer_module.decoder
+
+        state_dict = {
+            k: v
+            for k, v in (encoder_sd).items()
+            #  Variables captured by JIT
+            if k
+            not in (
+                "encoder.patcher3d.wavelets",
+                "encoder.patcher3d._arange",
+                "encoder.patcher3d.patch_size_buffer",
+                "quantizer._levels",
+                "quantizer._basis",
+                "quantizer.implicit_codebook",
+            )
+        }
+
+        tokenizer_module.load_state_dict(state_dict)
+
+        tokenizer_module.eval()
+        for param in tokenizer_module.parameters():
+            param.requires_grad = False
+        tokenizer_module.to(self.dtype)
+
+        self.tokenizer_module = tokenizer_module
+        self.encoder = self.tokenizer_module.encode
+
+    def reset_dtype(self, *args, **kwargs):
+        """
+        Resets the data type of the encoder and decoder to the model's default data type.
+
+        Args:
+            *args, **kwargs: Unused, present to allow flexibility in method calls.
+        """
+        del args, kwargs
+        self.decoder.to(self.dtype)
+        self.tokenizer_module.to(self.dtype)

ar_tokenizer_image_text_tokenizer.py

@@ -0,0 +1,318 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import transformers
+from transformers import AutoImageProcessor
+from transformers.image_utils import ImageInput, is_valid_image, load_image
+
+from .ar_tokenizer_text_tokenizer import TextTokenizer
+from .log import log
+
+# Configuration for different vision-language models
+IMAGE_CONFIGS = {
+    "pixtral": {
+        "patch_size": 16,
+        "image_token": "[IMG]",
+        "image_break_token": "[IMG_BREAK]",
+        "image_end_token": "[IMG_END]",
+    }
+}
+
+# Chat template for Pixtral-12B-Instruct
+PIXTRAL_CHAT_TEMPLATE = '{%- if messages[0]["role"] == "system" %}\n    {%- set system_message = messages[0]["content"] %}\n    {%- set loop_messages = messages[1:] %}\n{%- else %}\n    {%- set loop_messages = messages %}\n{%- endif %}\n\n{{- bos_token }}\n{%- for message in loop_messages %}\n    {%- if (message[\'role\'] == \'user\') != (loop.index0 % 2 == 0) %}\n        {{- raise_exception(\'After the optional system message, conversation roles must alternate user/assistant/user/assistant/...\') }}\n    {%- endif %}\n    {%- if message["role"] == "user" %}\n        {%- if loop.last and system_message is defined %}\n            {{- "[INST]" + system_message + "\n\n" }}\n        {%- else %}\n            {{- "[INST]" }}\n        {%- endif %}\n        {%- if message["content"] is not string %}\n            {%- for chunk in message["content"] %}\n                {%- if chunk["type"] == "text" %}\n                    {{- chunk["content"] }}\n                {%- elif chunk["type"] == "image" %}\n                    {{- "[IMG]" }}\n                {%- else %}\n                    {{- raise_exception("Unrecognized content type!") }}\n                {%- endif %}\n            {%- endfor %}\n        {%- else %}\n            {{- message["content"] }}\n        {%- endif %}\n        {{- "[/INST]" }}\n    {%- elif message["role"] == "assistant" %}\n        {{- message["content"] + eos_token}}\n    {%- else %}\n        {{- raise_exception("Only user and assistant roles are supported, with the exception of an initial optional system message!") }}\n    {%- endif %}\n{%- endfor %}'
+
+
+# Copied from transformers.models.pixtral.processing_pixtral.is_url
+def is_url(val) -> bool:
+    """Check if the given value is a URL."""
+    return isinstance(val, str) and val.startswith("http")
+
+
+# Copied from transformers.models.pixtral.processing_pixtral.is_image_or_image_url
+def is_image_or_image_url(elem):
+    """Check if the given element is an image or an image URL."""
+    return is_url(elem) or is_valid_image(elem)
+
+
+def load_image_list(
+    image_list: List[Union[str, "PIL.Image.Image"]], timeout: Optional[float] = None
+) -> List["PIL.Image.Image"]:
+    """
+    Load a list of images.
+
+    Args:
+        image_list (List[Union[str, PIL.Image.Image]]): The list of images to load.
+        timeout (Optional[float]): The timeout for loading the image.
+
+    Returns:
+        List[PIL.Image.Image]: The list of loaded images.
+    """
+    return [load_image(image, timeout=timeout) for image in image_list]
+
+
+class ImageTextTokenizer(TextTokenizer):
+    """
+    Image-text tokenizer class that extends the text tokenizer to support vision tokens as well.
+    """
+
+    def __init__(
+        self,
+        model_family: str,
+        is_instruct_model: bool,
+        tokenizer_path: str,
+        image_processor_path: str,
+    ):
+        """
+        Initialize the ImageTextTokenizer.
+
+        Args:
+            model_family (str): The model family.
+            is_instruct_model (bool): Whether the model is an instruct model.
+            s3_credential_path (str): The path to the s3 credential file. Defaults to "credentials/pbss_dir.secret".
+
+        Raises:
+            AssertionError: If the model family is not supported or if the transformers version is incompatible.
+        """
+        super().__init__(
+            model_family=model_family,
+            is_instruct_model=is_instruct_model,
+            local_path=tokenizer_path,
+        )
+        assert model_family in ["pixtral"], f"Unsupported model family: {model_family}"
+        if model_family == "pixtral":
+            # Need transformers>=4.45.0
+            assert transformers.__version__ >= "4.45.0", "Pixtral requires transformers>=4.45.0"
+            assert is_instruct_model, "Pixtral requires is_instruct_model=True"
+            if not hasattr(self.tokenizer, "chat_template") or self.tokenizer.chat_template is None:
+                setattr(self.tokenizer, "chat_template", PIXTRAL_CHAT_TEMPLATE)
+                log.debug(f"Pixtral tokenizer chat template set to: {PIXTRAL_CHAT_TEMPLATE}")
+
+        # Set up image-specific configurations
+        image_config = IMAGE_CONFIGS[model_family]
+        self.patch_size = image_config["patch_size"]
+        self.image_token = image_config["image_token"]
+        self.image_break_token = image_config["image_break_token"]
+        self.image_end_token = image_config["image_end_token"]
+
+        # Initialize the image processor
+        self.image_processor = AutoImageProcessor.from_pretrained(image_processor_path)
+
+    def encode(
+        self,
+        text: Union[str, List[str], List[int]],
+        *,  # Enforce keyword-only arguments
+        images: Optional[ImageInput] = None,
+        image_kwargs: Optional[Dict[str, Any]] = None,
+        **text_kwargs,
+    ) -> List[int]:
+        """
+        Process the images and return the tokenized images and text.
+
+        Args:
+            text (`str`, `List[str]`, `List[List[str]]`):
+                The sequence or batch of sequences to be encoded.
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+                The image or batch of images to be prepared.
+            image_kwargs (Optional[Dict[str, Any]]): Additional keyword arguments for image processing.
+            **text_kwargs: Additional keyword arguments for text processing.
+
+        Returns:
+            A dictionary with the following fields:
+            - **input_ids** -- List of token ids to be fed to a model.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model.
+            - **pixel_values** -- Pixel values to be fed to a model.
+
+        Raises:
+            ValueError: If the input images are in an invalid format.
+        """
+
+        output_dict, image_inputs = {}, {}
+        if images is not None:
+            # Preprocess images
+            if is_image_or_image_url(images):
+                images = [[images]]
+            elif isinstance(images, list) and is_image_or_image_url(images[0]):
+                images = [images]
+            elif (
+                not isinstance(images, list)
+                and not isinstance(images[0], list)
+                and not is_image_or_image_url(images[0][0])
+            ):
+                raise ValueError(
+                    "Invalid input images. Please provide a single image or a list of images or a list of list of images."
+                )
+
+            # Load and process images
+            images = [load_image_list(sample) for sample in images]
+            image_kwargs = image_kwargs or {}
+            image_inputs = self.image_processor(images, patch_size=self.patch_size, return_tensors="np", **image_kwargs)
+
+            # Validate image inputs
+            assert "pixel_values" in image_inputs, "pixel_values not found in image_inputs"
+            assert "image_sizes" in image_inputs, "image_sizes not found in image_inputs"
+            assert len(image_inputs.keys()) == 2, "Only one key is allowed in image_inputs, got {}".format(
+                image_inputs.keys()
+            )
+
+            # Extract pixel values and image sizes
+            pixel_values = image_inputs["pixel_values"][0]
+            image_sizes = image_inputs["image_sizes"][0]
+            unique_sizes = np.unique(image_sizes, axis=0)
+
+            assert len(unique_sizes) == 1, "All images must have the same size, got {}".format(unique_sizes)
+
+            # Convert pixel values to PyTorch tensor
+            pixel_values = np.asarray(pixel_values)
+            pixel_values = torch.from_numpy(pixel_values)
+            output_dict["pixel_values"] = pixel_values
+            output_dict["image_sizes"] = image_sizes
+
+        # Expand image tokens in text
+        if image_inputs.get("pixel_values") is not None:
+            replace_strings = []
+            # Calculate the number of tokens needed for each image and create a placeholder
+            for image_size in image_sizes:
+                height, width = image_size
+                num_height_tokens = height // self.patch_size
+                num_width_tokens = width // self.patch_size
+                replace_tokens = [[self.image_token] * num_width_tokens + [self.image_break_token]] * num_height_tokens
+                # Flatten list
+                replace_tokens = [item for sublist in replace_tokens for item in sublist]
+                replace_tokens[-1] = self.image_end_token
+                replace_str = "".join(replace_tokens)
+                replace_strings.append(replace_str)
+                text = text.replace(self.image_token, "<placeholder>", 1)
+
+            # Replace placeholders with actual image token sequences
+            while "<placeholder>" in text:
+                replace_str = replace_strings.pop(0)
+                text = text.replace("<placeholder>", replace_str, 1)
+
+        # Encode the text
+        text_inputs = super(ImageTextTokenizer, self).encode(text, **text_kwargs)
+
+        output_dict["input_ids"] = text_inputs
+        return output_dict
+
+    def apply_chat_template(
+        self,
+        conversation: List[Dict[str, Any]] | List[List[Dict[str, Any]]],
+        *,
+        images: Optional[ImageInput] = None,
+        image_kwargs: Optional[Dict[str, Any]] = None,
+        add_generation_prompt: bool = False,
+        tokenize: bool = True,
+        padding: bool = False,
+        truncation: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        return_dict: bool = True,
+        return_assistant_tokens_mask: bool = False,
+        generation_prefix: str = "",
+        tokenizer_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        """
+        Apply the chat template to the conversation.
+
+        Args:
+            conversation (List[Dict[str, Any]] | List[List[Dict[str, Any]]]): The conversation to process.
+            images (Optional[ImageInput]): Images to include in the conversation.
+            image_kwargs (Optional[Dict[str, Any]]): Additional keyword arguments for image processing.
+            add_generation_prompt (bool): Whether to add a generation prompt.
+            tokenize (bool): Whether to tokenize the output.
+            padding (bool): Whether to pad the output.
+            truncation (bool): Whether to truncate the output.
+            max_length (Optional[int]): Maximum length of the output.
+            return_tensors (Optional[str]): The type of tensors to return.
+            return_dict (bool): Whether to return a dictionary.
+            return_assistant_tokens_mask (bool): Whether to return the assistant tokens mask.
+            generation_prefix (str): Prefix to add before asking model to generate. Helpful to guide the generation. Defaults to "".
+            tokenizer_kwargs (Optional[Dict[str, Any]]): Additional keyword arguments for the tokenizer.
+            **kwargs: Additional keyword arguments.
+
+        Returns:
+            The processed conversation with applied chat template.
+
+        Raises:
+            AssertionError: If return_dict is False or if the conversation format is invalid.
+        """
+        assert return_dict, "return_dict must be True for ImageTextTokenizer"
+        assert isinstance(conversation, list), "conversation must be a list"
+        if isinstance(conversation[0], list):
+            assert len(conversation) == 1, "Only support single-conversation input, got {}".format(conversation)
+            conversation = conversation[0]
+
+        # Extract images from the conversation if not provided
+        if images is None:
+            images = []
+            for msg in conversation:
+                if msg.get("images", None) is not None:
+                    images = images + (msg["images"])
+            images = load_image_list(images)
+        # In case the input does not have images, will ignore
+        # Useful in feeding VLM inputs with and without images
+        if isinstance(images, list) and len(images) == 0:
+            images = None
+
+        # Apply the chat template to the text
+        text = super().apply_chat_template(
+            conversation,
+            tokenize=False,
+            add_generation_prompt=add_generation_prompt,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            return_dict=False,
+            return_assistant_tokens_mask=return_assistant_tokens_mask,
+            generation_prefix=generation_prefix,
+            tokenizer_kwargs=tokenizer_kwargs,
+            **kwargs,
+        )
+
+        if tokenizer_kwargs is None:
+            tokenizer_kwargs = {}
+
+        # Encode the text and images
+        output = self.encode(
+            text,
+            images=images,
+            image_kwargs=image_kwargs,
+            tokenize=tokenize,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            add_special_tokens=False,
+            return_tensors=return_tensors,
+            **tokenizer_kwargs,
+        )
+        return output
+
+    @property
+    def model_input_names(self):
+        """
+        Get the combined model input names from both the text tokenizer and image processor.
+
+        Returns:
+            List[str]: A list of unique input names.
+        """
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))

ar_tokenizer_modules.py

@@ -0,0 +1,560 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The model definition for 3D layers
+
+Adapted from: https://github.com/lucidrains/magvit2-pytorch/blob/9f49074179c912736e617d61b32be367eb5f993a/
+magvit2_pytorch/magvit2_pytorch.py#L889
+
+[MIT License Copyright (c) 2023 Phil Wang]
+https://github.com/lucidrains/magvit2-pytorch/blob/9f49074179c912736e617d61b32be367eb5f993a/LICENSE
+"""
+import math
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .ar_tokenizer_patching import Patcher3D, UnPatcher3D
+from .ar_tokenizer_utils import (
+    CausalNormalize,
+    batch2space,
+    batch2time,
+    cast_tuple,
+    is_odd,
+    nonlinearity,
+    replication_pad,
+    space2batch,
+    time2batch,
+)
+from .log import log
+
+
+class CausalConv3d(nn.Module):
+    def __init__(
+        self,
+        chan_in: int = 1,
+        chan_out: int = 1,
+        kernel_size: Union[int, Tuple[int, int, int]] = 3,
+        pad_mode: str = "constant",
+        **kwargs,
+    ):
+        super().__init__()
+        kernel_size = cast_tuple(kernel_size, 3)
+
+        time_kernel_size, height_kernel_size, width_kernel_size = kernel_size
+
+        assert is_odd(height_kernel_size) and is_odd(width_kernel_size)
+
+        dilation = kwargs.pop("dilation", 1)
+        stride = kwargs.pop("stride", 1)
+        time_stride = kwargs.pop("time_stride", 1)
+        time_dilation = kwargs.pop("time_dilation", 1)
+        padding = kwargs.pop("padding", 1)
+
+        self.pad_mode = pad_mode
+        time_pad = time_dilation * (time_kernel_size - 1) + (1 - time_stride)
+        self.time_pad = time_pad
+
+        self.spatial_pad = (padding, padding, padding, padding)
+
+        stride = (time_stride, stride, stride)
+        dilation = (time_dilation, dilation, dilation)
+        self.conv3d = nn.Conv3d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
+
+    def _replication_pad(self, x: torch.Tensor) -> torch.Tensor:
+        x_prev = x[:, :, :1, ...].repeat(1, 1, self.time_pad, 1, 1)
+        x = torch.cat([x_prev, x], dim=2)
+        padding = self.spatial_pad + (0, 0)
+        return F.pad(x, padding, mode=self.pad_mode, value=0.0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self._replication_pad(x)
+        return self.conv3d(x)
+
+
+class CausalHybridUpsample3d(nn.Module):
+    def __init__(self, in_channels: int, spatial_up: bool = True, temporal_up: bool = True, **ignore_kwargs) -> None:
+        super().__init__()
+        self.conv1 = (
+            CausalConv3d(in_channels, in_channels, kernel_size=(3, 1, 1), stride=1, time_stride=1, padding=0)
+            if temporal_up
+            else nn.Identity()
+        )
+        self.conv2 = (
+            CausalConv3d(in_channels, in_channels, kernel_size=(1, 3, 3), stride=1, time_stride=1, padding=1)
+            if spatial_up
+            else nn.Identity()
+        )
+        self.conv3 = (
+            CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1, time_stride=1, padding=0)
+            if spatial_up or temporal_up
+            else nn.Identity()
+        )
+        self.spatial_up = spatial_up
+        self.temporal_up = temporal_up
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.spatial_up and not self.temporal_up:
+            return x
+
+        # hybrid upsample temporally.
+        if self.temporal_up:
+            time_factor = 1.0 + 1.0 * (x.shape[2] > 1)
+            if isinstance(time_factor, torch.Tensor):
+                time_factor = time_factor.item()
+            x = x.repeat_interleave(int(time_factor), dim=2)
+            x = x[..., int(time_factor - 1) :, :, :]
+            x = self.conv1(x) + x
+
+        # hybrid upsample spatially.
+        if self.spatial_up:
+            x = x.repeat_interleave(2, dim=3).repeat_interleave(2, dim=4)
+            x = self.conv2(x) + x
+
+        # final 1x1x1 conv.
+        x = self.conv3(x)
+        return x
+
+
+class CausalHybridDownsample3d(nn.Module):
+    def __init__(
+        self, in_channels: int, spatial_down: bool = True, temporal_down: bool = True, **ignore_kwargs
+    ) -> None:
+        super().__init__()
+        self.conv1 = (
+            CausalConv3d(in_channels, in_channels, kernel_size=(1, 3, 3), stride=2, time_stride=1, padding=0)
+            if spatial_down
+            else nn.Identity()
+        )
+        self.conv2 = (
+            CausalConv3d(in_channels, in_channels, kernel_size=(3, 1, 1), stride=1, time_stride=2, padding=0)
+            if temporal_down
+            else nn.Identity()
+        )
+        self.conv3 = (
+            CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1, time_stride=1, padding=0)
+            if spatial_down or temporal_down
+            else nn.Identity()
+        )
+        self.spatial_down = spatial_down
+        self.temporal_down = temporal_down
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.spatial_down and not self.temporal_down:
+            return x
+
+        # hybrid downsample spatially.
+        if self.spatial_down:
+            pad = (0, 1, 0, 1, 0, 0)
+            x = F.pad(x, pad, mode="constant", value=0)
+            x1 = self.conv1(x)
+            x2 = F.avg_pool3d(x, kernel_size=(1, 2, 2), stride=(1, 2, 2))
+            x = x1 + x2
+
+        # hybrid downsample temporally.
+        if self.temporal_down:
+            x = replication_pad(x)
+            x1 = self.conv2(x)
+            x2 = F.avg_pool3d(x, kernel_size=(2, 1, 1), stride=(2, 1, 1))
+            x = x1 + x2
+
+        # final 1x1x1 conv.
+        x = self.conv3(x)
+        return x
+
+
+class CausalResnetBlockFactorized3d(nn.Module):
+    def __init__(self, *, in_channels: int, out_channels: int = None, dropout: float, num_groups: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+
+        self.norm1 = CausalNormalize(in_channels, num_groups=1)
+        self.conv1 = nn.Sequential(
+            CausalConv3d(in_channels, out_channels, kernel_size=(1, 3, 3), stride=1, padding=1),
+            CausalConv3d(out_channels, out_channels, kernel_size=(3, 1, 1), stride=1, padding=0),
+        )
+        self.norm2 = CausalNormalize(out_channels, num_groups=num_groups)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = nn.Sequential(
+            CausalConv3d(out_channels, out_channels, kernel_size=(1, 3, 3), stride=1, padding=1),
+            CausalConv3d(out_channels, out_channels, kernel_size=(3, 1, 1), stride=1, padding=0),
+        )
+        self.nin_shortcut = (
+            CausalConv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+        x = self.nin_shortcut(x)
+
+        return x + h
+
+
+class CausalAttnBlock(nn.Module):
+    def __init__(self, in_channels: int, num_groups: int) -> None:
+        super().__init__()
+
+        self.norm = CausalNormalize(in_channels, num_groups=num_groups)
+        self.q = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        q, batch_size = time2batch(q)
+        k, batch_size = time2batch(k)
+        v, batch_size = time2batch(v)
+
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h * w)
+        q = q.permute(0, 2, 1)
+        k = k.reshape(b, c, h * w)
+        w_ = torch.bmm(q, k)
+        w_ = w_ * (int(c) ** (-0.5))
+        w_ = F.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h * w)
+        w_ = w_.permute(0, 2, 1)
+        h_ = torch.bmm(v, w_)
+        h_ = h_.reshape(b, c, h, w)
+
+        h_ = batch2time(h_, batch_size)
+        h_ = self.proj_out(h_)
+        return x + h_
+
+
+class CausalTemporalAttnBlock(nn.Module):
+    def __init__(self, in_channels: int, num_groups: int) -> None:
+        super().__init__()
+
+        self.norm = CausalNormalize(in_channels, num_groups=num_groups)
+        self.q = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = CausalConv3d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        q, batch_size, height = space2batch(q)
+        k, _, _ = space2batch(k)
+        v, _, _ = space2batch(v)
+
+        bhw, c, t = q.shape
+        q = q.permute(0, 2, 1)  # (bhw, t, c)
+        k = k.permute(0, 2, 1)  # (bhw, t, c)
+        v = v.permute(0, 2, 1)  # (bhw, t, c)
+
+        w_ = torch.bmm(q, k.permute(0, 2, 1))  # (bhw, t, t)
+        w_ = w_ * (int(c) ** (-0.5))
+
+        # Apply causal mask
+        mask = torch.tril(torch.ones_like(w_))
+        w_ = w_.masked_fill(mask == 0, float("-inf"))
+        w_ = F.softmax(w_, dim=2)
+
+        # attend to values
+        h_ = torch.bmm(w_, v)  # (bhw, t, c)
+        h_ = h_.permute(0, 2, 1).reshape(bhw, c, t)  # (bhw, c, t)
+
+        h_ = batch2space(h_, batch_size, height)
+        h_ = self.proj_out(h_)
+        return x + h_
+
+
+class EncoderFactorized(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        channels: int,
+        channels_mult: list[int],
+        num_res_blocks: int,
+        attn_resolutions: list[int],
+        dropout: float,
+        resolution: int,
+        z_channels: int,
+        spatial_compression: int,
+        temporal_compression: int,
+        **ignore_kwargs,
+    ) -> None:
+        super().__init__()
+        self.num_resolutions = len(channels_mult)
+        self.num_res_blocks = num_res_blocks
+
+        # Patcher.
+        patch_size = ignore_kwargs.get("patch_size", 1)
+        self.patcher3d = Patcher3D(patch_size, ignore_kwargs.get("patch_method", "rearrange"))
+        in_channels = in_channels * patch_size * patch_size * patch_size
+
+        # calculate the number of downsample operations
+        self.num_spatial_downs = int(math.log2(spatial_compression)) - int(math.log2(patch_size))
+        assert (
+            self.num_spatial_downs <= self.num_resolutions
+        ), f"Spatially downsample {self.num_resolutions} times at most"
+
+        self.num_temporal_downs = int(math.log2(temporal_compression)) - int(math.log2(patch_size))
+        assert (
+            self.num_temporal_downs <= self.num_resolutions
+        ), f"Temporally downsample {self.num_resolutions} times at most"
+
+        # downsampling
+        self.conv_in = nn.Sequential(
+            CausalConv3d(in_channels, channels, kernel_size=(1, 3, 3), stride=1, padding=1),
+            CausalConv3d(channels, channels, kernel_size=(3, 1, 1), stride=1, padding=0),
+        )
+
+        curr_res = resolution // patch_size
+        in_ch_mult = (1,) + tuple(channels_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = channels * in_ch_mult[i_level]
+            block_out = channels * channels_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(
+                    CausalResnetBlockFactorized3d(
+                        in_channels=block_in, out_channels=block_out, dropout=dropout, num_groups=1
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(
+                        nn.Sequential(
+                            CausalAttnBlock(block_in, num_groups=1), CausalTemporalAttnBlock(block_in, num_groups=1)
+                        )
+                    )
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                spatial_down = i_level < self.num_spatial_downs
+                temporal_down = i_level < self.num_temporal_downs
+                down.downsample = CausalHybridDownsample3d(
+                    block_in, spatial_down=spatial_down, temporal_down=temporal_down
+                )
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = CausalResnetBlockFactorized3d(
+            in_channels=block_in, out_channels=block_in, dropout=dropout, num_groups=1
+        )
+        self.mid.attn_1 = nn.Sequential(
+            CausalAttnBlock(block_in, num_groups=1), CausalTemporalAttnBlock(block_in, num_groups=1)
+        )
+        self.mid.block_2 = CausalResnetBlockFactorized3d(
+            in_channels=block_in, out_channels=block_in, dropout=dropout, num_groups=1
+        )
+
+        # end
+        self.norm_out = CausalNormalize(block_in, num_groups=1)
+        self.conv_out = nn.Sequential(
+            CausalConv3d(block_in, z_channels, kernel_size=(1, 3, 3), stride=1, padding=1),
+            CausalConv3d(z_channels, z_channels, kernel_size=(3, 1, 1), stride=1, padding=0),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.patcher3d(x)
+
+        # downsampling
+        h = self.conv_in(x)
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](h)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+            if i_level != self.num_resolutions - 1:
+                h = self.down[i_level].downsample(h)
+
+        # middle
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class DecoderFactorized(nn.Module):
+    def __init__(
+        self,
+        out_channels: int,
+        channels: int,
+        channels_mult: list[int],
+        num_res_blocks: int,
+        attn_resolutions: list[int],
+        dropout: float,
+        resolution: int,
+        z_channels: int,
+        spatial_compression: int,
+        temporal_compression: int,
+        **ignore_kwargs,
+    ):
+        super().__init__()
+        self.num_resolutions = len(channels_mult)
+        self.num_res_blocks = num_res_blocks
+
+        # UnPatcher.
+        patch_size = ignore_kwargs.get("patch_size", 1)
+        self.unpatcher3d = UnPatcher3D(patch_size, ignore_kwargs.get("patch_method", "rearrange"))
+        out_ch = out_channels * patch_size * patch_size * patch_size
+
+        # calculate the number of upsample operations
+        self.num_spatial_ups = int(math.log2(spatial_compression)) - int(math.log2(patch_size))
+        assert self.num_spatial_ups <= self.num_resolutions, f"Spatially upsample {self.num_resolutions} times at most"
+        self.num_temporal_ups = int(math.log2(temporal_compression)) - int(math.log2(patch_size))
+        assert (
+            self.num_temporal_ups <= self.num_resolutions
+        ), f"Temporally upsample {self.num_resolutions} times at most"
+
+        block_in = channels * channels_mult[self.num_resolutions - 1]
+        curr_res = (resolution // patch_size) // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+        log.debug("Working with z of shape {} = {} dimensions.".format(self.z_shape, np.prod(self.z_shape)))
+
+        # z to block_in
+        self.conv_in = nn.Sequential(
+            CausalConv3d(z_channels, block_in, kernel_size=(1, 3, 3), stride=1, padding=1),
+            CausalConv3d(block_in, block_in, kernel_size=(3, 1, 1), stride=1, padding=0),
+        )
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = CausalResnetBlockFactorized3d(
+            in_channels=block_in, out_channels=block_in, dropout=dropout, num_groups=1
+        )
+        self.mid.attn_1 = nn.Sequential(
+            CausalAttnBlock(block_in, num_groups=1), CausalTemporalAttnBlock(block_in, num_groups=1)
+        )
+        self.mid.block_2 = CausalResnetBlockFactorized3d(
+            in_channels=block_in, out_channels=block_in, dropout=dropout, num_groups=1
+        )
+
+        legacy_mode = ignore_kwargs.get("legacy_mode", False)
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = channels * channels_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(
+                    CausalResnetBlockFactorized3d(
+                        in_channels=block_in, out_channels=block_out, dropout=dropout, num_groups=1
+                    )
+                )
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(
+                        nn.Sequential(
+                            CausalAttnBlock(block_in, num_groups=1), CausalTemporalAttnBlock(block_in, num_groups=1)
+                        )
+                    )
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                # The layer index for temporal/spatial downsampling performed in the encoder should correspond
+                # to the layer index, inreverse order, where upsampling is performed in the decoder.
+                # If you've a pre-trained model, you can simply finetune.
+                # For example:
+                #   Input tensor = (1, 3, 17, 32, 32)
+                #   Patch size = 4 for 3D wavelet transform
+                #   Compression rate = (8x16x16)
+                #
+                # We expect successive downsampling in the encoder and upsampling in the decoder to be mirrored.
+                # ENCODER: `(...,5,8,8) -> (...,3,4,4) -> (...,3,2,2)`
+                # DECODER: `(...,3,2,2) -> (...,3,4,4) -> (...,5,8,8)`
+                #
+                # if legacy_mode is True, the temporal upsampling is not perfectly mirrored.
+                # ENCODER: `(...,5,8,8) -> (...,3,4,4) -> (...,3,2,2)`
+                # DECODER: `(...,3,2,2) -> (...,5,4,4) -> (...,5,8,8)`
+                #
+                # Most of the CV and DV tokenizers were trained before 09/01/2024 with upsampling that's not mirrored.
+                # Going forward, new CV/DV tokenizers will adopt `legacy_mode=False`, i.e. use mirrored upsampling.
+                i_level_reverse = self.num_resolutions - i_level - 1
+                if legacy_mode:
+                    temporal_up = i_level_reverse < self.num_temporal_ups
+                else:
+                    temporal_up = 0 < i_level_reverse < self.num_temporal_ups + 1
+                spatial_up = temporal_up or (
+                    i_level_reverse < self.num_spatial_ups and self.num_spatial_ups > self.num_temporal_ups
+                )
+                up.upsample = CausalHybridUpsample3d(block_in, spatial_up=spatial_up, temporal_up=temporal_up)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = CausalNormalize(block_in, num_groups=1)
+        self.conv_out = nn.Sequential(
+            CausalConv3d(block_in, out_ch, kernel_size=(1, 3, 3), stride=1, padding=1),
+            CausalConv3d(out_ch, out_ch, kernel_size=(3, 1, 1), stride=1, padding=0),
+        )
+
+    def forward(self, z):
+        h = self.conv_in(z)
+
+        # middle block.
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+
+        # decoder blocks.
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        h = self.unpatcher3d(h)
+        return h

ar_tokenizer_networks.py

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import namedtuple
+
+import torch
+from torch import nn
+
+from .ar_tokenizer_modules import CausalConv3d, DecoderFactorized, EncoderFactorized
+from .ar_tokenizer_quantizers import FSQuantizer
+from .log import log
+
+NetworkEval = namedtuple("NetworkEval", ["reconstructions", "quant_loss", "quant_info"])
+
+
+class CausalDiscreteVideoTokenizer(nn.Module):
+    def __init__(self, z_channels: int, z_factor: int, embedding_dim: int, **kwargs) -> None:
+        super().__init__()
+        self.name = kwargs.get("name", "CausalDiscreteVideoTokenizer")
+        self.embedding_dim = embedding_dim
+        self.encoder = EncoderFactorized(z_channels=z_factor * z_channels, **kwargs)
+        self.decoder = DecoderFactorized(z_channels=z_channels, **kwargs)
+
+        self.quant_conv = CausalConv3d(z_factor * z_channels, embedding_dim, kernel_size=1, padding=0)
+        self.post_quant_conv = CausalConv3d(embedding_dim, z_channels, kernel_size=1, padding=0)
+
+        self.quantizer = FSQuantizer(**kwargs)
+
+        num_parameters = sum(param.numel() for param in self.parameters())
+        log.debug(f"model={self.name}, num_parameters={num_parameters:,}")
+        log.debug(f"z_channels={z_channels}, embedding_dim={self.embedding_dim}.")
+
+    def to(self, *args, **kwargs):
+        setattr(self.quantizer, "dtype", kwargs.get("dtype", torch.bfloat16))
+        return super(CausalDiscreteVideoTokenizer, self).to(*args, **kwargs)
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        return self.quantizer(h)
+
+    def decode(self, quant):
+        quant = self.post_quant_conv(quant)
+        return self.decoder(quant)
+
+    def forward(self, input):
+        quant_info, quant_codes, quant_loss = self.encode(input)
+        reconstructions = self.decode(quant_codes)
+        if self.training:
+            return dict(reconstructions=reconstructions, quant_loss=quant_loss, quant_info=quant_info)
+        return NetworkEval(reconstructions=reconstructions, quant_loss=quant_loss, quant_info=quant_info)

ar_tokenizer_patching.py

@@ -0,0 +1,279 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The patcher and unpatcher implementation for 2D and 3D data."""
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+_WAVELETS = {
+    "haar": torch.tensor([0.7071067811865476, 0.7071067811865476]),
+    "rearrange": torch.tensor([1.0, 1.0]),
+}
+_PERSISTENT = False
+
+
+class Patcher(torch.nn.Module):
+    """A module to convert image tensors into patches using torch operations.
+
+    The main difference from `class Patching` is that this module implements
+    all operations using torch, rather than python or numpy, for efficiency purpose.
+
+    It's bit-wise identical to the Patching module outputs, with the added
+    benefit of being torch.jit scriptable.
+    """
+
+    def __init__(self, patch_size=1, patch_method="haar"):
+        super().__init__()
+        self.patch_size = patch_size
+        self.patch_method = patch_method
+        self.register_buffer("wavelets", _WAVELETS[patch_method], persistent=_PERSISTENT)
+        self.range = range(int(torch.log2(torch.tensor(self.patch_size)).item()))
+        self.register_buffer("_arange", torch.arange(_WAVELETS[patch_method].shape[0]), persistent=_PERSISTENT)
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, x):
+        if self.patch_method == "haar":
+            return self._haar(x)
+        elif self.patch_method == "rearrange":
+            return self._arrange(x)
+        else:
+            raise ValueError("Unknown patch method: " + self.patch_method)
+
+    def _dwt(self, x, mode="reflect", rescale=False):
+        dtype = x.dtype
+        h = self.wavelets
+
+        n = h.shape[0]
+        g = x.shape[1]
+        hl = h.flip(0).reshape(1, 1, -1).repeat(g, 1, 1)
+        hh = (h * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1)
+        hh = hh.to(dtype=dtype)
+        hl = hl.to(dtype=dtype)
+
+        x = F.pad(x, pad=(n - 2, n - 1, n - 2, n - 1), mode=mode).to(dtype)
+        xl = F.conv2d(x, hl.unsqueeze(2), groups=g, stride=(1, 2))
+        xh = F.conv2d(x, hh.unsqueeze(2), groups=g, stride=(1, 2))
+        xll = F.conv2d(xl, hl.unsqueeze(3), groups=g, stride=(2, 1))
+        xlh = F.conv2d(xl, hh.unsqueeze(3), groups=g, stride=(2, 1))
+        xhl = F.conv2d(xh, hl.unsqueeze(3), groups=g, stride=(2, 1))
+        xhh = F.conv2d(xh, hh.unsqueeze(3), groups=g, stride=(2, 1))
+
+        out = torch.cat([xll, xlh, xhl, xhh], dim=1)
+        if rescale:
+            out = out / 2
+        return out
+
+    def _haar(self, x):
+        for _ in self.range:
+            x = self._dwt(x, rescale=True)
+        return x
+
+    def _arrange(self, x):
+        x = rearrange(x, "b c (h p1) (w p2) -> b (c p1 p2) h w", p1=self.patch_size, p2=self.patch_size).contiguous()
+        return x
+
+
+class Patcher3D(Patcher):
+    """A 3D discrete wavelet transform for video data, expects 5D tensor, i.e. a batch of videos."""
+
+    def __init__(self, patch_size=1, patch_method="haar"):
+        super().__init__(patch_method=patch_method, patch_size=patch_size)
+        self.register_buffer(
+            "patch_size_buffer", patch_size * torch.ones([1], dtype=torch.int32), persistent=_PERSISTENT
+        )
+
+    def _dwt(self, x, mode="reflect", rescale=False):
+        dtype = x.dtype
+        h = self.wavelets
+
+        n = h.shape[0]
+        g = x.shape[1]
+        hl = h.flip(0).reshape(1, 1, -1).repeat(g, 1, 1)
+        hh = (h * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1)
+        hh = hh.to(dtype=dtype)
+        hl = hl.to(dtype=dtype)
+
+        # Handles temporal axis.
+        x = F.pad(x, pad=(max(0, n - 2), n - 1, n - 2, n - 1, n - 2, n - 1), mode=mode).to(dtype)
+        xl = F.conv3d(x, hl.unsqueeze(3).unsqueeze(4), groups=g, stride=(2, 1, 1))
+        xh = F.conv3d(x, hh.unsqueeze(3).unsqueeze(4), groups=g, stride=(2, 1, 1))
+
+        # Handles spatial axes.
+        xll = F.conv3d(xl, hl.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1))
+        xlh = F.conv3d(xl, hh.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1))
+        xhl = F.conv3d(xh, hl.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1))
+        xhh = F.conv3d(xh, hh.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1))
+
+        xlll = F.conv3d(xll, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xllh = F.conv3d(xll, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xlhl = F.conv3d(xlh, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xlhh = F.conv3d(xlh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xhll = F.conv3d(xhl, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xhlh = F.conv3d(xhl, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xhhl = F.conv3d(xhh, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xhhh = F.conv3d(xhh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+
+        out = torch.cat([xlll, xllh, xlhl, xlhh, xhll, xhlh, xhhl, xhhh], dim=1)
+        if rescale:
+            out = out / (2 * torch.sqrt(torch.tensor(2.0)))
+        return out
+
+    def _haar(self, x):
+        xi, xv = torch.split(x, [1, x.shape[2] - 1], dim=2)
+        x = torch.cat([xi.repeat_interleave(self.patch_size, dim=2), xv], dim=2)
+        for _ in self.range:
+            x = self._dwt(x, rescale=True)
+        return x
+
+    def _arrange(self, x):
+        xi, xv = torch.split(x, [1, x.shape[2] - 1], dim=2)
+        x = torch.cat([xi.repeat_interleave(self.patch_size, dim=2), xv], dim=2)
+        x = rearrange(
+            x,
+            "b c (t p1) (h p2) (w p3) -> b (c p1 p2 p3) t h w",
+            p1=self.patch_size,
+            p2=self.patch_size,
+            p3=self.patch_size,
+        ).contiguous()
+        return x
+
+
+class UnPatcher(torch.nn.Module):
+    """A module to convert patches into image tensorsusing torch operations.
+
+    The main difference from `class Unpatching` is that this module implements
+    all operations using torch, rather than python or numpy, for efficiency purpose.
+
+    It's bit-wise identical to the Unpatching module outputs, with the added
+    benefit of being torch.jit scriptable.
+    """
+
+    def __init__(self, patch_size=1, patch_method="haar"):
+        super().__init__()
+        self.patch_size = patch_size
+        self.patch_method = patch_method
+        self.register_buffer("wavelets", _WAVELETS[patch_method], persistent=_PERSISTENT)
+        self.range = range(int(torch.log2(torch.tensor(self.patch_size)).item()))
+        self.register_buffer("_arange", torch.arange(_WAVELETS[patch_method].shape[0]), persistent=_PERSISTENT)
+        for param in self.parameters():
+            param.requires_grad = False
+
+    def forward(self, x):
+        if self.patch_method == "haar":
+            return self._ihaar(x)
+        elif self.patch_method == "rearrange":
+            return self._iarrange(x)
+        else:
+            raise ValueError("Unknown patch method: " + self.patch_method)
+
+    def _idwt(self, x, rescale=False):
+        dtype = x.dtype
+        h = self.wavelets
+        n = h.shape[0]
+
+        g = x.shape[1] // 4
+        hl = h.flip([0]).reshape(1, 1, -1).repeat([g, 1, 1])
+        hh = (h * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1)
+        hh = hh.to(dtype=dtype)
+        hl = hl.to(dtype=dtype)
+
+        xll, xlh, xhl, xhh = torch.chunk(x.to(dtype), 4, dim=1)
+
+        # Inverse transform.
+        yl = torch.nn.functional.conv_transpose2d(xll, hl.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        yl += torch.nn.functional.conv_transpose2d(xlh, hh.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        yh = torch.nn.functional.conv_transpose2d(xhl, hl.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        yh += torch.nn.functional.conv_transpose2d(xhh, hh.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        y = torch.nn.functional.conv_transpose2d(yl, hl.unsqueeze(2), groups=g, stride=(1, 2), padding=(0, n - 2))
+        y += torch.nn.functional.conv_transpose2d(yh, hh.unsqueeze(2), groups=g, stride=(1, 2), padding=(0, n - 2))
+
+        if rescale:
+            y = y * 2
+        return y
+
+    def _ihaar(self, x):
+        for _ in self.range:
+            x = self._idwt(x, rescale=True)
+        return x
+
+    def _iarrange(self, x):
+        x = rearrange(x, "b (c p1 p2) h w -> b c (h p1) (w p2)", p1=self.patch_size, p2=self.patch_size)
+        return x
+
+
+class UnPatcher3D(UnPatcher):
+    """A 3D inverse discrete wavelet transform for video wavelet decompositions."""
+
+    def __init__(self, patch_size=1, patch_method="haar"):
+        super().__init__(patch_method=patch_method, patch_size=patch_size)
+
+    def _idwt(self, x, rescale=False):
+        dtype = x.dtype
+        h = self.wavelets
+
+        g = x.shape[1] // 8  # split into 8 spatio-temporal filtered tesnors.
+        hl = h.flip([0]).reshape(1, 1, -1).repeat([g, 1, 1])
+        hh = (h * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1)
+        hl = hl.to(dtype=dtype)
+        hh = hh.to(dtype=dtype)
+
+        xlll, xllh, xlhl, xlhh, xhll, xhlh, xhhl, xhhh = torch.chunk(x, 8, dim=1)
+
+        # Height height transposed convolutions.
+        xll = F.conv_transpose3d(xlll, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xll += F.conv_transpose3d(xllh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+
+        xlh = F.conv_transpose3d(xlhl, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xlh += F.conv_transpose3d(xlhh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+
+        xhl = F.conv_transpose3d(xhll, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xhl += F.conv_transpose3d(xhlh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+
+        xhh = F.conv_transpose3d(xhhl, hl.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+        xhh += F.conv_transpose3d(xhhh, hh.unsqueeze(2).unsqueeze(3), groups=g, stride=(1, 1, 2))
+
+        # Handles width transposed convolutions.
+        xl = F.conv_transpose3d(xll, hl.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1))
+        xl += F.conv_transpose3d(xlh, hh.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1))
+        xh = F.conv_transpose3d(xhl, hl.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1))
+        xh += F.conv_transpose3d(xhh, hh.unsqueeze(2).unsqueeze(4), groups=g, stride=(1, 2, 1))
+
+        # Handles time axis transposed convolutions.
+        x = F.conv_transpose3d(xl, hl.unsqueeze(3).unsqueeze(4), groups=g, stride=(2, 1, 1))
+        x += F.conv_transpose3d(xh, hh.unsqueeze(3).unsqueeze(4), groups=g, stride=(2, 1, 1))
+
+        if rescale:
+            x = x * (2 * torch.sqrt(torch.tensor(2.0)))
+        return x
+
+    def _ihaar(self, x):
+        for _ in self.range:
+            x = self._idwt(x, rescale=True)
+        x = x[:, :, self.patch_size - 1 :, ...]
+        return x
+
+    def _iarrange(self, x):
+        x = rearrange(
+            x,
+            "b (c p1 p2 p3) t h w -> b c (t p1) (h p2) (w p3)",
+            p1=self.patch_size,
+            p2=self.patch_size,
+            p3=self.patch_size,
+        )
+        x = x[:, :, self.patch_size - 1 :, ...]
+        return x

ar_tokenizer_quantizers.py

@@ -0,0 +1,165 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Quantizers for discrete image and video tokenization."""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+from .ar_tokenizer_utils import default, pack_one, round_ste, unpack_one
+
+
+class FSQuantizer(nn.Module):
+    """Finite Scalar Quantization: VQ-VAE Made Simple - https://arxiv.org/abs/2309.15505
+
+    Adapted from: https://github.com/lucidrains/vector-quantize-pytorch/blob/9502a1f447876d53fd37685b226bf28f250dc4a3/
+    vector_quantize_pytorch/finite_scalar_quantization.py
+    [Copyright (c) 2020 Phil Wang]
+    https://github.com/lucidrains/vector-quantize-pytorch/blob/9502a1f447876d53fd37685b226bf28f250dc4a3/LICENSE
+    """
+
+    def __init__(
+        self,
+        levels: list[int],
+        dim: Optional[int] = None,
+        num_codebooks=1,
+        keep_num_codebooks_dim: Optional[bool] = None,
+        scale: Optional[float] = None,
+        **ignore_kwargs,
+    ):
+        super().__init__()
+        self.dtype = ignore_kwargs.get("dtype", torch.float32)
+        _levels = torch.tensor(levels, dtype=torch.int32)
+        self.register_buffer("_levels", _levels, persistent=False)
+
+        _basis = torch.cumprod(torch.tensor([1] + levels[:-1]), dim=0, dtype=torch.int32)
+        self.register_buffer("_basis", _basis, persistent=False)
+
+        self.scale = scale
+
+        codebook_dim = len(levels)
+        self.codebook_dim = codebook_dim
+
+        effective_codebook_dim = codebook_dim * num_codebooks
+        self.num_codebooks = num_codebooks
+        self.effective_codebook_dim = effective_codebook_dim
+
+        keep_num_codebooks_dim = default(keep_num_codebooks_dim, num_codebooks > 1)
+        assert not (num_codebooks > 1 and not keep_num_codebooks_dim)
+        self.keep_num_codebooks_dim = keep_num_codebooks_dim
+
+        self.dim = default(dim, len(_levels) * num_codebooks)
+
+        has_projections = self.dim != effective_codebook_dim
+        self.project_in = nn.Linear(self.dim, effective_codebook_dim) if has_projections else nn.Identity()
+        self.project_out = nn.Linear(effective_codebook_dim, self.dim) if has_projections else nn.Identity()
+        self.has_projections = has_projections
+
+        self.codebook_size = self._levels.prod().item()
+
+        implicit_codebook = self.indices_to_codes(torch.arange(self.codebook_size), project_out=False)
+        self.register_buffer("implicit_codebook", implicit_codebook, persistent=False)
+
+    def bound(self, z: torch.Tensor, eps: float = 1e-3) -> torch.Tensor:
+        """Bound `z`, an array of shape (..., d)."""
+        half_l = (self._levels - 1) * (1 + eps) / 2
+        offset = torch.where(self._levels % 2 == 0, 0.5, 0.0)
+        shift = (offset / half_l).atanh()
+        return (z + shift).tanh() * half_l - offset
+
+    def quantize(self, z: torch.Tensor) -> torch.Tensor:
+        """Quantizes z, returns quantized zhat, same shape as z."""
+        quantized = round_ste(self.bound(z))
+        half_width = self._levels // 2  # Renormalize to [-1, 1].
+        return quantized / half_width
+
+    def _scale_and_shift(self, zhat_normalized: torch.Tensor) -> torch.Tensor:
+        half_width = self._levels // 2
+        return (zhat_normalized * half_width) + half_width
+
+    def _scale_and_shift_inverse(self, zhat: torch.Tensor) -> torch.Tensor:
+        half_width = self._levels // 2
+        return (zhat - half_width) / half_width
+
+    def codes_to_indices(self, zhat: torch.Tensor) -> torch.Tensor:
+        """Converts a `code` to an index in the codebook."""
+        assert zhat.shape[-1] == self.codebook_dim
+        zhat = self._scale_and_shift(zhat).float()
+        return (zhat * self._basis).sum(dim=-1).to(torch.int32)
+
+    def indices_to_codes(self, indices: torch.Tensor, project_out=True) -> torch.Tensor:
+        """Inverse of `codes_to_indices`."""
+        is_img_or_video = indices.ndim >= (3 + int(self.keep_num_codebooks_dim))
+        indices = rearrange(indices, "... -> ... 1")
+        codes_non_centered = (indices // self._basis) % self._levels
+        codes = self._scale_and_shift_inverse(codes_non_centered)
+
+        if self.keep_num_codebooks_dim:
+            codes = rearrange(codes, "... c d -> ... (c d)")
+
+        if project_out:
+            codes = self.project_out(codes)
+
+        if is_img_or_video:
+            codes = rearrange(codes, "b ... d -> b d ...")
+
+        return codes.to(self.dtype)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        """
+        einstein notation
+        b - batch
+        n - sequence (or flattened spatial dimensions)
+        d - feature dimension, which is also log2(codebook size)
+        c - number of codebook dim
+        """
+        is_img_or_video = z.ndim >= 4
+
+        # standardize image or video into (batch, seq, dimension)
+
+        if is_img_or_video:
+            z = rearrange(z, "b d ... -> b ... d")
+            z, ps = pack_one(z, "b * d")
+
+        assert z.shape[-1] == self.dim, f"expected dimension of {self.dim} but found dimension of {z.shape[-1]}"
+
+        z = self.project_in(z)
+
+        z = rearrange(z, "b n (c d) -> b n c d", c=self.num_codebooks)
+
+        codes = self.quantize(z)
+        indices = self.codes_to_indices(codes)
+
+        codes = rearrange(codes, "b n c d -> b n (c d)")
+
+        out = self.project_out(codes)
+
+        # reconstitute image or video dimensions
+
+        if is_img_or_video:
+            out = unpack_one(out, ps, "b * d")
+            out = rearrange(out, "b ... d -> b d ...")
+            indices = unpack_one(indices, ps, "b * c")
+            dummy_loss = torch.zeros_like(out.mean(dim=[1, 2, 3], keepdim=True))
+        else:
+            dummy_loss = torch.zeros_like(out.mean(dim=[1, 2], keepdim=True)).unsqueeze(1)
+
+        if not self.keep_num_codebooks_dim:
+            indices = rearrange(indices, "... 1 -> ...")
+
+        return (indices, out.to(self.dtype), dummy_loss)

ar_tokenizer_text_tokenizer.py

@@ -0,0 +1,317 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import AutoTokenizer
+
+from .log import log
+
+
+def get_tokenizer_path(model_family: str, is_instruct_model: bool = False):
+    """
+    Get the tokenizer path from the model family and instruct model flag.
+    Args:
+        model_family (str): The model family.
+        is_instruct_model (bool): Whether the model is an instruct model.
+    Returns:
+        str: The tokenizer path in s3.
+    """
+    model_family = model_family.lower()
+    if model_family == "mistral":
+        return "mistralai/Mistral-Nemo-Instruct-2407"
+    else:
+        assert model_family in ["llama3", "llama3.1"]
+        if model_family == "llama3":
+            model_path = "meta-llama/Meta-Llama-3-8B"
+        elif model_family == "llama3.1":
+            model_path = "meta-llama/Llama-3.1-8B"
+        else:
+            raise ValueError(f"Unsupported model family: {model_family}")
+        suffix = "-Instruct" if is_instruct_model else ""
+        model_path = f"{model_path}{suffix}"
+        return model_path
+
+
+class TextTokenizer:
+    """
+    Text tokenizer class built on HuggingFace's Fast Tokenizer (Rust based).
+    """
+
+    def __init__(
+        self,
+        model_family: str,
+        is_instruct_model: bool,
+        local_path: Optional[str] = None,
+    ):
+        """
+        Initialize the TextTokenizer.
+        Args:
+            model_family (str): The model family.
+            is_instruct_model (bool): Whether the model is an instruct model.
+            local_path (Optional[str]): The local path to the tokenizer. If not provided, the tokenizer will be downloaded from the remote path.
+        """
+        if local_path is None:
+            tokenizer_path = get_tokenizer_path(model_family, is_instruct_model)
+        else:
+            tokenizer_path = local_path
+
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=True)
+        self.stop_tokens = {
+            self.tokenizer.eos_token_id,
+        }
+        self.model_family = model_family
+        self.is_instruct_model = is_instruct_model
+        self.eos_id = self.tokenizer.eos_token_id
+        if self.tokenizer.pad_token is None:
+            if model_family.startswith("llama"):
+                self.pad_id = 128004  # "<|finetune_right_pad_id|>"
+            elif model_family == "mistral":
+                self.pad_id = 10  # "<pad>"
+            elif model_family == "pixtral":
+                self.pad_id = 11  # "<pad>"
+            else:
+                raise ValueError(f"pad_id not defined for model_family {model_family}")
+        else:
+            self.pad_id = self.tokenizer.pad_token_id
+
+    def tokenize(self, text: str, *, add_special_tokens: bool = False, **kwargs) -> List[str]:
+        """
+        Converts a string into a sequence of tokens, replacing unknown tokens with the `unk_token`.
+
+        Args:
+            text (`str`):
+                The sequence to be encoded.
+            add_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not to add the special tokens associated with the corresponding model.
+        Returns:
+            `List[str]`: The list of tokens.
+        """
+        return self.tokenizer.tokenize(text, add_special_tokens=add_special_tokens, **kwargs)
+
+    def encode(
+        self,
+        text: Union[str, List[str], List[int]],
+        *,  # Enforce keyword-only arguments
+        add_special_tokens: bool = True,
+        padding: Union[bool, str] = False,
+        truncation: Union[bool, str] = None,
+        max_length: Optional[int] = None,
+        stride: int = 0,
+        return_tensors: Optional[str] = None,
+        **kwargs,
+    ) -> List[int]:
+        """
+        Converts a string to a sequence of ids (integer), using the tokenizer and vocabulary.
+
+        Args:
+            text (`str`, `List[str]` or `List[int]`):
+                The first sequence to be encoded. This can be a string, a list of strings (tokenized string using the
+                `tokenize` method) or a list of integers (tokenized string ids using the `convert_tokens_to_ids`
+                method).
+            add_special_tokens (`bool`, *optional*, defaults to `True`):
+                Whether or not to add special tokens when encoding the sequences. This will use the underlying
+                `PretrainedTokenizerBase.build_inputs_with_special_tokens` function, which defines which tokens are
+                automatically added to the input ids. This is usefull if you want to add `bos` or `eos` tokens
+                automatically.
+            padding (`bool`, `str`, *optional*, defaults to `False`):
+                Activates and controls padding. Accepts the following values:
+
+                - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+                  sequence if provided).
+                - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
+                  acceptable input length for the model if that argument is not provided.
+                - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
+                  lengths).
+            truncation (`bool`, `str`, *optional*, defaults to `False`):
+                Activates and controls truncation. Accepts the following values:
+
+                - `True` or `'longest_first'`: Truncate to a maximum length specified with the argument `max_length` or
+                  to the maximum acceptable input length for the model if that argument is not provided. This will
+                  truncate token by token, removing a token from the longest sequence in the pair if a pair of
+                  sequences (or a batch of pairs) is provided.
+                - `'only_first'`: Truncate to a maximum length specified with the argument `max_length` or to the
+                  maximum acceptable input length for the model if that argument is not provided. This will only
+                  truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
+                - `'only_second'`: Truncate to a maximum length specified with the argument `max_length` or to the
+                  maximum acceptable input length for the model if that argument is not provided. This will only
+                  truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
+                - `False` or `'do_not_truncate'` (default): No truncation (i.e., can output batch with sequence lengths
+                  greater than the model maximum admissible input size).
+            max_length (`int`, *optional*):
+                Controls the maximum length to use by one of the truncation/padding parameters.
+
+                If left unset or set to `None`, this will use the predefined model maximum length if a maximum length
+                is required by one of the truncation/padding parameters. If the model has no specific maximum input
+                length (like XLNet) truncation/padding to a maximum length will be deactivated.
+            stride (`int`, *optional*, defaults to 0):
+                If set to a number along with `max_length`, the overflowing tokens returned when
+                `return_overflowing_tokens=True` will contain some tokens from the end of the truncated sequence
+                returned to provide some overlap between truncated and overflowing sequences. The value of this
+                argument defines the number of overlapping tokens.
+            is_split_into_words (`bool`, *optional*, defaults to `False`):
+                Whether or not the input is already pre-tokenized (e.g., split into words). If set to `True`, the
+                tokenizer assumes the input is already split into words (for instance, by splitting it on whitespace)
+                which it will tokenize. This is useful for NER or token classification.
+            pad_to_multiple_of (`int`, *optional*):
+                If set will pad the sequence to a multiple of the provided value. Requires `padding` to be activated.
+                This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
+                `>= 7.5` (Volta).
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors instead of list of python integers. Acceptable values are:
+
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return Numpy `np.ndarray` objects.
+        """
+        return self.tokenizer.encode(
+            text,
+            add_special_tokens=add_special_tokens,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            stride=stride,
+            return_tensors=return_tensors,
+        )
+
+    def decode(
+        self,
+        token_ids: Union[int, List[int], "np.ndarray", "torch.Tensor"],
+        *,  # Enforce keyword-only arguments
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool = None,
+        **kwargs,
+    ) -> str:
+        """
+        Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special
+        tokens and clean up tokenization spaces.
+
+        Args:
+            token_ids (`Union[int, List[int], np.ndarray, torch.Tensor, tf.Tensor]`):
+                List of tokenized input ids. Can be obtained using the `__call__` method.
+            skip_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not to remove special tokens in the decoding.
+            clean_up_tokenization_spaces (`bool`, *optional*):
+                Whether or not to clean up the tokenization spaces. If `None`, will default to
+                `self.clean_up_tokenization_spaces`.
+            kwargs (additional keyword arguments, *optional*):
+                Will be passed to the underlying model specific decode method.
+
+        Returns:
+            `str`: The decoded sentence.
+        """
+        return self.tokenizer.decode(
+            token_ids,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+    def apply_chat_template(
+        self,
+        conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]]],
+        *,
+        add_generation_prompt: bool = False,
+        tokenize: bool = True,
+        padding: bool = False,
+        truncation: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        return_dict: bool = False,
+        return_assistant_tokens_mask: bool = False,
+        generation_prefix: str = "",
+        tokenizer_kwargs: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        """
+        Converts a list of dictionaries with `"role"` and `"content"` keys to a list of token
+        ids. This method is intended for use with chat models, and will read the tokenizer's chat_template attribute to determine the format and control tokens to use when converting.
+
+        More details can be found at https://huggingface.co/docs/transformers/main/en/internal/tokenization_utils#transformers.PreTrainedTokenizerBase.apply_chat_template
+
+        Args:
+            conversation (Union[List[Dict[str, str]], List[List[Dict[str, str]]]]): A list of dicts
+                with "role" and "content" keys, representing the chat history so far.
+            add_generation_prompt (bool, *optional*):
+                If this is set, a prompt with the token(s) that indicate
+                the start of an assistant message will be appended to the formatted output. This is useful when you want to generate a response from the model.
+                Note that this argument will be passed to the chat template, and so it must be supported in the
+                template for this argument to have any effect.
+            continue_final_message (bool, *optional*):
+                If this is set, the chat will be formatted so that the final
+                message in the chat is open-ended, without any EOS tokens. The model will continue this message
+                rather than starting a new one. This allows you to "prefill" part of
+                the model's response for it. Cannot be used at the same time as `add_generation_prompt`.
+            tokenize (`bool`, defaults to `True`):
+                Whether to tokenize the output. If `False`, the output will be a string.
+            padding (`bool`, defaults to `False`):
+                Whether to pad sequences to the maximum length. Has no effect if tokenize is `False`.
+            truncation (`bool`, defaults to `False`):
+                Whether to truncate sequences at the maximum length. Has no effect if tokenize is `False`.
+            max_length (`int`, *optional*):
+                Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is `False`. If
+                not specified, the tokenizer's `max_length` attribute will be used as a default.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Has no effect if tokenize is `False`. Acceptable
+                values are:
+                - `'tf'`: Return TensorFlow `tf.Tensor` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+            return_dict (`bool`, defaults to `False`):
+                Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`.
+            generation_prefix (str): Prefix to add before asking model to generate. Helpful to guide the generation. Defaults to "".
+            tokenizer_kwargs (`Dict[str: Any]`, *optional*): Additional kwargs to pass to the tokenizer.
+            return_assistant_tokens_mask (`bool`, defaults to `False`):
+                Whether to return a mask of the assistant generated tokens. For tokens generated by the assistant,
+                the mask will contain 1. For user and system tokens, the mask will contain 0.
+                This functionality is only available for chat templates that support it via the `{% generation %}` keyword.
+            **kwargs: Additional kwargs to pass to the template renderer. Will be accessible by the chat template.
+
+        Returns:
+            `Union[List[int], Dict]`: A list of token ids representing the tokenized chat so far, including control tokens. This
+            output is ready to pass to the model, either directly or via methods like `generate()`. If `return_dict` is
+            set, will return a dict of tokenizer outputs instead.
+        """
+        if not self.is_instruct_model:
+            raise ValueError(
+                "apply_chat_template is only supported for instruct models. You should pass argument is_instruct_model=True to the TextTokenizer constructor."
+            )
+        # Since generation_prefix is added to the text in the end, ensure that the setting is correct
+        if generation_prefix:
+            assert not tokenize, "tokenize must be False when generation_prefix is provided."
+            assert add_generation_prompt, "add_generation_prompt must be set when generation_prefix is provided."
+        formatted_text: Union[str, List[int]] = self.tokenizer.apply_chat_template(
+            conversation,
+            add_generation_prompt=add_generation_prompt,
+            tokenize=tokenize,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            return_dict=return_dict,
+            return_assistant_tokens_mask=return_assistant_tokens_mask,
+            tokenizer_kwargs=tokenizer_kwargs,
+            **kwargs,
+        )
+        if generation_prefix:
+            formatted_text: str = formatted_text + generation_prefix
+            log.debug(
+                f"Adding generation prefix: {generation_prefix} to the formatted text\n"
+                f"Formatted text: {formatted_text}"
+            )
+        return formatted_text

ar_tokenizer_tokenizer.py

@@ -0,0 +1,322 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import defaultdict
+from typing import Optional
+
+import torch
+from einops import rearrange
+
+from .ar_config_base_tokenizer import TokenizerConfig
+from .lazy_config_init import instantiate as lazy_instantiate
+
+
+def update_vocab_size(
+    existing_vocab_size,
+    to_be_added_vocab_size,
+    training_type,
+    add_special_tokens,
+    video_special_tokens={},
+):
+    # New vocab size
+    if add_special_tokens:
+        existing_vocab_size += to_be_added_vocab_size + len(video_special_tokens)
+    # For text_to_video, we add one <bov> special token at the beginning of the video
+    elif training_type == "text_to_video":
+        existing_vocab_size += to_be_added_vocab_size + 1
+    else:
+        existing_vocab_size += to_be_added_vocab_size
+    return existing_vocab_size
+
+
+class DiscreteMultimodalTokenizer:
+    def __init__(self, tokenizer_config: TokenizerConfig):
+        self.tokenizer_config = tokenizer_config
+        self.vocab_size = 0
+        self.total_seq_len = tokenizer_config.seq_len
+        self.pad_to_multiple_of = tokenizer_config.pad_to_multiple_of
+        self.training_type = tokenizer_config.training_type
+        assert self.training_type in [
+            "text_only",
+            "text_to_video",
+            "video_to_video",
+            "image_text_interleaved",
+        ], f"{self.training_type} not supported"
+
+        self._build_text_tokenizer()
+        self._build_video_tokenizer()
+
+    def _build_text_tokenizer(self):
+        r"""Function to initialize the text tokenizer model."""
+        if self.tokenizer_config.text_tokenizer is not None:
+            self.text_tokenizer = lazy_instantiate(self.tokenizer_config.text_tokenizer.config)
+            self.vocab_size += self.tokenizer_config.text_tokenizer.vocab_size
+        else:
+            self.text_tokenizer = None
+
+    def _build_video_tokenizer(self):
+        r"""Function to initialize the video tokenizer model."""
+        if self.tokenizer_config.video_tokenizer is not None:
+            self.video_tokenizer = lazy_instantiate(self.tokenizer_config.video_tokenizer.config)
+            self.video_tokenizer = self.video_tokenizer.to("cuda")
+            self.video_vocab_size = self.tokenizer_config.video_tokenizer.vocab_size
+            special_token_offset = (
+                self.tokenizer_config.video_tokenizer.tokenizer_offset
+                + self.tokenizer_config.video_tokenizer.vocab_size
+            )
+            self.video_special_tokens = {
+                "<|begin_of_video|>": special_token_offset,
+                "<|end_of_video|>": special_token_offset + 1,
+                "<|pad_token_video|>": special_token_offset + 2,
+            }
+
+            self.vocab_size = update_vocab_size(
+                existing_vocab_size=self.vocab_size,
+                to_be_added_vocab_size=self.tokenizer_config.video_tokenizer.vocab_size,
+                training_type=self.training_type,
+                add_special_tokens=self.tokenizer_config.add_special_tokens,
+                video_special_tokens=self.video_special_tokens,
+            )
+        else:
+            self.video_tokenizer = None
+
+    @property
+    def pad_id(self):
+        r"""Returns the pad_id."""
+
+        if self.training_type == "text_only" or self.training_type == "image_text_interleaved":
+            pad_id = self.text_tokenizer.pad_id
+        elif self.training_type in ["text_to_video", "video_to_video"]:
+            pad_id = self.video_special_tokens["<|pad_token_video|>"]
+        else:
+            raise ValueError(f"training_type {self.training_type} not defined")
+        return pad_id
+
+    @property
+    def ignore_index(self):
+        r"""Returns which token should be ignored during loss computation."""
+        if self.training_type == "text_only" or self.training_type == "image_text_interleaved":
+            if self.text_tokenizer.pad_id == self.text_tokenizer.eos_id:
+                # If the PAD token is the same as the EOS token, we do not ignore it during loss
+                # computation, since we want the model to be able to predict EOS tokens in inference.
+                # The PyTorch default ignore_index for the cross-entropy loss is -100.
+                ignore_index = -100
+            else:
+                ignore_index = self.text_tokenizer.pad_id
+        elif self.training_type in ["text_to_video", "video_to_video"]:
+            ignore_index = self.pad_id
+        else:
+            raise ValueError(f"training_type {self.training_type} not defined")
+        return ignore_index
+
+    @property
+    def stop_tokens(self):
+        r"""Returns the stop tokens."""
+        if self.training_type == "text_only" or self.training_type == "image_text_interleaved":
+            stop_tokens = self.text_tokenizer.stop_tokens
+        elif self.training_type in ["text_to_video", "video_to_video"]:
+            stop_tokens = set([self.video_special_tokens["<|end_of_video|>"]])
+        else:
+            raise ValueError(f"training_type {self.training_type} not defined")
+        return stop_tokens
+
+    def _tokenize_text(self, raw_text: list[str], max_text_seq_len: int = -1):
+        r"""Function to tokenize text.
+        Args:
+            raw_text (list[str]): List of input strings
+            max_text_seq_len (int): Maximum sequence length returned by text tokenizer
+        Returns:
+            text_tokens (list[list[int]]): List of text tokens
+        """
+
+        batch_size = len(raw_text)
+        text_tokens = [self.text_tokenizer.encode(raw_text[i], bos=True, eos=True) for i in range(batch_size)]
+
+        # Clipping the text tokens so that the sequence length does not exceed max_text_seq_len
+        if max_text_seq_len > -1:
+            for i in range(len(text_tokens)):
+                if len(text_tokens[i]) > max_text_seq_len:
+                    # Simply clip and add end of seq token
+                    text_tokens[i] = text_tokens[i][0 : max_text_seq_len - 1] + [self.text_tokenizer.eos_id]
+        return text_tokens
+
+    def _tokenize_class(self, cls_labels: list[str]):
+        r"""Function to tokenize the class label.
+        Args:
+            cls_labels (list[str]): List of class indices
+        Returns:
+            class_tokens (list[list[int]]): List of class tokens
+        """
+
+        # tokenizer_offset tells what offset should be added to the tokens.
+        # This is needed for vocab expansion.
+        class_tokens = [[int(x) + self.tokenizer_config.class_tokenizer.tokenizer_offset] for x in cls_labels]
+
+        return class_tokens
+
+    def _tokenize_video(self, videos: torch.Tensor, pixel_chunk_duration: Optional[int] = None):
+        r"""Function to tokenize video.
+        Args:
+            videos (torch.Tensor): Input video data tensor
+            pixel_chunk_duration (Optional[float]): Pixel chunk duration. If provided, we pass it to the video tokenizer.
+        Returns:
+            video_tokens (list[list[int]]): List of video tokens
+        """
+
+        video_tokens = []
+        batch_size = videos.shape[0]
+
+        quantized_out, _ = self.video_tokenizer.encode(videos, pixel_chunk_duration=pixel_chunk_duration)
+        indices = self.video_tokenizer.fsq_quantizer.codes_to_indices(quantized_out.permute(0, 2, 3, 4, 1))
+
+        # Flatten the indices
+        indices = rearrange(indices, "B T H W -> B (T H W)")
+
+        # tokenizer_offset tells what offset should be added to the tokens.
+        # This is needed for vocab expansion.
+        indices += self.tokenizer_config.video_tokenizer.tokenizer_offset
+
+        # Add begin and end of video tokens
+        bov_token = self.video_special_tokens["<|begin_of_video|>"]
+        eov_token = self.video_special_tokens["<|end_of_video|>"]
+
+        # Append bov and eov tokens
+        if self.tokenizer_config.add_special_tokens:
+            for i in range(batch_size):
+                video_tokens.append([bov_token] + indices[i].tolist() + [eov_token])
+        else:
+            if self.training_type == "text_to_video":
+                for i in range(batch_size):
+                    video_tokens.append([bov_token] + indices[i].tolist())
+            else:
+                for i in range(batch_size):
+                    video_tokens.append(indices[i].tolist())
+                    assert (
+                        len(video_tokens[-1]) == self.tokenizer_config.video_tokenizer.max_seq_len
+                    ), f"Expected {self.tokenizer_config.video_tokenizer.max_seq_len} tokens, got {len(video_tokens[-1])}; video shape: {videos.shape}"
+
+        return video_tokens
+
+    def tokenize(self, data_batch: dict):
+        r"""Function to tokenize data_dict.
+        Args:
+            data_batch (dict): Input data dict
+        Returns:
+            tokens (torch.LongTensor): Token tensor dict
+        """
+
+        if (
+            self.training_type in ["text_only", "image_text_interleaved"]
+            and not self.tokenizer_config.text_tokenizer.tokenize_here
+        ):
+            # In case of pre-computed tokens, just return the data_batch
+            return data_batch["tokens"], None
+
+        # Online tokenization
+        tokens = []
+        token_boundaries = defaultdict(list)
+
+        # Obtain maximum sequence length
+        max_text_seq_len = -1
+        max_visual_seq_len = -1
+
+        if self.training_type in ["text_to_video", "video_to_video"]:
+            max_visual_seq_len = self.tokenizer_config.video_tokenizer.max_seq_len
+
+        # If max visual sequence length is specified, make sure that text is clipped so that
+        # the full video/image is always seen.
+        if max_visual_seq_len > -1:
+            if self.tokenizer_config.add_special_tokens:
+                max_visual_seq_len = max_visual_seq_len + 2  # Two special tokens is for [bov, eov] or [boi, eoi] token
+            elif self.training_type == "text_to_video":
+                max_visual_seq_len = max_visual_seq_len + 1
+            else:
+                max_visual_seq_len = max_visual_seq_len
+            assert (
+                max_visual_seq_len <= self.total_seq_len
+            ), f"max_visual_seq_len ({max_visual_seq_len}) is greater that total sequence length ({self.total_seq_len})"
+            max_text_seq_len = self.total_seq_len - max_visual_seq_len
+
+        # Tokenize the text
+        if (
+            "text" in self.training_type
+            and self.text_tokenizer is not None
+            and self.tokenizer_config.text_tokenizer.tokenize_here
+        ):
+            key = self.tokenizer_config.text_tokenizer.data_key
+            batch_size = len(data_batch[key])
+            assert key in data_batch, f"Key {key} should be present in data for text tokenizer"
+            tokens = self._tokenize_text(data_batch["caption"], max_text_seq_len)
+
+            for i in range(batch_size):
+                token_boundaries["text"].append((0, len(tokens[i])))
+        else:
+            tokens = []
+            batch_size = None
+
+        # Tokenize the class label
+        if "class" in self.training_type and self.tokenizer_config.class_tokenizer is not None:
+            key = self.tokenizer_config.class_tokenizer.data_key
+            assert key in data_batch, f"Key {key} should be present in data for class tokenizer"
+            batch_size = len(data_batch[key]) if batch_size is None else batch_size
+            tokens_class = self._tokenize_class(data_batch[key])
+            if len(tokens) == 0:
+                tokens = tokens_class
+                for i in range(batch_size):
+                    token_boundaries["class"].append((0, len(tokens[i])))
+            else:
+                for i in range(batch_size):
+                    token_boundaries["class"].append((len(tokens[i]), len(tokens[i]) + len(tokens_class[i])))
+                    tokens[i] = tokens[i] + tokens_class[i]
+
+        # Tokenize the video
+        if self.video_tokenizer is not None and self.tokenizer_config.video_tokenizer.tokenize_here:
+            key = self.tokenizer_config.video_tokenizer.data_key
+            assert key in data_batch, f"Key {key} should be present in data for video tokenizer"
+            batch_size = len(data_batch[key]) if batch_size is None else batch_size
+
+            pixel_chunk_duration = (
+                None  # If not specified, we assume it's a video dataset and use the default chunk duration
+            )
+            dataset_name = data_batch.get("dataset_name", None)
+            if dataset_name is not None and dataset_name.startswith("image"):
+                # If it's an image dataset, we use a pixel chunk duration of 1
+                pixel_chunk_duration = 1
+            tokens_video = self._tokenize_video(data_batch[key], pixel_chunk_duration=pixel_chunk_duration)
+            if len(tokens) == 0:
+                tokens = tokens_video
+                for i in range(batch_size):
+                    token_boundaries["video"].append((0, len(tokens[i])))
+                    # [B,] each entry is ((0, len(tokens[i])))
+            else:
+                for i in range(batch_size):
+                    token_boundaries["video"].append((len(tokens[i]), len(tokens[i]) + len(tokens_video[i])))
+                    tokens[i] = tokens[i] + tokens_video[i]
+
+        # Combine the tokens and do padding
+        max_seq_len_in_batch = max([len(token) for token in tokens])
+        if self.pad_to_multiple_of is not None:
+            # Pad the sequence length to the nearest multiple of pad_to_multiple_of
+            max_seq_len_in_batch = ((max_seq_len_in_batch - 1) // self.pad_to_multiple_of + 1) * self.pad_to_multiple_of
+        pad_to_len = min(max_seq_len_in_batch, self.total_seq_len)
+        for i in range(len(tokens)):
+            if len(tokens[i]) < pad_to_len:
+                tokens[i] = tokens[i] + [self.pad_id] * (pad_to_len - len(tokens[i]))
+            else:
+                tokens[i] = tokens[i][0:pad_to_len]
+
+        # Convert it to long tensor
+        tokens = torch.LongTensor(tokens)
+        return tokens, token_boundaries

ar_tokenizer_utils.py

@@ -0,0 +1,101 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+from einops import pack, rearrange, unpack
+
+
+def time2batch(x: torch.Tensor) -> tuple[torch.Tensor, int]:
+    batch_size = x.shape[0]
+    return rearrange(x, "b c t h w -> (b t) c h w"), batch_size
+
+
+def batch2time(x: torch.Tensor, batch_size: int) -> torch.Tensor:
+    return rearrange(x, "(b t) c h w -> b c t h w", b=batch_size)
+
+
+def space2batch(x: torch.Tensor) -> tuple[torch.Tensor, int]:
+    batch_size, height = x.shape[0], x.shape[-2]
+    return rearrange(x, "b c t h w -> (b h w) c t"), batch_size, height
+
+
+def batch2space(x: torch.Tensor, batch_size: int, height: int) -> torch.Tensor:
+    return rearrange(x, "(b h w) c t -> b c t h w", b=batch_size, h=height)
+
+
+def cast_tuple(t: Any, length: int = 1) -> Any:
+    return t if isinstance(t, tuple) else ((t,) * length)
+
+
+def replication_pad(x):
+    return torch.cat([x[:, :, :1, ...], x], dim=2)
+
+
+def divisible_by(num: int, den: int) -> bool:
+    return (num % den) == 0
+
+
+def is_odd(n: int) -> bool:
+    return not divisible_by(n, 2)
+
+
+def nonlinearity(x):
+    return x * torch.sigmoid(x)
+
+
+class CausalNormalize(torch.nn.Module):
+    def __init__(self, in_channels, num_groups=1):
+        super().__init__()
+        self.norm = torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+        self.num_groups = num_groups
+
+    def forward(self, x):
+        # if num_groups !=1, we apply a spatio-temporal groupnorm for backward compatibility purpose.
+        # All new models should use num_groups=1, otherwise causality is not guaranteed.
+        if self.num_groups == 1:
+            x, batch_size = time2batch(x)
+            return batch2time(self.norm(x), batch_size)
+        return self.norm(x)
+
+
+def exists(v):
+    return v is not None
+
+
+def default(*args):
+    for arg in args:
+        if exists(arg):
+            return arg
+    return None
+
+
+def pack_one(t, pattern):
+    return pack([t], pattern)
+
+
+def unpack_one(t, ps, pattern):
+    return unpack(t, ps, pattern)[0]
+
+
+def round_ste(z: torch.Tensor) -> torch.Tensor:
+    """Round with straight through gradients."""
+    zhat = z.round()
+    return z + (zhat - z).detach()
+
+
+def log(t, eps=1e-5):
+    return t.clamp(min=eps).log()

ar_utils_checkpoint.py

@@ -0,0 +1,76 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, Optional
+
+import torch
+
+# Substrings to ignore when processing state dicts
+substrings_to_ignore = [
+    "_extra_state",  # Extra states (BytesIO type) added by TransformerEngine for FP8 handling
+]
+
+
+def get_partial_state_dict(
+    state_dict: Dict[str, torch.Tensor],
+    prefix: str,
+) -> Dict[str, torch.Tensor]:
+    """
+    Get a partial state dict with keys starting with the given prefix
+    """
+    return {k: v for k, v in state_dict.items() if k.startswith(prefix)}
+
+
+def process_state_dict(
+    state_dict: Dict[str, torch.Tensor],
+    device: str = None,
+    dtype: torch.dtype = None,
+    prefix_to_remove: Optional[str] = None,
+) -> Dict[str, torch.Tensor]:
+    """
+    - Remove items with substring "_extra_state" in keys (TransformerEngine adds these for FP8)
+    - Move tensors to specified device and dtype if provided
+
+    Args:
+        state_dict (Dict[str, torch.Tensor]): The state dict to process
+        device (str, optional): The device to move tensors to. Defaults to None.
+        dtype (torch.dtype, optional): The dtype to move tensors to. Defaults to None.
+        prefix_to_remove (str, optional): The prefix to remove from the keys of the state dict. Defaults to None.
+
+    Returns:
+        Dict[str, torch.Tensor]: The processed state dict
+    """
+    new_state_dict = {}
+    tensor_kwargs = {}
+    if device is not None:
+        tensor_kwargs["device"] = device
+    if dtype is not None:
+        tensor_kwargs["dtype"] = dtype
+
+    for key, value in state_dict.items():
+        # Check if any of the substrings to ignore are in the key
+        skip = False
+        for substr in substrings_to_ignore:
+            if substr in key:
+                skip = True
+                break
+        if skip:
+            continue
+        if len(tensor_kwargs) > 0:
+            value = value.to(**tensor_kwargs)
+        if prefix_to_remove is not None and key.startswith(prefix_to_remove):
+            key = key[len(prefix_to_remove) :]
+        new_state_dict[key] = value
+    return new_state_dict

ar_utils_inference.py

@@ -0,0 +1,360 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import json
+import math
+import os
+from pathlib import Path
+from typing import List
+
+import numpy as np
+import torch
+import torchvision
+from PIL import Image
+
+from .ar_config_inference_inference_config import SamplingConfig
+from .log import log
+
+_IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", "webp"]
+_VIDEO_EXTENSIONS = [".mp4"]
+_SUPPORTED_CONTEXT_LEN = [1, 9]  # Input frames
+NUM_TOTAL_FRAMES = 33
+
+
+def add_common_arguments(parser):
+    """Add common command line arguments.
+
+    Args:
+        parser (ArgumentParser): Argument parser to add arguments to
+    """
+    parser.add_argument(
+        "--checkpoint_dir", type=str, default="checkpoints", help="Base directory containing model checkpoints"
+    )
+    parser.add_argument(
+        "--video_save_name",
+        type=str,
+        default="output",
+        help="Output filename for generating a single video",
+    )
+    parser.add_argument("--video_save_folder", type=str, default="outputs/", help="Output folder for saving videos")
+    parser.add_argument(
+        "--input_image_or_video_path",
+        type=str,
+        help="Input path for input image or video",
+    )
+    parser.add_argument(
+        "--batch_input_path",
+        type=str,
+        help="Input folder containing all input images or videos",
+    )
+    parser.add_argument(
+        "--num_input_frames",
+        type=int,
+        default=9,
+        help="Number of input frames for world generation",
+        choices=_SUPPORTED_CONTEXT_LEN,
+    )
+    parser.add_argument("--temperature", type=float, default=1.0, help="Temperature for sampling")
+    parser.add_argument("--top_p", type=float, default=0.8, help="Top-p value for sampling")
+    parser.add_argument("--seed", type=int, default=0, help="Random seed")
+    parser.add_argument("--disable_diffusion_decoder", action="store_true", help="Disable diffusion decoder")
+    parser.add_argument(
+        "--offload_guardrail_models",
+        action="store_true",
+        help="Offload guardrail models after inference",
+    )
+    parser.add_argument(
+        "--offload_diffusion_decoder",
+        action="store_true",
+        help="Offload diffusion decoder after inference",
+    )
+    parser.add_argument(
+        "--offload_ar_model",
+        action="store_true",
+        help="Offload AR model after inference",
+    )
+    parser.add_argument(
+        "--offload_tokenizer",
+        action="store_true",
+        help="Offload discrete tokenizer model after inference",
+    )
+
+
+def validate_args(args: argparse.Namespace, inference_type: str):
+    """Validate command line arguments for base and video2world generation."""
+    assert inference_type in [
+        "base",
+        "video2world",
+    ], "Invalid inference_type, must be 'base' or 'video2world'"
+    if args.input_type in ["image", "text_and_image"] and args.num_input_frames != 1:
+        args.num_input_frames = 1
+        log.info(f"Set num_input_frames to 1 for {args.input_type} input")
+
+    if args.num_input_frames == 1:
+        if "4B" in args.ar_model_dir:
+            log.warning(
+                "The failure rate for 4B model with image input is ~15%. 12B / 13B model have a smaller failure rate. Please be cautious and refer to README.md for more details."
+            )
+        elif "5B" in args.ar_model_dir:
+            log.warning(
+                "The failure rate for 5B model with image input is ~7%. 12B / 13B model have a smaller failure rate. Please be cautious and refer to README.md for more details."
+            )
+
+    # Validate prompt/image/video args for single or batch generation
+    assert (
+        args.input_image_or_video_path or args.batch_input_path
+    ), "--input_image_or_video_path or --batch_input_path must be provided."
+    if inference_type == "video2world" and (not args.batch_input_path):
+        assert args.prompt, "--prompt is required for single video generation."
+    args.data_resolution = [640, 1024]
+
+    # Validate number of GPUs
+    num_gpus = int(os.getenv("WORLD_SIZE", 1))
+    assert num_gpus <= 1, "We support only single GPU inference for now"
+
+    # Create output folder
+    Path(args.video_save_folder).mkdir(parents=True, exist_ok=True)
+
+    sampling_config = SamplingConfig(
+        echo=True,
+        temperature=args.temperature,
+        top_p=args.top_p,
+        compile_sampling=True,
+    )
+    return sampling_config
+
+
+def resize_input(video: torch.Tensor, resolution: list[int]):
+    r"""
+    Function to perform aspect ratio preserving resizing and center cropping.
+    This is needed to make the video into target resolution.
+    Args:
+        video (torch.Tensor): Input video tensor
+        resolution (list[int]): Data resolution
+    Returns:
+        Cropped video
+    """
+
+    orig_h, orig_w = video.shape[2], video.shape[3]
+    target_h, target_w = resolution
+
+    scaling_ratio = max((target_w / orig_w), (target_h / orig_h))
+    resizing_shape = (int(math.ceil(scaling_ratio * orig_h)), int(math.ceil(scaling_ratio * orig_w)))
+    video_resized = torchvision.transforms.functional.resize(video, resizing_shape)
+    video_cropped = torchvision.transforms.functional.center_crop(video_resized, resolution)
+    return video_cropped
+
+
+def load_image_from_list(flist, data_resolution: List[int]) -> dict:
+    """
+    Function to load images from a list of image paths.
+    Args:
+        flist (List[str]): List of image paths
+        data_resolution (List[int]): Data resolution
+    Returns:
+        Dict containing input images
+    """
+    all_videos = dict()
+    for img_path in flist:
+        ext = os.path.splitext(img_path)[1]
+        if ext in _IMAGE_EXTENSIONS:
+            # Read the image
+            img = Image.open(img_path)
+
+            # Convert to tensor
+            img = torchvision.transforms.functional.to_tensor(img)
+            static_vid = img.unsqueeze(0).repeat(NUM_TOTAL_FRAMES, 1, 1, 1)
+            static_vid = static_vid * 2 - 1
+
+            log.debug(
+                f"Resizing input image of shape ({static_vid.shape[2]}, {static_vid.shape[3]}) -> ({data_resolution[0]}, {data_resolution[1]})"
+            )
+            static_vid = resize_input(static_vid, data_resolution)
+            fname = os.path.basename(img_path)
+            all_videos[fname] = static_vid.transpose(0, 1).unsqueeze(0)
+
+    return all_videos
+
+
+def read_input_images(batch_input_path: str, data_resolution: List[int]) -> dict:
+    """
+    Function to read input images from a JSONL file.
+
+    Args:
+        batch_input_path (str): Path to JSONL file containing visual input paths
+        data_resolution (list[int]): Data resolution
+
+    Returns:
+        Dict containing input images
+    """
+    # Read visual inputs from JSONL
+    flist = []
+    with open(batch_input_path, "r") as f:
+        for line in f:
+            data = json.loads(line.strip())
+            flist.append(data["visual_input"])
+
+    return load_image_from_list(flist, data_resolution=data_resolution)
+
+
+def read_input_image(input_path: str, data_resolution: List[int]) -> dict:
+    """
+    Function to read input image.
+    Args:
+        input_path (str): Path to input image
+        data_resolution (List[int]): Data resolution
+    Returns:
+        Dict containing input image
+    """
+    flist = [input_path]
+    return load_image_from_list(flist, data_resolution=data_resolution)
+
+
+def read_input_videos(batch_input_path: str, data_resolution: List[int], num_input_frames: int) -> dict:
+    r"""
+    Function to read input videos.
+    Args:
+        batch_input_path (str): Path to JSONL file containing visual input paths
+        data_resolution (list[int]): Data resolution
+    Returns:
+        Dict containing input videos
+    """
+    # Read visual inputs from JSONL
+    flist = []
+    with open(batch_input_path, "r") as f:
+        for line in f:
+            data = json.loads(line.strip())
+            flist.append(data["visual_input"])
+    return load_videos_from_list(flist, data_resolution=data_resolution, num_input_frames=num_input_frames)
+
+
+def read_input_video(input_path: str, data_resolution: List[int], num_input_frames: int) -> dict:
+    """
+    Function to read input video.
+    Args:
+        input_path (str): Path to input video
+        data_resolution (List[int]): Data resolution
+        num_input_frames (int): Number of frames in context
+    Returns:
+        Dict containing input video
+    """
+    flist = [input_path]
+    return load_videos_from_list(flist, data_resolution=data_resolution, num_input_frames=num_input_frames)
+
+
+def load_videos_from_list(flist: List[str], data_resolution: List[int], num_input_frames: int) -> dict:
+    """
+    Function to load videos from a list of video paths.
+    Args:
+        flist (List[str]): List of video paths
+        data_resolution (List[int]): Data resolution
+        num_input_frames (int): Number of frames in context
+    Returns:
+        Dict containing input videos
+    """
+    all_videos = dict()
+
+    for video_path in flist:
+        ext = os.path.splitext(video_path)[-1]
+        if ext in _VIDEO_EXTENSIONS:
+            video, _, _ = torchvision.io.read_video(video_path, pts_unit="sec")
+            video = video.float() / 255.0
+            video = video * 2 - 1
+
+            # Resize the videos to the required dimension
+            nframes_in_video = video.shape[0]
+            if nframes_in_video < num_input_frames:
+                fname = os.path.basename(video_path)
+                log.warning(
+                    f"Video {fname} has {nframes_in_video} frames, less than the requried {num_input_frames} frames. Skipping."
+                )
+                continue
+
+            video = video[-num_input_frames:, :, :, :]
+
+            # Pad the video to NUM_TOTAL_FRAMES (because the tokenizer expects inputs of NUM_TOTAL_FRAMES)
+            video = torch.cat(
+                (video, video[-1, :, :, :].unsqueeze(0).repeat(NUM_TOTAL_FRAMES - num_input_frames, 1, 1, 1)),
+                dim=0,
+            )
+
+            video = video.permute(0, 3, 1, 2)
+
+            log.debug(
+                f"Resizing input video of shape ({video.shape[2]}, {video.shape[3]}) -> ({data_resolution[0]}, {data_resolution[1]})"
+            )
+            video = resize_input(video, data_resolution)
+
+            fname = os.path.basename(video_path)
+            all_videos[fname] = video.transpose(0, 1).unsqueeze(0)
+
+    return all_videos
+
+
+def load_vision_input(
+    input_type: str,
+    batch_input_path: str,
+    input_image_or_video_path: str,
+    data_resolution: List[int],
+    num_input_frames: int,
+):
+    """
+    Function to load vision input.
+    Note: We pad the frames of the input image/video to NUM_TOTAL_FRAMES here, and feed the padded video tensors to the video tokenizer to obtain tokens. The tokens will be truncated based on num_input_frames when feeding to the autoregressive model.
+    Args:
+        input_type (str): Type of input
+        batch_input_path (str): Folder containing input images or videos
+        input_image_or_video_path (str): Path to input image or video
+        data_resolution (List[int]): Data resolution
+        num_input_frames (int): Number of frames in context
+    Returns:
+        Dict containing input videos
+    """
+    if batch_input_path:
+        log.info(f"Reading batch inputs from path: {batch_input_path}")
+        if input_type == "image" or input_type == "text_and_image":
+            input_videos = read_input_images(batch_input_path, data_resolution=data_resolution)
+        elif input_type == "video" or input_type == "text_and_video":
+            input_videos = read_input_videos(
+                batch_input_path,
+                data_resolution=data_resolution,
+                num_input_frames=num_input_frames,
+            )
+        else:
+            raise ValueError(f"Invalid input type {input_type}")
+    else:
+        if input_type == "image" or input_type == "text_and_image":
+            input_videos = read_input_image(input_image_or_video_path, data_resolution=data_resolution)
+        elif input_type == "video" or input_type == "text_and_video":
+            input_videos = read_input_video(
+                input_image_or_video_path,
+                data_resolution=data_resolution,
+                num_input_frames=num_input_frames,
+            )
+        else:
+            raise ValueError(f"Invalid input type {input_type}")
+    return input_videos
+
+
+def prepare_video_batch_for_saving(video_batch: List[torch.Tensor]) -> List[np.ndarray]:
+    """
+    Function to convert output tensors to numpy format for saving.
+    Args:
+        video_batch (List[torch.Tensor]): List of output tensors
+    Returns:
+        List of numpy arrays
+    """
+    return [(video * 255).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy() for video in video_batch]

ar_utils_misc.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from omegaconf import DictConfig, OmegaConf
+
+
+class CustomSimpleNamespace:
+    """
+    A simple namespace class that supports both attribute-style and dictionary-style access.
+    """
+
+    def __init__(self, d):
+        self._d = d
+
+    def __getattr__(self, attr):
+        # Attribute-style access: config.key
+        try:
+            return self._d[attr]
+        except KeyError:
+            raise AttributeError(f"'CustomSimpleNamespace' object has no attribute '{attr}'")
+
+    def __getitem__(self, key):
+        # Dictionary-style access: config['key']
+        return self._d[key]
+
+
+def maybe_convert_to_namespace(config):
+    """
+    This function cast a OmegaConf's DictConfig or a standard dict to CustomSimpleNamespace, which supports both
+    attribute-style and dictionary-style access.
+    Note: We need to convert OmegaConf's DictConfig since it is not compatible with torch.compile.
+    """
+    # If input is OmegaConf's DictConfig, convert to a standard dict
+    if isinstance(config, DictConfig):
+        config = OmegaConf.to_container(config, resolve=True)
+
+    if isinstance(config, dict):
+        return CustomSimpleNamespace(config)
+    else:
+        return config

ar_utils_sampling.py

@@ -0,0 +1,195 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Tuple
+
+import torch
+
+from .ar_network_transformer import Transformer
+
+
+def sample_top_p(logits, temperature, top_p, return_probs: bool = False):
+    """
+    Perform top-p (nucleus) sampling on a probability distribution.
+
+    Args:
+        logits (torch.Tensor): Logits of the probability distribution.
+        temperature (float): Temperature for sampling.
+        top_p (float): Probability threshold for top-p sampling.
+
+    Returns:
+        torch.Tensor: Sampled token indices.
+
+    Note:
+        Top-p sampling selects the smallest set of tokens whose cumulative probability mass
+        exceeds the threshold p. The distribution is renormalized based on the selected tokens.
+    """
+    probs = torch.softmax(logits[:, -1, :] / temperature, dim=-1)
+    # Sort the probabilities in descending order and get their indices.
+    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+    # Compute the cumulative sum of the sorted probabilities.
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    # Create a mask where the cumulative probability exceeds the threshold p.
+    mask = probs_sum - probs_sort > top_p
+    # Set the probabilities that exceed the threshold to 0.
+    probs_sort[mask] = 0.0
+    # Renormalize the remaining probabilities so they sum to 1.
+    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+    # Sample from the renormalized probability distribution.
+    # next_token = torch.multinomial(probs_sort, num_samples=1)
+    next_token = multinomial_sample_one_no_sync(probs_sort, dtype=torch.int64)
+    # Gather the indices of the sampled tokens.
+    next_token = torch.gather(probs_idx, -1, next_token)
+    if return_probs:
+        # Initialize a tensor for unsorted probabilities
+        probs_unsorted = torch.zeros_like(probs_sort)
+        # Scatter the sorted probabilities back to their original order
+        probs_unsorted.scatter_(-1, probs_idx, probs_sort)
+    else:
+        probs_unsorted = None
+    return next_token, probs_unsorted
+
+
+def multinomial_sample_one_no_sync(probs_sort, dtype=torch.int):
+    """
+    Multinomial sampling without a cuda synchronization.
+    Source: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py
+    """
+    q = torch.empty_like(probs_sort).exponential_(1)
+    return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=dtype)
+
+
+def logits_to_probs(
+    logits,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+):
+    logits = logits / max(temperature, 1e-5)
+
+    if top_k is not None:
+        v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+        pivot = v.select(-1, -1).unsqueeze(-1)
+        logits = torch.where(logits < pivot, -float("Inf"), logits)
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+    return probs
+
+
+def sample_top_k(logits, temperature: float = 1.0, top_k: Optional[int] = None):
+    """
+    Sample from the logits using top-k sampling.
+    Source: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py
+    """
+    # logits: [batch_size, seq_len, vocab_size]
+    if temperature == 0.0:
+        idx_next = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
+        probs = None
+    else:
+        probs = logits_to_probs(logits[:, -1, :], temperature, top_k)
+        idx_next = multinomial_sample_one_no_sync(probs)
+    return idx_next, probs
+
+
+def prefill(
+    model: Transformer,
+    input_pos: torch.Tensor,
+    tokens: torch.Tensor = None,
+    token_embeddings: torch.Tensor = None,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+    top_p: Optional[float] = None,
+    **kwargs,
+) -> torch.Tensor:
+    logits = model(tokens=tokens, token_embeddings=token_embeddings, input_pos=input_pos, **kwargs)
+    # Only top-p or top-k can be provided
+    assert (
+        top_p is None or top_k is None
+    ), "Only one of top-p or top-k can be provided, got top-p={top_p} and top-k={top_k}"
+    if top_p is not None:
+        return sample_top_p(logits, temperature=temperature, top_p=top_p)[0]
+    else:
+        return sample_top_k(logits, temperature=temperature, top_k=top_k)[0]
+
+
+def decode_one_token(
+    model: Transformer,
+    tokens: torch.Tensor,
+    input_pos: torch.Tensor,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+    top_p: Optional[float] = None,
+    **kwargs,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Decode a single token from the autoregressive model.
+    """
+    logits = model(tokens=tokens, input_pos=input_pos, **kwargs)
+    if top_p is not None:
+        return sample_top_p(logits, temperature=temperature, top_p=top_p)
+    else:
+        return sample_top_k(logits, temperature=temperature, top_k=top_k)
+
+
+def decode_n_tokens(
+    model: Transformer,
+    cur_token: torch.Tensor,
+    input_pos: torch.Tensor,
+    num_new_tokens: int,
+    stop_tokens: torch.Tensor = None,
+    temperature: float = 1.0,
+    top_p: Optional[float] = None,
+    top_k: Optional[int] = None,
+    return_probs: bool = False,
+    decode_one_token_function=decode_one_token,
+    **kwargs,
+):
+    """
+    Decode n tokens from the autoregressive model.
+    Adapted from https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py
+    """
+    new_tokens, new_probs = [], []
+    batch_size = cur_token.shape[0]
+    assert (
+        top_p is None or top_k is None
+    ), "Only one of top-p or top-k can be provided, got top-p={top_p} and top-k={top_k}"
+    if stop_tokens is not None:
+        # Indicator for whether the EOS token (stop token) has been reached for each sample in the batch
+        eos_reached = torch.tensor([False] * batch_size, device="cuda")
+    for t in range(num_new_tokens):
+        with torch.backends.cuda.sdp_kernel(
+            enable_flash=False, enable_mem_efficient=False, enable_math=True
+        ):  # Actually better for Inductor to codegen attention here
+            next_token, next_prob = decode_one_token_function(
+                model,
+                tokens=cur_token,
+                input_pos=input_pos,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                **kwargs,
+            )
+            input_pos += 1
+            if stop_tokens is not None and len(stop_tokens) > 0:
+                eos_reached = eos_reached | (torch.isin(next_token, stop_tokens))
+                if eos_reached.all():
+                    break
+            new_tokens.append(next_token.clone())
+            if return_probs:
+                new_probs.append(next_prob.clone())
+            cur_token = next_token.clone()
+
+    if return_probs:
+        return new_tokens, new_probs
+    else:
+        return new_tokens

base.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+
+import imageio
+import torch
+
+from .world_generation_pipeline import ARBaseGenerationPipeline
+from .ar_utils_inference import add_common_arguments, load_vision_input, validate_args
+from .log import log
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Video to world generation demo script")
+    # Add common arguments
+    add_common_arguments(parser)
+    parser.add_argument(
+        "--ar_model_dir",
+        type=str,
+        default="Cosmos-1.0-Autoregressive-4B",
+    )
+    parser.add_argument("--input_type", type=str, default="video", help="Type of input", choices=["image", "video"])
+    args = parser.parse_args()
+    return args
+
+
+def main(args):
+    """Run video-to-world generation demo.
+
+    This function handles the main video-to-world generation pipeline, including:
+    - Setting up the random seed for reproducibility
+    - Initializing the generation pipeline with the provided configuration
+    - Processing single or multiple images/videos from input
+    - Generating videos from images/videos
+    - Saving the generated videos to disk
+
+    Args:
+        cfg (argparse.Namespace): Configuration namespace containing:
+            - Model configuration (checkpoint paths, model settings)
+            - Generation parameters (temperature, top_p)
+            - Input/output settings (images/videos, save paths)
+            - Performance options (model offloading settings)
+
+    The function will save:
+        - Generated MP4 video files
+
+    If guardrails block the generation, a critical log message is displayed
+    and the function continues to the next prompt if available.
+    """
+    inference_type = "base"  # When the inference_type is "base", AR model does not take text as input, the world generation is purely based on the input video
+    sampling_config = validate_args(args, inference_type)
+
+    # Initialize base generation model pipeline
+    pipeline = ARBaseGenerationPipeline(
+        inference_type=inference_type,
+        checkpoint_dir=args.checkpoint_dir,
+        checkpoint_name=args.ar_model_dir,
+        disable_diffusion_decoder=args.disable_diffusion_decoder,
+        offload_guardrail_models=args.offload_guardrail_models,
+        offload_diffusion_decoder=args.offload_diffusion_decoder,
+        offload_network=args.offload_ar_model,
+        offload_tokenizer=args.offload_tokenizer,
+    )
+
+    # Load input image(s) or video(s)
+    input_videos = load_vision_input(
+        input_type=args.input_type,
+        batch_input_path=args.batch_input_path,
+        input_image_or_video_path=args.input_image_or_video_path,
+        data_resolution=args.data_resolution,
+        num_input_frames=args.num_input_frames,
+    )
+
+    for idx, input_filename in enumerate(input_videos):
+        inp_vid = input_videos[input_filename]
+        # Generate video
+        log.info(f"Run with image or video path: {input_filename}")
+        out_vid = pipeline.generate(
+            inp_vid=inp_vid,
+            num_input_frames=args.num_input_frames,
+            seed=args.seed,
+            sampling_config=sampling_config,
+        )
+        if out_vid is None:
+            log.critical("Guardrail blocked base generation.")
+            continue
+
+        # Save video
+        if args.input_image_or_video_path:
+            out_vid_path = os.path.join(args.video_save_folder, f"{args.video_save_name}.mp4")
+        else:
+            out_vid_path = os.path.join(args.video_save_folder, f"{idx}.mp4")
+
+        imageio.mimsave(out_vid_path, out_vid, fps=25)
+
+        log.info(f"Saved video to {out_vid_path}")
+
+
+if __name__ == "__main__":
+    torch._C._jit_set_texpr_fuser_enabled(False)
+    args = parse_args()
+    main(args)

base_world_generation_pipeline.py

@@ -0,0 +1,358 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import gc
+import os
+from abc import ABC
+from typing import Any
+
+import numpy as np
+import torch
+
+from .t5_text_encoder import CosmosT5TextEncoder
+from .guardrail_common_presets import presets as guardrail_presets
+
+
+class BaseWorldGenerationPipeline(ABC):
+    def __init__(
+        self,
+        inference_type: str | None = None,
+        checkpoint_dir: str | None = None,
+        checkpoint_name: str | None = None,
+        has_text_input: bool = False,
+        offload_network: bool = False,
+        offload_tokenizer: bool = False,
+        offload_text_encoder_model: bool = False,
+        offload_guardrail_models: bool = False,
+    ):
+        """Initialize base world generation pipeline.
+
+        This abstract base class provides core functionality for world generation models including:
+        - Model loading and initialization
+        - Text encoding and embedding
+        - Safety checks and content filtering
+        - Memory management through model offloading
+
+        Args:
+            inference_type: The type of inference pipeline ("text2world" or "video2world")
+            checkpoint_dir: Root directory containing model checkpoints
+            checkpoint_name: Name of the specific checkpoint file to load
+            has_text_input: Whether the pipeline takes text input for world generation
+            offload_network: If True, moves main model to CPU after inference
+            offload_tokenizer: If True, moves tokenizer to CPU after use
+            offload_text_encoder_model: If True, moves T5 encoder to CPU after encoding
+            offload_guardrail_models: If True, moves safety models to CPU after checks
+        """
+        self.inference_type = inference_type
+        self.checkpoint_dir = checkpoint_dir
+        self.checkpoint_name = checkpoint_name
+        self.guardrail_dir = "Cosmos-1.0-Guardrail"
+        self.has_text_input = has_text_input
+
+        # Add offloading flags
+        self.offload_network = offload_network
+        self.offload_tokenizer = offload_tokenizer
+        self.offload_text_encoder_model = offload_text_encoder_model
+        self.offload_guardrail_models = offload_guardrail_models
+
+        # Initialize model instances
+        self.text_guardrail = None
+        self.video_guardrail = None
+        self.text_encoder = None
+        self.model = None
+
+        self._load_model()
+
+        if not self.offload_text_encoder_model:
+            self._load_text_encoder_model()
+        if not self.offload_guardrail_models:
+            if self.has_text_input:
+                self._load_text_guardrail()
+            self._load_video_guardrail()
+        if not self.offload_network:
+            self._load_network()
+        if not self.offload_tokenizer:
+            self._load_tokenizer()
+
+    def _load_tokenizer(self):
+        pass
+
+    def _load_network(self):
+        pass
+
+    def _load_model(self, checkpoint_name: str) -> Any:
+        """Load the world generation model from a checkpoint.
+
+        This abstract method must be implemented by subclasses to load their specific
+        model architecture and weights.
+
+        Args:
+            checkpoint_name: Path to the model checkpoint file
+
+        Returns:
+            The loaded model instance
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses
+        """
+        pass
+
+    def _load_text_encoder_model(self):
+        """Load the T5 text encoder model.
+
+        Initializes and loads the T5 encoder model used for converting text prompts
+        into embeddings that condition the world generation model.
+
+        Returns:
+            Loaded T5 text encoder model instance
+        """
+        self.text_encoder = CosmosT5TextEncoder(cache_dir=self.checkpoint_dir)
+
+    def _load_text_guardrail(self):
+        """Load text safety classifier models.
+
+        Initializes models used for checking input prompts against safety policies.
+        Models are loaded from the specified guardrail directory.
+        """
+        self.text_guardrail = guardrail_presets.create_text_guardrail_runner(
+            checkpoint_dir=os.path.join(self.checkpoint_dir, self.guardrail_dir)
+        )
+
+    def _load_video_guardrail(self):
+        """Load video safety classifier models.
+
+        Initializes models used for validating generated video content against
+        safety policies. Models are loaded from the specified guardrail directory.
+        """
+        self.video_guardrail = guardrail_presets.create_video_guardrail_runner(
+            checkpoint_dir=os.path.join(self.checkpoint_dir, self.guardrail_dir)
+        )
+
+    def _offload_network(self):
+        if self.model.model:
+            del self.model.model
+            self.model.model = None
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    def _offload_tokenizer(self):
+        if self.model.tokenizer:
+            del self.model.tokenizer
+            self.model.tokenizer = None
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    def _offload_guardrail_models(self):
+        """Offload safety classifier models to reduce memory usage.
+
+        Moves safety models to CPU and clears GPU memory if they are no longer needed.
+        This helps manage memory when processing multiple inputs sequentially.
+        """
+        if self.text_guardrail:
+            del self.text_guardrail
+            self.text_guardrail = None
+        if self.video_guardrail:
+            del self.video_guardrail
+            self.video_guardrail = None
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def _offload_text_encoder_model(self):
+        """Offload T5 text encoder to reduce memory usage.
+
+        Moves the T5 encoder to CPU and clears GPU memory after text encoding is complete.
+        This helps manage memory when processing multiple inputs sequentially.
+        """
+        if self.text_encoder:
+            del self.text_encoder
+            self.text_encoder = None
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    def _run_model(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        """Generate world latents using the model.
+
+        This abstract method must be implemented by subclasses to define their specific
+        generation process.
+
+        Args:
+            *args: Variable positional arguments for model inference
+            **kwargs: Variable keyword arguments for model inference
+
+        Returns:
+            torch.Tensor: Generated world representation tensor
+        """
+        pass
+
+    def _run_model_with_offload(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        """Generate world representation with memory management.
+
+        Handles loading the model before inference and offloading afterward if enabled.
+        This helps minimize GPU memory usage during inference.
+
+        Args:
+            *args: Arguments passed to _run_model
+            **kwargs: Keyword arguments passed to _run_model
+
+        Returns:
+            np.ndarray: Generated world representation as numpy array
+        """
+        pass
+
+    def _run_guardrail_on_prompt(self, prompt: str) -> bool:
+        """Check if prompt meets safety requirements.
+
+        Validates the input prompt against safety policies using loaded guardrail models.
+
+        Args:
+            prompt: Raw text prompt to validate
+
+        Returns:
+            bool: True if prompt passes all safety checks, False otherwise
+        """
+        return guardrail_presets.run_text_guardrail(prompt, self.text_guardrail)
+
+    def _run_guardrail_on_prompt_with_offload(self, prompt: str) -> bool:
+        """Check prompt safety with memory management.
+
+        Validates prompt safety while handling model loading/offloading to manage memory.
+
+        Args:
+            prompt: Raw text prompt to validate
+
+        Returns:
+            bool: True if prompt passes all safety checks, False otherwise
+        """
+        if self.offload_guardrail_models:
+            self._load_text_guardrail()
+
+        is_safe = self._run_guardrail_on_prompt(prompt)
+
+        if self.offload_guardrail_models:
+            self._offload_guardrail_models()
+
+        return is_safe
+
+    def _run_guardrail_on_video(self, video: np.ndarray) -> np.ndarray | None:
+        """Check if video meets safety requirements.
+
+        Validates generated video content against safety policies using guardrail models.
+
+        Args:
+            video: Video frames to validate
+
+        Returns:
+            np.ndarray: Processed video if safe, None if unsafe
+        """
+        return guardrail_presets.run_video_guardrail(video, self.video_guardrail)
+
+    def _run_guardrail_on_video_with_offload(self, video: np.ndarray) -> np.ndarray | None:
+        """Check if generated video meets safety requirements.
+
+        Args:
+            video: Video frames to validate
+
+        Returns:
+            np.ndarray: Processed video frames if safe, None otherwise
+
+        Note:
+            Guardrail models are offloaded after checks if enabled.
+        """
+        if self.offload_guardrail_models:
+            self._load_video_guardrail()
+
+        video = self._run_guardrail_on_video(video)
+
+        if self.offload_guardrail_models:
+            self._offload_guardrail_models()
+        return video
+
+    def _run_text_embedding_on_prompt(
+        self, prompts: list[str], **kwargs: Any
+    ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+        """Convert text prompts to embeddings.
+
+        Processes text prompts into embedding tensors that condition the generation model.
+
+        Args:
+            prompts: List of text prompts to encode
+            **kwargs: Additional arguments for text encoding
+
+        Returns:
+            tuple containing:
+                - List of text embedding tensors for each prompt
+                - List of attention masks for each embedding
+        """
+
+        embeddings = []
+        masks = []
+        for prompt in prompts:
+            embedding, mask = self.text_encoder.encode_prompts(
+                [prompt],
+                **kwargs,
+            )
+            embeddings.append(embedding)
+            masks.append(mask)
+
+        return embeddings, masks
+
+    def _run_text_embedding_on_prompt_with_offload(
+        self, prompts: list[str], **kwargs: Any
+    ) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+        """Convert text prompt into embeddings using T5 encoder.
+
+        Args:
+            prompt: Processed and validated text prompt
+
+        Returns:
+            Text embedding tensor to condition diffusion model
+
+        Note:
+            T5 model is offloaded after encoding if enabled.
+        """
+        if self.offload_text_encoder_model:
+            self._load_text_encoder_model()
+
+        embeddings, masks = self._run_text_embedding_on_prompt(prompts, **kwargs)
+
+        if self.offload_text_encoder_model:
+            self._offload_text_encoder_model()
+        return embeddings, masks
+
+    def _run_tokenizer_decoding(self, samples: torch.Tensor) -> np.ndarray:
+        """Decode model outputs into final world representation.
+
+        This abstract method must be implemented by subclasses to convert raw model
+        outputs into their specific world representation format.
+
+        Args:
+            samples: Raw output tensor from the generation model
+
+        Returns:
+            np.ndarray: Decoded world representation
+        """
+        pass
+
+    def generate(self, *args: Any, **kwargs: Any):
+        """Generate world representation.
+
+        This abstract method must be implemented by subclasses to convert raw model
+        outputs into their specific world representation format.
+
+        Args:
+            *args: Variable positional arguments for model inference
+            **kwargs: Variable keyword arguments for model inference
+        """
+        pass

config.json

@@ -0,0 +1,10 @@
+{
+    "architectures": [
+      "ARVideo2World"
+    ],
+    "auto_map": {
+      "AutoConfig": "video2world_hf.ARVideo2WorldConfig",
+      "AutoModel": "video2world_hf.ARVideo2World"
+    },
+    "model_type": "AutoModel"
+}

config.py

@@ -0,0 +1,165 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Any, TypeVar
+
+import attrs
+
+from .lazy_config_init import LazyDict
+from .misc import Color
+
+T = TypeVar("T")
+
+
+def _is_attrs_instance(obj: object) -> bool:
+    """
+    Helper function to check if an object is an instance of an attrs-defined class.
+
+    Args:
+        obj: The object to check.
+
+    Returns:
+        bool: True if the object is an instance of an attrs-defined class, False otherwise.
+    """
+    return hasattr(obj, "__attrs_attrs__")
+
+
+def make_freezable(cls: T) -> T:
+    """
+    A decorator that adds the capability to freeze instances of an attrs-defined class.
+
+    NOTE: This requires the wrapped attrs to be defined with attrs.define(slots=False) because we need
+    to hack on a "_is_frozen" attribute.
+
+    This decorator enhances an attrs-defined class with the ability to be "frozen" at runtime.
+    Once an instance is frozen, its attributes cannot be changed. It also recursively freezes
+    any attrs-defined objects that are attributes of the class.
+
+    Usage:
+        @make_freezable
+        @attrs.define(slots=False)
+        class MyClass:
+            attribute1: int
+            attribute2: str
+
+        obj = MyClass(1, 'a')
+        obj.freeze()  # Freeze the instance
+        obj.attribute1 = 2  # Raises AttributeError
+
+    Args:
+        cls: The class to be decorated.
+
+    Returns:
+        The decorated class with added freezing capability.
+    """
+
+    if not hasattr(cls, "__dict__"):
+        raise TypeError(
+            "make_freezable cannot be used with classes that do not define __dict__. Make sure that the wrapped "
+            "class was defined with `@attrs.define(slots=False)`"
+        )
+
+    original_setattr = cls.__setattr__
+
+    def setattr_override(self, key, value) -> None:  # noqa: ANN001
+        """
+        Override __setattr__ to allow modifications during initialization
+        and prevent modifications once the instance is frozen.
+        """
+        if hasattr(self, "_is_frozen") and self._is_frozen and key != "_is_frozen":
+            raise AttributeError("Cannot modify frozen instance")
+        original_setattr(self, key, value)  # type: ignore
+
+    cls.__setattr__ = setattr_override  # type: ignore
+
+    def freeze(self: object) -> None:
+        """
+        Freeze the instance and all its attrs-defined attributes.
+        """
+        for _, value in attrs.asdict(self, recurse=False).items():
+            if _is_attrs_instance(value) and hasattr(value, "freeze"):
+                value.freeze()
+        self._is_frozen = True  # type: ignore
+
+    cls.freeze = freeze  # type: ignore
+
+    return cls
+
+
+def _pretty_print_attrs_instance(obj: object, indent: int = 0, use_color: bool = False) -> str:
+    """
+    Recursively pretty prints attrs objects with color.
+    """
+
+    assert attrs.has(obj.__class__)
+
+    lines: list[str] = []
+    for attribute in attrs.fields(obj.__class__):
+        value = getattr(obj, attribute.name)
+        if attrs.has(value.__class__):
+            if use_color:
+                lines.append("   " * indent + Color.cyan("* ") + Color.green(attribute.name) + ":")
+            else:
+                lines.append("   " * indent + "* " + attribute.name + ":")
+            lines.append(_pretty_print_attrs_instance(value, indent + 1, use_color))
+        else:
+            if use_color:
+                lines.append(
+                    "   " * indent + Color.cyan("* ") + Color.green(attribute.name) + ": " + Color.yellow(value)
+                )
+            else:
+                lines.append("   " * indent + "* " + attribute.name + ": " + str(value))
+    return "\n".join(lines)
+
+
+@make_freezable
+@attrs.define(slots=False)
+class JobConfig:
+    # Project name.
+    project: str = ""
+    # Experiment name.
+    group: str = ""
+    # Run/job name.
+    name: str = ""
+
+    @property
+    def path(self) -> str:
+        return f"{self.project}/{self.group}/{self.name}"
+
+
+@make_freezable
+@attrs.define(slots=False)
+class Config:
+    """Config for a job.
+
+    See /README.md/Configuration System for more info.
+    """
+
+    # Model configs.
+    model: LazyDict
+
+    # Training job configs.
+    job: JobConfig = attrs.field(factory=JobConfig)
+
+    def to_dict(self) -> dict[str, Any]:
+        return attrs.asdict(self)
+
+    def validate(self) -> None:
+        """Validate that the config has all required fields."""
+        assert self.job.project != "", "Project name is required."
+        assert self.job.group != "", "Group name is required."
+        assert self.job.name != "", "Job name is required."

config_helper.py

@@ -0,0 +1,197 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib
+import os
+import pkgutil
+import sys
+from dataclasses import fields as dataclass_fields
+from dataclasses import is_dataclass
+from typing import Any, Dict, Optional
+
+import attr
+import attrs
+from hydra import compose, initialize
+from hydra.core.config_store import ConfigStore
+from omegaconf import DictConfig, OmegaConf
+
+from .log import log
+from .config import Config
+
+
+def is_attrs_or_dataclass(obj) -> bool:
+    """
+    Check if the object is an instance of an attrs class or a dataclass.
+
+    Args:
+        obj: The object to check.
+
+    Returns:
+        bool: True if the object is an instance of an attrs class or a dataclass, False otherwise.
+    """
+    return is_dataclass(obj) or attr.has(type(obj))
+
+
+def get_fields(obj):
+    """
+    Get the fields of an attrs class or a dataclass.
+
+    Args:
+        obj: The object to get fields from. Must be an instance of an attrs class or a dataclass.
+
+    Returns:
+        list: A list of field names.
+
+    Raises:
+        ValueError: If the object is neither an attrs class nor a dataclass.
+    """
+    if is_dataclass(obj):
+        return [field.name for field in dataclass_fields(obj)]
+    elif attr.has(type(obj)):
+        return [field.name for field in attr.fields(type(obj))]
+    else:
+        raise ValueError("The object is neither an attrs class nor a dataclass.")
+
+
+def override(config: Config, overrides: Optional[list[str]] = None) -> Config:
+    """
+    :param config: the instance of class `Config` (usually from `make_config`)
+    :param overrides: list of overrides for config
+    :return: the composed instance of class `Config`
+    """
+    # Store the class of the config for reconstruction after overriding.
+    # config_class = type(config)
+
+    # Convert Config object to a DictConfig object
+    config_dict = attrs.asdict(config)
+    config_omegaconf = DictConfig(content=config_dict, flags={"allow_objects": True})
+    # Enforce "--" separator between the script arguments and overriding configs.
+    if overrides:
+        if overrides[0] != "--":
+            raise ValueError('Hydra config overrides must be separated with a "--" token.')
+        overrides = overrides[1:]
+    # Use Hydra to handle overrides
+    cs = ConfigStore.instance()
+    cs.store(name="config", node=config_omegaconf)
+    with initialize(version_base=None):
+        config_omegaconf = compose(config_name="config", overrides=overrides)
+        OmegaConf.resolve(config_omegaconf)
+
+    def config_from_dict(ref_instance: Any, kwargs: Any) -> Any:
+        """
+        Construct an instance of the same type as ref_instance using the provided dictionary or data or unstructured data
+
+        Args:
+            ref_instance: The reference instance to determine the type and fields when needed
+            kwargs: A dictionary of keyword arguments to use for constructing the new instance or primitive data or unstructured data
+
+        Returns:
+            Any: A new instance of the same type as ref_instance constructed using the provided kwargs or the primitive data or unstructured data
+
+        Raises:
+            AssertionError: If the fields do not match or if extra keys are found.
+            Exception: If there is an error constructing the new instance.
+        """
+        is_type = is_attrs_or_dataclass(ref_instance)
+        if not is_type:
+            return kwargs
+        else:
+            ref_fields = set(get_fields(ref_instance))
+            assert isinstance(kwargs, dict) or isinstance(
+                kwargs, DictConfig
+            ), "kwargs must be a dictionary or a DictConfig"
+            keys = set(kwargs.keys())
+
+            # ref_fields must equal to or include all keys
+            extra_keys = keys - ref_fields
+            assert ref_fields == keys or keys.issubset(
+                ref_fields
+            ), f"Fields mismatch: {ref_fields} != {keys}. Extra keys found: {extra_keys} \n \t when constructing {type(ref_instance)} with {keys}"
+
+            resolved_kwargs: Dict[str, Any] = {}
+            for f in keys:
+                resolved_kwargs[f] = config_from_dict(getattr(ref_instance, f), kwargs[f])
+            try:
+                new_instance = type(ref_instance)(**resolved_kwargs)
+            except Exception as e:
+                log.error(f"Error when constructing {type(ref_instance)} with {resolved_kwargs}")
+                log.error(e)
+                raise e
+            return new_instance
+
+    config = config_from_dict(config, config_omegaconf)
+
+    return config
+
+
+def get_config_module(config_file: str) -> str:
+    if not config_file.endswith(".py"):
+        log.error("Config file cannot be specified as module.")
+        log.error("Please provide the path to the Python config file (relative to the Cosmos root).")
+    assert os.path.isfile(config_file), f"Cosmos config file ({config_file}) not found."
+    # Convert to importable module format.
+    config_module = config_file.replace("/", ".").replace(".py", "")
+    return config_module
+
+
+def import_all_modules_from_package(package_path: str, reload: bool = False, skip_underscore: bool = True) -> None:
+    """
+    Import all modules from the specified package path recursively.
+
+    This function is typically used in conjunction with Hydra to ensure that all modules
+    within a specified package are imported, which is necessary for registering configurations.
+
+    Example usage:
+    ```python
+    import_all_modules_from_package("cosmos1.models.diffusion.config.inference", reload=True, skip_underscore=False)
+    ```
+
+    Args:
+        package_path (str): The dotted path to the package from which to import all modules.
+        reload (bool): Flag to determine whether to reload modules if they're already imported.
+        skip_underscore (bool): If True, skips importing modules that start with an underscore.
+    """
+    return
+    log.debug(f"{'Reloading' if reload else 'Importing'} all modules from package {package_path}")
+    package = importlib.import_module(package_path)
+    package_directory = package.__path__
+
+    def import_modules_recursively(directory: str, prefix: str) -> None:
+        """
+        Recursively imports or reloads all modules in the given directory.
+
+        Args:
+            directory (str): The file system path to the current package directory.
+            prefix (str): The module prefix (e.g., 'cosmos1.models.diffusion.config').
+        """
+        for _, module_name, is_pkg in pkgutil.iter_modules([directory]):
+            if skip_underscore and module_name.startswith("_"):
+                log.debug(f"Skipping module {module_name} as it starts with an underscore")
+                continue
+
+            full_module_name = f"{prefix}.{module_name}"
+            log.debug(f"{'Reloading' if reload else 'Importing'} module {full_module_name}")
+
+            if full_module_name in sys.modules and reload:
+                importlib.reload(sys.modules[full_module_name])
+            else:
+                importlib.import_module(full_module_name)
+
+            if is_pkg:
+                sub_package_directory = os.path.join(directory, module_name)
+                import_modules_recursively(sub_package_directory, full_module_name)
+
+    for directory in package_directory:
+        import_modules_recursively(directory, package_path)

convert_pixtral_ckpt.py

@@ -0,0 +1,209 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Convert pretrained Pixtral vision model weights to checkpoint and verify the checkpoint loading.
+
+    Usage:
+
+    PYTHONPATH=$(pwd) python cosmos1/scripts/convert_pixtral_ckpt.py
+"""
+
+import argparse
+import json
+import os
+import shutil
+from glob import glob
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+
+
+def convert_pixtral_checkpoint(checkpoint_dir: str, checkpoint_name: str, vit_type: str):
+    """
+    Main function to convert Pixtral vision model weights to checkpoint and optionally verify and save the converted checkpoint.
+
+    Args:
+        checkpoint_dir (str): Path to the checkpoint directory
+        checkpoint_name (str): Name of the checkpoint
+        vit_type (str): Type of ViT used in the Pixtral model
+
+    This function performs the following steps:
+    0. Download the checkpoint from Hugging Face
+    1. Loads the original Pixtral checkpoint
+    2. Splits the checkpoint into vision encoder, projector, and LLM weights
+    3. Reorganizes the weights to match the expected format
+    4. Extracts and verifies the vision encoder configuration
+    5. Optionally verifies the converted checkpoint by loading it into a VisionTransformer
+    6. Optionally saves the converted checkpoint and configuration
+    """
+
+    save_dir = os.path.join(checkpoint_dir, checkpoint_name)
+    os.makedirs(save_dir, exist_ok=True)
+    # Save the converted checkpoint
+    save_path = os.path.join(save_dir, "model.pt")
+    if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
+        print(f"Checkpoint {save_path} already exists and is not empty")
+        return
+
+    pixtral_ckpt_dir = os.path.join(checkpoint_dir, "Pixtral-12B-2409")
+    os.makedirs(pixtral_ckpt_dir, exist_ok=True)
+    repo_id = "mistralai/Pixtral-12B-2409"
+    print(f"Downloading {repo_id} to {pixtral_ckpt_dir}...")
+    snapshot_download(
+        repo_id=repo_id,
+        allow_patterns=["params.json", "consolidated.safetensors"],
+        local_dir=pixtral_ckpt_dir,
+        local_dir_use_symlinks=False,
+    )
+    orig_dtype = torch.get_default_dtype()
+    dtype = torch.bfloat16
+    torch.set_default_dtype(dtype)
+
+    # Load checkpoint file
+    ckpt_files = glob(os.path.join(pixtral_ckpt_dir, "*.safetensors"))
+    assert len(ckpt_files) == 1, "ckpt_dir should contain only one file"
+    ckpt_path = ckpt_files[0]
+    ckpt = load_file(ckpt_path)
+
+    # Split checkpoint into weights of vision encoder, projector, and LLM
+    vit_key_prefix = "vision_encoder."
+    vit_ckpt = {}
+    for key, value in ckpt.items():
+        if key.startswith(vit_key_prefix):
+            vit_ckpt[key.lstrip(vit_key_prefix)] = value
+
+    projector_key_prefix = "vision_language_adapter."
+    projector_ckpt = {}
+    substring_replacement_map = {
+        "w_in.": "projector.0.",
+        "w_out.": "projector.2.",
+    }
+    for key, value in ckpt.items():
+        if key.startswith(projector_key_prefix):
+            key = key.lstrip(projector_key_prefix)
+            for old, new in substring_replacement_map.items():
+                key = key.replace(old, new)
+            projector_ckpt[key] = value
+
+    llm_ckpt = {}
+    for key, value in ckpt.items():
+        if key.startswith(vit_key_prefix) or key.startswith(projector_key_prefix):
+            continue
+        llm_ckpt[key] = value
+
+    vlm_ckpt = {}
+    for key, value in llm_ckpt.items():
+        vlm_ckpt["model." + key] = value
+    for key, value in projector_ckpt.items():
+        vlm_ckpt["mm_projector." + key] = value
+    for key, value in vit_ckpt.items():
+        vlm_ckpt["vision_encoder." + key] = value
+
+    # Load config
+    config_path = os.path.join(pixtral_ckpt_dir, "params.json")
+    with open(config_path, "r") as f:
+        pixtral_config = json.load(f)
+
+    # Extract the vision encoder configuration
+    vision_encoder_config = {
+        "dim": pixtral_config["vision_encoder"]["hidden_size"],
+        "num_channels": pixtral_config["vision_encoder"]["num_channels"],
+        "image_size": pixtral_config["vision_encoder"]["image_size"],
+        "patch_size": pixtral_config["vision_encoder"]["patch_size"],
+        "rope_theta": pixtral_config["vision_encoder"]["rope_theta"],
+        "ffn_hidden_size": pixtral_config["vision_encoder"]["intermediate_size"],
+        "n_layers": pixtral_config["vision_encoder"]["num_hidden_layers"],
+        "n_heads": pixtral_config["vision_encoder"]["num_attention_heads"],
+        "n_kv_heads": pixtral_config["vision_encoder"]["num_attention_heads"],
+        "norm_type": "rmsnorm",
+        "norm_eps": pixtral_config["norm_eps"],
+        "image_token_id": pixtral_config["vision_encoder"]["image_token_id"],
+    }
+    # Configuration for the 400M ViT of Pixtral 12B VLM
+    vit_config = dict(
+        dim=1024,
+        num_channels=3,
+        image_size=1024,
+        patch_size=16,
+        rope_theta=10000,
+        ffn_hidden_size=4096,
+        n_layers=24,
+        n_heads=16,
+        n_kv_heads=16,
+        norm_type="rmsnorm",
+        norm_eps=1e-5,
+        image_token_id=10,
+    )
+    # Compare the two configurations
+    for key, value in vit_config.items():
+        assert vision_encoder_config[key] == value, f"Mismatch in {key}: {vision_encoder_config[key]} != {value}"
+
+    llm_config_keys = [
+        "dim",
+        "n_layers",
+        "head_dim",
+        "hidden_dim",
+        "n_heads",
+        "n_kv_heads",
+        "rope_theta",
+        "norm_eps",
+        "vocab_size",
+    ]
+    assert set(list(pixtral_config.keys())) == set(llm_config_keys + ["vision_encoder"]), "Config keys mismatch"
+    replace_map = {
+        "hidden_dim": "ffn_hidden_size",
+    }
+    llm_config = {}
+    for k, v in pixtral_config.items():
+        if k in llm_config_keys:
+            llm_config[replace_map.get(k, k)] = v
+        elif k == "vision_encoder":
+            llm_config["vision_encoder"] = vit_type
+        else:
+            raise ValueError(f"Unknown key: {k}")
+
+    ckpt_to_save = {"model": vlm_ckpt, "mm_projector": projector_ckpt, "vision_encoder": vit_ckpt}
+    torch.save(ckpt_to_save, save_path)
+    print(f"Model saved to {save_path}")
+
+    # Save config
+    config_path = os.path.join(save_dir, "config.json")
+    with open(config_path, "w") as f:
+        json.dump(llm_config, f)
+
+    torch.set_default_dtype(orig_dtype)  # Reset the default dtype
+
+    # Remove the original Pixtral checkpoint
+    shutil.rmtree(pixtral_ckpt_dir, ignore_errors=True)
+    print(f"Removed {pixtral_ckpt_dir}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Convert pretrained Pixtral vision model weights to checkpoint and verify accuracy"
+    )
+    parser.add_argument("--checkpoint_dir", type=str, default="checkpoints", help="Path to the checkpoint directory")
+    parser.add_argument(
+        "--checkpoint_name",
+        type=str,
+        default="Pixtral-12B",
+        help="Name of the checkpoint",
+    )
+    parser.add_argument("--vit_type", default="pixtral-12b-vit", help="Type of ViT used in the Pixtral model")
+    args = parser.parse_args()
+    convert_pixtral_checkpoint(
+        checkpoint_dir=args.checkpoint_dir, checkpoint_name=args.checkpoint_name, vit_type=args.vit_type
+    )

cosmos1/models/POST_TRAINING.md

@@ -0,0 +1,23 @@
+# Cosmos Post-training
+
+In the [Cosmos paper](https://research.nvidia.com/publication/2025-01_cosmos-world-foundation-model-platform-physical-ai), we discuss several post-training examples of Cosmos pre-trained World Foundation Models (WFMs) for various Physical AI tasks, including
+
+- General Post-Training: Fine-tune the WFM to generate a target distribution of videos based on the custom dataset. The target distribution could include a specific camera spec or a specific domain such as a factory.
+- Instruction Control: Post-trains models for robotic manipulation to predict videos based on textual instructions, enabling robots to visually simulate tasks like folding clothes or picking up objects.
+- Action Control: Post-trains models for robotic manipulation to predict the next visual frame based on action vectors, simulating robotic tasks like object handling or movement planning.
+- Camera Control: Adds camera pose conditioning to generate 3D-consistent video simulations from single images, enabling joystick-like navigation in virtual environments.
+- Multi-View Generation: Post-trains models for autonomous vehicles to generate synchronized multi-view videos from text prompts, simulating driving scenarios with multiple camera perspectives.
+- Multi-View Generation with Vehicle Trajectory Control: Extends multi-view generation by incorporating trajectory inputs, enabling precise simulation of driving environments for autonomous vehicles, adhering to specified paths.
+
+Except for the instruction control where the WFM is post-trained on a dataset of instruction-video pairs, all other cases require minor modifications of the network architectures. Post-training tasks will be supported by NeMo Framework. In this initial release, we provide post-training scripts for the general post-training of both diffusion and autorgressive WFMs. Scripts of the other post-training tasks will be provided in a future release.
+
+## Post-training Support Matrix
+
+| Post-training Task  | Diffusion WFM | Autoregressive WFM |
+|---------------------|---------------|--------------------|
+| General post-training | [Supported](../models/diffusion/nemo/post_training/README.md) | [Supported](../models/autoregressive/nemo/post_training/README.md) |
+| Instruction control | Coming soon | Coming soon |
+| Action control | Coming soon | Coming soon |
+| Camera control | Coming soon | Coming soon |
+| Multi-view generation | Coming soon | Coming soon |
+| Multi-view generation with vehicle trajectory control | Coming soon | Coming soon |

cosmos1/models/autoregressive/README.md

@@ -0,0 +1,427 @@
+# Cosmos Autoregressive-based World Foundation Models
+
+## Table of Contents
+- [Getting Started](#getting-started)
+  - [Set Up Docker Environment](#set-up-docker-environment)
+  - [Download Checkpoints](#download-checkpoints)
+- [Usage](#usage)
+  - [Model Types](#model-types)
+  - [Single and Batch Generation](#single-and-batch-generation)
+  - [Sample Commands](#sample-commands)
+    - [Base Models (4B/12B)](#base-basepy-4b-and-12b)
+    - [Video2World Models (5B/13B)](#video2world-video2worldpy-5b-and-13b)
+  - [Arguments](#arguments)
+    - [Common Parameters](#common-parameters)
+    - [Base Specific Parameters](#base-specific-parameters)
+    - [Video2World Specific Parameters](#video2world-specific-parameters)
+  - [Safety Features](#safety-features)
+
+This page details the steps for using the Cosmos autoregressive-based world foundation models.
+
+## Getting Started
+
+### Set Up Docker Environment
+
+Follow our [Installation Guide](../../../INSTALL.md) to set up the Docker environment. All commands on this page should be run inside Docker.
+
+### Download Checkpoints
+
+1. Generate a [Hugging Face](https://huggingface.co/settings/tokens) access token. Set the access token to 'Read' permission (default is 'Fine-grained').
+
+2. Log in to Hugging Face with the access token:
+
+```bash
+huggingface-cli login
+```
+
+3. Download the Cosmos model weights from [Hugging Face](https://huggingface.co/collections/nvidia/cosmos-6751e884dc10e013a0a0d8e6):
+
+```bash
+PYTHONPATH=$(pwd) python cosmos1/scripts/download_autoregressive.py --model_sizes 4B 5B 12B 13B
+```
+
+4. The downloaded files should be in the following structure:
+
+```
+checkpoints/
+├── Cosmos-1.0-Autoregressive-4B
+│   ├── model.pt
+│   └── config.json
+├── Cosmos-1.0-Autoregressive-5B-Video2World
+│   ├── model.pt
+│   └── config.json
+├── Cosmos-1.0-Autoregressive-12B
+│   ├── model.pt
+│   └── config.json
+├── Cosmos-1.0-Autoregressive-13B-Video2World
+│   ├── model.pt
+│   └── config.json
+├── Cosmos-1.0-Tokenizer-CV8x8x8
+│   ├── decoder.jit
+│   ├── encoder.jit
+│   └── mean_std.pt
+├── Cosmos-1.0-Tokenizer-DV8x16x16
+│   ├── decoder.jit
+│   └── encoder.jit
+├── Cosmos-1.0-Diffusion-7B-Decoder-DV8x16x16ToCV8x8x8
+│   ├── aux_vars.pt
+│   └── model.pt
+└── Cosmos-1.0-Guardrail
+    ├── aegis/
+    ├── blocklist/
+    ├── face_blur_filter/
+    └── video_content_safety_filter/
+```
+
+## Usage
+
+
+### Model Types
+
+There are two model types available for autoregressive world generation:
+
+1. **Base**: Supports world generation from image/video input
+
+* Models: `Cosmos-1.0-Autoregressive-4B` and `Cosmos-1.0-Autoregressive-12B`
+* Inference script: [base.py](/cosmos1/models/autoregressive/inference/base.py)
+
+2. **Video2World**: Supports world generation from image/video input and text input
+
+* Models: `Cosmos-1.0-Autoregressive-5B-Video2World` and `Cosmos-1.0-Autoregressive-13B-Video2World`
+* Inference script: [video2world.py](/cosmos1/models/autoregressive/inference/video2world.py)
+
+Our models now support video extension up to 33 frames. Starting from either a single image or a 9-frame video input, they can generate the remaining frames to reach the 33-frame length (generating 32 or 24 frames, respectively).
+
+We have evaluated all eight possible configurations (4 models × 2 vision input types: image or video) using 100 test videos on physical AI topics. Below are the failure rates for each configuration:
+
+| Model                                      | Image input | Video input (9 frames) |
+|:------------------------------------------|:--------------:|:-------------------------:|
+| Cosmos-1.0-Autoregressive-4B              | 15%           | 1%                       |
+| Cosmos-1.0-Autoregressive-5B-Video2World  | 7%            | 2%                       |
+| Cosmos-1.0-Autoregressive-12B             | 2%            | 1%                       |
+| Cosmos-1.0-Autoregressive-13B-Video2World | 3%            | 0%                       |
+
+We define failure cases as videos with severe distortions, such as:
+
+* Sudden appearance of large unexpected objects
+* Video degrading to a single solid color
+
+Note that the following are not considered failures in our analysis:
+
+* Static video frames
+* Minor object distortions or artifacts
+
+### Single and Batch Generation
+
+We support both single and batch video generation.
+
+For generating a single video, `base` mode requires the input argument `--input_image_or_video_path` (image/video input), while `video2world` mode requires both `--input_image_or_video_path` (image/video input) and `--prompt` (text input).
+
+Note that our model only works with 1024x640 resolution videos. If the input image/video is not in this resolution, it will be resized and cropped.
+
+For generating a batch of videos, both `base` and `video2world` require `--batch_input_path` (path to a JSONL file). For `base`, the JSONL file should contain one visual input per line in the following format, where each line must contain a "visual_input" field:
+
+```json
+{"visual_input": "path/to/video1.mp4"}
+{"visual_input": "path/to/video2.mp4"}
+```
+
+For `video2world`, each line in the JSONL file must contain both "prompt" and "visual_input" fields:
+
+```json
+{"prompt": "prompt1", "visual_input": "path/to/video1.mp4"}
+{"prompt": "prompt2", "visual_input": "path/to/video2.mp4"}
+```
+
+### Sample Commands
+
+There are two main demo scripts for autoregressive world generation: `base.py` and `video2world.py`. Below you will find sample commands for single and batch generation, as well as commands for running with low-memory GPUs using model offloading. We also provide a memory usage table comparing different offloading strategies to help with configuration.
+
+#### Base (base.py): 4B and 12B
+
+Generates world from image/video input.
+
+The `input_type` argument can be either `video` or `image`. We have tuned the sampling parameters `top_p` and `temperature` to achieve the best performance. Please use the provided values in the command examples.
+
+Note that the command examples below all use video input. If you want to use image input, please change the `input_type` to `image`.
+
+##### Single Generation
+
+```bash
+# Example using 4B model
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/base.py \
+    --input_type=video \
+    --input_image_or_video_path=cosmos1/models/autoregressive/assets/v1p0/input.mp4 \
+    --video_save_name=Cosmos-1.0-Autoregressive-4B \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-4B \
+    --top_p=0.8 \
+    --temperature=1.0
+
+# Example for low-memory GPUs using 4B model with model offloading
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/base.py \
+    --input_type=video \
+    --input_image_or_video_path=cosmos1/models/autoregressive/assets/v1p0/input.mp4 \
+    --video_save_name=Cosmos-1.0-Autoregressive-4B \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-4B \
+    --top_p=0.8 \
+    --temperature=1.0 \
+    --offload_guardrail_models \
+    --offload_diffusion_decoder \
+    --offload_ar_model \
+    --offload_tokenizer
+
+# Example using 12B model
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/base.py \
+    --input_type=video \
+    --input_image_or_video_path=cosmos1/models/autoregressive/assets/v1p0/input.mp4 \
+    --video_save_name=Cosmos-1.0-Autoregressive-12B \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-12B \
+    --top_p=0.9 \
+    --temperature=1.0
+
+# Example for low-memory GPUs using 12B model with model offloading
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/base.py \
+    --input_type=video \
+    --input_image_or_video_path=cosmos1/models/autoregressive/assets/v1p0/input.mp4 \
+    --video_save_name=Cosmos-1.0-Autoregressive-12B \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-12B \
+    --top_p=0.9 \
+    --temperature=1.0 \
+    --offload_guardrail_models \
+    --offload_diffusion_decoder \
+    --offload_ar_model \
+    --offload_tokenizer
+```
+
+##### Batch Generation
+
+```bash
+# Example using 4B model
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/base.py \
+    --input_type=video \
+    --batch_input_path=cosmos1/models/autoregressive/assets/v1p0/batch_inputs/base.jsonl \
+    --video_save_folder=outputs/Cosmos-1.0-Autoregressive-4B \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-4B \
+    --top_p=0.8 \
+    --temperature=1.0
+
+# Example using 12B model
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/base.py \
+    --input_type=video \
+    --batch_input_path=cosmos1/models/autoregressive/assets/v1p0/batch_inputs/base.jsonl \
+    --video_save_folder=outputs/Cosmos-1.0-Autoregressive-12B \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-12B \
+    --top_p=0.9 \
+    --temperature=1.0
+```
+
+##### Example Output
+
+Here is an example output video generated using base.py with image input, using `Cosmos-1.0-Autoregressive-12B`:
+
+<video src="https://github.com/user-attachments/assets/634403a5-1873-42d7-8dd0-eb7fb4ac8cf4">
+  Your browser does not support the video tag.
+</video>
+
+The input image used to generate this video can be found in `cosmos1/models/autoregressive/assets/v1p0/input.jpg`. The image is from [BDD dataset](http://bdd-data.berkeley.edu/).
+
+Here is an example output video generated using base.py with 9-frame video input, using `Cosmos-1.0-Autoregressive-12B`:
+
+<video src="https://github.com/user-attachments/assets/1a3ff099-87d7-41e8-b149-a25cfcd4f40b">
+  Your browser does not support the video tag.
+</video>
+
+The input video used to generate this video can be found in `cosmos1/models/autoregressive/assets/v1p0/input.mp4`.
+
+##### Inference Time and GPU Memory Usage
+
+These numbers may vary based on system specifications and are provided for reference only.
+
+| Offloading Strategy | Cosmos-1.0-Autoregressive-4B | Cosmos-1.0-Autoregressive-12B |
+|-------------|---------|---------|
+| No offloading | 31.3 GB | 47.5 GB |
+| Guardrails | 28.9 GB | 45.2 GB |
+| Guardrails & Diffusion decoder | 28.5 GB | 43.1 GB |
+| Guardrails & Diffusion decoder & Tokenizer | 27.3 GB | 42.9 GB |
+| Guardrails & Diffusion decoder & Tokenizer & AR model | 18.7 GB | 27.4 GB |
+
+End-to-end inference runtime on one H100 without offloading and after model initialization:
+
+| Cosmos-1.0-Autoregressive-4B | Cosmos-1.0-Autoregressive-12B |
+|---------|---------|
+| ~62 seconds | ~119 seconds |
+
+#### Video2World (video2world.py): 5B and 13B
+
+Generates world from image/video and text input.
+
+The `input_type` argument can be either `text_and_video` or `text_and_image`. We have tuned the sampling parameters `top_p` and `temperature` to achieve the best performance. Please use the provided values in the command examples.
+
+Note that the command examples below all use video input. If you want to use image input, please change the `input_type` to `text_and_image`.
+
+##### Single Generation
+
+```bash
+# Example using 5B model
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/video2world.py \
+    --input_type=text_and_video \
+    --input_image_or_video_path=cosmos1/models/autoregressive/assets/v1p0/input.mp4 \
+    --prompt="A video recorded from a moving vehicle's perspective, capturing roads, buildings, landscapes, and changing weather and lighting conditions." \
+    --video_save_name=Cosmos-1.0-Autoregressive-5B-Video2World \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-5B-Video2World \
+    --top_p=0.7 \
+    --temperature=1.0
+
+# Example for low-memory GPUs using 5B model with model offloading
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/video2world.py \
+    --input_type=text_and_video \
+    --input_image_or_video_path=cosmos1/models/autoregressive/assets/v1p0/input.mp4 \
+    --prompt="A video recorded from a moving vehicle's perspective, capturing roads, buildings, landscapes, and changing weather and lighting conditions." \
+    --video_save_name=Cosmos-1.0-Autoregressive-5B-Video2World \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-5B-Video2World \
+    --top_p=0.7 \
+    --temperature=1.0 \
+    --offload_guardrail_models \
+    --offload_diffusion_decoder \
+    --offload_ar_model \
+    --offload_tokenizer \
+    --offload_text_encoder_model
+
+# Example using 13B model
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/video2world.py \
+    --input_type=text_and_video \
+    --input_image_or_video_path=cosmos1/models/autoregressive/assets/v1p0/input.mp4 \
+    --prompt="A video recorded from a moving vehicle's perspective, capturing roads, buildings, landscapes, and changing weather and lighting conditions." \
+    --video_save_name=Cosmos-1.0-Autoregressive-13B-Video2World \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-13B-Video2World \
+    --top_p=0.8 \
+    --temperature=1.0 \
+    --offload_guardrail_models
+
+# Example for low-memory GPUs using 13B model with model offloading
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/video2world.py \
+    --input_type=text_and_video \
+    --input_image_or_video_path=cosmos1/models/autoregressive/assets/v1p0/input.mp4 \
+    --prompt="A video recorded from a moving vehicle's perspective, capturing roads, buildings, landscapes, and changing weather and lighting conditions." \
+    --video_save_name=Cosmos-1.0-Autoregressive-13B-Video2World \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-13B-Video2World \
+    --top_p=0.8 \
+    --temperature=1.0 \
+    --offload_guardrail_models \
+    --offload_diffusion_decoder \
+    --offload_ar_model \
+    --offload_tokenizer \
+    --offload_text_encoder_model
+```
+
+##### Batch Generation
+
+```bash
+# Example using 5B model
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/video2world.py \
+    --input_type=text_and_video \
+    --batch_input_path=cosmos1/models/autoregressive/assets/v1p0/batch_inputs/video2world.jsonl \
+    --video_save_folder=outputs/Cosmos-1.0-Autoregressive-5B-Video2World \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-5B-Video2World \
+    --top_p=0.7 \
+    --temperature=1.0
+
+# Example using 13B model
+CUDA_VISIBLE_DEVICES=0 PYTHONPATH=$(pwd) python cosmos1/models/autoregressive/inference/video2world.py \
+    --input_type=text_and_video \
+    --batch_input_path=cosmos1/models/autoregressive/assets/v1p0/batch_inputs/video2world.jsonl \
+    --video_save_folder=outputs/Cosmos-1.0-Autoregressive-13B-Video2World \
+    --ar_model_dir=Cosmos-1.0-Autoregressive-13B-Video2World \
+    --top_p=0.8 \
+    --temperature=1.0 \
+    --offload_guardrail_models
+```
+
+##### Example Output
+
+Here is an example output video generated using video2world.py with image input, using `Cosmos-1.0-Autoregressive-13B-Video2World`:
+
+<video src="https://github.com/user-attachments/assets/869f3b81-fabd-462e-a545-c04cdd9c1d22">
+  Your browser does not support the video tag.
+</video>
+
+The input image used to generate this video can be found in `cosmos1/models/autoregressive/assets/v1p0/input.jpg`. The prompt for generating the video is:
+
+```
+A driving video captures a serene urban street scene on a sunny day. The camera is mounted on the dashboard of a moving vehicle, providing a first-person perspective as it travels down a two-lane road. The street is lined with parked cars on both sides, predominantly black and silver sedans and SUVs. The road is flanked by a mix of residential and commercial buildings, with a prominent red-brick building on the left side, featuring multiple windows and a flat roof. The sky is clear with a few scattered clouds, casting soft shadows on the street. Trees with lush green foliage line the right side of the road, providing a natural contrast to the urban environment. The camera remains steady, maintaining a consistent forward motion, suggesting a leisurely drive. Traffic is light, with a few vehicles moving in the opposite direction, including a black sedan and a yellow taxi. Street signs are visible, including a no-parking sign on the right. The overall atmosphere is calm and peaceful, with no pedestrians visible, emphasizing the focus on the drive and the surrounding urban landscape.
+```
+
+Here is an example output video generated using video2world.py with 9-frame video input, using `Cosmos-1.0-Autoregressive-13B-Video2World`:
+
+<video src="https://github.com/user-attachments/assets/81840e1c-624b-4b01-9240-ab7db3722e58">
+  Your browser does not support the video tag.
+</video>
+
+The input video used to generate this video can be found in `cosmos1/models/autoregressive/assets/v1p0/input.mp4`. The prompt for generating the video is:
+
+```
+A video recorded from a moving vehicle's perspective, capturing roads, buildings, landscapes, and changing weather and lighting conditions.
+```
+
+##### Inference Time and GPU Memory Usage
+
+These numbers may vary based on system specifications and are provided for reference only.
+
+| Offloading Strategy | Cosmos-1.0-Autoregressive-5B-Video2World | Cosmos-1.0-Autoregressive-13B-Video2World |
+|-------------|---------|---------|
+| No offloading | 66.2 GB | > 80 GB |
+| Guardrails | 58.7 GB | 76.6 GB |
+| Guardrails & T5 encoder | 41.3 GB | 58.0 GB |
+| Guardrails & T5 encoder & Diffusion decoder | 29.0 GB | 46.9 GB |
+| Guardrails & T5 encoder & Diffusion decoder & Tokenizer | 28.8 GB | 46.7 GB |
+| Guardrails & T5 encoder & Diffusion decoder & Tokenizer & AR model | 21.1 GB | 30.9 GB |
+
+End-to-end inference runtime on one H100 with no offloading for 5B model and guardrail offloading for 13B, after model initialization:
+
+| Cosmos-1.0-Autoregressive-5B-Video2World | Cosmos-1.0-Autoregressive-13B-Video2World |
+|---------|---------|
+| ~73 seconds | ~150 seconds |
+
+### Arguments
+
+#### Common Parameters
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--checkpoint_dir` | Directory containing model weights | "checkpoints" |
+| `--video_save_name` | Output video filename for single video generation | "output" |
+| `--video_save_folder` | Folder where all output videos are stored | "outputs/" |
+| `--input_image_or_video_path` | Input image or video path. Required for single video generation | None |
+| `--batch_input_path` | Folder containing input images or videos. Required for batch video generation | None |
+| `--num_input_frames` | Number of input frames to use for Video2World prediction | 9 |
+| `--temperature` | Temperature used while sampling | 1.0 (recommend using values in sample commands provided) |
+| `--top_p` | Top-p value for top-p sampling | 0.8 (recommend using values in sample commands provided) |
+| `--seed` | Random seed | 0 |
+| `--disable_diffusion_decoder` | When set to True, use discrete tokenizer to decode discrete tokens to video. Otherwise, use diffusion decoder to decode video | False |
+| `--offload_guardrail_models` | Offload guardrail models after inference, used for low-memory GPUs | False |
+| `--offload_diffusion_decoder` | Offload diffusion decoder after inference, used for low-memory GPUs | False |
+| `--offload_ar_model` | Offload AR model after inference, used for low-memory GPUs | False |
+| `--offload_prompt_upsampler` | Offload prompt upsampler after inference, used for low-memory GPUs | False |
+
+#### Base Specific Parameters
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--ar_model_dir` | Directory containing AR model weight | "Cosmos-1.0-Autoregressive-4B" |
+| `--input_type` | Input type, either `video` or `image` | "video" |
+
+#### Video2World Specific Parameters
+
+| Parameter | Description | Default |
+|-----------|-------------|---------|
+| `--ar_model_dir` | Directory containing AR model weight | "Cosmos-1.0-Autoregressive-4B" |
+| `--input_type` | Input type, either `text_and_video` or `text_and_image` | "text_and_video" |
+| `--prompt` | Text prompt for single video generation. Required for single video generation | None |
+| `--input_prompts_path` | Path to JSONL file for batch video generation. Required for batch video generation | None |
+| `--offload_text_encoder_model` | Offload text encoder after inference, used for low-memory GPUs | False |
+
+### Safety Features
+
+The model uses a built-in safety guardrail system that cannot be disabled. Generating human faces is not allowed and will be blurred by the guardrail.
+
+For more information, check out the [Cosmos Guardrail Documentation](../guardrail/README.md).

cosmos1/models/autoregressive/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.