fix

Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>

.buildkite/pipeline.yml

@@ -117,11 +117,11 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
-                - "csrc/attn/vsa/**"
-                - "csrc/attn/tk/**"
-                - "csrc/attn/setup_vsa.py"
-                - "csrc/attn/config_vsa.py"
-                - "csrc/attn/vsa.cpp"
+                - "csrc/attn/video_sparse_attn/**"
+                - "csrc/attn/video_sparse_attn/tk/**"
+                - "csrc/attn/video_sparse_attn/setup.py"
+                - "csrc/attn/video_sparse_attn/config_vsa.py"
+                - "csrc/attn/video_sparse_attn/vsa.cpp"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -133,10 +133,10 @@ steps:
                   queue: "default"
             - path:
                 - "fastvideo/**"
-                - "csrc/attn/st_attn/**"
-                - "csrc/attn/setup_sta.py"
-                - "csrc/attn/config_sta.py"
-                - "csrc/attn/st_attn.cpp"
+                - "csrc/attn/sliding_tile_attn/**"
+                - "csrc/attn/sliding_tile_attn/setup.py"
+                - "csrc/attn/sliding_tile_attn/config_sta.py"
+                - "csrc/attn/sliding_tile_attn/st_attn.cpp"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -147,10 +147,10 @@ steps:
                 agents:
                   queue: "default"
             - path:
-                - "csrc/attn/st_attn/**"
-                - "csrc/attn/setup_sta.py"
-                - "csrc/attn/config_sta.py"
-                - "csrc/attn/st_attn.cpp"
+                - "csrc/attn/sliding_tile_attn/**"
+                - "csrc/attn/sliding_tile_attn/setup.py"
+                - "csrc/attn/sliding_tile_attn/config_sta.py"
+                - "csrc/attn/sliding_tile_attn/st_attn.cpp"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:
@@ -161,12 +161,12 @@ steps:
                 agents:
                   queue: "default"
             - path:
-                - "csrc/attn/vsa/**"
-                - "csrc/attn/tk/**"
+                - "csrc/attn/video_sparse_attn/**"
+                - "csrc/attn/video_sparse_attn/tk/**"
                 - "csrc/attn/tests/test_vsa.py"
-                - "csrc/attn/setup_vsa.py"
-                - "csrc/attn/config_vsa.py"
-                - "csrc/attn/vsa.cpp"
+                - "csrc/attn/video_sparse_attn/setup.py"
+                - "csrc/attn/video_sparse_attn/config_vsa.py"
+                - "csrc/attn/video_sparse_attn/vsa.cpp"
                 - "pyproject.toml"
                 - "docker/Dockerfile.python3.12"
               config:

.github/workflows/pr-test.yml

@@ -104,16 +104,17 @@ jobs:
               - 'pyproject.toml'
               - 'docker/Dockerfile.python3.12'
             sta-kernel-paths: &sta-kernel-paths
-              - 'csrc/attn/st_attn/**'
-              - 'csrc/attn/setup_sta.py'
-              - 'csrc/attn/config_sta.py'
-              - 'csrc/attn/st_attn.cpp'
+              - 'csrc/attn/sliding_tile_attn/**'
+              - 'csrc/attn/sliding_tile_attn/tk/**'
+              - 'csrc/attn/sliding_tile_attn/setup.py'
+              - 'csrc/attn/sliding_tile_attn/config_sta.py'
+              - 'csrc/attn/sliding_tile_attn/st_attn.cpp'
             vsa-kernel-paths: &vsa-kernel-paths
-              - 'csrc/attn/vsa/**'
-              - 'csrc/attn/tk/**'
-              - 'csrc/attn/setup_vsa.py'
-              - 'csrc/attn/config_vsa.py'
-              - 'csrc/attn/vsa.cpp'
+              - 'csrc/attn/video_sparse_attn/**'
+              - 'csrc/attn/video_sparse_attn/tk/**'
+              - 'csrc/attn/video_sparse_attn/setup.py'
+              - 'csrc/attn/video_sparse_attn/config_vsa.py'
+              - 'csrc/attn/video_sparse_attn/vsa.cpp'
             vsa-paths: &vsa-paths
               - 'fastvideo/**'
               - *common-paths

.github/workflows/sta-publish.yml

@@ -5,7 +5,7 @@ on:
     branches:
       - main
     paths:
-      - "csrc/attn/setup_sta.py"
+      - "csrc/attn/sliding_tile_attn/setup.py"
   workflow_dispatch:
 
 jobs:
@@ -23,13 +23,13 @@ jobs:
       - name: Check if version changed
         id: check-version
         run: |
-          cd csrc/attn
+          cd csrc/attn/sliding_tile_attn
           # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_sta.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup_sta.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
           echo "Old version: $OLD_VERSION"
           
           if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -144,13 +144,13 @@ jobs:
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/sliding_tile_attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_sta.py bdist_wheel --dist-dir=dist
+          python setup.py bdist_wheel --dist-dir=dist
 
       - name: Rename wheel file
         run: |
-          cd csrc/attn
+          cd csrc/attn/sliding_tile_attn
 
           CUDA_SHORT_VERSION=$(echo ${{ matrix.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
           TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-version }} | cut -d. -f1,2)
@@ -165,7 +165,7 @@ jobs:
         uses: actions/upload-artifact@v4
         with:
           name: ${{ env.wheel_name }}
