chekcpoint

.buildkite/pipeline.yml

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

.github/workflows/pr-test.yml

@@ -104,17 +104,16 @@ jobs:
               - 'pyproject.toml'
               - 'docker/Dockerfile.python3.12'
             sta-kernel-paths: &sta-kernel-paths
-              - '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'
+              - 'csrc/attn/st_attn/**'
+              - 'csrc/attn/setup_sta.py'
+              - 'csrc/attn/config_sta.py'
+              - 'csrc/attn/st_attn.cpp'
             vsa-kernel-paths: &vsa-kernel-paths
-              - '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'
+              - 'csrc/attn/vsa/**'
+              - 'csrc/attn/tk/**'
+              - 'csrc/attn/setup_vsa.py'
+              - 'csrc/attn/config_vsa.py'
+              - 'csrc/attn/vsa.cpp'
             vsa-paths: &vsa-paths
               - 'fastvideo/**'
               - *common-paths
@@ -235,7 +234,7 @@ jobs:
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
+  
   training-test:
     needs: change-filter
     if: >-
@@ -373,4 +372,4 @@ jobs:
           JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12", "training-test", "training-test-VSA", "inference-test-STA", "precision-test-STA", "precision-test-VSA"]'
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           GITHUB_RUN_ID: ${{ github.run_id }}
-        run: python .github/scripts/runpod_cleanup.py
+        run: python .github/scripts/runpod_cleanup.py

.github/workflows/sta-publish.yml

@@ -5,7 +5,7 @@ on:
     branches:
       - main
     paths:
-      - "csrc/attn/sliding_tile_attn/setup.py"
+      - "csrc/attn/setup_sta.py"
   workflow_dispatch:
 
 jobs:
@@ -23,13 +23,13 @@ jobs:
       - name: Check if version changed
         id: check-version
         run: |
-          cd csrc/attn/sliding_tile_attn
+          cd csrc/attn
           # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_sta.py)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup_sta.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/sliding_tile_attn  # Move into the correct folder
+          cd csrc/attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py bdist_wheel --dist-dir=dist
+          python setup_sta.py bdist_wheel --dist-dir=dist
 
       - name: Rename wheel file
         run: |
