ckpt

[bugfix]: _compile_conditions regression (#1077 )
[misc] readme small fix (#1076 )
2026-02-09 05:50:48 +00:00 · 2026-02-06 18:57:43 -08:00 · 2026-02-06 16:57:48 -08:00 · 2026-02-06 19:43:18 -05:00 · 2026-02-06 13:38:35 -08:00 · 2026-02-05 20:33:08 -08:00
166 changed files with 14437 additions and 6438 deletions
@@ -1,35 +1,29 @@
-<div align="center">
-<img src=assets/logos/logo.svg width="30%"/>
-</div>

-<p align="center">
-     | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><b> Quick Start</b></a> | <a href="https://github.com/hao-ai-lab/FastVideo/discussions/982"  target="_blank"><b>Weekly Dev Meeting</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/sv3MMKyv" target="_blank"> <b> WeChat </b> </a> |
-</p>
+| **[Documentation](https://hao-ai-lab.github.io/FastVideo)** | **[Quick Start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/)** | **[Weekly Dev Meeting](https://github.com/hao-ai-lab/FastVideo/discussions/982)** | 🟣💬 **[Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ)** |

 **FastVideo is a unified post-training and inference framework for accelerated video generation.**

 ## NEWS
- ```2025/11/19```: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py)
- ```2025/08/04```: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).

-<details>
-<summary>More</summary>
+- `2025/11/19`: Release [CausalWan2.2 I2V A14B Preview](https://huggingface.co/FastVideo/CausalWan2.2-I2V-A14B-Preview-Diffusers) models, [Blog](https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/) and [Inference Code!](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_self_forcing_causal_wan2_2_i2v.py)
+- `2025/08/04`: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).

- ```2025/06/14```: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
- ```2025/04/24```: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
- ```2025/02/18```: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
+### More News

-</details>
+- `2025/06/14`: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
+- `2025/04/24`: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
+- `2025/02/18`: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).

 ## Key Features

 FastVideo has the following features:
+
 - End-to-end post-training support for bidirectional and autoregressive models:
  - Support full finetuning and LoRA finetuning for state-of-the-art open video DiTs
  - Data preprocessing pipeline for video, image, and text data
  - Distribution Matching Distillation (DMD2) stepwise distillation.
  - Sparse attention with [Video Sparse Attention](https://arxiv.org/pdf/2505.13389)
-  - [Sparse distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/) to achineve >50x denoising speedup
+  - [Sparse distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/) to achieve >50x denoising speedup
  - Scalable training with FSDP2, sequence parallelism, and selective activation checkpointing.
  - Causal distillation through Self-Forcing
  - See this [page](https://hao-ai-lab.github.io/FastVideo/training/overview/) for full list of supported models and recipes.
@@ -44,6 +38,7 @@ FastVideo has the following features:
  - See this [page](https://hao-ai-lab.github.io/FastVideo/inference/hardware_support/) for full list of supported hardware and OS.

 ## Getting Started
+
 We recommend using an environment manager such as `Conda` to create a clean environment:

 ```bash
@@ -58,18 +53,20 @@ pip install fastvideo
 Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions.

 ## Sparse Distillation
+
 For our sparse distillation techniques, please see our [distillation docs](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/) and check out our [blog](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).

 See below for recipes and datasets:

-|                                            Model                                              |                                               Sparse Distillation                                                 |                                                  Dataset                                                  |
-|:-------------------------------------------------------------------------------------------:  |:---------------------------------------------------------------------------------------------------------------:  |:--------------------------------------------------------------------------------------------------------: |
-| [FastWan2.1-T2V-1.3B](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers)         |    [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P)      | [FastVideo Synthetic Wan2.1 480P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)      |
-| [FastWan2.1-T2V-14B-Preview](https://huggingface.co/FastVideo/FastWan2.1-T2V-14B-Diffusers)   |                                                   Coming soon!                                                    |   [FastVideo Synthetic Wan2.1 720P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x768x1280_250k)   |
-| [FastWan2.2-TI2V-5B](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-Diffusers)           | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free)   | [FastVideo Synthetic Wan2.2 720P](https://huggingface.co/datasets/FastVideo/Wan2.2-Syn-121x704x1280_32k)  |
+| Model                                                                                 | Sparse Distillation                                                                                             | Dataset                                                                                                  |
+| ------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------- |
+| [FastWan2.1-T2V-1.3B](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers) | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P)       | [FastVideo Synthetic Wan2.1 480P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)     |
+| [FastWan2.2-TI2V-5B](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-Diffusers)   | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free) | [FastVideo Synthetic Wan2.2 720P](https://huggingface.co/datasets/FastVideo/Wan2.2-Syn-121x704x1280_32k) |

 ## Inference
+
 ### Generating Your First Video
+
 Here's a minimal example to generate a video using the default settings. Make sure VSA kernels are [installed](https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation/). Create a file called `example.py` with the following code:

 ```python
@@ -108,55 +105,36 @@ python example.py

 For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/).

-### Other docs:
+## More Guides

 - [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview/)
- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/)
-
-## Distillation and Finetuning
 - [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/)
-<!-- - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html) -->
+- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/contributing/overview/)

 ## Awesome work using FastVideo or our research projects

- [SGLang](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen): SGLang's diffusion inference functionality is based  on a fork of FastVideo on Sept. 24, 2025. [![Star](https://img.shields.io/github/stars/sgl-project/sglang.svg?style=social&label=Star)](https://github.com/sgl-project/sglang)
-
- [DanceGRPO](https://github.com/XueZeyue/DanceGRPO): A  unified framework to adapt Group Relative Policy Optimization (GRPO) to visual generation paradigms. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/XueZeyue/DanceGRPO.svg?style=social&label=Star)](https://github.com/XueZeyue/DanceGRPO)
- [SRPO](https://github.com/Tencent-Hunyuan/SRPO): A method to directly align the full diffusion trajectory with fine-grained human preference. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/Tencent-Hunyuan/SRPO.svg?style=social&label=Star)](https://github.com/Tencent-Hunyuan/SRPO)
- [DCM](https://github.com/Vchitect/DCM): Dual-expert consistency model for efficient and high-quality video generation. Code based on FastVideo. [![Star](https://img.shields.io/github/stars/Vchitect/DCM.svg?style=social&label=Star)](https://github.com/Vchitect/DCM)
- [Hunyuan Video 1.5](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5): A leading lightweight video generation model, where they proposed SSTA based on Sliding Tile Attention. [![Star](https://img.shields.io/github/stars/Tencent-Hunyuan/HunyuanVideo-1.5.svg?style=social&label=Star)](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5)
- [Kandinsky-5.0](https://github.com/kandinskylab/kandinsky-5): A family of diffusion models for video & image generation, where their NABLA attention includes a Sliding Tile Attention branch. [![Star](https://img.shields.io/github/stars/kandinskylab/kandinsky-5.svg?style=social&label=Star)](https://github.com/kandinskylab/kandinsky-5)
- [LongCat Video](https://github.com/meituan-longcat/LongCat-Video): A foundational video generation model with 13.6B parameters with block-sparse attention similar to Video Sparse Attention. [![Star](https://img.shields.io/github/stars/meituan-longcat/LongCat-Video.svg?style=social&label=Star)](https://github.com/meituan-longcat/LongCat-Video)
+- [SGLang](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen): SGLang's diffusion inference functionality is based on a fork of FastVideo on Sept. 24, 2025.
+- [DanceGRPO](https://github.com/XueZeyue/DanceGRPO): A unified framework to adapt Group Relative Policy Optimization (GRPO) to visual generation paradigms. Code based on FastVideo.
+- [SRPO](https://github.com/Tencent-Hunyuan/SRPO): A method to directly align the full diffusion trajectory with fine-grained human preference. Code based on FastVideo.
+- [DCM](https://github.com/Vchitect/DCM): Dual-expert consistency model for efficient and high-quality video generation. Code based on FastVideo.
+- [Hunyuan Video 1.5](https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5): A leading lightweight video generation model, where they proposed SSTA based on Sliding Tile Attention.
+- [Kandinsky-5.0](https://github.com/kandinskylab/kandinsky-5): A family of diffusion models for video & image generation, where their NABLA attention includes a Sliding Tile Attention branch.
+- [LongCat Video](https://github.com/meituan-longcat/LongCat-Video): A foundational video generation model with 13.6B parameters with block-sparse attention similar to Video Sparse Attention.

 ## 🤝 Contributing

 We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview/).
 See details in [development roadmap](https://github.com/hao-ai-lab/FastVideo/issues/899).
-## Acknowledgement
-We learned and reused code from the following projects:
- [Wan-Video](https://github.com/Wan-Video)
- [ThunderKittens](https://github.com/HazyResearch/ThunderKittens)
- [Triton](https://github.com/triton-lang/triton)
- [DMD2](https://github.com/tianweiy/DMD2)
- [diffusers](https://github.com/huggingface/diffusers)
- [xDiT](https://github.com/xdit-project/xDiT)
- [vLLM](https://github.com/vllm-project/vllm)
- [SGLang](https://github.com/sgl-project/sglang)

-We thank [MBZUAI](https://ifm.mbzuai.ac.ae/), [Anyscale](https://www.anyscale.com/), and [GMI Cloud](https://www.gmicloud.ai/) for their support throughout this project.
+## Acknowledgement
+
+We learned the design and reused code from the following projects: [Wan-Video](https://github.com/Wan-Video), [ThunderKittens](https://github.com/HazyResearch/ThunderKittens), [DMD2](https://github.com/tianweiy/DMD2), [diffusers](https://github.com/huggingface/diffusers), [xDiT](https://github.com/xdit-project/xDiT), [vLLM](https://github.com/vllm-project/vllm), [SGLang](https://github.com/sgl-project/sglang). We thank [MBZUAI](https://ifm.mbzuai.ac.ae/), [Anyscale](https://www.anyscale.com/), and [GMI Cloud](https://www.gmicloud.ai/) for their support throughout this project.

 ## Citation
-If you find FastVideo useful, please considering citing our work:
+
+If you find FastVideo useful, please consider citing our research work:

 ```bibtex
-@software{fastvideo2024,
-  title        = {FastVideo: A Unified Framework for Accelerated Video Generation},
-  author       = {The FastVideo Team},
-  url          = {https://github.com/hao-ai-lab/FastVideo},
-  month        = apr,
-  year         = {2024},
-}
-
@article{zhang2025vsa,
  title={Vsa: Faster video diffusion with trainable sparse attention},
  author={Zhang, Peiyuan and Chen, Yongqi and Huang, Haofeng and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp310-cp310-linux_x86_64.whl
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp310-cp310-linux_x86_64.whl

 COPY . .

@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp311-cp311-linux_x86_64.whl
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp311-cp311-linux_x86_64.whl

 COPY . .

@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.5.4/flash_attn-2.8.3%2Bcu128torch2.9-cp312-cp312-linux_x86_64.whl
+    uv pip install --no-cache-dir https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.7.16/flash_attn-2.8.3+cu128torch2.10-cp312-cp312-linux_x86_64.whl

 COPY . .

@@ -121,6 +121,15 @@ This page contains the complete API reference for the FastVideo library.
      show_root_toc_entry: true
      heading_level: 4

+## fastvideo.registry
+
+::: fastvideo.registry
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 3
+
 ## fastvideo.pipelines

 ::: fastvideo.pipelines
@@ -6,6 +6,8 @@ FastVideo provides highly optimized custom attention kernels to accelerate video

 * **[Video Sparse Attention (VSA)](vsa/index.md)**: Sparse attention mechanism selecting top-k blocks.
 * **[Sliding Tile Attention (STA)](sta/index.md)**: Optimized attention for window-based video generation.
+* **Backend development guide**: See the developer guide at
+  [Attention Backend Development](../contributing/attention_backend.md).

 ## General Build Instructions

@@ -0,0 +1,176 @@
+# Attention Backend Development
+
+This guide is for contributors adding a new attention backend (kernel or
+implementation) to FastVideo. If you just want to use existing kernels or build
+fastvideo-kernel, see [Attention overview](../attention/index.md).
+
+## When you need this guide
+
+Use this guide when you are:
+
+- Adding a new attention kernel or algorithm.
+- Wiring an existing kernel into FastVideo's attention selection.
+- Extending attention support to a new platform.
+
+## 0) Choose a backend name and scope
+
+Pick a backend name in `UPPER_SNAKE_CASE` and decide where it should run.
+Example: `MY_NEW_ATTN`.
+
+You will use this name in:
+
+- `AttentionBackendEnum` (global list of backends).
+- `get_name()` in your backend class (match the enum name).
+- Platform selectors (CUDA/ROCm/MPS/NPU) to return your backend class.
+
+## 1) Add enum + platform selection
+
+1) Add your backend to `fastvideo/platforms/interface.py`:
+
+```python
+class AttentionBackendEnum(enum.Enum):
+    ...
+    MY_NEW_ATTN = enum.auto()
+```
+
+1) Register it in platform selection (example: CUDA). Update
+`fastvideo/platforms/cuda.py` inside `get_attn_backend_cls`:
+
+```python
+elif selected_backend == AttentionBackendEnum.MY_NEW_ATTN:
+    try:
+        from fastvideo.attention.backends.my_new_attn import MyNewAttnBackend
+        return "fastvideo.attention.backends.my_new_attn.MyNewAttnBackend"
+    except ImportError as e:
+        logger.error("Failed to import MY_NEW_ATTN backend: %s", str(e))
+        raise
+```
+
+If you want support on other platforms, add a similar branch in
+`fastvideo/platforms/rocm.py`, `fastvideo/platforms/mps.py`, or `fastvideo/platforms/npu.py`.
+
+## 2) Implement the backend
+
+Create `fastvideo/attention/backends/my_new_attn.py` and implement the required
+classes.
+
+Minimal skeleton (no custom metadata):
+
+```python
+import torch
+from dataclasses import dataclass
+from fastvideo.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+
+class MyNewAttnBackend(AttentionBackend):
+    @staticmethod
+    def get_name() -> str:
+        return "MY_NEW_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> type["MyNewAttnImpl"]:
+        return MyNewAttnImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        return AttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+        return MyNewAttnMetadataBuilder
+
+
+@dataclass
+class MyNewAttnMetadata(AttentionMetadata):
+    current_timestep: int
+
+
+class MyNewAttnMetadataBuilder(AttentionMetadataBuilder):
+    def __init__(self) -> None:
+        pass
+
+    def prepare(self) -> None:
+        pass
+
+    def build(self, current_timestep: int, **kwargs):
+        return MyNewAttnMetadata(current_timestep=current_timestep)
+
+
+class MyNewAttnImpl(AttentionImpl):
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        softmax_scale: float,
+        causal: bool = False,
+        num_kv_heads: int | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.softmax_scale = softmax_scale
+        self.causal = causal
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: MyNewAttnMetadata,
+    ) -> torch.Tensor:
+        # Implement attention
+        return torch.nn.functional.scaled_dot_product_attention(
+            query.transpose(1, 2),
+            key.transpose(1, 2),
+            value.transpose(1, 2),
+            is_causal=self.causal,
+            scale=self.softmax_scale,
+        ).transpose(1, 2)
+```
+
+Optional:
+
+- Implement `preprocess_qkv` / `postprocess_output` if your kernel needs tiling
+  or reshaping.
+- Use `fastvideo.forward_context.get_forward_context()` if you need dynamic
+  per-step data (e.g., window sizes).
+- Set `accept_output_buffer = True` if your backend writes into a provided
+  output buffer.
+
+## 3) Wire into attention layers
+
+Backends are used by `LocalAttention` and `DistributedAttention`. These layers
+accept a `supported_attention_backends` tuple. If your backend should be
+eligible, update the call sites that construct these layers (search for
+`supported_attention_backends=`).
+
+## 4) Add compiled kernels (optional)
+
+If you have a custom CUDA kernel:
+
+1) Add sources in `fastvideo-kernel/csrc/attention/`.
+2) Register bindings in `fastvideo-kernel/csrc/common_extension.cpp`.
+3) Add to `fastvideo-kernel/CMakeLists.txt` (and any feature flags).
+4) Expose in `fastvideo-kernel/python/fastvideo_kernel/ops.py`.
+5) Export in `fastvideo-kernel/python/fastvideo_kernel/__init__.py`.
+
+Keep a Python/Triton fallback so the backend runs even when the kernel is not
+available.
+
+## 5) Testing and debugging
+
+- Add a small parity test or microbenchmark comparing to SDPA.
+- Force your backend with the env var:
+  `FASTVIDEO_ATTENTION_BACKEND=MY_NEW_ATTN`.
+- Check logs from `fastvideo/attention/selector.py` to confirm selection.
+
+## Checklist
+
+- [ ] Added enum entry in `fastvideo/platforms/interface.py`.
+- [ ] Implemented backend in `fastvideo/attention/backends/`.
+- [ ] Registered selection in platform(s).
+- [ ] Updated layer call sites to include the backend where appropriate.
+- [ ] Added tests and documentation.
@@ -0,0 +1,480 @@
+# FastVideo + Coding Agents
+
+Coding agents are now strong at navigating large codebases and iterating fast
+with parity tests and examples. This guide shows how to use them to add new
+model pipelines and ship PRs in a production-grade video diffusion framework.
+
+FastVideo is a great project to contribute to, with production-grade
+infrastructure, active collaborations (including NVIDIA), and a pipeline design
+and inference architecture that has been forked by [SGLang’s
+multimodal generation stack](https://github.com/sgl-project/sglang/tree/main/python/sglang/multimodal_gen).
+
+Goal: run the new pipeline with a minimal script like
+`examples/inference/basic/basic.py`. In production, FastVideo can download
+models automatically via `HF_HOME`; for development, use local directories so
+agents can run scripts and tests deterministically. We standardize local paths
+as:
+
+- `official_weights/<model_name>/` for official checkpoints
+- `converted_weights/<model_name>/` if conversion is required
+
+## Tips when prompting the agent
+
+When prompting the agent, include:
+
+- This guide and the [FastVideo design overview](../design/overview.md).
+- Exact file paths to edit.
+- A closest reference example file in FastVideo.
+- Expected behavior and acceptance criteria.
+- Repro steps (command, inputs, logs).
+- Constraints (performance, memory, compatibility).
+- Local paths (e.g., `official_weights/<model_name>/` or
+  `converted_weights/<model_name>/`) for parity tests.
+
+## FastVideo structure at a glance
+
+Before diving in, scan these references:
+
+- [Contributing overview](overview.md) for environment/setup context.
+- [FastVideo design overview](../design/overview.md) for pipeline architecture, configs, and HF layout.
+
+FastVideo maps a Diffusers-style repo into a pipeline like:
+
+- `fastvideo/models/*`: model implementations (DiT, VAE, encoders, upsamplers).
+- `fastvideo/configs/models/*`: arch configs and `param_names_mapping` for
+  weight name translation.
+- `fastvideo/configs/pipelines/*`: pipeline wiring (component classes + names).
+- `fastvideo/configs/sample/*`: default runtime sampling parameters.
+- `fastvideo/pipelines/basic/*`: end-to-end pipeline logic built from stages.
+- `model_index.json`: the HF repo entrypoint that maps component names to
+  classes and weight files.
+- Component loading happens in `VideoGenerator.from_pretrained`, which reads
+  `model_index.json`, resolves configs, and loads weights.
+
+Minimal usage example (based on `examples/inference/basic/basic.py`):
+
+```python
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam
+
+model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # or official_weights/<model_name>/
+generator = VideoGenerator.from_pretrained(model_id, num_gpus=1)
+
+sampling = SamplingParam.from_pretrained(model_id)
+sampling.num_frames = 45
+video = generator.generate_video(
+    "A vibrant city street at sunset.",
+    sampling_param=sampling,
+    output_path="video_samples",
+    save_video=True,
+)
+```
+
+## Some questions to ask yourself before starting
+
+Answering these upfront clarifies the work and speeds up implementation.
+
+### Is the model already supported by SGLang's multimodal generation stack?
+If yes, you can port many components from SGLang. It is a FastVideo fork, so
+interfaces line up, but you still need to swap layers/modules to match
+FastVideo's architecture and attention stack.
+
+If not, implement the model directly in FastVideo.
+
+### Is there an official implementation of the model you are adding?
+
+If yes, use it as the numerical reference. For example, LTX‑2 has an official
+implementation here: https://github.com/Lightricks/LTX-2. Prefer official code
+even if Diffusers also has one.
+
+### Is there a HuggingFace repo for the model you are adding? Is it in Diffusers format?
+
+If yes, load it directly in FastVideo after setting tensor mapping rules in the
+config. Otherwise, convert the weights to Diffusers format. See [Weights and
+Diffusers format](../design/overview.md#weights-and-diffusers-format) for details.
+
+### Do I have official weights + local paths ready?
+
+Standardize local paths as:
+
+- `official_weights/<model_name>/` for official checkpoints
+- `converted_weights/<model_name>/` if conversion is required (can be created later)
+
+### What pipeline components are required for the model you are adding?
+
+Usually you need a transformer (DiT), VAE, text encoder, and tokenizer. Some
+models add extra components.
+
+### What tasks does the model support?
+
+Usually a video diffusion model supports text‑to‑video (T2V),
+image‑to‑video (I2V), and video‑to‑video (V2V). Some add extra tasks (two‑stage
+generation, keyframe interpolation), which require extra components.
+
+It's usually easiest to start with a T2V pipeline and add the other tasks later.
+
+You can refer to the [Pipeline system](../design/overview.md#pipeline-system)
+section for more details.
+
+### Am I able to generate videos with the official implementation?
+
+These videos and prompts are your reference. Once the FastVideo pipeline works,
+compare outputs to the official implementation. Due to seeding and other
+factors, outputs may not match exactly, but they should be comparable.
+
+## Workflow: adding a full pipeline
+
+This is an example workflow for adding a full model pipeline (model + configs +
+examples + tests). This guide is in active development; feedback is welcome.
+
+!!! note
+    If you get stuck, refer to existing models/pipelines in FastVideo or ask in Slack.
+
+### 0) Fetch official model's code and weights
+
+Purpose:
+
+- Keep official checkpoints and source code local so conversion, parity tests,
+  and reference runs are reproducible.
+- Clone the official repo so you can use it as a numerical reference.
+
+Action:
+
+- Download official weights into `official_weights/<model_name>/`
+  (Diffusers format or not).
+- Clone the official repo under the project root (e.g., `FastVideo/LTX-2/`).
+- If a Diffusers-format HF repo already exists, you can skip manual weight
+  handling and download it directly with
+  `scripts/huggingface/download_hf.py`.
+
+!!! note
+    This step is best done manually because large downloads can time out.
+    Example:
+    ```bash
+    python scripts/huggingface/download_hf.py \
+      --repo_id Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+      --local_dir official_weights/Wan2.1-T2V-1.3B-Diffusers \
+      --repo_type model
+    ```
+  
+### 1) Implement the model + config mapping
+
+Purpose:
+
+- Model weights are a dictionary of named tensors (`state_dict`). If the names
+  don’t line up with FastVideo’s module names, weights won’t load correctly (or
+  will silently load into the wrong layer).
+- Official checkpoints often use different prefixes or module layouts than
+  FastVideo, so we translate names via the mapping (during load or conversion).
+- Mapping aligns three things:
+  1. the official implementation’s module names,
+  2. the checkpoint `state_dict` keys,
+  3. FastVideo’s model classes and layer naming conventions.
+- If names don’t align, weights won’t load; implement the FastVideo model and
+  define mapping rules first.
+
+Action:
+
+- Implement the FastVideo model + config mapping.
+  - Add/extend the model in `fastvideo/models/...` and config in
+    `fastvideo/configs/models/...` (including `param_names_mapping`).
+  - Reuse existing FastVideo layers/modules where possible.
+  - Use FastVideo’s attention layers:
+    - `DistributedAttention` only for full‑sequence self‑attention in the DiT.
+    - `LocalAttention` for cross‑attention and other attention layers.
+  - See the “Configuration System” and “Weights and Diffusers format” sections
+    in `docs/design/overview.md` for how these pieces connect.
+  - If you are using an agent, ask it to implement the model, config mapping,
+    and a parity test together so you can validate numerics immediately.
+
+!!! note
+    After the first component is aligned and parity‑tested, open a **DRAFT PR**
+    on FastVideo so the rest of the pipeline work can build on top of it.
+
+!!! note
+    If a Diffusers-format HF repo already exists and loads correctly, you can
+    skip conversion entirely (no conversion script needed) and just download it
+    with `scripts/huggingface/download_hf.py`. Otherwise, you may need a
+    conversion script + a `converted_weights/<model>/` staging directory.
+
+Example (key renaming via arch config mapping, Wan2.1‑style):
+
+```python
+# Official model (simplified) in the upstream repo.
+class OfficialWanTransformer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.patch_embedding = torch.nn.Conv3d(16, 1536, kernel_size=2, padding=0)
+
+    def forward(self, x):
+        return self.patch_embedding(x)
+
+# FastVideo model (simplified) in fastvideo/models/dits/wanvideo.py
+class PatchEmbed(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = torch.nn.Conv3d(16, 1536, kernel_size=2, padding=0)
+
+    def forward(self, x):
+        return self.proj(x)
+
+class WanTransformer3DModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.patch_embedding = PatchEmbed()
+
+    def forward(self, x):
+        return self.patch_embedding(x)
+
+# Mapping defined in a config (simplified; see the real mapping in
+# fastvideo/configs/models/dits/wanvideo.py)
+param_names_mapping = {
+    r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1",
+    r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": r"blocks.\1.to_q.\2",
+}
+
+def apply_regex_map(state_dict, mapping):
+    # Pseudocode: apply regex substitutions in order
+    ...
+
+# Official checkpoint keys (example)
+official = {
+    "patch_embedding.weight": ...,
+    "blocks.0.attn1.to_q.weight": ...,
+}
+
+# Apply mapping so keys match FastVideo modules
+converted = apply_regex_map(official, param_names_mapping)
+
+```
+
+Optional helper (print a few checkpoint keys quickly):
+
+```bash
+python - <<'PY'
+import safetensors.torch as st
+keys = list(st.load_file("official_weights/<model>/transformer/diffusion_pytorch_model.safetensors").keys())
+print(keys[:20])
+PY
+```
+
+Example agent prompt (task request):
+
+```
+Please add the Wan2.1 T2V 1.3B Diffusers pipeline to FastVideo:
+- Add a FastVideo native Wan2.1 DiT implementation + config mapping.
+- Make sure to use the existing FastVideo layers and attention modules where possible.
+- Add a parity test that loads the official model alongside the FastVideo model and compares outputs numerically with fixed seeds and inputs.
+
+Paths:
+  - Official repo: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+  - Local download: official_weights/Wan2.1-T2V-1.3B-Diffusers
+Mapping steps:
+  - Load the official DiT weights from
+    official_weights/Wan2.1-T2V-1.3B-Diffusers/transformer/diffusion_pytorch_model.safetensors.
+  - Instantiate the FastVideo DiT (`WanTransformer3DModel`) and compare
+    its `state_dict().keys()` to the official keys.
+  - Update `param_names_mapping` in
+    fastvideo/configs/models/dits/wanvideo.py to resolve missing/unexpected keys.
+  - Use `load_state_dict(strict=False)` during iteration to surface mismatches.
+```
+
+External examples of the same pattern:
+- SGLang uses prefix-based routing in its weight loader to map checkpoint keys
+  into internal submodules (e.g., stripping a top-level prefix before delegating).
+- vLLM includes model-specific renamers for certain checkpoints that adjust
+  key prefixes so weights match its internal naming.
+
+### 2) Test numerical alignment with the official implementation
+
+Purpose:
+
+- Verify that the FastVideo component is numerically aligned with the official
+  implementation.
+
+Action:
+
+- Add or reuse a numerical parity test that loads the official model and the
+  FastVideo model and compares outputs.
+- See examples in `tests/local_tests/` (e.g., `tests/local_tests/upsamplers/`)
+  and the commands in `tests/local_tests/README.md`.
+- If there are discrepancies, add opt‑in logging to both models and compare
+  activation summaries (layer output sums, per‑stage logs).
+- First align the loaded weights (validate `param_names_mapping`).
+- Then align forward outputs using fixed seeds and inputs.
+  - Start with `atol=1e-4, rtol=1e-4` in `assert_close`.
+  - Keep dtype consistent (bf16 if available; otherwise fp32).
+  - If attention parity is unstable, align backends (e.g.,
+    `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`).
+
+### 3) Repeat the process for each component
+
+If the model requires additional components, repeat Steps 1–2 for each one.
+For example, implement the VAE in `fastvideo/models/vaes/` and its config in
+`fastvideo/configs/models/vaes/`, then add parity coverage for it.
+
+### 4) Add a pipeline config + sample defaults
+
+Purpose:
+
+- `fastvideo/configs/pipelines/` describes pipeline wiring and model module
+  names.
+- `fastvideo/configs/sample/` defines default runtime parameters.
+
+Action:
+
+- Add a new pipeline config + sampling params.
+- Register them in `fastvideo/registry.py` using explicit
+  `register_configs(...)` blocks (this file is the single source of truth now).
+
+### 5) Wire pipeline stages
+
+Purpose:
+
+- `fastvideo/pipelines/basic/<pipeline>/` contains the actual pipeline logic.
+- `fastvideo/pipelines/stages/` holds reusable, testable stages.
+
+Action:
+
+- Build the pipeline using stages; keep new stages isolated and documented.
+- Prefer opt‑in flags for expensive or optional steps.
+
+### 6) Add pipeline‑level tests
+
+Purpose:
+
+- Ensure the end‑to‑end pipeline works and stays aligned as pieces evolve.
+
+Action:
+
+- Add a pipeline parity test under `tests/local_tests/pipelines/`.
+- See the [Testing Guide](testing.md) for test conventions.
+
+### 7) Add user‑facing examples
+
+Purpose:
+
+- `examples/inference/basic/` is the entry point for simple, runnable scripts.
+
+Action:
+
+- Provide a minimal “hello world” example plus advanced variations.
+- Use fixed seeds and stable prompts.
+- Run the example locally to confirm end‑to‑end behavior.
+
+### 8) Add SSIM tests for CI checks
+
+Purpose:
+
+- Ensure visual similarity stays within expected bounds for regression testing.
+- SSIM tests act as a higher‑level guardrail beyond unit/parity tests.
+
+Action:
+
+- Add SSIM tests under `fastvideo/tests/ssim/` and include reference videos
+  (see the structure in the Testing Guide).
+- Use stable prompts/seeds and document any GPU‑specific requirements.
+- Follow the [Testing Guide](testing.md) for reference video placement and
+  execution details.
+
+### 9) Document it
+
+Purpose:
+
+- `docs/` is where users find the new pipeline usage and limitations.
+
+Action:
+
+- Add a short doc page or update an existing one.
+- Mention any caveats (memory, speed, constraints).
+
+## Common pitfalls when porting models
+
+- **Attention backend mismatch**: parity can fail if the official model uses a
+  different attention backend (e.g., SDPA vs custom). Align backends before
+  debugging deeper issues.
+- **Patchifier shape mistakes**: wrong patchification or reshape lengths can
+  silently corrupt outputs. Validate patch shapes early.
+- **Mask handling**: attention masks must match the official behavior (padding,
+  causal masks, and broadcast shapes).
+- **Scheduler / sigma schedule mismatch**: even small differences in schedules
+  or timestep shapes can cause noticeable drift.
+
+## Diffusers vs manual conversion
+
+If a model already ships in Diffusers format (with a proper `model_index.json`),
+prefer downloading it directly and loading it via FastVideo. In that case:
+
+- You usually **do not need** a conversion script.
+- You still need a correct `param_names_mapping` if the internal module names
+  differ from FastVideo’s implementation.
+
+If the model does **not** have a Diffusers-format repo:
+
+- You will need a conversion script to rewrite `state_dict` keys into FastVideo
+  naming and stage the result (e.g., under `converted_weights/<model>/`).
+- You may still use the official repo for reference parity and debugging.
+
+In both cases, parity testing is required to validate correctness.
+
+If you want to publish a Diffusers‑style repo after conversion, use
+`scripts/checkpoint_conversion/create_hf_repo.py` to assemble a HuggingFace‑ready
+directory before uploading.
+
+## FAQ
+
+**Q: Why do we implement the FastVideo model before conversion?**  
+A: You can’t define the key‑mapping rules until the FastVideo module names are
+known. The implementation determines the target `state_dict` schema.
+
+**Q: Do we always need a conversion script?**  
+A: No. If a Diffusers‑format repo exists and loads correctly, download it and
+skip conversion.
+
+**Q: How do I figure out `param_names_mapping` quickly?**  
+A: Load the official weights, instantiate the FastVideo model, and diff
+`state_dict().keys()` on both sides. Add regex rules until missing/unexpected
+keys are resolved. Agents can help you with this.
+
+**Q: What if parity fails even after mapping?**  
+A: Align attention backends, sigma schedules, and timestep shapes first. Then
+add opt‑in activation logging to locate the first divergent layer.
+
+## Case study: LTX‑2 port (from PLAN.md)
+
+The LTX‑2 port in `PLAN.md` shows the real sequence of steps and backtracking
+that happened during integration. Use it as a reference for how parity work
+actually unfolds:
+
+- Ported components first (transformer, VAE, audio, text encoder).
+- Added parity tests per component; used SDPA for reference parity.
+- Added debug logging to compare per‑block activations and isolate divergence.
+- Fixed cross‑attention reshape and patch grid bounds issues after logging.
+- Aligned sigma schedule and masking behavior to match the official pipeline.
+
+Recommendation:
+
+- Keep raw step‑by‑step logs in your own local `PLAN.md` for large ports.
+
+## Worked example: Wan2.1 T2V 1.3B pipeline
+
+The Wan2.1 T2V 1.3B Diffusers pipeline is a good “standard” example for
+FastVideo integration.
+
+1. Verify model config + mapping.
+   - DiT mapping: `fastvideo/configs/models/dits/wanvideo.py`
+   - VAE: `fastvideo/models/vaes/wanvae.py`
+   - Text encoder: `fastvideo/models/encoders/t5.py`
+
+2. Parity test the core components.
+   - Example tests: `fastvideo/tests/transformers/test_wanvideo.py`,
+     `fastvideo/tests/vaes/test_wan_vae.py`,
+     `fastvideo/tests/encoders/test_t5_encoder.py`
+
+3. Pipeline wiring.
+   - Pipeline: `fastvideo/pipelines/basic/wan/wan_pipeline.py`
+   - Pipeline config: `fastvideo/configs/pipelines/wan.py`
+   - Sampling defaults: `fastvideo/configs/sample/wan.py`
+
+4. Minimal example.
+   - Script: `examples/inference/basic/basic.py`
@@ -1,16 +1,55 @@

 # 📦 Developing FastVideo on RunPod

-You can easily use the FastVideo Docker image as a custom container on [RunPod](https://www.runpod.io) for development or experimentation.
+You can easily use the FastVideo Pod Template on [RunPod](https://www.runpod.io) for development or experimentation.

 ## Creating a new pod

-Choose a GPU that supports CUDA 12.8
-
-Pick 1 or 2 L40S GPU(s)
+- Make sure you are using the correct RunPod account.
+![RunPod Account Selection](../../assets/images/runpod_account.png)

+- Use "Additional Filters" to select CUDA 12.8.
 ![RunPod CUDA selection](../../assets/images/runpod_cuda.png)

+- Click "Deploy" and Pick a single A40 or RTX 4090 GPU.
+![RunPod GPU Selection](../../assets/images/runpod_deploy.png)
+
+- Select the "FastVideo" or "fastvideo-dev" Pod Template.
+![RunPod Pod Template Selection](../../assets/images/runpod_create.png)
+
+- Set the Pod name to "`<name>-<FastVideo>-<date>`".
+
+- Finally, once the pod is deployed (will take a few minutes as the image is being pulled), you can SSH into it using "SSH exposed over TCP". You'll need to use the matching private ssh key you provided.
+![RunPod SSH](../../assets/images/runpod_ssh.png)
+
+## Working with the pod
+
+After SSH'ing into your pod, you'll find the correct `uv` environment already activated and you should be in /FastVideo/ directory. Make sure to use /FastVideo/ for all your work.
+
+To pull in the latest changes from the GitHub repo:
+
+```bash
+cd /FastVideo
+git pull
+```
+
+Run your development workflows as usual:
+
+```bash
+# Run linters
+pre-commit run --all-files
+
+# Run tests
+pytest tests/
+```
+
+Make sure to push your changes back to the GitHub repo as nothing will be saved to the pod when it is terminated.
+
+After you are done with your work, you can terminate the pod by clicking the "Terminate" and "Delete" buttons. Remember if the pod is not completely deleted, Runpod will keep charging you for it.
+
+## Extra Information:
+If you need to customize the pod template this section has some useful information. For the most part you can leave the defaults of the FastVideo Pod Template.
+
 When creating your pod template, use this image:

 ```
@@ -24,30 +63,3 @@ bash -c "apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-serve
 ```

 ![RunPod template configuration](../../assets/images/runpod_template.png)
-
-After deploying, the pod will take a few minutes to pull the image and start the SSH service.
-
-![RunPod ssh](../../assets/images/runpod_ssh.png)
-
-## Working with the pod
-
-After SSH'ing into your pod, you'll find the `fastvideo-dev` Conda environment already activated.
-
-To pull in the latest changes from the GitHub repo:
-
-```bash
-cd /FastVideo
-git pull
-```
-
-`If you have a persistent volume and want to keep your code changes, you can move /FastVideo to /workspace/FastVideo, or simply clone the repository there.`
-
-Run your development workflows as usual:
-
-```bash
-# Run linters
-pre-commit run --all-files
-
-# Run tests
-pytest tests/
-```
@@ -1,76 +1,84 @@

 # 🛠️ Contributing to FastVideo

-Thank you for your interest in contributing to FastVideo. We want to make the process as smooth for you as possible and this is a guide to help get you started!
+Thank you for your interest in contributing to FastVideo. We want the process
+to be smooth and beginner‑friendly, whether you are adding a new pipeline,
+improving performance, or fixing a bug.

-Our community is open to everyone and welcomes any contributions no matter how large or small.
+## Quick prerequisites

-# Developer Environment:
-Do make sure you have CUDA 12.4 installed and supported. FastVideo currently only supports Linux and CUDA GPUs, but we hope to support other platforms in the future.
+- **OS**: Linux is the primary development target (WSL can work).
+- **GPU**: NVIDIA GPU recommended for inference and training workflows.
+- **CUDA**: Use a recent CUDA 12.x toolchain (see the installation guide for
+  the current recommendation).

-We recommend using a fresh Python 3.10 Conda environment to develop FastVideo:
+For a full install checklist, see `docs/getting_started/installation/gpu.md`.
+
+## Local development (Conda + editable install)

 Install Miniconda:

-```
+```bash
 wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
 bash Miniconda3-latest-Linux-x86_64.sh
 source ~/.bashrc
 ```

-Create and activate a Conda environment for FastVideo:
+Create and activate a Conda environment:

-```
+```bash
 conda create -n fastvideo python=3.12 -y
 conda activate fastvideo
 ```

 Install `uv` (optional, but recommended):

-From instructions on [uv](https://astral.sh/uv/):
-
-```
+```bash
 curl -LsSf https://astral.sh/uv/install.sh | sh
-# or 
+# or
 wget -qO- https://astral.sh/uv/install.sh | sh
 ```

-Clone the FastVideo repository and go to the FastVideo directory:
+Clone the repo:

-```
+```bash
 git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
-
 ```

-Now you can install FastVideo and setup git hooks for running linting. By using `pre-commit`, the linters will run and have to pass before you'll be able to make a commit.
+Install FastVideo in editable mode and set up hooks:

 ```bash
 uv pip install -e .[dev]

-# Can also install flash-attn (optional)
-uv pip install flash-attn --no-build-isolation 
+# Optional: FlashAttention (builds native kernels)
+uv pip install flash-attn --no-build-isolation

-# Linting, formatting and static type checking
+# Linting, formatting, static typing
 pre-commit install --hook-type pre-commit --hook-type commit-msg
-
-# You can manually run pre-commit with
 pre-commit run --all-files

 # Unit tests
 pytest tests/
 ```

-If you are on a Hopper GPU, you should also install [FA3](https://github.com/Dao-AILab/flash-attention) for much better performance:
+If you are on a Hopper GPU, installing FlashAttention 3 can improve
+performance (see `docs/inference/optimizations.md`).

-```
-git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention/hopper
+## Docker development (optional)

-# make sure you have ninja installed
-uv pip install ninja
-
-python setup.py install
-```
+If you prefer a containerized environment, use the dev image documented in
+`docs/contributing/developer_env/docker.md`.

 ## Testing

-Please refer to the [Testing Guide](testing.md) for more information on how to add and run tests in FastVideo.
+See the [Testing Guide](testing.md) for how to add and run tests in FastVideo.
+
+## Attention backend development
+
+If you are adding a new attention kernel or backend, follow
+[Attention Backend Development](attention_backend.md).
+
+## Contributing with coding agents
+
+For a step‑by‑step workflow on adding pipelines or components with coding
+agents, see `docs/contributing/coding_agents.md`.
@@ -1,424 +1,179 @@
-# 🔍 FastVideo Overview
+# FastVideo Architecture Overview

-This document outlines FastVideo's architecture for developers interested in framework internals or contributions. It serves as an onboarding guide for new contributors by providing an overview of the most important directories and files within the `fastvideo/` codebase.
+This document summarizes how FastVideo is structured and how a Diffusers-style
+model repo maps into a runnable pipeline. It is intended for contributors who
+need the high-level layout and key entrypoints, not every internal detail.

-## Table of Contents - Directory Structure and Files
+## FastVideo structure at a glance

- [`fastvideo/pipelines/`](#pipeline-system) - Core diffusion pipeline components
- [`fastvideo/models/`](#model-components) - Model implementations
-  - [`dits/`](#transformer-models) - Transformer-based diffusion models
-  - [`vaes/`](#vae-variational-auto-encoder) - Variational autoencoders
-  - [`encoders/`](#text-and-image-encoders) - Text and image encoders
-  - [`schedulers/`](#schedulers) - Diffusion schedulers
- [`fastvideo/attention/`](#optimized-attention) - Optimized attention implementations
- [`fastvideo/distributed/`](#distributed-processing) - Distributed computing utilities
- [`fastvideo/layers/`](#tensor-parallelism) - Custom neural network layers
- [`fastvideo/platforms/`](#platforms) - Hardware platform abstractions
- [`fastvideo/worker/`](#executor-and-worker-system) - Multi-GPU process management
- [`fastvideo/fastvideo_args.py`](#fastvideoargs) - Argument handling
- [`fastvideo/forward_context.py`](#forward-context-management) - Forward pass context management
- `fastvideo/utils.py` - Utility functions
- [`fastvideo/logger.py`](#logger) - Logging infrastructure
+FastVideo maps a Diffusers-style repo into a pipeline like this:

-## Core Architecture
+- `fastvideo/models/*`: model implementations (DiT, VAE, encoders, upsamplers).
+- `fastvideo/configs/models/*`: arch configs and `param_names_mapping` for
+  weight name translation.
+- `fastvideo/configs/pipelines/*`: pipeline wiring (component classes + names).
+- `fastvideo/configs/sample/*`: default runtime sampling parameters.
+- `fastvideo/pipelines/basic/*`: end-to-end pipelines.
+- `fastvideo/pipelines/stages/*`: reusable pipeline stages.
+- `fastvideo/models/loader/*`: component loaders for Diffusers-style repos.
+- `model_index.json`: HF repo entrypoint mapping component names to classes.

-FastVideo separates model components from execution logic with these principles:
+Flow:
+`model_index.json` -> component loaders -> model modules -> pipeline stages ->
+sampling params.

- **Component Isolation**: Models (encoders, VAEs, transformers) are isolated from execution (pipelines, stages, distributed processing)
- **Modular Design**: Components can be independently replaced
- **Distributed Execution**: Supports various parallelism strategies (Tensor, Sequence)
- **Custom Attention Backends**: Components can support and use different Attention implementations
- **Pipeline Abstraction**: Consistent interface across diffusion models
-
-## FastVideoArgs
-
-The `FastVideoArgs` class in `fastvideo/fastvideo_args.py` serves as the central configuration system for FastVideo. It contains all parameters needed to control model loading, inference configuration, performance optimization settings, and more.
-
-Key features include:
-
- **Command-line Interface**: Automatic conversion between CLI arguments and dataclass fields
- **Configuration Groups**: Organized by functional areas (model loading, video params, optimization settings)
- **Context Management**: Global access to current settings via `get_current_fastvideo_args()`
- **Parameter Validation**: Ensures valid combinations of settings
-
-Common configuration areas:
-
- **Model paths and loading options**: `model_path`, `trust_remote_code`, `revision`
- **Distributed execution settings**: `num_gpus`, `tp_size`, `sp_size`
- **Video generation parameters**: `height`, `width`, `num_frames`, `num_inference_steps`
- **Precision settings**: Control computation precision for different components
-
-Example usage:
+Minimal usage (from `examples/inference/basic/basic.py`):

 ```python
-# Load arguments from command line
-fastvideo_args = prepare_fastvideo_args(sys.argv[1:])
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam

-# Access parameters
-model = load_model(fastvideo_args.model_path)
+model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # or official_weights/<model_name>/
+generator = VideoGenerator.from_pretrained(model_id, num_gpus=1)

-# Set as global context
-with set_current_fastvideo_args(fastvideo_args):
-    # Code that requires access to these arguments
-    result = generate_video()
-```
-
-## Pipeline System
-
-### `ComposedPipelineBase`
-
-This foundational class provides:
-
- **Model Loading**: Automatically loads components from HuggingFace-Diffusers-compatible model directories
- **Stage Management**: Creates and orchestrates processing stages
- **Data Flow Coordination**: Ensures proper state flow between stages
-
-```python
-class MyCustomPipeline(ComposedPipelineBase):
-    _required_config_modules = [
-        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
-    ]
-    
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        # Pipeline-specific initialization
-        pass
-        
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        self.add_stage("input_validation_stage", InputValidationStage())
-        self.add_stage("text_encoding_stage", CLIPTextEncodingStage(
-            text_encoder=self.get_module("text_encoder"),
-            tokenizer=self.get_module("tokenizer")
-        ))
-        # Additional stages...
-```
-
-### Pipeline Stages
-
-Each stage handles a specific diffusion process component:
-
- **Input Validation**: Parameter verification
- **Text Encoding**: CLIP, LLaMA, or T5-based encoding
- **Image Encoding**: Image input processing
- **Timestep & Latent Preparation**: Setup for diffusion
- **Denoising**: Core diffusion loop
- **Decoding**: Latent-to-pixel conversion
-
-Each stage implements a standard interface:
-
-```python
-def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
-    # Process batch and update state
-    return batch
-```
-
-![Pipeline execution and data flow](../assets/images/pipeline.png)
-
-### ForwardBatch
-
-Defined in `fastvideo/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:
-
- **Input Data**: Prompts, images, generation parameters
- **Intermediate State**: Embeddings, latents, timesteps, accumulated during stage execution
- **Output Storage**: Generated results and metadata
- **Configuration**: Sampling parameters, precision settings
-
-This structure facilitates clear state transitions between stages.
-
-## Model Components
-
-The `fastvideo/models/` directory contains implementations of the core neural network models used in video diffusion:
-
-### Transformer Models
-
-Transformer networks perform the actual denoising during diffusion:
-
- **Location**: `fastvideo/models/dits/`
- **Examples**:
-  - `WanTransformer3DModel`
-  - `HunyuanVideoTransformer3DModel`
-
-Features include:
-
- Text/image conditioning
- Standardized interface for model-specific optimizations
-
-```python
-def forward(
-    self, 
-    latents,                    # [B, T, C, H, W]
-    encoder_hidden_states,      # Text embeddings
-    timestep,                   # Current diffusion timestep
-    encoder_hidden_states_image=None,  # Optional image embeddings
-    **kwargs
-):
-    # Perform denoising computation
-    return noise_pred  # Predicted noise residual
-```
-
-### VAE (Variational Auto-Encoder)
-
-VAEs handle conversion between pixel space and latent space:
-
- **Location**: `fastvideo/models/vaes/`
- **Examples**:
-  - `AutoencoderKLWan`
-  - `AutoencoderKLHunyuanVideo`
-
-These models compress image/video data to a more efficient latent representation (typically 4x-8x smaller in each dimension).
-
-FastVideo's VAE implementations include:
-
- Efficient video batch processing
- Memory optimization
- Optional tiling for large frames
- Distributed weight support
-
-### Text and Image Encoders
-
-Encoders process conditioning inputs into embeddings:
-
- **Location**: `fastvideo/models/encoders/`
- **Text Encoders**:
-  - `CLIPTextModel`
-  - `LlamaModel`
-  - `UMT5EncoderModel`
- **Image Encoders**:
-  - `CLIPVisionModel`
-
-FastVideo implements optimizations such as:
-
- Vocab parallelism for distributed processing
- Caching for common prompts
- Precision-tuned computation
-
-### Schedulers
-
-Schedulers manage the diffusion sampling process:
-
- **Location**: `fastvideo/models/schedulers/`
- **Examples**:
-  - `UniPCMultistepScheduler`
-  - `FlowMatchEulerDiscreteScheduler`
-
-These components control:
-
- Diffusion timestep sequences
- Noise prediction to latent update conversions
- Quality/speed trade-offs
-
-```python
-def step(
-    self, 
-    model_output: torch.Tensor,
-    timestep: torch.LongTensor,
-    sample: torch.Tensor,
-    **kwargs
-) -> torch.Tensor:
-    # Process model output and update latents
-    # Return updated latents
-    return prev_sample
-```
-
-This diagram shows how models are discovered, validated, and loaded across entrypoints, executors, pipelines, and model loaders.
-
-![Model loading flow](../assets/images/load_models.png)
-
-## Optimized Attention
-
-The `fastvideo/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:
-
-### Attention Backends
-
-Multiple implementations with automatic selection:
-
- **FLASH_ATTN**: Optimized for supporting hardware
- **TORCH_SDPA**: Built-in PyTorch scaled dot-product attention
- **SLIDING_TILE_ATTN**: For very long sequences
-
-```python
-# Configure available attention backends for this layer
-self.attn = LocalAttention(
-    num_heads=num_heads,
-    head_size=head_dim,
-    causal=False,
-    supported_attention_backends=(_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
-)
-
-# Override via environment variable
-# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-```
-
-![Attention backend selector design](../assets/images/attention_backend.png)
-
-### Attention Patterns
-
-Supports various patterns with memory optimization techniques:
-
- **Cross/Self/Temporal/Global-Local Attention**
- Chunking, progressive computation, optimized masking
-
-## Distributed Processing
-
-The `fastvideo/distributed/` directory contains implementations for distributed model execution:
-
-### Tensor Parallelism
-
-Tensor parallelism splits model weights across devices:
-
- **Implementation**: Through `RowParallelLinear` and `ColumnParallelLinear` layers
- **Use cases**: Will be used by encoder models as their sequence lengths are shorter and enables efficient sharding.
-
-```python
-# Tensor-parallel layers in a transformer block
-from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear
-
-# Split along output dimension
-self.qkv_proj = ColumnParallelLinear(
-    input_size=hidden_size,
-    output_size=3 * hidden_size,
-    bias=True,
-    gather_output=False
-)
-
-# Split along input dimension
-self.out_proj = RowParallelLinear(
-    input_size=hidden_size,
-    output_size=hidden_size,
-    bias=True,
-    input_is_parallel=True
+sampling = SamplingParam.from_pretrained(model_id)
+sampling.num_frames = 45
+video = generator.generate_video(
+    "A vibrant city street at sunset.",
+    sampling_param=sampling,
+    output_path="video_samples",
+    save_video=True,
 )
 ```

-### Sequence Parallelism
+## Configuration system

-Sequence parallelism splits sequences across devices:
+FastVideo uses typed configs to keep model definitions, pipeline wiring, and
+runtime parameters consistent:

- **Implementation**: Through `DistributedAttention` and sequence splitting
- **Use cases**: Long video sequences or high-resolution processing. Used by DiT models.
+- `fastvideo/configs/models/`: architecture definitions, layer shapes, and
+  `param_names_mapping` rules for key renaming.
+- `fastvideo/configs/pipelines/`: pipeline wiring and required components.
+- `fastvideo/configs/sample/`: default sampling parameters (steps, frames,
+  guidance scale, resolution, fps).
+- `fastvideo/registry.py`: unified registry for pipeline config + sampling
+  defaults and model metadata resolution, defined via explicit
+  `register_configs(...)` blocks (no separate dict registries).

-```python
-# Distributed attention for long sequences
-from fastvideo.attention import DistributedAttention
+`FastVideoArgs` (in `fastvideo/fastvideo_args.py`) provides runtime settings and
+is passed into pipeline construction and stages.

-self.attn = DistributedAttention(
-    num_heads=num_heads,
-    head_size=head_dim,
-    causal=False,
-    supported_attention_backends=(_Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN)
-)
+## Weights and Diffusers format
+
+FastVideo follows the HuggingFace Diffusers repo layout. This keeps loaders
+compatible with HF repos and makes it easy to add new components.
+
+Typical Diffusers repo:
+
+```
+<model-repo>/
+  model_index.json
+  scheduler/
+    scheduler_config.json
+  transformer/               # or unet/ for image models
+    config.json
+    diffusion_pytorch_model.safetensors
+  vae/
+    config.json
+    diffusion_pytorch_model.safetensors
+  text_encoder/
+    config.json
+    model.safetensors
+  tokenizer/
+    tokenizer_config.json
+    tokenizer.json
 ```

-### Communication Primitives
+Key points:

-Efficient distributed operations via AllGather, AllReduce, and synchronization mechanisms.
+- `model_index.json` is the root map that tells FastVideo which components to
+  load and which classes implement them.
+- Each component lives in its own folder with a `config.json` and weights.
+- Weights are usually in `diffusion_pytorch_model.safetensors`.

-Efficient communication primitives minimize distributed overhead:
+Note on tensor names:

- **Sequence-Parallel AllGather**: Collects sequence chunks
- **Tensor-Parallel AllReduce**: Combines partial results
- **Distributed Synchronization**: Coordinates execution
+Official checkpoints often use different `state_dict` names than FastVideo's
+module layout. We translate tensor names via the DiT arch config mapping
+(`param_names_mapping` under `fastvideo/configs/models/dits/`). This is similar
+in spirit to name-translation layers used in systems like vLLM and SGLang.

-## Forward Context Management
+Example HF repo (Wan 2.1 T2V 1.3B Diffusers):

-### ForwardContext
-
-Defined in `fastvideo/forward_context.py`, `ForwardContext` manages execution-specific state *within* a forward pass, particularly for low-level optimizations. It is accessed via `get_forward_context()`.
-
- **Attention Metadata**: Configuration for optimized attention kernels (`attn_metadata`)
- **Profiling Data**: Potential hooks for performance metrics collection
-
-This context-based approach enables:
-
- Dynamic optimization based on execution state (e.g., attention backend selection)
- Step-specific customizations within model components
-
-Usage example:
-
-```python
-with set_forward_context(current_timestep, attn_metadata, fastvideo_args):
-    # During this forward pass, components can access context
-    # through get_forward_context()
-    output = model(inputs)
+```
+https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers/tree/main
 ```

-## Executor and Worker System
+Example `model_index.json` from that repo:

-The `fastvideo/worker/` directory contains the distributed execution framework:
-
-### Executor Abstraction
-
-FastVideo implements a flexible execution model for distributed processing:
-
- **Executor Base Class**: An abstract base class defining the interface for all executors
- **MultiProcExecutor**: Primary implementation that spawns and manages worker processes
- **GPU Workers**: Handle actual model execution on individual GPUs
-
-The MultiProcExecutor implementation:
-
-1. Spawns worker processes for each GPU
-2. Establishes communication channels via pipes
-3. Coordinates distributed operations across workers
-4. Handles graceful startup and shutdown of the process group
-
-Each GPU worker:
-
-1. Initializes the distributed environment
-2. Builds the pipeline for the specified model
-3. Executes requested operations on its assigned GPU
-4. Manages local resources and communicates results back to the executor
-
-This design allows FastVideo to efficiently utilize multiple GPUs while providing a simple, unified interface for model execution.
-
-## Platforms
-
-The `fastvideo/platforms/` directory provides hardware platform abstractions that enable FastVideo to run efficiently on different hardware configurations:
-
-### Platform Abstraction
-
-FastVideo's platform abstraction layer enables:
-
- **Hardware Detection**: Automatic detection of available hardware
- **Backend Selection**: Appropriate selection of compute kernels
- **Memory Management**: Efficient utilization of hardware-specific memory features
-
-The primary components include:
-
- **Platform Interface**: Defines the common API for all platform implementations
- **CUDA Platform**: Optimized implementation for NVIDIA GPUs
- **Backend Enum**: Used throughout the codebase for feature selection
-
-Usage example:
-
-```python
-from fastvideo.platforms import current_platform, _Backend
-
-# Check hardware capabilities
-if current_platform.supports_backend(_Backend.FLASH_ATTN):
-    # Use FlashAttention implementation
-else:
-    # Fall back to standard implementation
+```json
+{
+  "_class_name": "WanPipeline",
+  "_diffusers_version": "0.33.0.dev0",
+  "scheduler": [
+    "diffusers",
+    "UniPCMultistepScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "UMT5EncoderModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "T5TokenizerFast"
+  ],
+  "transformer": [
+    "diffusers",
+    "WanTransformer3DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLWan"
+  ]
+}
 ```

-The platform system is designed to be extensible for future hardware targets.
+How this maps to FastVideo:

-## Logger
+- `WanPipeline` -> `fastvideo/pipelines/basic/wan/wan_pipeline.py`
+- `WanTransformer3DModel` -> `fastvideo/models/dits/wanvideo.py`
+- `AutoencoderKLWan` -> `fastvideo/models/vaes/wanvae.py`
+- `UMT5EncoderModel` -> `fastvideo/models/encoders/t5.py`
+- `T5TokenizerFast` -> loaded via HF in `fastvideo/models/loader/`
+- `UniPCMultistepScheduler` -> loaded via Diffusers scheduler utilities
+- Pipeline defaults -> `fastvideo/configs/pipelines/wan.py`
+- Sampling defaults -> `fastvideo/configs/sample/wan.py`

-See [PR](https://github.com/hao-ai-lab/FastVideo/pull/356)
+## Pipeline system

-*TODO*: (help wanted) Add an environment variable that disables process-aware logging.
+- `fastvideo/pipelines/basic/*` contains end-to-end pipelines for each model
+  family.
+- `fastvideo/pipelines/stages/*` contains reusable, testable stages.
+- Pipelines subclass `ComposedPipelineBase` and declare required components via
+  `_required_config_modules`.
+- `ForwardBatch` (in `fastvideo/pipelines/pipeline_batch_info.py`) carries
+  prompts, latents, timesteps, and intermediate state across stages.

-## Contributing to FastVideo
+## Model components

-If you're a new contributor, here are some common areas to explore:
+- DiT models: `fastvideo/models/dits/`
+- VAEs: `fastvideo/models/vaes/`
+- Text/image encoders: `fastvideo/models/encoders/`
+- Schedulers: `fastvideo/models/schedulers/`
+- Upsamplers: `fastvideo/models/upsamplers/`
+- Optional audio models: `fastvideo/models/audio/`

-1. **Adding a new model**: Implement new model types in the appropriate subdirectory of `fastvideo/models/`
-2. **Optimizing performance**: Look at attention implementations or memory management
-3. **Adding a new pipeline**: Create a new pipeline subclass in `fastvideo/pipelines/`
-4. **Hardware support**: Extend the `platforms` module for new hardware targets
+## Attention and distributed execution

-When adding code, follow these practices:
+- Attention backends live in `fastvideo/attention/` and can be selected via
+  `FASTVIDEO_ATTENTION_BACKEND`.
+- `LocalAttention` is used for cross-attention and most attention layers.
+- `DistributedAttention` is used for full-sequence self-attention in the DiT.
+- Tensor-parallel layers live in `fastvideo/layers/`.
+- Sequence/tensor parallel utilities live in `fastvideo/distributed/`.

- Use type hints for better code readability
- Add appropriate docstrings
- Maintain the separation between model components and execution logic
- Follow existing patterns for distributed processing
+## Related docs
+
+- [Contributing overview](../contributing/overview.md)
+- [Coding agents workflow](../contributing/coding_agents.md)
+- [Testing guide](../contributing/testing.md)
@@ -131,17 +131,41 @@ self.attn = DistributedAttention(

 ### Registering Models

-Register implemented modules in the model registry:
+Register implemented modules for auto‑discovery by adding `EntryClass` in each
+model module (the registry scans for it):

 ```python
-# In fastvideo/models/registry.py
-_TEXT_TO_VIDEO_DIT_MODELS = {
-    "YourTransformerModel": ("dits", "yourmodule", "YourTransformerClass"),
-}
+# In fastvideo/models/dits/your_module.py
+class YourTransformerModel(...):
+    ...

-_VAE_MODELS = {
-    "YourVAEModel": ("vaes", "yourvae", "YourVAEClass"),
-}
+# Entry point for model registry
+EntryClass = YourTransformerModel
+```
+
+```python
+# In fastvideo/models/vaes/your_vae.py
+class YourVAEModel(...):
+    ...
+
+# Entry point for model registry
+EntryClass = YourVAEModel
+```
+
+Register pipeline config + sampling defaults in the unified registry:
+
+```python
+# In fastvideo/registry.py
+register_configs(
+    sampling_param_cls=YourSamplingParam,
+    pipeline_config_cls=YourPipelineConfig,
+    hf_model_paths=[
+        "org/your-model-id",
+    ],
+    model_detectors=[
+        lambda path: "your-model" in path.lower(),
+    ],
+)
 ```

 ## Step 2: Directory Structure
@@ -0,0 +1,140 @@
+# Offloading
+
+This page describes how to use offloading techniques for inference to reduce GPU memory usage while maintaining acceptable performance.
+
+## Default Behavior
+
+```python
+dit_cpu_offload: bool = True
+use_fsdp_inference: bool = False
+dit_layerwise_offload: bool = True
+text_encoder_cpu_offload: bool = True
+image_encoder_cpu_offload: bool = True
+vae_cpu_offload: bool = True
+pin_cpu_memory: bool = True
+```
+
+## Behavior Explanation
+
+!!! note
+    For CLI usage, replace underscores (`_`) with hyphens (`-`).
+
+### `use_fsdp_inference`
+
+Enables [FSDP](https://docs.pytorch.org/tutorials/intermediate/FSDP_tutorial.html) for inference. The model weights are sharded across multiple GPUs to reduce memory usage per GPU, and weights are broadcast to all GPUs layer by layer during inference.
+
+#### Performance Impact
+
+FSDP inference introduces negligible performance overhead due to weight prefetching. Performance overhead may be visible when GPU interconnect is slow (e.g., multiple consumer-level GPUs connected by slow PCIe without GPU P2P support).
+
+#### Usage Recommendation
+
+We recommend enabling this option when multiple GPUs are available.
+
+### `dit_cpu_offload`
+
+Enables CPU offloading for FSDP inference. When enabled, the model weights are offloaded to CPU memory, and the weight of each layer is moved to GPU memory only when that layer is being computed.
+
+#### Performance Impact
+
+The PyTorch FSDP implementation does not overlap computation and data transfer perfectly for inference, so enabling this option will harm performance.
+
+#### Usage Recommendation
+
+This option only takes effect when FSDP is enabled. For single GPU usage, we recommend using `dit_layerwise_offload` instead.
+
+### `dit_layerwise_offload`
+
+This option is similar to `dit_cpu_offload`, but with two key differences:
+
+1. It overlaps computation and PCIe data transfer
+2. It only works for single GPU inference
+
+#### Performance Impact
+
+This option introduces negligible performance overhead.
+
+#### Usage Recommendation
+
+We recommend enabling this option for single GPU usage. This option is not compatible with FSDP.
+
+### `text_encoder_cpu_offload`
+
+When enabled, the text encoder model weights are offloaded to CPU memory, and text encoding is computed on CPU.
+
+#### Performance Impact
+
+This option significantly slows down text encoding computation, but text encoding is usually not the bottleneck.
+
+#### Usage Recommendation
+
+We recommend enabling this option only when OOM happens.
+
+### `image_encoder_cpu_offload` and `vae_cpu_offload`
+
+When enabled, the weights are stored in CPU memory and moved to GPU memory when the corresponding module is being computed. After computation, the weights are moved back to CPU memory.
+
+#### Performance Impact
+
+These options introduce performance overhead due to PCIe data transfer.
+
+#### Usage Recommendation
+
+We recommend enabling these options when OOM happens.
+
+## General Recommendations
+
+### Single GPU Inference
+
+We recommend enabling `dit_layerwise_offload`. If OOM happens, also enable `image_encoder_cpu_offload` and `vae_cpu_offload`. If OOM still happens, consider enabling `text_encoder_cpu_offload`.
+
+### Multi-GPU Inference
+
+We recommend enabling `use_fsdp_inference` and disabling both `dit_layerwise_offload` and `dit_cpu_offload`. If OOM happens, consider enabling `text_encoder_cpu_offload`, `image_encoder_cpu_offload`, and `vae_cpu_offload`. If OOM still happens, consider enabling `dit_cpu_offload`.
+
+## Examples
+
+### Single GPU with Layerwise Offloading
+
+```python
+from fastvideo import VideoGenerator
+
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=1,
+    # Recommended for single GPU
+    dit_layerwise_offload=True,
+    # Enable if OOM happens
+    vae_cpu_offload=True,
+    image_encoder_cpu_offload=True,
+    text_encoder_cpu_offload=True,
+    # Speeds up CPU-GPU transfer
+    pin_cpu_memory=True,
+)
+
+prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers."
+video = generator.generate_video(prompt, output_path="output/", save_video=True)
+```
+
+### Multi-GPU with FSDP
+
+```python
+from fastvideo import VideoGenerator
+
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=2,
+    # Recommended for multi-GPU
+    use_fsdp_inference=True,
+    dit_layerwise_offload=False,
+    dit_cpu_offload=False,
+    # Enable if OOM happens
+    vae_cpu_offload=True,
+    image_encoder_cpu_offload=True,
+    text_encoder_cpu_offload=True,
+    pin_cpu_memory=True,
+)
+
+prompt = "A majestic lion strides across the golden savanna."
+video = generator.generate_video(prompt, output_path="output/", save_video=True)
+```
@@ -103,7 +103,7 @@ python setup.py install  # or pip install -e .

 **`SAGE_ATTN_THREE`**

-[SageAttention 3](https://huggingface.co/jt-zhang/SageAttention3) is an advanced attention mechanism that leverages FP4 quantization and Blackwell GPU Tensor Cores for significant performance improvements.
+[SageAttention 3](https://github.com/thu-ml/SageAttention/tree/main/sageattention3_blackwell) is an advanced attention mechanism that leverages FP4 quantization and Blackwell GPU Tensor Cores for significant performance improvements.

 #### Hardware Requirements

@@ -113,11 +113,7 @@ python setup.py install  # or pip install -e .

 Note that Sage Attention 3 requires `python>=3.13`, `torch>=2.8.0`, `CUDA >=12.8`. If you are using `uv` and using `torch==2.8.0` make sure that `sentencepiece==0.2.1` in the pyproject.toml file.

-To use Sage Attention 3 in FastVideo, first get access to the SageAttention3 code, then move `sageattn/` and `setup.py` to the directory `fastvideo/attention/backends`, then install from using:
-
-```bash
-python setup.py install
-```
+To use Sage Attention 3 in FastVideo, follow the `README.md` in the linked repository to install the package from source.

 ## Teacache

@@ -1,169 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=dmd_3333
-#SBATCH --partition=main
-#SBATCH --nodes=4
-#SBATCH --ntasks=4
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=dmd_3333_output/dmd_3333_%j.out
-#SBATCH --error=dmd_3333_output/dmd_3333_%j.err
-#SBATCH --exclusive
-
-# Basic Info
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-export NCCL_DEBUG_SUBSYS=INIT,NET
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export TOKENIZERS_PARALLELISM=false
-# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde"
-export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
-# export WANDB_API_KEY='your_wandb_api_key_here'
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-
-source ~/conda/miniconda/bin/activate
-conda activate wei-fv-distill
-export HOME="/mnt/weka/home/hao.zhang/wei"
-
-# Configs
-NUM_GPUS=8
-NUM_TOTAL_GPUS=32
-
-# Model paths for Self-Forcing DMD distillation with Hunyuan1.5:
-GENERATOR_MODEL_PATH="weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v"  # Updated to Wan2.2
-REAL_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" # Teacher model
-FAKE_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
-# REAL_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled"
-# FAKE_SCORE_MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled" # Critic model
-
-# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
-# DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
-DATA_DIR="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/vidprom_16k_text_embed"
-DATA_DIR_2="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-121"
-# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
-# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
-# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
-VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
-
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-training_args=(
-  --tracker_project_name SFhy1.5_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
-  --output_dir "/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_3333_1e-5_1e-5_cfg6_corrected_scheduler"
-  --wandb_run_name "self_forcing_3333_context_forcing"
-  --max_train_steps 4000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 480
-  --num_width 848
-  --enable_gradient_checkpointing_type "full"
-  --simulate_generator_forward
-  --log_visualization
-  --num_frames 81
-  --num_frame_per_block 3  # Frame generation block size for self-forcing
-  # --resume-from-checkpoint "/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_1333/checkpoint-300"
-  --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_worldplay_df_init_3333/checkpoint-2000/transformer/diffusion_pytorch_model.safetensors
-  # --init_weights_from_safetensors /mnt/weka/home/hao.zhang/wei/hy15_distilled_ode_init_3333/checkpoint-2400/transformer/diffusion_pytorch_model.safetensors
-)
-
-parallel_args=(
-  --num_gpus $NUM_TOTAL_GPUS # 64
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1 # 64
-  --hsdp_shard_dim $NUM_TOTAL_GPUS
-)
-
-model_args=(
-  --model_path $GENERATOR_MODEL_PATH
-  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
-  --real_score_model_path $REAL_SCORE_MODEL_PATH
-  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
-)
-
-dataset_args=(
-  --data_path "$DATA_DIR"
-  # --data_path_2 "$DATA_DIR_2"
-  --dataloader_num_workers 4
-)
-
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 50
-  --validation_sampling_steps "8"
-  --validation_guidance_scale "1.0" # not used for dmd inference
-  --text-encoder-cpu-offload
-  # --vae_cpu_offload True
-)
-
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 100
-  --weight_only_checkpointing_steps 500
-  --weight_decay 0.01
-  --betas '0.0,0.999'
-  --max_grad_norm 1.0
-)
-
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --flow_shift 5
-  --seed 1000
-  --use_ema True
-  --ema_decay 0.99
-  --ema_start_step 200
-)
-
-dmd_args=(
-  --warp_denoising_step
-  --dmd_denoising_steps '1000,875,750,625,500,375,250,125'
-  # --dmd_denoising_steps '1000,760,520,280'
-  # --dmd_denoising_steps '1000,750,500,250'
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.98
-  --dfake_gen_update_ratio 5
-  --real_score_guidance_scale 3.5
-  --fake_score_learning_rate 5e-6
-  --fake_score_betas '0.0,0.999'
-)
-
-self_forcing_args=(
-  --independent_first_frame False  # Whether to treat first frame independently
-  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
-  --last_step_only False  # Whether to only use the last denoising step
-  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
-  # --use-context-forcing True
-)
-
-srun torchrun \
--nnodes $SLURM_JOB_NUM_NODES \
--nproc_per_node $NUM_GPUS \
--node_rank $SLURM_PROCID \
--rdzv_backend=c10d \
--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/hy15_self_forcing_distillation_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}" \
-    "${self_forcing_args[@]}"
@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam
+
+
+def main():
+    # Point this to your local diffusers model dir (or replace with a HF model ID).
+    model_path = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
+
+    generator = VideoGenerator.from_pretrained(
+        model_path,
+        num_gpus=1,
+        use_fsdp_inference=False,  # set True if GPU is out of memory
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+    )
+
+    sampling_param = SamplingParam.from_pretrained(model_path)
+
+    # image2world example from official repo
+    image_path = "images/bus_terminal.jpg"
+
+    prompt = (
+        "A nighttime city bus terminal gradually shifts from stillness to subtle movement. "
+        "At first, multiple double-decker buses are parked under the glow of overhead lights, "
+        "with a central bus labeled '87D' facing forward and stationary. "
+        "As the video progresses, the bus in the middle moves ahead slowly, its headlights brightening the surrounding area "
+        "and casting reflections onto adjacent vehicles. "
+        "The motion creates space in the lineup, signaling activity within the otherwise quiet station. "
+        "It then comes to a smooth stop, resuming its position in line. "
+        "Overhead signage in Chinese characters remains illuminated, enhancing the vibrant, urban night scene."
+    )
+
+    generator.generate_video(
+        prompt,
+        sampling_param=sampling_param,
+        image_path=str(image_path),
+        num_cond_frames=1,
+        output_path="outputs_video/cosmos2_5_i2w.mp4",
+        save_video=True,
+    )
+
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()
+
+
@@ -1,4 +1,5 @@
 from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam


 def main():
@@ -15,19 +16,27 @@ def main():
        pin_cpu_memory=True,
    )

+    # Load default sampling parameters (negative_prompt, resolution, steps, etc.)
+    sampling_param = SamplingParam.from_pretrained(model_path)
+
    prompt = (
-       "A high-definition video captures the precision of robotic welding in an industrial setting. The first frame showcases a robotic arm, equipped with a welding torch, positioned over a large metal structure. The welding process is in full swing, with bright sparks and intense light illuminating the scene, creating a vivid display of blue and white hues. A significant amount of smoke billows around the welding area, partially obscuring the view but emphasizing the heat and activity. The background reveals parts of the workshop environment, including a ventilation system and various pieces of machinery, indicating a busy and functional industrial workspace. As the video progresses, the robotic arm maintains its steady position, continuing the welding process and moving to its left. The welding torch consistently emits sparks and light, and the smoke continues to rise, diffusing slightly as it moves upward. The metal surface beneath the torch shows ongoing signs of heating and melting. The scene retains its industrial ambiance, with the welding sparks and smoke dominating the visual field, underscoring the ongoing nature of the welding operation."
+        "A high-definition video captures the precision of robotic welding in an industrial setting. "
+        "The first frame showcases a robotic arm, equipped with a welding torch, positioned over a large metal structure. "
+        "The welding process is in full swing, with bright sparks and intense light illuminating the scene, "
+        "creating a vivid display of blue and white hues. "
+        "A significant amount of smoke billows around the welding area, partially obscuring the view but emphasizing the heat and activity. "
+        "The background reveals parts of the workshop environment, including a ventilation system and various pieces of machinery, "
+        "indicating a busy and functional industrial workspace. "
+        "As the video progresses, the robotic arm maintains its steady position, continuing the welding process and moving to its left. "
+        "The welding torch consistently emits sparks and light, and the smoke continues to rise, diffusing slightly as it moves upward. "
+        "The metal surface beneath the torch shows ongoing signs of heating and melting. "
+        "The scene retains its industrial ambiance, with the welding sparks and smoke dominating the visual field, "
+        "underscoring the ongoing nature of the welding operation."
    )

-    video = generator.generate_video(
+    generator.generate_video(
        prompt,
-        negative_prompt="",
-        height=704,
-        width=1280,
-        num_frames=77,
-        num_inference_steps=35,
-        guidance_scale=7.0,
-        fps=24,
+        sampling_param=sampling_param,
        output_path="outputs_video/cosmos2_5_t2w.mp4",
        save_video=True,
    )
@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample import SamplingParam
+
+
+def main():
+    # Point this to your local diffusers model dir (or replace with a HF model ID).
+    model_path = "KyleShao/Cosmos-Predict2.5-2B-Diffusers"
+
+    generator = VideoGenerator.from_pretrained(
+        model_path,
+        num_gpus=1,
+        use_fsdp_inference=False,  # set True if GPU is out of memory
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+    )
+
+    sampling_param = SamplingParam.from_pretrained(model_path)
+
+    # video2world example from official repo
+    video_path = "videos/robot_pouring.mp4"
+
+    prompt = (
+        "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. "
+        "The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. "
+        "The robotic arm is in the process of pouring a transparent liquid into the glass. "
+        "To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. "
+        "In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. "
+        "The lighting is bright, casting soft shadows on the table. "
+        "The robotic arm's movements are smooth and controlled, demonstrating precision in its task. "
+        "As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. "
+        "The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. "
+        "The other robotic arm on the right side also stays stationary throughout the video. "
+        "The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper."
+    )
+
+    generator.generate_video(
+        prompt,
+        sampling_param=sampling_param,
+        video_path=str(video_path),
+        num_cond_frames=1,
+        output_path="outputs_video/cosmos2_5_v2w.mp4",
+        save_video=True,
+    )
+
+    generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()
+
+
@@ -2,14 +2,14 @@ from fastvideo import VideoGenerator
 import json
 # from fastvideo.configs.sample import SamplingParam

-OUTPUT_PATH = "video_samples_hy15_t2v_distilled"
+OUTPUT_PATH = "video_samples_hy15"
 def main():
    # FastVideo will automatically use the optimal default arguments for the
    # model.
    # If a local path is provided, FastVideo will make a best effort
    # attempt to identify the optimal arguments.
    generator = VideoGenerator.from_pretrained(
-        "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled",
+        "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
        # FastVideo will automatically handle distributed setup
        num_gpus=1,
        use_fsdp_inference=False, # set to True if GPU is out of memory
@@ -18,35 +18,15 @@ def main():
        text_encoder_cpu_offload=True,
        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
        # image_encoder_cpu_offload=False,
-        # init_weights_from_safetensors="/mnt/weka/home/hao.zhang/wei/SFhy1.5_distill_1333_5e-5_2e-6_cfg3.5/checkpoint-500/ema/generator_ema.safetensors"
    )

-    # json_path = "/mnt/weka/home/hao.zhang/wei/FastVideo/data/mixkit_i2v_full_720p.json"
-    # with open(json_path, 'r') as f:
-    #     data_list = json.load(f)["data"]
-    
-    # # Now you can index into data_list however you like
-    # # For example: data_list[0], data_list[1:3], etc.
-    # for data in data_list:
-    #     prompt = data["prompt"]
-    #     image_path = data["image_path"]
-    #     generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path, num_frames=121, fps=24)
-    # return
-
-    # prompt = (
-    #     "In a brightly lit studio, a photographer wearing a denim jacket focuses intently, capturing shots with a professional camera. Facing him, a model stands gracefully, adjusting her long, flowing hair with delicate movements. The scene is characterized by strong contrasts; the model's soft pink attire and gentle gestures complement the rugged, precise demeanor of the photographer. Positioned against a minimalist backdrop, the pair work seamlessly, with the camera\u2019s lens pointed directly at the model, capturing her elegance. The soft, diffused lighting casts a gentle glow on both subjects, creating an airy and ethereal atmosphere perfect for a high-fashion photo shoot."
-    # )
-
-    # video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=121, fps=24, image_path="data/1.png")
-    # return
-
    prompt = (
        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
        "wide with interest. The playful yet serene atmosphere is complemented by soft "
        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
    )

-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=81, fps=16)

    prompt2 = (
        "A majestic lion strides across the golden savanna, its powerful frame "
@@ -55,8 +35,8 @@ def main():
        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
        "cinematic.")

-    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, negative_prompt="", num_frames=81, fps=16)


 if __name__ == "__main__":
-    main()
+    main()
@@ -0,0 +1,43 @@
+from fastvideo import VideoGenerator
+import json
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_hy15_1080p"
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    generator = VideoGenerator.from_pretrained(
+        "weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR", # 480p -> 720p -> 1080p
+        # or "weizhou03/HunyuanVideo-1.5-Diffusers-1080p" # 720p -> 1080p
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=False, # set to True if GPU is out of memory
+        dit_cpu_offload=True,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        # image_encoder_cpu_offload=False,
+    )
+
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, negative_prompt="")
+
+    prompt2 = (
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, negative_prompt="")
+
+
+if __name__ == "__main__":
+    main()
@@ -1,165 +1,11 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Basic example for HYWorld (HY-WorldPlay) video generation using FastVideo.
-
-This example replicates the same functionality as HY-WorldPlay/run.sh,
-demonstrating image-to-video generation with camera trajectory control.
-"""
-
-import time
-import math
-import numpy as np
-import imageio
-import torchvision
-from einops import rearrange
-
 from fastvideo import VideoGenerator
-from fastvideo.pipelines import ForwardBatch
-from fastvideo.utils import shallow_asdict, align_to
-from fastvideo.logger import init_logger
-
-from fastvideo.models.dits.hyworld.pose import pose_to_input, compute_latent_num
 from fastvideo.models.dits.hyworld.resolution_utils import get_resolution_from_image

-logger = init_logger(__name__)
-
-class HYWorldVideoGenerator(VideoGenerator):
-    """Extended VideoGenerator that adds HYWorld-specific parameters to batch.extra."""
-
-    def _generate_single_video(self, prompt: str, sampling_param=None, **kwargs):
-        """Override to add viewmats, Ks, and action to batch.extra."""
-        fastvideo_args = self.fastvideo_args
-        pipeline_config = fastvideo_args.pipeline_config
-
-        if sampling_param is None:
-            from fastvideo.configs.sample import SamplingParam
-            sampling_param = SamplingParam.from_pretrained(fastvideo_args.model_path)
-
-        # Update sampling param with kwargs
-        if kwargs:
-            for key, value in kwargs.items():
-                if hasattr(sampling_param, key):
-                    setattr(sampling_param, key, value)
-
-        # Get pose string from sampling_param or kwargs
-        pose = kwargs.get('pose', getattr(sampling_param, 'POSE', 'w-31'))
-        num_frames = kwargs.get('num_frames', getattr(sampling_param, 'num_frames', 125))
-
-        # Calculate number of latents
-        latent_num = compute_latent_num(num_frames)
-
-        # Convert pose to viewmats, Ks, and action
-        viewmats, Ks, action = pose_to_input(pose, latent_num)
-
-        # Convert to tensors and add batch dimension
-        viewmats = viewmats.unsqueeze(0)  # (1, T, 4, 4)
-        Ks = Ks.unsqueeze(0)  # (1, T, 3, 3)
-        action = action.unsqueeze(0)  # (1, T)
-
-        # Validate inputs
-        prompt = prompt.strip()
-        sampling_param = sampling_param.__class__(**shallow_asdict(sampling_param))
-        output_path = kwargs.get("output_path", sampling_param.output_path)
-        sampling_param.prompt = prompt
-
-        if sampling_param.negative_prompt is not None:
-            sampling_param.negative_prompt = sampling_param.negative_prompt.strip()
-
-        # Validate dimensions
-        if (sampling_param.height <= 0 or sampling_param.width <= 0 or
-            sampling_param.num_frames <= 0):
-            raise ValueError(
-                f"Height, width, and num_frames must be positive integers")
-
-        temporal_scale_factor = pipeline_config.vae_config.arch_config.temporal_compression_ratio
-        num_frames = sampling_param.num_frames
-        num_gpus = fastvideo_args.num_gpus
-        use_temporal_scaling_frames = pipeline_config.vae_config.use_temporal_scaling_frames
-
-        # Adjust number of frames based on number of GPUs
-        if use_temporal_scaling_frames:
-            orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1
-        else:
-            orig_latent_num_frames = sampling_param.num_frames // 17 * 3
-
-        if orig_latent_num_frames % fastvideo_args.num_gpus != 0:
-            if use_temporal_scaling_frames:
-                new_num_frames = (orig_latent_num_frames - 1) * temporal_scale_factor + 1
-            else:
-                divisor = math.lcm(3, num_gpus)
-                orig_latent_num_frames = (
-                    (orig_latent_num_frames + divisor - 1) // divisor) * divisor
-                new_num_frames = orig_latent_num_frames // 3 * 17
-
-            logger.info(
-                "Adjusting number of frames from %s to %s based on number of GPUs (%s)",
-                sampling_param.num_frames, new_num_frames, fastvideo_args.num_gpus)
-            sampling_param.num_frames = new_num_frames
-
-        # Calculate sizes
-        target_height = align_to(sampling_param.height, 16)
-        target_width = align_to(sampling_param.width, 16)
-
-        # Calculate latent sizes
-        latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
-                        sampling_param.height // 8, sampling_param.width // 8]
-        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
-
-        # Prepare batch
-        batch = ForwardBatch(
-            **shallow_asdict(sampling_param),
-            eta=0.0,
-            n_tokens=n_tokens,
-            VSA_sparsity=fastvideo_args.VSA_sparsity,
-        )
-
-        # Add HYWorld-specific parameters to batch.extra
-        batch.extra['viewmats'] = viewmats
-        batch.extra['Ks'] = Ks
-        batch.extra['action'] = action
-        batch.extra['chunk_latent_frames'] = 16  # For bidirectional model
-
-        # Run inference
-        start_time = time.perf_counter()
-        output_batch = self.executor.execute_forward(batch, fastvideo_args)
-        samples = output_batch.output
-        logging_info = output_batch.logging_info
-
-        gen_time = time.perf_counter() - start_time
-        logger.info("Generated successfully in %.2f seconds", gen_time)
-
-        # Process outputs
-        videos = rearrange(samples, "b c t h w -> t b c h w")
-        frames = []
-        for x in videos:
-            x = torchvision.utils.make_grid(x, nrow=6)
-            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-            frames.append((x * 255).numpy().astype(np.uint8))
-
-        # Save video if requested
-        if batch.save_video:
-            imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4")
-            logger.info("Saved video to %s", output_path)
-
-        if batch.return_frames:
-            return frames
-        else:
-            return {
-                "samples": samples,
-                "frames": frames,
-                "prompts": prompt,
-                "size": (target_height, target_width, batch.num_frames),
-                "generation_time": gen_time,
-                "logging_info": logging_info,
-                "trajectory": output_batch.trajectory_latents,
-                "trajectory_timesteps": output_batch.trajectory_timesteps,
-                "trajectory_decoded": output_batch.trajectory_decoded,
-            }
-
 # Default prompt from HY-WorldPlay run.sh
 DEFAULT_PROMPT = 'A paved pathway leads towards a stone arch bridge spanning a calm body of water.  Lush green trees and foliage line the path and the far bank of the water. A traditional-style pavilion with a tiered, reddish-brown roof sits on the far shore. The water reflects the surrounding greenery and the sky.  The scene is bathed in soft, natural light, creating a tranquil and serene atmosphere. The pathway is composed of large, rectangular stones, and the bridge is constructed of light gray stone.  The overall composition emphasizes the peaceful and harmonious nature of the landscape.'
 DEFAULT_IMAGE = 'https://raw.githubusercontent.com/Tencent-Hunyuan/HY-WorldPlay/main/assets/img/test.png'

+OUTPUT_PATH = "video_samples_hyworld"
 def main():
    import argparse

@@ -169,7 +15,7 @@ def main():
    parser.add_argument("--prompt", type=str, default=DEFAULT_PROMPT, help="Text prompt for video generation")
    parser.add_argument("--image", type=str, default=DEFAULT_IMAGE, help="Path or URL to input image")
    parser.add_argument("--pose", type=str, default='w-31', help="Pose string (e.g., 'a-31', 'w-31', 's-31', 'd-31')")
-    parser.add_argument("--output_path", type=str, default='video_samples_hyworld', help="Output video path")
+    parser.add_argument("--output_path", type=str, default=OUTPUT_PATH, help="Output video path")
    parser.add_argument("--num-frames", type=int, default=125, help="Number of frames")
    parser.add_argument("--seed", type=int, default=1, help="Random seed")
    parser.add_argument("--resolution", type=str, default="480p", help="Only support 480p for now")
@@ -185,8 +31,7 @@ def main():

    # Initialize generator
    print("\nInitializing VideoGenerator for HYWorld...")
-    
-    generator = HYWorldVideoGenerator.from_pretrained(
+    generator = VideoGenerator.from_pretrained(
        "FastVideo/HY-WorldPlay-Bidirectional-Diffusers",
        num_gpus=1,
        use_fsdp_inference=True,
@@ -198,11 +43,12 @@ def main():
    )

    # Generate video
+    # The pose string is automatically converted to camera matrices by the pipeline
    print("\nGenerating video...")
-    start_time = time.time()
-    video = generator.generate_video(
+    generator.generate_video(
        prompt=args.prompt,
        image_path=args.image,
+        pose=args.pose,  # Camera trajectory control
        output_path=args.output_path,
        save_video=True,
        negative_prompt="",
@@ -211,13 +57,9 @@ def main():
        height=HEIGHT,
        width=WIDTH,
        seed=args.seed,
-        pose=args.pose,
    )
-    elapsed = time.time() - start_time

-    print(f"\nVideo generated successfully!")
-    print(f"Saved to: {args.output_path}")
-    print(f"Time: {elapsed:.2f}s")
+    print(f"\nVideo saved to: {args.output_path}")


 if __name__ == "__main__":
@@ -1,6 +1,5 @@
 from fastvideo import VideoGenerator
-from fastvideo.configs.pipelines.wan import MatrixGameI2V480PConfig
-from fastvideo.models.dits.matrix_game.utils import create_action_presets
+from fastvideo.models.dits.matrixgame.utils import create_action_presets

 import torch

@@ -1,5 +1,5 @@
 from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator
-from fastvideo.models.dits.matrix_game.utils import get_current_action_async, expand_action_to_frames
+from fastvideo.models.dits.matrixgame.utils import get_current_action_async, expand_action_to_frames

 import torch
 import asyncio
@@ -10,7 +10,7 @@ from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import HTMLResponse, FileResponse

 from fastvideo.entrypoints.streaming_generator import StreamingVideoGenerator
-from fastvideo.models.dits.matrix_game.utils import expand_action_to_frames
+from fastvideo.models.dits.matrixgame.utils import expand_action_to_frames


 VARIANT_CONFIG = {
@@ -1,42 +1,39 @@
 #!/bin/bash

 export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=offline
+export WANDB_MODE=online
 export TOKENIZERS_PARALLELISM=false

-MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
-DATA_DIR="data/ode-preprocessing-hy15-test/"
-VALIDATION_DATASET_FILE="data/validation_64.json"
+MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
 NUM_GPUS=1
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name "wan_ode_init"
-  --output_dir "ode_init_hy15_test"
-  --wandb_run_name "vidprom_bz128_1e-5"
-  # --resume_from_checkpoint "ode_init_diffusers/"
-  # --warp_denoising_step
-  # --log_visualization
-  --max_train_steps 6001
+  --output_dir "wan_ode_init_crush_smol"
+  --wandb_run_name "wan_ode_init_crush_smol"
+  --max_train_steps 6000
  --train_batch_size 1
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
-  --num_latent_t 19
+  --num_latent_t 21
  --num_height 480
  --num_width 832
-  --num_frames 73
-  --dmd_denoising_steps "1000,750,500,250"
+  --num_frames 77
+  --warp_denoising_step
  --enable_gradient_checkpointing_type "full"
 )

 # Parallel arguments
 parallel_args=(
-  --num_gpus 1
+  --num_gpus $NUM_GPUS
  --sp_size 1
  --tp_size 1
  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 1
+  --hsdp_shard_dim $NUM_GPUS
 )

 # Model arguments
@@ -53,19 +50,17 @@ dataset_args=(

 # Validation arguments
 validation_args=(
-  # --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 200
-  --validation_sampling_steps "50"
-  --validation_guidance_scale "6.0"
+  --log-visualization
+  --visualization-steps 100
 )

 # Optimizer arguments
 optimizer_args=(
  --learning_rate 1e-5
  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 1000
  --training_state_checkpointing_steps 1000
-  --weight_decay 1e-4
+  --weight_decay 0.01
  --max_grad_norm 1.0
 )

@@ -79,7 +74,6 @@ miscellaneous_args=(
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
-  # --enable_gradient_checkpointing_type "full"
 )

 # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
@@ -1,132 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=hy15_ode_3333
-#SBATCH --partition=main
-#SBATCH --nodes=8
-#SBATCH --ntasks=8
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=hy15_ode_3333_output/hy15_ode_3333_%j.out
-#SBATCH --error=hy15_ode_3333_output/hy15_ode_3333_%j.err
-#SBATCH --exclusive
-
-# Basic Info
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-export NCCL_DEBUG_SUBSYS=INIT,NET
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export TOKENIZERS_PARALLELISM=false
-# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde"
-export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
-# export WANDB_API_KEY='your_wandb_api_key_here'
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-
-source ~/conda/miniconda/bin/activate
-conda activate wei-fv-distill
-export HOME="/mnt/weka/home/hao.zhang/wei"
-
-# Configs
-NUM_GPUS=8
-NUM_TOTAL_GPUS=64
-
-MODEL_PATH="weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v"
-DATA_DIR="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-121"
-VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name "hy15_ode_init"
-  --output_dir "/mnt/weka/home/hao.zhang/wei/hy15_ode_init_1333_new"
-  --wandb_run_name "hy15_ode_init_1333_new"
-  --warp_denoising_step
-  --dmd_denoising_steps '1000,760,520,280,0'
-  --log_visualization
-  --visualization_steps 100
-  --max_train_steps 4000
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 2
-  --num_latent_t 31
-  --num_height 480
-  --num_width 848
-  --num_frames 121
-  --enable_gradient_checkpointing_type "full"
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_TOTAL_GPUS # 64
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1 # 64
-  --hsdp_shard_dim $NUM_TOTAL_GPUS
-)
-
-# Model arguments
-model_args=(
-  --model_path $MODEL_PATH
-  --pretrained_model_name_or_path $MODEL_PATH
-  --text-encoder-cpu-offload
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR" 
-  --dataloader_num_workers 4
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 400
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-validation_args=(
-  # --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 10
-  --validation_sampling_steps "4"
-  --validation_guidance_scale "1.0" # not used for dmd inference
-  # --vae_cpu_offload True
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 6
-  --training_cfg_rate 0.1
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-  # --enable_gradient_checkpointing_type "full"
-)
-
-# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
-srun torchrun \
--nnodes $SLURM_JOB_NUM_NODES \
--nproc_per_node $NUM_GPUS \
--node_rank $SLURM_PROCID \
--rdzv_backend=c10d \
--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/ode_causal_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}"
@@ -0,0 +1,93 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+MODEL_PATH="FastVideo/Matrix-Game-2.0-Foundation-Diffusers"
+DATA_DIR="footsies-dataset/preprocessed/combined_parquet_dataset"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=8
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "matrixgame_finetune"
+  --output_dir "checkpoints/matrixgame_finetune"
+  --max_train_steps 1500
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 4
+  --num_latent_t 20
+  --num_height 352
+  --num_width 640
+  --num_frames 77
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 2
+  --tp_size 1
+  --hsdp_replicate_dim 4
+  --hsdp_shard_dim 2
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 100
+  --validation_sampling_steps "40"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-5
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 100
+  --training_state_checkpointing_steps 100
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/matrixgame_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -0,0 +1,27 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="Matrix-Game-2.0-Foundation-Diffusers"
+DATA_MERGE_PATH="footsies-dataset/merge.txt"
+OUTPUT_DIR="footsies-dataset/preprocessed/"
+
+# export CUDA_VISIBLE_DEVICES=0
+export MASTER_ADDR=localhost
+export MASTER_PORT=29500
+export RANK=0
+export WORLD_SIZE=1
+
+python fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 4 \
+    --seed 42 \
+    --max_height 352 \
+    --max_width 640 \
+    --num_frames 77 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --samples_per_file 4 \
+    --train_fps 25 \
+    --flush_frequency 4 \
+    --preprocess_task matrixgame
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0

 import torch
-from fastvideo.attention.backends.sageattn.api import sageattn_blackwell
+from sageattn3 import sageattn3_blackwell

 from fastvideo.attention.backends.abstract import (AttentionBackend,
                                                   AttentionImpl,
@@ -67,6 +67,6 @@ class SageAttention3Impl(AttentionImpl):
        query = query.transpose(1, 2)
        key = key.transpose(1, 2)
        value = value.transpose(1, 2)
-        output = sageattn_blackwell(query, key, value, is_causal=self.causal)
+        output = sageattn3_blackwell(query, key, value, is_causal=self.causal)
        output = output.transpose(1, 2)
        return output
@@ -5,6 +5,7 @@ from fastvideo.configs.models.vaes.base import VAEConfig
 from fastvideo.configs.models.audio import (LTX2AudioDecoderConfig,
                                            LTX2AudioEncoderConfig,
                                            LTX2VocoderConfig)
+from fastvideo.configs.models.upsamplers.base import UpsamplerConfig

 __all__ = [
    "ModelConfig",
@@ -14,4 +15,5 @@ __all__ = [
    "LTX2AudioEncoderConfig",
    "LTX2AudioDecoderConfig",
    "LTX2VocoderConfig",
+    "UpsamplerConfig",
 ]
@@ -26,9 +26,6 @@ class HunyuanVideo15ArchConfig(DiTArchConfig):

    param_names_mapping: dict = field(
        default_factory=lambda: {
-            r"^cond_type_embed\.(.*)$":
-            r"cond_type_embed.\1",
-
            # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
            r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
            r"txt_in.t_embedder.mlp.fc_in.\1",
@@ -60,14 +57,6 @@ class HunyuanVideo15ArchConfig(DiTArchConfig):
            r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.self_attn_qkv\.(.*)$":
            r"txt_in.refiner_blocks.\1.self_attn_qkv.\2",
-            r"^image_embedder\.linear_1\.(.*)$":
-            r"image_embedder.linear_1.\1",
-            r"^image_embedder\.linear_2\.(.*)$":
-            r"image_embedder.linear_2.\1",
-            r"^image_embedder\.norm_in\.(.*)$":
-            r"image_embedder.norm_in.\1",
-            r"^image_embedder\.norm_out\.(.*)$":
-            r"image_embedder.norm_out.\1",

            # 2. txt_in_2 mapping:
            r"^context_embedder_2\.(.*)$":
@@ -157,12 +146,6 @@ class HunyuanVideo15ArchConfig(DiTArchConfig):
    exclude_lora_layers: list[str] = field(
        default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])

-    # Causal HunyuanVideo1.5
-    local_attn_size: int = -1  # Window size for temporal local attention (-1 indicates global attention)
-    sink_size: int = 0  # Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
-    num_frames_per_block: int = 3
-    sliding_window_num_frames: int = 31
-
    def __post_init__(self):
        super().__post_init__()
        self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
@@ -8,12 +8,8 @@ def is_double_block(n: str, m) -> bool:
    return "double" in n and str.isdigit(n.split(".")[-1])


-def is_single_block(n: str, m) -> bool:
-    return "single" in n and str.isdigit(n.split(".")[-1])
-
-
-# def is_refiner_block(n: str, m) -> bool:
-#     return "refiner" in n and str.isdigit(n.split(".")[-1])
+def is_refiner_block(n: str, m) -> bool:
+    return "refiner" in n and str.isdigit(n.split(".")[-1])


 def is_txt_in(n: str, m) -> bool:
@@ -23,10 +19,10 @@ def is_txt_in(n: str, m) -> bool:
@dataclass
 class HYWorldArchConfig(DiTArchConfig):
    _fsdp_shard_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_single_block])
+        default_factory=lambda: [is_double_block, is_refiner_block])

    _compile_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_single_block, is_txt_in])
+        default_factory=lambda: [is_double_block, is_refiner_block, is_txt_in])

    param_names_mapping: dict = field(
        default_factory=lambda: {
@@ -54,8 +50,7 @@ class HYWorldArchConfig(DiTArchConfig):
            r"^txt_in\.individual_token_refiner\.blocks\.(\d+)\.adaLN_modulation\.1\.(.*)$":
            r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",

-            # 2. time_in mappings (HYWorld uses TimestepEmbedder directly,
-            #    but FastVideo model inherits HunyuanVideo15TimeEmbedding with timestep_embedder):
+            # 2. time_in mappings:
            r"^time_in\.mlp\.0\.(.*)$":
            r"time_in.timestep_embedder.mlp.fc_in.\1",
            r"^time_in\.mlp\.2\.(.*)$":
@@ -1,4 +1,6 @@
 from dataclasses import dataclass, field
+
+import torch
 from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig


@@ -8,8 +10,8 @@ class MatrixGameWanVideoArchConfig(WanVideoArchConfig):
    # because MatrixGame checkpoints already have patch_embedding.proj format
    param_names_mapping: dict = field(
        default_factory=lambda: {
-            # Removed: r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1"
-            # because checkpoint already has correct format
+            r"^patch_embedding\.(?!proj\.)(.*)$":
+            r"patch_embedding.proj.\1",
            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$":
            r"condition_embedder.text_embedder.fc_in.\1",
            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$":
@@ -76,8 +78,14 @@ class MatrixGameWanVideoArchConfig(WanVideoArchConfig):
    image_dim: int = 1280


+def _is_transformer_block(param_name: str, module: torch.nn.Module) -> bool:
+    return bool("blocks" in param_name and param_name.split(".")[-1].isdigit())
+
+
@dataclass
 class MatrixGameWanVideoConfig(WanVideoConfig):
    arch_config: MatrixGameWanVideoArchConfig = field(
        default_factory=MatrixGameWanVideoArchConfig)
    prefix: str = "Wan"
+    _compile_conditions: list = field(
+        default_factory=lambda: [_is_transformer_block])
@@ -87,10 +87,10 @@ class CLIPVisionConfig(ImageEncoderConfig):
    arch_config: ImageEncoderArchConfig = field(
        default_factory=CLIPVisionArchConfig)

-    num_hidden_layers_override: int | None = None
+    num_hidden_layers_override: int | None = 31
    require_post_norm: bool | None = None
-    enable_scale: bool = True
-    is_causal: bool = True
+    enable_scale: bool = False
+    is_causal: bool = False
    prefix: str = "clip"


@@ -0,0 +1,6 @@
+from fastvideo.configs.models.upsamplers.hunyuan15 import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig
+from fastvideo.configs.models.upsamplers.base import UpsamplerConfig
+
+__all__ = [
+    "SRTo720pUpsamplerConfig", "SRTo1080pUpsamplerConfig", "UpsamplerConfig"
+]
@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+from fastvideo.configs.models.base import ModelConfig
+
+
+@dataclass
+class UpsamplerConfig(ModelConfig):
+    pass
@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from fastvideo.configs.models.upsamplers.base import UpsamplerConfig
+
+
+@dataclass
+class SRTo720pUpsamplerConfig(UpsamplerConfig):
+    in_channels: int = 0
+    out_channels: int = 0
+    hidden_channels: int = 64
+    num_blocks: int = 6
+    global_residual: bool = False
+
+
+@dataclass
+class SRTo1080pUpsamplerConfig(UpsamplerConfig):
+    z_channels: int = 0
+    out_channels: int = 0
+    block_out_channels: tuple[int, ...] = (0, 0)
+    num_res_blocks: int = 2
+    is_residual: bool = False
@@ -6,8 +6,7 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
-from fastvideo.configs.pipelines.registry import (
-    get_pipeline_config_cls_from_name)
+from fastvideo.registry import get_pipeline_config_cls_from_name
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
 from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig,
                                             WanI2V480PConfig, WanI2V720PConfig,
@@ -8,7 +8,7 @@ from typing import Any, cast
 import torch

 from fastvideo.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
-                                      VAEConfig)
+                                      VAEConfig, UpsamplerConfig)
 from fastvideo.configs.models.encoders import BaseEncoderOutput
 from fastvideo.configs.utils import update_config_from_args
 from fastvideo.logger import init_logger
@@ -44,12 +44,15 @@ class PipelineConfig:
    # Video generation parameters
    embedded_cfg_scale: float = 6.0
    flow_shift: float | None = None
+    flow_shift_sr: float | None = None
    disable_autocast: bool = False
    is_causal: bool = False

    # Model configuration
    dit_config: DiTConfig = field(default_factory=DiTConfig)
    dit_precision: str = "bf16"
+    upsampler_config: UpsamplerConfig = field(default_factory=UpsamplerConfig)
+    upsampler_precision: str = "fp32"

    # VAE configuration
    vae_config: VAEConfig = field(default_factory=VAEConfig)
@@ -216,6 +219,24 @@ class PipelineConfig:
            "Comma-separated list of denoising steps (e.g., '1000,757,522')",
        )

+        # STA (Sliding Tile Attention) parameters
+        parser.add_argument(
+            f"--{prefix_with_dot}STA-mode",
+            type=str,
+            dest=f"{prefix_with_dot.replace('-', '_')}STA_mode",
+            default=PipelineConfig.STA_mode.value,
+            choices=[mode.value for mode in STA_Mode],
+            help=
+            "STA mode: STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None",
+        )
+        parser.add_argument(
+            f"--{prefix_with_dot}skip-time-steps",
+            type=int,
+            dest=f"{prefix_with_dot.replace('-', '_')}skip_time_steps",
+            default=PipelineConfig.skip_time_steps,
+            help="Number of time steps to warmup (full attention) for STA",
+        )
+
        # Add VAE configuration arguments
        from fastvideo.configs.models.vaes.base import VAEConfig
        VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")
@@ -245,8 +266,7 @@ class PipelineConfig:
        """
        use the pipeline class setting from model_path to match the pipeline config
        """
-        from fastvideo.configs.pipelines.registry import (
-            get_pipeline_config_cls_from_name)
+        from fastvideo.registry import get_pipeline_config_cls_from_name
        pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)

        return cast(PipelineConfig, pipeline_config_cls(model_path=model_path))
@@ -260,8 +280,7 @@ class PipelineConfig:
        kwargs: dictionary of kwargs
        config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance
        """
-        from fastvideo.configs.pipelines.registry import (
-            get_pipeline_config_cls_from_name)
+        from fastvideo.registry import get_pipeline_config_cls_from_name

        prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip()
                                                      != "") else ""
@@ -298,6 +317,12 @@ class PipelineConfig:
        # 4. Update PipelineConfig from CLI arguments if provided
        kwargs[prefix_with_dot + 'model_path'] = model_path
        pipeline_config.update_config_from_dict(kwargs, config_cli_prefix)
+
+        # Convert STA_mode string to enum if necessary
+        if isinstance(pipeline_config.STA_mode, str) and not isinstance(
+                pipeline_config.STA_mode, STA_Mode):
+            pipeline_config.STA_mode = STA_Mode(pipeline_config.STA_mode)
+
        return pipeline_config

    def check_pipeline_config(self) -> None:
@@ -11,7 +11,8 @@ from fastvideo.configs.models.dits import HunyuanVideo15Config
 from fastvideo.configs.models.encoders import (BaseEncoderOutput,
                                               Qwen2_5_VLConfig, T5Config)
 from fastvideo.configs.models.vaes import Hunyuan15VAEConfig
-from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models.upsamplers import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig
+from fastvideo.configs.pipelines.base import PipelineConfig, UpsamplerConfig

 PROMPT_TEMPLATE_TOKEN_LENGTH = 108

@@ -127,12 +128,12 @@ class Hunyuan15T2V480PConfig(PipelineConfig):
    vae_tiling: bool = True

    def __post_init__(self):
-        self.vae_config.load_encoder = True
+        self.vae_config.load_encoder = False
        self.vae_config.load_decoder = True


@dataclass
-class Hunyuan15DistilledI2V480PConfig(Hunyuan15T2V480PConfig):
+class Hunyuan15I2V480PStepDistilledConfig(Hunyuan15T2V480PConfig):
    flow_shift: int = 7


@@ -145,11 +146,21 @@ class Hunyuan15T2V720PConfig(Hunyuan15T2V480PConfig):


@dataclass
-class SelfForcingHunyuan15T2V480PConfig(Hunyuan15T2V480PConfig):
-    flow_shift: int = 5
-    is_causal: bool = True
-    # dmd_denoising_steps: list[int] | None = field(
-    #     default_factory=lambda: [1000, 750, 500, 250])
-    dmd_denoising_steps: list[int] | None = field(
-        default_factory=lambda: [1000, 875, 750, 625, 500, 375, 250, 125])
-    warp_denoising_step: bool = True
+class Hunyuan15I2V720PConfig(Hunyuan15T2V720PConfig):
+    """Base configuration for HunYuan pipeline architecture."""
+
+    # HunyuanConfig-specific parameters with defaults
+    flow_shift: int = 7
+
+
+@dataclass
+class Hunyuan15SR1080PConfig(Hunyuan15T2V720PConfig):
+    """Base configuration for HunYuan pipeline architecture."""
+
+    # HunyuanConfig-specific parameters with defaults
+    flow_shift: int = 7
+    flow_shift_sr: int = 2
+    upsampler_config: tuple[UpsamplerConfig, ...] = field(
+        default_factory=lambda:
+        (SRTo720pUpsamplerConfig(), SRTo1080pUpsamplerConfig()))
+    upsampler_precision: str = "fp32"
@@ -1,218 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Registry for pipeline weight-specific configurations."""
-
-import os
-from collections.abc import Callable
-
-from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.configs.pipelines.cosmos import CosmosConfig
-from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
-from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
-from fastvideo.configs.pipelines.hyworld import HYWorldConfig
-from fastvideo.configs.pipelines.ltx2 import LTX2T2VConfig
-from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig, SelfForcingHunyuan15T2V480PConfig, Hunyuan15DistilledI2V480PConfig
-from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
-from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
-from fastvideo.configs.pipelines.turbodiffusion import (
-    TurboDiffusionT2V_1_3B_Config, TurboDiffusionT2V_14B_Config,
-    TurboDiffusionI2V_A14B_Config)
-
-# isort: off
-from fastvideo.configs.pipelines.wan import (
-    FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
-    Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config,
-    WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig,
-    SelfForcingWanT2V480PConfig, WANV2VConfig, SelfForcingWan2_2_T2V480PConfig,
-    MatrixGameI2V480PConfig)
-# isort: on
-from fastvideo.logger import init_logger
-from fastvideo.utils import (maybe_download_model_index,
-                             verify_model_config_and_directory)
-
-logger = init_logger(__name__)
-
-# Registry maps specific model weights to their config classes
-PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
-    "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
-    "hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
-    "weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v":
-    SelfForcingHunyuan15T2V480PConfig,
-    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled":
-    Hunyuan15DistilledI2V480PConfig,
-    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v":
-    Hunyuan15T2V480PConfig,
-    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v":
-    Hunyuan15T2V720PConfig,
-    "FastVideo/HY-WorldPlay-Bidirectional-Diffusers": HYWorldConfig,
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
-    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": WanI2V480PConfig,
-    "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers": WANV2VConfig,
-    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig,
-    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
-    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig,
-    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWan2_1_T2V_480P_Config,
-    "FastVideo/FastWan2.1-T2V-14B-480P-Diffusers": FastWan2_1_T2V_480P_Config,
-    "FastVideo/FastWan2.2-TI2V-5B-Diffusers": FastWan2_2_TI2V_5B_Config,
-    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
-    "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
-    "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
-    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers": SelfForcingWan2_2_T2V480PConfig,
-    "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers":
-    SelfForcingWan2_2_T2V480PConfig,
-    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config,
-    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
-    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,
-    "nvidia/Cosmos-Predict2-2B-Video2World": CosmosConfig,
-    "KyleShao/Cosmos-Predict2.5-2B-Diffusers": Cosmos25Config,
-    "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGameI2V480PConfig,
-    "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGameI2V480PConfig,
-    "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGameI2V480PConfig,
-    # LongCat Video models
-    "FastVideo/LongCat-Video-T2V-Diffusers": LongCatT2V480PConfig,
-    "FastVideo/LongCat-Video-I2V-Diffusers": LongCatT2V480PConfig,
-    "FastVideo/LongCat-Video-VC-Diffusers": LongCatT2V480PConfig,
-    # LTX-2 models
-    "Lightricks/LTX-2": LTX2T2VConfig,
-    "converted/ltx2_diffusers": LTX2T2VConfig,
-    # TurboDiffusion models
-    "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers": TurboDiffusionT2V_1_3B_Config,
-    "loayrashid/TurboWan2.1-T2V-14B-Diffusers": TurboDiffusionT2V_14B_Config,
-    "loayrashid/TurboWan2.2-I2V-A14B-Diffusers": TurboDiffusionI2V_A14B_Config,
-    # Add other specific weight variants
-}
-
-# For determining pipeline type from model ID
-PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
-    "longcatimagetovideo":
-    lambda id: "longcatimagetovideo" in id.lower(),
-    "longcatvideocontinuation":
-    lambda id: "longcatvideocontinuation" in id.lower(),
-    "longcat":
-    lambda id: "longcat" in id.lower(),
-    "hunyuan":
-    lambda id: "hunyuan" in id.lower(),
-    "hunyuan15":
-    lambda id: "hunyuan15" in id.lower(),
-    "hyworld":
-    lambda id: "hyworld" in id.lower(),
-    "matrixgame":
-    lambda id: "matrix-game" in id.lower() or "matrixgame" in id.lower(),
-    "wanpipeline":
-    lambda id: "wanpipeline" in id.lower(),
-    "wanimagetovideo":
-    lambda id: "wanimagetovideo" in id.lower(),
-    "wandmdpipeline":
-    lambda id: "wandmdpipeline" in id.lower(),
-    "wancausaldmdpipeline":
-    lambda id: "wancausaldmdpipeline" in id.lower(),
-    "stepvideo":
-    lambda id: "stepvideo" in id.lower(),
-    "cosmos":
-    lambda id: "cosmos" in id.lower() and ("2.5" not in id.lower(
-    ) and "2_5" not in id.lower() and "25" not in id.lower()),
-    "cosmos25":
-    lambda id: "cosmos25" in id.lower(),
-    "turbodiffusion":
-    lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(),
-    "ltx2":
-    lambda id: "ltx2" in id.lower() or "ltx-2" in id.lower(),
-    # Add other pipeline architecture detectors
-}
-
-# Fallback configs when exact match isn't found but architecture is detected
-PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
-    "longcatimagetovideo": LongCatT2V480PConfig,
-    "longcatvideocontinuation": LongCatT2V480PConfig,
-    "longcat": LongCatT2V480PConfig,
-    "cosmos25": Cosmos25Config,
-    "hunyuan":
-    HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
-    "matrixgame": MatrixGameI2V480PConfig,
-    "hunyuan15":
-    Hunyuan15T2V480PConfig,  # Base Hunyuan15 config as fallback for any Hunyuan15 variant
-    "hyworld":
-    HYWorldConfig,  # HYWorld-specific config as fallback for any HYWorld variant
-    "wanpipeline":
-    WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
-    "wanimagetovideo": WanI2V480PConfig,
-    "wandmdpipeline": FastWan2_1_T2V_480P_Config,
-    "wancausaldmdpipeline": SelfForcingWanT2V480PConfig,
-    "stepvideo": StepVideoT2VConfig,
-    "turbodiffusion": TurboDiffusionT2V_1_3B_Config,
-    "ltx2": LTX2T2VConfig,
-    # Other fallbacks by architecture
-}
-
-
-def get_pipeline_config_cls_from_name(
-        pipeline_name_or_path: str) -> type[PipelineConfig]:
-    """Get the appropriate configuration class for a given pipeline name or path.
-
-    This function implements a multi-step lookup process to find the most suitable
-    configuration class for a given pipeline. It follows this order:
-    1. Exact match in the PIPE_NAME_TO_CONFIG
-    2. Partial match in the PIPE_NAME_TO_CONFIG
-    3. Fallback to class name in the model_index.json
-    4. else raise an error
-
-    Args:
-        pipeline_name_or_path (str): The name or path of the pipeline. This can be:
-            - A registered model ID (e.g., "FastVideo/FastHunyuan-diffusers")
-            - A local path to a model directory
-            - A model ID that will be downloaded
-
-    Returns:
-        Type[PipelineConfig]: The configuration class that best matches the pipeline.
-            This will be one of:
-            - A specific weight configuration class if an exact match is found
-            - A fallback configuration class based on the pipeline architecture
-            - The base PipelineConfig class if no matches are found
-
-    Note:
-        - For local paths, the function will verify the model configuration
-        - For remote models, it will attempt to download the model index
-        - Warning messages are logged when falling back to less specific configurations
-    """
-
-    pipeline_config_cls: type[PipelineConfig] | None = None
-
-    # First try exact match for specific weights
-    if pipeline_name_or_path in PIPE_NAME_TO_CONFIG:
-        pipeline_config_cls = PIPE_NAME_TO_CONFIG[pipeline_name_or_path]
-        return pipeline_config_cls
-
-    # Try partial matches (for local paths that might include the weight ID)
-    for registered_id, config_class in PIPE_NAME_TO_CONFIG.items():
-        if registered_id in pipeline_name_or_path:
-            pipeline_config_cls = config_class
-            break
-
-    # If no match, try to use the fallback config
-    if pipeline_config_cls is None:
-        if os.path.exists(pipeline_name_or_path):
-            config = verify_model_config_and_directory(pipeline_name_or_path)
-        else:
-            config = maybe_download_model_index(pipeline_name_or_path)
-        logger.warning(
-            "Trying to use the config from the model_index.json. FastVideo may not correctly identify the optimal config for this model in this situation."
-        )
-
-        pipeline_name = config["_class_name"]
-        # Try to determine pipeline architecture for fallback
-        for pipeline_type, detector in PIPELINE_DETECTOR.items():
-            if detector(pipeline_name.lower()):
-                pipeline_config_cls = PIPELINE_FALLBACK_CONFIG.get(
-                    pipeline_type)
-                break
-
-        if pipeline_config_cls is not None:
-            logger.warning(
-                "No match found for pipeline %s, using fallback config %s.",
-                pipeline_name_or_path, pipeline_config_cls)
-
-    if pipeline_config_cls is None:
-        raise ValueError(
-            f"No match found for pipeline {pipeline_name_or_path}, please check the pipeline name or path."
-        )
-
-    return pipeline_config_cls
@@ -196,6 +196,12 @@ class SelfForcingWan2_2_T2V480PConfig(Wan2_2_T2V_A14B_Config):
 # =============================================
 # ============= Matrix Game ===================
 # =============================================
+@dataclass
+class MatrixGameBaseI2V480PConfig(WanI2V480PConfig):
+    dit_config: DiTConfig = field(default_factory=MatrixGameWanVideoConfig)
+    flow_shift: float | None = 5.0
+
+
@dataclass
 class MatrixGameI2V480PConfig(WanI2V480PConfig):
    dit_config: DiTConfig = field(default_factory=MatrixGameWanVideoConfig)
@@ -28,6 +28,9 @@ class SamplingParam:
    keyboard_cond: Any | None = None  # Shape: (B, T, K)
    grid_sizes: Any | None = None  # Shape: (3,) [F,H,W]

+    # Camera control inputs (HYWorld)
+    pose: str | None = None  # Camera trajectory: pose string (e.g., 'w-31') or JSON file path
+
    # Refine inputs (LongCat 480p->720p upscaling)
    # Path-based refine (load stage1 video from disk, e.g. MP4)
    refine_from: str | None = None  # Path to stage1 video (480p output from distill)
@@ -55,10 +58,13 @@ class SamplingParam:
    num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
    height: int = 720
    width: int = 1280
+    height_sr: int = 1072
+    width_sr: int = 1920
    fps: int = 24

    # Denoising parameters
    num_inference_steps: int = 50
+    num_inference_steps_sr: int = 50
    guidance_scale: float = 1.0
    guidance_rescale: float = 0.0
    boundary_ratio: float | None = None
@@ -92,8 +98,7 @@ class SamplingParam:

    @classmethod
    def from_pretrained(cls, model_path: str) -> "SamplingParam":
-        from fastvideo.configs.sample.registry import (
-            get_sampling_param_cls_for_name)
+        from fastvideo.registry import get_sampling_param_cls_for_name
        sampling_cls = get_sampling_param_cls_for_name(model_path)
        if sampling_cls is not None:
            sampling_param: SamplingParam = sampling_cls()
@@ -5,15 +5,19 @@ from fastvideo.configs.sample.base import SamplingParam


@dataclass
-class Cosmos_Predict2_5_2B_Diffusers_SamplingParam(SamplingParam):
-    """Defaults for Cosmos 2.5 (Predict2.5) text-to-video diffusers-format model."""
-
-    height: int = 480
-    width: int = 832
-    num_frames: int = 121
+class Cosmos25SamplingParamBase(SamplingParam):
+    height: int = 704
+    width: int = 1280
+    num_frames: int = 77
    fps: int = 24
+    seed: int = 0

    guidance_scale: float = 7.0
-    # Official Cosmos2.5 sampling uses empty string as unconditional.
-    negative_prompt: str = ""
+    negative_prompt: str = (
+        "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, "
+        "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, "
+        "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, "
+        "low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, "
+        "unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. "
+        "Overall, the video is of poor quality.")
    num_inference_steps: int = 35
@@ -17,8 +17,8 @@ class Hunyuan15_480P_SamplingParam(SamplingParam):
    guidance_scale: float = 6.0
    prompt_attention_mask: list = field(default_factory=list)
    negative_attention_mask: list = field(default_factory=list)
-    # sigmas: list[float] | None = field(
-    #     default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1]))
+    sigmas: list[float] | None = field(
+        default_factory=lambda: list(np.linspace(1.0, 0.0, 50 + 1)[:-1]))

    negative_prompt: str = ""

@@ -29,10 +29,12 @@ class Hunyuan15_480P_SamplingParam(SamplingParam):


@dataclass
-class Hunyuan15_480P_Distilled_SamplingParam(Hunyuan15_480P_SamplingParam):
-    num_inference_steps: int = 8
-    fps: int = 24
+class Hunyuan15_480P_StepDistilled_I2V_SamplingParam(
+        Hunyuan15_480P_SamplingParam):
+    num_inference_steps: int = 12

+    height: int = 720
+    width: int = 1280
    guidance_scale: float = 1.0


@@ -40,3 +42,19 @@ class Hunyuan15_480P_Distilled_SamplingParam(Hunyuan15_480P_SamplingParam):
 class Hunyuan15_720P_SamplingParam(Hunyuan15_480P_SamplingParam):
    height: int = 720
    width: int = 1280
+
+
+@dataclass
+class Hunyuan15_720P_Distilled_I2V_SamplingParam(Hunyuan15_720P_SamplingParam):
+    guidance_scale: float = 1.0
+
+
+@dataclass
+class Hunyuan15_SR_1080P_SamplingParam(Hunyuan15_480P_SamplingParam):
+    height_sr: int = 1072
+    width_sr: int = 1920
+
+    num_inference_steps: int = 12
+    num_inference_steps_sr: int = 8
+
+    guidance_scale: float = 1.0
@@ -1,208 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-from collections.abc import Callable
-from typing import Any
-
-from fastvideo.configs.sample.hunyuan import (FastHunyuanSamplingParam,
-                                              HunyuanSamplingParam)
-from fastvideo.configs.sample.hyworld import HYWorld_SamplingParam
-from fastvideo.configs.sample.hunyuan15 import Hunyuan15_480P_SamplingParam, Hunyuan15_720P_SamplingParam, Hunyuan15_480P_Distilled_SamplingParam
-from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
-
-from fastvideo.configs.sample.cosmos import Cosmos_Predict2_2B_Video2World_SamplingParam
-from fastvideo.configs.sample.cosmos2_5 import Cosmos_Predict2_5_2B_Diffusers_SamplingParam
-from fastvideo.configs.sample.ltx2 import LTX2SamplingParam
-
-# isort: off
-from fastvideo.configs.sample.wan import (
-    FastWanT2V480P_SamplingParam,
-    Wan2_1_Fun_1_3B_InP_SamplingParam,
-    Wan2_2_I2V_A14B_SamplingParam,
-    Wan2_2_T2V_A14B_SamplingParam,
-    Wan2_2_TI2V_5B_SamplingParam,
-    WanI2V_14B_480P_SamplingParam,
-    WanI2V_14B_720P_SamplingParam,
-    WanT2V_1_3B_SamplingParam,
-    WanT2V_14B_SamplingParam,
-    Wan2_1_Fun_1_3B_Control_SamplingParam,
-    SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam,
-    SelfForcingWan2_2_T2V_A14B_480P_SamplingParam,
-    MatrixGame2_SamplingParam,
-)
-from fastvideo.configs.sample.turbodiffusion import (
-    TurboDiffusionT2V_1_3B_SamplingParam,
-    TurboDiffusionT2V_14B_SamplingParam,
-    TurboDiffusionI2V_A14B_SamplingParam,
-)
-# isort: on
-from fastvideo.logger import init_logger
-from fastvideo.utils import (maybe_download_model_index,
-                             verify_model_config_and_directory)
-
-logger = init_logger(__name__)
-# Registry maps specific model weights to their config classes
-SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
-    "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
-    "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
-    "weizhou03/SFHunyuanVideo-1.5-Diffusers-480p_t2v":
-    Hunyuan15_480P_SamplingParam,
-    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled":
-    Hunyuan15_480P_Distilled_SamplingParam,
-    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v":
-    Hunyuan15_480P_SamplingParam,
-    "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v":
-    Hunyuan15_720P_SamplingParam,
-    "FastVideo/HY-WorldPlay-Bidirectional-Diffusers": HYWorld_SamplingParam,
-    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
-
-    # Wan2.1
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam,
-    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam,
-    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
-    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
-    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers":
-    Wan2_1_Fun_1_3B_InP_SamplingParam,
-    "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers":
-    Wan2_1_Fun_1_3B_Control_SamplingParam,
-
-    # Wan2.2
-    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
-    "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers":
-    Wan2_2_TI2V_5B_SamplingParam,
-    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam,
-    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam,
-
-    # FastWan2.1
-    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480P_SamplingParam,
-
-    # FastWan2.2
-    "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
-
-    # Causal Self-Forcing Wan2.1
-    "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers":
-    SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam,
-
-    # Causal Self-Forcing Wan2.2
-    "rand0nmr/SFWan2.2-T2V-A14B-Diffusers":
-    SelfForcingWan2_2_T2V_A14B_480P_SamplingParam,
-    "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers":
-    SelfForcingWan2_2_T2V_A14B_480P_SamplingParam,
-
-    # Cosmos2
-    "nvidia/Cosmos-Predict2-2B-Video2World":
-    Cosmos_Predict2_2B_Video2World_SamplingParam,
-
-    # Cosmos2.5
-    "KyleShao/Cosmos-Predict2.5-2B-Diffusers":
-    Cosmos_Predict2_5_2B_Diffusers_SamplingParam,
-
-    # MatrixGame2.0 models
-    "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGame2_SamplingParam,
-    "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGame2_SamplingParam,
-    "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGame2_SamplingParam,
-
-    # TurboDiffusion models
-    "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers":
-    TurboDiffusionT2V_1_3B_SamplingParam,
-    "loayrashid/TurboWan2.1-T2V-14B-Diffusers":
-    TurboDiffusionT2V_14B_SamplingParam,
-    "loayrashid/TurboWan2.2-I2V-A14B-Diffusers":
-    TurboDiffusionI2V_A14B_SamplingParam,
-
-    # LTX-2 models
-    "Lightricks/LTX-2": LTX2SamplingParam,
-    "FastVideo/LTX2-Distilled-Diffusers": LTX2SamplingParam,
-
-    # Add other specific weight variants
-}
-
-# For determining pipeline type from model ID
-SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
-    "hunyuan":
-    lambda id: "hunyuan" in id.lower(),
-    "hunyuan15":
-    lambda id: "hunyuan15" in id.lower(),
-    "hyworld":
-    lambda id: "hyworld" in id.lower(),
-    "wanpipeline":
-    lambda id: "wanpipeline" in id.lower(),
-    "wanimagetovideo":
-    lambda id: "wanimagetovideo" in id.lower(),
-    "stepvideo":
-    lambda id: "stepvideo" in id.lower(),
-    "wandmdpipeline":
-    lambda id: "wandmdpipeline" in id.lower(),
-    "wancausaldmdpipeline":
-    lambda id: "wancausaldmdpipeline" in id.lower(),
-    "matrixgame":
-    lambda id: "matrixgame" in id.lower() or "matrix-game" in id.lower(),
-    "turbodiffusion":
-    lambda id: "turbodiffusion" in id.lower() or "turbowan" in id.lower(),
-    "cosmos25":
-    lambda id: "cosmos2_5" in id.lower(),
-    "cosmos":
-    lambda id: "cosmos" in id.lower() and "2_5" not in id.lower(),
-    "ltx2":
-    lambda id: "ltx2" in id.lower() or "ltx-2" in id.lower(),
-    # Add other pipeline architecture detectors
-}
-
-# Fallback configs when exact match isn't found but architecture is detected
-SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
-    "hunyuan":
-    HunyuanSamplingParam,  # Base Hunyuan config as fallback for any Hunyuan variant
-    "hunyuan15":
-    Hunyuan15_480P_SamplingParam,  # Base Hunyuan15 config as fallback for any Hunyuan15 variant
-    "hyworld":
-    HYWorld_SamplingParam,  # HYWorld-specific config as fallback for any HYWorld variant
-    "wanpipeline":
-    WanT2V_1_3B_SamplingParam,  # Base Wan config as fallback for any Wan variant
-    "wanimagetovideo": WanI2V_14B_480P_SamplingParam,
-    "wandmdpipeline": FastWanT2V480P_SamplingParam,
-    "wancausaldmdpipeline": SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam,
-    "stepvideo": StepVideoT2VSamplingParam,
-    "matrixgame": MatrixGame2_SamplingParam,
-    "turbodiffusion":
-    TurboDiffusionT2V_1_3B_SamplingParam,  # Default to T2V for fallback
-    "cosmos25": Cosmos_Predict2_5_2B_Diffusers_SamplingParam,
-    "cosmos": Cosmos_Predict2_2B_Video2World_SamplingParam,
-    "ltx2": LTX2SamplingParam,
-    # Other fallbacks by architecture
-}
-
-
-def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
-    """Get the appropriate sampling param for specific pretrained weights."""
-    # First try exact match for specific weights
-    if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY:
-        return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path]
-
-    # Try partial matches (for local paths that might include the weight ID)
-    for registered_id, config_class in SAMPLING_PARAM_REGISTRY.items():
-        if registered_id in pipeline_name_or_path:
-            return config_class
-
-    matrixgame_patterns = ["Matrix-Game", "Skywork--Matrix-Game", "matrixgame"]
-    for pattern in matrixgame_patterns:
-        if pattern.lower() in pipeline_name_or_path.lower():
-            return MatrixGame2_SamplingParam
-
-    if os.path.exists(pipeline_name_or_path):
-        config = verify_model_config_and_directory(pipeline_name_or_path)
-    else:
-        config = maybe_download_model_index(pipeline_name_or_path)
-
-    pipeline_name = config["_class_name"]
-
-    # If no match, try to use the fallback config
-    fallback_config = None
-    # Try to determine pipeline architecture for fallback
-    for pipeline_type, detector in SAMPLING_PARAM_DETECTOR.items():
-        if detector(pipeline_name.lower()):
-            fallback_config = SAMPLING_FALLBACK_PARAM.get(pipeline_type)
-            break
-
-    logger.warning(
-        "No match found for pipeline %s, using fallback sampling param %s.",
-        pipeline_name_or_path, fallback_config)
-    return fallback_config
@@ -127,10 +127,7 @@ def i2v_record_creator(batch: PreprocessBatch) -> list[dict[str, Any]]:
 def ode_text_only_record_creator(
        video_name: str, text_embedding: np.ndarray, caption: str,
        trajectory_latents: np.ndarray,
-        trajectory_timesteps: np.ndarray,
-        text_mask: np.ndarray | None = None,
-        text_embedding_2: np.ndarray | None = None,
-        text_mask_2: np.ndarray | None = None) -> dict[str, Any]:
+        trajectory_timesteps: np.ndarray) -> dict[str, Any]:
    """Create a text-only ODE trajectory record matching pyarrow_schema_ode_trajectory_text_only.

    Args:
@@ -168,25 +165,6 @@ def ode_text_only_record_creator(
        "trajectory_timesteps_dtype": str(trajectory_timesteps.dtype),
    })

-    if text_embedding_2 is not None:
-        record.update({
-            "text_embedding_2_bytes": text_embedding_2.tobytes(),
-            "text_embedding_2_shape": list(text_embedding_2.shape),
-            "text_embedding_2_dtype": str(text_embedding_2.dtype),
-        })
-    if text_mask is not None:
-        record.update({
-            "text_mask_bytes": text_mask.tobytes(),
-            "text_mask_shape": list(text_mask.shape),
-            "text_mask_dtype": str(text_mask.dtype),
-        })
-    if text_mask_2 is not None:
-        record.update({
-            "text_mask_2_bytes": text_mask_2.tobytes(),
-            "text_mask_2_shape": list(text_mask_2.shape),
-            "text_mask_2_dtype": str(text_mask_2.dtype),
-        })
-        
    return record


@@ -209,4 +187,4 @@ def text_only_record_creator(text_name: str, text_embedding: np.ndarray,
        "text_embedding_dtype": str(text_embedding.dtype),
        "caption": caption,
    }
-    return record
+    return record
@@ -90,15 +90,6 @@ pyarrow_schema_ode_trajectory_text_only = pa.schema([
    pa.field("text_embedding_shape", pa.list_(pa.int64())),
    # e.g., 'bfloat16' or 'float32'
    pa.field("text_embedding_dtype", pa.string()),
-    pa.field("text_embedding_2_bytes", pa.binary()),
-    pa.field("text_embedding_2_shape", pa.list_(pa.int64())),
-    pa.field("text_embedding_2_dtype", pa.string()),
-    pa.field("text_mask_bytes", pa.binary()),
-    pa.field("text_mask_shape", pa.list_(pa.int64())),
-    pa.field("text_mask_dtype", pa.string()),
-    pa.field("text_mask_2_bytes", pa.binary()),
-    pa.field("text_mask_2_shape", pa.list_(pa.int64())),
-    pa.field("text_mask_2_dtype", pa.string()),
    # --- ODE Trajectory ---
    pa.field("trajectory_latents_bytes", pa.binary()),
    pa.field("trajectory_latents_shape", pa.list_(pa.int64())),
@@ -124,4 +115,45 @@ pyarrow_schema_text_only = pa.schema([
    pa.field("text_embedding_dtype", pa.string()),
    # --- Metadata ---
    pa.field("caption", pa.string()),
-])
+])
+
+
+pyarrow_schema_matrixgame = pa.schema([
+    pa.field("id", pa.string()),
+    # --- Image/Video VAE latents ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("vae_latent_bytes", pa.binary()),
+    # e.g., [C, T, H, W] or [C, H, W]
+    pa.field("vae_latent_shape", pa.list_(pa.int64())),
+    # e.g., 'float32'
+    pa.field("vae_latent_dtype", pa.string()),
+    #I2V
+    pa.field("clip_feature_bytes", pa.binary()),
+    pa.field("clip_feature_shape", pa.list_(pa.int64())),
+    pa.field("clip_feature_dtype", pa.string()),
+    pa.field("first_frame_latent_bytes", pa.binary()),
+    pa.field("first_frame_latent_shape", pa.list_(pa.int64())),
+    pa.field("first_frame_latent_dtype", pa.string()),
+    # --- Action ---
+    pa.field("mouse_cond_bytes", pa.binary()),
+    pa.field("mouse_cond_shape", pa.list_(pa.int64())),  # [T, 2]
+    pa.field("mouse_cond_dtype", pa.string()),
+    pa.field("keyboard_cond_bytes", pa.binary()),
+    pa.field("keyboard_cond_shape", pa.list_(pa.int64())),  # [T, 4]
+    pa.field("keyboard_cond_dtype", pa.string()),
+    # I2V Validation
+    pa.field("pil_image_bytes", pa.binary()),
+    pa.field("pil_image_shape", pa.list_(pa.int64())),
+    pa.field("pil_image_dtype", pa.string()),
+    # --- Metadata ---
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),  # 'image' or 'video'
+    pa.field("width", pa.int64()),
+    pa.field("height", pa.int64()),
+    # -- Video-specific (can be null/default for images) ---
+    # Number of frames processed (e.g., 1 for image, N for video)
+    pa.field("num_frames", pa.int64()),
+    pa.field("duration_sec", pa.float64()),
+    pa.field("fps", pa.float64()),
+])
@@ -17,7 +17,6 @@ from PIL import Image
 from transformers import AutoTokenizer

 from fastvideo.logger import init_logger
-from fastvideo.utils import maybe_download_model

 logger = init_logger(__name__)

@@ -41,6 +40,7 @@ class PreprocessBatch:
    num_frames: int | None = None
    sample_frame_index: list[int] | None = None
    sample_num_frames: int | None = None
+    action_path: str | None = None

    # Processed data
    pixel_values: torch.Tensor | None = None
@@ -148,62 +148,6 @@ class DataValidationStage(DatasetFilterStage):
        return batch


-class ResolutionFilterStage(DatasetFilterStage):
-    """Stage for filtering data items based on resolution constraints."""
-
-    def __init__(self,
-                 max_h_div_w_ratio: float = 17 / 16,
-                 min_h_div_w_ratio: float = 8 / 16,
-                 max_height: int = 1024,
-                 max_width: int = 1024):
-        self.max_h_div_w_ratio = max_h_div_w_ratio
-        self.min_h_div_w_ratio = min_h_div_w_ratio
-        self.max_height = max_height
-        self.max_width = max_width
-
-    def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
-        """
-        Check if data item passes resolution filtering.
-        
-        Args:
-            batch: Dataset batch with resolution information
-            
-        Returns:
-            True if passes filter, False otherwise
-        """
-        # Only apply to videos
-        if not batch.is_video:
-            return True
-
-        if batch.resolution is None:
-            return False
-
-        height = batch.resolution.get("height", None)
-        width = batch.resolution.get("width", None)
-        if height is None or width is None:
-            return False
-
-        # Check aspect ratio
-        aspect = self.max_height / self.max_width
-        hw_aspect_thr = 1.5
-
-        return self.filter_resolution(
-            height,
-            width,
-            max_h_div_w_ratio=hw_aspect_thr * aspect,
-            min_h_div_w_ratio=1 / hw_aspect_thr * aspect,
-        )
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """Process does nothing for resolution filtering - filtering is handled by should_keep."""
-        return batch
-
-    def filter_resolution(self, h: int, w: int, max_h_div_w_ratio: float,
-                          min_h_div_w_ratio: float) -> bool:
-        """Filter based on height/width ratio."""
-        return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio
-
-
 class FrameSamplingStage(DatasetFilterStage):
    """Stage for temporal frame sampling and indexing."""

@@ -328,11 +272,6 @@ class VideoTransformStage(DatasetStage):
        video = rearrange(video, "t c h w -> c t h w")
        video = video.to(torch.uint8)

-        h, w = video.shape[-2:]
-        assert (
-            h / w <= 17 / 16 and h / w >= 8 / 16
-        ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({batch.path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}"
-
        video = video.float() / 127.5 - 1.0
        batch.pixel_values = video
        return batch
@@ -366,9 +305,11 @@ class ImageTransformStage(DatasetStage):
            image = self.transform_topcrop(image)
        elif self.transform is not None:
            image = self.transform(image)
+            image = image.float() / 127.5 - 1.0
+        else:
+            image = image.float() / 127.5 - 1.0

        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
-        image = image.float() / 127.5 - 1.0
        batch.pixel_values = image
        return batch

@@ -471,8 +412,11 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,

        # Initialize tokenizer
        tokenizer_path = os.path.join(args.model_path, "tokenizer")
-        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
-                                                  cache_dir=args.cache_dir)
+        if os.path.exists(tokenizer_path):
+            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
+                                                      cache_dir=args.cache_dir)
+        else:
+            tokenizer = None

        # Initialize processing stages
        self._init_stages(args, transform, transform_topcrop, tokenizer)
@@ -484,8 +428,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
                     tokenizer) -> None:
        """Initialize all processing stages."""
        self.validation_stage = DataValidationStage()
-        self.resolution_filter_stage = ResolutionFilterStage(
-            max_height=args.max_height, max_width=args.max_width)
        self.frame_sampling_stage = FrameSamplingStage(
            num_frames=args.num_frames,
            train_fps=args.train_fps,
@@ -496,11 +438,14 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
        self.video_transform_stage = VideoTransformStage(transform)
        self.image_transform_stage = ImageTransformStage(
            transform, transform_topcrop)
-        self.text_encoding_stage = TextEncodingStage(
-            tokenizer=tokenizer,
-            text_max_length=args.text_max_length,
-            cfg_rate=args.training_cfg_rate,
-            seed=self.seed)
+        if tokenizer is not None:
+            self.text_encoding_stage = TextEncodingStage(
+                tokenizer=tokenizer,
+                text_max_length=args.text_max_length,
+                cfg_rate=args.training_cfg_rate,
+                seed=self.seed)
+        else:
+            self.text_encoding_stage = None

    def _load_raw_data(self) -> list[dict]:
        """Load raw data from JSON files."""
@@ -522,6 +467,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
        # Update paths with folder prefix
        for item in data_items:
            item["path"] = opj(folder, item["path"])
+            if "action_path" in item and item["action_path"]:
+                item["action_path"] = opj(folder, item["action_path"])

        return data_items

@@ -533,7 +480,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
        # Initialize counters
        filter_counts = {
            "validation_failed": 0,
-            "resolution_failed": 0,
            "frame_sampling_failed": 0
        }
        sample_num_frames: list[int] = []
@@ -543,7 +489,8 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
                                    cap=item["cap"],
                                    resolution=item.get("resolution"),
                                    fps=item.get("fps"),
-                                    duration=item.get("duration"))
+                                    duration=item.get("duration"),
+                                    action_path=item.get("action_path"))

            # Apply filtering stages
            if not self._apply_filter_stages(batch, filter_counts):
@@ -568,10 +515,6 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
            filter_counts["validation_failed"] += 1
            return False

-        if not self.resolution_filter_stage.should_keep(batch):
-            filter_counts["resolution_failed"] += 1
-            return False
-
        if not self.frame_sampling_stage.should_keep(batch):
            filter_counts["frame_sampling_failed"] += 1
            return False
@@ -583,10 +526,9 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
                             after_count: int):
        """Log filtering statistics."""
        logger.info(
-            "validation_failed: %d, resolution_failed: %d, frame_sampling_failed: %d, "
+            "validation_failed: %d, frame_sampling_failed: %d, "
            "Counter(sample_num_frames): %s, before filter: %d, after filter: %d",
            filter_counts['validation_failed'],
-            filter_counts['resolution_failed'],
            filter_counts['frame_sampling_failed'], Counter(sample_num_frames),
            before_count, after_count)

@@ -605,20 +547,27 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
        # Apply transformation stages
        batch = self.video_transform_stage.process(batch)
        batch = self.image_transform_stage.process(batch)
-        batch = self.text_encoding_stage.process(batch)
+        if self.text_encoding_stage is not None:
+            batch = self.text_encoding_stage.process(batch)

        # Build result dictionary
        result = {
            "pixel_values": batch.pixel_values,
-            "text": batch.text,
-            "input_ids": batch.input_ids,
-            "cond_mask": batch.cond_mask,
            "path": batch.path,
        }

+        if batch.text is not None:
+            result["text"] = batch.text
+            result["input_ids"] = batch.input_ids
+            result["cond_mask"] = batch.cond_mask
+
        # Add video-specific fields
        if batch.is_video:
            result.update({"fps": batch.fps, "duration": batch.duration})
+        
+        # Add action_path
+        if batch.action_path:
+            result["action_path"] = batch.action_path

        return result

@@ -656,7 +605,7 @@ class TextDataset(torch.utils.data.IterableDataset,
        self.seed = seed

        # Initialize tokenizer
-        tokenizer_path = os.path.join(maybe_download_model(args.model_path), "tokenizer")
+        tokenizer_path = os.path.join(args.model_path, "tokenizer")
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
                                                  cache_dir=args.cache_dir)

@@ -759,4 +708,4 @@ class TextDataset(torch.utils.data.IterableDataset,

    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        """Load state dict from checkpoint."""
-        self.processed_batches = state_dict["processed_batches"]
+        self.processed_batches = state_dict["processed_batches"]
@@ -154,14 +154,8 @@ def collate_rows_from_parquet_schema(rows,
                    ) if rng else random.random()) < cfg_rate:
                        data = np.zeros((512, 4096), dtype=np.float32)
                    else:
-                        if row[f"{tensor_name}_dtype"] == "float32":
-                            data = np.frombuffer(
-                                bytes_data, dtype=np.float32).reshape(shape).copy()
-                        elif row[f"{tensor_name}_dtype"] == "int64":
-                            data = np.frombuffer(
-                                bytes_data, dtype=np.int64).reshape(shape).copy()
-                        else:
-                            raise ValueError(f"Unsupported dtype: {row[f"{tensor_name}_dtype"]}")
+                        data = np.frombuffer(
+                            bytes_data, dtype=np.float32).reshape(shape).copy()
                    tensor = torch.from_numpy(data)
                    # if len(data.shape) == 3:
                    #     B, L, D = tensor.shape
@@ -174,7 +168,7 @@ def collate_rows_from_parquet_schema(rows,
                tensor_list.append(torch.zeros(0, dtype=torch.bfloat16))

        # Stack tensors with special handling for text embeddings
-        if tensor_name == 'null':
+        if tensor_name == 'text_embedding':
            # Handle text embeddings with padding
            padded_tensors = []
            attention_masks = []
@@ -18,6 +18,7 @@ __all__ = [
    "cleanup_dist_env_and_memory",
    "model_parallel_is_initialized",
    "maybe_init_distributed_environment_and_model_parallel",
+    "warmup_sequence_parallel_communication",

    # World group
    "get_world_group",
@@ -7,10 +7,17 @@ import torch.distributed
 from fastvideo.distributed.parallel_state import (get_sp_group,
                                                  get_sp_parallel_rank,
                                                  get_sp_world_size,
-                                                  get_tp_group)
+                                                  get_tp_group,
+                                                  model_parallel_is_initialized)
 from fastvideo.distributed.utils import (unpad_sequence_tensor,
                                         compute_padding_for_sp,
                                         pad_sequence_tensor)
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+# Track if SP communication has been warmed up
+_sp_warmup_done = False


 def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
@@ -101,3 +108,86 @@ def sequence_model_parallel_shard(input_: torch.Tensor,
    input_ = input_.movedim(0, dim)

    return input_, original_seq_len
+
+
+def warmup_sequence_parallel_communication(
+        device: torch.device | None = None) -> None:
+    """Warmup NCCL communicators for sequence parallel all-to-all operations.
+    
+    The first NCCL collective operation is slow due to lazy communicator
+    initialization. This function runs dummy all-to-all operations to
+    trigger the initialization upfront, before the first real forward pass.
+    
+    Args:
+        device: Device to use for warmup tensors. If None, uses CUDA device 0.
+    """
+    global _sp_warmup_done
+
+    if _sp_warmup_done:
+        return
+
+    if not model_parallel_is_initialized():
+        return
+
+    sp_world_size = get_sp_world_size()
+    if sp_world_size <= 1:
+        _sp_warmup_done = True
+        return
+
+    if device is None:
+        device = torch.device("cuda")
+
+    logger.info("Warming up sequence parallel communication (SP=%d)...",
+                sp_world_size)
+
+    # Use small but representative tensor shapes for warmup
+    # Shape: [batch, seq_len, num_heads, head_dim]
+    # The all-to-all patterns used in attention:
+    #   1. scatter_dim=2 (heads), gather_dim=1 (seq) - before attention
+    #   2. scatter_dim=1 (seq), gather_dim=2 (heads) - after attention
+    batch_size = 1
+    seq_len_per_rank = 16  # Small sequence per rank
+    num_heads = sp_world_size * 4  # Must be divisible by sp_world_size
+    head_dim = 64
+
+    # Create dummy tensor for warmup
+    dummy = torch.zeros(batch_size,
+                        seq_len_per_rank,
+                        num_heads,
+                        head_dim,
+                        device=device,
+                        dtype=torch.bfloat16)
+
+    # Warmup pattern 1: scatter heads, gather sequence (before attention)
+    _ = sequence_model_parallel_all_to_all_4D(dummy,
+                                              scatter_dim=2,
+                                              gather_dim=1)
+
+    # Warmup pattern 2: scatter sequence, gather heads (after attention)
+    dummy2 = torch.zeros(batch_size,
+                         seq_len_per_rank * sp_world_size,
+                         num_heads // sp_world_size,
+                         head_dim,
+                         device=device,
+                         dtype=torch.bfloat16)
+    _ = sequence_model_parallel_all_to_all_4D(dummy2,
+                                              scatter_dim=1,
+                                              gather_dim=2)
+
+    # Warmup all-gather (used for replicated tokens)
+    dummy3 = torch.zeros(batch_size,
+                         8,
+                         num_heads // sp_world_size,
+                         head_dim,
+                         device=device,
+                         dtype=torch.bfloat16)
+    _ = sequence_model_parallel_all_gather(dummy3, dim=2)
+
+    # Synchronize to ensure warmup completes
+    torch.cuda.synchronize(device)
+
+    # Clean up
+    del dummy, dummy2, dummy3
+
+    _sp_warmup_done = True
+    logger.info("Sequence parallel communication warmup complete.")
@@ -9,6 +9,7 @@ diffusion models.
 import math
 import os
 import re
+import threading
 import time
 from copy import deepcopy
 from typing import Any
@@ -31,6 +32,21 @@ from fastvideo.worker.executor import Executor
 logger = init_logger(__name__)


+def _infer_latent_batch_size(batch: ForwardBatch) -> int:
+    if isinstance(batch.prompt, list):
+        latent_batch_size = len(batch.prompt)
+    elif batch.prompt is not None:
+        latent_batch_size = 1
+    elif batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0:
+        latent_batch_size = batch.prompt_embeds[0].shape[0]
+    else:
+        raise ValueError(
+            "Cannot infer batch size from batch; no prompt or prompt_embeds found"
+        )
+    latent_batch_size *= batch.num_videos_per_prompt
+    return latent_batch_size
+
+
 class VideoGenerator:
    """
    A unified class for generating videos using diffusion models.
@@ -372,8 +388,31 @@ class VideoGenerator:

        # Run inference
        start_time = time.perf_counter()
-        output_batch = self.executor.execute_forward(batch, fastvideo_args)
-        samples = output_batch.output
+
+        # Execute forward pass in a new thread for non-blocking tensor allocation
+        result_container = {}
+
+        def execute_forward_thread():
+            result_container['output_batch'] = self.executor.execute_forward(
+                batch, fastvideo_args)
+
+        thread = threading.Thread(target=execute_forward_thread)
+        thread.start()
+        latent_batch_size = _infer_latent_batch_size(batch)
+        samples = torch.empty((latent_batch_size, 3, sampling_param.num_frames,
+                               sampling_param.height, sampling_param.width),
+                              device='cpu',
+                              pin_memory=fastvideo_args.pin_cpu_memory)
+        thread.join()
+
+        output_batch = result_container['output_batch']
+        if output_batch.output.shape == samples.shape:
+            samples.copy_(output_batch.output)
+        else:
+            logger.warning(
+                "Output shape %s does not match expected shape %s; use slow path",
+                output_batch.output.shape, samples.shape)
+            samples = output_batch.output.cpu()
        logging_info = output_batch.logging_info

        gen_time = time.perf_counter() - start_time
@@ -10,7 +10,7 @@ from enum import Enum
 from typing import Any, TYPE_CHECKING

 from fastvideo.configs.configs import PreprocessConfig
-from fastvideo.configs.pipelines.base import PipelineConfig, STA_Mode
+from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.utils import clean_cli_args
 from fastvideo.layers.quantization import QUANTIZATION_METHODS, QuantizationMethods
 from fastvideo.logger import init_logger
@@ -141,8 +141,6 @@ class FastVideoArgs:

    # STA (Sliding Tile Attention) parameters
    mask_strategy_file_path: str | None = None
-    STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
-    skip_time_steps: int = 15

    # Compilation
    enable_torch_compile: bool = False
@@ -179,6 +177,7 @@ class FastVideoArgs:
    model_loaded: dict[str, bool] = field(default_factory=lambda: {
        "transformer": True,
        "vae": True,
+        "upsampler": True,
    })

    override_text_encoder_safetensors: str | None = None  # path to safetensors file for text encoder override
@@ -464,20 +463,6 @@ class FastVideoArgs:
        )

        # STA (Sliding Tile Attention) parameters
-        parser.add_argument(
-            "--STA-mode",
-            type=str,
-            default=FastVideoArgs.STA_mode.value,
-            choices=[mode.value for mode in STA_Mode],
-            help=
-            "STA mode contains STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None",
-        )
-        parser.add_argument(
-            "--skip-time-steps",
-            type=int,
-            default=FastVideoArgs.skip_time_steps,
-            help="Number of time steps to warmup (full attention) for STA",
-        )
        parser.add_argument(
            "--mask-strategy-file-path",
            type=str,
@@ -830,7 +815,6 @@ class TrainingArgs(FastVideoArgs):
    precedence.
    """
    data_path: str = ""
-    data_path_2: str | None = None
    dataloader_num_workers: int = 0
    num_height: int = 0
    num_width: int = 0
@@ -860,7 +844,6 @@ class TrainingArgs(FastVideoArgs):
    validation_sampling_steps: str = ""
    validation_guidance_scale: str = ""
    validation_steps: float = 0.0
-    visualization_steps: float = 0.0
    log_validation: bool = False
    trackers: list[str] = dataclasses.field(default_factory=list)
    tracker_project_name: str = ""
@@ -876,7 +859,6 @@ class TrainingArgs(FastVideoArgs):
    num_train_epochs: int = 0
    max_train_steps: int = 0
    gradient_accumulation_steps: int = 0
-    optimizer_type: str = "adamw"
    learning_rate: float = 0.0
    scale_lr: bool = False
    lr_scheduler: str = "constant"
@@ -934,6 +916,7 @@ class TrainingArgs(FastVideoArgs):
    training_state_checkpointing_steps: int = 0  # for resuming training
    weight_only_checkpointing_steps: int = 0  # for inference
    log_visualization: bool = False
+    visualization_steps: int = 0
    # simulate generator forward to match inference
    simulate_generator_forward: bool = False
    warp_denoising_step: bool = False
@@ -946,8 +929,6 @@ class TrainingArgs(FastVideoArgs):
    same_step_across_blocks: bool = False  # Use same exit timestep for all blocks
    last_step_only: bool = False  # Only use the last timestep for training
    context_noise: int = 0  # Context noise level for cache updates
-    use_context_forcing: bool = False
-    use_ode_init: bool = False

    @classmethod
    def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -1003,10 +984,6 @@ class TrainingArgs(FastVideoArgs):
                            type=str,
                            required=True,
                            help="Path to parquet files")
-        parser.add_argument("--data-path-2",
-                            type=str,
-                            required=False,
-                            help="Path to parquet files")
        parser.add_argument("--dataloader-num-workers",
                            type=int,
                            required=True,
@@ -1100,12 +1077,12 @@ class TrainingArgs(FastVideoArgs):
        parser.add_argument("--validation-steps",
                            type=float,
                            help="Number of validation steps")
-        parser.add_argument("--visualization-steps",
-                            type=float,
-                            help="Number of visualization steps")
        parser.add_argument("--log-validation",
                            action=StoreBoolean,
                            help="Whether to log validation results")
+        parser.add_argument("--visualization-steps",
+                            type=int,
+                            help="Number of visualization steps")
        parser.add_argument("--tracker-project-name",
                            type=str,
                            help="Project name for tracking")
@@ -1151,10 +1128,6 @@ class TrainingArgs(FastVideoArgs):
        parser.add_argument("--gradient-accumulation-steps",
                            type=int,
                            help="Number of steps to accumulate gradients")
-        parser.add_argument("--optimizer-type",
-                            type=str,
-                            choices=["adamw", "muon"],
-                            help="Optimizer type")
        parser.add_argument("--learning-rate",
                            type=float,
                            required=True,
@@ -1381,12 +1354,6 @@ class TrainingArgs(FastVideoArgs):
                            type=int,
                            default=TrainingArgs.context_noise,
                            help="Context noise level for cache updates")
-        parser.add_argument("--use-context-forcing",
-                            action=StoreBoolean,
-                            help="Whether to use context forcing")
-        parser.add_argument("--use-ode-init",
-                            action=StoreBoolean,
-                            help="Whether to use ODE init")

        return parser

@@ -193,3 +193,46 @@ class ImageProcessor:
        tensor = 2.0 * tensor - 1.0

        return tensor
+
+    def _preprocess_cosmos25(
+        self,
+        image: PIL.Image.Image | np.ndarray | torch.Tensor,
+        height: int,
+        width: int,
+    ) -> torch.Tensor:
+        """
+        Cosmos-Predict2.5-style preprocessing for image2world:
+        - aspect-preserving resize (scale so both dims >= target)
+        - center crop to (height, width)
+        - normalize to [-1, 1]
+
+        Returns:
+            torch.Tensor: (1, 3, height, width) in [-1, 1]
+        """
+        if not isinstance(image, PIL.Image.Image):
+            # Reuse the generic preprocess for non-PIL inputs.
+            return self.preprocess(image, height=height, width=width)
+
+        # Ensure RGB (official uses torchvision.to_tensor on PIL, which yields C=3 for RGB inputs)
+        image = image.convert("RGB")
+
+        # Make sure target dims are multiples of VAE scale factor (FastVideo convention)
+        height = height - (height % self.vae_scale_factor)
+        width = width - (width % self.vae_scale_factor)
+
+        orig_w, orig_h = image.size
+        # Scale so that resized image fully covers the target box (like official resize_input()).
+        scale = max(width / float(orig_w), height / float(orig_h))
+        resized_w = int(np.ceil(scale * orig_w))
+        resized_h = int(np.ceil(scale * orig_h))
+
+        image = image.resize((resized_w, resized_h),
+                             resample=PIL.Image.Resampling.BILINEAR)
+
+        # Center crop
+        left = max(0, (resized_w - width) // 2)
+        top = max(0, (resized_h - height) // 2)
+        image = image.crop((left, top, left + width, top + height))
+
+        image_np = np.array(image, dtype=np.float32) / 255.0  # [0,1]
+        return self._normalize_to_tensor(image_np)  # -> [-1,1], (1,3,H,W)
@@ -27,7 +27,7 @@ RESET = '\033[0;0m'
 _warned_local_main_process = False
 _warned_main_process = False

-_FORMAT = (f"{FASTVIDEO_LOGGING_PREFIX}%(levelname)s %(asctime)s "
+_FORMAT = (f"{FASTVIDEO_LOGGING_PREFIX}%(levelname)s %(asctime)s.%(msecs)03d "
           "[%(filename)s:%(lineno)d] %(message)s")
 _DATE_FORMAT = "%m-%d %H:%M:%S"

@@ -102,6 +102,7 @@ def _info(logger: Logger,
        - When both are False, the message will be logged from all processes
        - By default, only logs from processes with LOCAL_RANK=0
    """
+    is_distributed = int(os.environ.get("WORLD_SIZE", 1)) > 1
    try:
        local_rank = int(os.environ["LOCAL_RANK"])
        rank = int(os.environ["RANK"])
@@ -118,20 +119,22 @@ def _info(logger: Logger,

    global _warned_local_main_process, _warned_main_process

-    if not _warned_local_main_process and local_main_process_only:
-        logger.warning(
-            '%s By default, logger.info(..) will only log from the local main process. Set logger.info(..., is_local_main_process=False) to log from all processes.%s',
-            GREEN,
-            RESET,
-        )
-        _warned_local_main_process = True
-    if not _warned_main_process and main_process_only:
-        logger.warning(
-            '%s is_main_process_only is set to True, logging only from the main (RANK==0) process.%s',
-            GREEN,
-            RESET,
-        )
-        _warned_main_process = True
+    # Only show process-awareness warnings when actually running distributed
+    if is_distributed:
+        if not _warned_local_main_process and local_main_process_only:
+            logger.warning(
+                '%s By default, logger.info(..) will only log from the local main process. Set logger.info(..., is_local_main_process=False) to log from all processes.%s',
+                GREEN,
+                RESET,
+            )
+            _warned_local_main_process = True
+        if not _warned_main_process and main_process_only:
+            logger.warning(
+                '%s is_main_process_only is set to True, logging only from the main (RANK==0) process.%s',
+                GREEN,
+                RESET,
+            )
+            _warned_main_process = True

    if not main_process_only and not local_main_process_only:
        logger.log(logging.INFO, msg, *args, stacklevel=2, **kwargs)
@@ -1325,3 +1325,6 @@ def decode_audio(
    decoded_audio = audio_decoder(latent)
    decoded_audio = vocoder(decoded_audio).squeeze(0).float()
    return decoded_audio
+
+# Entry point for model registry
+EntryClass = [LTX2AudioEncoder, LTX2AudioDecoder, LTX2Vocoder]
@@ -1,808 +0,0 @@
-# Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved.
-#
-# 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, List
-import math
-
-import torch
-# import torch._dynamo
-# torch._dynamo.config.cache_size_limit = 128
-# try:
-#     torch._dynamo.config.recompile_limit = 128
-# except AttributeError:
-#     pass
-import torch.nn as nn
-import torch.nn.functional as F
-
-from torch.nn.attention.flex_attention import create_block_mask, flex_attention
-from torch.nn.attention.flex_attention import BlockMask
-flex_attention = torch.compile(
-    flex_attention, dynamic=False, mode="default")
-import torch.distributed as dist
-
-from fastvideo.attention import LocalAttention
-from fastvideo.forward_context import set_forward_context
-from fastvideo.configs.models.dits import HunyuanVideo15Config
-from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual,
-                                        ScaleResidualLayerNormScaleShift)
-from fastvideo.layers.linear import ReplicatedLinear
-# TODO(will-PY-refactor): RMSNorm ....
-from fastvideo.layers.mlp import MLP
-from fastvideo.layers.rotary_embedding import get_rotary_pos_embed, _apply_rotary_emb
-from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed,
-                                               unpatchify)
-from fastvideo.models.dits.base import CachableDiT
-from fastvideo.platforms import AttentionBackendEnum
-from fastvideo.logger import init_logger
-
-from fastvideo.models.dits.hunyuanvideo15 import (
-    HunyuanRMSNorm, 
-    HunyuanVideo15TimeEmbedding, 
-    HunyuanVideo15ByT5TextProjection,
-    HunyuanVideo15ImageProjection,
-    SingleTokenRefiner,
-    FinalLayer)
-
-logger = init_logger(__name__)
-
-class MMDoubleStreamBlock(nn.Module):
-    """
-    A multimodal DiT block with separate modulation for text and image/video,
-    using distributed attention and linear layers.
-    """
-
-    def __init__(
-        self,
-        hidden_size: int,
-        num_attention_heads: int,
-        mlp_ratio: float,
-        local_attn_size: int = -1,
-        sink_size: int = 0,
-        dtype: torch.dtype | None = None,
-        supported_attention_backends: tuple[AttentionBackendEnum, ...]
-        | None = None,
-        prefix: str = "",
-    ):
-        super().__init__()
-
-        self.deterministic = False
-        self.num_attention_heads = num_attention_heads
-        self.head_dim = hidden_size // num_attention_heads
-        self.local_attn_size = local_attn_size
-        self.sink_size = sink_size
-        mlp_hidden_dim = int(hidden_size * mlp_ratio)
-
-        # Image modulation components
-        self.img_mod = ModulateProjection(
-            hidden_size,
-            factor=6,
-            act_layer="silu",
-            dtype=dtype,
-            prefix=f"{prefix}.img_mod",
-        )
-
-        # Fused operations for image stream
-        self.img_attn_norm = LayerNormScaleShift(hidden_size,
-                                                 norm_type="layer",
-                                                 elementwise_affine=False,
-                                                 dtype=dtype)
-        self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
-            hidden_size,
-            norm_type="layer",
-            elementwise_affine=False,
-            dtype=dtype)
-        self.img_mlp_residual = ScaleResidual()
-
-        # Image attention components
-        self.img_attn_qkv = ReplicatedLinear(hidden_size,
-                                             hidden_size * 3,
-                                             bias=True,
-                                             params_dtype=dtype,
-                                             prefix=f"{prefix}.img_attn_qkv")
-
-        self.img_attn_q_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype)
-        self.img_attn_k_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype)
-
-        self.img_attn_proj = ReplicatedLinear(hidden_size,
-                                              hidden_size,
-                                              bias=True,
-                                              params_dtype=dtype,
-                                              prefix=f"{prefix}.img_attn_proj")
-
-        self.img_mlp = MLP(hidden_size,
-                           mlp_hidden_dim,
-                           bias=True,
-                           dtype=dtype,
-                           prefix=f"{prefix}.img_mlp")
-
-        # Text modulation components
-        self.txt_mod = ModulateProjection(
-            hidden_size,
-            factor=6,
-            act_layer="silu",
-            dtype=dtype,
-            prefix=f"{prefix}.txt_mod",
-        )
-
-        # Fused operations for text stream
-        self.txt_attn_norm = LayerNormScaleShift(hidden_size,
-                                                 norm_type="layer",
-                                                 elementwise_affine=False,
-                                                 dtype=dtype)
-        self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift(
-            hidden_size,
-            norm_type="layer",
-            elementwise_affine=False,
-            dtype=dtype)
-        self.txt_mlp_residual = ScaleResidual()
-
-        # Text attention components
-        self.txt_attn_qkv = ReplicatedLinear(hidden_size,
-                                             hidden_size * 3,
-                                             bias=True,
-                                             params_dtype=dtype)
-
-        # QK norm layers for text
-        self.txt_attn_q_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype)
-        self.txt_attn_k_norm = HunyuanRMSNorm(self.head_dim, eps=1e-6, dtype=dtype)
-
-        self.txt_attn_proj = ReplicatedLinear(hidden_size,
-                                              hidden_size,
-                                              bias=True,
-                                              params_dtype=dtype)
-
-        self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype)
-
-        self.max_attention_size = 21 * 1590 if local_attn_size == -1 else local_attn_size * 1590
-
-        self.attn = LocalAttention(
-            num_heads=self.num_attention_heads,
-            head_size=self.head_dim,
-            supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.attn"
-        )
-
-    def forward_txt(
-        self,
-        txt: torch.Tensor,
-        vec: torch.Tensor,
-        cache_txt: bool = False,
-    ):
-        txt_mod_outputs = self.txt_mod(vec)
-        (
-            txt_attn_shift,
-            txt_attn_scale,
-            txt_attn_gate,
-            txt_mlp_shift,
-            txt_mlp_scale,
-            txt_mlp_gate,
-        ) = torch.chunk(txt_mod_outputs, 6, dim=-1)
-
-        # Prepare text for attention using fused operation
-        txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale)
-
-        # Get QKV for text
-        txt_qkv, _ = self.txt_attn_qkv(txt_attn_input)
-        batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1]
-
-        # Split QKV
-        txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3,
-                               self.num_attention_heads, -1)
-        txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :,
-                                                                          2]
-        # Apply QK-Norm if needed
-        txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype)
-        txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype)
-
-        t_kv = {}
-        if cache_txt:
-            t_kv["k_txt"] = txt_k
-            t_kv["v_txt"] = txt_v
-
-        txt_attn = self.attn(txt_q, txt_k, txt_v)
-        # Process text attention output
-        txt_attn_out, _ = self.txt_attn_proj(
-            txt_attn.reshape(batch_size, text_seq_len, -1))
-
-        # Use fused operation for residual connection, normalization, and modulation
-        txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm(
-            txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale)
-
-        # Process text MLP
-        txt_mlp_out = self.txt_mlp(txt_mlp_input)
-        txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate)
-
-        return txt, t_kv
-
-    def forward_vision(
-        self,
-        img: torch.Tensor,
-        vec: torch.Tensor,
-        freqs_cis: tuple,
-        block_mask: BlockMask,
-        kv_cache: dict | None = None,
-        txt_kv_cache: list | None = None,
-        current_start: int = 0,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        # Process modulation vectors
-        if vec.dim() == 3:
-            img_mod_outputs = self.img_mod(vec).unflatten(dim=-1, sizes=(6, -1))
-            (
-                img_attn_shift,
-                img_attn_scale,
-                img_attn_gate,
-                img_mlp_shift,
-                img_mlp_scale,
-                img_mlp_gate,
-            ) = torch.chunk(img_mod_outputs, 6, dim=2)
-        else:
-            img_mod_outputs = self.img_mod(vec)
-            (
-                img_attn_shift,
-                img_attn_scale,
-                img_attn_gate,
-                img_mlp_shift,
-                img_mlp_scale,
-                img_mlp_gate,
-            ) = torch.chunk(img_mod_outputs, 6, dim=-1)
-
-        # Prepare image for attention using fused operation
-        img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale)
-        # Get QKV for image
-        img_qkv, _ = self.img_attn_qkv(img_attn_input)
-        batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1]
-
-        # Split QKV
-        img_qkv = img_qkv.view(batch_size, image_seq_len, 3,
-                               self.num_attention_heads, -1)
-        img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :,
-                                                                          2]
-
-        # Apply QK-Norm if needed
-        img_q = self.img_attn_q_norm(img_q).to(img_v)
-        img_k = self.img_attn_k_norm(img_k).to(img_v)
-
-        # Apply rotary embeddings
-        cos, sin = freqs_cis
-        img_q = _apply_rotary_emb(img_q, cos, sin, is_neox_style=False)
-        img_k = _apply_rotary_emb(img_k, cos, sin, is_neox_style=False)
-
-        # Apply flex_attention
-        # Does not support SP padding for now
-        if kv_cache is None:
-            q = img_q
-            k = torch.cat([img_k, txt_kv_cache["k_txt"]], dim=1)
-            v = torch.cat([img_v, txt_kv_cache["v_txt"]], dim=1)
-            # Padding for flex attention
-            padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
-            padded_kv_length = math.ceil(k.shape[1] / 128) * 128 - k.shape[1]
-            padded_roped_query = torch.cat(
-                [q,
-                    torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],
-                                device=q.device, dtype=v.dtype)],
-                dim=1
-            )
-
-            padded_roped_key = torch.cat(
-                [k, torch.zeros([k.shape[0], padded_kv_length, k.shape[2], k.shape[3]],
-                                        device=k.device, dtype=v.dtype)],
-                dim=1
-            )
-
-            padded_v = torch.cat(
-                [v, torch.zeros([v.shape[0], padded_kv_length, v.shape[2], v.shape[3]],
-                                device=v.device, dtype=v.dtype)],
-                dim=1
-            )
-
-            img_attn = flex_attention(
-                query=padded_roped_query.transpose(2, 1),
-                key=padded_roped_key.transpose(2, 1),
-                value=padded_v.transpose(2, 1),
-                block_mask=block_mask
-            )[:, :, :-padded_length].transpose(2, 1)
-
-            assert img_attn.shape[1] == image_seq_len
-            updated_kv_cache = None
-        else:
-            current_end = current_start + img_q.shape[1]
-            num_new_tokens = img_q.shape[1]
-            sink_tokens = self.sink_size * 1590
-            # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache
-            kv_cache_size = self.max_attention_size
-            
-            # Clone cache to avoid in-place modification during gradient checkpointing
-            k_cache = kv_cache["k"].clone()
-            v_cache = kv_cache["v"].clone()
-            
-            if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
-                    num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
-                # Calculate the number of new tokens added in this step
-                # Shift existing cache content left to discard oldest tokens
-                # Clone the source slice to avoid overlapping memory error
-                num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
-                num_rolled_tokens = kv_cache_size - num_new_tokens - sink_tokens
-                k_cache[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                    k_cache[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                v_cache[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                    v_cache[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                # Insert the new keys/values at the end
-                local_end_index = kv_cache["local_end_index"].item() + current_end - \
-                    kv_cache["global_end_index"].item() - num_evicted_tokens
-                local_start_index = local_end_index - num_new_tokens
-                assert local_end_index == self.max_attention_size
-            else:
-                # Assign new keys/values directly up to current_end
-                local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
-                local_start_index = local_end_index - num_new_tokens
-
-            assert local_start_index >= 0
-            q = img_q
-            k = torch.cat([k_cache[:, :local_start_index], img_k, txt_kv_cache["k_txt"]], dim=1)
-            v = torch.cat([v_cache[:, :local_start_index], img_v, txt_kv_cache["v_txt"]], dim=1)
-            img_attn = self.attn(q, k, v)
-
-            k_cache[:, local_start_index:local_end_index] = img_k
-            v_cache[:, local_start_index:local_end_index] = img_v
-            
-            updated_kv_cache = {
-                "k": k_cache,
-                "v": v_cache,
-                "global_end_index": torch.tensor([current_end], dtype=torch.long, device=k_cache.device),
-                "local_end_index": torch.tensor([local_end_index], dtype=torch.long, device=k_cache.device)
-            }
-        
-        img_attn_out, _ = self.img_attn_proj(
-            img_attn.view(batch_size, image_seq_len, -1))
-        # Use fused operation for residual connection, normalization, and modulation
-        img_mlp_input, img_residual = self.img_attn_residual_mlp_norm(
-            img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale)
-
-        # Process image MLP
-        img_mlp_out = self.img_mlp(img_mlp_input)
-        img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate)
-
-        return img, updated_kv_cache
-
-    def forward(
-        self,
-        txt_inference=False,
-        vision_inference=False,
-        **kwargs
-    ):
-        if txt_inference:
-            return self.forward_txt(**kwargs)
-        elif vision_inference:
-            return self.forward_vision(**kwargs)
-        else:
-            raise ValueError("txt_inference and vision_inference cannot be both False")
-
-
-class CausalHunyuanVideo15Transformer3DModel(CachableDiT):
-    r"""
-    A Transformer model for video-like data used in [HunyuanVideo1.5](https://huggingface.co/tencent/HunyuanVideo1.5).
-    """
-
-    # shard single stream, double stream blocks, and refiner_blocks
-    _fsdp_shard_conditions = HunyuanVideo15Config()._fsdp_shard_conditions
-    _compile_conditions = HunyuanVideo15Config()._compile_conditions
-    _supported_attention_backends = HunyuanVideo15Config(
-    )._supported_attention_backends
-    param_names_mapping = HunyuanVideo15Config().param_names_mapping
-    reverse_param_names_mapping = HunyuanVideo15Config(
-    ).reverse_param_names_mapping
-    lora_param_names_mapping = HunyuanVideo15Config().lora_param_names_mapping
-
-    def __init__(
-        self,
-        config: HunyuanVideo15Config,
-        hf_config: dict[str, Any],
-    ) -> None:
-        super().__init__(config=config, hf_config=hf_config)
-
-        self.hidden_size = config.hidden_size
-        self.num_attention_heads = config.num_attention_heads
-        self.num_channels_latents = config.num_channels_latents
-        self.out_channels = config.out_channels or config.in_channels
-        self.patch_size = (config.patch_size_t, config.patch_size, config.patch_size)
-
-        # 1. Latent and condition embedders
-        self.img_in = PatchEmbed(self.patch_size, 
-                                    config.in_channels, 
-                                    self.hidden_size,
-                                    prefix=f"{config.prefix}.img_in")
-        self.image_embedder = HunyuanVideo15ImageProjection(config.image_embed_dim, self.hidden_size)
-
-        self.txt_in = SingleTokenRefiner(config.text_embed_dim,
-                                                   self.hidden_size,
-                                                   config.num_attention_heads,
-                                                   depth=config.num_refiner_layers,
-                                                   dtype=None,
-                                                   prefix=f"{config.prefix}.txt_in")
-
-        self.txt_in_2 = HunyuanVideo15ByT5TextProjection(config.text_embed_2_dim, 2048, self.hidden_size)
-
-        self.time_in = HunyuanVideo15TimeEmbedding(self.hidden_size, use_meanflow=config.use_meanflow)
-
-        self.cond_type_embed = nn.Embedding(3, self.hidden_size)
-
-        # 3. Dual stream transformer blocks
-
-        self.double_blocks = nn.ModuleList(
-            [
-                MMDoubleStreamBlock(
-                    hidden_size=self.hidden_size,
-                    num_attention_heads=config.num_attention_heads,
-                    mlp_ratio=config.mlp_ratio,
-                    local_attn_size=config.local_attn_size,
-                    sink_size=config.sink_size,
-                    dtype=None,
-                    supported_attention_backends=self._supported_attention_backends,
-                    prefix=f"{config.prefix}.double_blocks.{i}"
-                )
-                for i in range(config.num_layers)
-            ]
-        )
-
-        # 5. Output projection
-        self.final_layer = FinalLayer(self.hidden_size,
-                                self.patch_size,
-                                self.out_channels,
-                                prefix=f"{config.prefix}.final_layer")
-
-        self.gradient_checkpointing = False
-
-        self.num_frame_per_block = config.num_frames_per_block
-        self.local_attn_size = config.local_attn_size
-        self.block_mask = None
-
-        self.__post_init__()
-
-    def get_text_and_mask(
-        self,
-        encoder_hidden_states: torch.Tensor,
-        encoder_hidden_states_2: torch.Tensor,
-        encoder_attention_mask: torch.Tensor,
-        encoder_attention_mask_2: torch.Tensor,
-        encoder_hidden_states_image: torch.Tensor,
-        timestep: torch.Tensor,
-    ):
-        batch_size, txt_seq_len = encoder_hidden_states.shape[0], encoder_hidden_states.shape[1]
-        # qwen text embedding
-        encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep, encoder_attention_mask)
-
-        encoder_hidden_states_cond_emb = self.cond_type_embed(
-            torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long)
-        )
-        encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb
-
-        # byt5 text embedding
-        encoder_hidden_states_2 = self.txt_in_2(encoder_hidden_states_2)
-
-        encoder_hidden_states_2_cond_emb = self.cond_type_embed(
-            torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long)
-        )
-        encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb
-
-        # image embed
-        encoder_hidden_states_3 = self.image_embedder(encoder_hidden_states_image)
-        is_t2v = torch.all(encoder_hidden_states_image == 0)
-        if is_t2v:
-            encoder_hidden_states_3 = encoder_hidden_states_3 * 0.0
-            encoder_attention_mask_3 = torch.zeros(
-                (batch_size, encoder_hidden_states_3.shape[1]),
-                dtype=encoder_attention_mask.dtype,
-                device=encoder_attention_mask.device,
-            )
-        else:
-            encoder_attention_mask_3 = torch.ones(
-                (batch_size, encoder_hidden_states_3.shape[1]),
-                dtype=encoder_attention_mask.dtype,
-                device=encoder_attention_mask.device,
-            )
-        encoder_hidden_states_3_cond_emb = self.cond_type_embed(
-            2
-            * torch.ones_like(
-                encoder_hidden_states_3[:, :, 0],
-                dtype=torch.long,
-            )
-        )
-        encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb
-
-        # reorder and combine text tokens: combine valid tokens first, then padding
-        encoder_attention_mask = encoder_attention_mask.bool()
-        encoder_attention_mask_2 = encoder_attention_mask_2.bool()
-        encoder_attention_mask_3 = encoder_attention_mask_3.bool()
-        new_encoder_hidden_states = []
-        new_encoder_attention_mask = []
-
-        for text, text_mask, text_2, text_mask_2, image, image_mask in zip(
-            encoder_hidden_states,
-            encoder_attention_mask,
-            encoder_hidden_states_2,
-            encoder_attention_mask_2,
-            encoder_hidden_states_3,
-            encoder_attention_mask_3,
-        ):
-            # Concatenate: [valid_image, valid_byt5, valid_mllm, invalid_image, invalid_byt5, invalid_mllm]
-            new_encoder_hidden_states.append(
-                torch.cat(
-                    [
-                        image[image_mask],  # valid image
-                        text_2[text_mask_2],  # valid byt5
-                        text[text_mask],  # valid mllm
-                        image[~image_mask],  # invalid image (zeroed)
-                        torch.zeros_like(text_2[~text_mask_2]),  # invalid byt5 (zeroed)
-                        torch.zeros_like(text[~text_mask]),  # invalid mllm (zeroed)
-                    ],
-                    dim=0,
-                )
-            )
-            # Apply same reordering to attention masks
-            new_encoder_attention_mask.append(
-                torch.cat(
-                    [
-                        image_mask[image_mask],
-                        text_mask_2[text_mask_2],
-                        text_mask[text_mask],
-                        image_mask[~image_mask],
-                        text_mask_2[~text_mask_2],
-                        text_mask[~text_mask],
-                    ],
-                    dim=0,
-                )
-            )
-
-        encoder_hidden_states = torch.stack(new_encoder_hidden_states)
-        encoder_attention_mask = torch.stack(new_encoder_attention_mask)
-        assert encoder_hidden_states.shape[0] == 1
-        return encoder_hidden_states, encoder_attention_mask
-
-    
-    @staticmethod
-    def _prepare_blockwise_causal_attn_mask(
-        device: torch.device | str, num_frames: int = 21,
-        frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1,
-        text_seq_len: int = 0
-    ) -> BlockMask:
-        """
-        we will divide the token sequence into the following format
-        [1 latent frame] [1 latent frame] ... [1 latent frame]
-        We use flexattention to construct the attention mask
-        """
-        total_length = num_frames * frame_seqlen
-        total_kv_length = total_length + text_seq_len
-
-        total_length_tensor = torch.tensor(total_length, device=device)
-        total_kv_length_tensor = torch.tensor(total_kv_length, device=device)
-
-        # we do right padding to get to a multiple of 128
-        padded_length = math.ceil(total_length / 128) * 128 - total_length
-        kv_padded_length = math.ceil(total_kv_length / 128) * 128 - total_kv_length
-
-        ends = torch.zeros(total_length + padded_length,
-                           device=device, dtype=torch.long)
-
-        # Block-wise causal mask will attend to all elements that are before the end of the current chunk
-        frame_indices = torch.arange(
-            # start=frame_seqlen,
-            start=0,
-            end=total_length,
-            step=frame_seqlen * num_frame_per_block,
-            device=device
-        )
-        # frame_indices = torch.cat([torch.tensor([0], device=device), frame_indices])
-
-        for i, tmp in enumerate(frame_indices):
-            # if i == 0:
-            #     ends[tmp:tmp + frame_seqlen] = tmp + frame_seqlen
-            # else:
-            ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \
-                frame_seqlen * num_frame_per_block
-
-        def attention_mask(b, h, q_idx, kv_idx):
-            if local_attn_size == -1:
-                return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) | ((kv_idx >= total_length_tensor) & (kv_idx < total_kv_length_tensor))
-            else:
-                return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) | ((kv_idx >= total_length_tensor) & (kv_idx < total_kv_length_tensor))
-            # return ((kv_idx < total_length) & (q_idx < total_length))  | (q_idx == kv_idx) # bidirectional mask
-
-        block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length,
-                                       KV_LEN=total_kv_length + kv_padded_length, _compile=False, device=device)
-
-        # if not dist.is_initialized() or dist.get_rank() == 0:
-        #     print(
-        #         f" cache a block wise causal mask with block size of {num_frame_per_block} frames")
-        #     print(block_mask)
-
-        # import imageio
-        # import numpy as np
-        # from torch.nn.attention.flex_attention import create_mask
-
-        # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length +
-        #                    padded_length, KV_LEN=total_length + padded_length, device=device)
-        # import cv2
-        # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024))
-        # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask))
-
-        return block_mask
-
-    def forward_txt(
-        self,
-        encoder_hidden_states: List[torch.Tensor],
-        timestep: torch.LongTensor,
-        encoder_hidden_states_image: Optional[List[torch.Tensor]] = None,
-        encoder_attention_mask: Optional[List[torch.Tensor]] = None,
-        guidance: Optional[torch.Tensor] = None,
-        timestep_r: Optional[torch.LongTensor] = None,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        cache_txt: bool = False,
-        **kwargs
-    ):
-        # Check that the timestep is only consisted of 0s
-        assert torch.all(timestep == 0), "Timestep for txt must be only consisted of 0s"
-
-        if cache_txt:
-            _kv_cache_new = []
-            transformer_num_layers = len(self.double_blocks)
-            for _ in range(transformer_num_layers):
-                _kv_cache_new.append(
-                    {"k_vision": None, "v_vision": None, "k_txt": None, "v_txt": None}
-                )
-
-        encoder_hidden_states_image = encoder_hidden_states_image[0]
-        encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states
-        encoder_attention_mask, encoder_attention_mask_2 = encoder_attention_mask
-
-        # 2. Conditional embeddings
-        if timestep.dim() == 2:
-            ts_seq_len = timestep.shape[1]
-            temb = self.time_in(timestep.flatten(), timestep_r=timestep_r, timestep_seq_len=ts_seq_len)
-        else:
-            temb = self.time_in(timestep, timestep_r=timestep_r)
-        # temb is [bs, seq_len, inner_dim] if ts_seq_len is not None, otherwise [bs, inner_dim]
-
-        encoder_hidden_states, encoder_attention_mask = self.get_text_and_mask(
-            encoder_hidden_states,
-            encoder_hidden_states_2,
-            encoder_attention_mask,
-            encoder_attention_mask_2,
-            encoder_hidden_states_image,
-            timestep
-        )
-        
-        encoder_hidden_states = encoder_hidden_states[encoder_attention_mask.bool().to(encoder_hidden_states.device)].unsqueeze(0)
-
-        # 4. Transformer blocks
-        for index, block in enumerate(self.double_blocks):
-            encoder_hidden_states, t_kv = block(
-                txt_inference=True,
-                vision_inference=False,
-                txt=encoder_hidden_states,
-                vec=temb,
-                cache_txt=cache_txt,
-            )
-
-            if cache_txt:
-                _kv_cache_new[index]["k_txt"] = t_kv["k_txt"]
-                _kv_cache_new[index]["v_txt"] = t_kv["v_txt"]
-
-        if cache_txt:
-            return _kv_cache_new
-
-    def forward_vision(
-        self,
-        hidden_states: torch.Tensor,
-        timestep: torch.LongTensor,
-        guidance: Optional[torch.Tensor] = None,
-        timestep_r: Optional[torch.LongTensor] = None,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        kv_cache: dict | None = None,
-        txt_kv_cache: list | None = None,
-        current_start: int = 0,
-        rope_start_idx: int = 0,
-    ):
-        assert txt_kv_cache is not None, "txt_kv_cache must be provided"
-
-        batch_size, num_channels, num_frames, height, width = hidden_states.shape
-        p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size
-        post_patch_num_frames = num_frames // p_t
-        post_patch_height = height // p_h
-        post_patch_width = width // p_w
-
-        # 1. RoPE
-        # Get rotary embeddings
-        freqs_cos, freqs_sin = get_rotary_pos_embed(
-            (post_patch_num_frames, post_patch_height, post_patch_width), self.hidden_size,
-            self.num_attention_heads, self.config.rope_axes_dim, self.config.rope_theta, start_frame=rope_start_idx)
-        freqs_cos = freqs_cos.to(hidden_states.device)
-        freqs_sin = freqs_sin.to(hidden_states.device)
-        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None
-
-        # 2. Conditional embeddings
-        if timestep.dim() == 2:
-            ts_seq_len = timestep.shape[1]
-            temb = self.time_in(timestep.flatten(), timestep_r=timestep_r, timestep_seq_len=ts_seq_len)
-        else:
-            temb = self.time_in(timestep, timestep_r=timestep_r)
-        # temb is [bs, seq_len, inner_dim] if ts_seq_len is not None, otherwise [bs, inner_dim]
-
-        hidden_states = self.img_in(hidden_states)
-
-        # Prepare block-wise causal attention mask
-        if kv_cache is None:
-            self.block_mask = self._prepare_blockwise_causal_attn_mask(
-                device=hidden_states.device,
-                num_frames=num_frames,
-                frame_seqlen=post_patch_height * post_patch_width,
-                num_frame_per_block=self.num_frame_per_block,
-                local_attn_size=self.local_attn_size,
-                text_seq_len=txt_kv_cache[0]["k_txt"].shape[1]
-            )
-
-        # 4. Transformer blocks
-        if torch.is_grad_enabled() and self.gradient_checkpointing:
-            for block_index, block in enumerate(self.double_blocks):
-                hidden_states, new_cache = self._gradient_checkpointing_func(
-                    block,
-                    txt_inference=False,
-                    vision_inference=True,
-                    img=hidden_states,
-                    vec=temb,
-                    freqs_cis=freqs_cis,
-                    block_mask=self.block_mask,
-                    kv_cache=kv_cache[block_index] if kv_cache is not None else None,
-                    txt_kv_cache=txt_kv_cache[block_index],
-                    current_start=current_start
-                )
-                if new_cache is not None and kv_cache is not None:
-                    for k in new_cache.keys():
-                        kv_cache[block_index][k] = new_cache[k].clone()
-
-        else:
-            for block_index, block in enumerate(self.double_blocks):
-                hidden_states, new_cache = block(
-                    txt_inference=False,
-                    vision_inference=True,
-                    img=hidden_states,
-                    vec=temb,
-                    freqs_cis=freqs_cis,
-                    block_mask=self.block_mask,
-                    kv_cache=kv_cache[block_index] if kv_cache is not None else None,
-                    txt_kv_cache=txt_kv_cache[block_index],
-                    current_start=current_start
-                )
-                if new_cache is not None and kv_cache is not None:
-                    for k in new_cache.keys():
-                        kv_cache[block_index][k] = new_cache[k].clone()
-
-        # Final layer processing
-        hidden_states = self.final_layer(hidden_states, temb)
-        # Unpatchify to get original shape
-        hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, self.patch_size, self.out_channels)
-
-        return hidden_states, kv_cache
-
-    def forward(
-        self,
-        txt_inference=False,
-        vision_inference=False,
-        **kwargs,
-    ):
-        if txt_inference:
-            return self.forward_txt(**kwargs)
-        elif vision_inference:
-            return self.forward_vision(**kwargs)
-        else:
-            raise ValueError("txt_inference and vision_inference cannot be both False")
@@ -33,7 +33,7 @@ from fastvideo.layers.visual_embedding import (PatchEmbed)
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import BaseDiT
 from fastvideo.models.dits.wanvideo import WanT2VCrossAttention, WanTimeTextImageEmbedding
-from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.platforms import AttentionBackendEnum, current_platform

 logger = init_logger(__name__)
 class CausalWanSelfAttention(nn.Module):
@@ -286,6 +286,8 @@ class CausalWanTransformerBlock(nn.Module):
        null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
            hidden_states, attn_output, gate_msa, null_shift, null_scale)
+        norm_hidden_states, hidden_states = norm_hidden_states.to(
+            orig_dtype), hidden_states.to(orig_dtype)

        # 2. Cross-attention
        attn_output = self.attn2(norm_hidden_states,
@@ -452,7 +454,6 @@ class CausalWanTransformer3DModel(BaseDiT):
        This function will be run for num_frame times.
        Process the latent frames one by one (1560 tokens each)
        """
-        from fastvideo.platforms import current_platform

        orig_dtype = hidden_states.dtype
        if not isinstance(encoder_hidden_states, torch.Tensor):
@@ -676,4 +677,7 @@ class CausalWanTransformer3DModel(BaseDiT):
            # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous())
            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
            out.append(u)
-        return out
+        return out
+
+# Entry point for model registry
+EntryClass = CausalWanTransformer3DModel
@@ -723,4 +723,7 @@ class CosmosTransformer3DModel(BaseDiT):
        hidden_states = hidden_states.permute(0, 7, 1, 6, 2, 4, 3, 5)
        hidden_states = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

-        return hidden_states
+        return hidden_states
+
+# Entry point for model registry
+EntryClass = CosmosTransformer3DModel
@@ -959,3 +959,5 @@ class Cosmos25Transformer3DModel(BaseDiT):

        return hidden_states

+# Entry point for model registry
+EntryClass = Cosmos25Transformer3DModel
@@ -940,3 +940,6 @@ class FinalLayer(nn.Module):
        x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        x, _ = self.linear(x)
        return x
+
+# Entry point for model registry
+EntryClass = HunyuanVideoTransformer3DModel
@@ -127,9 +127,8 @@ class HunyuanVideo15TimeEmbedding(nn.Module):
        self,
        timestep: torch.Tensor,
        timestep_r: Optional[torch.Tensor] = None,
-        timestep_seq_len: int | None = None,
    ) -> torch.Tensor:
-        timesteps_emb = self.timestep_embedder(timestep, timestep_seq_len)
+        timesteps_emb = self.timestep_embedder(timestep)

        if timestep_r is not None:
            timesteps_emb_r = self.timestep_embedder_r(timestep_r)
@@ -474,7 +473,6 @@ class HunyuanVideo15Transformer3DModel(CachableDiT):
        guidance: Optional[torch.Tensor] = None,
        timestep_r: Optional[torch.LongTensor] = None,
        attention_kwargs: Optional[Dict[str, Any]] = None,
-        **kwargs
    ):
        encoder_hidden_states_image = encoder_hidden_states_image[0]
        encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states
@@ -496,12 +494,7 @@ class HunyuanVideo15Transformer3DModel(CachableDiT):
        freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None

        # 2. Conditional embeddings
-        if timestep.dim() == 2:
-            ts_seq_len = timestep.shape[1]
-            temb = self.time_in(timestep.flatten(), timestep_r=timestep_r, timestep_seq_len=ts_seq_len)
-        else:
-            temb = self.time_in(timestep, timestep_r=timestep_r)
-        # temb is [bs, seq_len, inner_dim] if ts_seq_len is not None, otherwise [bs, inner_dim]
+        temb = self.time_in(timestep, timestep_r=timestep_r)

        hidden_states = self.img_in(hidden_states)
        hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1)
@@ -705,7 +698,6 @@ class SingleTokenRefiner(nn.Module):
        else:
            mask_float = mask.float().unsqueeze(-1)  # [B, L, 1]
            context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1)
-        context_aware_representations = context_aware_representations.to(original_dtype)

        context_aware_representations = self.c_embedder(
            context_aware_representations)
@@ -858,13 +850,9 @@ class FinalLayer(nn.Module):
    def forward(self, x, c):
        # What the heck HF? Why you change the scale and shift order here???
        scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1)
-        if c.dim() == 3:
-            # [bs, seq_len, inner_dim]
-            num_frames = scale.shape[1]
-            frame_seqlen = x.shape[1] // num_frames
-            x = (self.norm_final(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale.unsqueeze(2)) + shift.unsqueeze(2)).flatten(1, 2)
-        else:
-            # [bs, inner_dim]
-            x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        x, _ = self.linear(x)
-        return x
+        return x
+
+# Entry point for model registry
+EntryClass = HunyuanVideo15Transformer3DModel
@@ -22,3 +22,6 @@ __all__ = [
    # Inference utilities (used by examples)
    "get_resolution_from_image",
 ]
+
+# Entry point for model registry
+EntryClass = HYWorldTransformer3DModel
@@ -536,8 +536,8 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel):
                    temb_txt,
                    freqs_cis,
                    seq_attention_mask,
-                    viewmats_seq, # hyworld
-                    Ks_seq, # hyworld
+                    viewmats_seq,
+                    Ks_seq,
                )
        else:
            for block in self.double_blocks:
@@ -549,8 +549,8 @@ class HYWorldTransformer3DModel(HunyuanVideo15Transformer3DModel):
                    temb_txt,
                    freqs_cis,
                    seq_attention_mask,
-                    viewmats=viewmats_seq, # hyworld
-                    Ks=Ks_seq, # hyworld
+                    viewmats_seq,
+                    Ks_seq,
                )


@@ -1134,3 +1134,5 @@ class LongCatTransformer3DModel(CachableDiT):
        )
        return x

+# Entry point for model registry
+EntryClass = LongCatTransformer3DModel
@@ -18,12 +18,23 @@ import torch.nn as nn
 from einops import rearrange, repeat

 from fastvideo.attention.backends.sdpa import SDPAMetadata
-from fastvideo.attention.layer import LocalAttention
+from fastvideo.attention.layer import DistributedAttention, LocalAttention
 from fastvideo.configs.models.dits import LTX2VideoConfig
-from fastvideo.forward_context import get_forward_context, set_forward_context
+from fastvideo.distributed.communication_op import (
+    sequence_model_parallel_all_gather,
+    sequence_model_parallel_all_gather_with_unpad,
+    sequence_model_parallel_all_to_all_4D,
+    sequence_model_parallel_shard,
+)
+from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
+from fastvideo.distributed.utils import create_attention_mask_for_padding
+from fastvideo.forward_context import ForwardContext, get_forward_context, set_forward_context
+from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import CachableDiT
 from fastvideo.platforms import AttentionBackendEnum

+logger = init_logger(__name__)
+

 def get_timestep_embedding(
    timesteps: torch.Tensor,
@@ -295,6 +306,12 @@ class VideoLatentPatchifier:
        output_shape: VideoLatentShape,
        device: Optional[torch.device] = None,
    ) -> torch.Tensor:
+        """Get patch grid bounds for RoPE computation.
+
+        Args:
+            output_shape: Shape of the video latent tensor
+            device: Device to create tensors on
+        """
        frames = output_shape.frames
        height = output_shape.height
        width = output_shape.width
@@ -366,6 +383,12 @@ class AudioLatentPatchifier:
        output_shape: AudioLatentShape,
        device: Optional[torch.device] = None,
    ) -> torch.Tensor:
+        """Get patch grid bounds for audio RoPE computation.
+
+        Args:
+            output_shape: Shape of the audio latent tensor
+            device: Device to create tensors on
+        """
        start_timings = self._get_audio_latent_time_in_sec(
            self.shift,
            output_shape.frames + self.shift,
@@ -515,6 +538,44 @@ def apply_ltx_rotary_emb(
    raise ValueError(f"Invalid rope type: {rope_type}")


+def apply_ltx_rotary_emb_4d(
+    input_tensor: torch.Tensor,
+    freqs_cis: tuple[torch.Tensor, torch.Tensor],
+    rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+) -> torch.Tensor:
+    """Apply LTX-2 rotary embeddings to a 4D tensor [B, H, T, D].
+    
+    This is used for applying RoPE after all-to-all in distributed attention,
+    where the tensor is already in [B, H, T, D] format.
+    """
+    cos_freqs, sin_freqs = freqs_cis
+    if rope_type == LTXRopeType.INTERLEAVED:
+        # For interleaved, cos/sin have shape [B, T, inner_dim]
+        # Need to reshape to [B, H, T, D] format
+        # Actually interleaved doesn't have per-head rotations, so we broadcast
+        t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2)
+        t1, t2 = t_dup.unbind(dim=-1)
+        t_dup = torch.stack((-t2, t1), dim=-1)
+        input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)")
+        return input_tensor * cos_freqs + input_tensor_rot * sin_freqs
+    if rope_type == LTXRopeType.SPLIT:
+        # For split, cos/sin already have shape [B, H, T, D/2]
+        # input_tensor is [B, H, T, D]
+        split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
+        first_half_input = split_input[..., :1, :]
+        second_half_input = split_input[..., 1:, :]
+
+        output = split_input * cos_freqs.unsqueeze(-2)
+        first_half_output = output[..., :1, :]
+        second_half_output = output[..., 1:, :]
+
+        first_half_output.addcmul_(-sin_freqs.unsqueeze(-2), second_half_input)
+        second_half_output.addcmul_(sin_freqs.unsqueeze(-2), first_half_input)
+
+        return rearrange(output, "... d r -> ... (d r)")
+    raise ValueError(f"Invalid rope type: {rope_type}")
+
+
 def _apply_ltx_interleaved_rotary_emb(
    input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
 ) -> torch.Tensor:
@@ -918,6 +979,205 @@ class TransformerConfig:
    context_dim: int


+class LTXDistributedAttention(DistributedAttention):
+    """LTX-2 specialized DistributedAttention that handles LTX-style RoPE internally."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        rope_type: LTXRopeType,
+        num_kv_heads: int | None = None,
+        softmax_scale: float | None = None,
+        causal: bool = False,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        super().__init__(
+            num_heads=num_heads,
+            head_size=head_size,
+            num_kv_heads=num_kv_heads,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            supported_attention_backends=supported_attention_backends,
+            prefix=prefix,
+            **extra_impl_args,
+        )
+        self.rope_type = rope_type
+
+    @torch.compiler.disable
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        replicated_q: torch.Tensor | None = None,
+        replicated_k: torch.Tensor | None = None,
+        replicated_v: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Forward pass with LTX-2 style RoPE application.
+
+        Args:
+            q: Query tensor [batch_size, seq_len, num_heads, head_dim]
+            k: Key tensor [batch_size, seq_len, num_heads, head_dim]
+            v: Value tensor [batch_size, seq_len, num_heads, head_dim]
+            replicated_q: Replicated query tensor for text tokens
+            replicated_k: Replicated key tensor
+            replicated_v: Replicated value tensor
+            attention_mask: Attention mask [batch_size, seq_len]
+            ltx_freqs_cis: LTX-2 style RoPE (cos, sin) with shape [B, H, T, D]
+
+        Returns:
+            Tuple of output tensor and optional replicated output
+        """
+        assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"
+        batch_size, seq_len, num_heads, head_dim = q.shape
+        local_rank = get_sp_parallel_rank()
+        world_size = get_sp_world_size()
+
+        forward_context: ForwardContext = get_forward_context()
+        ctx_attn_metadata = forward_context.attn_metadata
+
+        # Stack QKV
+        qkv = torch.cat([q, k, v], dim=0)  # [3*batch, seq_len, num_heads, head_dim]
+
+        # Redistribute heads across sequence dimension
+        qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1)
+
+        # After all-to-all, each rank has the full sequence but only a subset of heads
+        valid_seq_len = None
+        if attention_mask is not None:
+            valid_seq_len = (attention_mask[0] == 1).sum().item()
+            qkv = qkv[:, :valid_seq_len, :, :]
+
+        # Apply LTX-2 style RoPE after all-to-all (when we have full sequence)
+        if ltx_freqs_cis is not None:
+            cos, sin = ltx_freqs_cis
+            heads_per_rank = num_heads // world_size
+            head_start = local_rank * heads_per_rank
+            head_end = head_start + heads_per_rank
+            # Slice to this rank's heads: [B, H/SP, T, D]
+            cos_local = cos[:, head_start:head_end, :valid_seq_len, :]
+            sin_local = sin[:, head_start:head_end, :valid_seq_len, :]
+
+            # Apply RoPE to Q and K together (first 2*batch_size in dim 0)
+            qk_part = qkv[:batch_size * 2]
+            # Transpose to [2*B, H/SP, T, D] for LTX RoPE application
+            qk_part = qk_part.transpose(1, 2)
+            qk_part = apply_ltx_rotary_emb_4d(qk_part, (cos_local, sin_local), self.rope_type)
+            # Transpose back to [2*B, T, H/SP, D]
+            qkv[:batch_size * 2] = qk_part.transpose(1, 2)
+
+        # Apply backend-specific preprocess_qkv
+        qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata)
+
+        # Concatenate with replicated QKV if provided
+        if replicated_q is not None:
+            assert replicated_k is not None and replicated_v is not None
+            replicated_qkv = torch.cat(
+                [replicated_q, replicated_k, replicated_v],
+                dim=0)  # [3, seq_len, num_heads, head_dim]
+            heads_per_rank = num_heads // world_size
+            replicated_qkv = replicated_qkv[:, :, local_rank *
+                                            heads_per_rank:(local_rank + 1) *
+                                            heads_per_rank]
+            qkv = torch.cat([qkv, replicated_qkv], dim=1)
+
+        q, k, v = qkv.chunk(3, dim=0)
+
+        output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
+
+        # Redistribute back if using sequence parallelism
+        replicated_output = None
+        if replicated_q is not None:
+            split_idx = seq_len * world_size if valid_seq_len is None else valid_seq_len
+            replicated_output = output[:, split_idx:]
+            output = output[:, :split_idx]
+            replicated_output = sequence_model_parallel_all_gather(
+                replicated_output.contiguous(), dim=2)
+
+        # Apply backend-specific postprocess_output
+        output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
+
+        if attention_mask is not None:
+            pad_len = (attention_mask[0] == 0).sum().item()
+            output = torch.nn.functional.pad(output, (0, 0, 0, 0, 0, pad_len))
+
+        output = sequence_model_parallel_all_to_all_4D(output, scatter_dim=1, gather_dim=2)
+
+        return output, replicated_output
+
+
+class LTXLocalAttention(LocalAttention):
+    """LTX-2 specialized LocalAttention that handles LTX-style RoPE internally."""
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        rope_type: LTXRopeType,
+        num_kv_heads: int | None = None,
+        softmax_scale: float | None = None,
+        causal: bool = False,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None,
+        **extra_impl_args,
+    ) -> None:
+        super().__init__(
+            num_heads=num_heads,
+            head_size=head_size,
+            num_kv_heads=num_kv_heads,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            supported_attention_backends=supported_attention_backends,
+            **extra_impl_args,
+        )
+        self.rope_type = rope_type
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        ltx_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
+        ltx_k_freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        """Forward pass with LTX-2 style RoPE application.
+
+        Args:
+            q: Query tensor [batch_size, seq_len, num_heads, head_dim]
+            k: Key tensor [batch_size, seq_len, num_heads, head_dim]
+            v: Value tensor [batch_size, seq_len, num_heads, head_dim]
+            ltx_freqs_cis: LTX-2 style RoPE (cos, sin) for Q with shape [B, H, T, D]
+            ltx_k_freqs_cis: LTX-2 style RoPE (cos, sin) for K (if different from Q)
+
+        Returns:
+            Output tensor after local attention
+        """
+        assert q.dim() == 4 and k.dim() == 4 and v.dim() == 4, "Expected 4D tensors"
+
+        forward_context: ForwardContext = get_forward_context()
+        ctx_attn_metadata = forward_context.attn_metadata
+
+        # Apply LTX-2 style RoPE
+        if ltx_freqs_cis is not None:
+            # Use separate K RoPE if provided (for cross-attention), otherwise use same as Q
+            k_freqs = ltx_k_freqs_cis if ltx_k_freqs_cis is not None else ltx_freqs_cis
+            # Transpose to [B, H, T, D] for RoPE application
+            q = q.transpose(1, 2)
+            k = k.transpose(1, 2)
+            q = apply_ltx_rotary_emb_4d(q, ltx_freqs_cis, self.rope_type)
+            k = apply_ltx_rotary_emb_4d(k, k_freqs, self.rope_type)
+            # Transpose back to [B, T, H, D]
+            q = q.transpose(1, 2)
+            k = k.transpose(1, 2)
+
+        output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
+        return output
+
+
 class LTXSelfAttention(nn.Module):
    """LTX-2 attention block with RMSNorm + FastVideo LocalAttention."""
    def __init__(
@@ -945,17 +1205,19 @@ class LTXSelfAttention(nn.Module):
        self.to_v = nn.Linear(context_dim, inner_dim, bias=True)
        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity())

-        self.attn = LocalAttention(
+        self.attn = LTXLocalAttention(
            num_heads=heads,
            head_size=dim_head,
+            rope_type=rope_type,
            dropout_rate=0,
            softmax_scale=None,
            causal=False,
            supported_attention_backends=supported_attention_backends,
        )
-        self.attn_masked = LocalAttention(
+        self.attn_masked = LTXLocalAttention(
            num_heads=heads,
            head_size=dim_head,
+            rope_type=rope_type,
            dropout_rate=0,
            softmax_scale=None,
            causal=False,
@@ -978,15 +1240,13 @@ class LTXSelfAttention(nn.Module):
        q = self.q_norm(q)
        k = self.k_norm(k)

-        if pe is not None:
-            q = apply_ltx_rotary_emb(q, pe, self.rope_type)
-            k = apply_ltx_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
-
+        # RoPE is applied inside LTXLocalAttention
        b, q_len, _ = q.shape
        k_len = k.shape[1]
        q = q.view(b, q_len, self.heads, self.dim_head)
        k = k.view(b, k_len, self.heads, self.dim_head)
        v = v.view(b, k_len, self.heads, self.dim_head)
+
        if mask is not None:
            if mask.ndim == 2:
                mask = mask.unsqueeze(0)
@@ -1008,9 +1268,93 @@ class LTXSelfAttention(nn.Module):
                attn_metadata=attn_metadata,
                forward_batch=forward_batch,
            ):
-                out = self.attn_masked(q, k, v)
+                out = self.attn_masked(q, k, v, ltx_freqs_cis=pe, ltx_k_freqs_cis=k_pe)
        else:
-            out = self.attn(q, k, v)
+            out = self.attn(q, k, v, ltx_freqs_cis=pe, ltx_k_freqs_cis=k_pe)
+        out = out.reshape(b, q_len, -1)
+        return self.to_out(out)
+
+
+class LTXDistributedSelfAttention(nn.Module):
+    """LTX-2 attention block with RMSNorm + LTXDistributedAttention for SP."""
+
+    def __init__(
+        self,
+        query_dim: int,
+        context_dim: int | None,
+        heads: int,
+        dim_head: int,
+        norm_eps: float,
+        rope_type: LTXRopeType,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...],
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        inner_dim = dim_head * heads
+        context_dim = query_dim if context_dim is None else context_dim
+
+        self.heads = heads
+        self.dim_head = dim_head
+        self.rope_type = rope_type
+
+        self.q_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
+        self.k_norm = torch.nn.RMSNorm(inner_dim, eps=norm_eps)
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=True)
+        self.to_k = nn.Linear(context_dim, inner_dim, bias=True)
+        self.to_v = nn.Linear(context_dim, inner_dim, bias=True)
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim, bias=True), nn.Identity())
+
+        self.attn = LTXDistributedAttention(
+            num_heads=heads,
+            head_size=dim_head,
+            rope_type=rope_type,
+            causal=False,
+            supported_attention_backends=supported_attention_backends,
+            prefix=f"{prefix}.attn",
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor | None = None,
+        pe: tuple[torch.Tensor, torch.Tensor] | None = None,
+        k_pe: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Forward pass for distributed self-attention.
+
+        Args:
+            x: Query tensor [B, L, C]
+            context: Key/Value tensor [B, L, C], or None for self-attention
+            pe: Rotary position embeddings for Q (cos, sin) with shape [B, H, T_full, D]
+                NOTE: For SP, this must be the FULL sequence RoPE, not sharded!
+            k_pe: Rotary position embeddings for K (cos, sin), or None to use pe
+            attention_mask: Attention mask for padding [B, padded_seq_len]
+        """
+        q = self.to_q(x)
+        context = x if context is None else context
+        k = self.to_k(context)
+        v = self.to_v(context)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # RoPE is applied inside LTXDistributedAttention AFTER the all-to-all,
+        # when each rank has the full sequence (but subset of heads).
+        b, q_len, _ = q.shape
+        k_len = k.shape[1]
+        q = q.view(b, q_len, self.heads, self.dim_head)
+        k = k.view(b, k_len, self.heads, self.dim_head)
+        v = v.view(b, k_len, self.heads, self.dim_head)
+
+        # Pass full RoPE to distributed attention - it will apply after all-to-all
+        # and slice to this rank's heads
+        out, _ = self.attn(
+            q, k, v,
+            attention_mask=attention_mask,
+            ltx_freqs_cis=pe,
+        )
+
        out = out.reshape(b, q_len, -1)
        return self.to_out(out)

@@ -1025,12 +1369,31 @@ class BasicAVTransformerBlock(torch.nn.Module):
        audio: TransformerConfig | None = None,
        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
        norm_eps: float = 1e-6,
+        use_distributed_attention: bool = False,
+        prefix: str = "",
    ):
        super().__init__()
        self.idx = idx
+        self.use_distributed_attention = use_distributed_attention
+
+        # Choose attention class based on SP mode
+        # Self-attention and audio-video cross-attention use DistributedAttention when SP > 1
+        # Text cross-attention uses LocalAttention (text embeddings are replicated)
+        SelfAttnCls = LTXDistributedSelfAttention if use_distributed_attention else LTXSelfAttention
+        CrossAttnCls = LTXSelfAttention  # Text cross-attention is always local

        if video is not None:
-            self.attn1 = LTXSelfAttention(
+            # Video self-attention - use distributed when SP > 1
+            self.attn1 = SelfAttnCls(
+                query_dim=video.dim,
+                context_dim=None,
+                heads=video.heads,
+                dim_head=video.d_head,
+                norm_eps=norm_eps,
+                rope_type=rope_type,
+                supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA),
+                prefix=f"{prefix}.blocks.{idx}.attn1" if use_distributed_attention else "",
+            ) if use_distributed_attention else LTXSelfAttention(
                query_dim=video.dim,
                context_dim=None,
                heads=video.heads,
@@ -1039,7 +1402,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                rope_type=rope_type,
                supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA),
            )
-            self.attn2 = LTXSelfAttention(
+            # Text cross-attention - always local (text is replicated)
+            self.attn2 = CrossAttnCls(
                query_dim=video.dim,
                context_dim=video.context_dim,
                heads=video.heads,
@@ -1052,7 +1416,17 @@ class BasicAVTransformerBlock(torch.nn.Module):
            self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim))

        if audio is not None:
-            self.audio_attn1 = LTXSelfAttention(
+            # Audio self-attention - use distributed when SP > 1
+            self.audio_attn1 = SelfAttnCls(
+                query_dim=audio.dim,
+                context_dim=None,
+                heads=audio.heads,
+                dim_head=audio.d_head,
+                norm_eps=norm_eps,
+                rope_type=rope_type,
+                supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA),
+                prefix=f"{prefix}.blocks.{idx}.audio_attn1" if use_distributed_attention else "",
+            ) if use_distributed_attention else LTXSelfAttention(
                query_dim=audio.dim,
                context_dim=None,
                heads=audio.heads,
@@ -1061,7 +1435,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                rope_type=rope_type,
                supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA),
            )
-            self.audio_attn2 = LTXSelfAttention(
+            # Text cross-attention - always local (text is replicated)
+            self.audio_attn2 = CrossAttnCls(
                query_dim=audio.dim,
                context_dim=audio.context_dim,
                heads=audio.heads,
@@ -1074,6 +1449,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
            self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim))

        if audio is not None and video is not None:
+            # Audio-to-video cross-attention
+            # Uses local attention - context is gathered from all SP ranks in forward()
            self.audio_to_video_attn = LTXSelfAttention(
                query_dim=video.dim,
                context_dim=audio.dim,
@@ -1083,6 +1460,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                rope_type=rope_type,
                supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA),
            )
+            # Video-to-audio cross-attention
+            # Uses local attention - context is gathered from all SP ranks in forward()
            self.video_to_audio_attn = LTXSelfAttention(
                query_dim=audio.dim,
                context_dim=video.dim,
@@ -1131,7 +1510,17 @@ class BasicAVTransformerBlock(torch.nn.Module):
        self,
        video: TransformerArgs | None,
        audio: TransformerArgs | None,
+        video_attention_mask: torch.Tensor | None = None,
+        audio_attention_mask: torch.Tensor | None = None,
    ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
+        """Forward pass for transformer block.
+
+        Args:
+            video: Video transformer args
+            audio: Audio transformer args
+            video_attention_mask: SP padding attention mask for video [B, padded_seq_len]
+            audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len]
+        """
        vx = video.x if video is not None else None
        ax = audio.x if audio is not None else None

@@ -1145,7 +1534,12 @@ class BasicAVTransformerBlock(torch.nn.Module):
                self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
            )
            norm_vx = torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
-            vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa
+            # Self-attention: pass SP attention mask for distributed attention
+            if self.use_distributed_attention:
+                vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings, attention_mask=video_attention_mask) * vgate_msa
+            else:
+                vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa
+            # Text cross-attention: no SP mask needed (text is replicated, uses local attention)
            vx = vx + self.attn2(
                torch.nn.functional.rms_norm(vx, (vx.shape[-1],), eps=self.norm_eps),
                context=video.context,
@@ -1157,7 +1551,12 @@ class BasicAVTransformerBlock(torch.nn.Module):
                self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
            )
            norm_ax = torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
-            ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa
+            # Self-attention: pass SP attention mask for distributed attention
+            if self.use_distributed_attention:
+                ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings, attention_mask=audio_attention_mask) * agate_msa
+            else:
+                ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa
+            # Text cross-attention: no SP mask needed (text is replicated, uses local attention)
            ax = ax + self.audio_attn2(
                torch.nn.functional.rms_norm(ax, (ax.shape[-1],), eps=self.norm_eps),
                context=audio.context,
@@ -1197,28 +1596,91 @@ class BasicAVTransformerBlock(torch.nn.Module):
            if run_a2v:
                vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v
                ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
-                vx = vx + (
-                    self.audio_to_video_attn(
-                        vx_scaled,
-                        context=ax_scaled,
-                        pe=video.cross_positional_embeddings,
-                        k_pe=audio.cross_positional_embeddings,
+                # Audio-to-video cross-attention: Q from video (sharded), K/V from audio
+                # For SP, gather audio context from all ranks before cross-attention
+                if self.use_distributed_attention:
+                    # Gather audio context from all SP ranks
+                    ax_context = sequence_model_parallel_all_gather(ax_scaled, dim=1)
+                    # Video Q: video tokens divide evenly (adjusted upstream), so simple slicing works
+                    video_q_pe = None
+                    if video.cross_positional_embeddings is not None:
+                        sp_rank = get_sp_parallel_rank()
+                        local_seq_len = vx_scaled.shape[1]
+                        pe_tuple = video.cross_positional_embeddings
+                        start_idx = sp_rank * local_seq_len
+                        end_idx = start_idx + local_seq_len
+                        video_q_pe = tuple(pe[:, :, start_idx:end_idx, :] for pe in pe_tuple)
+                    # Audio K: gathered context may have padding, trim to original length
+                    audio_k_pe = audio.cross_positional_embeddings
+                    if audio_k_pe is not None:
+                        original_audio_len = audio_k_pe[0].shape[2]
+                        ax_context = ax_context[:, :original_audio_len, :]
+                    vx = vx + (
+                        self.audio_to_video_attn(
+                            vx_scaled,
+                            context=ax_context,
+                            pe=video_q_pe,
+                            k_pe=audio_k_pe,
+                        )
+                        * gate_out_a2v
+                    )
+                else:
+                    ax_context = ax_scaled
+                    video_q_pe = video.cross_positional_embeddings
+                    audio_k_pe = audio.cross_positional_embeddings
+                    vx = vx + (
+                        self.audio_to_video_attn(
+                            vx_scaled,
+                            context=ax_context,
+                            pe=video_q_pe,
+                            k_pe=audio_k_pe,
+                        )
+                        * gate_out_a2v
                    )
-                    * gate_out_a2v
-                )

            if run_v2a:
                ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
                vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
-                ax = ax + (
-                    self.video_to_audio_attn(
-                        ax_scaled,
-                        context=vx_scaled,
-                        pe=audio.cross_positional_embeddings,
-                        k_pe=video.cross_positional_embeddings,
+                # Video-to-audio cross-attention: Q from audio, K/V from video
+                # For SP, gather both modalities and run cross-attention on full sequences
+                # (audio is short, so no need to keep it sharded for cross-attention)
+                if self.use_distributed_attention:
+                    # Gather both modalities to full sequences
+                    ax_full = sequence_model_parallel_all_gather(ax_scaled, dim=1)
+                    vx_context = sequence_model_parallel_all_gather(vx_scaled, dim=1)
+                    # Use full positional embeddings directly
+                    audio_q_pe = audio.cross_positional_embeddings
+                    video_k_pe = video.cross_positional_embeddings
+                    # Trim gathered tensors to original lengths (may have padding from sharding)
+                    if audio_q_pe is not None:
+                        original_audio_len = audio_q_pe[0].shape[2]
+                        ax_full = ax_full[:, :original_audio_len, :]
+                    if video_k_pe is not None:
+                        original_video_len = video_k_pe[0].shape[2]
+                        vx_context = vx_context[:, :original_video_len, :]
+                    # Run cross-attention on full audio
+                    v2a_out_full = self.video_to_audio_attn(
+                        ax_full,
+                        context=vx_context,
+                        pe=audio_q_pe,
+                        k_pe=video_k_pe,
+                    )
+                    # Shard the output back to match ax's sharding
+                    v2a_out, _ = sequence_model_parallel_shard(v2a_out_full, dim=1)
+                    ax = ax + v2a_out * gate_out_v2a
+                else:
+                    vx_context = vx_scaled
+                    audio_q_pe = audio.cross_positional_embeddings
+                    video_k_pe = video.cross_positional_embeddings
+                    ax = ax + (
+                        self.video_to_audio_attn(
+                            ax_scaled,
+                            context=vx_context,
+                            pe=audio_q_pe,
+                            k_pe=video_k_pe,
+                        )
+                        * gate_out_v2a
                    )
-                    * gate_out_v2a
-                )

        if run_vx:
            vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
@@ -1289,6 +1751,8 @@ class LTXModel(torch.nn.Module):
        av_ca_timestep_scale_multiplier: int = 1,
        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
        double_precision_rope: bool = False,
+        use_distributed_attention: bool = False,
+        prefix: str = "",
    ):
        super().__init__()
        self._enable_gradient_checkpointing = False
@@ -1340,6 +1804,8 @@ class LTXModel(torch.nn.Module):
            audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0,
            audio_cross_attention_dim=audio_cross_attention_dim,
            norm_eps=norm_eps,
+            use_distributed_attention=use_distributed_attention,
+            prefix=prefix,
        )

    def _init_video(
@@ -1469,6 +1935,8 @@ class LTXModel(torch.nn.Module):
        audio_attention_head_dim: int,
        audio_cross_attention_dim: int,
        norm_eps: float,
+        use_distributed_attention: bool = False,
+        prefix: str = "",
    ) -> None:
        video_config = (
            TransformerConfig(
@@ -1490,6 +1958,7 @@ class LTXModel(torch.nn.Module):
            if self.model_type.is_audio_enabled()
            else None
        )
+        self.use_distributed_attention = use_distributed_attention
        self.transformer_blocks = torch.nn.ModuleList(
            [
                BasicAVTransformerBlock(
@@ -1498,6 +1967,8 @@ class LTXModel(torch.nn.Module):
                    audio=audio_config,
                    rope_type=self.rope_type,
                    norm_eps=norm_eps,
+                    use_distributed_attention=use_distributed_attention,
+                    prefix=prefix,
                )
                for idx in range(num_layers)
            ]
@@ -1507,9 +1978,16 @@ class LTXModel(torch.nn.Module):
        self,
        video: TransformerArgs | None,
        audio: TransformerArgs | None,
+        video_attention_mask: torch.Tensor | None = None,
+        audio_attention_mask: torch.Tensor | None = None,
    ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
        for block in self.transformer_blocks:
-            video, audio = block(video=video, audio=audio)
+            video, audio = block(
+                video=video,
+                audio=audio,
+                video_attention_mask=video_attention_mask,
+                audio_attention_mask=audio_attention_mask,
+            )
        return video, audio

    def _process_output(
@@ -1533,7 +2011,17 @@ class LTXModel(torch.nn.Module):
        self,
        video: Modality | None,
        audio: Modality | None,
+        video_attention_mask: torch.Tensor | None = None,
+        audio_attention_mask: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        """Forward pass through the LTX model.
+
+        Args:
+            video: Video modality input
+            audio: Audio modality input
+            video_attention_mask: SP padding attention mask for video [B, padded_seq_len]
+            audio_attention_mask: SP padding attention mask for audio [B, padded_seq_len]
+        """
        if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
            _debug_block_log_line(
                "fastvideo:patchify_proj"
@@ -1551,7 +2039,12 @@ class LTXModel(torch.nn.Module):
        audio_args = self.audio_args_preprocessor.prepare(audio) if audio is not None else None
        _debug_transformer_args("fastvideo:prep_video", video_args)
        _debug_transformer_args("fastvideo:prep_audio", audio_args)
-        video_out, audio_out = self._process_transformer_blocks(video_args, audio_args)
+        video_out, audio_out = self._process_transformer_blocks(
+            video_args,
+            audio_args,
+            video_attention_mask=video_attention_mask,
+            audio_attention_mask=audio_attention_mask,
+        )

        vx = (
            self._process_output(
@@ -1588,6 +2081,24 @@ class LTX2Transformer3DModel(CachableDiT):

        arch = config.arch_config

+        # Get SP world size for distributed attention
+        sp_world_size = get_sp_world_size()
+        use_distributed_attention = sp_world_size > 1
+
+        # Validate that attention heads are divisible by SP world size
+        if sp_world_size > 1:
+            assert arch.num_attention_heads % sp_world_size == 0, (
+                f"The number of video attention heads ({arch.num_attention_heads}) "
+                f"must be divisible by the sequence parallel size ({sp_world_size})"
+            )
+            assert arch.audio_num_attention_heads % sp_world_size == 0, (
+                f"The number of audio attention heads ({arch.audio_num_attention_heads}) "
+                f"must be divisible by the sequence parallel size ({sp_world_size})"
+            )
+            logger.info(
+                f"LTX2 sequence parallelism enabled with SP world size {sp_world_size}"
+            )
+
        model_type = LTXModelType.AudioVideo
        self.model = LTXModel(
            model_type=model_type,
@@ -1612,6 +2123,8 @@ class LTX2Transformer3DModel(CachableDiT):
            audio_cross_attention_dim=arch.audio_cross_attention_dim,
            audio_positional_embedding_max_pos=arch.audio_positional_embedding_max_pos,
            av_ca_timestep_scale_multiplier=arch.av_ca_timestep_scale_multiplier,
+            use_distributed_attention=use_distributed_attention,
+            prefix=config.prefix,
        )

        self.patchifier = VideoLatentPatchifier(patch_size=arch.patch_size[1])
@@ -1627,6 +2140,7 @@ class LTX2Transformer3DModel(CachableDiT):
        self.hidden_size = arch.num_attention_heads * arch.attention_head_dim
        self.num_attention_heads = arch.num_attention_heads
        self.num_channels_latents = arch.num_channels_latents
+        self._logged_attention_mask = False

        if os.getenv("LTX2_DEBUG_DETAIL", "0") == "1":
            detail_path = os.getenv("LTX2_PIPELINE_DEBUG_DETAIL_PATH", "")
@@ -1702,10 +2216,12 @@ class LTX2Transformer3DModel(CachableDiT):
        if isinstance(encoder_hidden_states, list):
            encoder_hidden_states = encoder_hidden_states[0]

-        video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape)
-        positions = self.patchifier.get_patch_grid_bounds(
-            video_shape, device=hidden_states.device
-        )
+        # Get SP parameters
+        sp_world_size = get_sp_world_size()
+        sp_rank = get_sp_parallel_rank()
+        batch_size = hidden_states.shape[0]
+
+        # Get fps for position computation
        fps = None
        try:
            forward_ctx = get_forward_context()
@@ -1717,17 +2233,60 @@ class LTX2Transformer3DModel(CachableDiT):
                fps_value = fps_value[0] if fps_value else None
            if fps_value is not None:
                fps = float(fps_value)
+
+        # Patchify video latents
+        video_shape = VideoLatentShape.from_torch_shape(hidden_states.shape)
+        latents = self.patchifier.patchify(hidden_states)
+
+        # Shard video latents and timestep across SP ranks
+        video_original_seq_len = latents.shape[1]
+        video_attention_mask = None
+        video_timestep = timestep
+        if sp_world_size > 1:
+            latents, video_original_seq_len = sequence_model_parallel_shard(latents, dim=1)
+            # Shard timestep along sequence dimension (timestep has shape [batch, seq_len])
+            video_timestep, _ = sequence_model_parallel_shard(timestep, dim=1)
+            # Create attention mask for padded tokens
+            current_seq_len = latents.shape[1]
+            padded_seq_len = current_seq_len * sp_world_size
+            if padded_seq_len > video_original_seq_len:
+                if not self._logged_attention_mask:
+                    logger.info(
+                        f"Video padding applied, original seq len: {video_original_seq_len}, "
+                        f"padded seq len: {padded_seq_len}"
+                    )
+                    self._logged_attention_mask = True
+                video_attention_mask = create_attention_mask_for_padding(
+                    seq_len=video_original_seq_len,
+                    padded_seq_len=padded_seq_len,
+                    batch_size=batch_size,
+                    device=hidden_states.device,
+                )
+
+        # Compute RoPE positions for the FULL sequence (before sharding)
+        # This is necessary because sequence sharding may not align to frame boundaries.
+        # The full RoPE will be passed to DistributedAttention which applies it after
+        # the all-to-all (when each rank has the full sequence but subset of heads).
+        positions = self.patchifier.get_patch_grid_bounds(
+            video_shape, device=hidden_states.device
+        )
        positions = _get_pixel_coords(
            positions,
            DEFAULT_LTX2_SCALE_FACTORS,
            fps=fps,
            causal_fix=True,
        ).to(hidden_states.dtype)
-        latents = self.patchifier.patchify(hidden_states)
+
+        # Pad positions to match padded sequence length for SP cross-attention
+        if sp_world_size > 1 and padded_seq_len > video_original_seq_len:
+            padding_needed = padded_seq_len - video_original_seq_len
+            last_pos = positions[:, :, -1:, :].expand(-1, -1, padding_needed, -1)
+            positions = torch.cat([positions, last_pos], dim=2)
+
        video_modality = Modality(
            enabled=True,
            latent=latents,
-            timesteps=timestep,
+            timesteps=video_timestep,
            positions=positions,
            context=encoder_hidden_states,
            context_mask=encoder_attention_mask,
@@ -1745,19 +2304,45 @@ class LTX2Transformer3DModel(CachableDiT):
                f"latent_head={video_head} "
                f"latent_checksum={video_checksum:.6f}"
            )
+
+        # Process audio modality if provided
        audio_modality = None
        audio_shape = None
+        audio_original_seq_len = 0
+        audio_attention_mask = None
+
        if audio_hidden_states is not None and audio_encoder_hidden_states is not None and audio_timestep is not None:
-            audio_shape = AudioLatentShape.from_torch_shape(
-                audio_hidden_states.shape)
+            audio_shape = AudioLatentShape.from_torch_shape(audio_hidden_states.shape)
+            audio_latents = self.audio_patchifier.patchify(audio_hidden_states)
+
+            # Shard audio latents and timestep across SP ranks
+            audio_original_seq_len = audio_latents.shape[1]
+            sharded_audio_timestep = audio_timestep
+            if sp_world_size > 1:
+                audio_latents, audio_original_seq_len = sequence_model_parallel_shard(audio_latents, dim=1)
+                # Shard audio timestep along sequence dimension
+                sharded_audio_timestep, _ = sequence_model_parallel_shard(audio_timestep, dim=1)
+                # Create attention mask for padded tokens
+                audio_current_seq_len = audio_latents.shape[1]
+                audio_padded_seq_len = audio_current_seq_len * sp_world_size
+                if audio_padded_seq_len > audio_original_seq_len:
+                    audio_attention_mask = create_attention_mask_for_padding(
+                        seq_len=audio_original_seq_len,
+                        padded_seq_len=audio_padded_seq_len,
+                        batch_size=batch_size,
+                        device=audio_hidden_states.device,
+                    )
+
+            # Compute audio RoPE positions for the FULL sequence (before sharding)
+            # Same as video: full RoPE is applied after all-to-all in DistributedAttention
            audio_positions = self.audio_patchifier.get_patch_grid_bounds(
-                audio_shape, device=audio_hidden_states.device)
-            audio_latents = self.audio_patchifier.patchify(
-                audio_hidden_states)
+                audio_shape, device=audio_hidden_states.device
+            )
+
            audio_modality = Modality(
                enabled=True,
                latent=audio_latents,
-                timesteps=audio_timestep,
+                timesteps=sharded_audio_timestep,
                positions=audio_positions,
                context=audio_encoder_hidden_states,
                context_mask=audio_encoder_attention_mask,
@@ -1775,10 +2360,16 @@ class LTX2Transformer3DModel(CachableDiT):
                    f"latent_head={audio_head} "
                    f"latent_checksum={audio_checksum:.6f}"
                )
+
+        # Run transformer with attention masks
        video_out, audio_out = self.model(
            video=video_modality,
            audio=audio_modality,
+            video_attention_mask=video_attention_mask,
+            audio_attention_mask=audio_attention_mask,
        )
+
+        # Denoised prediction
        if video_out is not None and video_modality is not None:
            video_out = _to_denoised(
                video_modality.latent,
@@ -1791,6 +2382,20 @@ class LTX2Transformer3DModel(CachableDiT):
                audio_out,
                audio_modality.timesteps,
            )
+
+        # Gather and unpad video output
+        if sp_world_size > 1 and video_out is not None:
+            video_out = sequence_model_parallel_all_gather_with_unpad(
+                video_out, video_original_seq_len, dim=1
+            )
+
+        # Gather and unpad audio output
+        if sp_world_size > 1 and audio_out is not None:
+            audio_out = sequence_model_parallel_all_gather_with_unpad(
+                audio_out, audio_original_seq_len, dim=1
+            )
+
+        # Unpatchify
        video_out = self.patchifier.unpatchify(
            video_out, output_shape=video_shape)
        if audio_out is None or audio_shape is None:
@@ -1798,3 +2403,6 @@ class LTX2Transformer3DModel(CachableDiT):
        audio_out = self.audio_patchifier.unpatchify(
            audio_out, output_shape=audio_shape)
        return video_out, audio_out
+
+# Entry point for model registry
+EntryClass = LTX2Transformer3DModel
@@ -1,567 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from Matrix-Game: https://github.com/SkyworkAI/Matrix-Game/blob/main/Matrix-Game-2/wan/modules/action_module.py
-
-from einops import rearrange
-import torch
-import torch.nn as nn
-import math
-from torch.nn.attention.flex_attention import flex_attention
-
-from fastvideo.attention import LocalAttention
-from fastvideo.layers.linear import ReplicatedLinear
-from fastvideo.layers.layernorm import FP32LayerNorm, RMSNorm
-from fastvideo.layers.rotary_embedding import (
-    get_nd_rotary_pos_embed as _fv_get_nd_rotary_pos_embed,
-    _apply_rotary_emb,
-)
-from fastvideo.platforms import AttentionBackendEnum
-
-
-DISABLE_COMPILE = False
-flex_attention = torch.compile(
-    flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs")
-
-
-def _get_nd_rotary_pos_embed_matrixgame(
-    rope_dim_list,
-    rope_sizes,
-    theta: float = 10000.0,
-    theta_rescale_factor: float = 1.0,
-):
-    cos, sin = _fv_get_nd_rotary_pos_embed(
-        rope_dim_list,
-        rope_sizes,
-        theta=theta,
-        theta_rescale_factor=theta_rescale_factor,
-        dtype=torch.float32,
-    )
-    # convert from [S, D/2] to [S, D] format
-    cos = cos.repeat_interleave(2, dim=1)
-    sin = sin.repeat_interleave(2, dim=1)
-    return cos, sin
-
-
-def _apply_rotary_emb_qk(
-    xq: torch.Tensor,
-    xk: torch.Tensor,
-    freqs_cos: torch.Tensor,
-    freqs_sin: torch.Tensor,
-    start_offset: int = 0,
-) -> tuple[torch.Tensor, torch.Tensor]:
-    seq_len = xq.shape[1]
-    
-    # Slice frequencies based on offset
-    cos = freqs_cos[start_offset:start_offset + seq_len]  # [S, D]
-    sin = freqs_sin[start_offset:start_offset + seq_len]  # [S, D]
-    
-    # Move to device
-    cos = cos.to(xq.device)
-    sin = sin.to(xq.device)
-    
-    # Convert from [S, D] (interleaved) back to [S, D/2]
-    cos_half = cos[:, ::2]  # [S, D/2]
-    sin_half = sin[:, ::2]  # [S, D/2]
-
-    # xq/xk are [B, S, H, D], need to reshape for each batch
-    B, S, H, D = xq.shape
-
-    xq_out = _apply_rotary_emb(xq, cos_half, sin_half, is_neox_style=False)
-    xk_out = _apply_rotary_emb(xk, cos_half, sin_half, is_neox_style=False)
-    
-    return xq_out, xk_out
-
-class ActionModule(nn.Module):
-    """
-    action module from https://arxiv.org/pdf/2501.08325
-    """
-
-    def __init__(
-        self, 
-        mouse_dim_in: int = 2,
-        keyboard_dim_in: int = 6,
-        hidden_size: int = 128,
-        img_hidden_size: int = 1536,
-        keyboard_hidden_dim: int = 1024,
-        mouse_hidden_dim: int = 1024,
-        vae_time_compression_ratio: int = 4, 
-        windows_size: int = 3,
-        heads_num: int = 16,
-        patch_size: list | None = None,
-        qk_norm: bool = True,
-        qkv_bias: bool = False,
-        rope_dim_list: list | None = None,
-        rope_theta = 256,
-        mouse_qk_dim_list: list | None = None,
-        enable_mouse = True,
-        enable_keyboard = True,
-        local_attn_size = 6,
-        blocks: list | None = None,
-    ):
-        super().__init__()
-        # Initialize mutable defaults
-        patch_size = patch_size if patch_size is not None else [1, 2, 2]
-        rope_dim_list = rope_dim_list if rope_dim_list is not None else [8, 28, 28]
-        mouse_qk_dim_list = mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28]
-        blocks = blocks if blocks is not None else []
-        self.local_attn_size = local_attn_size
-        self.enable_mouse = enable_mouse
-        self.enable_keyboard = enable_keyboard
-
-        self.rope_dim_list = rope_dim_list
-        self.rope_theta = rope_theta
-        if self.enable_keyboard:
-            self.keyboard_embed = nn.Sequential(
-                nn.Linear(keyboard_dim_in, hidden_size, bias=True),
-                nn.SiLU(),
-                nn.Linear(hidden_size, hidden_size, bias=True)
-            )
-
-        self.mouse_qk_dim_list = mouse_qk_dim_list
-        self.heads_num = heads_num
-        if self.enable_mouse:
-            c = mouse_hidden_dim
-            self.mouse_mlp = nn.Sequential(
-                nn.Linear(mouse_dim_in * vae_time_compression_ratio * windows_size + img_hidden_size, c, bias=True),
-                nn.GELU(approximate="tanh"),
-                nn.Linear(c, c),
-                FP32LayerNorm(c, elementwise_affine=True),
-            )
-            
-            head_dim = c // heads_num
-            self.t_qkv = ReplicatedLinear(c, c*3, bias=qkv_bias)
-            self.img_attn_q_norm = (
-                RMSNorm(head_dim, eps=1e-6)
-                if qk_norm
-                else nn.Identity()
-            )
-            self.img_attn_k_norm = (
-                RMSNorm(head_dim, eps=1e-6)
-                if qk_norm
-                else nn.Identity()
-            )
-            self.proj_mouse = ReplicatedLinear(c, img_hidden_size, bias=qkv_bias)
-
-        if self.enable_keyboard:
-            head_dim_key = keyboard_hidden_dim // heads_num
-            self.key_attn_q_norm = (
-                RMSNorm(head_dim_key, eps=1e-6)
-                if qk_norm
-                else nn.Identity()
-            )
-            self.key_attn_k_norm = (
-                RMSNorm(head_dim_key, eps=1e-6)
-                if qk_norm
-                else nn.Identity()
-            )
-            
-            self.mouse_attn_q = ReplicatedLinear(img_hidden_size, keyboard_hidden_dim, bias=qkv_bias)
-            self.keyboard_attn_kv = ReplicatedLinear(hidden_size * windows_size * vae_time_compression_ratio, keyboard_hidden_dim * 2, bias=qkv_bias)
-            self.proj_keyboard = ReplicatedLinear(keyboard_hidden_dim, img_hidden_size, bias=qkv_bias)
-
-        self.mouse_attn_layer = LocalAttention(
-            num_heads=heads_num,
-            head_size=mouse_hidden_dim // heads_num,
-            causal=False,
-            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
-                                          AttentionBackendEnum.TORCH_SDPA)
-        ) if self.enable_mouse else None
-
-        self.keyboard_attn_layer = LocalAttention(
-            num_heads=heads_num,
-            head_size=keyboard_hidden_dim // heads_num,
-            causal=False,
-            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
-                                          AttentionBackendEnum.TORCH_SDPA)
-        ) if self.enable_keyboard else None
-
-        self.vae_time_compression_ratio = vae_time_compression_ratio
-        self.windows_size = windows_size
-        self.patch_size = patch_size
-        # Lazy initialization: freqs will be created on first forward pass
-        self._freqs_cos = None
-        self._freqs_sin = None
-
-    def patchify(self, x, patch_size):
-        """
-        x : (N C T H W)
-        """
-        pt, ph, pw = self.patch_size
-        t, h, w = x.shape[2] //  pt, x.shape[3] // ph, x.shape[4] // pw
-        c = x.shape[1]
-        x = x.reshape(shape=(x.shape[0], c, t , pt, h , ph, w , pw))
-        x = torch.einsum("nctohpwq->nthwcopq", x)
-        x = x.reshape(shape=(x.shape[0], t*h*w,  c*pt*ph*pw))
-        return x
-
-    def unpatchify(self, x, t, h, w, patch_size):
-        """
-        x: (N, T, patch_size**2 * C)
-        imgs: (N, H, W, C)
-        """
-        c =  x.shape[2] // patch_size #self.unpatchify_channels
-        pt, ph, pw = self.patch_size
-        assert t * h * w == x.shape[1]
-
-        x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw))
-        x = torch.einsum("nthwcopq->nctohpwq", x)
-        imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))
-
-        return imgs
-
-    def get_rotary_pos_embed(self, video_length, height, width, head_dim, rope_dim_list = None, start_offset=0):
-        target_ndim = 3
-        ndim = 5 - 2
-        
-        latents_size = [video_length+start_offset, height, width]
-
-        if isinstance(self.patch_size, int):
-            assert all(s % self.patch_size == 0 for s in latents_size), (
-                f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), "
-                f"but got {latents_size}."
-            )
-            rope_sizes = [s // self.patch_size for s in latents_size]
-        elif isinstance(self.patch_size, list):
-            assert all(
-                s % self.patch_size[idx] == 0
-                for idx, s in enumerate(latents_size)
-            ), (
-                f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), "
-                f"but got {latents_size}."
-            )
-            rope_sizes = [
-                s // self.patch_size[idx] for idx, s in enumerate(latents_size)
-            ]
-
-        if len(rope_sizes) != target_ndim:
-            rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes  # time axis
-        
-        if rope_dim_list is None:
-            rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
-        assert (
-            sum(rope_dim_list) == head_dim
-        ), "sum(rope_dim_list) should equal to head_dim of attention layer"
-        # Use Matrix-Game wrapper for FastVideo's function
-        freqs_cos, freqs_sin = _get_nd_rotary_pos_embed_matrixgame(
-            rope_dim_list,
-            rope_sizes,
-            theta=self.rope_theta,
-            theta_rescale_factor=1,
-        )
-        return freqs_cos[-video_length*rope_sizes[1]*rope_sizes[2]//self.patch_size[0]:], freqs_sin[-video_length*rope_sizes[1]*rope_sizes[2]//self.patch_size[0]:]
-
-    def forward(self, x, tt, th, tw, mouse_condition=None, keyboard_condition=None, block_mask_mouse=None, block_mask_keyboard=None, is_causal=False, kv_cache_mouse=None, kv_cache_keyboard=None, start_frame=0, use_rope_keyboard=True, num_frame_per_block=3):
-        '''
-        hidden_states: B, tt*th*tw, C
-        mouse_condition: B, N_frames, C1
-        keyboard_condition: B, N_frames, C2
-        '''
-        assert use_rope_keyboard
-
-        B, N_frames, C = keyboard_condition.shape
-        assert tt*th*tw == x.shape[1]
-        assert ((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0
-        N_feats = int((N_frames - 1) / self.vae_time_compression_ratio) + 1
-        
-        # Lazy initialization of freqs on first forward pass
-        if self._freqs_cos is None or self._freqs_sin is None:
-            self._freqs_cos, self._freqs_sin = self.get_rotary_pos_embed(
-                7500, self.patch_size[1], self.patch_size[2], 64, 
-                self.mouse_qk_dim_list, start_offset=0
-            )
-        
-        # Defined freqs_cis early so it's available for both mouse and keyboard
-        freqs_cis = (self._freqs_cos, self._freqs_sin)
-
-        assert (N_feats == tt and ((is_causal and kv_cache_mouse is None) or not is_causal)) or ((N_frames - 1) // self.vae_time_compression_ratio + 1 == start_frame + num_frame_per_block and is_causal)
-        
-        if self.enable_mouse and mouse_condition is not None:
-            hidden_states = rearrange(x, "B (T S) C -> (B S) T C", T=tt, S=th*tw) # 65*272*480 -> 17*(272//16)*(480//16) -> 8670
-            B, N_frames, C = mouse_condition.shape
-        else:
-            hidden_states = x
-        # padding
-        
-        pad_t = self.vae_time_compression_ratio * self.windows_size
-        if self.enable_mouse and mouse_condition is not None:
-            pad = mouse_condition[:, 0:1, :].expand(-1, pad_t, -1)
-            mouse_condition = torch.cat([pad, mouse_condition], dim=1)
-            if is_causal and kv_cache_mouse is not None: 
-                mouse_condition = mouse_condition[:, self.vae_time_compression_ratio*(N_feats - num_frame_per_block - self.windows_size) + pad_t:, :] 
-                group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(num_frame_per_block)]
-            else:
-                group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(N_feats)]
-                
-            group_mouse = torch.stack(group_mouse, dim = 1)
-
-            S = th * tw 
-            group_mouse = group_mouse.unsqueeze(-1).expand(B, num_frame_per_block, pad_t, C, S)
-            group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(B * S, num_frame_per_block, pad_t * C) 
-
-            group_mouse = torch.cat([hidden_states, group_mouse], dim = -1)
-            group_mouse = self.mouse_mlp(group_mouse)
-            # qkv
-            mouse_qkv, _ = self.t_qkv(group_mouse)
-            q, k, v = rearrange(mouse_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num) # BHW F H C
-            q = self.img_attn_q_norm(q).to(v)
-            k = self.img_attn_k_norm(k).to(v)        
-            # rope embd
-
-            
-            # freqs_cis = (self.freqs_cos, self.freqs_sin)
-            
-            
-            q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame)
-            ## TODO: adding cache here
-            if is_causal:
-                if kv_cache_mouse is None:
-                    assert q.shape[0] ==  k.shape[0] and q.shape[0] % 880 == 0 # == 880, f"{q.shape[0]},{k.shape[0]}"
-                    padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
-                    padded_q = torch.cat(
-                        [q,
-                            torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],
-                                        device=q.device, dtype=v.dtype)],
-                        dim=1
-                    )
-                    padded_k = torch.cat(
-                        [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]],
-                                                device=k.device, dtype=v.dtype)],
-                        dim=1
-                    )
-                    padded_v = torch.cat(
-                        [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]],
-                                        device=v.device, dtype=v.dtype)],
-                        dim=1
-                    )
-                    attn = flex_attention(
-                        query=padded_q.transpose(2, 1), # after: B, HW, F, C
-                        key=padded_k.transpose(2, 1),
-                        value=padded_v.transpose(2, 1),
-                        block_mask=block_mask_mouse
-                    )[:, :, :-padded_length].transpose(2, 1)
-                else:
-                    current_start = start_frame
-                    current_end = current_start + q.shape[1]
-                    
-                    assert q.shape[1] == num_frame_per_block
-                    sink_size = 0
-                    max_attention_size = self.local_attn_size
-                    sink_tokens = sink_size * 1
-                    kv_cache_size = kv_cache_mouse["k"].shape[1]
-                    num_new_tokens = q.shape[1]
-                    
-
-                    if (current_end > kv_cache_mouse["global_end_index"].item()) and (
-                        num_new_tokens + kv_cache_mouse["local_end_index"].item() > kv_cache_size):
-                        num_evicted_tokens = num_new_tokens + kv_cache_mouse["local_end_index"].item() - kv_cache_size
-                        num_rolled_tokens = kv_cache_mouse["local_end_index"].item() - num_evicted_tokens - sink_tokens
-                        kv_cache_mouse["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                            kv_cache_mouse["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                        kv_cache_mouse["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                            kv_cache_mouse["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                        # Insert the new keys/values at the end
-                        local_end_index = kv_cache_mouse["local_end_index"].item() + current_end - \
-                            kv_cache_mouse["global_end_index"].item() - num_evicted_tokens
-                        local_start_index = local_end_index - num_new_tokens
-                    else:
-                        local_end_index = kv_cache_mouse["local_end_index"].item() + current_end - kv_cache_mouse["global_end_index"].item()
-                        local_start_index = local_end_index - num_new_tokens
-                    kv_cache_mouse["k"][:, local_start_index:local_end_index] = k
-                    kv_cache_mouse["v"][:, local_start_index:local_end_index] = v
-
-                    attn = self.mouse_attn_layer(
-                        q,
-                        kv_cache_mouse["k"][:, max(0, local_end_index - max_attention_size):local_end_index],
-                        kv_cache_mouse["v"][:, max(0, local_end_index - max_attention_size):local_end_index],
-                    )
-                    kv_cache_mouse["global_end_index"].fill_(current_end)
-                    kv_cache_mouse["local_end_index"].fill_(local_end_index)
-            else:
-                attn = self.mouse_attn_layer(q, k, v)
-            # Compute cu_squlens and max_seqlen for flash attention
-            # qk norm
-            attn = rearrange(attn, '(b S) T h d -> b (T S) (h d)',b=B)
-            
-            hidden_states = rearrange(x, "(B S) T C -> B (T S) C", B=B)
-            attn, _ = self.proj_mouse(attn)
-            
-            hidden_states = hidden_states + attn
-        
-        if self.enable_keyboard and keyboard_condition is not None:
-            pad = keyboard_condition[:, 0:1, :].expand(-1, pad_t, -1)
-            keyboard_condition = torch.cat([pad, keyboard_condition], dim=1)
-            if is_causal and kv_cache_keyboard is not None:
-                keyboard_condition = keyboard_condition[:, self.vae_time_compression_ratio*(N_feats - num_frame_per_block - self.windows_size) + pad_t:, :] # keyboard_condition[:, self.vae_time_compression_ratio*(start_frame - self.windows_size) + pad_t:start_frame * self.vae_time_compression_ratio + pad_t,:]
-                keyboard_condition = self.keyboard_embed(keyboard_condition)
-                group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(num_frame_per_block)]
-            else:
-                keyboard_condition = self.keyboard_embed(keyboard_condition)
-                group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(N_feats)]
-            group_keyboard = torch.stack(group_keyboard, dim = 1) # B F RW C
-            group_keyboard = group_keyboard.reshape(shape=(group_keyboard.shape[0],group_keyboard.shape[1],-1))
-            # apply cross attn
-            mouse_q, _ = self.mouse_attn_q(hidden_states)
-            keyboard_kv, _ = self.keyboard_attn_kv(group_keyboard)
-
-            B, L, HD = mouse_q.shape
-            D = HD // self.heads_num
-            q = mouse_q.view(B, L, self.heads_num, D)
-
-            B, L, KHD = keyboard_kv.shape
-            k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute(2, 0, 1, 3, 4)
-        
-            # Compute cu_squlens and max_seqlen for flash attention
-            # qk norm
-            
-            q = self.key_attn_q_norm(q).to(v)
-            k = self.key_attn_k_norm(k).to(v)
-            S = th * tw
-            assert S == 880
-            # position embed 
-            if use_rope_keyboard: 
-                B, TS, H, D = q.shape
-                T_ = TS // S 
-                q = q.view(B, T_, S, H, D).transpose(1, 2).reshape(B * S, T_, H, D)
-                q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame)
-
-                k1, k2, k3, k4 = k.shape
-                k = k.expand(S, k2, k3, k4)
-                v = v.expand(S, k2, k3, k4)
-
-
-                if is_causal:
-                    if kv_cache_keyboard is None:
-                        assert q.shape[0] == k.shape[0] and q.shape[0] % 880 == 0 
-
-                        padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
-                        padded_q = torch.cat(
-                            [q,
-                                torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],
-                                            device=q.device, dtype=v.dtype)],
-                            dim=1
-                        )
-                        padded_k = torch.cat(
-                            [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]],
-                                                    device=k.device, dtype=v.dtype)],
-                            dim=1
-                        )
-                        padded_v = torch.cat(
-                            [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]],
-                                            device=v.device, dtype=v.dtype)],
-                            dim=1
-                        )
-                        attn = flex_attention(
-                            query=padded_q.transpose(2, 1), # after: B, HW, F, C
-                            key=padded_k.transpose(2, 1),
-                            value=padded_v.transpose(2, 1),
-                            block_mask=block_mask_keyboard
-                        )[:, :, :-padded_length].transpose(2, 1)
-                    else:
-                        current_start = start_frame
-                        current_end = current_start + k.shape[1]
-                        assert k.shape[1] == num_frame_per_block
-                        sink_size = 0
-                        max_attention_size = self.local_attn_size
-                        sink_tokens = sink_size * 1
-                        kv_cache_size = kv_cache_keyboard["k"].shape[1]
-                        num_new_tokens = k.shape[1]
-
-                        if (current_end > kv_cache_keyboard["global_end_index"].item()) and (
-                            num_new_tokens + kv_cache_keyboard["local_end_index"].item() > kv_cache_size):
-                            num_evicted_tokens = num_new_tokens + kv_cache_keyboard["local_end_index"].item() - kv_cache_size
-                            num_rolled_tokens = kv_cache_keyboard["local_end_index"].item() - num_evicted_tokens - sink_tokens
-                            kv_cache_keyboard["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                                kv_cache_keyboard["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                            kv_cache_keyboard["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                                kv_cache_keyboard["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                            # Insert the new keys/values at the end
-                            local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - \
-                                kv_cache_keyboard["global_end_index"].item() - num_evicted_tokens
-                            local_start_index = local_end_index - num_new_tokens
-                        else:
-                            local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - kv_cache_keyboard["global_end_index"].item()
-                            local_start_index = local_end_index - num_new_tokens
-                        assert k.shape[0] == 880 # BS == 1 or the cache should not be saved/ load method should be modified
-                        kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k[:1]
-                        kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v[:1]
-
-                        attn = self.keyboard_attn_layer(
-                            q,
-                            kv_cache_keyboard["k"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat(S, 1, 1, 1),
-                            kv_cache_keyboard["v"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat(S, 1, 1, 1),
-                        )
-
-                        kv_cache_keyboard["global_end_index"].fill_(current_end)
-                        kv_cache_keyboard["local_end_index"].fill_(local_end_index)
-                else:
-                    attn = self.keyboard_attn_layer(q, k, v)
-                attn = rearrange(attn, '(B S) T H D -> B (T S) (H D)', S=S)
-            else:
-                if is_causal:
-                    if kv_cache_keyboard is None:
-                        
-                        padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
-                        padded_q = torch.cat(
-                            [q,
-                                torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],
-                                            device=q.device, dtype=v.dtype)],
-                            dim=1
-                        )
-                        padded_k = torch.cat(
-                            [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]],
-                                                    device=k.device, dtype=v.dtype)],
-                            dim=1
-                        )
-                        padded_v = torch.cat(
-                            [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]],
-                                            device=v.device, dtype=v.dtype)],
-                            dim=1
-                        )
-                        attn = flex_attention(
-                            query=padded_q.transpose(2, 1), # after: B, HW, F, C
-                            key=padded_k.transpose(2, 1),
-                            value=padded_v.transpose(2, 1),
-                            block_mask=block_mask_keyboard
-                        )[:, :, :-padded_length].transpose(2, 1)
-                    else:
-                        current_start = start_frame
-                        current_end = current_start + k.shape[1]
-                        assert k.shape[1] == num_frame_per_block
-                        sink_size = 0
-                        max_attention_size = self.local_attn_size
-                        sink_tokens = sink_size * 1
-                        kv_cache_size = kv_cache_keyboard["k"].shape[1]
-                        num_new_tokens = k.shape[1]
-
-
-                        if (current_end > kv_cache_keyboard["global_end_index"].item()) and (
-                            num_new_tokens + kv_cache_keyboard["local_end_index"].item() > kv_cache_size):
-                            num_evicted_tokens = num_new_tokens + kv_cache_keyboard["local_end_index"].item() - kv_cache_size
-                            num_rolled_tokens = kv_cache_keyboard["local_end_index"].item() - num_evicted_tokens - sink_tokens
-                            kv_cache_keyboard["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                                kv_cache_keyboard["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                            kv_cache_keyboard["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
-                                kv_cache_keyboard["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
-                            # Insert the new keys/values at the end
-                            local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - \
-                                kv_cache_keyboard["global_end_index"].item() - num_evicted_tokens
-                            local_start_index = local_end_index - num_new_tokens
-
-                            
-                        else:
-                            local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - kv_cache_keyboard["global_end_index"].item()
-                            local_start_index = local_end_index - num_new_tokens
-                        kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k
-                        kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v
-                        attn = self.keyboard_attn_layer(
-                            q,
-                            kv_cache_keyboard["k"][:, max(0, local_end_index - max_attention_size):local_end_index],
-                            kv_cache_keyboard["v"][:, max(0, local_end_index - max_attention_size):local_end_index],
-                        )
-                        kv_cache_keyboard["global_end_index"].fill_(current_end)
-                        kv_cache_keyboard["local_end_index"].fill_(local_end_index)
-                else:
-                    attn = self.keyboard_attn_layer(q, k, v)
-                attn = rearrange(attn, 'B L H D -> B L (H D)')
-            attn, _ = self.proj_keyboard(attn)
-            hidden_states = hidden_states + attn
-        return hidden_states
@@ -9,3 +9,6 @@ __all__ = [
    "CausalMatrixGameTransformerBlock",
    "ActionModule",
 ]
+
+# Entry point for model registry
+EntryClass = [MatrixGameWanModel, CausalMatrixGameWanModel]
@@ -0,0 +1,822 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from Matrix-Game: https://github.com/SkyworkAI/Matrix-Game/blob/main/Matrix-Game-2/wan/modules/action_module.py
+
+from einops import rearrange
+import torch
+import torch.nn as nn
+import math
+from torch.nn.attention.flex_attention import flex_attention, BlockMask
+
+from fastvideo.attention import LocalAttention
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.layernorm import FP32LayerNorm, RMSNorm
+from fastvideo.layers.rotary_embedding import (
+    get_nd_rotary_pos_embed as _fv_get_nd_rotary_pos_embed,
+    _apply_rotary_emb,
+)
+from fastvideo.platforms import AttentionBackendEnum
+
+
+DISABLE_COMPILE = False
+flex_attention = torch.compile(
+    flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs"
+)
+
+
+def _get_nd_rotary_pos_embed_matrixgame(
+    rope_dim_list,
+    rope_sizes,
+    theta: float = 10000.0,
+    theta_rescale_factor: float = 1.0,
+):
+    cos, sin = _fv_get_nd_rotary_pos_embed(
+        rope_dim_list,
+        rope_sizes,
+        theta=theta,
+        theta_rescale_factor=theta_rescale_factor,
+        dtype=torch.float32,
+    )
+    # convert from [S, D/2] to [S, D] format
+    cos = cos.repeat_interleave(2, dim=1)
+    sin = sin.repeat_interleave(2, dim=1)
+    return cos, sin
+
+
+def _apply_rotary_emb_qk(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor,
+    start_offset: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    seq_len = xq.shape[1]
+
+    # Slice frequencies based on offset
+    cos = freqs_cos[start_offset : start_offset + seq_len]  # [S, D]
+    sin = freqs_sin[start_offset : start_offset + seq_len]  # [S, D]
+
+    # Move to device
+    cos = cos.to(xq.device)
+    sin = sin.to(xq.device)
+
+    # Convert from [S, D] (interleaved) back to [S, D/2]
+    cos_half = cos[:, ::2]  # [S, D/2]
+    sin_half = sin[:, ::2]  # [S, D/2]
+
+    # xq/xk are [B, S, H, D], need to reshape for each batch
+    B, S, H, D = xq.shape
+
+    xq_out = _apply_rotary_emb(xq, cos_half, sin_half, is_neox_style=False)
+    xk_out = _apply_rotary_emb(xk, cos_half, sin_half, is_neox_style=False)
+
+    return xq_out, xk_out
+
+
+def _padding_q_k_v(tensor: torch.Tensor, padded_length: int) -> torch.Tensor:
+    return torch.cat(
+        [
+            tensor,
+            torch.zeros(
+                [
+                    tensor.shape[0],
+                    padded_length,
+                    tensor.shape[2],
+                    tensor.shape[3],
+                ],
+                device=tensor.device,
+                dtype=tensor.dtype,
+            ),
+        ],
+        dim=1,
+    )
+
+
+def _update_kv_cache_and_attend(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    kv_cache: dict[str, torch.Tensor],
+    attn_layer: LocalAttention,
+    start_frame: int,
+    num_frame_per_block: int,
+    local_attn_size: int,
+    *,
+    use_k_for_num_tokens: bool = False,
+    store_first_only: bool = False,
+    repeat_factor: int | None = None,
+) -> torch.Tensor:
+    """
+    Update KV cache with new tokens and perform attention with cached values.
+
+    Args:
+        q: Query tensor
+        k: Key tensor
+        v: Value tensor
+        kv_cache: Dictionary containing cache tensors and indices
+        attn_layer: Attention layer to use
+        start_frame: Starting frame index
+        num_frame_per_block: Number of frames per block
+        local_attn_size: Maximum attention window size
+        use_k_for_num_tokens: If True, use k.shape[1] for num_new_tokens, else use q.shape[1]
+        store_first_only: If True, only store k[:1] and v[:1] in cache (for keyboard with rope)
+        repeat_factor: If provided, repeat cached k,v by this factor when retrieving (for keyboard with rope)
+
+    Returns:
+        Attention output tensor
+    """
+    current_start = start_frame
+    current_end = current_start + (
+        k.shape[1] if use_k_for_num_tokens else q.shape[1]
+    )
+
+    assert (
+        k.shape[1] if use_k_for_num_tokens else q.shape[1]
+    ) == num_frame_per_block
+
+    sink_size = 0
+    max_attention_size = local_attn_size
+    sink_tokens = sink_size * 1
+    kv_cache_size = kv_cache["k"].shape[1]
+    num_new_tokens = k.shape[1] if use_k_for_num_tokens else q.shape[1]
+
+    # Check if we need to evict tokens
+    if (current_end > kv_cache["global_end_index"].item()) and (
+        num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size
+    ):
+        num_evicted_tokens = (
+            num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
+        )
+        num_rolled_tokens = (
+            kv_cache["local_end_index"].item()
+            - num_evicted_tokens
+            - sink_tokens
+        )
+        # Roll k cache
+        kv_cache["k"][:, sink_tokens : sink_tokens + num_rolled_tokens] = (
+            kv_cache["k"][
+                :,
+                sink_tokens + num_evicted_tokens : sink_tokens
+                + num_evicted_tokens
+                + num_rolled_tokens,
+            ].clone()
+        )
+        # Roll v cache
+        kv_cache["v"][:, sink_tokens : sink_tokens + num_rolled_tokens] = (
+            kv_cache["v"][
+                :,
+                sink_tokens + num_evicted_tokens : sink_tokens
+                + num_evicted_tokens
+                + num_rolled_tokens,
+            ].clone()
+        )
+        # Calculate indices with eviction adjustment
+        local_end_index = (
+            kv_cache["local_end_index"].item()
+            + current_end
+            - kv_cache["global_end_index"].item()
+            - num_evicted_tokens
+        )
+        local_start_index = local_end_index - num_new_tokens
+    else:
+        # Calculate indices without eviction
+        local_end_index = (
+            kv_cache["local_end_index"].item()
+            + current_end
+            - kv_cache["global_end_index"].item()
+        )
+        local_start_index = local_end_index - num_new_tokens
+
+    # Store new k, v in cache
+    if store_first_only:
+        kv_cache["k"][:, local_start_index:local_end_index] = k[:1]
+        kv_cache["v"][:, local_start_index:local_end_index] = v[:1]
+    else:
+        kv_cache["k"][:, local_start_index:local_end_index] = k
+        kv_cache["v"][:, local_start_index:local_end_index] = v
+
+    # Retrieve from cache and perform attention
+    cache_start = max(0, local_end_index - max_attention_size)
+    cached_k = kv_cache["k"][:, cache_start:local_end_index]
+    cached_v = kv_cache["v"][:, cache_start:local_end_index]
+
+    if repeat_factor is not None:
+        cached_k = cached_k.repeat(repeat_factor, 1, 1, 1)
+        cached_v = cached_v.repeat(repeat_factor, 1, 1, 1)
+
+    attn = attn_layer(q, cached_k, cached_v)
+
+    # Update indices
+    kv_cache["global_end_index"].fill_(current_end)
+    kv_cache["local_end_index"].fill_(local_end_index)
+
+    return attn
+
+
+class ActionModule(nn.Module):
+    """
+    action module from https://arxiv.org/pdf/2501.08325
+    """
+
+    def __init__(
+        self,
+        mouse_dim_in: int = 2,
+        keyboard_dim_in: int = 6,
+        hidden_size: int = 128,
+        img_hidden_size: int = 1536,
+        keyboard_hidden_dim: int = 1024,
+        mouse_hidden_dim: int = 1024,
+        vae_time_compression_ratio: int = 4,
+        windows_size: int = 3,
+        heads_num: int = 16,
+        patch_size: list | None = None,
+        qk_norm: bool = True,
+        qkv_bias: bool = False,
+        rope_dim_list: list | None = None,
+        rope_theta=256,
+        mouse_qk_dim_list: list | None = None,
+        enable_mouse=True,
+        enable_keyboard=True,
+        local_attn_size=6,
+        blocks: list | None = None,
+    ):
+        super().__init__()
+        # Initialize mutable defaults
+        patch_size = patch_size if patch_size is not None else [1, 2, 2]
+        rope_dim_list = (
+            rope_dim_list if rope_dim_list is not None else [8, 28, 28]
+        )
+        mouse_qk_dim_list = (
+            mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28]
+        )
+        blocks = blocks if blocks is not None else []
+        self.local_attn_size = local_attn_size
+        self.enable_mouse = enable_mouse
+        self.enable_keyboard = enable_keyboard
+
+        self.rope_dim_list = rope_dim_list
+        self.rope_theta = rope_theta
+        if self.enable_keyboard:
+            self.keyboard_embed = nn.Sequential(
+                nn.Linear(keyboard_dim_in, hidden_size, bias=True),
+                nn.SiLU(),
+                nn.Linear(hidden_size, hidden_size, bias=True),
+            )
+
+        self.mouse_qk_dim_list = mouse_qk_dim_list
+        self.heads_num = heads_num
+        if self.enable_mouse:
+            c = mouse_hidden_dim
+            self.mouse_mlp = nn.Sequential(
+                nn.Linear(
+                    mouse_dim_in * vae_time_compression_ratio * windows_size
+                    + img_hidden_size,
+                    c,
+                    bias=True,
+                ),
+                nn.GELU(approximate="tanh"),
+                nn.Linear(c, c),
+                FP32LayerNorm(c, elementwise_affine=True),
+            )
+
+            head_dim = c // heads_num
+            self.t_qkv = ReplicatedLinear(c, c * 3, bias=qkv_bias)
+            self.img_attn_q_norm = (
+                RMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity()
+            )
+            self.img_attn_k_norm = (
+                RMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity()
+            )
+            self.proj_mouse = ReplicatedLinear(
+                c, img_hidden_size, bias=qkv_bias
+            )
+
+        if self.enable_keyboard:
+            head_dim_key = keyboard_hidden_dim // heads_num
+            self.key_attn_q_norm = (
+                RMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity()
+            )
+            self.key_attn_k_norm = (
+                RMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity()
+            )
+
+            self.mouse_attn_q = ReplicatedLinear(
+                img_hidden_size, keyboard_hidden_dim, bias=qkv_bias
+            )
+            self.keyboard_attn_kv = ReplicatedLinear(
+                hidden_size * windows_size * vae_time_compression_ratio,
+                keyboard_hidden_dim * 2,
+                bias=qkv_bias,
+            )
+            self.proj_keyboard = ReplicatedLinear(
+                keyboard_hidden_dim, img_hidden_size, bias=qkv_bias
+            )
+
+        self.mouse_attn_layer = (
+            LocalAttention(
+                num_heads=heads_num,
+                head_size=mouse_hidden_dim // heads_num,
+                causal=False,
+                supported_attention_backends=(
+                    AttentionBackendEnum.FLASH_ATTN,
+                    AttentionBackendEnum.TORCH_SDPA,
+                ),
+            )
+            if self.enable_mouse
+            else None
+        )
+
+        self.keyboard_attn_layer = (
+            LocalAttention(
+                num_heads=heads_num,
+                head_size=keyboard_hidden_dim // heads_num,
+                causal=False,
+                supported_attention_backends=(
+                    AttentionBackendEnum.FLASH_ATTN,
+                    AttentionBackendEnum.TORCH_SDPA,
+                ),
+            )
+            if self.enable_keyboard
+            else None
+        )
+
+        self.vae_time_compression_ratio = vae_time_compression_ratio
+        self.windows_size = windows_size
+        self.patch_size = patch_size
+        # Lazy initialization: freqs will be created on first forward pass
+        self._freqs_cos = None
+        self._freqs_sin = None
+
+    def patchify(self, x, patch_size):
+        """
+        x : (N C T H W)
+        """
+        pt, ph, pw = self.patch_size
+        t, h, w = x.shape[2] // pt, x.shape[3] // ph, x.shape[4] // pw
+        c = x.shape[1]
+        x = x.reshape(shape=(x.shape[0], c, t, pt, h, ph, w, pw))
+        x = torch.einsum("nctohpwq->nthwcopq", x)
+        x = x.reshape(shape=(x.shape[0], t * h * w, c * pt * ph * pw))
+        return x
+
+    def unpatchify(self, x, t, h, w, patch_size):
+        """
+        x: (N, T, patch_size**2 * C)
+        imgs: (N, H, W, C)
+        """
+        c = x.shape[2] // patch_size  # self.unpatchify_channels
+        pt, ph, pw = self.patch_size
+        assert t * h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw))
+        x = torch.einsum("nthwcopq->nctohpwq", x)
+        imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))
+
+        return imgs
+
+    def get_rotary_pos_embed(
+        self,
+        video_length,
+        height,
+        width,
+        head_dim,
+        rope_dim_list=None,
+        start_offset=0,
+    ):
+        target_ndim = 3
+        ndim = 5 - 2
+
+        latents_size = [video_length + start_offset, height, width]
+
+        if isinstance(self.patch_size, int):
+            assert all(s % self.patch_size == 0 for s in latents_size), (
+                f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), "
+                f"but got {latents_size}."
+            )
+            rope_sizes = [s // self.patch_size for s in latents_size]
+        elif isinstance(self.patch_size, list):
+            assert all(
+                s % self.patch_size[idx] == 0
+                for idx, s in enumerate(latents_size)
+            ), (
+                f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), "
+                f"but got {latents_size}."
+            )
+            rope_sizes = [
+                s // self.patch_size[idx] for idx, s in enumerate(latents_size)
+            ]
+
+        if len(rope_sizes) != target_ndim:
+            rope_sizes = [1] * (
+                target_ndim - len(rope_sizes)
+            ) + rope_sizes  # time axis
+
+        if rope_dim_list is None:
+            rope_dim_list = [
+                head_dim // target_ndim for _ in range(target_ndim)
+            ]
+        assert sum(rope_dim_list) == head_dim, (
+            "sum(rope_dim_list) should equal to head_dim of attention layer"
+        )
+        # Use Matrix-Game wrapper for FastVideo's function
+        freqs_cos, freqs_sin = _get_nd_rotary_pos_embed_matrixgame(
+            rope_dim_list,
+            rope_sizes,
+            theta=self.rope_theta,
+            theta_rescale_factor=1,
+        )
+        return freqs_cos[
+            -video_length
+            * rope_sizes[1]
+            * rope_sizes[2]
+            // self.patch_size[0] :
+        ], freqs_sin[
+            -video_length
+            * rope_sizes[1]
+            * rope_sizes[2]
+            // self.patch_size[0] :
+        ]
+
+    def _forward_mouse(
+        self,
+        hidden_states: torch.Tensor,
+        mouse_condition: torch.Tensor,
+        *,
+        is_causal: bool,
+        kv_cache_mouse: dict[str, torch.Tensor] | None,
+        pad_t: int,
+        num_frame_per_block: int,
+        block_mask_mouse: BlockMask | None,
+        start_frame: int,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        N_feats: int,
+        B: int,
+        C: int,
+        tt: int,
+        th: int,
+        tw: int,
+    ) -> torch.Tensor:
+        pad = mouse_condition[:, 0:1, :].expand(-1, pad_t, -1)
+        mouse_condition = torch.cat([pad, mouse_condition], dim=1)
+        if is_causal and kv_cache_mouse is not None:
+            mouse_condition = mouse_condition[
+                :,
+                self.vae_time_compression_ratio
+                * (N_feats - num_frame_per_block - self.windows_size)
+                + pad_t :,
+                :,
+            ]
+            group_mouse = [
+                mouse_condition[
+                    :,
+                    self.vae_time_compression_ratio * (i - self.windows_size)
+                    + pad_t : i * self.vae_time_compression_ratio + pad_t,
+                    :,
+                ]
+                for i in range(num_frame_per_block)
+            ]
+        else:
+            local_num_frames = tt
+            group_mouse = [
+                mouse_condition[
+                    :,
+                    self.vae_time_compression_ratio * (i - self.windows_size)
+                    + pad_t : i * self.vae_time_compression_ratio + pad_t,
+                    :,
+                ]
+                for i in range(local_num_frames)
+            ]
+
+        group_mouse = torch.stack(group_mouse, dim=1)
+        actual_num_frames = group_mouse.shape[
+            1
+        ]  # Use actual stacked frame count
+
+        S = th * tw
+        group_mouse = group_mouse.unsqueeze(-1).expand(
+            B, actual_num_frames, pad_t, C, S
+        )
+        group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(
+            B * S, actual_num_frames, pad_t * C
+        )
+
+        group_mouse = torch.cat([hidden_states, group_mouse], dim=-1)
+        group_mouse = self.mouse_mlp(group_mouse)
+        # qkv
+        mouse_qkv, _ = self.t_qkv(group_mouse)
+        q, k, v = rearrange(
+            mouse_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num
+        )  # BHW F H C
+        q = self.img_attn_q_norm(q).to(v)
+        k = self.img_attn_k_norm(k).to(v)
+        # rope embd
+
+        # freqs_cis = (self.freqs_cos, self.freqs_sin)
+
+        q, k = _apply_rotary_emb_qk(
+            q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame
+        )
+        ## TODO: adding cache here
+        if is_causal:
+            if kv_cache_mouse is None:
+                assert (
+                    q.shape[0] == k.shape[0] and q.shape[0] % S == 0
+                )  # == 880, f"{q.shape[0]},{k.shape[0]}"
+                padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
+                padded_q = _padding_q_k_v(q, padded_length)
+                padded_k = _padding_q_k_v(k, padded_length)
+                padded_v = _padding_q_k_v(v, padded_length)
+                attn = flex_attention(
+                    query=padded_q.transpose(2, 1),  # after: B, HW, F, C
+                    key=padded_k.transpose(2, 1),
+                    value=padded_v.transpose(2, 1),
+                    block_mask=block_mask_mouse,
+                )[:, :, :-padded_length].transpose(2, 1)
+            else:
+                attn = _update_kv_cache_and_attend(
+                    q,
+                    k,
+                    v,
+                    kv_cache_mouse,
+                    self.mouse_attn_layer,
+                    start_frame,
+                    num_frame_per_block,
+                    self.local_attn_size,
+                    use_k_for_num_tokens=False,
+                )
+        else:
+            attn = self.mouse_attn_layer(q, k, v)
+        # Compute cu_squlens and max_seqlen for flash attention
+        # qk norm
+        attn = rearrange(attn, "(b S) T h d -> b (T S) (h d)", b=B)
+        attn, _ = self.proj_mouse(attn)
+        return attn
+
+    def _forward_keyboard(
+        self,
+        hidden_states: torch.Tensor,
+        keyboard_condition: torch.Tensor,
+        *,
+        is_causal: bool,
+        use_rope_keyboard: bool,
+        kv_cache_keyboard: dict[str, torch.Tensor] | None,
+        pad_t: int,
+        num_frame_per_block: int,
+        block_mask_keyboard: BlockMask | None,
+        start_frame: int,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        N_feats: int,
+        B: int,
+        tt: int,
+        th: int,
+        tw: int,
+    ) -> torch.Tensor:
+        pad = keyboard_condition[:, 0:1, :].expand(-1, pad_t, -1)
+        keyboard_condition = torch.cat([pad, keyboard_condition], dim=1)
+        if is_causal and kv_cache_keyboard is not None:
+            keyboard_condition = keyboard_condition[
+                :,
+                self.vae_time_compression_ratio
+                * (N_feats - num_frame_per_block - self.windows_size)
+                + pad_t :,
+                :,
+            ]  # keyboard_condition[:, self.vae_time_compression_ratio*(start_frame - self.windows_size) + pad_t:start_frame * self.vae_time_compression_ratio + pad_t,:]
+            keyboard_condition = self.keyboard_embed(keyboard_condition)
+            group_keyboard = [
+                keyboard_condition[
+                    :,
+                    self.vae_time_compression_ratio * (i - self.windows_size)
+                    + pad_t : i * self.vae_time_compression_ratio + pad_t,
+                    :,
+                ]
+                for i in range(num_frame_per_block)
+            ]
+        else:
+            keyboard_condition = self.keyboard_embed(keyboard_condition)
+            local_num_frames = tt
+            group_keyboard = [
+                keyboard_condition[
+                    :,
+                    self.vae_time_compression_ratio * (i - self.windows_size)
+                    + pad_t : i * self.vae_time_compression_ratio + pad_t,
+                    :,
+                ]
+                for i in range(local_num_frames)
+            ]
+        group_keyboard = torch.stack(group_keyboard, dim=1)  # B F RW C
+        group_keyboard = group_keyboard.reshape(
+            shape=(group_keyboard.shape[0], group_keyboard.shape[1], -1)
+        )
+        # apply cross attn
+        mouse_q, _ = self.mouse_attn_q(hidden_states)
+        keyboard_kv, _ = self.keyboard_attn_kv(group_keyboard)
+
+        B, L, HD = mouse_q.shape
+        D = HD // self.heads_num
+        q = mouse_q.view(B, L, self.heads_num, D)
+
+        B, L, KHD = keyboard_kv.shape
+        k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute(
+            2, 0, 1, 3, 4
+        )
+
+        # Compute cu_squlens and max_seqlen for flash attention
+        # qk norm
+
+        q = self.key_attn_q_norm(q).to(v)
+        k = self.key_attn_k_norm(k).to(v)
+        S = th * tw
+        # assert S == 880
+        # position embed
+        if use_rope_keyboard:
+            B, TS, H, D = q.shape
+            T_ = TS // S
+            q = q.view(B, T_, S, H, D).transpose(1, 2).reshape(B * S, T_, H, D)
+            q, k = _apply_rotary_emb_qk(
+                q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame
+            )
+
+            k = k.repeat_interleave(S, dim=0)
+            v = v.repeat_interleave(S, dim=0)
+
+            if is_causal:
+                if kv_cache_keyboard is None:
+                    assert q.shape[0] == k.shape[0] and q.shape[0] % S == 0
+
+                    padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
+                    padded_q = _padding_q_k_v(q, padded_length)
+                    padded_k = _padding_q_k_v(k, padded_length)
+                    padded_v = _padding_q_k_v(v, padded_length)
+                    attn = flex_attention(
+                        query=padded_q.transpose(2, 1),  # after: B, HW, F, C
+                        key=padded_k.transpose(2, 1),
+                        value=padded_v.transpose(2, 1),
+                        block_mask=block_mask_keyboard,
+                    )[:, :, :-padded_length].transpose(2, 1)
+                else:
+                    assert (
+                        k.shape[0] == S
+                    )  # BS == 1 or the cache should not be saved/ load method should be modified
+                    attn = _update_kv_cache_and_attend(
+                        q,
+                        k,
+                        v,
+                        kv_cache_keyboard,
+                        self.keyboard_attn_layer,
+                        start_frame,
+                        num_frame_per_block,
+                        self.local_attn_size,
+                        use_k_for_num_tokens=True,
+                        store_first_only=True,
+                        repeat_factor=S,
+                    )
+            else:
+                attn = self.keyboard_attn_layer(q, k, v)
+            attn = rearrange(attn, "(B S) T H D -> B (T S) (H D)", S=S)
+        else:
+            if is_causal:
+                if kv_cache_keyboard is None:
+                    padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1]
+                    padded_q = _padding_q_k_v(q, padded_length)
+                    padded_k = _padding_q_k_v(k, padded_length)
+                    padded_v = _padding_q_k_v(v, padded_length)
+                    attn = flex_attention(
+                        query=padded_q.transpose(2, 1),  # after: B, HW, F, C
+                        key=padded_k.transpose(2, 1),
+                        value=padded_v.transpose(2, 1),
+                        block_mask=block_mask_keyboard,
+                    )[:, :, :-padded_length].transpose(2, 1)
+                else:
+                    attn = _update_kv_cache_and_attend(
+                        q,
+                        k,
+                        v,
+                        kv_cache_keyboard,
+                        self.keyboard_attn_layer,
+                        start_frame,
+                        num_frame_per_block,
+                        self.local_attn_size,
+                        use_k_for_num_tokens=True,
+                    )
+            else:
+                attn = self.keyboard_attn_layer(q, k, v)
+            attn = rearrange(attn, "B L H D -> B L (H D)")
+        attn, _ = self.proj_keyboard(attn)
+        return attn
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        tt: int,
+        th: int,
+        tw: int,
+        mouse_condition: torch.Tensor | None = None,
+        keyboard_condition: torch.Tensor | None = None,
+        block_mask_mouse: BlockMask | None = None,
+        block_mask_keyboard: BlockMask | None = None,
+        is_causal: bool = False,
+        kv_cache_mouse: dict[str, torch.Tensor] | None = None,
+        kv_cache_keyboard: dict[str, torch.Tensor] | None = None,
+        start_frame: int = 0,
+        use_rope_keyboard: bool = True,
+        num_frame_per_block: int = 3,
+    ):
+        """
+        hidden_states: B, tt*th*tw, C
+        mouse_condition: B, N_frames, C1
+        keyboard_condition: B, N_frames, C2
+        """
+        assert use_rope_keyboard
+
+        target_device = x.device
+        target_dtype = x.dtype
+        if mouse_condition is not None:
+            mouse_condition = mouse_condition.to(
+                device=target_device, dtype=target_dtype
+            )
+        if keyboard_condition is not None:
+            keyboard_condition = keyboard_condition.to(
+                device=target_device, dtype=target_dtype
+            )
+        else:
+            return x
+
+        B, N_frames, C = keyboard_condition.shape
+        assert tt * th * tw == x.shape[1]
+        assert (
+            (N_frames - 1) + self.vae_time_compression_ratio
+        ) % self.vae_time_compression_ratio == 0
+        N_feats = int((N_frames - 1) / self.vae_time_compression_ratio) + 1
+
+        # Lazy initialization of freqs on first forward pass
+        if self._freqs_cos is None or self._freqs_sin is None:
+            self._freqs_cos, self._freqs_sin = self.get_rotary_pos_embed(
+                7500,
+                self.patch_size[1],
+                self.patch_size[2],
+                64,
+                self.mouse_qk_dim_list,
+                start_offset=0,
+            )
+
+        # Defined freqs_cis early so it's available for both mouse and keyboard
+        freqs_cis = (self._freqs_cos, self._freqs_sin)
+
+        if is_causal:
+            assert (N_feats == tt and kv_cache_mouse is None) or (
+                (N_frames - 1) // self.vae_time_compression_ratio + 1
+                == start_frame + num_frame_per_block
+            )
+        # For non-causal (training), we trust that the caller provides correctly shaped inputs
+
+        if self.enable_mouse and mouse_condition is not None:
+            hidden_states = rearrange(
+                x, "B (T S) C -> (B S) T C", T=tt, S=th * tw
+            )  # 65*272*480 -> 17*(272//16)*(480//16) -> 8670
+            B, N_frames, C = mouse_condition.shape
+        else:
+            hidden_states = x
+        # padding
+
+        pad_t = self.vae_time_compression_ratio * self.windows_size
+        if self.enable_mouse and mouse_condition is not None:
+            attn = self._forward_mouse(
+                hidden_states,
+                mouse_condition,
+                is_causal=is_causal,
+                kv_cache_mouse=kv_cache_mouse,
+                pad_t=pad_t,
+                num_frame_per_block=num_frame_per_block,
+                block_mask_mouse=block_mask_mouse,
+                start_frame=start_frame,
+                freqs_cis=freqs_cis,
+                N_feats=N_feats,
+                B=B,
+                C=C,
+                tt=tt,
+                th=th,
+                tw=tw,
+            )
+            hidden_states = rearrange(x, "(B S) T C -> B (T S) C", B=B)
+
+            hidden_states = hidden_states + attn
+
+        if self.enable_keyboard and keyboard_condition is not None:
+            attn = self._forward_keyboard(
+                hidden_states,
+                keyboard_condition,
+                is_causal=is_causal,
+                use_rope_keyboard=use_rope_keyboard,
+                kv_cache_keyboard=kv_cache_keyboard,
+                pad_t=pad_t,
+                num_frame_per_block=num_frame_per_block,
+                block_mask_keyboard=block_mask_keyboard,
+                start_frame=start_frame,
+                freqs_cis=freqs_cis,
+                N_feats=N_feats,
+                B=B,
+                tt=tt,
+                th=th,
+                tw=tw,
+            )
+            hidden_states = hidden_states + attn
+        return hidden_states
@@ -9,20 +9,32 @@ import torch.nn as nn
 from fastvideo.attention import DistributedAttention
 from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig
 from fastvideo.distributed.parallel_state import get_sp_world_size
-from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift,
-                                        RMSNorm, ScaleResidual,
-                                        ScaleResidualLayerNormScaleShift)
+from fastvideo.layers.layernorm import (
+    FP32LayerNorm,
+    LayerNormScaleShift,
+    RMSNorm,
+    ScaleResidual,
+    ScaleResidualLayerNormScaleShift,
+)
 from fastvideo.layers.linear import ReplicatedLinear
 from fastvideo.layers.mlp import MLP
-from fastvideo.layers.rotary_embedding import (_apply_rotary_emb,
-                                               get_rotary_pos_embed)
-from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection
+from fastvideo.layers.rotary_embedding import (
+    _apply_rotary_emb,
+    get_rotary_pos_embed,
+)
+from fastvideo.layers.visual_embedding import (
+    PatchEmbed,
+    TimestepEmbedder,
+    ModulateProjection,
+)
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import BaseDiT
-from fastvideo.models.dits.wanvideo import (WanSelfAttention,
-                                           WanI2VCrossAttention,
-                                           WanT2VCrossAttention,
-                                           WanImageEmbedding)
+from fastvideo.models.dits.wanvideo import (
+    WanSelfAttention,
+    WanI2VCrossAttention,
+    WanT2VCrossAttention,
+    WanImageEmbedding,
+)
 from fastvideo.platforms import AttentionBackendEnum, current_platform

 # Import ActionModule
@@ -41,11 +53,12 @@ class MatrixGameTimeImageEmbedding(nn.Module):
        super().__init__()

        self.time_embedder = TimestepEmbedder(
-            dim, frequency_embedding_size=time_freq_dim, act_layer="silu")
-        self.time_modulation = ModulateProjection(dim,
-                                                  factor=6,
-                                                  act_layer="silu")
-        
+            dim, frequency_embedding_size=time_freq_dim, act_layer="silu"
+        )
+        self.time_modulation = ModulateProjection(
+            dim, factor=6, act_layer="silu"
+        )
+
        self.image_embedder = None
        if image_embed_dim is not None:
            self.image_embedder = WanImageEmbedding(image_embed_dim, dim)
@@ -53,7 +66,7 @@ class MatrixGameTimeImageEmbedding(nn.Module):
    def forward(
        self,
        timestep: torch.Tensor,
-        encoder_hidden_states: torch.Tensor, # Kept for interface compatibility
+        encoder_hidden_states: torch.Tensor,  # Kept for interface compatibility
        encoder_hidden_states_image: torch.Tensor | None = None,
        timestep_seq_len: int | None = None,
    ):
@@ -61,14 +74,25 @@ class MatrixGameTimeImageEmbedding(nn.Module):
        timestep_proj = self.time_modulation(temb)

        # MatrixGame does not use text embeddings, so we ignore encoder_hidden_states
-        # and return None for the text embedding part
-        
+
        if encoder_hidden_states_image is not None:
            assert self.image_embedder is not None
            encoder_hidden_states_image = self.image_embedder(
-                encoder_hidden_states_image)
+                encoder_hidden_states_image
+            )

-        return temb, timestep_proj, None, encoder_hidden_states_image
+        encoder_hidden_states = torch.zeros(
+            (timestep.shape[0], 0, temb.shape[-1]),
+            device=temb.device,
+            dtype=temb.dtype,
+        )
+
+        return (
+            temb,
+            timestep_proj,
+            encoder_hidden_states,
+            encoder_hidden_states_image,
+        )


 class MatrixGameCrossAttention(WanSelfAttention):
@@ -84,7 +108,7 @@ class MatrixGameCrossAttention(WanSelfAttention):

        # compute query, key, value
        q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
-        
+
        if crossattn_cache is not None:
            if not crossattn_cache["is_init"]:
                crossattn_cache["is_init"] = True
@@ -98,7 +122,7 @@ class MatrixGameCrossAttention(WanSelfAttention):
        else:
            k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d)
            v = self.to_v(context)[0].view(b, -1, n, d)
-        
+
        # compute attention
        x = self.attn(q, k, v)

@@ -109,19 +133,20 @@ class MatrixGameCrossAttention(WanSelfAttention):


 class MatrixGameTransformerBlock(nn.Module):
-    
-    def __init__(self,
-                 dim: int,
-                 ffn_dim: int,
-                 num_heads: int,
-                 qk_norm: str = "rms_norm_across_heads",
-                 cross_attn_norm: bool = False,
-                 eps: float = 1e-6,
-                 added_kv_proj_dim: int | None = None,
-                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
-                 | None = None,
-                 prefix: str = "",
-                 action_config: dict | None = None):
+    def __init__(
+        self,
+        dim: int,
+        ffn_dim: int,
+        num_heads: int,
+        qk_norm: str = "rms_norm_across_heads",
+        cross_attn_norm: bool = False,
+        eps: float = 1e-6,
+        added_kv_proj_dim: int | None = None,
+        supported_attention_backends: tuple[AttentionBackendEnum, ...]
+        | None = None,
+        prefix: str = "",
+        action_config: dict | None = None,
+    ):
        super().__init__()
        action_config = action_config or {}

@@ -137,7 +162,8 @@ class MatrixGameTransformerBlock(nn.Module):
            head_size=dim // num_heads,
            causal=False,
            supported_attention_backends=supported_attention_backends,
-            prefix=f"{prefix}.attn1")
+            prefix=f"{prefix}.attn1",
+        )
        self.hidden_dim = dim
        self.num_attention_heads = num_heads
        dim_head = dim // num_heads
@@ -158,28 +184,28 @@ class MatrixGameTransformerBlock(nn.Module):
            eps=eps,
            elementwise_affine=True,
            dtype=torch.float32,
-            compute_dtype=torch.float32)
+            compute_dtype=torch.float32,
+        )

        # 2. Cross-attention
        if added_kv_proj_dim is not None:
            # I2V
-            self.attn2 = WanI2VCrossAttention(dim,
-                                              num_heads,
-                                              qk_norm=qk_norm,
-                                              eps=eps)
+            self.attn2 = WanI2VCrossAttention(
+                dim, num_heads, qk_norm=qk_norm, eps=eps
+            )
        else:
            # T2V
-            self.attn2 = WanT2VCrossAttention(dim,
-                                              num_heads,
-                                              qk_norm=qk_norm,
-                                              eps=eps)
+            self.attn2 = WanT2VCrossAttention(
+                dim, num_heads, qk_norm=qk_norm, eps=eps
+            )
        self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
            dim,
            norm_type="layer",
            eps=eps,
            elementwise_affine=False,
            dtype=torch.float32,
-            compute_dtype=torch.float32)
+            compute_dtype=torch.float32,
+        )

        # 2.1. Action Module Integration
        self.use_action_module = len(action_config) > 0
@@ -201,7 +227,7 @@ class MatrixGameTransformerBlock(nn.Module):
        temb: torch.Tensor,
        freqs_cis: tuple[torch.Tensor, torch.Tensor],
        # Action Module specific args
-        grid_sizes: torch.Tensor | None = None,
+        grid_sizes: torch.Tensor,
        mouse_cond: torch.Tensor | None = None,
        keyboard_cond: torch.Tensor | None = None,
    ) -> torch.Tensor:
@@ -211,9 +237,16 @@ class MatrixGameTransformerBlock(nn.Module):

        if temb.dim() == 4:
            # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v)
-            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-                self.scale_shift_table.unsqueeze(0) + temb.float()
-            ).chunk(6, dim=2)
+            (
+                shift_msa,
+                scale_msa,
+                gate_msa,
+                c_shift_msa,
+                c_scale_msa,
+                c_gate_msa,
+            ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(
+                6, dim=2
+            )
            # batch_size, seq_len, 1, inner_dim
            shift_msa = shift_msa.squeeze(2)
            scale_msa = scale_msa.squeeze(2)
@@ -224,13 +257,20 @@ class MatrixGameTransformerBlock(nn.Module):
        else:
            # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B)
            e = self.scale_shift_table + temb.float()
-            shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
-                6, dim=1)
+            (
+                shift_msa,
+                scale_msa,
+                gate_msa,
+                c_shift_msa,
+                c_scale_msa,
+                c_gate_msa,
+            ) = e.chunk(6, dim=1)
        assert shift_msa.dtype == torch.float32

        # 1. Self-attention
-        norm_hidden_states = (self.norm1(hidden_states.float()) *
-                              (1 + scale_msa) + shift_msa).to(orig_dtype)
+        norm_hidden_states = (
+            self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa
+        ).to(orig_dtype)
        query, _ = self.to_q(norm_hidden_states)
        key, _ = self.to_k(norm_hidden_states)
        value, _ = self.to_v(norm_hidden_states)
@@ -246,9 +286,10 @@ class MatrixGameTransformerBlock(nn.Module):

        # Apply rotary embeddings
        cos, sin = freqs_cis
-        query, key = _apply_rotary_emb(query, cos, sin,
-                                       is_neox_style=False), _apply_rotary_emb(
-                                           key, cos, sin, is_neox_style=False)
+        query, key = (
+            _apply_rotary_emb(query, cos, sin, is_neox_style=False),
+            _apply_rotary_emb(key, cos, sin, is_neox_style=False),
+        )

        attn_output, _ = self.attn1(query, key, value)
        attn_output = attn_output.flatten(2)
@@ -257,18 +298,24 @@ class MatrixGameTransformerBlock(nn.Module):

        null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
-            hidden_states, attn_output, gate_msa, null_shift, null_scale)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
+            hidden_states, attn_output, gate_msa, null_shift, null_scale
+        )
+        norm_hidden_states, hidden_states = (
+            norm_hidden_states.to(orig_dtype),
+            hidden_states.to(orig_dtype),
+        )

        # 2. Cross-attention
-        attn_output = self.attn2(norm_hidden_states,
-                                 context=encoder_hidden_states,
-                                 context_lens=None)
+        attn_output = self.attn2(
+            norm_hidden_states, context=encoder_hidden_states, context_lens=None
+        )
        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
-            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
-        norm_hidden_states, hidden_states = norm_hidden_states.to(
-            orig_dtype), hidden_states.to(orig_dtype)
+            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa
+        )
+        norm_hidden_states, hidden_states = (
+            norm_hidden_states.to(orig_dtype),
+            hidden_states.to(orig_dtype),
+        )

        # ================= Action Module =================
        if self.action_model is not None:
@@ -277,9 +324,12 @@ class MatrixGameTransformerBlock(nn.Module):
                # ActionModule implementation takes hidden_states directly
                hidden_states = self.action_model(
                    hidden_states,
-                    grid_sizes[0], grid_sizes[1], grid_sizes[2],
-                    mouse_cond, keyboard_cond,
-                    num_frame_per_block=grid_sizes[0],
+                    int(grid_sizes[0]),
+                    int(grid_sizes[1]),
+                    int(grid_sizes[2]),
+                    mouse_cond,
+                    keyboard_cond,
+                    num_frame_per_block=int(grid_sizes[0]),
                )
        # =================================================

@@ -290,6 +340,7 @@ class MatrixGameTransformerBlock(nn.Module):

        return hidden_states

+
 _DEFAULT_MATRIXGAME_CONFIG = MatrixGameWanVideoConfig()


@@ -299,27 +350,38 @@ class MatrixGameWanModel(BaseDiT):

    _fsdp_shard_conditions = _DEFAULT_MATRIXGAME_CONFIG._fsdp_shard_conditions
    _compile_conditions = _DEFAULT_MATRIXGAME_CONFIG._compile_conditions
-    _supported_attention_backends = _DEFAULT_MATRIXGAME_CONFIG._supported_attention_backends
+    _supported_attention_backends = (
+        _DEFAULT_MATRIXGAME_CONFIG._supported_attention_backends
+    )
    param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.param_names_mapping
-    reverse_param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.reverse_param_names_mapping
-    lora_param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.lora_param_names_mapping
+    reverse_param_names_mapping = (
+        _DEFAULT_MATRIXGAME_CONFIG.reverse_param_names_mapping
+    )
+    lora_param_names_mapping = (
+        _DEFAULT_MATRIXGAME_CONFIG.lora_param_names_mapping
+    )

-    def __init__(self,
-                 config: MatrixGameWanVideoConfig,
-                 hf_config: dict[str, Any],
-                 **kwargs) -> None:
+    def __init__(
+        self,
+        config: MatrixGameWanVideoConfig,
+        hf_config: dict[str, Any],
+        **kwargs,
+    ) -> None:
        super().__init__(config=config, hf_config=hf_config)

        inner_dim = config.num_attention_heads * config.attention_head_dim
        self.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
+        self.num_channels_latents = config.out_channels
        self.patch_size = config.patch_size

        # 1. Patch & position embedding
-        self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
-                                          embed_dim=inner_dim,
-                                          patch_size=config.patch_size,
-                                          flatten=False)
+        self.patch_embedding = PatchEmbed(
+            in_chans=config.in_channels,
+            embed_dim=inner_dim,
+            patch_size=config.patch_size,
+            flatten=False,
+        )

        # 2. Condition embeddings
        self.condition_embedder = MatrixGameTimeImageEmbedding(
@@ -329,57 +391,73 @@ class MatrixGameWanModel(BaseDiT):
        )

        # 2.1. Get action config
-        self.action_config = getattr(config, 'action_config', {})
+        self.action_config = getattr(config, "action_config", {})

        # 3. Transformer blocks
-        self.blocks = nn.ModuleList([
-            MatrixGameTransformerBlock(
-                inner_dim,
-                config.ffn_dim,
-                config.num_attention_heads,
-                config.qk_norm,
-                config.cross_attn_norm,
-                config.eps,
-                config.added_kv_proj_dim,
-                self._supported_attention_backends,
-                prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}",
-                action_config=self.action_config)
-            for i in range(config.num_layers)
-        ])
+        self.blocks = nn.ModuleList(
+            [
+                MatrixGameTransformerBlock(
+                    inner_dim,
+                    config.ffn_dim,
+                    config.num_attention_heads,
+                    config.qk_norm,
+                    config.cross_attn_norm,
+                    config.eps,
+                    config.added_kv_proj_dim,
+                    self._supported_attention_backends,
+                    prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}",
+                    action_config=self.action_config,
+                )
+                for i in range(config.num_layers)
+            ]
+        )

        # 4. Output norm & projection
-        self.norm_out = LayerNormScaleShift(inner_dim,
-                                            norm_type="layer",
-                                            eps=config.eps,
-                                            elementwise_affine=False,
-                                            dtype=torch.float32,
-                                            compute_dtype=torch.float32)
+        self.norm_out = LayerNormScaleShift(
+            inner_dim,
+            norm_type="layer",
+            eps=config.eps,
+            elementwise_affine=False,
+            dtype=torch.float32,
+            compute_dtype=torch.float32,
+        )
        self.proj_out = nn.Linear(
-            inner_dim, config.out_channels * math.prod(config.patch_size))
+            inner_dim, config.out_channels * math.prod(config.patch_size)
+        )
        self.scale_shift_table = nn.Parameter(
-            torch.randn(1, 2, inner_dim) / inner_dim**0.5)
- 
+            torch.randn(1, 2, inner_dim) / inner_dim**0.5
+        )
+
        self.gradient_checkpointing = False

-    def forward(self,
-                hidden_states: torch.Tensor,
-                encoder_hidden_states: torch.Tensor | list[torch.Tensor],
-                timestep: torch.LongTensor,
-                encoder_hidden_states_image: torch.Tensor
-                | list[torch.Tensor] | None = None,
-                # Action inputs
-                mouse_cond: torch.Tensor | None = None,
-                keyboard_cond: torch.Tensor | None = None,
-                **kwargs) -> torch.Tensor:
-        if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: torch.Tensor
+        | list[torch.Tensor]
+        | None = None,
+        # Action inputs
+        mouse_cond: torch.Tensor | None = None,
+        keyboard_cond: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        if encoder_hidden_states is not None and not isinstance(
+            encoder_hidden_states, torch.Tensor
+        ):
            encoder_hidden_states = encoder_hidden_states[0]
-        if isinstance(encoder_hidden_states_image,
-                      list) and len(encoder_hidden_states_image) > 0:
+        if (
+            isinstance(encoder_hidden_states_image, list)
+            and len(encoder_hidden_states_image) > 0
+        ):
            encoder_hidden_states_image = encoder_hidden_states_image[0]
        else:
            encoder_hidden_states_image = None

-        batch_size, num_channels, num_frames, height, width = hidden_states.shape
+        batch_size, num_channels, num_frames, height, width = (
+            hidden_states.shape
+        )
        p_t, p_h, p_w = self.patch_size
        post_patch_num_frames = num_frames // p_t
        post_patch_height = height // p_h
@@ -389,17 +467,25 @@ class MatrixGameWanModel(BaseDiT):
        d = self.hidden_size // self.num_attention_heads
        rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
        freqs_cos, freqs_sin = get_rotary_pos_embed(
-            (post_patch_num_frames * get_sp_world_size(), post_patch_height,
-             post_patch_width),
+            (
+                post_patch_num_frames * get_sp_world_size(),
+                post_patch_height,
+                post_patch_width,
+            ),
            self.hidden_size,
            self.num_attention_heads,
            rope_dim_list,
            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
-            rope_theta=10000)
+            rope_theta=10000,
+            do_sp_sharding=True,
+        )
        freqs_cos = freqs_cos.to(hidden_states.device)
        freqs_sin = freqs_sin.to(hidden_states.device)
-        freqs_cis = (freqs_cos.float(),
-                     freqs_sin.float()) if freqs_cos is not None else None
+        freqs_cis = (
+            (freqs_cos.float(), freqs_sin.float())
+            if freqs_cos is not None
+            else None
+        )

        hidden_states = self.patch_embedding(hidden_states)
        hidden_states = hidden_states.flatten(2).transpose(1, 2)
@@ -407,8 +493,14 @@ class MatrixGameWanModel(BaseDiT):
        if timestep.dim() == 2:
            timestep = timestep.flatten()

-        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep, encoder_hidden_states, encoder_hidden_states_image)
+        (
+            temb,
+            timestep_proj,
+            encoder_hidden_states,
+            encoder_hidden_states_image,
+        ) = self.condition_embedder(
+            timestep, encoder_hidden_states, encoder_hidden_states_image
+        )
        timestep_proj = timestep_proj.unflatten(1, (6, -1))

        if encoder_hidden_states is not None:
@@ -424,25 +516,30 @@ class MatrixGameWanModel(BaseDiT):
        if encoder_hidden_states_image is not None:
            if encoder_hidden_states is not None:
                encoder_hidden_states = torch.concat(
-                    [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+                    [encoder_hidden_states_image, encoder_hidden_states], dim=1
+                )
            else:
                encoder_hidden_states = encoder_hidden_states_image

        # This is [F, H, W] for the ActionModule
-        grid_sizes = torch.tensor([
-            post_patch_num_frames, post_patch_height, post_patch_width
-        ],
-                                 device=hidden_states.device)
+        grid_sizes = torch.tensor(
+            [post_patch_num_frames, post_patch_height, post_patch_width],
+            device=hidden_states.device,
+        )

        # Blocks
        for block in self.blocks:
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                hidden_states = self._gradient_checkpointing_func(
-                    block, hidden_states, encoder_hidden_states, timestep_proj,
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    timestep_proj,
                    freqs_cis,
                    grid_sizes=grid_sizes,
                    mouse_cond=mouse_cond,
-                    keyboard_cond=keyboard_cond)
+                    keyboard_cond=keyboard_cond,
+                )
            else:
                hidden_states = block(
                    hidden_states,
@@ -451,19 +548,27 @@ class MatrixGameWanModel(BaseDiT):
                    freqs_cis,
                    grid_sizes=grid_sizes,
                    mouse_cond=mouse_cond,
-                    keyboard_cond=keyboard_cond)
+                    keyboard_cond=keyboard_cond,
+                )

        # Output
-        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
-                                                                          dim=1)
+        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(
+            2, dim=1
+        )
        hidden_states = self.norm_out(hidden_states, shift, scale)
        hidden_states = self.proj_out(hidden_states)

-        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
-                                              post_patch_height,
-                                              post_patch_width, p_t, p_h, p_w,
-                                              -1)
+        hidden_states = hidden_states.reshape(
+            batch_size,
+            post_patch_num_frames,
+            post_patch_height,
+            post_patch_width,
+            p_t,
+            p_h,
+            p_w,
+            -1,
+        )
        hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

-        return output
+        return output
@@ -157,8 +157,8 @@ def load_initial_image(image_path: str = None) -> Image.Image:
    return Image.new("RGB", (640, 352), (128, 128, 128))

 def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None):
-    if keyboard_dim not in (2, 4, 7):
-        raise ValueError(f"keyboard_dim must be 2, 4, or 7, got {keyboard_dim}")
+    if keyboard_dim not in (2, 4, 6, 7):
+        raise ValueError(f"keyboard_dim must be 2, 4, 6, or 7, got {keyboard_dim}")
    if num_frames % 4 != 1:
        raise ValueError("Matrix-Game conditioning expects num_frames to be 4k+1.")

@@ -181,6 +181,11 @@ def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = No
        actions_double_action = []
        actions_single_camera = ["camera_l", "camera_r"]
        keyboard_idx = {"forward": 0, "back": 1}
+    elif keyboard_dim == 6:
+        actions_single_action = ["forward", "back", "left", "right"]
+        actions_double_action = ["forward_left", "forward_right"]
+        actions_single_camera = ["camera_l", "camera_r"]
+        keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3, "t1": 4, "t2": 5}
    else:  # keyboard_dim == 7
        # Temple Run model: still, w, s, left, right, a, d (no mouse)
        actions_single_action = ["forward", "back", "left", "right"]
@@ -685,3 +685,6 @@ class StepVideoModel(BaseDiT):

        output = rearrange(output, '(b f) c h w -> b c f h w', f=frame)
        return output
+
+# Entry point for model registry
+EntryClass = StepVideoModel
@@ -76,7 +76,7 @@ class WanTimeTextImageEmbedding(nn.Module):
                                 dim,
                                 dim,
                                 bias=True,
-                                 act_type="gelu_pytorch_tanh")
+                                 act_type="gelu_pytorch_tanh") if text_embed_dim > 0 else None

        self.image_embedder = None
        if image_embed_dim is not None:
@@ -92,7 +92,10 @@ class WanTimeTextImageEmbedding(nn.Module):
        temb = self.time_embedder(timestep, timestep_seq_len)
        timestep_proj = self.time_modulation(temb)

-        encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+        if self.text_embedder is not None:
+            encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+        else:
+            encoder_hidden_states = torch.zeros((timestep.shape[0], 0, temb.shape[-1]), device=temb.device, dtype=temb.dtype)
        if encoder_hidden_states_image is not None:
            assert self.image_embedder is not None
            encoder_hidden_states_image = self.image_embedder(
@@ -179,7 +182,10 @@ class WanT2VCrossAttention(WanSelfAttention):
            v = self.to_v(context)[0].view(b, -1, n, d)

        # compute attention
-        x = self.attn(q, k, v)
+        if k.size(1) > 0:
+            x = self.attn(q, k, v)
+        else:
+            x = torch.zeros_like(q)

        # output
        x = x.flatten(2)
@@ -227,7 +233,10 @@ class WanI2VCrossAttention(WanSelfAttention):
        v_img = self.add_v_proj(context_img)[0].view(b, -1, n, d)
        img_x = self.attn(q, k_img, v_img)
        # compute attention
-        x = self.attn(q, k, v)
+        if k.size(1) > 0:
+            x = self.attn(q, k, v)
+        else:
+            x = torch.zeros_like(q)

        # output
        x = x.flatten(2)
@@ -633,7 +642,7 @@ class WanTransformer3DModel(CachableDiT):
        enable_teacache = forward_batch is not None and forward_batch.enable_teacache

        orig_dtype = hidden_states.dtype
-        if not isinstance(encoder_hidden_states, torch.Tensor):
+        if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]
        if isinstance(encoder_hidden_states_image,
                      list) and len(encoder_hidden_states_image) > 0:
@@ -705,8 +714,11 @@ class WanTransformer3DModel(CachableDiT):
            timestep_proj = timestep_proj.unflatten(1, (6, -1))

        if encoder_hidden_states_image is not None:
-            encoder_hidden_states = torch.concat(
-                [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+            if encoder_hidden_states is not None:
+                encoder_hidden_states = torch.concat(
+                    [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+            else:
+                encoder_hidden_states = encoder_hidden_states_image

        if current_platform.is_mps() or current_platform.is_npu():
            encoder_hidden_states = encoder_hidden_states.to(orig_dtype)
@@ -859,4 +871,6 @@ class WanTransformer3DModel(CachableDiT):
            return hidden_states + self.previous_residual_even
        else:
            return hidden_states + self.previous_residual_odd
-        
+
+# Entry point for model registry
+EntryClass = WanTransformer3DModel
@@ -631,3 +631,6 @@ class CLIPVisionModel(ImageEncoder):
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params
+
+# Entry point for model registry
+EntryClass = [CLIPTextModel, CLIPVisionModel]
@@ -561,3 +561,6 @@ def _norm_and_concat_padded_batch(
    mask_flattened = mask.reshape(b, t, 1).expand(-1, -1, d * l)
    normed = normed.masked_fill(~mask_flattened, 0.0)
    return normed
+
+# Entry point for model registry
+EntryClass = LTX2GemmaTextEncoderModel
@@ -429,3 +429,6 @@ class LlamaModel(TextEncoder):
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params
+
+# Entry point for model registry
+EntryClass = LlamaModel
@@ -385,3 +385,6 @@ class Qwen2_5_VLTextModel(TextEncoder):
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params
+
+# Entry point for model registry
+EntryClass = Qwen2_5_VLTextModel
@@ -351,3 +351,5 @@ class Reason1TextEncoder(TextEncoder):
            tensor.std(dim=-1, keepdim=True) + 1e-8
        )

+# Entry point for model registry
+EntryClass = Reason1TextEncoder
@@ -417,3 +417,6 @@ class SiglipVisionModel(ImageEncoder):
            loaded_params.add(name)
            
        return loaded_params
+
+# Entry point for model registry
+EntryClass = SiglipVisionModel
@@ -583,3 +583,6 @@ class STEP1TextEncoder(torch.nn.Module):
                                      self.device) if with_mask else None)
            y_mask = txt_tokens.attention_mask
        return y.transpose(0, 1), y_mask
+
+# Entry point for model registry
+EntryClass = STEP1TextEncoder
@@ -719,3 +719,6 @@ class UMT5EncoderModel(TextEncoder):
                weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params
+
+# Entry point for model registry
+EntryClass = [UMT5EncoderModel, T5EncoderModel]
@@ -79,6 +79,7 @@ class ComponentLoader(ABC):
            "scheduler": (SchedulerLoader, "diffusers"),
            "transformer": (TransformerLoader, "diffusers"),
            "transformer_2": (TransformerLoader, "diffusers"),
+            "transformer_3": (TransformerLoader, "diffusers"),
            "vae": (VAELoader, "diffusers"),
            "audio_vae": (AudioDecoderLoader, "diffusers"),
            "audio_decoder": (AudioDecoderLoader, "diffusers"),
@@ -90,6 +91,8 @@ class ComponentLoader(ABC):
            "image_processor": (ImageProcessorLoader, "transformers"),
            "feature_extractor": (ImageProcessorLoader, "transformers"),
            "image_encoder": (ImageEncoderLoader, "transformers"),
+            "upsampler": (UpsamplerLoader, "diffusers"),
+            "upsampler_2": (UpsamplerLoader, "diffusers"),
        }

        if module_type in module_loaders:
@@ -391,8 +394,6 @@ class TextEncoderLoader(ComponentLoader):

            from fastvideo.platforms import current_platform

-            logger.info("Loading text encoder with cpu_offload: %s", use_cpu_offload)
-
            if use_cpu_offload:
                pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available()
                # Disable FSDP for MPS as it's not compatible
@@ -560,9 +561,11 @@ class VAELoader(ComponentLoader):
    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
        """Load the VAE based on the model path, and inference args."""
        config = get_diffusers_config(model=model_path)
-        config.pop("_name_or_path", None)
        class_name = config.pop("_class_name")
-        assert class_name is not None, "Model config does not contain a _class_name attribute. Only diffusers format is supported."
+        config.pop("_name_or_path", None)
+        assert class_name is not None, (
+            "Model config does not contain a _class_name attribute. Only diffusers format is supported."
+        )
        fastvideo_args.model_paths["vae"] = model_path

        from fastvideo.platforms import current_platform
@@ -730,9 +733,9 @@ class TransformerLoader(ComponentLoader):
    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
        """Load the transformer based on the model path, and inference args."""
        config = get_diffusers_config(model=model_path)
-        config.pop("_name_or_path", None)
        hf_config = deepcopy(config)
        cls_name = config.pop("_class_name")
+        config.pop("_name_or_path", None)
        if cls_name is None:
            raise ValueError(
                "Model config does not contain a _class_name attribute. "
@@ -747,7 +750,7 @@ class TransformerLoader(ComponentLoader):
        fastvideo_args.model_paths["transformer"] = model_path

        # Config from Diffusers supersedes fastvideo's model config
-        dit_config = fastvideo_args.pipeline_config.dit_config
+        dit_config = deepcopy(fastvideo_args.pipeline_config.dit_config)
        dit_config.update_model_arch(config)

        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
@@ -831,7 +834,7 @@ class TransformerLoader(ComponentLoader):
        )

        total_params = sum(p.numel() for p in model.parameters())
-        logger.info("Loaded model with %.2fB parameters, with cpu_offload: %s", total_params / 1e9, fastvideo_args.dit_cpu_offload)
+        logger.info("Loaded model with %.2fB parameters", total_params / 1e9)

        assert next(model.parameters()).dtype == default_dtype, (
            "Model dtype does not match default dtype"
@@ -879,6 +882,50 @@ class SchedulerLoader(ComponentLoader):
        return scheduler


+class UpsamplerLoader(ComponentLoader):
+    """Loader for upsamplers."""
+
+    def load(self, model_path: str, fastvideo_args: FastVideoArgs):
+        """Load the upsampler based on the model path, and inference args."""
+        config_dict = get_diffusers_config(model=model_path)
+        class_name = config_dict.pop("_class_name", None)
+
+        if class_name is None:
+            raise ValueError(
+                "Model config does not contain a _class_name attribute. "
+                "Only diffusers format is supported."
+            )
+
+        try:
+            upsampler_cfg = deepcopy(fastvideo_args.pipeline_config.upsampler_config[0])
+            upsampler_cfg.update_model_config(config_dict)
+        except Exception as e:
+            upsampler_cfg = deepcopy(fastvideo_args.pipeline_config.upsampler_config[1])
+            upsampler_cfg.update_model_config(config_dict)
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+        model = model_cls(upsampler_cfg)
+
+        target_device = get_local_torch_device()
+        model = model.to(target_device, dtype=PRECISION_TO_TYPE[fastvideo_args.pipeline_config.upsampler_precision])
+
+        # Find all safetensors files
+        safetensors_list = glob.glob(
+            os.path.join(str(model_path), "*.safetensors"))
+        if not safetensors_list:
+            raise ValueError(f"No safetensors files found in {model_path}")
+        
+        if len(safetensors_list) == 1:
+            loaded = safetensors_load_file(safetensors_list[0])
+        else:
+            loaded = {}
+            for sf_file in safetensors_list:
+                loaded.update(safetensors_load_file(sf_file))
+        
+        model.load_state_dict(loaded, strict=True)
+
+        return model.eval()
+
 class GenericComponentLoader(ComponentLoader):
    """Generic loader for components that don't have a specific loader."""

@@ -78,10 +78,9 @@ def hf_to_custom_state_dict(
    for source_param_name, full_tensor in hf_param_sd:  # type: ignore
        target_param_name, merge_index, num_params_to_merge = param_names_mapping(
            source_param_name)
-        if merge_index is None:
-            reverse_param_names_mapping[target_param_name] = (source_param_name,
-                                                            merge_index,
-                                                            num_params_to_merge)
+        reverse_param_names_mapping[target_param_name] = (source_param_name,
+                                                          merge_index,
+                                                          num_params_to_merge)
        if merge_index is not None:
            to_merge_params[target_param_name][merge_index] = full_tensor
            if len(to_merge_params[target_param_name]) == num_params_to_merge:
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/models/registry.py

+import ast
 import importlib
 import os
 import pickle
@@ -28,8 +29,6 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
    ("dits", "hunyuanvideo15", "HunyuanVideo15Transformer3DModel"),
    "HYWorldTransformer3DModel":
    ("dits", "hyworld", "HYWorldTransformer3DModel"),
-    "CausalHunyuanVideo15Transformer3DModel":
-    ("dits", "causal_hunyuanvideo15", "CausalHunyuanVideo15Transformer3DModel"),
    "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
    "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
    "StepVideoModel": ("dits", "stepvideo", "StepVideoModel"),
@@ -44,8 +43,8 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
    # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
    "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
    "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
-    "MatrixGameWanModel": ("dits", "matrix_game", "MatrixGameWanModel"),
-    "CausalMatrixGameWanModel": ("dits", "matrix_game", "CausalMatrixGameWanModel"),
+    "MatrixGameWanModel": ("dits", "matrixgame", "MatrixGameWanModel"),
+    "CausalMatrixGameWanModel": ("dits", "matrixgame", "CausalMatrixGameWanModel"),
 }

 _TEXT_ENCODER_MODELS = {
@@ -100,7 +99,12 @@ _SCHEDULERS = {
    ("schedulers", "scheduling_rcm", "RCMScheduler"),
 }

-_FAST_VIDEO_MODELS = {
+_UPSAMPLERS = {
+    "SRTo720pUpsampler": ("upsamplers", "hunyuan15", "SRTo720pUpsampler"),
+    "SRTo1080pUpsampler": ("upsamplers", "hunyuan15", "SRTo1080pUpsampler"),
+}
+
+_LEGACY_FAST_VIDEO_MODELS = {
    **_TEXT_TO_VIDEO_DIT_MODELS,
    **_IMAGE_TO_VIDEO_DIT_MODELS,
    **_TEXT_ENCODER_MODELS,
@@ -108,8 +112,102 @@ _FAST_VIDEO_MODELS = {
    **_VAE_MODELS,
    **_AUDIO_MODELS,
    **_SCHEDULERS,
+    **_UPSAMPLERS,
 }

+MODELS_PATH = os.path.dirname(__file__)
+
+
+@lru_cache(maxsize=None)
+def _discover_and_register_models() -> dict[str, tuple[str, str, str]]:
+    discovered_models: dict[str, tuple[str, str, str]] = {}
+    for root, dirs, files in os.walk(MODELS_PATH):
+        dirs[:] = [
+            d for d in dirs
+            if not d.startswith(".") and d != "__pycache__"
+        ]
+
+        for filename in files:
+            if not filename.endswith(".py"):
+                continue
+
+            filepath = os.path.join(root, filename)
+            try:
+                with open(filepath, "r", encoding="utf-8") as f:
+                    source = f.read()
+                tree = ast.parse(source, filename=filename)
+
+                entry_class_node = None
+                first_class_def = None
+
+                for node in ast.walk(tree):
+                    if isinstance(node, ast.Assign):
+                        for target in node.targets:
+                            if isinstance(target, ast.Name) and target.id == "EntryClass":
+                                entry_class_node = node
+                                break
+                    if first_class_def is None and isinstance(node, ast.ClassDef):
+                        first_class_def = node
+
+                if not entry_class_node or not first_class_def:
+                    continue
+
+                model_cls_name_list: list[str] = []
+                value_node = entry_class_node.value
+
+                if isinstance(value_node, ast.Name):
+                    model_cls_name_list.append(value_node.id)
+                elif isinstance(value_node, (ast.List, ast.Tuple)):
+                    for elt in value_node.elts:
+                        if isinstance(elt, ast.Constant) and isinstance(
+                                elt.value, str):
+                            model_cls_name_list.append(elt.value)
+                        elif isinstance(elt, ast.Name):
+                            model_cls_name_list.append(elt.id)
+
+                if not model_cls_name_list:
+                    continue
+
+                rel_dir = os.path.relpath(root, MODELS_PATH)
+                if rel_dir == ".":
+                    continue
+
+                rel_parts = rel_dir.split(os.sep)
+                component_name = rel_parts[0]
+                sub_parts = rel_parts[1:]
+
+                if filename == "__init__.py":
+                    mod_relname = ".".join(sub_parts)
+                else:
+                    mod_base = filename[:-3]
+                    mod_relname = ".".join(sub_parts +
+                                           [mod_base]) if sub_parts else mod_base
+
+                for model_cls_str in model_cls_name_list:
+                    if model_cls_str in discovered_models:
+                        logger.warning(
+                            "Duplicate architecture found: %s. Overwriting.",
+                            model_cls_str)
+                    discovered_models[model_cls_str] = (
+                        component_name,
+                        mod_relname,
+                        model_cls_str,
+                    )
+
+            except Exception as e:
+                logger.warning("Could not parse %s to find models: %s",
+                               filepath, e)
+
+    return discovered_models
+
+
+_DISCOVERED_MODELS = _discover_and_register_models()
+_FAST_VIDEO_MODELS = dict(_DISCOVERED_MODELS)
+for model_arch, spec in _LEGACY_FAST_VIDEO_MODELS.items():
+    if model_arch in _FAST_VIDEO_MODELS:
+        continue
+    _FAST_VIDEO_MODELS[model_arch] = spec
+
 _SUBPROCESS_COMMAND = [sys.executable, "-m", "fastvideo.models.dits.registry"]

 _T = TypeVar("_T")
@@ -341,7 +439,8 @@ class _ModelRegistry:
 ModelRegistry = _ModelRegistry({
    model_arch:
    _LazyRegisteredModel(
-        module_name=f"fastvideo.models.{component_name}.{mod_relname}",
+        module_name=(f"fastvideo.models.{component_name}.{mod_relname}"
+                     if mod_relname else f"fastvideo.models.{component_name}"),
        component_name=component_name,
        class_name=cls_name,
    )
@@ -155,8 +155,8 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,

        self.sigmas = sigmas.to(
            "cpu")  # to avoid too much CPU/GPU communication
-        self.sigma_min = sigma_min if sigma_min is not None else self.sigmas[-1].item()
-        self.sigma_max = sigma_max if sigma_max is not None else self.sigmas[0].item()
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()

        BaseScheduler.__init__(self)

@@ -289,7 +289,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
        sigmas: list[float] | None = None,
        mu: float | None = None,
        timesteps: list[float] | None = None,
-        extra_one_step: bool = False,
    ) -> None:
        """
        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
@@ -351,10 +350,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
            if timesteps_array is None:
                t_max = self._sigma_to_t(self.sigma_max)
                t_min = self._sigma_to_t(self.sigma_min)
-                if extra_one_step:
-                    timesteps_array = np.linspace(t_max, t_min, num_inference_steps + 1)[:-1]
-                else:
-                    timesteps_array = np.linspace(t_max, t_min, num_inference_steps)
+                timesteps_array = np.linspace(t_max, t_min, num_inference_steps)
            sigmas_array = timesteps_array / self.config.num_train_timesteps
        else:
            sigmas_array = np.array(sigmas).astype(np.float32)
@@ -648,7 +644,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
        self,
        clean_latent: torch.Tensor,
        noise: torch.Tensor,
-        timestep: torch.Tensor,
+        timestep: torch.IntTensor,
    ) -> torch.Tensor:

        """
@@ -689,3 +685,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,

    def __len__(self) -> int:
        return self.config.num_train_timesteps
+
+# Entry point for model registry
+EntryClass = FlowMatchEulerDiscreteScheduler
@@ -853,3 +853,6 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):

    def __len__(self):
        return self.config.num_train_timesteps
+
+# Entry point for model registry
+EntryClass = FlowUniPCMultistepScheduler
@@ -321,3 +321,6 @@ class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):

    def __len__(self) -> int:
        return self.config.num_train_timesteps
+
+# Entry point for model registry
+EntryClass = RCMScheduler
@@ -160,4 +160,6 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):

    def set_shift(self, shift: float) -> None:
        self.shift = shift
-    
+
+# Entry point for model registry
+EntryClass = SelfForcingFlowMatchScheduler
@@ -1094,4 +1094,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
        return noisy_samples

    def __len__(self):
-        return self.config.num_train_timesteps
+        return self.config.num_train_timesteps
+
+# Entry point for model registry
+EntryClass = UniPCMultistepScheduler
@@ -0,0 +1,170 @@
+# Licensed under the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/main/LICENSE
+#
+# Unless and only to the extent required by applicable law, the Tencent Hunyuan works and any
+# output and results therefrom are provided "AS IS" without any express or implied warranties of
+# any kind including any warranties of title, merchantability, noninfringement, course of dealing,
+# usage of trade, or fitness for a particular purpose. You are solely responsible for determining the
+# appropriateness of using, reproducing, modifying, performing, displaying or distributing any of
+# the Tencent Hunyuan works or outputs and assume any and all risks associated with your or a
+# third party's use or distribution of any of the Tencent Hunyuan works or outputs and your exercise
+# of rights and permissions under this agreement.
+# See the License for the specific language governing permissions and limitations under the License.
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from enum import Enum
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from torch import Tensor
+
+from fastvideo.models.vaes.hunyuan15vae import (
+    HunyuanVideo15CausalConv3d,
+    HunyuanVideo15RMS_norm,
+)
+
+from fastvideo.layers.activation import get_act_fn
+from fastvideo.configs.models.upsamplers import SRTo720pUpsamplerConfig, SRTo1080pUpsamplerConfig
+
+class HunyuanVideo15ResnetBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int | None = None,
+        non_linearity: str = "swish",
+    ) -> None:
+        super().__init__()
+        out_channels = out_channels or in_channels
+
+        self.nonlinearity = get_act_fn(non_linearity)
+
+        self.norm1 = HunyuanVideo15RMS_norm(in_channels, images=False)
+        self.conv1 = HunyuanVideo15CausalConv3d(in_channels, out_channels, kernel_size=3)
+
+        self.norm2 = HunyuanVideo15RMS_norm(out_channels, images=False)
+        self.conv2 = HunyuanVideo15CausalConv3d(out_channels, out_channels, kernel_size=3)
+
+        self.nin_shortcut = None
+        if in_channels != out_channels:
+            self.nin_shortcut = nn.Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv1(hidden_states)
+
+        hidden_states = self.norm2(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.nin_shortcut is not None:
+            residual = self.nin_shortcut(residual)
+
+        return hidden_states + residual
+
+class SRResidualCausalBlock3D(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.block = nn.Sequential(
+            HunyuanVideo15CausalConv3d(channels, channels, kernel_size=3),
+            nn.SiLU(inplace=True),
+            HunyuanVideo15CausalConv3d(channels, channels, kernel_size=3),
+            nn.SiLU(inplace=True),
+            HunyuanVideo15CausalConv3d(channels, channels, kernel_size=3),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.block(x)
+
+
+class SRTo720pUpsampler(nn.Module):
+
+    def __init__(
+        self,
+        config: SRTo720pUpsamplerConfig,
+    ):
+        super().__init__()
+        self.in_conv = HunyuanVideo15CausalConv3d(config.in_channels, config.hidden_channels, kernel_size=3)
+        self.blocks = nn.ModuleList([SRResidualCausalBlock3D(config.hidden_channels) for _ in range(config.num_blocks)])
+        self.out_conv = HunyuanVideo15CausalConv3d(config.hidden_channels, config.out_channels, kernel_size=3)
+        self.global_residual = bool(config.global_residual)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        y = self.in_conv(x)
+        for blk in self.blocks:
+            y = blk(y)
+        y = self.out_conv(y)
+        if self.global_residual and (y.shape == residual.shape):
+            y = y + residual
+        return y
+
+
+class SRTo1080pUpsampler(nn.Module):
+
+    def __init__(
+        self,
+        config: SRTo1080pUpsamplerConfig,
+    ):
+        super().__init__()
+        self.num_res_blocks = config.num_res_blocks
+        self.block_out_channels = config.block_out_channels
+        self.z_channels = config.z_channels
+
+        block_in = config.block_out_channels[0]
+        self.conv_in = HunyuanVideo15CausalConv3d(config.z_channels, block_in, kernel_size=3)
+
+        self.up = nn.ModuleList()
+        for i_level, ch in enumerate(config.block_out_channels):
+            block = nn.ModuleList()
+            block_out = ch
+            for _ in range(self.num_res_blocks + 1):
+                block.append(HunyuanVideo15ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+
+            self.up.append(up)
+
+        self.norm_out = HunyuanVideo15RMS_norm(block_in, images=False)
+        self.conv_out = HunyuanVideo15CausalConv3d(block_in, config.out_channels, kernel_size=3)
+
+        self.gradient_checkpointing = False
+        self.is_residual = config.is_residual
+
+    def forward(self, z: Tensor, target_shape: Sequence[int] = None) -> Tensor:
+        """
+        Args:
+            z: (B, C, T, H, W)
+            target_shape: (H, W)
+        """
+        if target_shape is not None and z.shape[-2:] != target_shape:
+            bsz = z.shape[0]
+            z = rearrange(z, "b c f h w -> (b f) c h w")
+            z = F.interpolate(z, size=target_shape, mode="bilinear", align_corners=False)
+            z = rearrange(z, "(b f) c h w -> b c f h w", b=bsz)
+
+        # z to block_in
+        repeats = self.block_out_channels[0] // (self.z_channels)
+        h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1)
+
+        # upsampling
+        for i_level in range(len(self.block_out_channels)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+            if hasattr(self.up[i_level], "upsample"):
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = get_act_fn("swish")(h)
+        h = self.conv_out(h)
+        return h
@@ -5,9 +5,6 @@ from typing import Any

 import torch

-from fastvideo.logger import init_logger
-
-logger = init_logger(__name__)

 # TODO(PY): move it elsewhere
 def auto_attributes(init_func):
@@ -663,7 +663,7 @@ class AutoencoderKLHunyuanVideo15(nn.Module, ParallelTiledVAE):
        # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
        # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
        # intermediate tiles together, the memory requirement can be lowered.
-        self.use_tiling = True
+        self.use_tiling = False

        # The minimal tile height and width for spatial tiling to be used
        self.tile_sample_min_height = 256
@@ -700,4 +700,7 @@ class AutoencoderKLHunyuanVideo15(nn.Module, ParallelTiledVAE):
        else:
            z = posterior.mode()
        dec = self.decode(z)
-        return dec
+        return dec
+
+# Entry point for model registry
+EntryClass = AutoencoderKLHunyuanVideo15
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
JerryZhou54	8cb2ae9d27	ckpt	2026-02-09 05:50:48 +00:00
XOR-op	f76efe798e	[bugfix]: `_compile_conditions` regression (#1077 )	2026-02-06 18:57:43 -08:00
Zhang Peiyuan	09f455233e	[misc] readme small fix (#1076 )	2026-02-06 16:57:48 -08:00
XOR-op	aea300f690	[perf]: use CUDA IPC in multiproc executor to avoid serialization overhead (#1061 )	2026-02-06 19:43:18 -05:00
Hao Zhang	b92219f6a6	more fix and relocate STA arguments to pipeline config (#1073 )	2026-02-06 13:38:35 -08:00
Jinzhe Pan	a321b95a8a	[Fix] remove video ratio limitation (#1069 )	2026-02-05 20:33:08 -08:00
Wei Zhou	98308db7e0	[Feature] [Hy1.5] Support HY1.5 super-resolution pipeline for 1080p videos (#1046 )	2026-02-05 16:39:20 -08:00
Hao Zhang	c1e18f6722	Some minor fixes (#1068 )	2026-02-05 16:28:00 -08:00
William Lin	75e193a2c9	[core] Refactor and centralize our registry for models, pipelines, and sampling params (#1066 )	2026-02-05 14:40:30 -08:00
XOR-op	d6e0a7d0dd	[refactor] Action module (#1065 )	2026-02-05 13:56:35 -08:00
William Linandgemini-code-assist[bot]	7fc5f241da	[misc] Fix naming instruction in runpod.md (#1067 ) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2026-02-05 13:45:49 -08:00
Mingjia Huo	aae48a7e90	[feat] HYworld VAE with cache (#1057 )	2026-02-05 04:18:17 -08:00
William Lin	88f38eb0f4	[misc] upgrade torch to 2.10 (#1048 )	2026-02-05 04:15:22 -08:00
Shao Duan	d750b463dc	Added Sequence Parallelism for LTX-2 Distilled (#1036 )	2026-02-04 15:25:30 -08:00
KyleShaoandWill Lin	e10b26a3d8	[feat] Add Cosmos 2.5 I2W/V2W support (staged pipeline + examples) (#1021 ) Co-authored-by: Will Lin <wlsaidhi@gmail.com>	2026-02-04 14:19:30 -08:00
XOR-op	74636ba246	[chore]: use higher precision timestamp in logging (#1062 )	2026-02-04 13:22:59 -08:00
Wei Zhou	7e2f3f14e7	[Bugfix] [Wan I2V] Fix CLIP Image encoder config (#1063 )	2026-02-04 13:20:19 -08:00
Kaiqin Kong	caa1c402ba	[bugfix] Double Normalization in Preprocessing Dataset (#1055 )	2026-01-31 11:11:59 -08:00
XOR-op	38a6bd93d3	[chore]: update sageattn3 installation instructions (#1050 )	2026-01-29 15:45:52 -08:00
alexzms	b867ef7e7c	[SP Sharding] Fix SP loss sharding on token axis (thw) with padding; add distributed correctness tests (#1045 ) Fixes the sequence parallel sharding on t, now SP shards on thw	2026-01-27 22:53:04 -08:00
William Lin	3ae58c277a	[docs] Update design overview and add agents tutorial (#1044 )	2026-01-27 15:56:58 -08:00
Kaiqin Kong	0c6862ca55	[feature] Add Matrix Game 2.0 training (#1017 ) The CI tests are quite unstable, but since multiple CI tests indicates that each individual tests are passed, I think we can merge this.	2026-01-26 19:21:43 -08:00
XOR-opandWill Lin	e8c854bcf1	[docs] Offloading instruction (#1022 ) Co-authored-by: Will Lin <wlsaidhi@gmail.com>	2026-01-26 13:40:28 -08:00
William Lin	06860e96fe	[docs] Update runpod instructions (#1043 )	2026-01-26 13:13:54 -08:00