-          path: csrc/attn/dist/*.whl
+          path: csrc/attn/sliding_tile_attn/dist/*.whl
           retention-days: 90
 
   publish_package:
@@ -239,11 +239,11 @@ jobs:
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/sliding_tile_attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_sta.py sdist --dist-dir=dist
+          python setup.py sdist --dist-dir=dist
 
       - name: Publish release distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: csrc/attn/dist/
+          packages-dir: csrc/attn/sliding_tile_attn/dist/

.github/workflows/vsa-publish.yml

@@ -5,7 +5,7 @@ on:
     branches:
       - main
     paths:
-      - "csrc/attn/setup_vsa.py"
+      - "csrc/attn/video_sparse_attn/setup.py"
   workflow_dispatch:
 
 jobs:
@@ -23,13 +23,13 @@ jobs:
       - name: Check if version changed
         id: check-version
         run: |
-          cd csrc/attn
+          cd csrc/attn/video_sparse_attn
           # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_vsa.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup_vsa.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
           echo "Old version: $OLD_VERSION"
           
           if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -152,13 +152,13 @@ jobs:
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/video_sparse_attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_vsa.py bdist_wheel --dist-dir=dist
+          python setup.py bdist_wheel --dist-dir=dist
 
       - name: Rename wheel file
         run: |
-          cd csrc/attn
+          cd csrc/attn/video_sparse_attn
 
           CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
           TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2)
@@ -173,7 +173,7 @@ jobs:
         uses: actions/upload-artifact@v4
         with:
           name: ${{ env.wheel_name }}
-          path: csrc/attn/dist/*.whl
+          path: csrc/attn/video_sparse_attn/dist/*.whl
           retention-days: 90
 
   publish_package:
@@ -247,11 +247,11 @@ jobs:
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/video_sparse_attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_vsa.py sdist --dist-dir=dist
+          python setup.py sdist --dist-dir=dist
 
       - name: Publish release distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: csrc/attn/dist/
+          packages-dir: csrc/attn/video_sparse_attn/dist/

.gitmodules

@@ -1,3 +1,7 @@
-[submodule "csrc/attn/tk"]
-	path = csrc/attn/tk
+[submodule "csrc/attn/video_sparse_attn/tk"]
+	path = csrc/attn/video_sparse_attn/tk
+	url = https://github.com/HazyResearch/ThunderKittens.git
+
+[submodule "csrc/attn/sliding_tile_attn/tk"]
+	path = csrc/attn/sliding_tile_attn/tk
 	url = https://github.com/HazyResearch/ThunderKittens.git

csrc/attn/README.md

@@ -25,9 +25,9 @@ sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave
 sudo apt update
 sudo apt install clang-11
 ```
-(If you use CUDA12.4)
+(If you use CUDA12.8)
 ```bash
-export CUDA_HOME=/usr/local/cuda-12.4
+export CUDA_HOME=/usr/local/cuda-12.8
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 ```

csrc/attn/scratch.sh

@@ -1,4 +0,0 @@
-    off_hz = tl.program_id(2)
-    b = off_hz // H
-    h = off_hz % H
-    meta_base = ((b * H + h) * q_tiles + q_blk)

csrc/attn/sliding_tile_attn/MANIFEST.in

@@ -1,2 +1,2 @@
 recursive-include tk *
-include config.py
+include config_sta.py

csrc/attn/sliding_tile_attn/README.md

@@ -0,0 +1,87 @@
+
+# Attention Kernel Used in FastVideo
+
+## Sliding Tile Attention (STA)
+We only support H100 for STA.
+
+### Installation
+```bash
+pip install st_attn
+``` 
+
+Install from source:
+
+```bash
+git submodule update --init --recursive
+python setup.py install
+```
+
+If you encounter error during installation, try below:
+Install C++20 for ThunderKittens:
+```bash
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+(If you use CUDA12.8)
+```bash
+export CUDA_HOME=/usr/local/cuda-12.8
+export PATH=${CUDA_HOME}/bin:${PATH} 
+export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+###  Usage
+End-2-end inference with FastVideo:
+```bash
+bash scripts/inference/v1_inference_wan_STA.sh
+```
+
+If you want to use sliding tile attention in your custom model:
+```python
+from st_attn import sliding_tile_attention
+# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
+# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256
+# a tile is a cube of size (6, 8, 8)
+# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video.
+# text_length: int ranging from 0 to 256
+# If your attention contains text token (Hunyuan)
+out = sliding_tile_attention(q, k, v, window_size, text_length)
+# If your attention does not contain text token (StepVideo)
+out = sliding_tile_attention(q, k, v, window_size, 0, False)
+```
+
+
+### Test
+```bash
+python ../tests/test_sta.py # test STA
+python ../tests/test_vsa.py # test VSA
+```
+### Benchmark
+```bash
+python ../benchmarks/bench_sta.py
+```
+
+
+### How Does STA Work?
+We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 
+
+
+https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747
+
+## Why is STA Fast?
+2D/3D Sliding Window Attention (SWA) creates many mixed blocks in the attention map. Even though mixed blocks have less output value,a mixed block is significantly slower than a dense block due to the GPU-unfriendly masking operation. 
+
+STA removes mixed blocks.
+
+
+<div align="center">
+<img src=../../../assets/sliding_tile_attn_map.png width="80%"/>
+</div>
+
+## Acknowledgement
+
+We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.

csrc/attn/sliding_tile_attn/setup.py

@@ -1,7 +1,7 @@
 import os
 import subprocess
 
-from csrc.attn.config_sta import kernels, sources, target
+from config_sta import kernels, sources, target
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
@@ -9,7 +9,7 @@ target = target.lower()
 
 # Package metadata
 PACKAGE_NAME = "st_attn"
-VERSION = "0.0.4"
+VERSION = "0.0.6"
 AUTHOR = "Hao AI Lab"
 DESCRIPTION = "Sliding Tile Atteniton Kernel Used in FastVideo"
 URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/sliding_tile_attention"

csrc/attn/video_sparse_attn/MANIFEST.in

@@ -0,0 +1,2 @@
+recursive-include tk *
+include config_vsa.py

csrc/attn/video_sparse_attn/README.md

@@ -0,0 +1,61 @@
+
+
+# Attention Kernel Used in FastVideo
+
+## Video Sparse Attention (VSA)
+
+### Installation
+We support H100 (via TK) and any other GPU (via triton) for VSA.
+
+```bash
+pip install vsa
+``` 
+
+Install from source:
+
+```bash
+git submodule update --init --recursive
+python setup.py install
+```
+
+
+If you encounter error during installation, try below:
+Install C++20 for ThunderKittens:
+```bash
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+(If you use CUDA12.8)
+```bash
+export CUDA_HOME=/usr/local/cuda-12.8
+export PATH=${CUDA_HOME}/bin:${PATH} 
+export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+### Verify if you have successfully installed
+
+```bash
+# test numerical
+python ../tests/test_vsa.py
+# (For H100) test speed
+python ../benchmarks/bench_vsa_hopper.py
+```
+
+bench_vsa_hopper.py should print something like this:
+
+```bash
+Using topk=76 kv blocks per q block (out of 768 total kv blocks)
+
+=== BLOCK SPARSE ATTENTION BENCHMARK ===
+Block Sparse Forward  - TFLOPS: 5622.26
+Block Sparse Backward - TFLOPS: 3865.68
+```
+
+## Acknowledgement
+
+We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.

csrc/attn/video_sparse_attn/setup.py

@@ -9,10 +9,10 @@ target = target.lower()
 
 # Package metadata
 PACKAGE_NAME = "vsa"
-VERSION = "0.0.1"
+VERSION = "0.0.3"
 AUTHOR = "Hao AI Lab"
 DESCRIPTION = "Video Sparse Attention Kernel Used in FastVideo"
-URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn"
+URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn/video_sparse_attn"
 
 # Set environment variables
 tk_root = os.getenv('THUNDERKITTENS_ROOT', os.path.abspath(os.path.join(os.getcwd(), 'tk/')))

docker/Dockerfile.python3.10

@@ -58,15 +58,15 @@ RUN source $HOME/.local/bin/env && \
 # Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/sliding_tile_attn && \
     git submodule update --init --recursive && \
-    python setup_sta.py install
+    python setup.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
     git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install
 
-EXPOSE 22
+EXPOSE 22

docker/Dockerfile.python3.11

@@ -58,15 +58,15 @@ RUN source $HOME/.local/bin/env && \
 # Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/sliding_tile_attn && \
     git submodule update --init --recursive && \
-    python setup_sta.py install
+    python setup.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
     git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install
 
-EXPOSE 22
+EXPOSE 22

docker/Dockerfile.python3.12

@@ -58,15 +58,15 @@ RUN source $HOME/.local/bin/env && \
 # Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/sliding_tile_attn && \
     git submodule update --init --recursive && \
-    python setup_sta.py install
+    python setup.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
     git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install
 
-EXPOSE 22
+EXPOSE 22

docker/Dockerfile.python3.12.cuda12.9.1

@@ -58,15 +58,15 @@ RUN source $HOME/.local/bin/env && \
 # Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/sliding_tile_attn && \
     git submodule update --init --recursive && \
-    python setup_sta.py install
+    python setup.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
     git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install
 
-EXPOSE 22
+EXPOSE 22

docs/source/sliding_tile_attention/installation.md

@@ -4,7 +4,7 @@
 You can install the Sliding Tile Attention package using
 
 ```
-pip install st_attn==0.0.4
+pip install st_attn
 ``` 
 
 # Building from Source
@@ -12,7 +12,6 @@ We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have impleme
 First, install C++20 for ThunderKittens:
 
 ```bash
-cd csrc/sliding_tile_attention/
 sudo apt update
 sudo apt install gcc-11 g++-11
 
@@ -22,14 +21,20 @@ sudo apt update
 sudo apt install clang-11
 ```
 
-Install STA:
+Set up CUDA environment (if using CUDA 12.4):
 
 ```bash
 export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+Install STA:
+
+```bash
+cd csrc/attn/sliding_tile_attn/
 git submodule update --init --recursive
-python setup_sta.py install
+python setup.py install
 ```
 
 # 🧪 Test

docs/source/video_sparse_attention/installation.md

@@ -4,8 +4,7 @@
 You can install the Video Sparse Attention package using
 
 ```bash
-git submodule update --init --recursive
-python setup_vsa.py install
+pip install vsa
 ```
 
 # Building from Source
@@ -34,9 +33,9 @@ export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 Install VSA:
 
 ```bash
-cd csrc/attn/
+cd csrc/attn/video_sparse_attn/
 git submodule update --init --recursive
-python setup_vsa.py install
+python setup.py install
 ```
 
 # 🧪 Test

examples/datasets/crush_smol/download_dataset.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"

examples/datasets/crush_smol/validation.json

@@ -0,0 +1,76 @@
+{
+  "data": [
+    {
+      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
+      "image_path": null,
+      "video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
+      "image_path": null,
+      "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
+      "image_path": null,
+      "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
+      "image_path": null,
+      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "The video shows a green and orange object being flattened as if it were under a hydraulic press, with the press moving down and compressing the object.",
+      "image_path": null,
+      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-056.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "The video shows a cylindrical object with a cityscape image being flattened as if it were under a hydraulic press. The object is placed on a metal platform, and a large, striped cylinder presses down on it, causing it to collapse and release a liquid inside. The background features a green wall with a yellow and red warning sign.",
+      "image_path": null,
+      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-059.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "The video shows a close-up of an orange being flattened as if it were under a hydraulic press, with the press moving down and compressing the fruit until it is completely flattened.",
+      "image_path": null,
+      "video_path": "validation_dataset/EJqsC21GSBY-Scene-059.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.",
+      "image_path": null,
+      "video_path": "validation_dataset/GBSfpTcKegk-Scene-003.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    }
+  ]
+}

examples/datasets/crush_smol/validation_single.json

@@ -0,0 +1,13 @@
+{
+  "data": [
+    {
+      "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
+      "image_path": null,
+      "video_path": "validation_dataset/1gGQy4nxyUo-Scene-016.mp4",
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    }
+  ]
+}

examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/README.md

@@ -4,9 +4,7 @@ These are end-to-end example scripts for distilling Wan2.1 T2V 1.3B model using
 ### 0. Make sure you have installed VSA
 
 ```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
+pip install vsa
 ```
 
 ### 1. Download dataset:

examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/README.md

@@ -4,9 +4,7 @@ These are end-to-end example scripts for distilling Wan2.2 TI2V 5B model DMD+VSA
 ### 0. Make sure you have installed VSA
 
 ```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
+pip install vsa
 ```
 
 ### Data-free Distillation

examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/README.md

@@ -4,9 +4,7 @@ These are end-to-end example scripts for distilling Wan2.2 TI2V 5B model DMD+VSA
 ### 0. Make sure you have installed VSA
 
 ```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
+pip install vsa
 ```
 
 ### 1. Download dataset:

examples/inference/basic/basic.py

@@ -1,6 +1,6 @@
 from fastvideo import VideoGenerator
 
-# from fastvideo.configs.sample import SamplingParam
+from fastvideo.configs.sample import SamplingParam
 
 OUTPUT_PATH = "video_samples"
 def main():
@@ -17,30 +17,29 @@ def main():
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,
         pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        ti2v_task=True,
         # image_encoder_cpu_offload=False,
     )
 
-    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
     # sampling_param.num_frames = 45
-    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    sampling_param.image_path = "test.jpg"
     # Generate videos with the same simple API, regardless of GPU count
     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."
+        "A girl is packing a suitcase when stuff suddently starts flying around the room."
     )
     video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
     # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
 
     # Generate another video with a different prompt, without reloading the
     # model!
-    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)
+    # 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)
 
 
 if __name__ == "__main__":

examples/inference/basic/basic_self_forcing_causal.py

@@ -0,0 +1,33 @@
+import os
+import time
+from fastvideo import VideoGenerator, SamplingParam
+
+OUTPUT_PATH = "video_samples_causal"
+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.
+    model_name = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=True,
+        text_encoder_cpu_offload=False,
+        dit_cpu_offload=False,
+    )
+
+    sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.num_frames = 81
+
+    prompts = [
+        "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
+        "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.",
+    ]
+
+    for prompt in prompts:
+        video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_wan2_2_ti2v.py

@@ -0,0 +1,41 @@
+from fastvideo import VideoGenerator
+
+OUTPUT_PATH = "video_samples_wan2_2_5B_ti2v"
+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.
+    model_name = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=True,
+        vae_cpu_offload=False,
+        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,
+    )
+
+    # I2V is triggered just by passing in an image_path argument
+    prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
+    image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path)
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+
+    # T2V mode
+    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)
+
+
+if __name__ == "__main__":
+    main()

examples/training/consistency_finetune/ode_init/finetune_ode_init.sh

@@ -0,0 +1,99 @@
+#!/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="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+# DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-extended-t2v-1-3b/"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=1
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_70k"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "fixed_wan_ode_init_70k_6e-6"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --max_train_steps 6000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --warp_denoising_step
+  # --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim 1
+)
+
+# 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 50
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 6e-6
+  --mixed_precision "bf16"
+  --checkpointing_steps 1000
+  --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
+  --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
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

examples/training/consistency_finetune/ode_init/finetune_ode_init_16k_vidprom.slurm

@@ -0,0 +1,132 @@
+#!/bin/bash
+#SBATCH --job-name=4e6B8_16kFV_no_warp_ode_vidprom
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=ode_vidprom16k_no_warp/ode_vidprom8b16k_4e-6.out
+#SBATCH --error=ode_vidprom16k_no_warp/ode_vidprom8b16k_4e-6.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate will-fv2
+
+export WANDB_MODE="online"
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# 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 CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+
+# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing-16k-t2v-1-3b/"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
+NUM_GPUS=2
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "vidprom_8b16k_test_no_warp_4e-6"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "vidprom_8b16k_wan_ode_init_4e-6"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --max_train_steps 6001
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --dmd_denoising_steps "1000,750,500,0"
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim $NUM_GPUS
+  --hsdp_shard_dim 1
+)
+
+# 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 50
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 4e-6
+  --mixed_precision "bf16"
+  --checkpointing_steps 500
+  --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
+  # --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[@]}"

examples/training/consistency_finetune/ode_init/finetune_ode_init_2k_vidprom.slurm

@@ -0,0 +1,131 @@
+#!/bin/bash
+#SBATCH --job-name=ode_vidprom2k
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=ode_vidprom2k_output/ode_vidprom2k.out
+#SBATCH --error=ode_vidprom2k_output/ode_vidprom2k.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate will-fv2
+
+export WANDB_MODE="online"
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# 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 CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+
+# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="/mnt/sharefs/users/hao.zhang/klin/preproc/data/test-ode-preprocessing/"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
+NUM_GPUS=8
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_vidprom2k"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "vidprom2k_wan_ode_init_5e-6"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --max_train_steps 6001
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  # --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 8
+  --hsdp_shard_dim 1
+)
+
+# 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 "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 5e-6
+  --mixed_precision "bf16"
+  --checkpointing_steps 2000
+  --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
+  --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[@]}"

examples/training/consistency_finetune/ode_init/finetune_ode_init_crush.slurm

@@ -0,0 +1,132 @@
+#!/bin/bash
+#SBATCH --job-name=ode_crush
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=ode_crush_output/ode_crush.out
+#SBATCH --error=ode_crush_output/ode_crush.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate will-fv2
+
+export WANDB_MODE="online"
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+# 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 CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
+export TOKENIZERS_PARALLELISM=false
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+
+# MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_5/combined_parquet_dataset/"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/ode_init/validation.json"
+NUM_GPUS=2
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_warp_2"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "2warp_fixed_wan_ode_init_5e-6"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  # --warp_denoising_step
+  --max_train_steps 6001
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  # --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim $NUM_GPUS
+  --hsdp_shard_dim 1
+)
+
+# 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 20
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 5e-6
+  --mixed_precision "bf16"
+  --checkpointing_steps 2000
+  --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
+  --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[@]}"

examples/training/consistency_finetune/ode_init/finetune_ode_init_overfit.sh

@@ -0,0 +1,98 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+export WANDB_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="/mnt/weka/home/hao.zhang/wl/FastVideo2/data/crush-smol_processed_t2v_1_3b_ode_init_single"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=1
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_crush_smol"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "overfitwan_ode_init_crush_smol"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --max_train_steps 2001
+  # --warp_denoising_step
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  # --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim 1
+)
+
+# 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 20
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --checkpointing_steps 500
+  --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
+  --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
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

examples/training/consistency_finetune/ode_init/preprocess_wan_data_ode.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_TYPE="wan"
+DATA_MERGE_PATH="data/crush-smol_single/merge.txt"
+OUTPUT_DIR="data/crush-smol_processed_t2v_1_3b_ode_init_single/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 1 \
+    --seed 42 \
+    --max_height 480 \
+    --max_width 832 \
+    --num_frames 81 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --train_fps 16 \
+    --samples_per_file 1 \
+    --flush_frequency 1 \
+    --video_length_tolerance_range 5 \
+    --preprocess_task "ode_trajectory" 

examples/training/consistency_finetune/ode_init/validation.json

@@ -0,0 +1,40 @@
+{
+  "data": [
+    {
+      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about.",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A white and orange tabby cat is seen happily darting through a dense garden, as if chasing something. Its eyes are wide and happy as it jogs forward, scanning the branches, flowers, and leaves as it walks. The path is narrow as it makes its way between all the plants. the scene is captured from a ground-level angle, following the cat closely, giving a low and intimate perspective. The image is cinematic with warm tones and a grainy texture. The scattered daylight between the leaves and plants above creates a warm contrast, accentuating the cat’s orange fur. The shot is clear and sharp, with a shallow depth of field.",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    },
+    {
+      "caption": "A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.",
+      "image_path": null,
+      "video_path": null,
+      "num_inference_steps": 40,
+      "height": 480,
+      "width": 832,
+      "num_frames": 77
+    }
+  ]
+}

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/README.md

@@ -7,9 +7,7 @@ These are e2e example scripts for finetuning Wan2.1 T2V with VSA to accelerate i
 ## Make sure you have installed VSA
 
 ```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
+pip install vsa
 ```
 
 ### Download the synthetic dataset:

examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh

@@ -6,8 +6,8 @@ export TOKENIZERS_PARALLELISM=false
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+DATA_DIR="data/crush-smol_processed_t2v_old"
+VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
 NUM_GPUS=4
 # export CUDA_VISIBLE_DEVICES=4,5
 
@@ -52,7 +52,7 @@ dataset_args=(
 validation_args=(
   --log_validation 
   --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 200
+  --validation_steps 50
   --validation_sampling_steps "50" 
   --validation_guidance_scale "6.0"
 )

examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh

@@ -4,7 +4,7 @@ GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_t2v/"
+OUTPUT_DIR="data/crush-smol_processed_t2v_old/"
 
 torchrun --nproc_per_node=$GPU_NUM \
     fastvideo/pipelines/preprocess/v1_preprocess.py \

examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v_new.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-GPU_NUM=2 # 2,4,8
+GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATASET_PATH="data/crush-smol/"
 OUTPUT_DIR="data/crush-smol_processed_t2v/"
@@ -13,7 +13,7 @@ torchrun --nproc_per_node=$GPU_NUM \
     --preprocess.dataset_type merged \
     --preprocess.dataset_path $DATASET_PATH \
     --preprocess.dataset_output_dir $OUTPUT_DIR \
-    --preprocess.preprocess_video_batch_size 2 \
+    --preprocess.preprocess_video_batch_size 8 \
     --preprocess.dataloader_num_workers 0 \
     --preprocess.max_height 480 \
     --preprocess.max_width 832 \

examples/training/t2v_to_i2v/wan_t2v/finetune_t2v_to_i2v.sh

@@ -0,0 +1,94 @@
+#!/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="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v_old"
+VALIDATION_DATASET_FILE="examples/datasets/crush_smol/validation.json"
+NUM_GPUS=8
+# export CUDA_VISIBLE_DEVICES=4,5
+
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_t2v_i2v_finetune"
+  --output_dir "checkpoints/wan_t2v_i2v_finetune"
+  --max_train_steps 5000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 2
+  --num_latent_t 20
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS 
+  --sp_size 4
+  --tp_size 1
+  --hsdp_replicate_dim 2
+  --hsdp_shard_dim 4
+)
+
+# 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 50
+  --validation_sampling_steps "50" 
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 5e-5
+  --mixed_precision "bf16"
+  --checkpointing_steps 1000
+  --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
+  --enable_gradient_checkpointing_type "full"
+  --t2v_as_i2v_task True
+  # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500"
+)
+
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_t2v_i2v_training_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"

examples/training/t2v_to_i2v/wan_t2v/preprocess_wan_data_t2v.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_TYPE="wan"
+DATA_MERGE_PATH="data/crush-smol/merge.txt"
+OUTPUT_DIR="data/crush-smol_processed_t2v_i2v_1_3b/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 2 \
+    --seed 42 \
+    --max_height 480 \
+    --max_width 832 \
+    --num_frames 77 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --train_fps 16 \
+    --samples_per_file 8 \
+    --flush_frequency 8 \
+    --video_length_tolerance_range 5 \
+    --preprocess_task "t2v_ode_trajectory" 

fastvideo/configs/models/dits/wanvideo.py

@@ -92,6 +92,12 @@ class WanVideoArchConfig(DiTArchConfig):
     pos_embed_seq_len: int | None = None
     exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
 
+    # Causal Wan
+    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 = 21
+
     def __post_init__(self):
         super().__post_init__()
         self.out_channels = self.out_channels or self.in_channels

fastvideo/configs/pipelines/base.py

@@ -45,6 +45,8 @@ class PipelineConfig:
     embedded_cfg_scale: float = 6.0
     flow_shift: float | None = None
     disable_autocast: bool = False
+    ti2v_task: bool = False
+    t2v_as_i2v_task: bool = False
 
     # Model configuration
     dit_config: DiTConfig = field(default_factory=DiTConfig)
@@ -211,6 +213,24 @@ class PipelineConfig:
             "Comma-separated list of denoising steps (e.g., '1000,757,522')",
         )
 
+        # TI2V task
+        parser.add_argument(
+            f"--{prefix_with_dot}ti2v-task",
+            action=StoreBoolean,
+            dest=f"{prefix_with_dot.replace('-', '_')}ti2v_task",
+            default=PipelineConfig.ti2v_task,
+            help="Enable TI2V",
+        )
+
+        # T2V to I2V task
+        parser.add_argument(
+            f"--{prefix_with_dot}t2v-as-i2v-task",
+            action=StoreBoolean,
+            dest=f"{prefix_with_dot.replace('-', '_')}t2v_as_i2v_task",
+            default=PipelineConfig.t2v_as_i2v_task,
+            help="Enable T2V to I2V task",
+        )
+
         # Add VAE configuration arguments
         from fastvideo.configs.models.vaes.base import VAEConfig
         VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")
@@ -242,7 +262,9 @@ class PipelineConfig:
         """
         from fastvideo.configs.pipelines.registry import (
             get_pipeline_config_cls_from_name)
+        logger.info("WTF model_path: %s", model_path)
         pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)
+        logger.info("pipeline_config_cls: %s", pipeline_config_cls)
 
         return cast(PipelineConfig, pipeline_config_cls(model_path=model_path))
 

fastvideo/configs/pipelines/registry.py

@@ -7,10 +7,14 @@ from collections.abc import Callable
 from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
-from fastvideo.configs.pipelines.wan import (FastWan2_1_T2V_480P_Config,
-                                             FastWan2_2_TI2V_5B_Config,
-                                             WanI2V480PConfig, WanI2V720PConfig,
-                                             WanT2V480PConfig, WanT2V720PConfig)
+
+# isort: off
+from fastvideo.configs.pipelines.wan import (
+    FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
+    SelfForcingWanT2V480PConfig, Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config,
+    Wan2_2_TI2V_5B_Config, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig,
+    WanT2V720PConfig, SelfForcingWanT2V480PConfig)
+# isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
                              verify_model_config_and_directory)
@@ -31,9 +35,10 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "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,
-    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig,
-    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": WanT2V480PConfig,
-    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": WanI2V480PConfig,
+    "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+    "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,
     # Add other specific weight variants
 }
 
@@ -44,6 +49,7 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
     "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
     "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(),
     "stepvideo": lambda id: "stepvideo" in id.lower(),
+    "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(),
     # Add other pipeline architecture detectors
 }
 
@@ -55,7 +61,8 @@ PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = {
     WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V480PConfig,
     "wandmdpipeline": FastWan2_1_T2V_480P_Config,
-    "stepvideo": StepVideoT2VConfig
+    "stepvideo": StepVideoT2VConfig,
+    "wancausaldmdpipeline": SelfForcingWanT2V480PConfig,
     # Other fallbacks by architecture
 }
 

fastvideo/configs/pipelines/wan.py

@@ -39,7 +39,7 @@ class WanT2V480PConfig(PipelineConfig):
     vae_sp: bool = False
 
     # Denoising stage
-    flow_shift: int = 3
+    flow_shift: float | None = 8.0
 
     # Text encoding stage
     text_encoder_configs: tuple[EncoderConfig, ...] = field(
@@ -138,3 +138,14 @@ class Wan2_2_T2V_A14B_Config(WanT2V480PConfig):
 @dataclass
 class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
     pass
+
+
+# =============================================
+# ============= Causal Self-Forcing =============
+# =============================================
+@dataclass
+class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
+    is_causal: bool = True
+    flow_shift: float | None = 5.0
+    dmd_denoising_steps: list[int] | None = field(
+        default_factory=lambda: [1000, 750, 500, 250])

fastvideo/configs/sample/base.py

@@ -47,6 +47,8 @@ class SamplingParam:
     # Misc
     save_video: bool = True
     return_frames: bool = False
+    return_trajectory_latents: bool = False  # returns all latents for each timestep
+    return_trajectory_decoded: bool = False  # returns decoded latents for each timestep
 
     def __post_init__(self) -> None:
         self.data_type = "video" if self.num_frames > 1 else "image"
@@ -191,6 +193,18 @@ class SamplingParam:
             default=SamplingParam.image_path,
             help="Path to input image for image-to-video generation",
         )
+        parser.add_argument(
+            "--return-trajectory-latents",
+            action="store_true",
+            default=SamplingParam.return_trajectory_latents,
+            help="Whether to return the trajectory",
+        )
+        parser.add_argument(
+            "--return-trajectory-decoded",
+            action="store_true",
+            default=SamplingParam.return_trajectory_decoded,
+            help="Whether to return the decoded trajectory",
+        )
         return parser
 
 

fastvideo/configs/sample/registry.py

@@ -10,6 +10,7 @@ from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
 # isort: off
 from fastvideo.configs.sample.wan import (
     FastWanT2V480PConfig,
+    Wan2_1_Fun_1_3B_InP_SamplingParam,
     Wan2_2_I2V_A14B_SamplingParam,
     Wan2_2_T2V_A14B_SamplingParam,
     Wan2_2_TI2V_5B_SamplingParam,
@@ -17,6 +18,7 @@ from fastvideo.configs.sample.wan import (
     WanI2V_14B_720P_SamplingParam,
     WanT2V_1_3B_SamplingParam,
     WanT2V_14B_SamplingParam,
+    SelfForcingWanT2V480PConfig,
 )
 # isort: on
 from fastvideo.logger import init_logger
@@ -28,17 +30,31 @@ logger = init_logger(__name__)
 SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
     "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
+    "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,
-    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
-    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
+    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers":
+    Wan2_1_Fun_1_3B_InP_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": FastWanT2V480PConfig,
+
+    # 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": SelfForcingWanT2V480PConfig,
     # Add other specific weight variants
 }
 

fastvideo/configs/sample/wan.py

@@ -107,6 +107,21 @@ class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
     fps: int = 16
 
 
+# =============================================
+# ============= Wan2.1 Fun Models =============
+# =============================================
+@dataclass
+class Wan2_1_Fun_1_3B_InP_SamplingParam(SamplingParam):
+    """Sampling parameters for Wan2.1 Fun 1.3B InP model."""
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 16
+    negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+    guidance_scale: float = 6.0
+    num_inference_steps: int = 50
+
+
 # =============================================
 # ============= Wan2.2 TI2V Models =============
 # =============================================
@@ -141,3 +156,11 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
     guidance_scale_2: float = 3.5
     num_inference_steps: int = 40
     fps: int = 16
+
+
+# =============================================
+# ============= Causal Self-Forcing =============
+# =============================================
+@dataclass
+class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
+    pass

fastvideo/dataset/dataloader/schema.py

@@ -78,3 +78,43 @@ pyarrow_schema_t2v = pa.schema([
     pa.field("duration_sec", pa.float64()),
     pa.field("fps", pa.float64()),
 ])
+
+pyarrow_schema_ode_trajectory = 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()),
+    # --- Text encoder output tensor ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("text_embedding_bytes", pa.binary()),
+    # e.g., [SeqLen, Dim]
+    pa.field("text_embedding_shape", pa.list_(pa.int64())),
+    # e.g., 'bfloat16' or 'float32'
+    pa.field("text_embedding_dtype", pa.string()),
+    # I2V
+    pa.field("image_condition_latents_bytes", pa.binary()),
+    pa.field("image_condition_latents_shape", pa.list_(pa.int64())),
+    pa.field("image_condition_latents_dtype", pa.string()),
+    # --- ODE Trajectory ---
+    pa.field("trajectory_latents_bytes", pa.binary()),
+    pa.field("trajectory_latents_shape", pa.list_(pa.int64())),
+    pa.field("trajectory_latents_dtype", pa.string()),
+    pa.field("trajectory_timesteps_bytes", pa.binary()),
+    pa.field("trajectory_timesteps_shape", pa.list_(pa.int64())),
+    pa.field("trajectory_timesteps_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()),
+])

fastvideo/dataset/utils.py

@@ -3,9 +3,12 @@ from typing import Any, cast
 
 import numpy as np
 import torch
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
 
 
-def pad(t: torch.Tensor, padding_length: int) -> torch.Tensor:
+def pad(t: torch.Tensor, padding_length: int) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Pad or crop an embedding [L, D] to exactly padding_length tokens.
     Return:

fastvideo/entrypoints/video_generator.py

@@ -344,6 +344,9 @@ class VideoGenerator:
                 "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,
             }
 
     def set_lora_adapter(self,

fastvideo/fastvideo_args.py

@@ -154,6 +154,7 @@ class FastVideoArgs:
         "transformer": True,
         "vae": True,
     })
+    override_transformer_cls_name: str | None = None
 
     # # DMD parameters
     # dmd_denoising_steps: List[int] | None = field(default=None)
@@ -382,6 +383,12 @@ class FastVideoArgs:
             default=FastVideoArgs.enable_stage_verification,
             help="Enable input/output verification for pipeline stages",
         )
+        parser.add_argument(
+            "--override-transformer-cls-name",
+            type=str,
+            default=FastVideoArgs.override_transformer_cls_name,
+            help="Override transformer cls name",
+        )
         # Add pipeline configuration arguments
         PipelineConfig.add_cli_args(parser)
 
@@ -674,6 +681,7 @@ class TrainingArgs(FastVideoArgs):
     log_visualization: bool = False
     # simulate generator forward to match inference
     simulate_generator_forward: bool = False
+    warp_denoising_step: bool = False
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -1018,6 +1026,10 @@ class TrainingArgs(FastVideoArgs):
             "--simulate-generator-forward",
             action=StoreBoolean,
             help="Whether to simulate generator forward to match inference")
+        parser.add_argument(
+            "--warp-denoising-step",
+            action=StoreBoolean,
+            help="Whether to warp denoising step according to the scheduler time shift")
 
         return parser
 

fastvideo/layers/layernorm.py

@@ -9,6 +9,9 @@ import torch.nn.functional as F
 from fastvideo.layers.custom_op import CustomOp
 from fastvideo.platforms import current_platform
 
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
 
 @CustomOp.register("rms_norm")
 class RMSNorm(CustomOp):
@@ -100,7 +103,13 @@ class ScaleResidual(nn.Module):
     def forward(self, residual: torch.Tensor, x: torch.Tensor,
                 gate: torch.Tensor) -> torch.Tensor:
         """Apply gated residual connection."""
-        return residual + x * gate
+        # logger.info("x.shape: %s", x.shape)
+        # if isinstance(gate, torch.Tensor):
+            # logger.info("gate.shape: %s", gate.shape)
+
+        num_frames = gate.shape[1]
+        frame_seqlen = x.shape[1] // num_frames
+        return residual + (x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2)
 
 
 # adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py
@@ -172,11 +181,35 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
               but before normalization)
         """
         # Apply residual connection with gating
-        residual_output = residual + x * gate
+        # logger.info("x.shape: %s", x.shape)
+        if isinstance(gate, int): 
+            # used by cross-attention, should be 1
+            assert gate == 1
+            residual_output = residual + x * gate
+        elif isinstance(gate, torch.Tensor):
+            # logger.info("gate.shape: %s", gate.shape)
+            if gate.dim() == 3:
+                # used by bidirectional self attention
+                residual_output = residual + x * gate
+            else:
+                assert gate.dim() == 4 
+                num_frames = gate.shape[1]
+                frame_seqlen = x.shape[1] // num_frames
+                residual_output = residual + (x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2)
+                # residual_output = residual + x * gate
+        else:
+            raise ValueError(f"Gate type {type(gate)} not supported")
+        # logger.info("residual_output.shape: %s", residual_output.shape)
+
         # Apply normalization
         normalized = self.norm(residual_output)
         # Apply scale and shift
-        modulated = normalized * (1.0 + scale) + shift
+        if isinstance(scale, torch.Tensor) and scale.dim() == 4:
+            num_frames = scale.shape[1]
+            frame_seqlen = normalized.shape[1] // num_frames
+            modulated = (normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2)
+        else:
+            modulated = normalized * (1.0 + scale) + shift
         return modulated, residual_output
 
 
@@ -219,7 +252,15 @@ class LayerNormScaleShift(nn.Module):
                 scale: torch.Tensor) -> torch.Tensor:
         """Apply ln followed by scale and shift in a single fused operation."""
         normalized = self.norm(x)
-        if self.compute_dtype == torch.float32:
-            return (normalized.float() * (1.0 + scale) + shift).to(x.dtype)
+        if scale.dim() == 4:
+            num_frames = scale.shape[1]
+            frame_seqlen = normalized.shape[1] // num_frames
+            if self.compute_dtype == torch.float32:
+                return (normalized.float().unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2).to(x.dtype)
+            else:
+                return (normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2)
         else:
-            return normalized * (1.0 + scale) + shift
+            if self.compute_dtype == torch.float32:
+                return (normalized.float() * (1.0 + scale) + shift).to(x.dtype)
+            else:
+                return normalized * (1.0 + scale) + shift

fastvideo/layers/rotary_embedding.py

@@ -29,6 +29,9 @@ import torch
 
 from fastvideo.distributed.parallel_state import get_sp_group
 from fastvideo.layers.custom_op import CustomOp
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
 
 
 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
@@ -267,6 +270,7 @@ def get_nd_rotary_pos_embed(
     sp_rank: int = 0,
     sp_world_size: int = 1,
     dtype: torch.dtype = torch.float32,
+    start_frame: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure.
@@ -292,6 +296,9 @@ def get_nd_rotary_pos_embed(
     full_grid = get_meshgrid_nd(
         start, *args, dim=len(rope_dim_list))  # [3, W, H, D] / [2, W, H]
 
+    if start_frame > 0:
+        full_grid[0] += start_frame
+
     # Shard the grid if using sequence parallelism (sp_world_size > 1)
     assert shard_dim < len(
         rope_dim_list
@@ -370,6 +377,7 @@ def get_rotary_pos_embed(
     interpolation_factor=1.0,
     shard_dim: int = 0,
     dtype: torch.dtype = torch.float32,
+    start_frame: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Generate rotary positional embeddings for the given sizes.
@@ -413,6 +421,7 @@ def get_rotary_pos_embed(
         sp_rank=sp_rank,
         sp_world_size=sp_world_size,
         dtype=dtype,
+        start_frame=start_frame,
     )
     return freqs_cos, freqs_sin
 

fastvideo/layers/visual_embedding.py

@@ -86,12 +86,16 @@ class TimestepEmbedder(nn.Module):
                        dtype=dtype)
         self.freq_dtype = freq_dtype
 
-    def forward(self, t: torch.Tensor) -> torch.Tensor:
+    def forward(self,
+                t: torch.Tensor,
+                timestep_seq_len: int | None = None) -> torch.Tensor:
         t_freq = timestep_embedding(t,
                                     self.frequency_embedding_size,
                                     self.max_period,
                                     dtype=self.freq_dtype).to(
                                         self.mlp.fc_in.weight.dtype)
+        if timestep_seq_len is not None:
+            t_freq = t_freq.unflatten(0, (1, timestep_seq_len))
         # t_freq = t_freq.to(self.mlp.fc_in.weight.dtype)
         t_emb = self.mlp(t_freq)
         return t_emb

fastvideo/models/dits/causal_wanvideo.py

@@ -0,0 +1,699 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+from torch.nn.attention.flex_attention import BlockMask
+# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention
+# see https://github.com/pytorch/pytorch/issues/133254
+# change to default for other models
+flex_attention = torch.compile(
+    flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs")
+import torch.distributed as dist
+
+import fastvideo.envs as envs
+from fastvideo.attention import (DistributedAttention,
+                                 LocalAttention)
+from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.distributed.parallel_state import get_sp_world_size
+from fastvideo.forward_context import get_forward_context
+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)
+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, current_platform
+
+logger = init_logger(__name__)
+class CausalWanSelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 local_attn_size: int = -1,
+                 sink_size: int = 0,
+                 qk_norm=True,
+                 eps=1e-6,
+                 parallel_attention=False) -> None:
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.local_attn_size = local_attn_size
+        self.sink_size = sink_size
+        self.qk_norm = qk_norm
+        self.eps = eps
+        self.parallel_attention = parallel_attention
+        self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560
+
+        # Scaled dot product attention
+        self.attn = LocalAttention(
+            num_heads=num_heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
+                                          AttentionBackendEnum.TORCH_SDPA))
+
+    def forward(self, 
+                q: torch.Tensor,
+                k: torch.Tensor,
+                v: torch.Tensor,
+                freqs_cis: tuple[torch.Tensor, torch.Tensor],
+                block_mask: BlockMask,
+                kv_cache: dict | None = None,
+                current_start: int = 0,
+                cache_start: int | None = None):
+        r"""
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        if cache_start is None:
+            cache_start = current_start
+
+        cos, sin = freqs_cis
+        roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v)
+        roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v)
+
+        if kv_cache is None:
+            # Padding for flex attention
+            padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
+            padded_roped_query = torch.cat(
+                [roped_query,
+                    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(
+                [roped_key, 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
+            )
+
+            x = 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)
+        else:
+            frame_seqlen = q.shape[1]
+            current_end = current_start + roped_query.shape[1]
+            sink_tokens = self.sink_size * frame_seqlen
+            # 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 = kv_cache["k"].shape[1]
+            num_new_tokens = roped_query.shape[1]
+            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["local_end_index"].item() - num_evicted_tokens - sink_tokens
+                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()
+                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()
+                # 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
+                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+                kv_cache["v"][:, local_start_index:local_end_index] = v
+            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
+                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+                kv_cache["v"][:, local_start_index:local_end_index] = v
+            x = self.attn(
+                roped_query,
+                kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index],
+                kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index]
+            )
+            kv_cache["global_end_index"].fill_(current_end)
+            kv_cache["local_end_index"].fill_(local_end_index)
+
+        return x
+
+class CausalWanTransformerBlock(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 ffn_dim: int,
+                 num_heads: int,
+                 local_attn_size: int = -1,
+                 sink_size: int = 0,
+                 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 = ""):
+        super().__init__()
+
+        # 1. Self-attention
+        self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
+        self.to_q = ReplicatedLinear(dim, dim, bias=True)
+        self.to_k = ReplicatedLinear(dim, dim, bias=True)
+        self.to_v = ReplicatedLinear(dim, dim, bias=True)
+
+        self.to_out = ReplicatedLinear(dim, dim, bias=True)
+        self.attn1 = CausalWanSelfAttention(
+            dim,
+            num_heads,
+            local_attn_size=local_attn_size,
+            sink_size=sink_size,
+            qk_norm=qk_norm,
+            eps=eps)
+        self.hidden_dim = dim
+        self.num_attention_heads = num_heads
+        self.local_attn_size = local_attn_size
+        dim_head = dim // num_heads
+        if qk_norm == "rms_norm":
+            self.norm_q = RMSNorm(dim_head, eps=eps)
+            self.norm_k = RMSNorm(dim_head, eps=eps)
+        elif qk_norm == "rms_norm_across_heads":
+            # LTX applies qk norm across all heads
+            self.norm_q = RMSNorm(dim, eps=eps)
+            self.norm_k = RMSNorm(dim, eps=eps)
+        else:
+            print("QK Norm type not supported")
+            raise Exception
+        assert cross_attn_norm is True
+        self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift(
+            dim,
+            norm_type="layer",
+            eps=eps,
+            elementwise_affine=True,
+            dtype=torch.float32,
+            compute_dtype=torch.float32)
+
+        # 2. Cross-attention
+        # Only T2V for now
+        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)
+
+        # 3. Feed-forward
+        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.mlp_residual = ScaleResidual()
+
+        self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        freqs_cis: tuple[torch.Tensor, torch.Tensor],
+        block_mask: BlockMask,
+        kv_cache: dict | None = None,
+        crossattn_cache: dict | None = None,
+        current_start: int = 0,
+        cache_start: int | None = None,
+    ) -> torch.Tensor:
+        # logger.info("temb.shape: %s", temb.shape)
+        num_frames = temb.shape[1]
+        # logger.info("first hidden_states.shape: %s", hidden_states.shape)
+        # logger.info("num_frames: %s", num_frames)
+        if hidden_states.dim() == 4:
+            hidden_states = hidden_states.squeeze(1)
+        frame_seqlen = hidden_states.shape[1] // temb.shape[1]
+        # logger.info("frame_seqlen: %s", frame_seqlen)
+        bs, seq_length, _ = hidden_states.shape
+        orig_dtype = hidden_states.dtype
+        # assert orig_dtype != torch.float32
+        e = self.scale_shift_table + temb.float()
+        # logger.info("e.shape: %s", e.shape)
+        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
+            6, dim=2)
+        assert shift_msa.dtype == torch.float32
+
+        # 1. Self-attention
+        # print(f"hidden_states: {hidden_states.shape}")
+        # print(f"hidden_states: {scale_msa.shape}")
+        # print(f"hidden_states: {shift_msa.shape}")
+        # norm_hidden_states = (self.norm1(hidden_states.float()) *
+        #                       (1 + scale_msa) + shift_msa).to(orig_dtype)
+        norm_hidden_states = (self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) *
+                              (1 + scale_msa) + shift_msa).flatten(1, 2).to(orig_dtype)
+        # logger.info("1 norm_hidden_states.shape: %s", norm_hidden_states.shape)
+        query, _ = self.to_q(norm_hidden_states)
+        key, _ = self.to_k(norm_hidden_states)
+        value, _ = self.to_v(norm_hidden_states)
+
+        if self.norm_q is not None:
+            query = self.norm_q(query)
+        if self.norm_k is not None:
+            key = self.norm_k(key)
+
+        query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+        value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
+
+        attn_output = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start)
+        attn_output = attn_output.flatten(2)
+        attn_output, _ = self.to_out(attn_output)
+        attn_output = attn_output.squeeze(1)
+
+        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)
+        # logger.info("after self_attn_residual_norm norm_hidden_states.shape: %s", norm_hidden_states.shape)
+        # logger.info("after self_attn_residual_norm hidden_states.shape: %s", hidden_states.shape)
+        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,
+                                 crossattn_cache=crossattn_cache)
+        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
+            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
+        # logger.info("after cross_attn_residual_norm norm_hidden_states.shape: %s", norm_hidden_states.shape)
+        # logger.info("after cross_attn_residual_norm hidden_states.shape: %s", hidden_states.shape)
+        norm_hidden_states, hidden_states = norm_hidden_states.to(
+            orig_dtype), hidden_states.to(orig_dtype)
+
+        # 3. Feed-forward
+        ff_output = self.ffn(norm_hidden_states)
+        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
+        # logger.info("after mlp_residual norm_hidden_states.shape: %s", norm_hidden_states.shape)
+        # logger.info("after mlp_residual hidden_states.shape: %s", hidden_states.shape)
+        hidden_states = hidden_states.to(orig_dtype)
+
+        return hidden_states
+
+class CausalWanTransformer3DModel(BaseDiT):
+    _fsdp_shard_conditions = WanVideoConfig()._fsdp_shard_conditions
+    _compile_conditions = WanVideoConfig()._compile_conditions
+    _supported_attention_backends = WanVideoConfig(
+    )._supported_attention_backends
+    param_names_mapping = WanVideoConfig().param_names_mapping
+    reverse_param_names_mapping = WanVideoConfig().reverse_param_names_mapping
+    lora_param_names_mapping = WanVideoConfig().lora_param_names_mapping
+
+    def __init__(self, config: WanVideoConfig, hf_config: dict[str,
+                                                               Any]) -> 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.attention_head_dim = config.attention_head_dim
+        self.in_channels = config.in_channels
+        self.out_channels = config.out_channels
+        self.num_channels_latents = config.num_channels_latents
+        self.patch_size = config.patch_size
+        self.text_len = config.text_len
+        self.local_attn_size = config.local_attn_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)
+
+        # 2. Condition embeddings
+        self.condition_embedder = WanTimeTextImageEmbedding(
+            dim=inner_dim,
+            time_freq_dim=config.freq_dim,
+            text_embed_dim=config.text_dim,
+            image_embed_dim=config.image_dim,
+        )
+
+        # 3. Transformer blocks
+        self.blocks = nn.ModuleList([
+            CausalWanTransformerBlock(inner_dim,
+                              config.ffn_dim,
+                              config.num_attention_heads,
+                              config.local_attn_size,
+                              config.sink_size,
+                              config.qk_norm,
+                              config.cross_attn_norm,
+                              config.eps,
+                              config.added_kv_proj_dim,
+                              self._supported_attention_backends,
+                              prefix=f"{config.prefix}.blocks.{i}")
+            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)
+        # Debug: Log configuration values
+        proj_out_dim = config.out_channels * math.prod(config.patch_size)
+        
+        self.proj_out = nn.Linear(inner_dim, proj_out_dim)
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(1, 2, inner_dim) / inner_dim**0.5)
+
+        self.gradient_checkpointing = False
+
+        # Causal-specific
+        self.block_mask = None
+        self.num_frame_per_block = 1
+        self.independent_first_frame = False
+
+        self.__post_init__()
+
+    @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
+    ) -> 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
+
+        # we do right padding to get to a multiple of 128
+        padded_length = math.ceil(total_length / 128) * 128 - total_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=0,
+            end=total_length,
+            step=frame_seqlen * num_frame_per_block,
+            device=device
+        )
+
+        for tmp in frame_indices:
+            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)
+            else:
+                return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx)
+            # 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_length + 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_inference(
+                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,
+                kv_cache: dict = None,
+                crossattn_cache: dict = None,
+                current_start: int = 0,
+                cache_start: int = 0,
+                start_frame: int = 0,
+                **kwargs) -> torch.Tensor:
+        r"""
+        Run the diffusion model with kv caching.
+        See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details.
+        This function will be run for num_frame times.
+        Process the latent frames one by one (1560 tokens each)
+        """
+        # logger.info("forward inference hidden_states.shape: %s", hidden_states.shape)
+
+        orig_dtype = hidden_states.dtype
+        if 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:
+            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
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # Get rotary embeddings
+        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),
+            self.hidden_size,
+            self.num_attention_heads,
+            rope_dim_list,
+            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+            rope_theta=10000,
+            start_frame=start_frame # Assume that start_frame is 0 when kv_cache is None
+        )
+        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
+
+        hidden_states = self.patch_embedding(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        # logger.info("forward inference flattened and transposed hidden_states.shape: %s", hidden_states.shape)
+
+        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
+            timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
+        timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
+
+        if encoder_hidden_states_image is not None:
+            encoder_hidden_states = torch.concat(
+                [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+
+        encoder_hidden_states = encoder_hidden_states.to(
+            orig_dtype) if current_platform.is_mps(
+            ) else encoder_hidden_states  # cast to orig_dtype for MPS
+
+        assert encoder_hidden_states.dtype == orig_dtype
+
+        # 4. Transformer blocks
+        for block_index, block in enumerate(self.blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                causal_kwargs = {
+                    "kv_cache": kv_cache[block_index],
+                    "current_start": current_start,
+                    "cache_start": cache_start,
+                    "block_mask": self.block_mask
+                }
+                hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, encoder_hidden_states,
+                    timestep_proj, freqs_cis,
+                    **causal_kwargs)
+            else:
+                causal_kwargs = {
+                    "kv_cache": kv_cache[block_index],
+                    "crossattn_cache": crossattn_cache[block_index],
+                    "current_start": current_start,
+                    "cache_start": cache_start,
+                    "block_mask": self.block_mask
+                }
+                hidden_states = block(hidden_states, encoder_hidden_states,
+                                        timestep_proj, freqs_cis,
+                                        **causal_kwargs)
+
+        # 5. Output norm, projection & unpatchify
+        # logger.info("===== INFERENCE 5. Output norm, projection & unpatchify")
+        # logger.info("hidden_states.shape: %s", hidden_states.shape)
+        # logger.info("temb.shape: %s", temb.shape)
+        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.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+        return output
+
+    def _forward_train(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,
+                start_frame: int = 0,
+                **kwargs) -> torch.Tensor:
+
+        # logger.info("===== forward train hidden_states.shape: %s", hidden_states.shape)
+        # logger.info("===== forward train timestep.shape: %s", timestep.shape)
+        orig_dtype = hidden_states.dtype
+        if 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:
+            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
+        p_t, p_h, p_w = self.patch_size
+        post_patch_num_frames = num_frames // p_t
+        post_patch_height = height // p_h
+        post_patch_width = width // p_w
+
+        # Get rotary embeddings
+        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),
+            self.hidden_size,
+            self.num_attention_heads,
+            rope_dim_list,
+            dtype=torch.float32 if current_platform.is_mps() else torch.float64,
+            rope_theta=10000,
+            start_frame=start_frame
+        )
+        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
+
+        # Construct blockwise causal attn mask
+        if self.block_mask 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
+            )
+
+        hidden_states = self.patch_embedding(hidden_states)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        # logger.info("forward train flattened and transposed hidden_states.shape: %s", hidden_states.shape)
+
+        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
+            timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
+        # logger.info("forward train timestep_proj.shape: %s", timestep_proj.shape)
+        # logger.info("forward train timestep.shape: %s", timestep.shape)
+        # logger.info("forward train temb.shape: %s", temb.shape)
+        timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
+
+        if encoder_hidden_states_image is not None:
+            encoder_hidden_states = torch.concat(
+                [encoder_hidden_states_image, encoder_hidden_states], dim=1)
+
+        encoder_hidden_states = encoder_hidden_states.to(
+            orig_dtype) if current_platform.is_mps(
+            ) else encoder_hidden_states  # cast to orig_dtype for MPS
+
+        assert encoder_hidden_states.dtype == orig_dtype
+
+        # 4. Transformer blocks
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            for block in self.blocks:
+                hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, encoder_hidden_states,
+                    timestep_proj, freqs_cis,
+                    block_mask=self.block_mask)
+        else:
+            for block_index, block in enumerate(self.blocks):
+                # logger.info("===== TRAIN block %d", block_index)
+                # logger.info("hidden_states.shape: %s", hidden_states.shape)
+                # logger.info("encoder_hidden_states.shape: %s", encoder_hidden_states.shape)
+                # logger.info("timestep_proj.shape: %s", timestep_proj.shape)
+                # logger.info("freqs_cis.shape: %s", freqs_cis.shape)
+                # logger.info("block_mask.shape: %s", self.block_mask.shape)
+                hidden_states = block(hidden_states, encoder_hidden_states,
+                                        timestep_proj, freqs_cis,
+                                        block_mask=self.block_mask)
+
+        # 5. Output norm, projection & unpatchify
+        # logger.info("===== TRAIN 5. Output norm, projection & unpatchify")
+        # logger.info("hidden_states.shape: %s", hidden_states.shape)
+        # logger.info("temb.shape: %s", temb.shape)
+        # shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
+        temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
+        # logger.info("WTFWTF train temb.shape: %s", temb.shape)
+        # logger.info("WTFWTF train self.scale_shift_table.shape: %s", self.scale_shift_table.shape)
+        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2,
+                                                                          dim=2)
+        # logger.info("DEBUG scale.shape: %s", scale.shape)
+        # logger.info("DEBUG shift.shape: %s", shift.shape)
+        hidden_states = self.norm_out(hidden_states, shift, scale)
+        hidden_states = self.proj_out(hidden_states)
+        # logger.info("DEBUG after proj_out hidden_states.shape: %s", hidden_states.shape)
+        # logger.info(f"DEBUG reshape dimensions: batch_size={batch_size}, post_patch_num_frames={post_patch_num_frames}")
+        # logger.info(f"DEBUG reshape dimensions: post_patch_height={post_patch_height}, post_patch_width={post_patch_width}")
+        # logger.info(f"DEBUG patch dimensions: p_t={p_t}, p_h={p_h}, p_w={p_w}")
+
+        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
+
+    def forward(
+        self,
+        *args,
+        **kwargs
+    ):
+        if kwargs.get('kv_cache', None) is not None:
+            return self._forward_inference(*args, **kwargs)
+        else:
+            return self._forward_train(*args, **kwargs)

fastvideo/models/dits/wanvideo.py

@@ -81,8 +81,9 @@ class WanTimeTextImageEmbedding(nn.Module):
         timestep: torch.Tensor,
         encoder_hidden_states: torch.Tensor,
         encoder_hidden_states_image: torch.Tensor | None = None,
+        timestep_seq_len: int | None = None,
     ):
-        temb = self.time_embedder(timestep)
+        temb = self.time_embedder(timestep, timestep_seq_len)
         timestep_proj = self.time_modulation(temb)
 
         encoder_hidden_states = self.text_embedder(encoder_hidden_states)
@@ -145,7 +146,7 @@ class WanSelfAttention(nn.Module):
 
 class WanT2VCrossAttention(WanSelfAttention):
 
-    def forward(self, x, context, context_lens):
+    def forward(self, x, context, context_lens, crossattn_cache=None):
         r"""
         Args:
             x(Tensor): Shape [B, L1, C]
@@ -156,8 +157,20 @@ class WanT2VCrossAttention(WanSelfAttention):
 
         # compute query, key, value
         q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
-        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)
+
+        if crossattn_cache is not None:
+            if not crossattn_cache["is_init"]:
+                crossattn_cache["is_init"] = True
+                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)
+                crossattn_cache["k"] = k
+                crossattn_cache["v"] = v
+            else:
+                k = crossattn_cache["k"]
+                v = crossattn_cache["v"]
+        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)
@@ -307,9 +320,24 @@ class WanTransformerBlock(nn.Module):
         bs, seq_length, _ = hidden_states.shape
         orig_dtype = hidden_states.dtype
         # assert orig_dtype != torch.float32
-        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)
+
+        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)
+            # batch_size, seq_len, 1, inner_dim
+            shift_msa = shift_msa.squeeze(2)
+            scale_msa = scale_msa.squeeze(2)
+            gate_msa = gate_msa.squeeze(2)
+            c_shift_msa = c_shift_msa.squeeze(2)
+            c_scale_msa = c_scale_msa.squeeze(2)
+            c_gate_msa = c_gate_msa.squeeze(2)
+        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)
         assert shift_msa.dtype == torch.float32
 
         # 1. Self-attention
@@ -637,9 +665,21 @@ class WanTransformer3DModel(CachableDiT):
         hidden_states = self.patch_embedding(hidden_states)
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
+        # timestep shape: batch_size, or batch_size, seq_len (wan 2.2 ti2v)
+        if timestep.dim() == 2:
+            ts_seq_len = timestep.shape[1]
+            timestep = timestep.flatten()  # batch_size * seq_len
+        else:
+            ts_seq_len = None
+
         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))
+            timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len)
+        if ts_seq_len is not None:
+            # batch_size, seq_len, 6, inner_dim
+            timestep_proj = timestep_proj.unflatten(2, (6, -1))
+        else:
+            # batch_size, 6, inner_dim
+            timestep_proj = timestep_proj.unflatten(1, (6, -1))
 
         if encoder_hidden_states_image is not None:
             encoder_hidden_states = torch.concat(
@@ -676,8 +716,15 @@ class WanTransformer3DModel(CachableDiT):
             if enable_teacache:
                 self.maybe_cache_states(hidden_states, original_hidden_states)
         # 5. Output norm, projection & unpatchify
-        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
-                                                                          dim=1)
+        if temb.dim() == 3:
+            # batch_size, seq_len, inner_dim (wan 2.2 ti2v)
+            shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
+            shift = shift.squeeze(2)
+            scale = scale.squeeze(2)
+        else:
+            # batch_size, inner_dim
+            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)
 
@@ -781,3 +828,4 @@ class WanTransformer3DModel(CachableDiT):
             return hidden_states + self.previous_residual_even
         else:
             return hidden_states + self.previous_residual_odd
+        

fastvideo/models/loader/component_loader.py

@@ -415,6 +415,10 @@ class TransformerLoader(ComponentLoader):
             raise ValueError(
                 "Model config does not contain a _class_name attribute. "
                 "Only diffusers format is supported.")
+        logger.info("transformer cls_name: %s", cls_name)
+        if fastvideo_args.override_transformer_cls_name is not None:
+            cls_name = fastvideo_args.override_transformer_cls_name
+            logger.info("Overriding transformer cls_name to %s", cls_name)
 
         fastvideo_args.model_paths["transformer"] = model_path
 

fastvideo/models/registry.py

@@ -25,12 +25,14 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "HunyuanVideoTransformer3DModel":
     ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
+    "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
     "StepVideoModel": ("dits", "stepvideo", "StepVideoModel")
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {
     # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
     "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
+    "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
 }
 
 _TEXT_ENCODER_MODELS = {

fastvideo/models/utils.py

@@ -6,6 +6,9 @@ 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):
     """
@@ -137,3 +140,24 @@ def modulate(x: torch.Tensor,
     else:
         return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(
             1)  # type: ignore[union-attr]
+
+
+def pred_noise_to_pred_video(pred_noise: torch.Tensor,
+                             noise_input_latent: torch.Tensor,
+                             timestep: torch.Tensor,
+                             scheduler: Any) -> torch.Tensor:
+    """
+    Convert predicted noise to clean latent.
+    """
+    timestep = timestep.expand(noise_input_latent.shape[0])
+    dtype = pred_noise.dtype
+    device = pred_noise.device
+    pred_noise = pred_noise.float().to(device)
+    noise_input_latent = noise_input_latent.float().to(device)
+    sigmas = scheduler.sigmas.float().to(device)
+    timesteps = scheduler.timesteps.float().to(device)
+    timestep_id = torch.argmin(
+        (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+    sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1)
+    pred_video = noise_input_latent - sigma_t * pred_noise
+    return pred_video.to(dtype)

fastvideo/pipelines/basic/wan/wan_causal_dmd_pipeline.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Wan causal DMD pipeline implementation.
+
+This module wires the causal DMD denoising stage into the modular pipeline.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+
+# isort: off
+from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage,
+                                        CausalDMDDenosingStage,
+                                        InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        TimestepPreparationStage)
+# isort: on
+
+logger = init_logger(__name__)
+
+
+class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift)
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+
+        self.add_stage(stage_name="conditioning_stage",
+                       stage=ConditioningStage())
+
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer", None)))
+
+        self.add_stage(stage_name="denoising_stage",
+                       stage=CausalDMDDenosingStage(
+                           transformer=self.get_module("transformer"),
+                           scheduler=self.get_module("scheduler")))
+
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+
+EntryClass = WanCausalDMDPipeline

fastvideo/pipelines/basic/wan/wan_pipeline.py

@@ -63,6 +63,7 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase):
                            transformer=self.get_module("transformer"),
                            transformer_2=self.get_module("transformer_2", None),
                            scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae"),
                            pipeline=self))
 
         self.add_stage(stage_name="decoding_stage",

fastvideo/pipelines/composed_pipeline_base.py

@@ -121,14 +121,25 @@ class ComposedPipelineBase(ABC):
                         model_path: str,
                         device: str | None = None,
                         torch_dtype: torch.dtype | None = None,
-                        pipeline_config: str | PipelineConfig | None = None,
+                        pipeline_config: PipelineConfig | None = None,
                         args: argparse.Namespace | None = None,
                         required_config_modules: list[str] | None = None,
                         loaded_modules: dict[str, torch.nn.Module]
                         | None = None,
                         **kwargs) -> "ComposedPipelineBase":
         """
-        Load a pipeline from a pretrained model.
+        Load a pipeline from a pretrained model. 
+        Few different patterns are supported:
+        - Only provide model_path:
+            - This will load the pipeline in inference mode.
+            - The pipeline will be initialized with the default config.
+            - The pipeline will be initialized with the default modules.
+            - The pipeline will be initialized with the default stages.
+            - The pipeline will be initialized with the default stages.
+        - override the default config using pipeline_config or args or kwargs
+        - override the default modules using loaded_modules
+        - override the pipelineconfig
+
         loaded_modules: Optional[Dict[str, torch.nn.Module]] = None,
         If provided, loaded_modules will be used instead of loading from config/pretrained weights.
         """
@@ -136,9 +147,18 @@ class ComposedPipelineBase(ABC):
 
             kwargs['model_path'] = model_path
             fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
+            if pipeline_config is not None:
+                fastvideo_args.pipeline_config = pipeline_config
+            if fastvideo_args.override_transformer_cls_name is not None:
+                pipeline_config = PipelineConfig.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers")
+                fastvideo_args.pipeline_config = pipeline_config
         else:
             assert args is not None, "args must be provided for training mode"
             fastvideo_args = TrainingArgs.from_cli_args(args)
+            if fastvideo_args.override_transformer_cls_name is not None:
+                pipeline_config = PipelineConfig.from_pretrained("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers")
+                fastvideo_args.pipeline_config = pipeline_config
+                logger.info("in 2 Overriding transformer cls name to %s", fastvideo_args.override_transformer_cls_name)
             # TODO(will): fix this so that its not so ugly
             fastvideo_args.model_path = model_path
             for key, value in kwargs.items():
@@ -149,7 +169,8 @@ class ComposedPipelineBase(ABC):
             # model is loaded with the correct precision. Subsequently we will
             # use FSDP2's MixedPrecisionPolicy to set the precision for the
             # fwd, bwd, and other operations' precision.
-            assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
+            fastvideo_args.pipeline_config.dit_precision = 'fp32'
+            # assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
 
         logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
 

fastvideo/pipelines/pipeline_batch_info.py

@@ -147,7 +147,12 @@ class ForwardBatch:
     modules: dict[str, Any] = field(default_factory=dict)
 
     # Final output (after pipeline completion)
-    output: Any = None
+    output: torch.Tensor | None = None
+    return_trajectory_latents: bool = False
+    return_trajectory_decoded: bool = False
+    trajectory_timesteps: list[int] | None = None
+    trajectory_latents: torch.Tensor | None = None
+    trajectory_decoded: list[torch.Tensor] | None = None
 
     # Extra parameters that might be needed by specific pipeline implementations
     extra: dict[str, Any] = field(default_factory=dict)
@@ -206,6 +211,10 @@ class TrainingBatch:
     infos: list[dict[str, Any]] | None = None
     mask_lat_size: torch.Tensor | None = None
 
+    # ODE trajectory supervision
+    trajectory_latents: torch.Tensor | None = None
+    trajectory_timesteps: torch.Tensor | None = None
+
     # Transformer inputs
     noisy_model_input: torch.Tensor | None = None
     timesteps: torch.Tensor | None = None

fastvideo/pipelines/pipeline_registry.py

@@ -21,6 +21,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "WanPipeline": "wan",
     "WanDMDPipeline": "wan",
     "WanImageToVideoPipeline": "wan",
+    "WanCausalDMDPipeline": "wan",
     "StepVideoPipeline": "stepvideo",
     "HunyuanVideoPipeline": "hunyuan",
 }

fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py

@@ -0,0 +1,654 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+ODE Trajectory Data Preprocessing pipeline implementation.
+
+This module contains an implementation of the ODE Trajectory Data Preprocessing pipeline
+using the modular pipeline architecture.
+
+Sec 4.3 of CausVid paper: https://arxiv.org/pdf/2412.07772
+"""
+
+import os
+from collections.abc import Iterator
+from typing import Any
+
+import numpy as np
+import pyarrow as pa
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm import tqdm
+
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset import getdataset
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
+    BasePreprocessPipeline)
+from fastvideo.pipelines.stages import (DecodingStage, DenoisingStage,
+                                        ImageVAEEncodingStage,
+                                        InputValidationStage,
+                                        LatentPreparationStage,
+                                        TextEncodingStage,
+                                        TimestepPreparationStage)
+from fastvideo.utils import save_decoded_latents_as_video, shallow_asdict
+
+logger = init_logger(__name__)
+
+
+class FlowMatchScheduler:
+
+    order = 1
+
+    def __init__(self,
+                 num_inference_steps=100,
+                 num_train_timesteps=1000,
+                 shift=3.0,
+                 sigma_max=1.0,
+                 sigma_min=0.003 / 1.002,
+                 inverse_timesteps=False,
+                 extra_one_step=False,
+                 reverse_sigmas=False):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self,
+                      num_inference_steps=100,
+                      denoising_strength=1.0,
+                      training=False,
+                      device=None):
+        sigma_start = self.sigma_min + \
+            (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min,
+                                         num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min,
+                                         num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        self.sigmas = self.shift * self.sigmas / \
+            (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(
+                -2 * ((x - num_inference_steps / 2) / num_inference_steps)**2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * \
+                (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+
+    def step(self,
+             model_output,
+             timestep,
+             sample,
+             to_final=False,
+             return_dict=False,
+             **kwargs):
+        assert return_dict is False
+        assert kwargs == {}
+        self.sigmas = self.sigmas.to(model_output.device)
+        self.timesteps = self.timesteps.to(model_output.device)
+        logger.info('step timestep: %s', timestep)
+        logger.info('step timestep: %s', timestep.shape)
+        # timestep  is [num_frames]
+        # timestep_id = torch.argmin(
+        #     (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        # assert timestep.ndim == 1
+        # assert timestep.shape[0] == 1
+        timestep_id = torch.argmin((self.timesteps - timestep).abs(), dim=0)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
+            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return (prev_sample, )
+
+    def scale_model_input(self, sample: torch.Tensor, *args,
+                          **kwargs) -> torch.Tensor:
+        """
+        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+        current timestep.
+
+        Args:
+            sample (`torch.Tensor`):
+                The input sample.
+
+        Returns:
+            `torch.Tensor`:
+                A scaled input sample.
+        """
+        return sample
+
+    def add_noise(self, original_samples, noise, timestep):
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B, C, H, W]
+            - noise: the noise with shape [B, C, H, W]
+            - timestep: the timestep with shape [B]
+        Output: the corrupted latent with shape [B, C, H, W]
+        """
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        timestep_id = torch.argmin(
+            (self.timesteps - timestep.to(self.timesteps.device)).abs())
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights
+
+
+class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline):
+    """ODE Trajectory preprocessing pipeline implementation."""
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+    preprocess_dataloader: StatefulDataLoader
+    preprocess_loader_iter: Iterator[dict[str, Any]]
+
+    def get_schema_fields(self):
+        """Get the schema fields for ODE Trajectory pipeline."""
+        return [f.name for f in pyarrow_schema_ode_trajectory]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+        fastvideo_args.pipeline_config.flow_shift = 5
+        logger.info('WTF flow_shift: %s',
+                    fastvideo_args.pipeline_config.flow_shift)
+
+        assert fastvideo_args.pipeline_config.flow_shift == 5
+        # self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
+        # shift=fastvideo_args.pipeline_config.flow_shift)
+        self.modules["scheduler"] = FlowMatchScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift,
+            sigma_min=0.0,
+            extra_one_step=True)
+        self.modules["scheduler"].set_timesteps(num_inference_steps=48,
+                                                denoising_strength=1.0)
+        logger.info('WTF scheduler timesteps: %s',
+                    self.modules["scheduler"].timesteps)
+
+        self.add_stage(stage_name="input_validation_stage",
+                       stage=InputValidationStage())
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+        self.add_stage(stage_name="vae_encoding_stage",
+                       stage=ImageVAEEncodingStage(
+                           vae=self.get_module("vae"), ))
+        self.add_stage(stage_name="timestep_preparation_stage",
+                       stage=TimestepPreparationStage(
+                           scheduler=self.get_module("scheduler")))
+        self.add_stage(stage_name="latent_preparation_stage",
+                       stage=LatentPreparationStage(
+                           scheduler=self.get_module("scheduler"),
+                           transformer=self.get_module("transformer", None)))
+        self.add_stage(stage_name="denoising_stage",
+                       stage=DenoisingStage(
+                           transformer=self.get_module("transformer"),
+                           transformer_2=self.get_module("transformer_2", None),
+                           scheduler=self.get_module("scheduler"),
+                           pipeline=self,
+                       ))
+        self.add_stage(stage_name="decoding_stage",
+                       stage=DecodingStage(vae=self.get_module("vae")))
+
+    def preprocess_video_and_text_and_trajectory(self,
+                                                 fastvideo_args: FastVideoArgs,
+                                                 args):
+
+        for batch_idx, data in enumerate(self.pbar):
+            if data is None:
+                continue
+
+            with torch.inference_mode():
+                # Filter out invalid samples (those with all zeros)
+                valid_indices = []
+                for i, pixel_values in enumerate(data["pixel_values"]):
+                    if not torch.all(
+                            pixel_values == 0):  # Check if all values are zero
+                        valid_indices.append(i)
+                self.num_processed_samples += len(valid_indices)
+
+                if not valid_indices:
+                    continue
+
+                # Create new batch with only valid samples
+                valid_data = {
+                    "pixel_values":
+                    torch.stack(
+                        [data["pixel_values"][i] for i in valid_indices]),
+                    "text": [data["text"][i] for i in valid_indices],
+                    "path": [data["path"][i] for i in valid_indices],
+                    "fps": [data["fps"][i] for i in valid_indices],
+                    "duration": [data["duration"][i] for i in valid_indices],
+                }
+
+                # VAE
+                with torch.autocast("cuda", dtype=torch.float32):
+                    latents = self.get_module("vae").encode(
+                        valid_data["pixel_values"].to(
+                            get_local_torch_device())).mean
+
+                # Get extra features if needed
+                extra_features = self.get_extra_features(
+                    valid_data, fastvideo_args)
+
+                batch_captions = valid_data["text"]
+                logger.info(f"===== batch_captions: {batch_captions}")
+                # Encode text using the standalone TextEncodingStage API
+                prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                    batch_captions,
+                    fastvideo_args,
+                    encoder_index=[0],
+                    return_attention_mask=True,
+                )
+                prompt_embeds = prompt_embeds_list[0]
+                prompt_attention_masks = prompt_masks_list[0]
+                assert prompt_embeds.shape[0] == prompt_attention_masks.shape[0]
+
+                # # Get sequence lengths from attention masks (number of 1s)
+                # seq_lens = prompt_attention_mask.sum(dim=1)
+
+                # non_padded_embeds = []
+                # non_padded_masks = []
+
+                # # Process each item in the batch
+                # for i in range(prompt_embeds.size(0)):
+                #     seq_len = seq_lens[i].item()
+                #     # Slice the embeddings and masks to keep only non-padding parts
+                #     non_padded_embeds.append(prompt_embeds[i, :seq_len])
+                #     non_padded_masks.append(prompt_attention_mask[i, :seq_len])
+
+                # Update the tensors with non-padded versions
+                # prompt_embeds = non_padded_embeds
+                # prompt_attention_masks = non_padded_masks
+                # prompt_embeds = prompt_embeds
+
+                # logger.info(f"===== prompt_embeds: {prompt_embeds[0].shape}")
+                # logger.info(f"===== prompt_attention_masks: {prompt_attention_masks[0].shape}")
+
+                sampling_params = SamplingParam.from_pretrained(args.model_path)
+
+                # encode negative prompt for trajectory collection
+                if sampling_params.guidance_scale > 1 and sampling_params.negative_prompt is not None:
+                    negative_prompt_embeds_list, negative_prompt_masks_list = self.prompt_encoding_stage.encode_text(
+                        sampling_params.negative_prompt,
+                        fastvideo_args,
+                        encoder_index=[0],
+                        return_attention_mask=True,
+                    )
+                    negative_prompt_embed = negative_prompt_embeds_list[0][0]
+                    negative_prompt_attention_mask = negative_prompt_masks_list[
+                        0][0]
+                else:
+                    negative_prompt_embed = None
+                    negative_prompt_attention_mask = None
+
+                trajectory_latents = []
+                trajectory_timesteps = []
+                trajectory_decoded = []
+                for i, (prompt_embed, prompt_attention_mask) in enumerate(
+                        zip(prompt_embeds, prompt_attention_masks, strict=False)):
+                    prompt_embed = prompt_embed.unsqueeze(0)
+                    prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
+                    logger.info("what")
+                    logger.info(f"===== prompt_embed: {prompt_embed.shape}")
+                    logger.info(
+                        f"===== prompt_attention_mask: {prompt_attention_mask.shape}"
+                    )
+                    # Collect the trajectory data
+                    batch = ForwardBatch(
+                        **shallow_asdict(sampling_params),
+                        # data_type="video",
+                        # seed=args.seed,
+                        # prompt=batch_captions[i],
+                        # prompt_embeds=[prompt_embed],
+                        # prompt_attention_mask=[prompt_attention_mask],
+                        # height=args.max_height,
+                        # width=args.max_width,
+                        # num_frames=81,
+                        # fps=args.train_fps,
+                        # return_trajectory_latents=True,
+                        # guidance_scale=3.0,
+                        # do_classifier_free_guidance=True,
+                    )
+                    batch.prompt_embeds = [prompt_embed]
+                    batch.prompt_attention_mask = [prompt_attention_mask]
+                    batch.negative_prompt_embeds = [negative_prompt_embed]
+                    batch.negative_attention_mask = [
+                        negative_prompt_attention_mask
+                    ]
+                    batch.return_trajectory_latents = True
+                    batch.return_trajectory_decoded = False
+                    batch.height = args.max_height
+                    batch.width = args.max_width
+                    batch.num_inference_steps = 48
+                    # batch.num_frames = 81
+                    batch.fps = args.train_fps
+                    batch.guidance_scale = 6.0
+                    batch.do_classifier_free_guidance = True
+                    # fastvideo_args.pipeline_config.ti2v_task = True
+
+                    result_batch = self.input_validation_stage(
+                        batch, fastvideo_args)
+                    # result_batch = self.prompt_encoding_stage(result_batch, fastvideo_args)
+                    # result_batch = self.vae_encoding_stage(result_batch, fastvideo_args)
+                    result_batch = self.timestep_preparation_stage(
+                        batch, fastvideo_args)
+                    result_batch = self.latent_preparation_stage(
+                        result_batch, fastvideo_args)
+                    result_batch = self.denoising_stage(result_batch,
+                                                        fastvideo_args)
+                    result_batch = self.decoding_stage(result_batch,
+                                                       fastvideo_args)
+                    # trajectory_latents = result_batch.trajectory_latents
+                    trajectory_latents.append(
+                        result_batch.trajectory_latents.cpu())
+                    trajectory_timesteps.append(
+                        result_batch.trajectory_timesteps.cpu())
+                    trajectory_decoded.append(result_batch.trajectory_decoded)
+
+            extra_features["trajectory_latents"] = trajectory_latents
+            extra_features["trajectory_timesteps"] = trajectory_timesteps
+            logger.info(
+                f"===== trajectory_latents: {trajectory_latents[0].shape}")
+            logger.info(
+                f"===== trajectory_latents len: {len(trajectory_latents)}")
+            logger.info(f"===== trajectory_timesteps: {trajectory_timesteps}")
+            logger.info(
+                f"===== trajectory_timesteps len: {len(trajectory_timesteps)}")
+
+            if batch.return_trajectory_decoded:
+                logger.info("===== SAVING TRAJECTORY DECODED")
+                for i, decoded_frames in enumerate(trajectory_decoded):
+                    for j, decoded_frame in enumerate(decoded_frames):
+                        logger.info(
+                            f"===== SAVING TRAJECTORY DECODED {i} for prompt {batch_captions[i]}"
+                        )
+                        save_decoded_latents_as_video(
+                            decoded_frame,
+                            f"decoded_videos/trajectory_decoded_{i}_{j}.mp4",
+                            args.train_fps)
+            # assert False
+            # Prepare batch data for Parquet dataset
+            batch_data = []
+
+            # Add progress bar for saving outputs
+            save_pbar = tqdm(enumerate(valid_data["path"]),
+                             desc="Saving outputs",
+                             unit="item",
+                             leave=False)
+            for idx, video_path in save_pbar:
+                # Get the corresponding latent and info using video name
+                latent = latents[idx].cpu()
+                video_name = os.path.basename(video_path).split(".")[0]
+
+                # Convert tensors to numpy arrays
+                vae_latent = latent.cpu().numpy()
+                text_embedding = prompt_embeds[idx].cpu().numpy()
+
+                # Get extra features for this sample if needed
+                sample_extra_features = {}
+                if extra_features:
+                    for key, value in extra_features.items():
+                        logger.info(f"===== key: {key}")
+                        if isinstance(value, torch.Tensor):
+                            logger.info(f"===== value: {value[idx].shape}")
+                            sample_extra_features[key] = value[idx].cpu().numpy(
+                            )
+                        else:
+                            assert isinstance(value, list)
+                            if isinstance(value[idx], torch.Tensor):
+                                logger.info(
+                                    f"===== value in list: {value[idx].shape}")
+                                sample_extra_features[key] = value[idx].cpu(
+                                ).float().numpy()
+                            else:
+                                logger.info("===== value in list: not tensor")
+                                sample_extra_features[key] = value[idx]
+                            # logger.info(f"===== value: not tensor")
+                            # sample_extra_features[key] = value[idx]
+
+                # Create record for Parquet dataset
+                record = self.create_record(
+                    video_name=video_name,
+                    vae_latent=vae_latent,
+                    text_embedding=text_embedding,
+                    valid_data=valid_data,
+                    idx=idx,
+                    extra_features=sample_extra_features)
+                batch_data.append(record)
+
+            if batch_data:
+                # Add progress bar for writing to Parquet dataset
+                write_pbar = tqdm(total=1,
+                                  desc="Writing to Parquet dataset",
+                                  unit="batch")
+                # Convert batch data to PyArrow arrays
+                arrays = []
+                for field in self.get_schema_fields():
+                    if field.endswith('_bytes'):
+                        arrays.append(
+                            pa.array([record[field] for record in batch_data],
+                                     type=pa.binary()))
+                    elif field.endswith('_shape'):
+                        arrays.append(
+                            pa.array([record[field] for record in batch_data],
+                                     type=pa.list_(pa.int32())))
+                    elif field in ['width', 'height', 'num_frames']:
+                        arrays.append(
+                            pa.array([record[field] for record in batch_data],
+                                     type=pa.int32()))
+                    elif field in ['duration_sec', 'fps']:
+                        arrays.append(
+                            pa.array([record[field] for record in batch_data],
+                                     type=pa.float32()))
+                    else:
+                        arrays.append(
+                            pa.array([record[field] for record in batch_data]))
+
+                table = pa.Table.from_arrays(arrays,
+                                             names=self.get_schema_fields())
+                write_pbar.update(1)
+                write_pbar.close()
+
+                # Store the table in a list for later processing
+                if not hasattr(self, 'all_tables'):
+                    self.all_tables = []
+                self.all_tables.append(table)
+
+                logger.info("Collected batch with %s samples", len(table))
+
+            if self.num_processed_samples >= args.flush_frequency:
+                self._flush_tables(self.num_processed_samples, args,
+                                   self.combined_parquet_dir)
+                self.num_processed_samples = 0
+                self.all_tables = []
+
+    def get_extra_features(self, valid_data: dict[str, Any],
+                           fastvideo_args: FastVideoArgs) -> dict[str, Any]:
+
+        # TODO(will): move these to cpu at some point
+        self.get_module("vae").to(get_local_torch_device())
+
+        # generator = torch.Generator(device=get_local_torch_device(), seed=42)
+        generator = torch.Generator("cpu").manual_seed(42)
+
+        features = {}
+        """Get CLIP features from the first frame of each video."""
+        first_frame = valid_data["pixel_values"][:, :, 0, :, :].permute(
+            0, 2, 3, 1)  # (B, C, T, H, W) -> (B, H, W, C)
+        _, _, num_frames, height, width = valid_data["pixel_values"].shape
+        # latent_height = height // self.get_module(
+        #     "vae").spatial_compression_ratio
+        # latent_width = width // self.get_module("vae").spatial_compression_ratio
+
+        unprocessed_images = []
+        pil_images = []
+        # Frame has values between -1 and 1
+        for frame in first_frame:
+            frame = (frame + 1) * 127.5
+            frame_pil = Image.fromarray(frame.cpu().numpy().astype(np.uint8))
+            pil_images.append(frame_pil)
+            # processed_img = self.get_module("image_processor")(
+            #     images=frame_pil, return_tensors="pt")
+            unprocessed_images.append(frame_pil)
+        """Get VAE features from the first frame of each video"""
+        video_conditions = []
+        for frame in unprocessed_images:
+
+            latent = self.vae_encoding_stage.encode_image(
+                frame, height, width, fastvideo_args, generator)
+            video_conditions.append(latent)
+
+        features["image_condition_latents"] = video_conditions
+        features["pil_images"] = pil_images
+        return features
+
+    def create_record(
+            self,
+            video_name: str,
+            vae_latent: np.ndarray,
+            text_embedding: np.ndarray,
+            valid_data: dict[str, Any],
+            idx: int,
+            extra_features: dict[str, Any] | None = None) -> dict[str, Any]:
+        """Create a record for the Parquet dataset with CLIP features."""
+        record = super().create_record(video_name=video_name,
+                                       vae_latent=vae_latent,
+                                       text_embedding=text_embedding,
+                                       valid_data=valid_data,
+                                       idx=idx,
+                                       extra_features=extra_features)
+
+        if extra_features and "image_condition_latents" in extra_features:
+            image_condition_latents = extra_features["image_condition_latents"]
+            record.update({
+                "image_condition_latents_bytes":
+                image_condition_latents.tobytes(),
+                "image_condition_latents_shape":
+                list(image_condition_latents.shape),
+                "image_condition_latents_dtype":
+                str(image_condition_latents.dtype),
+            })
+        else:
+            record.update({
+                "image_condition_latents_bytes": b"",
+                "image_condition_latents_shape": [],
+                "image_condition_latents_dtype": "",
+            })
+
+        if extra_features and "trajectory_latents" in extra_features:
+            trajectory_latents = extra_features["trajectory_latents"]
+            record.update({
+                "trajectory_latents_bytes":
+                trajectory_latents.tobytes(),
+                "trajectory_latents_shape":
+                list(trajectory_latents.shape),
+                "trajectory_latents_dtype":
+                str(trajectory_latents.dtype),
+            })
+        else:
+            record.update({
+                "trajectory_latents_bytes": b"",
+                "trajectory_latents_shape": [],
+                "trajectory_latents_dtype": "",
+            })
+
+        if extra_features and "trajectory_timesteps" in extra_features:
+            trajectory_timesteps = extra_features["trajectory_timesteps"]
+            record.update({
+                "trajectory_timesteps_bytes":
+                trajectory_timesteps.tobytes(),
+                "trajectory_timesteps_shape":
+                list(trajectory_timesteps.shape),
+                "trajectory_timesteps_dtype":
+                str(trajectory_timesteps.dtype),
+            })
+        else:
+            record.update({
+                "trajectory_timesteps_bytes": b"",
+                "trajectory_timesteps_shape": [],
+                "trajectory_timesteps_dtype": "",
+            })
+
+        if extra_features and "pil_image" in extra_features:
+            pil_image = extra_features["pil_image"]
+            record.update({
+                "pil_image_bytes": pil_image.tobytes(),
+                "pil_image_shape": list(pil_image.shape),
+                "pil_image_dtype": str(pil_image.dtype),
+            })
+        else:
+            record.update({
+                "pil_image_bytes": b"",
+                "pil_image_shape": [],
+                "pil_image_dtype": "",
+            })
+
+        return record
+
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs, args):
+        if not self.post_init_called:
+            self.post_init()
+
+        self.local_rank = int(os.getenv("RANK", 0))
+        os.makedirs(args.output_dir, exist_ok=True)
+        # Create directory for combined data
+        self.combined_parquet_dir = os.path.join(args.output_dir,
+                                                 "combined_parquet_dataset")
+        os.makedirs(self.combined_parquet_dir, exist_ok=True)
+
+        # Loading dataset
+        train_dataset = getdataset(args)
+
+        self.preprocess_dataloader = DataLoader(
+            train_dataset,
+            batch_size=args.preprocess_video_batch_size,
+            num_workers=args.dataloader_num_workers,
+        )
+
+        self.preprocess_loader_iter = iter(self.preprocess_dataloader)
+
+        self.num_processed_samples = 0
+        # Add progress bar for video preprocessing
+        self.pbar = tqdm(self.preprocess_loader_iter,
+                         desc="Processing videos",
+                         unit="batch",
+                         disable=self.local_rank != 0)
+
+        # Initialize class variables for data sharing
+        self.video_data: dict[str, Any] = {}  # Store video metadata and paths
+        self.latent_data: dict[str, Any] = {}  # Store latent tensors
+        self.preprocess_video_and_text_and_trajectory(fastvideo_args, args)
+
+
+EntryClass = PreprocessPipeline_ODE_Trajectory

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -9,6 +9,8 @@ from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.preprocess.preprocess_pipeline_i2v import (
     PreprocessPipeline_I2V)
+from fastvideo.pipelines.preprocess.preprocess_pipeline_ode_trajectory import (
+    PreprocessPipeline_ODE_Trajectory)
 from fastvideo.pipelines.preprocess.preprocess_pipeline_t2v import (
     PreprocessPipeline_T2V)
 from fastvideo.utils import maybe_download_model
@@ -24,7 +26,8 @@ def main(args) -> None:
     pipeline_config = PipelineConfig.from_pretrained(args.model_path)
     kwargs = {
         "vae_precision": "fp32",
-        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
+        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
+        "flow_shift": 5,
     }
     pipeline_config.update_config_from_dict(kwargs)
     fastvideo_args = FastVideoArgs(
@@ -35,7 +38,19 @@ def main(args) -> None:
         text_encoder_cpu_offload=False,
         pipeline_config=pipeline_config,
     )
-    PreprocessPipeline = PreprocessPipeline_I2V if args.preprocess_task == "i2v" else PreprocessPipeline_T2V
+    if args.preprocess_task == "t2v":
+        PreprocessPipeline = PreprocessPipeline_T2V
+    elif args.preprocess_task == "i2v":
+        PreprocessPipeline = PreprocessPipeline_I2V
+    elif args.preprocess_task == "ode_trajectory":
+        PreprocessPipeline = PreprocessPipeline_ODE_Trajectory
+    else:
+        raise ValueError(f"Invalid preprocess task: {args.preprocess_task}")
+
+    logger.info(
+        f"Preprocess task: {args.preprocess_task} using {PreprocessPipeline.__name__}"
+    )
+
     pipeline = PreprocessPipeline(args.model_path, fastvideo_args)
     pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
 

fastvideo/pipelines/stages/__init__.py

@@ -7,6 +7,7 @@ complete diffusion pipelines.
 """
 
 from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.causal_denoising import CausalDMDDenosingStage
 from fastvideo.pipelines.stages.conditioning import ConditioningStage
 from fastvideo.pipelines.stages.decoding import DecodingStage
 from fastvideo.pipelines.stages.denoising import (DenoisingStage,
@@ -30,6 +31,7 @@ __all__ = [
     "ConditioningStage",
     "DenoisingStage",
     "DmdDenoisingStage",
+    "CausalDMDDenosingStage",
     "EncodingStage",
     "DecodingStage",
     "ImageEncodingStage",

fastvideo/pipelines/stages/causal_denoising.py

@@ -0,0 +1,414 @@
+import torch  # type: ignore
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler)
+from fastvideo.models.utils import pred_noise_to_pred_video
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising import DenoisingStage
+
+try:
+    from fastvideo.attention.backends.sliding_tile_attn import (
+        SlidingTileAttentionBackend)
+    st_attn_available = True
+except ImportError:
+    st_attn_available = False
+    SlidingTileAttentionBackend = None  # type: ignore
+
+try:
+    from fastvideo.attention.backends.video_sparse_attn import (
+        VideoSparseAttentionBackend)
+    vsa_available = True
+except ImportError:
+    vsa_available = False
+    VideoSparseAttentionBackend = None  # type: ignore
+
+logger = init_logger(__name__)
+
+
+class CausalDMDDenosingStage(DenoisingStage):
+    """
+    Denoising stage for causal diffusion.
+    """
+
+    def __init__(self, transformer, scheduler) -> None:
+        super().__init__(transformer, scheduler)
+        self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0)
+        # KV and cross-attention cache state (initialized on first forward)
+        self.kv_cache1: list | None = None
+        self.crossattn_cache: list | None = None
+        # Model-dependent constants (aligned with causal_inference.py assumptions)
+        self.num_transformer_blocks = self.transformer.config.arch_config.num_layers
+        self.num_frames_per_block = self.transformer.config.arch_config.num_frames_per_block
+        self.sliding_window_num_frames = self.transformer.config.arch_config.sliding_window_num_frames
+
+        try:
+            self.local_attn_size = getattr(self.transformer.model,
+                                           "local_attn_size",
+                                           -1)  # type: ignore
+        except Exception:
+            self.local_attn_size = -1
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        target_dtype = torch.bfloat16
+        autocast_enabled = (target_dtype != torch.float32
+                            ) and not fastvideo_args.disable_autocast
+
+        latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2]
+        patch_ratio = self.transformer.config.arch_config.patch_size[
+            -1] * self.transformer.config.arch_config.patch_size[-2]
+        self.frame_seq_length = latent_seq_length // patch_ratio
+        # TODO(will): make this a parameter once we add i2v support
+        independent_first_frame = self.transformer.independent_first_frame
+
+        # Timesteps for DMD
+        timesteps = torch.tensor(
+            fastvideo_args.pipeline_config.dmd_denoising_steps,
+            dtype=torch.long,
+            device=get_local_torch_device())
+
+        # Image kwargs (kept empty unless caller provides compatible args)
+        image_kwargs: dict = {}
+
+        pos_cond_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                # "encoder_hidden_states_2": batch.clip_embedding_pos,
+                "encoder_attention_mask": batch.prompt_attention_mask,
+            },
+        )
+
+        # STA
+        if st_attn_available and self.attn_backend == SlidingTileAttentionBackend:
+            self.prepare_sta_param(batch, fastvideo_args)
+
+        # Latents and prompts
+        assert batch.latents is not None, "latents must be provided"
+        latents = batch.latents  # [B, C, T, H, W]
+        b, c, t, h, w = latents.shape
+        prompt_embeds = batch.prompt_embeds
+        assert torch.isnan(prompt_embeds[0]).sum() == 0
+
+        # Initialize or reset caches
+        if self.kv_cache1 is None:
+            self._initialize_kv_cache(batch_size=latents.shape[0],
+                                      dtype=target_dtype,
+                                      device=latents.device)
+            self._initialize_crossattn_cache(
+                batch_size=latents.shape[0],
+                max_text_len=fastvideo_args.pipeline_config.
+                text_encoder_configs[0].arch_config.text_len,
+                dtype=target_dtype,
+                device=latents.device)
+        else:
+            assert self.crossattn_cache is not None
+            # reset cross-attention cache
+            for block_index in range(self.num_transformer_blocks):
+                self.crossattn_cache[block_index][
+                    "is_init"] = False  # type: ignore
+            # reset kv cache pointers
+            for block_index in range(len(self.kv_cache1)):
+                self.kv_cache1[block_index][
+                    "global_end_index"] = torch.tensor(  # type: ignore
+                        [0],
+                        dtype=torch.long,
+                        device=latents.device)
+                self.kv_cache1[block_index][
+                    "local_end_index"] = torch.tensor(  # type: ignore
+                        [0],
+                        dtype=torch.long,
+                        device=latents.device)
+
+        # Optional: cache context features from provided image latents prior to generation
+        current_start_frame = 0
+        if getattr(batch, "image_latent", None) is not None:
+            image_latent = batch.image_latent
+            assert image_latent is not None
+            input_frames = image_latent.shape[2]
+            # timestep zero (or configured context noise) for cache warm-up
+            t_zero = torch.zeros([latents.shape[0]],
+                                 device=latents.device,
+                                 dtype=torch.long)
+            if independent_first_frame and input_frames >= 1:
+                # warm-up with the very first frame independently
+                image_first_btchw = image_latent[:, :, :1, :, :].to(
+                    target_dtype).permute(0, 2, 1, 3, 4)
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled):
+                    _ = self.transformer(
+                        image_first_btchw,
+                        prompt_embeds,
+                        t_zero,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame *
+                        self.frame_seq_length,
+                        **image_kwargs,
+                        **pos_cond_kwargs,
+                    )
+                current_start_frame += 1
+                remaining_frames = input_frames - 1
+            else:
+                remaining_frames = input_frames
+
+            # process remaining input frames in blocks of num_frame_per_block
+            while remaining_frames > 0:
+                block = min(self.num_frames_per_block, remaining_frames)
+                ref_btchw = image_latent[:, :, current_start_frame:
+                                         current_start_frame +
+                                         block, :, :].to(target_dtype).permute(
+                                             0, 2, 1, 3, 4)
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled):
+                    _ = self.transformer(
+                        ref_btchw,
+                        prompt_embeds,
+                        t_zero,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=current_start_frame *
+                        self.frame_seq_length,
+                        **image_kwargs,
+                        **pos_cond_kwargs,
+                    )
+                current_start_frame += block
+                remaining_frames -= block
+
+        # Base position offset from any cache warm-up
+        pos_start_base = current_start_frame
+
+        # Determine block sizes
+        if not independent_first_frame or (independent_first_frame
+                                           and batch.image_latent is not None):
+            if t % self.num_frames_per_block != 0:
+                raise ValueError(
+                    "num_frames must be divisible by num_frames_per_block for causal DMD denoising"
+                )
+            num_blocks = t // self.num_frames_per_block
+            block_sizes = [self.num_frames_per_block] * num_blocks
+            start_index = 0
+        else:
+            if (t - 1) % self.num_frames_per_block != 0:
+                raise ValueError(
+                    "(num_frames - 1) must be divisible by num_frame_per_block when independent_first_frame=True"
+                )
+            num_blocks = (t - 1) // self.num_frames_per_block
+            block_sizes = [1] + [self.num_frames_per_block] * num_blocks
+            start_index = 0
+
+        # DMD loop in causal blocks
+        with self.progress_bar(total=len(block_sizes) *
+                               len(timesteps)) as progress_bar:
+            for current_num_frames in block_sizes:
+                current_latents = latents[:, :, start_index:start_index +
+                                          current_num_frames, :, :]
+                # use BTCHW for DMD conversion routines
+                noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4)
+                video_raw_latent_shape = noise_latents_btchw.shape
+
+                for i, t_cur in enumerate(timesteps):
+                    # Copy for pred conversion
+                    noise_latents = noise_latents_btchw.clone()
+                    latent_model_input = current_latents.to(target_dtype)
+
+                    if batch.image_latent is not None and independent_first_frame and start_index == 0:
+                        latent_model_input = torch.cat([
+                            latent_model_input,
+                            batch.image_latent.to(target_dtype)
+                        ],
+                                                       dim=2)
+
+                    # Prepare inputs
+                    t_expand = t_cur.expand(latent_model_input.shape[0])
+                    # t_expand = t_cur * torch.ones((latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long)
+                    # t_expand = t_expand.repeat(1, self.sliding_window_num_frames)
+
+
+                    # Attention metadata if needed
+                    if (vsa_available and self.attn_backend
+                            == VideoSparseAttentionBackend):
+                        self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls(
+                        )
+                        if self.attn_metadata_builder_cls is not None:
+                            self.attn_metadata_builder = self.attn_metadata_builder_cls(
+                            )
+                            attn_metadata = self.attn_metadata_builder.build(  # type: ignore
+                                current_timestep=i,  # type: ignore
+                                raw_latent_shape=(current_num_frames, h,
+                                                  w),  # type: ignore
+                                patch_size=fastvideo_args.pipeline_config.
+                                dit_config.patch_size,  # type: ignore
+                                STA_param=batch.STA_param,  # type: ignore
+                                VSA_sparsity=fastvideo_args.
+                                VSA_sparsity,  # type: ignore
+                                device=get_local_torch_device(),  # type: ignore
+                            )  # type: ignore
+                            assert attn_metadata is not None, "attn_metadata cannot be None"
+                        else:
+                            attn_metadata = None
+                    else:
+                        attn_metadata = None
+
+                    with torch.autocast(device_type="cuda",
+                                        dtype=target_dtype,
+                                        enabled=autocast_enabled), \
+                        set_forward_context(current_timestep=i,
+                                            attn_metadata=attn_metadata,
+                                            forward_batch=batch):
+                        # Run transformer; follow DMD stage pattern
+                        t_expanded_noise= t_cur * torch.ones((latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long)
+                        pred_noise_btchw = self.transformer(
+                            latent_model_input,
+                            prompt_embeds,
+                            t_expanded_noise,
+                            kv_cache=self.kv_cache1,
+                            crossattn_cache=self.crossattn_cache,
+                            current_start=(pos_start_base + start_index) *
+                            self.frame_seq_length,
+                            start_frame=start_index,
+                            **image_kwargs,
+                            **pos_cond_kwargs,
+                        ).permute(0, 2, 1, 3, 4)
+
+                    # Convert pred noise to pred video with FM Euler scheduler utilities
+                    pred_video_btchw = pred_noise_to_pred_video(
+                        pred_noise=pred_noise_btchw.flatten(0, 1),
+                        noise_input_latent=noise_latents.flatten(0, 1),
+                        timestep=t_expand,
+                        scheduler=self.scheduler).unflatten(
+                            0, pred_noise_btchw.shape[:2])
+
+                    if i < len(timesteps) - 1:
+                        next_timestep = timesteps[i + 1] * torch.ones(
+                            [1],
+                            dtype=torch.long,
+                            device=pred_video_btchw.device)
+                        noise = torch.randn(
+                            video_raw_latent_shape,
+                            dtype=pred_video_btchw.dtype,
+                            generator=(batch.generator[0] if isinstance(
+                                batch.generator, list) else
+                                       batch.generator)).to(self.device)
+                        noise_btchw = noise
+                        noise_latents_btchw = self.scheduler.add_noise(
+                            pred_video_btchw.flatten(0, 1),
+                            noise_btchw.flatten(0, 1),
+                            next_timestep).unflatten(0,
+                                                     pred_video_btchw.shape[:2])
+                        current_latents = noise_latents_btchw.permute(
+                            0, 2, 1, 3, 4)
+                    else:
+                        current_latents = pred_video_btchw.permute(
+                            0, 2, 1, 3, 4)
+
+                    if progress_bar is not None:
+                        progress_bar.update()
+
+                # Write back and advance
+                latents[:, :, start_index:start_index +
+                        current_num_frames, :, :] = current_latents
+
+                # Re-run with context timestep to update KV cache using clean context
+                context_noise = getattr(fastvideo_args.pipeline_config,
+                                        "context_noise", 0)
+                t_context = torch.ones([latents.shape[0]],
+                                       device=latents.device,
+                                       dtype=torch.long) * int(context_noise)
+                context_bcthw = current_latents.to(target_dtype)
+                with torch.autocast(device_type="cuda",
+                                    dtype=target_dtype,
+                                    enabled=autocast_enabled), \
+                    set_forward_context(current_timestep=0,
+                                        attn_metadata=attn_metadata,
+                                        forward_batch=batch):
+                    t_expanded_context = t_context * torch.ones((context_bcthw.shape[0], 1), device=context_bcthw.device, dtype=torch.long)
+                    _ = self.transformer(
+                        context_bcthw,
+                        prompt_embeds,
+                        t_expanded_context,
+                        kv_cache=self.kv_cache1,
+                        crossattn_cache=self.crossattn_cache,
+                        current_start=(pos_start_base + start_index) *
+                        self.frame_seq_length,
+                        start_frame=start_index,
+                        **image_kwargs,
+                        **pos_cond_kwargs,
+                    )
+                start_index += current_num_frames
+
+        batch.latents = latents
+        return batch
+
+    def _initialize_kv_cache(self, batch_size, dtype, device) -> None:
+        """
+        Initialize a Per-GPU KV cache aligned with the Wan model assumptions.
+        """
+        kv_cache1 = []
+        num_attention_heads = self.transformer.num_attention_heads
+        attention_head_dim = self.transformer.attention_head_dim
+        if self.local_attn_size != -1:
+            kv_cache_size = self.local_attn_size * self.frame_seq_length
+        else:
+            kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames
+
+        for _ in range(self.num_transformer_blocks):
+            kv_cache1.append({
+                "k":
+                torch.zeros([
+                    batch_size, kv_cache_size, num_attention_heads,
+                    attention_head_dim
+                ],
+                            dtype=dtype,
+                            device=device),
+                "v":
+                torch.zeros([
+                    batch_size, kv_cache_size, num_attention_heads,
+                    attention_head_dim
+                ],
+                            dtype=dtype,
+                            device=device),
+                "global_end_index":
+                torch.tensor([0], dtype=torch.long, device=device),
+                "local_end_index":
+                torch.tensor([0], dtype=torch.long, device=device),
+            })
+
+        self.kv_cache1 = kv_cache1
+
+    def _initialize_crossattn_cache(self, batch_size, max_text_len, dtype,
+                                    device) -> None:
+        """
+        Initialize a Per-GPU cross-attention cache aligned with the Wan model assumptions.
+        """
+        crossattn_cache = []
+        num_attention_heads = self.transformer.num_attention_heads
+        attention_head_dim = self.transformer.attention_head_dim
+        for _ in range(self.num_transformer_blocks):
+            crossattn_cache.append({
+                "k":
+                torch.zeros([
+                    batch_size, max_text_len, num_attention_heads,
+                    attention_head_dim
+                ],
+                            dtype=dtype,
+                            device=device),
+                "v":
+                torch.zeros([
+                    batch_size, max_text_len, num_attention_heads,
+                    attention_head_dim
+                ],
+                            dtype=dtype,
+                            device=device),
+                "is_init":
+                False,
+            })
+        self.crossattn_cache = crossattn_cache

fastvideo/pipelines/stages/decoding.py

@@ -50,6 +50,50 @@ class DecodingStage(PipelineStage):
         result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)])
         return result
 
+    @torch.no_grad()
+    def decode(self, latents: torch.Tensor,
+               fastvideo_args: FastVideoArgs) -> torch.Tensor:
+        """Decode latents into pixel space."""
+        self.vae = self.vae.to(get_local_torch_device())
+        latents = latents.to(get_local_torch_device())
+
+        # Setup VAE precision
+        vae_dtype = PRECISION_TO_TYPE[
+            fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (
+            vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        if isinstance(self.vae.scaling_factor, torch.Tensor):
+            latents = latents / self.vae.scaling_factor.to(
+                latents.device, latents.dtype)
+        else:
+            latents = latents / self.vae.scaling_factor
+
+        # Apply shifting if needed
+        if (hasattr(self.vae, "shift_factor")
+                and self.vae.shift_factor is not None):
+            if isinstance(self.vae.shift_factor, torch.Tensor):
+                latents += self.vae.shift_factor.to(latents.device,
+                                                    latents.dtype)
+            else:
+                latents += self.vae.shift_factor
+
+        # Decode latents
+        with torch.autocast(device_type="cuda",
+                            dtype=vae_dtype,
+                            enabled=vae_autocast_enabled):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+            # if fastvideo_args.vae_sp:
+            #     self.vae.enable_parallel()
+            if not vae_autocast_enabled:
+                latents = latents.to(vae_dtype)
+            image = self.vae.decode(latents)
+
+        # Normalize image to [0, 1] range
+        image = (image / 2 + 0.5).clamp(0, 1)
+        return image
+
     @torch.no_grad()
     def forward(
         self,
@@ -66,6 +110,7 @@ class DecodingStage(PipelineStage):
         Returns:
             The batch with decoded outputs.
         """
+        # load vae if not already loaded (used for memory constrained devices)
         pipeline = self.pipeline() if self.pipeline else None
         if not fastvideo_args.model_loaded["vae"]:
             loader = VAELoader()
@@ -75,58 +120,31 @@ class DecodingStage(PipelineStage):
                 pipeline.add_module("vae", self.vae)
             fastvideo_args.model_loaded["vae"] = True
 
-        self.vae = self.vae.to(get_local_torch_device())
-
-        latents = batch.latents
-        # TODO(will): remove this once we add input/output validation for stages
-        if latents is None:
-            raise ValueError("Latents must be provided")
-
-        # Skip decoding if output type is latent
         if fastvideo_args.output_type == "latent":
-            image = latents
+            frames = batch.latents
         else:
-            # Setup VAE precision
-            vae_dtype = PRECISION_TO_TYPE[
-                fastvideo_args.pipeline_config.vae_precision]
-            vae_autocast_enabled = (vae_dtype != torch.float32
-                                    ) and not fastvideo_args.disable_autocast
+            frames = self.decode(batch.latents, fastvideo_args)
 
-            if isinstance(self.vae.scaling_factor, torch.Tensor):
-                latents = latents / self.vae.scaling_factor.to(
-                    latents.device, latents.dtype)
-            else:
-                latents = latents / self.vae.scaling_factor
-
-            # Apply shifting if needed
-            if (hasattr(self.vae, "shift_factor")
-                    and self.vae.shift_factor is not None):
-                if isinstance(self.vae.shift_factor, torch.Tensor):
-                    latents += self.vae.shift_factor.to(latents.device,
-                                                        latents.dtype)
-                else:
-                    latents += self.vae.shift_factor
-
-            # Decode latents
-            with torch.autocast(device_type="cuda",
-                                dtype=vae_dtype,
-                                enabled=vae_autocast_enabled):
-                if fastvideo_args.pipeline_config.vae_tiling:
-                    self.vae.enable_tiling()
-                # if fastvideo_args.vae_sp:
-                #     self.vae.enable_parallel()
-                if not vae_autocast_enabled:
-                    latents = latents.to(vae_dtype)
-                image = self.vae.decode(latents)
-
-        # Normalize image to [0, 1] range
-        image = (image / 2 + 0.5).clamp(0, 1)
+        # decode trajectory latents if needed
+        if batch.return_trajectory_decoded:
+            batch.trajectory_decoded = []
+            logger.info(f"batch.trajectory_latents.shape: {batch.trajectory_latents.shape}")
+            assert batch.trajectory_latents is not None, "batch should have trajectory latents"
+            for idx in range(batch.trajectory_latents.shape[1]):
+                # bathc.trajectory_latents is [batch_size, timesteps, channels, frames, height, width]
+                cur_latent = batch.trajectory_latents[:, idx, :, :, :, :]
+                logger.info(f"cur_latent.shape: {cur_latent.shape}")
+                cur_timestep = batch.trajectory_timesteps[idx]
+                logger.info(
+                    f"decoding trajectory latent for timestep: {cur_timestep}")
+                decoded_frames = self.decode(cur_latent, fastvideo_args)
+                batch.trajectory_decoded.append(decoded_frames.cpu().float())
 
         # Convert to CPU float32 for compatibility
-        image = image.cpu().float()
+        frames = frames.cpu().float()
 
         # Update batch with decoded image
-        batch.output = image
+        batch.output = frames
 
         # Offload models if needed
         if hasattr(self, 'maybe_free_model_hooks'):

fastvideo/pipelines/stages/denoising.py

@@ -4,6 +4,7 @@ Denoising stage for diffusion pipelines.
 """
 
 import inspect
+import math
 import weakref
 from collections.abc import Iterable
 from typing import Any
@@ -24,12 +25,13 @@ from fastvideo.logger import init_logger
 from fastvideo.models.loader.component_loader import TransformerLoader
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler)
+from fastvideo.models.utils import pred_noise_to_pred_video
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.platforms import AttentionBackendEnum
-from fastvideo.utils import dict_to_3d_list
+from fastvideo.utils import dict_to_3d_list, masks_like
 
 try:
     from fastvideo.attention.backends.sliding_tile_attn import (
@@ -60,11 +62,13 @@ class DenoisingStage(PipelineStage):
                  transformer,
                  scheduler,
                  pipeline=None,
-                 transformer_2=None) -> None:
+                 transformer_2=None,
+                 vae=None) -> None:
         super().__init__()
         self.transformer = transformer
         self.transformer_2 = transformer_2
         self.scheduler = scheduler
+        self.vae = vae
         self.pipeline = weakref.ref(pipeline) if pipeline else None
         attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
         self.attn_backend = get_attn_backend(
@@ -128,11 +132,12 @@ class DenoisingStage(PipelineStage):
             latents = latents[:, :, rank_in_sp_group, :, :, :]
             batch.latents = latents
             if batch.image_latent is not None:
-                image_latent = rearrange(batch.image_latent,
-                                         "b c (n t) h w -> b c n t h w",
-                                         n=sp_world_size).contiguous()
-                image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
-                batch.image_latent = image_latent
+                if not fastvideo_args.pipeline_config.ti2v_task and not fastvideo_args.pipeline_config.t2v_as_i2v_task:
+                    image_latent = rearrange(batch.image_latent,
+                                             "b c (n t) h w -> b c n t h w",
+                                             n=sp_world_size).contiguous()
+                    image_latent = image_latent[:, :, rank_in_sp_group, :, :, :]
+                    batch.image_latent = image_latent
         # Get timesteps and calculate warmup steps
         timesteps = batch.timesteps
         # TODO(will): remove this once we add input/output validation for stages
@@ -193,6 +198,47 @@ class DenoisingStage(PipelineStage):
             boundary_timestep = fastvideo_args.boundary_ratio * self.scheduler.num_train_timesteps
         else:
             boundary_timestep = None
+        latent_model_input = latents.to(target_dtype)
+        assert latent_model_input.shape[0] == 1, "only support batch size 1"
+
+        if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+            # TI2V directly replaces the first frame of the latent with
+            # the image latent instead of appending along the channel dim
+            assert batch.image_latent is None, "TI2V task should not have image latents"
+            assert self.vae is not None, "VAE is not provided for TI2V task"
+            z = self.vae.encode(batch.pil_image).mean.float()
+            if (hasattr(self.vae, "shift_factor")
+                    and self.vae.shift_factor is not None):
+                if isinstance(self.vae.shift_factor, torch.Tensor):
+                    z -= self.vae.shift_factor.to(z.device, z.dtype)
+                else:
+                    z -= self.vae.shift_factor
+
+            if isinstance(self.vae.scaling_factor, torch.Tensor):
+                z = z * self.vae.scaling_factor.to(z.device, z.dtype)
+            else:
+                z = z * self.vae.scaling_factor
+
+            latent_model_input = latent_model_input.squeeze(0)
+            _, mask2 = masks_like([latent_model_input], zero=True)
+
+            latent_model_input = (1. -
+                                  mask2[0]) * z + mask2[0] * latent_model_input
+            # latent_model_input = latent_model_input.unsqueeze(0)
+            latent_model_input = latent_model_input.to(get_local_torch_device())
+            latents = latent_model_input
+            F = batch.num_frames
+            temporal_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_temporal
+            spatial_scale = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
+            patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
+            seq_len = ((F - 1) // temporal_scale +
+                       1) * (batch.height // spatial_scale) * (
+                           batch.width // spatial_scale) // (patch_size[1] *
+                                                             patch_size[2])
+            seq_len = int(math.ceil(seq_len / sp_world_size)) * sp_world_size
+
+        trajectory_timesteps: list[int] = []
+        trajectory_latents: list[torch.Tensor] = []
 
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
@@ -217,19 +263,52 @@ class DenoisingStage(PipelineStage):
                         self.transformer.to('cpu')
                     current_model = self.transformer_2
                     current_guidance_scale = batch.guidance_scale_2
+                assert current_model is not None, "current_model is None"
 
                 # Expand latents for I2V
                 latent_model_input = latents.to(target_dtype)
-                if batch.image_latent is not None:
+                if batch.image_latent is not None and not fastvideo_args.pipeline_config.t2v_as_i2v_task:
+                    assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task"
                     latent_model_input = torch.cat(
                         [latent_model_input, batch.image_latent],
                         dim=1).to(target_dtype)
+                elif batch.image_latent is not None and fastvideo_args.pipeline_config.t2v_as_i2v_task:
+                    assert batch.image_latent is not None, "image latents should be provided for T2V to I2V task"
+                    if rank_in_sp_group == 0:
+                        logger.info("latent_model_input.shape: %s",
+                                    latent_model_input.shape)
+                        latent_model_input = torch.cat([
+                            batch.image_latent,
+                            latent_model_input[:, :, 1:, :, :],
+                        ],
+                                                       dim=2).to(target_dtype)
+                        logger.info("latent_model_input.shape: %s",
+                                    latent_model_input.shape)
+
+                if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+                    timestep = torch.stack([t]).to(get_local_torch_device())
+                    temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
+                    temp_ts = torch.cat([
+                        temp_ts,
+                        temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep
+                    ])
+                    timestep = temp_ts.unsqueeze(0)
+                    t_expand = timestep.repeat(latent_model_input.shape[0], 1)
+                else:
+                    t_expand = t.repeat(latent_model_input.shape[0])
+
                 assert torch.isnan(latent_model_input).sum() == 0
                 latent_model_input = self.scheduler.scale_model_input(
                     latent_model_input, t)
+                if fastvideo_args.pipeline_config.t2v_as_i2v_task:
+                    if rank_in_sp_group == 0:
+                        latent_model_input = torch.cat([
+                            batch.image_latent,
+                            latent_model_input[:, :, 1:, :, :],
+                        ],
+                                                       dim=2).to(target_dtype)
 
                 # Prepare inputs for transformer
-                t_expand = t.repeat(latent_model_input.shape[0])
                 guidance_expand = (
                     torch.tensor(
                         [fastvideo_args.pipeline_config.embedded_cfg_scale] *
@@ -329,6 +408,17 @@ class DenoisingStage(PipelineStage):
                                                   latents,
                                                   **extra_step_kwargs,
                                                   return_dict=False)[0]
+                    if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+                        latents = latents.squeeze(0)
+                        latents = (1. - mask2[0]) * z + mask2[0] * latents
+                        # latents = latents.unsqueeze(0)
+
+                # save trajectory latents if needed
+                if batch.return_trajectory_latents:
+                    trajectory_timesteps.append(t)
+                    # trajectory_latents.append(latents.cpu())
+                    trajectory_latents.append(latents)
+
                 # Update progress bar
                 if i == len(timesteps) - 1 or (
                     (i + 1) > num_warmup_steps and
@@ -337,8 +427,32 @@ class DenoisingStage(PipelineStage):
                     progress_bar.update()
 
         # Gather results if using sequence parallelism
+        trajectory_tensor: torch.Tensor | None = None
+        if trajectory_latents:
+            trajectory_tensor = torch.stack(trajectory_latents, dim=1)
+        else:
+            trajectory_tensor = None
+
         if sp_group:
             latents = sequence_model_parallel_all_gather(latents, dim=2)
+            if batch.return_trajectory_latents:
+                # logger.info("before stack trajectory_latents.shape: %s", trajectory_latents[0].shape)
+                logger.info("after stack trajectory_latents.shape: %s", trajectory_tensor.shape)
+                trajectory_tensor = trajectory_tensor.to(
+                    get_local_torch_device())
+                trajectory_tensor = sequence_model_parallel_all_gather(
+                    trajectory_tensor, dim=3)
+
+        if trajectory_tensor is not None:
+            batch.trajectory_timesteps = torch.tensor(trajectory_timesteps).cpu()
+            batch.trajectory_latents = trajectory_tensor.cpu()
+
+        if fastvideo_args.pipeline_config.t2v_as_i2v_task:
+            latents = torch.cat([
+                batch.image_latent,
+                latents[:, :, 1:, :, :],
+            ],
+                                dim=2)
 
         # Update batch with final latents
         batch.latents = latents
@@ -778,8 +892,6 @@ class DmdDenoisingStage(DenoisingStage):
                             **pos_cond_kwargs,
                         ).permute(0, 2, 1, 3, 4)
 
-                    from fastvideo.training.training_utils import (
-                        pred_noise_to_pred_video)
                     pred_video = pred_noise_to_pred_video(
                         pred_noise=pred_noise.flatten(0, 1),
                         noise_input_latent=noise_latents.flatten(0, 1),

fastvideo/pipelines/stages/image_encoding.py

@@ -105,6 +105,81 @@ class ImageVAEEncodingStage(PipelineStage):
     def __init__(self, vae: ParallelTiledVAE) -> None:
         self.vae: ParallelTiledVAE = vae
 
+    def encode_image(self,
+                     image: PIL.Image.Image,
+                     height: int,
+                     width: int,
+                     fastvideo_args: FastVideoArgs,
+                     generator: torch.Generator | None = None) -> torch.Tensor:
+        """
+        Encode image into latent space.
+        """
+        image = self.preprocess(
+            image,
+            vae_scale_factor=self.vae.spatial_compression_ratio,
+            height=height,
+            width=width).to(get_local_torch_device(), dtype=torch.float32)
+
+        # (B, C, H, W) -> (B, C, 1, H, W)
+        print(f"image.shape: {image.shape}")
+        image = image.unsqueeze(2)
+        print(f"after unsqueeze image.shape: {image.shape}")
+        return self.encode_tensor(image, fastvideo_args, generator)
+
+    def encode_tensor(self,
+                      video_condition: torch.Tensor,
+                      fastvideo_args: FastVideoArgs,
+                      generator: torch.Generator | None = None) -> torch.Tensor:
+        """
+        Encode frames into latent space.
+        """
+        self.vae = self.vae.to(get_local_torch_device())
+        video_condition = video_condition.to(device=get_local_torch_device(),
+                                             dtype=torch.float32)
+
+        # Setup VAE precision
+        vae_dtype = PRECISION_TO_TYPE[
+            fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (
+            vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        # Encode Image
+        with torch.autocast(device_type="cuda",
+                            dtype=vae_dtype,
+                            enabled=vae_autocast_enabled):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+            # if fastvideo_args.vae_sp:
+            #     self.vae.enable_parallel()
+            if not vae_autocast_enabled:
+                video_condition = video_condition.to(vae_dtype)
+            encoder_output = self.vae.encode(video_condition)
+
+        if fastvideo_args.mode == ExecutionMode.PREPROCESS:
+            latent_condition = encoder_output.mean
+        else:
+            generator = generator
+            if generator is None:
+                raise ValueError("Generator must be provided")
+            latent_condition = self.retrieve_latents(encoder_output, generator)
+
+        # Apply shifting if needed
+        if (hasattr(self.vae, "shift_factor")
+                and self.vae.shift_factor is not None):
+            if isinstance(self.vae.shift_factor, torch.Tensor):
+                latent_condition -= self.vae.shift_factor.to(
+                    latent_condition.device, latent_condition.dtype)
+            else:
+                latent_condition -= self.vae.shift_factor
+
+        if isinstance(self.vae.scaling_factor, torch.Tensor):
+            latent_condition = latent_condition * self.vae.scaling_factor.to(
+                latent_condition.device, latent_condition.dtype)
+        else:
+            latent_condition = latent_condition * self.vae.scaling_factor
+
+        return latent_condition
+
     def forward(
         self,
         batch: ForwardBatch,
@@ -157,58 +232,29 @@ class ImageVAEEncodingStage(PipelineStage):
         # (B, C, H, W) -> (B, C, 1, H, W)
         image = image.unsqueeze(2)
 
-        video_condition = torch.cat([
-            image,
-            image.new_zeros(image.shape[0], image.shape[1], num_frames - 1,
-                            image.shape[3], image.shape[4])
-        ],
-                                    dim=2)
-        video_condition = video_condition.to(device=get_local_torch_device(),
-                                             dtype=torch.float32)
+        if fastvideo_args.pipeline_config.t2v_as_i2v_task:
+            # repeat the image self.vae.temporal_compression_ratio times
+            video_condition = image.repeat(1, 1,
+                                           self.vae.temporal_compression_ratio,
+                                           1, 1)
+            # video_condition = image
+            logger.info("video_condition.shape: %s", video_condition.shape)
+        else:
+            video_condition = torch.cat([
+                image,
+                image.new_zeros(image.shape[0], image.shape[1], num_frames - 1,
+                                image.shape[3], image.shape[4])
+            ],
+                                        dim=2)
 
-        # Setup VAE precision
-        vae_dtype = PRECISION_TO_TYPE[
-            fastvideo_args.pipeline_config.vae_precision]
-        vae_autocast_enabled = (
-            vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
-
-        # Encode Image
-        with torch.autocast(device_type="cuda",
-                            dtype=vae_dtype,
-                            enabled=vae_autocast_enabled):
-            if fastvideo_args.pipeline_config.vae_tiling:
-                self.vae.enable_tiling()
-            # if fastvideo_args.vae_sp:
-            #     self.vae.enable_parallel()
-            if not vae_autocast_enabled:
-                video_condition = video_condition.to(vae_dtype)
-            encoder_output = self.vae.encode(video_condition)
-
-        if fastvideo_args.mode == ExecutionMode.PREPROCESS:
-            latent_condition = encoder_output.mean
-        else:
-            generator = batch.generator
-            if generator is None:
-                raise ValueError("Generator must be provided")
-            latent_condition = self.retrieve_latents(encoder_output, generator)
-
-        # Apply shifting if needed
-        if (hasattr(self.vae, "shift_factor")
-                and self.vae.shift_factor is not None):
-            if isinstance(self.vae.shift_factor, torch.Tensor):
-                latent_condition -= self.vae.shift_factor.to(
-                    latent_condition.device, latent_condition.dtype)
-            else:
-                latent_condition -= self.vae.shift_factor
-
-        if isinstance(self.vae.scaling_factor, torch.Tensor):
-            latent_condition = latent_condition * self.vae.scaling_factor.to(
-                latent_condition.device, latent_condition.dtype)
-        else:
-            latent_condition = latent_condition * self.vae.scaling_factor
+        latent_condition = self.encode_tensor(video_condition, fastvideo_args,
+                                              batch.generator)
 
         if fastvideo_args.mode == ExecutionMode.PREPROCESS:
             batch.image_latent = latent_condition
+        elif fastvideo_args.pipeline_config.t2v_as_i2v_task:
+            logger.info("latent_condition.shape: %s", latent_condition.shape)
+            batch.image_latent = latent_condition
         else:
             mask_lat_size = torch.ones(1, 1, num_frames, latent_height,
                                        latent_width)

fastvideo/pipelines/stages/input_validation.py

@@ -4,6 +4,8 @@ Input validation stage for diffusion pipelines.
 """
 
 import torch
+import torchvision.transforms.functional as TF
+from PIL import Image
 
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
@@ -12,6 +14,7 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import (StageValidators,
                                                    VerificationResult)
+from fastvideo.utils import best_output_size
 
 logger = init_logger(__name__)
 
@@ -32,9 +35,15 @@ class InputValidationStage(PipelineStage):
         """Generate seeds for the inference"""
         seed = batch.seed
         num_videos_per_prompt = batch.num_videos_per_prompt
+        if isinstance(batch.prompt, list):
+            num_prompts = len(batch.prompt)
+        else:
+            num_prompts = 1
+
+        total_num_videos = num_prompts * num_videos_per_prompt
 
         assert seed is not None
-        seeds = [seed + i for i in range(num_videos_per_prompt)]
+        seeds = [seed + i for i in range(total_num_videos)]
         batch.seeds = seeds
         # Peiyuan: using GPU seed will cause A100 and H100 to generate different results...
         batch.generator = [
@@ -94,6 +103,7 @@ class InputValidationStage(PipelineStage):
             )
 
         # for i2v, get image from image_path
+        # @TODO(Wei) hard-coded for wan2.2 5b ti2v for now. Should put this in image_encoding stage
         if batch.image_path is not None:
             if batch.image_path.endswith(".mp4"):
                 image = load_video(batch.image_path)[0]
@@ -101,6 +111,36 @@ class InputValidationStage(PipelineStage):
                 image = load_image(batch.image_path)
             batch.pil_image = image
 
+        # further processing for ti2v task
+        if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+            img = batch.pil_image
+            ih, iw = img.height, img.width
+            patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
+            vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
+            dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride
+            max_area = 704 * 1280
+            ow, oh = best_output_size(iw, ih, dw, dh, max_area)
+
+            scale = max(ow / iw, oh / ih)
+            img = img.resize((round(iw * scale), round(ih * scale)),
+                             Image.LANCZOS)
+            logger.info("resized img height: %s, img width: %s", img.height,
+                        img.width)
+
+            # center-crop
+            x1 = (img.width - ow) // 2
+            y1 = (img.height - oh) // 2
+            img = img.crop((x1, y1, x1 + ow, y1 + oh))
+            assert img.width == ow and img.height == oh
+
+            # to tensor
+            img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(
+                self.device).unsqueeze(1)
+            img = img.unsqueeze(0)
+            batch.height = oh
+            batch.width = ow
+            batch.pil_image = img
+
         return batch
 
     def verify_input(self, batch: ForwardBatch,

fastvideo/pipelines/stages/latent_preparation.py

@@ -59,6 +59,10 @@ class LatentPreparationStage(PipelineStage):
 
         # Adjust batch size for number of videos per prompt
         batch_size *= batch.num_videos_per_prompt
+        logger.info(f"===== batch_size: {batch_size}")
+        logger.info(f"===== batch.prompt: {batch.prompt}")
+        logger.info(f"===== batch.prompt_embeds: {batch.prompt_embeds}")
+        logger.info(f"===== batch.prompt_attention_mask: {batch.prompt_attention_mask}")
 
         # Get required parameters
         dtype = batch.prompt_embeds[0].dtype

fastvideo/pipelines/stages/text_encoding.py

@@ -59,58 +59,37 @@ class TextEncodingStage(PipelineStage):
         assert len(self.text_encoders) == len(
             fastvideo_args.pipeline_config.text_encoder_configs)
 
-        for tokenizer, text_encoder, encoder_config, preprocess_func, postprocess_func in zip(
-                self.tokenizers,
-                self.text_encoders,
-                fastvideo_args.pipeline_config.text_encoder_configs,
-                fastvideo_args.pipeline_config.preprocess_text_funcs,
-                fastvideo_args.pipeline_config.postprocess_text_funcs,
-                strict=True):
+        # Encode positive prompt with all available encoders
+        assert batch.prompt is not None
+        prompt_text: str | list[str] = batch.prompt
+        all_indices: list[int] = list(range(len(self.text_encoders)))
+        prompt_embeds_list, prompt_masks_list = self.encode_text(
+            prompt_text,
+            fastvideo_args,
+            encoder_index=all_indices,
+            return_attention_mask=True,
+        )
+        for pe in prompt_embeds_list:
+            batch.prompt_embeds.append(pe)
+        if batch.prompt_attention_mask is not None:
+            for am in prompt_masks_list:
+                batch.prompt_attention_mask.append(am)
 
-            assert isinstance(batch.prompt, str | list)
-            if isinstance(batch.prompt, str):
-                batch.prompt = [batch.prompt]
-            texts = []
-            for prompt_str in batch.prompt:
-                texts.append(preprocess_func(prompt_str))
-            text_inputs = tokenizer(texts,
-                                    **encoder_config.tokenizer_kwargs).to(
-                                        get_local_torch_device())
-            input_ids = text_inputs["input_ids"]
-            attention_mask = text_inputs["attention_mask"]
-            with set_forward_context(current_timestep=0, attn_metadata=None):
-                outputs = text_encoder(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    output_hidden_states=True,
-                )
-            prompt_embeds = postprocess_func(outputs)
-            batch.prompt_embeds.append(prompt_embeds)
-            if batch.prompt_attention_mask is not None:
-                batch.prompt_attention_mask.append(attention_mask)
-
-            if batch.do_classifier_free_guidance:
-                assert isinstance(batch.negative_prompt, str)
-                negative_text = preprocess_func(batch.negative_prompt)
-                negative_text_inputs = tokenizer(
-                    negative_text, **encoder_config.tokenizer_kwargs).to(
-                        get_local_torch_device())
-                negative_input_ids = negative_text_inputs["input_ids"]
-                negative_attention_mask = negative_text_inputs["attention_mask"]
-                with set_forward_context(current_timestep=0,
-                                         attn_metadata=None):
-                    negative_outputs = text_encoder(
-                        input_ids=negative_input_ids,
-                        attention_mask=negative_attention_mask,
-                        output_hidden_states=True,
-                    )
-                negative_prompt_embeds = postprocess_func(negative_outputs)
-
-                assert batch.negative_prompt_embeds is not None
-                batch.negative_prompt_embeds.append(negative_prompt_embeds)
-                if batch.negative_attention_mask is not None:
-                    batch.negative_attention_mask.append(
-                        negative_attention_mask)
+        # Encode negative prompt if CFG is enabled
+        if batch.do_classifier_free_guidance:
+            assert isinstance(batch.negative_prompt, str)
+            neg_embeds_list, neg_masks_list = self.encode_text(
+                batch.negative_prompt,
+                fastvideo_args,
+                encoder_index=all_indices,
+                return_attention_mask=True,
+            )
+            assert batch.negative_prompt_embeds is not None
+            for ne in neg_embeds_list:
+                batch.negative_prompt_embeds.append(ne)
+            if batch.negative_attention_mask is not None:
+                for nm in neg_masks_list:
+                    batch.negative_attention_mask.append(nm)
 
         return batch
 
@@ -129,6 +108,171 @@ class TextEncodingStage(PipelineStage):
                          V.none_or_list)
         return result
 
+    @torch.no_grad()
+    def encode_text(
+        self,
+        text: str | list[str],
+        fastvideo_args: FastVideoArgs,
+        encoder_index: int | list[int] | None = None,
+        return_attention_mask: bool = False,
+        return_type: str = "list",  # one of: "list", "dict", "stack"
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+        max_length: int | None = None,
+        truncation: bool | None = None,
+        padding: bool | str | None = None,
+    ):
+        """
+        Encode plain text using selected text encoder(s) and return embeddings.
+
+        Args:
+            text: A single string or a list of strings to encode.
+            fastvideo_args: The inference arguments providing pipeline config,
+                including tokenizer and encoder settings, preprocess and postprocess
+                functions.
+            encoder_index: Encoder selector by index. Accepts an int or list of ints.
+            return_attention_mask: If True, also return attention masks for each
+                selected encoder.
+            return_type: "list" (default) returns a list aligned with selection;
+                "dict" returns a dict keyed by encoder index as a string; "stack" stacks along a
+                new first dimension (requires matching shapes).
+            device: Optional device override for inputs; defaults to local torch device.
+            dtype: Optional dtype to cast returned embeddings to.
+            max_length: Optional per-call tokenizer override.
+            truncation: Optional per-call tokenizer override.
+            padding: Optional per-call tokenizer override.
+
+        Returns:
+            Depending on return_type and return_attention_mask:
+            - list: List[Tensor] or (List[Tensor], List[Tensor])
+            - dict: Dict[str, Tensor] or (Dict[str, Tensor], Dict[str, Tensor])
+            - stack: Tensor of shape [num_encoders, ...] or a tuple with stacked
+              attention masks
+        """
+
+        assert len(self.tokenizers) == len(self.text_encoders)
+        assert len(self.text_encoders) == len(
+            fastvideo_args.pipeline_config.text_encoder_configs)
+
+        # Resolve selection into indices
+        encoder_cfgs = fastvideo_args.pipeline_config.text_encoder_configs
+        if encoder_index is None:
+            indices: list[int] = [0]
+        elif isinstance(encoder_index, int):
+            indices = [encoder_index]
+        else:
+            indices = list(encoder_index)
+        # validate range
+        num_encoders = len(self.text_encoders)
+        for idx in indices:
+            if idx < 0 or idx >= num_encoders:
+                raise IndexError(
+                    f"encoder index {idx} out of range [0, {num_encoders-1}]")
+
+        # Validate indices are within range
+        num_encoders = len(self.text_encoders)
+
+        # Normalize input to list[str]
+        assert isinstance(text, str | list)
+        if isinstance(text, str):
+            texts: list[str] = [text]
+        else:
+            texts = text
+
+        embeds_list: list[torch.Tensor] = []
+        attn_masks_list: list[torch.Tensor] = []
+
+        preprocess_funcs = fastvideo_args.pipeline_config.preprocess_text_funcs
+        postprocess_funcs = fastvideo_args.pipeline_config.postprocess_text_funcs
+        encoder_cfgs = fastvideo_args.pipeline_config.text_encoder_configs
+
+        if return_type not in ("list", "dict", "stack"):
+            raise ValueError(
+                f"Invalid return_type '{return_type}'. Expected one of: 'list', 'dict', 'stack'"
+            )
+
+        target_device = device if device is not None else get_local_torch_device(
+        )
+
+        for i in indices:
+            tokenizer = self.tokenizers[i]
+            text_encoder = self.text_encoders[i]
+            encoder_config = encoder_cfgs[i]
+            preprocess_func = preprocess_funcs[i]
+            postprocess_func = postprocess_funcs[i]
+
+            processed_texts: list[str] = []
+            for prompt_str in texts:
+                processed_texts.append(preprocess_func(prompt_str))
+
+            tok_kwargs = dict(encoder_config.tokenizer_kwargs)
+            if max_length is not None:
+                tok_kwargs["max_length"] = max_length
+            if truncation is not None:
+                tok_kwargs["truncation"] = truncation
+            if padding is not None:
+                tok_kwargs["padding"] = padding
+
+            text_inputs = tokenizer(processed_texts,
+                                    **tok_kwargs).to(target_device)
+
+            input_ids = text_inputs["input_ids"]
+            attention_mask = text_inputs["attention_mask"]
+
+            with set_forward_context(current_timestep=0, attn_metadata=None):
+                outputs = text_encoder(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    output_hidden_states=True,
+                )
+
+            prompt_embeds = postprocess_func(outputs)
+            if dtype is not None:
+                prompt_embeds = prompt_embeds.to(dtype=dtype)
+            embeds_list.append(prompt_embeds)
+            if return_attention_mask:
+                attn_masks_list.append(attention_mask)
+
+        # Shape results according to return_type
+        if return_type == "list":
+            if return_attention_mask:
+                return embeds_list, attn_masks_list
+            return embeds_list
+
+        if return_type == "dict":
+            key_strs = [str(i) for i in indices]
+            embeds_dict = {
+                k: v
+                for k, v in zip(key_strs, embeds_list, strict=False)
+            }
+            if return_attention_mask:
+                attn_dict = {
+                    k: v
+                    for k, v in zip(key_strs, attn_masks_list, strict=False)
+                }
+                return embeds_dict, attn_dict
+            return embeds_dict
+
+        # return_type == "stack"
+        # Validate shapes are compatible
+        base_shape = list(embeds_list[0].shape)
+        for t in embeds_list[1:]:
+            if list(t.shape) != base_shape:
+                raise ValueError(
+                    f"Cannot stack embeddings with differing shapes: {[list(t.shape) for t in embeds_list]}"
+                )
+        stacked_embeds = torch.stack(embeds_list, dim=0)
+        if return_attention_mask:
+            base_mask_shape = list(attn_masks_list[0].shape)
+            for m in attn_masks_list[1:]:
+                if list(m.shape) != base_mask_shape:
+                    raise ValueError(
+                        f"Cannot stack attention masks with differing shapes: {[list(m.shape) for m in attn_masks_list]}"
+                    )
+            stacked_masks = torch.stack(attn_masks_list, dim=0)
+            return stacked_embeds, stacked_masks
+        return stacked_embeds
+
     def verify_output(self, batch: ForwardBatch,
                       fastvideo_args: FastVideoArgs) -> VerificationResult:
         """Verify text encoding stage outputs."""

fastvideo/platforms/__init__.py

@@ -82,8 +82,8 @@ def rocm_platform_plugin() -> str | None:
                 logger.info("ROCm platform is available")
         finally:
             amdsmi.amdsmi_shut_down()
-    except Exception as e:
-        logger.info("ROCm platform is unavailable: %s", e)
+    except Exception:
+        pass
 
     return "fastvideo.platforms.rocm.RocmPlatform" if is_rocm else None
 

fastvideo/tests/stages/test_text_encoding.py

@@ -0,0 +1,128 @@
+import torch
+import types
+import pytest
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig, BaseEncoderOutput
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+class TensorDict(dict):
+    def to(self, device):
+        return TensorDict({k: v.to(device) for k, v in self.items()})
+
+class FakeTokenizer:
+    def __call__(self, texts, **kwargs):
+        B = len(texts)
+        seq_len = int(kwargs.get("max_length", 4))
+        return TensorDict({
+            "input_ids": torch.arange(B * seq_len).view(B, seq_len),
+            "attention_mask": torch.ones(B, seq_len, dtype=torch.long),
+        })
+
+class FakeTextEncoder(torch.nn.Module):
+    def __init__(self, hidden_size=8):
+        super().__init__()
+        self.hidden_size = hidden_size
+
+    def forward(self, input_ids, attention_mask, output_hidden_states=True):
+        B, T = input_ids.shape
+        last_hidden_state = torch.arange(B * T * self.hidden_size, dtype=torch.float32).view(B, T, self.hidden_size)
+        return types.SimpleNamespace(last_hidden_state=last_hidden_state)
+
+def id_preprocess(x: str) -> str:
+    return x
+
+def take_mean_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    # [B, T, H] -> [B, H]
+    return outputs.last_hidden_state.mean(dim=1)
+
+def make_args(num_encoders=2, text_len=4, hidden_size=8):
+    enc_cfgs = []
+    preprocess_fns = []
+    postprocess_fns = []
+    for _ in range(num_encoders):
+        arch = TextEncoderArchConfig(text_len=text_len)
+        enc_cfgs.append(TextEncoderConfig(arch_config=arch))
+        preprocess_fns.append(id_preprocess)
+        postprocess_fns.append(take_mean_postprocess)
+    pipe_cfg = PipelineConfig(
+        text_encoder_configs=tuple(enc_cfgs),
+        text_encoder_precisions=tuple(["fp32"] * num_encoders),
+        preprocess_text_funcs=tuple(preprocess_fns),
+        postprocess_text_funcs=tuple(postprocess_fns),
+    )
+    return FastVideoArgs(model_path="", pipeline_config=pipe_cfg), hidden_size
+
+def make_stage(num_encoders=2, hidden_size=8):
+    tokenizers = [FakeTokenizer() for _ in range(num_encoders)]
+    encoders = [FakeTextEncoder(hidden_size=hidden_size) for _ in range(num_encoders)]
+    return TextEncodingStage(text_encoders=encoders, tokenizers=tokenizers)
+
+def test_encode_text_selection_and_shapes():
+    fastvideo_args, hidden = make_args(num_encoders=2, text_len=4, hidden_size=8)
+    stage = make_stage(num_encoders=2, hidden_size=hidden)
+
+    # list return, two encoders
+    embeds = stage.encode_text(["a", "b"], fastvideo_args, encoder_index=[0, 1])
+    assert isinstance(embeds, list) and len(embeds) == 2
+    for e in embeds:
+        assert e.shape == (2, hidden)
+
+    # with masks
+    embeds2, masks2 = stage.encode_text("a", fastvideo_args, encoder_index=[1], return_attention_mask=True)
+    assert len(embeds2) == 1 and len(masks2) == 1
+    assert embeds2[0].shape == (1, hidden)
+    assert masks2[0].shape == (1, 4)
+
+    # dict return
+    d = stage.encode_text(["a","b"], fastvideo_args, encoder_index=[0,1], return_type="dict")
+    assert set(d.keys()) == {"0", "1"}
+    assert d["0"].shape == (2, hidden)
+
+    # stack return
+    s = stage.encode_text(["a","b"], fastvideo_args, encoder_index=[0,1], return_type="stack")
+    assert s.shape == (2, 2, hidden)  # [encoders, batch, hidden]
+
+    # overrides: dtype + max_length
+    e3, m3 = stage.encode_text(["a"], fastvideo_args, encoder_index=[0], dtype=torch.float16, return_attention_mask=True, max_length=3)
+    assert e3[0].dtype == torch.float16
+    assert m3[0].shape[1] == 3
+
+def test_forward_integration_cfg_off_and_on():
+    fastvideo_args, hidden = make_args(num_encoders=2, text_len=4, hidden_size=8)
+    stage = make_stage(num_encoders=2, hidden_size=hidden)
+
+    # CFG off
+    batch = ForwardBatch(
+        data_type="video",
+        prompt="a cat",
+        negative_prompt="",
+        do_classifier_free_guidance=False,
+        prompt_embeds=[],
+        negative_prompt_embeds=None,
+        prompt_attention_mask=[],
+        negative_attention_mask=None,
+    )
+    out = stage.forward(batch, fastvideo_args)
+    assert len(out.prompt_embeds) == 2
+    for e in out.prompt_embeds:
+        assert e.shape[1] == hidden
+
+    # CFG on
+    batch2 = ForwardBatch(
+        data_type="video",
+        prompt=["a cat", "a dog"],
+        negative_prompt="bad picture",
+        do_classifier_free_guidance=True,
+        prompt_embeds=[],
+        negative_prompt_embeds=[],
+        prompt_attention_mask=[],
+        negative_attention_mask=[],
+    )
+    out2 = stage.forward(batch2, fastvideo_args)
+    assert len(out2.prompt_embeds) == 2
+    assert len(out2.negative_prompt_embeds) == 2
+    assert len(out2.prompt_attention_mask) == 2
+    assert len(out2.negative_attention_mask) == 2

fastvideo/training/distillation_pipeline.py

@@ -29,6 +29,7 @@ from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler)
+from fastvideo.models.utils import pred_noise_to_pred_video
 from fastvideo.pipelines import (ComposedPipelineBase, ForwardBatch,
                                  TrainingBatch)
 from fastvideo.training.activation_checkpoint import (
@@ -36,8 +37,7 @@ from fastvideo.training.activation_checkpoint import (
 from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
     clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler,
-    load_distillation_checkpoint, pred_noise_to_pred_video,
-    save_distillation_checkpoint, shift_timestep)
+    load_distillation_checkpoint, save_distillation_checkpoint, shift_timestep)
 from fastvideo.utils import is_vsa_available, set_random_seed
 
 import wandb  # isort: skip