-          cd csrc/attn/sliding_tile_attn
+          cd csrc/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/sliding_tile_attn/dist/*.whl
+          path: csrc/attn/dist/*.whl
           retention-days: 90
 
   publish_package:
@@ -239,11 +239,11 @@ jobs:
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn/sliding_tile_attn  # Move into the correct folder
+          cd csrc/attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py sdist --dist-dir=dist
+          python setup_sta.py sdist --dist-dir=dist
 
       - name: Publish release distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: csrc/attn/sliding_tile_attn/dist/
+          packages-dir: csrc/attn/dist/

.github/workflows/vsa-publish.yml

@@ -5,7 +5,7 @@ on:
     branches:
       - main
     paths:
-      - "csrc/attn/video_sparse_attn/setup.py"
+      - "csrc/attn/setup_vsa.py"
   workflow_dispatch:
 
 jobs:
@@ -23,13 +23,13 @@ jobs:
       - name: Check if version changed
         id: check-version
         run: |
-          cd csrc/attn/video_sparse_attn
+          cd csrc/attn
           # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_vsa.py)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup_vsa.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/video_sparse_attn  # Move into the correct folder
+          cd csrc/attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py bdist_wheel --dist-dir=dist
+          python setup_vsa.py bdist_wheel --dist-dir=dist
 
       - name: Rename wheel file
         run: |
-          cd csrc/attn/video_sparse_attn
+          cd csrc/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/video_sparse_attn/dist/*.whl
+          path: csrc/attn/dist/*.whl
           retention-days: 90
 
   publish_package:
@@ -247,11 +247,11 @@ jobs:
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn/video_sparse_attn  # Move into the correct folder
+          cd csrc/attn  # Move into the correct folder
           git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py sdist --dist-dir=dist
+          python setup_vsa.py sdist --dist-dir=dist
 
       - name: Publish release distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: csrc/attn/video_sparse_attn/dist/
+          packages-dir: csrc/attn/dist/

.gitmodules

@@ -1,7 +1,3 @@
-[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
+[submodule "csrc/attn/tk"]
+	path = csrc/attn/tk
 	url = https://github.com/HazyResearch/ThunderKittens.git

README.md

@@ -7,7 +7,7 @@
 FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.
 
 <p align="center">
-    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/S7HLCSTh" target="_blank"> <b> WeChat </b> </a> |
+    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/wZPZTLKg" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 <div align="center">

csrc/attn/MANIFEST.in

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

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.8)
+(If you use CUDA12.4)
 ```bash
-export CUDA_HOME=/usr/local/cuda-12.8
+export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 ```

csrc/attn/scratch.sh

@@ -0,0 +1,4 @@
+    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/setup_sta.py

@@ -1,7 +1,7 @@
 import os
 import subprocess
 
-from config_sta import kernels, sources, target
+from csrc.attn.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.6"
+VERSION = "0.0.4"
 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/setup_vsa.py

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

csrc/attn/sliding_tile_attn/README.md

@@ -1,87 +0,0 @@
-
-# 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/video_sparse_attn/MANIFEST.in

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

csrc/attn/video_sparse_attn/README.md

@@ -1,61 +0,0 @@
-
-
-# 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.

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/sliding_tile_attn && \
+    cd csrc/attn && \
     git submodule update --init --recursive && \
-    python setup.py install
+    python setup_sta.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
+    cd csrc/attn && \
     git submodule update --init --recursive && \
-    python setup.py install
+    python setup_vsa.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/sliding_tile_attn && \
+    cd csrc/attn && \
     git submodule update --init --recursive && \
-    python setup.py install
+    python setup_sta.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
+    cd csrc/attn && \
     git submodule update --init --recursive && \
-    python setup.py install
+    python setup_vsa.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/sliding_tile_attn && \
+    cd csrc/attn && \
     git submodule update --init --recursive && \
-    python setup.py install
+    python setup_sta.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
+    cd csrc/attn && \
     git submodule update --init --recursive && \
-    python setup.py install
+    python setup_vsa.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/sliding_tile_attn && \
+    cd csrc/attn && \
     git submodule update --init --recursive && \
-    python setup.py install
+    python setup_sta.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
+    cd csrc/attn && \
     git submodule update --init --recursive && \
-    python setup.py install
+    python setup_vsa.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
+pip install st_attn==0.0.4
 ``` 
 
 # Building from Source
@@ -12,6 +12,7 @@ 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
 
@@ -21,20 +22,14 @@ sudo apt update
 sudo apt install clang-11
 ```
 
-Set up CUDA environment (if using CUDA 12.4):
+Install STA:
 
 ```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.py install
+python setup_sta.py install
 ```
 
 # 🧪 Test

docs/source/video_sparse_attention/installation.md

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

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

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

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

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

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

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

examples/inference/basic/basic_self_forcing_causal.py

@@ -1,31 +0,0 @@
-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)
-
-    prompt = (
-        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
-        "wide with interest. The playful yet serene atmosphere is complemented by soft "
-        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
-    )
-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_wan2_2_ti2v.py

@@ -1,41 +0,0 @@
-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/finetune/Wan2.1-VSA/Wan-Syn-Data/README.md

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

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

@@ -10,7 +10,6 @@ torchrun --nproc_per_node=$GPU_NUM \
     --model_path $MODEL_PATH \
     --mode preprocess \
     --workload_type t2v \
-    --preprocess.video_loader_type torchvision \
     --preprocess.dataset_type merged \
     --preprocess.dataset_path $DATASET_PATH \
     --preprocess.dataset_output_dir $OUTPUT_DIR \

fastvideo/configs/configs.py

@@ -32,29 +32,6 @@ class DatasetType(str, Enum):
         return [dataset_type.value for dataset_type in cls]
 
 
-class VideoLoaderType(str, Enum):
-    """
-    Enumeration for different video loaders.
-    """
-    TORCHCODEC = "torchcodec"
-    TORCHVISION = "torchvision"
-
-    @classmethod
-    def from_string(cls, value: str) -> "VideoLoaderType":
-        """Convert string to VideoLoader enum."""
-        try:
-            return cls(value.lower())
-        except ValueError:
-            raise ValueError(
-                f"Invalid video loader: {value}. Must be one of: {', '.join([m.value for m in cls])}"
-            ) from None
-
-    @classmethod
-    def choices(cls) -> list[str]:
-        """Get all available choices as strings for argparse."""
-        return [video_loader.value for video_loader in cls]
-
-
 @dataclasses.dataclass
 class PreprocessConfig:
     """Configuration for preprocessing operations."""
@@ -74,7 +51,6 @@ class PreprocessConfig:
     flush_frequency: int = 256
 
     # Video processing parameters
-    video_loader_type: VideoLoaderType = VideoLoaderType.TORCHCODEC
     max_height: int = 480
     max_width: int = 848
     num_frames: int = 163
@@ -144,12 +120,6 @@ class PreprocessConfig:
                                      help="How often to save to parquet files")
 
         # Video processing parameters
-        preprocess_args.add_argument(
-            f"--{prefix_with_dot}video-loader-type",
-            type=str,
-            choices=VideoLoaderType.choices(),
-            default=PreprocessConfig.video_loader_type.value,
-            help="Type of the video loader")
         preprocess_args.add_argument(f"--{prefix_with_dot}max-height",
                                      type=int,
                                      default=PreprocessConfig.max_height,
@@ -204,10 +174,6 @@ class PreprocessConfig:
         if 'dataset_type' in kwargs and isinstance(kwargs['dataset_type'], str):
             kwargs['dataset_type'] = DatasetType.from_string(
                 kwargs['dataset_type'])
-        if 'video_loader_type' in kwargs and isinstance(
-                kwargs['video_loader_type'], str):
-            kwargs['video_loader_type'] = VideoLoaderType.from_string(
-                kwargs['video_loader_type'])
 
         preprocess_config = cls()
         if not update_config_from_args(

fastvideo/configs/models/dits/wanvideo.py

@@ -92,12 +92,6 @@ 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

@@ -85,9 +85,6 @@ class PipelineConfig:
     # DMD parameters
     dmd_denoising_steps: list[int] | None = field(default=None)
 
-    # Wan2.2 TI2V parameters
-    ti2v_task: bool = False
-
     # Compilation
     # enable_torch_compile: bool = False
 

fastvideo/configs/pipelines/registry.py

@@ -7,14 +7,10 @@ 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
-
-# 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)
-# isort: on
+from fastvideo.configs.pipelines.wan import (FastWan2_1_T2V_480P_Config,
+                                             FastWan2_2_TI2V_5B_Config,
+                                             WanI2V480PConfig, WanI2V720PConfig,
+                                             WanT2V480PConfig, WanT2V720PConfig)
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
                              verify_model_config_and_directory)
@@ -35,10 +31,9 @@ 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,
-    "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,
+    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig,
+    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": WanT2V480PConfig,
+    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": WanI2V480PConfig,
     # Add other specific weight variants
 }
 

fastvideo/configs/pipelines/wan.py

@@ -138,13 +138,3 @@ 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
-    dmd_denoising_steps: list[int] | None = field(
-        default_factory=lambda: [1000, 750, 500, 250])

fastvideo/configs/sample/registry.py

@@ -10,7 +10,6 @@ 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,
@@ -18,7 +17,6 @@ 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
@@ -30,31 +28,16 @@ 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,
-    "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers":
-    Wan2_1_Fun_1_3B_InP_SamplingParam,
-
-    # Wan2.2
+    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
     "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
-    "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers":
-    Wan2_2_TI2V_5B_SamplingParam,
+    "FastVideo/FastWan2.2-TI2V-5B-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,21 +107,6 @@ 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 =============
 # =============================================
@@ -156,11 +141,3 @@ 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/__init__.py

@@ -4,7 +4,7 @@ from torchvision.transforms import Lambda
 
 from fastvideo.dataset.parquet_dataset_map_style import (
     build_parquet_map_style_dataloader)
-from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset, TextDataset
+from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset
 from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255,
                                          TemporalRandomCrop)
 from fastvideo.dataset.validation_dataset import ValidationDataset
@@ -39,13 +39,7 @@ def getdataset(args) -> VideoCaptionMergedDataset:
                                      seed=args.seed)
 
 
-def gettextdataset(args) -> TextDataset:
-    return TextDataset(data_merge_path=args.data_merge_path,
-                       args=args,
-                       seed=args.seed)
-
-
 __all__ = [
     "build_parquet_map_style_dataloader", "ValidationDataset",
-    "VideoCaptionMergedDataset", "TextDataset"
-]
+    "VideoCaptionMergedDataset"
+]

fastvideo/dataset/dataloader/schema.py

@@ -78,76 +78,3 @@ 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()),
-])
-
-pyarrow_schema_ode_trajectory_text_only = pa.schema([
-    pa.field("id", 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()),
-    # --- 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()),  # Always 'text' for text-only
-])
-
-pyarrow_schema_text_only = pa.schema([
-    pa.field("id", 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()),
-])

fastvideo/dataset/preprocessing_datasets.py

@@ -628,134 +628,3 @@ class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
         """Load state dict from checkpoint."""
         self.processed_batches = state_dict["processed_batches"]
-
-
-class TextDataset(torch.utils.data.IterableDataset,
-                  torch.distributed.checkpoint.stateful.Stateful):
-    """
-    Text-only dataset for processing prompts from a simple text file.
-    
-    Assumes that data_merge_path is a text file with one prompt per line:
-    A cat playing with a ball
-    A dog running in the park
-    A person cooking dinner
-    ...
-    
-    This dataset processes text data through text encoding stages only.
-    """
-
-    def __init__(self,
-                 data_merge_path: str,
-                 args,
-                 start_idx: int = 0,
-                 seed: int = 42):
-        self.data_merge_path = data_merge_path
-        self.start_idx = start_idx
-        self.args = args
-        self.seed = seed
-
-        # Initialize tokenizer
-        tokenizer_path = os.path.join(args.model_path, "tokenizer")
-        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
-                                                  cache_dir=args.cache_dir)
-
-        # Initialize text encoding stage
-        self.text_encoding_stage = TextEncodingStage(
-            tokenizer=tokenizer,
-            text_max_length=args.text_max_length,
-            cfg_rate=getattr(args, 'training_cfg_rate', 0.0),
-            seed=self.seed)
-
-        # Process text data
-        self.processed_batches = self._process_text_data()
-
-    def _load_text_data(self) -> list[str]:
-        """Load text prompts from file."""
-        prompts = []
-        with open(self.data_merge_path, 'r', encoding='utf-8') as f:
-            for line in f:
-                line = line.strip()
-                if line:  # Skip empty lines
-                    prompts.append(line)
-        
-        logger.info(f"Loaded {len(prompts)} text prompts from {self.data_merge_path}")
-        return prompts
-
-    def _process_text_data(self) -> list[PreprocessBatch]:
-        """Process the text prompts through text encoding stage."""
-        raw_prompts = self._load_text_data()
-        processed_batches = []
-
-        for idx, prompt in enumerate(raw_prompts):
-            # Create a text-only batch with dummy path
-            batch = PreprocessBatch(
-                path=f"text_prompt_{idx}",
-                cap=[prompt],  # TextEncodingStage expects a list
-                resolution=None,
-                fps=None,
-                duration=None,
-                num_frames=0,
-                sample_frame_index=None,
-                sample_num_frames=0
-            )
-            
-            processed_batches.append(batch)
-
-        logger.info(f"Processed {len(processed_batches)} text batches")
-        return processed_batches
-
-    def __iter__(self):
-        """Iterator for the dataset."""
-        # Set up distributed sampling if needed
-        if torch.distributed.is_available() and torch.distributed.is_initialized():
-            rank = torch.distributed.get_rank()
-            world_size = torch.distributed.get_world_size()
-        else:
-            rank = 0
-            world_size = 1
-
-        # Calculate chunk for this rank
-        total_items = len(self.processed_batches)
-        items_per_rank = math.ceil(total_items / world_size)
-        start_idx = rank * items_per_rank + self.start_idx
-        end_idx = min(start_idx + items_per_rank, total_items)
-
-        # Yield items for this rank
-        for idx in range(start_idx, end_idx):
-            if idx < len(self.processed_batches):
-                yield self._get_item(idx)
-
-    def _get_item(self, idx: int) -> dict:
-        """Get a single processed text item."""
-        batch = self.processed_batches[idx]
-
-        # Apply text encoding stage
-        batch = self.text_encoding_stage.process(batch)
-
-        # Build result dictionary for text-only processing with required schema fields
-        result = {
-            "text": batch.text,
-            "input_ids": batch.input_ids,
-            "cond_mask": batch.cond_mask,
-            "path": batch.path,
-            # Required schema fields for ODE trajectory processing
-            "id": f"text_{idx}",
-            "file_name": batch.path,
-            "caption": batch.text,
-            "media_type": "text",
-            "width": 1,
-            "height": 1,
-            "num_frames": 0,
-            "duration_sec": 0.0,
-            "fps": 0.0,
-        }
-
-        return result
-
-    def state_dict(self) -> dict[str, Any]:
-        """Return state dict for checkpointing."""
-        return {"processed_batches": self.processed_batches}
-
-    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
-        """Load state dict from checkpoint."""
-        self.processed_batches = state_dict["processed_batches"]

fastvideo/layers/rotary_embedding.py

@@ -29,9 +29,6 @@ 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:
@@ -270,7 +267,6 @@ 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.
@@ -296,9 +292,6 @@ 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
@@ -377,7 +370,6 @@ 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.
@@ -421,7 +413,6 @@ 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,16 +86,12 @@ class TimestepEmbedder(nn.Module):
                        dtype=dtype)
         self.freq_dtype = freq_dtype
 
-    def forward(self,
-                t: torch.Tensor,
-                timestep_seq_len: int | None = None) -> torch.Tensor:
+    def forward(self, t: torch.Tensor) -> 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

@@ -1,648 +0,0 @@
-# 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:
-        if hidden_states.dim() == 4:
-            hidden_states = hidden_states.squeeze(1)
-        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)
-        assert shift_msa.dtype == torch.float32
-
-        # 1. Self-attention
-        norm_hidden_states = (self.norm1(hidden_states.float()) *
-                              (1 + scale_msa) + shift_msa).to(orig_dtype)
-        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)
-        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)
-        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)
-        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)
-        self.proj_out = nn.Linear(
-            inner_dim, config.out_channels * math.prod(config.patch_size))
-        self.scale_shift_table = nn.Parameter(
-            torch.randn(1, 2, inner_dim) / inner_dim**0.5)
-
-        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)
-        """
-
-        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)
-
-        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep, encoder_hidden_states, encoder_hidden_states_image)
-        timestep_proj = timestep_proj.unflatten(1, (6, -1))
-
-        if encoder_hidden_states_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
-        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:
-
-        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)
-
-        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep, encoder_hidden_states, encoder_hidden_states_image)
-        timestep_proj = timestep_proj.unflatten(1, (6, -1))
-
-        if encoder_hidden_states_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 in self.blocks:
-                hidden_states = block(hidden_states, encoder_hidden_states,
-                                        timestep_proj, freqs_cis,
-                                        block_mask=self.block_mask)
-
-        # 5. Output norm, projection & unpatchify
-        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(
-        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,9 +81,8 @@ 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, timestep_seq_len)
+        temb = self.time_embedder(timestep)
         timestep_proj = self.time_modulation(temb)
 
         encoder_hidden_states = self.text_embedder(encoder_hidden_states)
@@ -146,7 +145,7 @@ class WanSelfAttention(nn.Module):
 
 class WanT2VCrossAttention(WanSelfAttention):
 
-    def forward(self, x, context, context_lens, crossattn_cache=None):
+    def forward(self, x, context, context_lens):
         r"""
         Args:
             x(Tensor): Shape [B, L1, C]
@@ -157,20 +156,8 @@ class WanT2VCrossAttention(WanSelfAttention):
 
         # compute query, key, value
         q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d)
-
-        if crossattn_cache is not None:
-            if not crossattn_cache["is_init"]:
-                crossattn_cache["is_init"] = True
-                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)
+        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)
@@ -320,24 +307,9 @@ class WanTransformerBlock(nn.Module):
         bs, seq_length, _ = hidden_states.shape
         orig_dtype = hidden_states.dtype
         # assert orig_dtype != torch.float32
-
-        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)
+        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
@@ -665,21 +637,9 @@ 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_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))
+            timestep, encoder_hidden_states, encoder_hidden_states_image)
+        timestep_proj = timestep_proj.unflatten(1, (6, -1))
 
         if encoder_hidden_states_image is not None:
             encoder_hidden_states = torch.concat(
@@ -716,15 +676,8 @@ class WanTransformer3DModel(CachableDiT):
             if enable_teacache:
                 self.maybe_cache_states(hidden_states, original_hidden_states)
         # 5. Output norm, projection & unpatchify
-        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)
-            
+        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)
 
@@ -828,4 +781,3 @@ class WanTransformer3DModel(CachableDiT):
             return hidden_states + self.previous_residual_even
         else:
             return hidden_states + self.previous_residual_odd
-        

fastvideo/models/registry.py

@@ -25,14 +25,12 @@ _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

@@ -137,24 +137,3 @@ 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

@@ -1,69 +0,0 @@
-# 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,7 +63,6 @@ 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/pipeline_batch_info.py

@@ -241,5 +241,5 @@ class TrainingBatch:
 
 @dataclass
 class PreprocessBatch(ForwardBatch):
-    video_loader: list["VideoDecoder"] | list[str] = field(default_factory=list)
+    video_loader: list["VideoDecoder"] = field(default_factory=list)
     video_file_name: list[str] = field(default_factory=list)

fastvideo/pipelines/pipeline_registry.py

@@ -21,7 +21,6 @@ _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

@@ -1,762 +0,0 @@
-# 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, gettextdataset
-from fastvideo.dataset.dataloader.schema import pyarrow_schema_ode_trajectory, pyarrow_schema_ode_trajectory_text_only
-from fastvideo.distributed import get_local_torch_device
-from fastvideo.fastvideo_args import FastVideoArgs
-from fastvideo.logger import init_logger
-from fastvideo.utils import shallow_asdict, save_decoded_latents_as_video
-from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
-from fastvideo.pipelines.preprocess.preprocess_pipeline_base import (
-    BasePreprocessPipeline)
-from fastvideo.pipelines.stages import (DenoisingStage, ImageVAEEncodingStage,
-                                        InputValidationStage,
-                                        LatentPreparationStage,
-                                        TextEncodingStage,
-                                        TimestepPreparationStage,
-                                        DecodingStage)
-
-logger = init_logger(__name__)
-
-
-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."""
-        # Check if we're using text dataset by checking if the dataset is TextDataset
-        if hasattr(self, 'preprocess_dataloader') and hasattr(self.preprocess_dataloader.dataset, '_process_text_data'):
-            return [f.name for f in pyarrow_schema_ode_trajectory_text_only]
-        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."""
-        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"]
-                # 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)):
-                    prompt_embed = prompt_embed.unsqueeze(0)
-                    prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
-                    logger.info(f"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_frames = 81
-                    batch.fps = args.train_fps
-                    batch.guidance_scale = 3.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(f"===== 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(f"===== 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 preprocess_text_and_trajectory(self,
-                                     fastvideo_args: FastVideoArgs,
-                                     args):
-        """Preprocess text-only data and generate trajectory information."""
-        
-        for batch_idx, data in enumerate(self.pbar):
-            if data is None:
-                continue
-
-            with torch.inference_mode():
-                # For text-only processing, we only need text data
-                # Filter out samples without text
-                valid_indices = []
-                for i, text in enumerate(data["text"]):
-                    if text and text.strip():  # Check if text is not empty
-                        valid_indices.append(i)
-                self.num_processed_samples += len(valid_indices)
-
-                if not valid_indices:
-                    continue
-
-                # Create new batch with only valid samples (text-only)
-                valid_data = {
-                    "text": [data["text"][i] for i in valid_indices],
-                    "path": [data["path"][i] for i in valid_indices],
-                }
-                
-                # Add fps and duration if available in data
-                if "fps" in data:
-                    valid_data["fps"] = [data["fps"][i] for i in valid_indices]
-                if "duration" in data:
-                    valid_data["duration"] = [data["duration"][i] for i in valid_indices]
-
-                batch_captions = valid_data["text"]
-                # 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]
-
-                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)):
-                    prompt_embed = prompt_embed.unsqueeze(0)
-                    prompt_attention_mask = prompt_attention_mask.unsqueeze(0)
-                    
-                    # Collect the trajectory data (text-to-video generation)
-                    batch = ForwardBatch(
-                        **shallow_asdict(sampling_params),
-                    )
-                    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.fps = args.train_fps
-                    batch.guidance_scale = 3.0
-                    batch.do_classifier_free_guidance = True
-
-                    result_batch = self.input_validation_stage(
-                        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.append(result_batch.trajectory_latents.cpu())
-                    trajectory_timesteps.append(result_batch.trajectory_timesteps.cpu())
-                    trajectory_decoded.append(result_batch.trajectory_decoded)
-
-                # Prepare extra features for text-only processing
-                extra_features = {
-                    "trajectory_latents": trajectory_latents,
-                    "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(f"===== 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)
-
-                # 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:
-                    video_name = os.path.basename(video_path).split(".")[0]
-
-                    # Convert tensors to numpy arrays
-                    text_embedding = prompt_embeds[idx].cpu().numpy()
-
-                    # Get extra features for this sample
-                    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(f"===== value in list: not tensor")
-                                    sample_extra_features[key] = value[idx]
-
-                    # Create record for Parquet dataset (without VAE latents for text-only)
-                    record = self.create_text_only_record(
-                        args,
-                        video_name=video_name,
-                        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 = []
-        
-        # Final flush for any remaining samples
-        if hasattr(self, 'all_tables') and self.all_tables and self.num_processed_samples > 0:
-            logger.info(f"Final flush with {self.num_processed_samples} remaining samples")
-            self._flush_tables(self.num_processed_samples, args, self.combined_parquet_dir)
-            self.num_processed_samples = 0
-            self.all_tables = []
-
-    def create_text_only_record(
-            self,
-            args,
-            video_name: str,
-            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 text-only preprocessing using text-only schema."""
-        
-        # Create base record using only fields from text-only schema
-        record = {
-            "id": f"text_{video_name}_{idx}",
-            "text_embedding_bytes": text_embedding.tobytes(),
-            "text_embedding_shape": list(text_embedding.shape),
-            "text_embedding_dtype": str(text_embedding.dtype),
-            "file_name": video_name,
-            "caption": valid_data["text"][idx],
-            "media_type": "text",
-        }
-
-        # Add trajectory data if available
-        if extra_features and "trajectory_latents" in extra_features:
-            trajectory_latents = extra_features["trajectory_latents"][idx] if isinstance(extra_features["trajectory_latents"], list) else 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"][idx] if isinstance(extra_features["trajectory_timesteps"], list) else 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": "",
-            })
-
-        return record
-
-
-    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)
-        train_dataset = gettextdataset(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)
-        self.preprocess_text_and_trajectory(fastvideo_args, args)
-
-
-EntryClass = PreprocessPipeline_ODE_Trajectory

fastvideo/pipelines/preprocess/preprocess_stages.py

@@ -4,11 +4,9 @@ from typing import cast
 
 import numpy as np
 import torch
-import torchvision
 from einops import rearrange
 from torchvision import transforms
 
-from fastvideo.configs.configs import VideoLoaderType
 from fastvideo.dataset.transform import (CenterCropResizeVideo,
                                          TemporalRandomCrop)
 from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
@@ -63,16 +61,7 @@ class VideoTransformStage(PipelineStage):
                 else:
                     frame_indices = frame_indices[:self.num_frames]
 
-            if fastvideo_args.preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC:
-                video = batch.video_loader[i].get_frames_at(frame_indices).data
-            elif fastvideo_args.preprocess_config.video_loader_type == VideoLoaderType.TORCHVISION:
-                video, _, _ = torchvision.io.read_video(batch.video_loader[i],
-                                                        output_format="TCHW")
-                video = video[frame_indices]
-            else:
-                raise ValueError(
-                    f"Invalid video loader type: {fastvideo_args.preprocess_config.video_loader_type}"
-                )
+            video = batch.video_loader[i].get_frames_at(frame_indices).data
             video = self.video_transform(video)
             video_pixel_batch.append(video)
 

fastvideo/pipelines/preprocess/v1_preprocess.py

@@ -9,11 +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.pipelines.preprocess_text import PreprocessPipeline_Text
 from fastvideo.utils import maybe_download_model
 
 logger = init_logger(__name__)
@@ -24,22 +21,12 @@ def main(args) -> None:
     maybe_init_distributed_environment_and_model_parallel(1, 1)
     num_gpus = int(os.environ["WORLD_SIZE"])
     assert num_gpus == 1, "Only support 1 GPU"
-    
-    if args.preprocess_task == "text_only":
-        pipeline_config = PipelineConfig.from_pretrained(args.model_path)
-        kwargs = {
-            "text_encoder_cpu_offload": False,
-        }
-        pipeline_config.update_config_from_dict(kwargs)
-    else:
-        # Full config for video/image processing
-        pipeline_config = PipelineConfig.from_pretrained(args.model_path)
-        kwargs = {
-            "vae_precision": "fp32",
-            "vae_config": WanVAEConfig(load_encoder=True, load_decoder=True),
-        }
-        pipeline_config.update_config_from_dict(kwargs)
-    
+    pipeline_config = PipelineConfig.from_pretrained(args.model_path)
+    kwargs = {
+        "vae_precision": "fp32",
+        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
+    }
+    pipeline_config.update_config_from_dict(kwargs)
     fastvideo_args = FastVideoArgs(
         model_path=args.model_path,
         num_gpus=get_world_size(),
@@ -48,25 +35,7 @@ def main(args) -> None:
         text_encoder_cpu_offload=False,
         pipeline_config=pipeline_config,
     )
-    
-    if args.preprocess_task == "t2v":
-        PreprocessPipeline = PreprocessPipeline_T2V
-    elif args.preprocess_task == "i2v":
-        PreprocessPipeline = PreprocessPipeline_I2V
-    elif args.preprocess_task == "ode_trajectory":
-        print("Preprocess pipeline...")
-        PreprocessPipeline = PreprocessPipeline_ODE_Trajectory
-    elif args.preprocess_task == "text_only":
-        print("Text-only preprocessing pipeline...")
-        PreprocessPipeline = PreprocessPipeline_Text
-    else:
-        raise ValueError(f"Invalid preprocess task: {args.preprocess_task}. "
-                        f"Valid options: t2v, i2v, ode_trajectory, text_only")
-
-    logger.info(
-        f"Preprocess task: {args.preprocess_task} using {PreprocessPipeline.__name__}"
-    )
-
+    PreprocessPipeline = PreprocessPipeline_I2V if args.preprocess_task == "i2v" else PreprocessPipeline_T2V
     pipeline = PreprocessPipeline(args.model_path, fastvideo_args)
     pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
 
@@ -105,11 +74,7 @@ if __name__ == "__main__":
     parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
     parser.add_argument("--group_frame", action="store_true")  # TODO
     parser.add_argument("--group_resolution", action="store_true")  # TODO
-    parser.add_argument("--preprocess_task", 
-                        type=str, 
-                        default="t2v",
-                        choices=["t2v", "i2v", "ode_trajectory", "text_only"],
-                        help="Type of preprocessing task to run")
+    parser.add_argument("--preprocess_task", type=str, default="t2v")
     parser.add_argument("--train_fps", type=int, default=30)
     parser.add_argument("--use_image_num", type=int, default=0)
     parser.add_argument("--text_max_length", type=int, default=256)
@@ -133,4 +98,4 @@ if __name__ == "__main__":
     )
 
     args = parser.parse_args()
-    main(args)
+    main(args)

fastvideo/pipelines/preprocess_text.py

@@ -1,216 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""
-Text-only Data Preprocessing pipeline implementation.
-
-This module contains an implementation of the Text-only Data Preprocessing pipeline
-using the modular pipeline architecture, based on the ODE Trajectory preprocessing.
-"""
-
-import os
-from collections.abc import Iterator
-from typing import Any
-
-import numpy as np
-import pyarrow as pa
-import torch
-from torch.utils.data import DataLoader
-from torchdata.stateful_dataloader import StatefulDataLoader
-from tqdm import tqdm
-
-from fastvideo.dataset import gettextdataset
-from fastvideo.dataset.dataloader.schema import pyarrow_schema_text_only
-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 (TextEncodingStage)
-
-logger = init_logger(__name__)
-
-
-class PreprocessPipeline_Text(BasePreprocessPipeline):
-    """Text-only preprocessing pipeline implementation."""
-
-    _required_config_modules = [
-        "text_encoder", "tokenizer"
-    ]
-    preprocess_dataloader: StatefulDataLoader
-    preprocess_loader_iter: Iterator[dict[str, Any]]
-
-    def get_schema_fields(self):
-        """Get the schema fields for text-only pipeline."""
-        return [f.name for f in pyarrow_schema_text_only]
-
-    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
-        """Set up pipeline stages with proper dependency injection."""
-        self.add_stage(stage_name="prompt_encoding_stage",
-                       stage=TextEncodingStage(
-                           text_encoders=[self.get_module("text_encoder")],
-                           tokenizers=[self.get_module("tokenizer")],
-                       ))
-
-    def preprocess_text_only(self,
-                            fastvideo_args: FastVideoArgs,
-                            args):
-        """Preprocess text-only data."""
-        
-        for batch_idx, data in enumerate(self.pbar):
-            if data is None:
-                continue
-
-            with torch.inference_mode():
-                # For text-only processing, we only need text data
-                # Filter out samples without text
-                valid_indices = []
-                for i, text in enumerate(data["text"]):
-                    if text and text.strip():  # Check if text is not empty
-                        valid_indices.append(i)
-                self.num_processed_samples += len(valid_indices)
-
-                if not valid_indices:
-                    continue
-
-                # Create new batch with only valid samples (text-only)
-                valid_data = {
-                    "text": [data["text"][i] for i in valid_indices],
-                    "path": [data["path"][i] for i in valid_indices],
-                }
-
-                batch_captions = valid_data["text"]
-                # 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]
-
-                logger.info(f"===== prompt_embeds: {prompt_embeds.shape}")
-                logger.info(f"===== prompt_attention_masks: {prompt_attention_masks.shape}")
-
-                # 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, text_path in save_pbar:
-                    text_name = os.path.basename(text_path).split(".")[0]
-
-                    # Convert tensors to numpy arrays
-                    text_embedding = prompt_embeds[idx].cpu().numpy()
-
-                    # Create record for Parquet dataset (text-only)
-                    record = self.create_text_only_record(
-                        text_name=text_name,
-                        text_embedding=text_embedding,
-                        valid_data=valid_data,
-                        idx=idx)
-                    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())))
-                        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 = []
-        
-        # Final flush for any remaining samples
-        if hasattr(self, 'all_tables') and self.all_tables and self.num_processed_samples > 0:
-            logger.info(f"Final flush with {self.num_processed_samples} remaining samples")
-            self._flush_tables(self.num_processed_samples, args, self.combined_parquet_dir)
-            self.num_processed_samples = 0
-            self.all_tables = []
-
-    def create_text_only_record(
-            self,
-            text_name: str,
-            text_embedding: np.ndarray,
-            valid_data: dict[str, Any],
-            idx: int) -> dict[str, Any]:
-        """Create a record for text-only preprocessing using text-only schema."""
-        
-        # Create base record using only fields from text-only schema
-        record = {
-            "id": f"text_{text_name}_{idx}",
-            "text_embedding_bytes": text_embedding.tobytes(),
-            "text_embedding_shape": list(text_embedding.shape),
-            "text_embedding_dtype": str(text_embedding.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 text dataset
-        train_dataset = gettextdataset(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 text preprocessing
-        self.pbar = tqdm(self.preprocess_loader_iter,
-                         desc="Processing text",
-                         unit="batch",
-                         disable=self.local_rank != 0)
-
-        # Initialize class variables for data sharing
-        self.text_data: dict[str, Any] = {}  # Store text metadata and paths
-        
-        self.preprocess_text_only(fastvideo_args, args)
-
-
-EntryClass = PreprocessPipeline_Text

fastvideo/pipelines/stages/__init__.py

@@ -7,7 +7,6 @@ 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,
@@ -31,7 +30,6 @@ __all__ = [
     "ConditioningStage",
     "DenoisingStage",
     "DmdDenoisingStage",
-    "CausalDMDDenosingStage",
     "EncodingStage",
     "DecodingStage",
     "ImageEncodingStage",

fastvideo/pipelines/stages/causal_denoising.py

@@ -1,409 +0,0 @@
-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.repeat(latent_model_input.shape[0])
-
-                    # 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
-                        pred_noise_btchw = self.transformer(
-                            latent_model_input,
-                            prompt_embeds,
-                            t_expand,
-                            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):
-                    _ = self.transformer(
-                        context_bcthw,
-                        prompt_embeds,
-                        t_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/denoising.py

@@ -4,7 +4,6 @@ Denoising stage for diffusion pipelines.
 """
 
 import inspect
-import math
 import weakref
 from collections.abc import Iterable
 from typing import Any
@@ -25,13 +24,12 @@ 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, masks_like
+from fastvideo.utils import dict_to_3d_list
 
 try:
     from fastvideo.attention.backends.sliding_tile_attn import (
@@ -62,13 +60,11 @@ class DenoisingStage(PipelineStage):
                  transformer,
                  scheduler,
                  pipeline=None,
-                 transformer_2=None,
-                 vae=None) -> None:
+                 transformer_2=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(
@@ -197,44 +193,6 @@ 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
 
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
@@ -259,32 +217,19 @@ 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:
-                    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)
-                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)
 
                 # Prepare inputs for transformer
+                t_expand = t.repeat(latent_model_input.shape[0])
                 guidance_expand = (
                     torch.tensor(
                         [fastvideo_args.pipeline_config.embedded_cfg_scale] *
@@ -384,11 +329,6 @@ 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)
-
                 # Update progress bar
                 if i == len(timesteps) - 1 or (
                     (i + 1) > num_warmup_steps and
@@ -838,6 +778,8 @@ 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/encoding.py

@@ -92,6 +92,24 @@ class EncodingStage(PipelineStage):
                 latents = latents.to(vae_dtype)
             latents = self.vae.encode(latents).mean
 
+        # Apply shifting if needed (reverse of decoding)
+        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
+
+        # Apply scaling factor
+        if (hasattr(self.vae, "scaling_factor")
+                and self.vae.scaling_factor is not None):
+            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
+
         # Update batch with encoded latents
         batch.latents = latents
 

fastvideo/pipelines/stages/input_validation.py

@@ -4,8 +4,6 @@ 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
@@ -14,7 +12,6 @@ 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__)
 
@@ -97,7 +94,6 @@ 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]
@@ -105,36 +101,6 @@ 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/text_encoding.py

@@ -59,37 +59,58 @@ class TextEncodingStage(PipelineStage):
         assert len(self.text_encoders) == len(
             fastvideo_args.pipeline_config.text_encoder_configs)
 
-        # 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)
+        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 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)
+            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)
 
         return batch
 
@@ -108,171 +129,6 @@ 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/tests/ssim/README.md

@@ -4,8 +4,7 @@ The reference videos in the `*_reference_videos` directory are used as part of a
 
 run `bash update_reference_videos.sh` from inside the `fastvideo/tests/ssim/` directory after running `test_inference_similarity.py` to update reference videos. Note: make sure to update the path to the corresponding device.
 
-reference videos were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
-causal videos were generated on commit b318063c0a4618f1d5d99ea82ca67a06aad0d19d
+all reference videos are were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
 
 ## Generation Details
 
@@ -77,4 +76,4 @@ Wan2.1-I2V-14B-480P-Diffusers: {
 ### Image-to-Video Prompts
 
 1. "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
-    Image path: "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    Image path: "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"

fastvideo/tests/ssim/test_causal_similarity.py

@@ -1,152 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import json
-import os
-
-import torch
-import pytest
-
-from fastvideo import VideoGenerator
-from fastvideo.logger import init_logger
-from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
-from fastvideo.worker.multiproc_executor import MultiprocExecutor
-
-logger = init_logger(__name__)
-
-device_name = torch.cuda.get_device_name()
-device_reference_folder_suffix = '_reference_videos'
-
-if "A40" in device_name:
-    device_reference_folder = "A40" + device_reference_folder_suffix
-elif "L40S" in device_name:
-    device_reference_folder = "L40S" + device_reference_folder_suffix
-
-# Base parameters from the shell script
-
-SF_WAN_T2V_PARAMS = {
-    "num_gpus": 1,
-    "model_path": "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",
-    "height": 480,
-    "width": 832,
-    "num_frames": 81,
-    "num_inference_steps": 4,
-    "seed": 1024,
-    "sp_size": 1,
-    "tp_size": 1,
-}
-
-
-MODEL_TO_PARAMS = {
-    "SFWan2.1-T2V-1.3B-Diffusers": SF_WAN_T2V_PARAMS,
-}
-
-I2V_MODEL_TO_PARAMS = {
-}
-
-TEST_PROMPTS = [
-    "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
-    # "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature."
-]
-
-I2V_TEST_PROMPTS = [
-    "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.",
-]
-
-I2V_IMAGE_PATHS = [
-    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg",
-]
-
-@pytest.mark.parametrize("prompt", TEST_PROMPTS)
-@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
-@pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys()))
-def test_causal_similarity(prompt, ATTENTION_BACKEND, model_id):
-    """
-    Test that runs inference with different parameters and compares the output
-    to reference videos using SSIM.
-    """
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
-
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-
-    base_output_dir = os.path.join(script_dir, 'generated_videos', model_id)
-    output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
-    output_video_name = f"{prompt[:100]}.mp4"
-
-    os.makedirs(output_dir, exist_ok=True)
-
-    BASE_PARAMS = MODEL_TO_PARAMS[model_id]
-    num_inference_steps = BASE_PARAMS["num_inference_steps"]
-
-    init_kwargs = {
-        "num_gpus": BASE_PARAMS["num_gpus"],
-        "sp_size": BASE_PARAMS["sp_size"],
-        "tp_size": BASE_PARAMS["tp_size"],
-        "dit_cpu_offload": True,
-    }
-    if BASE_PARAMS.get("vae_sp"):
-        init_kwargs["vae_sp"] = True
-        init_kwargs["vae_tiling"] = True
-    #if "text-encoder-precision" in BASE_PARAMS:
-    #    init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"]
-
-    generation_kwargs = {
-        "num_inference_steps": num_inference_steps,
-        "output_path": output_dir,
-        "height": BASE_PARAMS["height"],
-        "width": BASE_PARAMS["width"],
-        "num_frames": BASE_PARAMS["num_frames"],
-        "seed": BASE_PARAMS["seed"],
-    }
-    if "neg_prompt" in BASE_PARAMS:
-        generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"]
-
-    generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs)
-    generator.generate_video(prompt, **generation_kwargs)
-
-    if isinstance(generator.executor, MultiprocExecutor):
-        generator.executor.shutdown()
-
-    assert os.path.exists(
-        output_dir), f"Output video was not generated at {output_dir}"
-
-    reference_folder = os.path.join(script_dir, device_reference_folder, model_id, ATTENTION_BACKEND)
-
-    if not os.path.exists(reference_folder):
-        logger.error("Reference folder missing")
-        raise FileNotFoundError(
-            f"Reference video folder does not exist: {reference_folder}")
-
-    # Find the matching reference video based on the prompt
-    reference_video_name = None
-
-    for filename in os.listdir(reference_folder):
-        if filename.endswith('.mp4') and prompt[:100] in filename:
-            reference_video_name = filename
-            break
-
-    if not reference_video_name:
-        logger.error(f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}")
-        raise FileNotFoundError(f"Reference video missing")
-
-    reference_video_path = os.path.join(reference_folder, reference_video_name)
-    generated_video_path = os.path.join(output_dir, output_video_name)
-
-    logger.info(
-        f"Computing SSIM between {reference_video_path} and {generated_video_path}"
-    )
-    ssim_values = compute_video_ssim_torchvision(reference_video_path,
-                                                 generated_video_path,
-                                                 use_ms_ssim=True)
-
-    mean_ssim = ssim_values[0]
-    logger.info(f"SSIM mean value: {mean_ssim}")
-    logger.info(f"Writing SSIM results to directory: {output_dir}")
-
-    success = write_ssim_results(output_dir, ssim_values, reference_video_path,
-                                 generated_video_path, num_inference_steps,
-                                 prompt)
-
-    if not success:
-        logger.error("Failed to write SSIM results to file")
-
-    min_acceptable_ssim = 0.98
-    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"

fastvideo/tests/ssim/test_inference_similarity.py

@@ -101,7 +101,7 @@ I2V_IMAGE_PATHS = [
 
 
 @pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS)
-@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
+@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"])
 @pytest.mark.parametrize("model_id", list(I2V_MODEL_TO_PARAMS.keys()))
 def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     """

fastvideo/tests/stages/test_text_encoding.py

@@ -1,128 +0,0 @@
-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,7 +29,6 @@ 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 (
@@ -37,7 +36,8 @@ 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, save_distillation_checkpoint, shift_timestep)
+    load_distillation_checkpoint, pred_noise_to_pred_video,
+    save_distillation_checkpoint, shift_timestep)
 from fastvideo.utils import is_vsa_available, set_random_seed
 
 import wandb  # isort: skip

fastvideo/training/training_utils.py

@@ -826,6 +826,27 @@ def custom_to_hf_state_dict(
     return new_state_dict
 
 
+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)
+
+
 def shift_timestep(timestep: torch.Tensor, shift: float,
                    num_train_timestep: float) -> torch.Tensor:
     if shift == 1:

fastvideo/utils.py

@@ -812,66 +812,3 @@ def set_random_seed(seed: int) -> None:
 @lru_cache(maxsize=1)
 def is_vsa_available() -> bool:
     return importlib.util.find_spec("vsa") is not None
-
-
-# adapted from: https://github.com/Wan-Video/Wan2.2/blob/main/wan/utils/utils.py
-def masks_like(tensor,
-               zero=False,
-               generator=None,
-               p=0.2) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
-    assert isinstance(tensor, list)
-    out1 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensor]
-
-    out2 = [torch.ones(u.shape, dtype=u.dtype, device=u.device) for u in tensor]
-
-    if zero:
-        if generator is not None:
-            for u, v in zip(out1, out2, strict=False):
-                random_num = torch.rand(1,
-                                        generator=generator,
-                                        device=generator.device).item()
-                if random_num < p:
-                    u[:, 0] = torch.normal(mean=-3.5,
-                                           std=0.5,
-                                           size=(1, ),
-                                           device=u.device,
-                                           generator=generator).expand_as(
-                                               u[:, 0]).exp()
-                    v[:, 0] = torch.zeros_like(v[:, 0])
-                else:
-                    u[:, 0] = u[:, 0]
-                    v[:, 0] = v[:, 0]
-
-        else:
-            for u, v in zip(out1, out2, strict=False):
-                u[:, 0] = torch.zeros_like(u[:, 0])
-                v[:, 0] = torch.zeros_like(v[:, 0])
-
-    return out1, out2
-
-
-# adapted from: https://github.com/Wan-Video/Wan2.2/blob/main/wan/utils/utils.py
-def best_output_size(w, h, dw, dh, expected_area):
-    # float output size
-    ratio = w / h
-    ow = (expected_area * ratio)**0.5
-    oh = expected_area / ow
-
-    # process width first
-    ow1 = int(ow // dw * dw)
-    oh1 = int(expected_area / ow1 // dh * dh)
-    assert ow1 % dw == 0 and oh1 % dh == 0 and ow1 * oh1 <= expected_area
-    ratio1 = ow1 / oh1
-
-    # process height first
-    oh2 = int(oh // dh * dh)
-    ow2 = int(expected_area / oh2 // dw * dw)
-    assert oh2 % dh == 0 and ow2 % dw == 0 and ow2 * oh2 <= expected_area
-    ratio2 = ow2 / oh2
-
-    # compare ratios
-    if max(ratio / ratio1, ratio1 / ratio) < max(ratio / ratio2,
-                                                 ratio2 / ratio):
-        return ow1, oh1
-    else:
-        return ow2, oh2

fastvideo/version.py

@@ -1 +1 @@
-__version__ = "0.1.6"
+__version__ = "0.1.5"

fastvideo/workflow/preprocess/components.py

@@ -13,8 +13,7 @@ import pyarrow.parquet as pq
 import torch
 from datasets import Dataset, Video, load_dataset
 
-from fastvideo.configs.configs import (DatasetType, PreprocessConfig,
-                                       VideoLoaderType)
+from fastvideo.configs.configs import DatasetType, PreprocessConfig
 from fastvideo.distributed.parallel_state import get_world_rank, get_world_size
 from fastvideo.logger import init_logger
 from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
@@ -432,8 +431,7 @@ def build_dataset(preprocess_config: PreprocessConfig, split: str,
             return item
 
         dataset = dataset.map(add_video_column)
-        if preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC:
-            dataset = dataset.cast_column("video", Video())
+        dataset = dataset.cast_column("video", Video())
     else:
         raise ValueError(
             f"Invalid dataset type: {preprocess_config.dataset_type}")

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "fastvideo"
-version = "0.1.6"
+version = "0.1.5"
 description = "FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"

scripts/checkpoint_conversion/wan_to_diffusers.py

@@ -1,161 +0,0 @@
-from huggingface_hub import save_torch_state_dict, load_state_dict_from_file
-# from safetensors import safetensors
-from safetensors.torch import save_file
-import torch
-import re
-from collections import OrderedDict
-
-_param_names_mapping: dict = {
-    r"^text_embedding\.0\.(.*)$":
-    r"condition_embedder.text_embedder.linear_1.\1",
-    r"^text_embedding\.2\.(.*)$":
-    r"condition_embedder.text_embedder.linear_2.\1",
-    r"^time_embedding\.0\.(.*)$":
-    r"condition_embedder.time_embedder.linear_1.\1",
-    r"^time_embedding\.2\.(.*)$":
-    r"condition_embedder.time_embedder.linear_2.\1",
-    r"^time_projection\.1\.(.*)$":
-    r"condition_embedder.time_proj.\1",
-    r"^img_emb\.proj\.0\.(.*)$":
-    r"condition_embedder.image_embedder.norm1.\1",
-    r"^img_emb\.proj\.1\.(.*)$":
-    r"condition_embedder.image_embedder.ff.net.0.proj.\1",
-    r"^img_emb\.proj\.3\.(.*)$":
-    r"condition_embedder.image_embedder.ff.net.2.\1",
-    r"^img_emb\.proj\.4\.(.*)$":
-    r"condition_embedder.image_embedder.norm2.\1",
-    r"^head\.modulation":
-    r"scale_shift_table",
-    r"^head\.head\.(.*)$":
-    r"proj_out.\1",
-    r"^blocks\.(\d+)\.self_attn\.q\.(.*)$":
-    r"blocks.\1.attn1.to_q.\2",
-    r"^blocks\.(\d+)\.self_attn\.k\.(.*)$":
-    r"blocks.\1.attn1.to_k.\2",
-    r"^blocks\.(\d+)\.self_attn\.v\.(.*)$":
-    r"blocks.\1.attn1.to_v.\2",
-    r"^blocks\.(\d+)\.self_attn\.o\.(.*)$":
-    r"blocks.\1.attn1.to_out.0.\2",
-    r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$":
-    r"blocks.\1.attn1.norm_q.\2",
-    r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$":
-    r"blocks.\1.attn1.norm_k.\2",
-    r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$":
-    r"blocks.\1.attn2.to_q.\2",
-    r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$":
-    r"blocks.\1.attn2.to_k.\2",
-    r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$":
-    r"blocks.\1.attn2.add_k_proj.\2",
-    r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$":
-    r"blocks.\1.attn2.to_v.\2",
-    r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$":
-    r"blocks.\1.attn2.add_v_proj.\2",
-    r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$":
-    r"blocks.\1.attn2.to_out.0.\2",
-    r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
-    r"blocks.\1.attn2.norm_q.\2",
-    r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
-    r"blocks.\1.attn2.norm_k.\2",
-    r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$":
-    r"blocks.\1.attn2.norm_added_k.\2",
-    r"^blocks\.(\d+)\.ffn\.0\.(.*)$":
-    r"blocks.\1.ffn.net.0.proj.\2",
-    r"^blocks\.(\d+)\.ffn\.2\.(.*)$":
-    r"blocks.\1.ffn.net.2.\2",
-    r"^blocks\.(\d+)\.modulation":
-    r"blocks.\1.scale_shift_table",
-    r"^blocks\.(\d+)\.norm3\.(.*)$":
-    r"blocks.\1.norm2.\2",
-}
-
-# The following mapping has an extra 'patch_embedding' field and also contains
-# the 'model' prefixes
-_self_forcing_to_diffusers_param_names_mapping: dict = {
-    r"^model.patch_embedding\.(.*)$":
-    r"patch_embedding.\1",
-    r"^model.text_embedding\.0\.(.*)$":
-    r"condition_embedder.text_embedder.linear_1.\1",
-    r"^model.text_embedding\.2\.(.*)$":
-    r"condition_embedder.text_embedder.linear_2.\1",
-    r"^model.time_embedding\.0\.(.*)$":
-    r"condition_embedder.time_embedder.linear_1.\1",
-    r"^model.time_embedding\.2\.(.*)$":
-    r"condition_embedder.time_embedder.linear_2.\1",
-    r"^model.time_projection\.1\.(.*)$":
-    r"condition_embedder.time_proj.\1",
-    r"^model.img_emb\.proj\.0\.(.*)$":
-    r"condition_embedder.image_embedder.norm1.\1",
-    r"^model.img_emb\.proj\.1\.(.*)$":
-    r"condition_embedder.image_embedder.ff.net.0.proj.\1",
-    r"^model.img_emb\.proj\.3\.(.*)$":
-    r"condition_embedder.image_embedder.ff.net.2.\1",
-    r"^model.img_emb\.proj\.4\.(.*)$":
-    r"condition_embedder.image_embedder.norm2.\1",
-    r"^model.head\.modulation":
-    r"scale_shift_table",
-    r"^model.head\.head\.(.*)$":
-    r"proj_out.\1",
-    r"^model.blocks\.(\d+)\.self_attn\.q\.(.*)$":
-    r"blocks.\1.attn1.to_q.\2",
-    r"^model.blocks\.(\d+)\.self_attn\.k\.(.*)$":
-    r"blocks.\1.attn1.to_k.\2",
-    r"^model.blocks\.(\d+)\.self_attn\.v\.(.*)$":
-    r"blocks.\1.attn1.to_v.\2",
-    r"^model.blocks\.(\d+)\.self_attn\.o\.(.*)$":
-    r"blocks.\1.attn1.to_out.0.\2",
-    r"^model.blocks\.(\d+)\.self_attn\.norm_q\.(.*)$":
-    r"blocks.\1.attn1.norm_q.\2",
-    r"^model.blocks\.(\d+)\.self_attn\.norm_k\.(.*)$":
-    r"blocks.\1.attn1.norm_k.\2",
-    r"^model.blocks\.(\d+)\.cross_attn\.q\.(.*)$":
-    r"blocks.\1.attn2.to_q.\2",
-    r"^model.blocks\.(\d+)\.cross_attn\.k\.(.*)$":
-    r"blocks.\1.attn2.to_k.\2",
-    r"^model.blocks\.(\d+)\.cross_attn\.k_img\.(.*)$":
-    r"blocks.\1.attn2.add_k_proj.\2",
-    r"^model.blocks\.(\d+)\.cross_attn\.v\.(.*)$":
-    r"blocks.\1.attn2.to_v.\2",
-    r"^model.blocks\.(\d+)\.cross_attn\.v_img\.(.*)$":
-    r"blocks.\1.attn2.add_v_proj.\2",
-    r"^model.blocks\.(\d+)\.cross_attn\.o\.(.*)$":
-    r"blocks.\1.attn2.to_out.0.\2",
-    r"^model.blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$":
-    r"blocks.\1.attn2.norm_q.\2",
-    r"^model.blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$":
-    r"blocks.\1.attn2.norm_k.\2",
-    r"^model.blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$":
-    r"blocks.\1.attn2.norm_added_k.\2",
-    r"^model.blocks\.(\d+)\.ffn\.0\.(.*)$":
-    r"blocks.\1.ffn.net.0.proj.\2",
-    r"^model.blocks\.(\d+)\.ffn\.2\.(.*)$":
-    r"blocks.\1.ffn.net.2.\2",
-    r"^model.blocks\.(\d+)\.modulation":
-    r"blocks.\1.scale_shift_table",
-    r"^model.blocks\.(\d+)\.norm3\.(.*)$":
-    r"blocks.\1.norm2.\2",
-}
-
-state_dict = load_state_dict_from_file("checkpoints/self_forcing_dmd.pt")
-state_dict = state_dict["generator_ema"]
-new_state_dict = OrderedDict()
-for k, v in state_dict.items():
-    new_key = k
-    for pattern, replacement in _self_forcing_to_diffusers_param_names_mapping.items():
-        if re.match(pattern, k):
-            new_key = re.sub(pattern, replacement, k)
-            break  # Stop at the first match
-    else:
-        # print(f"No match found for {k}")
-        raise ValueError(f"No match found for {k}")
-    new_state_dict[new_key] = v
-    if "norm_added_k" in new_key:
-        dummy_key = new_key.replace("norm_added_k", "norm_added_q")
-        dummy_value = torch.zeros_like(v)
-        new_state_dict[dummy_key] = dummy_value
-del state_dict
-
-save_torch_state_dict(
-    new_state_dict,
-    "new2/",
-    max_shard_size="10GB"
-)

scripts/data/analyze.py

@@ -0,0 +1,32 @@
+import ray  # type: ignore
+import ray.data as rd  # type: ignore
+
+ray.init(ignore_reinit_error=True)
+
+jsonl_path = "/mnt/fast-disks/nfs/hao_lab/FastVideo/logs/user_prompts.jsonl"
+prompts_file = "/mnt/fast-disks/nfs/hao_lab/FastVideo/prompts/prompts_final.txt"
+video_outputs_dir = "/mnt/fast-disks/nfs/hao_lab/FastVideo/outputs/"
+
+# Load blocklist of prompts (exact line match, ignoring only newline chars)
+with open(prompts_file, "r", encoding="utf-8") as f:
+    blocklist = {line.rstrip("\r\n") for line in f if line.rstrip("\r\n")}
+
+ds = rd.read_json(jsonl_path)
+
+# Filter out any rows whose 'prompt' is in the text file
+# (batchwise Pandas filtering is efficient for large datasets)
+ds_filtered = ds.map_batches(
+    lambda df: df[~df["prompt"].isin(blocklist)],
+    batch_format="pandas",
+)
+
+# Keep one row per prompt; choose deterministically by latest timestamp
+ds_dedup = ds_filtered.groupby("prompt").map_groups(
+    lambda df: df.sort_values("timestamp", ascending=False).head(1),
+    batch_format="pandas",
+)
+
+print("before:", ds.count(), "after filter:", ds_filtered.count(), "after dedup:", ds_dedup.count())
+# Write a single JSONL file (not sharded)
+out_path = "/mnt/fast-disks/nfs/hao_lab/FastVideo/logs/user_prompts_dedup.jsonl"
+ds_dedup.to_pandas().to_json(out_path, orient="records", lines=True, force_ascii=False)

scripts/preprocess/launch_slurm.sh

@@ -1,21 +0,0 @@
-#!/bin/bash
-
-# Create output directory if it doesn't exist
-mkdir -p preprocess_output
-
-# Launch 8 jobs, one for each node
-# Each node processes 8 consecutive files (64 total files / 8 nodes = 8 files per node)
-for node_id in {0..1}; do
-    # Calculate the starting file number for this node
-    start_file=$((node_id * 8 + 1))
-    
-    echo "Launching node $node_id with files v2m_${start_file}.txt to v2m_$((start_file + 7)).txt"
-    echo "sbatch --job-name=ode-${node_id} --output=preprocess_output/preprocess-node-${node_id}.out --error=preprocess_output/preprocess-node-${node_id}.err /mnt/weka/home/hao.zhang/wl/kev/preproc/slurms/syn.slurm $start_file $node_id"
-    
-    sbatch --job-name=ode-${node_id} \
-           --output=preprocess_output/preprocess-node-${node_id}.out \
-           --error=preprocess_output/preprocess-node-${node_id}.err \
-           /mnt/weka/home/hao.zhang/wl/kev/preproc/slurms/syn.slurm $start_file $node_id
-done
-
-echo "All 8 nodes launched successfully!"

scripts/preprocess/launch_slurm_text.sh

@@ -1,21 +0,0 @@
-#!/bin/bash
-
-# Create output directory if it doesn't exist
-mkdir -p preprocess_output_text
-
-# Launch 8 jobs, one for each node
-# Each node processes 8 consecutive files (64 total files / 8 nodes = 8 files per node)
-for node_id in {0..1}; do
-    # Calculate the starting file number for this node
-    start_file=$((node_id * 8 + 1))
-    
-    echo "Launching text-only node $node_id with files v2m_${start_file}.txt to v2m_$((start_file + 7)).txt"
-    echo "sbatch --job-name=text-${node_id} --output=preprocess_output_text/preprocess-text-node-${node_id}.out --error=preprocess_output_text/preprocess-text-node-${node_id}.err /mnt/weka/home/hao.zhang/matthew/FastVideo/scripts/preprocess/syn_text.slurm $start_file $node_id"
-    
-    sbatch --job-name=text-${node_id} \
-           --output=preprocess_output_text/preprocess-text-node-${node_id}.out \
-           --error=preprocess_output_text/preprocess-text-node-${node_id}.err \
-           /mnt/weka/home/hao.zhang/matthew/FastVideo/scripts/preprocess/syn_text.slurm $start_file $node_id
-done
-
-echo "All 8 text-only nodes launched successfully!" 

scripts/preprocess/syn.slurm

@@ -1,88 +0,0 @@
-#!/bin/bash
-#SBATCH --partition=main
-#SBATCH --qos=hao
-#SBATCH --nodes=1
-#SBATCH --ntasks-per-node=8
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=16
-#SBATCH --mem=960G
-#SBATCH --exclusive
-#SBATCH --time=72:00:00
-
-# conda init
-# source ~/conda/miniconda/bin/activate
-# PYTHON_VIRTUAL_ENVIRONMENT=fastvideo-train-yq
-# conda activate $PYTHON_VIRTUAL_ENVIRONMENT
-nvidia-smi
-nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv
-
-echo " "
-echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
-echo " GPUs per node:= " $SLURM_JOB_GPUS
-echo " Running on multiple nodes/GPU devices"
-echo ""
-echo " Run started at:- "
-date
-
-# Accept parameters from launch script
-START_FILE=${1:-1}  # Starting file number for this node
-NODE_ID=${2:-0}  # Node identifier (0-7)
-
-num_gpus=1
-export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-# Start port number - we'll increment for each job
-base_port=$((29603 + NODE_ID * 100))  # Different port range per node
-
-# Create an array of CUDA device IDs
-gpu_ids=(0 1 2 3 4 5 6 7)
-
-GPU_NUM=1
-MODEL_TYPE="wan"
-
-echo "NODE_ID: $NODE_ID"
-echo "START_FILE: $START_FILE"
-echo "Base port for this node: $base_port"
-
-# Run 8 parallel preprocessing jobs on this node
-for i in {1..8}; do
-    # Calculate port for this job
-    port=$((base_port + i))
-    
-    # Get GPU ID using modulo to cycle through available GPUs
-    gpu=${gpu_ids[((i-1))]}
-    
-    # Calculate which file this GPU should process
-    file_num=$((START_FILE + i - 1))
-    DATA_MERGE_PATH="/mnt/weka/home/hao.zhang/wl/kev/preproc/prompts/v2m_${file_num}.txt"
-    
-    # Create unique output directory based on node and GPU
-    OUTPUT_DIR="data/test-ode-preprocessing/Node_${NODE_ID}_GPU_${i}_File_${file_num}"
-    
-    start_cpu=$(( (i-1)*2 ))  # Reduced CPU allocation for 8 nodes
-    end_cpu=$(( start_cpu+1 ))
-    
-    echo "Starting GPU $gpu processing file v2m_${file_num}.txt on port $port, output: $OUTPUT_DIR"
-    
-    # Run the preprocessing command in background
-    CUDA_VISIBLE_DEVICES=$gpu taskset -c ${start_cpu}-${end_cpu} torchrun --nnodes=1 --nproc_per_node=$GPU_NUM --master_port $port \
-        /mnt/weka/home/hao.zhang/wl/kev/FastVideo/fastvideo/pipelines/preprocess/v1_preprocess.py \
-        --model_path $MODEL_BASE \
-        --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 "ode_trajectory" &
-done
-
-# Wait for all jobs on this node to complete
-wait
-
-echo "All processing blocks completed!"

scripts/preprocess/syn_text.slurm

@@ -1,89 +0,0 @@
-#!/bin/bash
-#SBATCH --partition=main
-#SBATCH --qos=hao
-#SBATCH --nodes=1
-#SBATCH --ntasks-per-node=8
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=16
-#SBATCH --mem=960G
-#SBATCH --exclusive
-#SBATCH --time=72:00:00
-
-# conda init
-# source ~/conda/miniconda/bin/activate
-# PYTHON_VIRTUAL_ENVIRONMENT=fastvideo-train-yq
-# conda activate $PYTHON_VIRTUAL_ENVIRONMENT
-nvidia-smi
-nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv
-
-echo " "
-echo " Number of nodes:= " $SLURM_JOB_NUM_NODES
-echo " GPUs per node:= " $SLURM_JOB_GPUS
-echo " Running on multiple nodes/GPU devices for TEXT-ONLY preprocessing"
-echo ""
-echo " Run started at:- "
-date
-
-# Accept parameters from launch script
-START_FILE=${1:-1}  # Starting file number for this node
-NODE_ID=${2:-0}  # Node identifier (0-7)
-
-num_gpus=1
-export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-# Start port number - we'll increment for each job
-base_port=$((29603 + NODE_ID * 100))  # Different port range per node
-
-# Create an array of CUDA device IDs
-gpu_ids=(0 1 2 3 4 5 6 7)
-
-GPU_NUM=1
-MODEL_TYPE="wan"
-
-echo "NODE_ID: $NODE_ID"
-echo "START_FILE: $START_FILE"
-echo "Base port for this node: $base_port"
-echo "Processing TEXT-ONLY data"
-
-# Run 8 parallel preprocessing jobs on this node
-for i in {1..8}; do
-    # Calculate port for this job
-    port=$((base_port + i))
-    
-    # Get GPU ID using modulo to cycle through available GPUs
-    gpu=${gpu_ids[((i-1))]}
-    
-    # Calculate which file this GPU should process
-    file_num=$((START_FILE + i - 1))
-    DATA_MERGE_PATH="/mnt/weka/home/hao.zhang/wl/kev/preproc/prompts/v2m_${file_num}.txt"
-    
-    # Create unique output directory based on node and GPU for text-only processing
-    OUTPUT_DIR="data/test-text-preprocessing/Node_${NODE_ID}_GPU_${i}_File_${file_num}"
-    
-    start_cpu=$(( (i-1)*2 ))  # Reduced CPU allocation for 8 nodes
-    end_cpu=$(( start_cpu+1 ))
-    
-    echo "Starting GPU $gpu processing text-only file v2m_${file_num}.txt on port $port, output: $OUTPUT_DIR"
-    
-    # Run the text-only preprocessing command in background
-    CUDA_VISIBLE_DEVICES=$gpu taskset -c ${start_cpu}-${end_cpu} torchrun --nnodes=1 --nproc_per_node=$GPU_NUM --master_port $port \
-        /mnt/weka/home/hao.zhang/matthew/FastVideo/fastvideo/pipelines/preprocess/v1_preprocess.py \
-        --model_path $MODEL_BASE \
-        --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 "text" &
-done
-
-# Wait for all jobs on this node to complete
-wait
-
-echo "All text-only processing blocks completed!"