preprocessing text

[misc] Update Slack invite link (#786 )
Update WeChat Link
2025-09-06 05:56:13 +00:00 · 2025-09-05 12:16:18 -07:00 · 2025-09-05 11:40:47 -07:00 · 2025-09-05 01:23:01 -07:00 · 2025-09-04 17:51:27 -07:00 · 2025-09-03 20:47:33 -07:00
99 changed files with 2577 additions and 2422 deletions
@@ -117,11 +117,11 @@ steps:
                  queue: "default"
            - path:
                - "fastvideo/**"
-                - "csrc/attn/vsa/**"
-                - "csrc/attn/tk/**"
-                - "csrc/attn/setup_vsa.py"
-                - "csrc/attn/config_vsa.py"
-                - "csrc/attn/vsa.cpp"
+                - "csrc/attn/video_sparse_attn/**"
+                - "csrc/attn/video_sparse_attn/tk/**"
+                - "csrc/attn/video_sparse_attn/setup.py"
+                - "csrc/attn/video_sparse_attn/config_vsa.py"
+                - "csrc/attn/video_sparse_attn/vsa.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -133,10 +133,10 @@ steps:
                  queue: "default"
            - path:
                - "fastvideo/**"
-                - "csrc/attn/st_attn/**"
-                - "csrc/attn/setup_sta.py"
-                - "csrc/attn/config_sta.py"
-                - "csrc/attn/st_attn.cpp"
+                - "csrc/attn/sliding_tile_attn/**"
+                - "csrc/attn/sliding_tile_attn/setup.py"
+                - "csrc/attn/sliding_tile_attn/config_sta.py"
+                - "csrc/attn/sliding_tile_attn/st_attn.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -147,10 +147,10 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "csrc/attn/st_attn/**"
-                - "csrc/attn/setup_sta.py"
-                - "csrc/attn/config_sta.py"
-                - "csrc/attn/st_attn.cpp"
+                - "csrc/attn/sliding_tile_attn/**"
+                - "csrc/attn/sliding_tile_attn/setup.py"
+                - "csrc/attn/sliding_tile_attn/config_sta.py"
+                - "csrc/attn/sliding_tile_attn/st_attn.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -161,12 +161,12 @@ steps:
                agents:
                  queue: "default"
            - path:
-                - "csrc/attn/vsa/**"
-                - "csrc/attn/tk/**"
+                - "csrc/attn/video_sparse_attn/**"
+                - "csrc/attn/video_sparse_attn/tk/**"
                - "csrc/attn/tests/test_vsa.py"
-                - "csrc/attn/setup_vsa.py"
-                - "csrc/attn/config_vsa.py"
-                - "csrc/attn/vsa.cpp"
+                - "csrc/attn/video_sparse_attn/setup.py"
+                - "csrc/attn/video_sparse_attn/config_vsa.py"
+                - "csrc/attn/video_sparse_attn/vsa.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
@@ -104,16 +104,17 @@ jobs:
              - 'pyproject.toml'
              - 'docker/Dockerfile.python3.12'
            sta-kernel-paths: &sta-kernel-paths
-              - 'csrc/attn/st_attn/**'
-              - 'csrc/attn/setup_sta.py'
-              - 'csrc/attn/config_sta.py'
-              - 'csrc/attn/st_attn.cpp'
+              - 'csrc/attn/sliding_tile_attn/**'
+              - 'csrc/attn/sliding_tile_attn/tk/**'
+              - 'csrc/attn/sliding_tile_attn/setup.py'
+              - 'csrc/attn/sliding_tile_attn/config_sta.py'
+              - 'csrc/attn/sliding_tile_attn/st_attn.cpp'
            vsa-kernel-paths: &vsa-kernel-paths
-              - 'csrc/attn/vsa/**'
-              - 'csrc/attn/tk/**'
-              - 'csrc/attn/setup_vsa.py'
-              - 'csrc/attn/config_vsa.py'
-              - 'csrc/attn/vsa.cpp'
+              - 'csrc/attn/video_sparse_attn/**'
+              - 'csrc/attn/video_sparse_attn/tk/**'
+              - 'csrc/attn/video_sparse_attn/setup.py'
+              - 'csrc/attn/video_sparse_attn/config_vsa.py'
+              - 'csrc/attn/video_sparse_attn/vsa.cpp'
            vsa-paths: &vsa-paths
              - 'fastvideo/**'
              - *common-paths
@@ -234,7 +235,7 @@ jobs:
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-  
+
  training-test:
    needs: change-filter
    if: >-
@@ -372,4 +373,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
@@ -5,7 +5,7 @@ on:
    branches:
      - main
    paths:
-      - "csrc/attn/setup_sta.py"
+      - "csrc/attn/sliding_tile_attn/setup.py"
  workflow_dispatch:

 jobs:
@@ -23,13 +23,13 @@ jobs:
      - name: Check if version changed
        id: check-version
        run: |
-          cd csrc/attn
+          cd csrc/attn/sliding_tile_attn
          # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_sta.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
          echo "New version: $NEW_VERSION"
          
          # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup_sta.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
          echo "Old version: $OLD_VERSION"
          
          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -144,13 +144,13 @@ jobs:
          pip install setuptools
          pip install ninja packaging wheel

-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/sliding_tile_attn  # Move into the correct folder
          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_sta.py bdist_wheel --dist-dir=dist
+          python setup.py bdist_wheel --dist-dir=dist

      - name: Rename wheel file
        run: |
-          cd csrc/attn
+          cd csrc/attn/sliding_tile_attn

          CUDA_SHORT_VERSION=$(echo ${{ matrix.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
          TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-version }} | cut -d. -f1,2)
@@ -165,7 +165,7 @@ jobs:
        uses: actions/upload-artifact@v4
        with:
          name: ${{ env.wheel_name }}
-          path: csrc/attn/dist/*.whl
+          path: csrc/attn/sliding_tile_attn/dist/*.whl
          retention-days: 90

  publish_package:
@@ -239,11 +239,11 @@ jobs:
          pip install setuptools
          pip install ninja packaging wheel

-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/sliding_tile_attn  # Move into the correct folder
          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_sta.py sdist --dist-dir=dist
+          python setup.py sdist --dist-dir=dist

      - name: Publish release distributions to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
-          packages-dir: csrc/attn/dist/
+          packages-dir: csrc/attn/sliding_tile_attn/dist/
@@ -5,7 +5,7 @@ on:
    branches:
      - main
    paths:
-      - "csrc/attn/setup_vsa.py"
+      - "csrc/attn/video_sparse_attn/setup.py"
  workflow_dispatch:

 jobs:
@@ -23,13 +23,13 @@ jobs:
      - name: Check if version changed
        id: check-version
        run: |
-          cd csrc/attn
+          cd csrc/attn/video_sparse_attn
          # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_vsa.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
          echo "New version: $NEW_VERSION"
          
          # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup_vsa.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
          echo "Old version: $OLD_VERSION"
          
          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -152,13 +152,13 @@ jobs:
          pip install setuptools
          pip install ninja packaging wheel

-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/video_sparse_attn  # Move into the correct folder
          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_vsa.py bdist_wheel --dist-dir=dist
+          python setup.py bdist_wheel --dist-dir=dist

      - name: Rename wheel file
        run: |
-          cd csrc/attn
+          cd csrc/attn/video_sparse_attn

          CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
          TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2)
@@ -173,7 +173,7 @@ jobs:
        uses: actions/upload-artifact@v4
        with:
          name: ${{ env.wheel_name }}
-          path: csrc/attn/dist/*.whl
+          path: csrc/attn/video_sparse_attn/dist/*.whl
          retention-days: 90

  publish_package:
@@ -247,11 +247,11 @@ jobs:
          pip install setuptools
          pip install ninja packaging wheel

-          cd csrc/attn  # Move into the correct folder
+          cd csrc/attn/video_sparse_attn  # Move into the correct folder
          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_vsa.py sdist --dist-dir=dist
+          python setup.py sdist --dist-dir=dist

      - name: Publish release distributions to PyPI
        uses: pypa/gh-action-pypi-publish@release/v1
        with:
-          packages-dir: csrc/attn/dist/
+          packages-dir: csrc/attn/video_sparse_attn/dist/
@@ -64,5 +64,3 @@ docs/source/distillation/examples/
 !docs/source/_static/images/**/*.png
 !comfyui/assets/**/*.png
 !comfyui/assets/**/*.gif
-
-dmd_t2v_output/
@@ -1,3 +1,7 @@
-[submodule "csrc/attn/tk"]
-	path = csrc/attn/tk
+[submodule "csrc/attn/video_sparse_attn/tk"]
+	path = csrc/attn/video_sparse_attn/tk
+	url = https://github.com/HazyResearch/ThunderKittens.git
+
+[submodule "csrc/attn/sliding_tile_attn/tk"]
+	path = csrc/attn/sliding_tile_attn/tk
 	url = https://github.com/HazyResearch/ThunderKittens.git
@@ -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-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/rG0QpZdw" 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-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/S7HLCSTh" target="_blank"> <b> WeChat </b> </a> |
 </p>

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

-from csrc.attn.config_sta import kernels, sources, target
+from config_sta import kernels, sources, target
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension

@@ -9,7 +9,7 @@ target = target.lower()

 # Package metadata
 PACKAGE_NAME = "st_attn"
-VERSION = "0.0.4"
+VERSION = "0.0.6"
 AUTHOR = "Hao AI Lab"
 DESCRIPTION = "Sliding Tile Atteniton Kernel Used in FastVideo"
 URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/sliding_tile_attention"
@@ -0,0 +1,2 @@
+recursive-include tk *
+include config_vsa.py
@@ -0,0 +1,61 @@
+
+
+# Attention Kernel Used in FastVideo
+
+## Video Sparse Attention (VSA)
+
+### Installation
+We support H100 (via TK) and any other GPU (via triton) for VSA.
+
+```bash
+pip install vsa
+``` 
+
+Install from source:
+
+```bash
+git submodule update --init --recursive
+python setup.py install
+```
+
+
+If you encounter error during installation, try below:
+Install C++20 for ThunderKittens:
+```bash
+sudo apt update
+sudo apt install gcc-11 g++-11
+
+sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+sudo apt update
+sudo apt install clang-11
+```
+(If you use CUDA12.8)
+```bash
+export CUDA_HOME=/usr/local/cuda-12.8
+export PATH=${CUDA_HOME}/bin:${PATH} 
+export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+### Verify if you have successfully installed
+
+```bash
+# test numerical
+python ../tests/test_vsa.py
+# (For H100) test speed
+python ../benchmarks/bench_vsa_hopper.py
+```
+
+bench_vsa_hopper.py should print something like this:
+
+```bash
+Using topk=76 kv blocks per q block (out of 768 total kv blocks)
+
+=== BLOCK SPARSE ATTENTION BENCHMARK ===
+Block Sparse Forward  - TFLOPS: 5622.26
+Block Sparse Backward - TFLOPS: 3865.68
+```
+
+## Acknowledgement
+
+We learned or reuse code from FlexAtteniton, NATEN, and ThunderKittens.
@@ -9,10 +9,10 @@ target = target.lower()

 # Package metadata
 PACKAGE_NAME = "vsa"
-VERSION = "0.0.1"
+VERSION = "0.0.3"
 AUTHOR = "Hao AI Lab"
 DESCRIPTION = "Video Sparse Attention Kernel Used in FastVideo"
-URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn"
+URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn/video_sparse_attn"

 # Set environment variables
 tk_root = os.getenv('THUNDERKITTENS_ROOT', os.path.abspath(os.path.join(os.getcwd(), 'tk/')))
@@ -58,15 +58,15 @@ RUN source $HOME/.local/bin/env && \
 # Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/sliding_tile_attn && \
    git submodule update --init --recursive && \
-    python setup_sta.py install
+    python setup.py install

 # Install VSA 
 RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
    git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install

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

 # Install VSA 
 RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
    git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install

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

 # Install VSA 
 RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
    git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install

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

 # Install VSA 
 RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
-    cd csrc/attn && \
+    cd csrc/attn/video_sparse_attn && \
    git submodule update --init --recursive && \
-    python setup_vsa.py install
+    python setup.py install

-EXPOSE 22
+EXPOSE 22
@@ -4,7 +4,7 @@
 You can install the Sliding Tile Attention package using

 ```
-pip install st_attn==0.0.4
+pip install st_attn
 ``` 

 # Building from Source
@@ -12,7 +12,6 @@ We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have impleme
 First, install C++20 for ThunderKittens:

 ```bash
-cd csrc/sliding_tile_attention/
 sudo apt update
 sudo apt install gcc-11 g++-11

@@ -22,14 +21,20 @@ sudo apt update
 sudo apt install clang-11
 ```

-Install STA:
+Set up CUDA environment (if using CUDA 12.4):

 ```bash
 export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+```
+
+Install STA:
+
+```bash
+cd csrc/attn/sliding_tile_attn/
 git submodule update --init --recursive
-python setup_sta.py install
+python setup.py install
 ```

 # 🧪 Test
@@ -4,8 +4,7 @@
 You can install the Video Sparse Attention package using

 ```bash
-git submodule update --init --recursive
-python setup_vsa.py install
+pip install vsa
 ```

 # Building from Source
@@ -34,9 +33,9 @@ export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 Install VSA:

 ```bash
-cd csrc/attn/
+cd csrc/attn/video_sparse_attn/
 git submodule update --init --recursive
-python setup_vsa.py install
+python setup.py install
 ```

 # 🧪 Test
@@ -1,165 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=t2v
-#SBATCH --partition=main
-#SBATCH --nodes=1
-#SBATCH --ntasks=1
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:1
-#SBATCH --cpus-per-task=128
-#SBATCH --mem=1440G
-#SBATCH --output=dmd_t2v_output/t2v_%j.out
-#SBATCH --error=dmd_t2v_output/t2v_%j.err
-#SBATCH --exclusive
-set -e -x
-
-# Environment Setup
-source ~/conda/miniconda/bin/activate
-conda activate wei-fv
-
-# Basic Info
-export WANDB_MODE="online"
-export NCCL_P2P_DISABLE=1
-export TORCH_NCCL_ENABLE_MONITORING=0
-# different cache dir for different processes
-export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
-export MASTER_PORT=29500
-export NODE_RANK=$SLURM_PROCID
-nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
-export MASTER_ADDR=${nodes[0]}
-export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
-export TOKENIZERS_PARALLELISM=false
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=1
-
-# Model paths for Self-Forcing DMD distillation:
-GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
-REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
-FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
-
-DATA_DIR="data/crush-smol-single_processed_t2v/combined_parquet_dataset/"
-VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
-# export CUDA_VISIBLE_DEVICES=4,5
-# IP=[MASTER NODE IP]
-
-# Training arguments
-training_args=(
-  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  # Updated for self-forcing DMD
-  --output_dir "checkpoints/SFwan_t2v_finetune"
-  --max_train_steps 500
-  --train_batch_size 1
-  --train_sp_batch_size 1
-  --gradient_accumulation_steps 1
-  --num_latent_t 21
-  --num_height 480
-  --num_width 832
-  --num_frames 81  # Must be divisible by num_frame_per_block (81 % 3 = 0 ✓)
-  --enable_gradient_checkpointing_type "full"
-  --log_visualization
-  --simulate_generator_forward
-  --num_frame_per_block 3  # Frame generation block size for self-forcing
-  --enable_gradient_masking
-  --gradient_mask_last_n_frames 21
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus 1 # 64
-  --sp_size 1
-  --tp_size 1
-  --hsdp_replicate_dim 1 # 64
-  --hsdp_shard_dim 1
-)
-
-# Model arguments
-model_args=(
-  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
-  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
-  --generator_model_path $GENERATOR_MODEL_PATH
-  --real_score_model_path $REAL_SCORE_MODEL_PATH
-  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 4
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_dataset_file "$VALIDATION_DATASET_FILE"
-  --validation_steps 10
-  --validation_sampling_steps "4"
-  --validation_guidance_scale "6.0" # not used for dmd inference
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --training_state_checkpointing_steps 50
-  --weight_only_checkpointing_steps 50
-  --weight_decay 0.01
-  --betas '0.0,0.999'
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --checkpoints_total_limit 3
-  --training_cfg_rate 0.0
-  --dit_precision "fp32"
-  --flow_shift 5
-  --seed 1000
-  --use_ema True
-  --ema_decay 0.99
-  --ema_start_step 100
-  --init_weights_from_safetensors "/mnt/weka/home/hao.zhang/wl/Self-Forcing/diffusers_ode_init/model.safetensors"
-)
-
-# Self-forcing DMD arguments
-dmd_args=(
-  --dmd_denoising_steps '1000,750,500,250'
-  --min_timestep_ratio 0.02
-  --max_timestep_ratio 0.98
-  --dfake_gen_update_ratio 5
-  --real_score_guidance_scale 3.0
-  --fake_score_learning_rate 8e-6
-  --fake_score_betas '0.0,0.999'
-)
-
-# Self-forcing specific arguments
-self_forcing_args=(
-  --independent_first_frame False  # Whether to treat first frame independently
-  --same_step_across_blocks False  # Whether to use same denoising step across all blocks
-  --last_step_only False  # Whether to only use the last denoising step
-  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
-  --validate_cache_structure False  # Set to True for debugging KV cache issues
-)
-
-srun torchrun \
--nnodes $SLURM_JOB_NUM_NODES \
--nproc_per_node $NUM_GPUS \
--node_rank $SLURM_PROCID \
--rdzv_backend=c10d \
--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
-    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
-    "${parallel_args[@]}" \
-    "${model_args[@]}" \
-    "${dataset_args[@]}" \
-    "${training_args[@]}" \
-    "${optimizer_args[@]}" \
-    "${validation_args[@]}" \
-    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}" \
-    "${self_forcing_args[@]}"
@@ -1,3 +0,0 @@
-#!/bin/bash
-
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -1,43 +0,0 @@
-#!/bin/bash
-
-# Download the full dataset first
-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
-
-# Create a single-example dataset for debugging
-SINGLE_EXAMPLE_DIR="data/crush-smol-single"
-mkdir -p "$SINGLE_EXAMPLE_DIR/videos"
-
-# Copy the specific video that matches the validation.json style (macaron crushing)
-cp "data/crush-smol/videos/7P02AihYkCU-Scene-005.mp4" "$SINGLE_EXAMPLE_DIR/videos/"
-
-# Create a single-line videos.txt
-echo "videos/7P02AihYkCU-Scene-005.mp4" > "$SINGLE_EXAMPLE_DIR/videos.txt"
-
-# Create a single-line prompt.txt with the macaron crushing prompt
-echo "PIKA_CRUSH A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out." > "$SINGLE_EXAMPLE_DIR/prompt.txt"
-
-# Generate the JSON file and merge.txt for the single example
-python scripts/dataset_preparation/prepare_json_file.py --data_folder "$SINGLE_EXAMPLE_DIR" --output "videos2caption.json"
-
-# Create a validation.json that uses the same example for consistency
-cat > "$SINGLE_EXAMPLE_DIR/validation.json" << 'EOF'
-{
-    "data": [
-      {
-        "caption": "A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out.",
-        "image_path": null,
-        "video_path": null,
-        "num_inference_steps": 50,
-        "height": 480,
-        "width": 832,
-        "num_frames": 81
-      }
-    ]
-  }
-EOF
-
-echo "Single example dataset created at $SINGLE_EXAMPLE_DIR"
-echo "Contains:"
-echo "- 1 video: $(cat $SINGLE_EXAMPLE_DIR/videos.txt)"
-echo "- 1 prompt: $(cat $SINGLE_EXAMPLE_DIR/prompt.txt)"
-echo "- Validation file created with the same example for consistency" 
@@ -1,24 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=1 # 2,4,8
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_TYPE="wan"
-DATA_MERGE_PATH="data/crush-smol/merge.txt"
-OUTPUT_DIR="data/crush-smol_processed_t2v/"
-
-torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/pipelines/preprocess/v1_preprocess.py \
-    --model_path $MODEL_PATH \
-    --data_merge_path $DATA_MERGE_PATH \
-    --preprocess_video_batch_size 8 \
-    --seed 42 \
-    --max_height 480 \
-    --max_width 832 \
-    --num_frames 81 \
-    --dataloader_num_workers 0 \
-    --output_dir=$OUTPUT_DIR \
-    --train_fps 16 \
-    --samples_per_file 8 \
-    --flush_frequency 8 \
-    --video_length_tolerance_range 5 \
-    --preprocess_task "t2v" 
@@ -1,29 +0,0 @@
-#!/bin/bash
-
-GPU_NUM=1 # 2,4,8
-MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_TYPE="wan"
-DATA_MERGE_PATH="data/crush-smol-single/merge.txt"
-OUTPUT_DIR="data/crush-smol-single_processed_t2v/"
-
-torchrun --nproc_per_node=$GPU_NUM \
-    fastvideo/pipelines/preprocess/v1_preprocess.py \
-    --model_path $MODEL_PATH \
-    --data_merge_path $DATA_MERGE_PATH \
-    --preprocess_video_batch_size 1 \
-    --seed 42 \
-    --max_height 480 \
-    --max_width 832 \
-    --num_frames 81 \
-    --dataloader_num_workers 0 \
-    --output_dir=$OUTPUT_DIR \
-    --train_fps 16 \
-    --samples_per_file 1 \
-    --flush_frequency 1 \
-    --video_length_tolerance_range 5 \
-    --preprocess_task "t2v"
-
-# Copy the validation.json to the output directory for consistency
-cp "data/crush-smol-single/validation.json" "$OUTPUT_DIR/"
-
-echo "Preprocessing completed. Validation file copied to $OUTPUT_DIR/" 
@@ -1,31 +0,0 @@
-{
-    "data": [
-      {
-        "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
-        "image_path": null,
-        "video_path": null,
-        "num_inference_steps": 50,
-        "height": 480,
-        "width": 832,
-        "num_frames": 77
-      },
-      {
-        "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
-        "image_path": null,
-        "video_path": null,
-        "num_inference_steps": 50,
-        "height": 480,
-        "width": 832,
-        "num_frames": 77
-      },
-      {
-        "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
-        "image_path": null,
-        "video_path": null,
-        "num_inference_steps": 50,
-        "height": 480,
-        "width": 832,
-        "num_frames": 77
-      }
-    ]
-  }
@@ -4,9 +4,7 @@ These are end-to-end example scripts for distilling Wan2.1 T2V 1.3B model using
 ### 0. Make sure you have installed VSA

 ```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
+pip install vsa
 ```

 ### 1. Download dataset:
@@ -4,9 +4,7 @@ These are end-to-end example scripts for distilling Wan2.2 TI2V 5B model DMD+VSA
 ### 0. Make sure you have installed VSA

 ```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
+pip install vsa
 ```

 ### Data-free Distillation
@@ -4,9 +4,7 @@ These are end-to-end example scripts for distilling Wan2.2 TI2V 5B model DMD+VSA
 ### 0. Make sure you have installed VSA

 ```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
+pip install vsa
 ```

 ### 1. Download dataset:
@@ -28,4 +28,4 @@ def main():
    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)

 if __name__ == "__main__":
-    main()
+    main()
@@ -0,0 +1,41 @@
+from fastvideo import VideoGenerator
+
+OUTPUT_PATH = "video_samples_wan2_2_5B_ti2v"
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    model_name = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=True,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        # image_encoder_cpu_offload=False,
+    )
+
+    # I2V is triggered just by passing in an image_path argument
+    prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
+    image_path = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, image_path=image_path)
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+
+    # T2V mode
+    prompt2 = (
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
+
+
+if __name__ == "__main__":
+    main()
@@ -7,9 +7,7 @@ These are e2e example scripts for finetuning Wan2.1 T2V with VSA to accelerate i
 ## Make sure you have installed VSA

 ```bash
-cd csrc/attn
-git submodule update --init --recursive
-python setup_vsa.py install
+pip install vsa
 ```

 ### Download the synthetic dataset:
@@ -10,6 +10,7 @@ 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 \
@@ -32,6 +32,29 @@ 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."""
@@ -51,6 +74,7 @@ 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
@@ -120,6 +144,12 @@ 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,
@@ -174,6 +204,10 @@ 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(
@@ -109,4 +109,4 @@ class WanVideoArchConfig(DiTArchConfig):
 class WanVideoConfig(DiTConfig):
    arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig)

-    prefix: str = "Wan"
+    prefix: str = "Wan"
@@ -85,6 +85,9 @@ 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

@@ -7,11 +7,14 @@ from collections.abc import Callable
 from fastvideo.configs.pipelines.base import PipelineConfig
 from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
 from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig
-from fastvideo.configs.pipelines.wan import (FastWan2_1_T2V_480P_Config,
-                                             FastWan2_2_TI2V_5B_Config,
-                                             SelfForcingWanT2V480PConfig,
-                                             WanI2V480PConfig, WanI2V720PConfig,
-                                             WanT2V480PConfig, WanT2V720PConfig)
+
+# isort: off
+from fastvideo.configs.pipelines.wan import (
+    FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config,
+    SelfForcingWanT2V480PConfig, Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config,
+    Wan2_2_TI2V_5B_Config, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig,
+    WanT2V720PConfig)
+# isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
                             verify_model_config_and_directory)
@@ -32,10 +35,10 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
    "FastVideo/FastWan2.2-TI2V-5B-Diffusers": FastWan2_2_TI2V_5B_Config,
    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
    "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
-    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig,
-    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": WanT2V480PConfig,
-    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": WanI2V480PConfig,
    "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig,
+    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_Config,
+    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config,
+    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config,
    # Add other specific weight variants
 }

@@ -57,7 +57,7 @@ class WanT2V480PConfig(PipelineConfig):
    # WanConfig-specific added parameters

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


@@ -146,6 +146,5 @@ class Wan2_2_I2V_A14B_Config(WanT2V480PConfig):
@dataclass
 class SelfForcingWanT2V480PConfig(WanT2V480PConfig):
    is_causal: bool = True
-    flow_shift: int = 5
    dmd_denoising_steps: list[int] | None = field(
-        default_factory=lambda: [1000, 750, 500, 250])
+        default_factory=lambda: [1000, 750, 500, 250])
@@ -10,6 +10,7 @@ from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
 # isort: off
 from fastvideo.configs.sample.wan import (
    FastWanT2V480PConfig,
+    Wan2_1_Fun_1_3B_InP_SamplingParam,
    Wan2_2_I2V_A14B_SamplingParam,
    Wan2_2_T2V_A14B_SamplingParam,
    Wan2_2_TI2V_5B_SamplingParam,
@@ -36,6 +37,8 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
    "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
    "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
@@ -107,6 +107,21 @@ class FastWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
    fps: int = 16


+# =============================================
+# ============= Wan2.1 Fun Models =============
+# =============================================
+@dataclass
+class Wan2_1_Fun_1_3B_InP_SamplingParam(SamplingParam):
+    """Sampling parameters for Wan2.1 Fun 1.3B InP model."""
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 16
+    negative_prompt: str | None = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+    guidance_scale: float = 6.0
+    num_inference_steps: int = 50
+
+
 # =============================================
 # ============= Wan2.2 TI2V Models =============
 # =============================================
@@ -148,4 +163,4 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
 # =============================================
@dataclass
 class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParam):
-    pass
+    pass
@@ -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
+from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset, TextDataset
 from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255,
                                         TemporalRandomCrop)
 from fastvideo.dataset.validation_dataset import ValidationDataset
@@ -39,7 +39,13 @@ 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"
-]
+    "VideoCaptionMergedDataset", "TextDataset"
+]
@@ -78,3 +78,76 @@ 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()),
+])
@@ -628,3 +628,134 @@ 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"]
@@ -591,11 +591,6 @@ class TrainingArgs(FastVideoArgs):
    pretrained_model_name_or_path: str = ""
    dit_model_name_or_path: str = ""

-    # DMD model paths - separate paths for each network
-    generator_model_path: str = ""  # path for generator (student) model
-    real_score_model_path: str = ""  # path for real score (teacher) model
-    fake_score_model_path: str = ""  # path for fake score (critic) model
-
    # diffusion setting
    ema_decay: float = 0.0
    ema_start_step: int = 0
@@ -618,7 +613,6 @@ class TrainingArgs(FastVideoArgs):
    checkpoints_total_limit: int = 0
    checkpointing_steps: int = 0
    resume_from_checkpoint: str = ""  # specify the checkpoint folder to resume from
-    init_weights_from_safetensors: str = ""  # path to safetensors file for initial weight loading

    # optimizer & scheduler
    num_train_epochs: int = 0
@@ -650,7 +644,6 @@ class TrainingArgs(FastVideoArgs):
    linear_quadratic_threshold: float = 0.0
    linear_range: float = 0.0
    weight_decay: float = 0.0
-    betas: str = "0.9,0.999"  # betas for optimizer, format: "beta1,beta2"
    use_ema: bool = False
    multi_phased_distill_schedule: str = ""
    pred_decay_weight: float = 0.0
@@ -671,29 +664,16 @@ class TrainingArgs(FastVideoArgs):

    # distillation args
    generator_update_interval: int = 5
-    dfake_gen_update_ratio: int = 5  # self-forcing: how often to train generator vs critic
    min_timestep_ratio: float = 0.2
    max_timestep_ratio: float = 0.98
    real_score_guidance_scale: float = 3.5
    fake_score_learning_rate: float = 0.0  # separate learning rate for fake_score_transformer, if 0.0, use learning_rate
    fake_score_lr_scheduler: str = "constant"  # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler
-    fake_score_betas: str = "0.9,0.999"  # betas for fake score optimizer, format: "beta1,beta2"
    training_state_checkpointing_steps: int = 0  # for resuming training
    weight_only_checkpointing_steps: int = 0  # for inference
    log_visualization: bool = False
    # simulate generator forward to match inference
    simulate_generator_forward: bool = False
-    warp_denoising_step: bool = False
-
-    # Self-forcing specific arguments
-    num_frame_per_block: int = 3
-    independent_first_frame: bool = False
-    enable_gradient_masking: bool = True
-    gradient_mask_last_n_frames: int = 21
-    validate_cache_structure: bool = False  # Debug flag for cache validation
-    same_step_across_blocks: bool = False  # Use same exit timestep for all blocks
-    last_step_only: bool = False  # Only use the last timestep for training
-    context_noise: int = 0  # Context noise level for cache updates

    @classmethod
    def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -795,20 +775,6 @@ class TrainingArgs(FastVideoArgs):
                            type=str,
                            help="Directory to cache models")

-        # DMD model paths - separate paths for each network
-        parser.add_argument(
-            "--generator-model-path",
-            type=str,
-            help="Path to generator (student) model for DMD distillation")
-        parser.add_argument(
-            "--real-score-model-path",
-            type=str,
-            help="Path to real score (teacher) model for DMD distillation")
-        parser.add_argument(
-            "--fake-score-model-path",
-            type=str,
-            help="Path to fake score (critic) model for DMD distillation")
-
        # Diffusion settings
        parser.add_argument("--ema-decay",
                            type=float,
@@ -879,10 +845,6 @@ class TrainingArgs(FastVideoArgs):
        parser.add_argument("--resume-from-checkpoint",
                            type=str,
                            help="Path to checkpoint to resume from")
-        parser.add_argument(
-            "--init-weights-from-safetensors",
-            type=str,
-            help="Path to safetensors file for initial weight loading")
        parser.add_argument("--logging-dir",
                            type=str,
                            help="Directory for logging")
@@ -987,10 +949,6 @@ class TrainingArgs(FastVideoArgs):
                            help="Linear quadratic threshold")
        parser.add_argument("--linear-range", type=float, help="Linear range")
        parser.add_argument("--weight-decay", type=float, help="Weight decay")
-        parser.add_argument("--betas",
-                            type=str,
-                            default=TrainingArgs.betas,
-                            help="Betas for optimizer (format: 'beta1,beta2')")
        parser.add_argument("--use-ema",
                            action=StoreBoolean,
                            help="Whether to use EMA")
@@ -1032,13 +990,6 @@ class TrainingArgs(FastVideoArgs):
                            type=int,
                            default=TrainingArgs.generator_update_interval,
                            help="Ratio of student updates to critic updates.")
-        parser.add_argument(
-            "--dfake-gen-update-ratio",
-            type=int,
-            default=TrainingArgs.dfake_gen_update_ratio,
-            help=
-            "Self-forcing: How often to train generator vs critic (train generator every N steps)."
-        )
        parser.add_argument("--min-timestep-ratio",
                            type=float,
                            default=TrainingArgs.min_timestep_ratio,
@@ -1055,11 +1006,6 @@ class TrainingArgs(FastVideoArgs):
                            type=float,
                            default=TrainingArgs.fake_score_learning_rate,
                            help="Learning rate for fake score transformer")
-        parser.add_argument(
-            "--fake-score-betas",
-            type=str,
-            default=TrainingArgs.fake_score_betas,
-            help="Betas for fake score optimizer (format: 'beta1,beta2')")
        parser.add_argument(
            "--fake-score-lr-scheduler",
            type=str,
@@ -1072,49 +1018,6 @@ class TrainingArgs(FastVideoArgs):
            "--simulate-generator-forward",
            action=StoreBoolean,
            help="Whether to simulate generator forward to match inference")
-        parser.add_argument(
-            "--warp-denoising-step",
-            action=StoreBoolean,
-            help=
-            "Whether to warp denoising step according to the scheduler time shift"
-        )
-
-        # Self-forcing specific arguments
-        parser.add_argument(
-            "--num-frame-per-block",
-            type=int,
-            default=TrainingArgs.num_frame_per_block,
-            help="Number of frames per block for causal generation")
-        parser.add_argument(
-            "--independent-first-frame",
-            action=StoreBoolean,
-            help="Whether the first frame is independent in causal generation")
-        parser.add_argument(
-            "--enable-gradient-masking",
-            action=StoreBoolean,
-            help="Whether to enable frame-level gradient masking")
-        parser.add_argument(
-            "--gradient-mask-last-n-frames",
-            type=int,
-            default=TrainingArgs.gradient_mask_last_n_frames,
-            help="Number of last frames to enable gradients for")
-        parser.add_argument(
-            "--validate-cache-structure",
-            action=StoreBoolean,
-            help="Whether to validate KV cache structure (debug flag)")
-        parser.add_argument(
-            "--same-step-across-blocks",
-            action=StoreBoolean,
-            help="Whether to use the same exit timestep for all blocks")
-        parser.add_argument(
-            "--last-step-only",
-            action=StoreBoolean,
-            help="Whether to only use the last timestep for training")
-        parser.add_argument(
-            "--context-noise",
-            type=int,
-            default=TrainingArgs.context_noise,
-            help="Context noise level for cache updates")

        return parser

@@ -9,9 +9,6 @@ import torch.nn.functional as F
 from fastvideo.layers.custom_op import CustomOp
 from fastvideo.platforms import current_platform

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

@CustomOp.register("rms_norm")
 class RMSNorm(CustomOp):
@@ -103,13 +100,7 @@ class ScaleResidual(nn.Module):
    def forward(self, residual: torch.Tensor, x: torch.Tensor,
                gate: torch.Tensor) -> torch.Tensor:
        """Apply gated residual connection."""
-        # logger.info("x.shape: %s", x.shape)
-        # if isinstance(gate, torch.Tensor):
-            # logger.info("gate.shape: %s", gate.shape)
-
-        num_frames = gate.shape[1]
-        frame_seqlen = x.shape[1] // num_frames
-        return residual + (x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2)
+        return residual + x * gate


 # adapted from Diffusers: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/normalization.py
@@ -181,35 +172,11 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
              but before normalization)
        """
        # Apply residual connection with gating
-        # logger.info("x.shape: %s", x.shape)
-        if isinstance(gate, int): 
-            # used by cross-attention, should be 1
-            assert gate == 1
-            residual_output = residual + x * gate
-        elif isinstance(gate, torch.Tensor):
-            # logger.info("gate.shape: %s", gate.shape)
-            if gate.dim() == 3:
-                # used by bidirectional self attention
-                residual_output = residual + x * gate
-            else:
-                assert gate.dim() == 4 
-                num_frames = gate.shape[1]
-                frame_seqlen = x.shape[1] // num_frames
-                residual_output = residual + (x.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * gate).flatten(1, 2)
-                # residual_output = residual + x * gate
-        else:
-            raise ValueError(f"Gate type {type(gate)} not supported")
-        # logger.info("residual_output.shape: %s", residual_output.shape)
-
+        residual_output = residual + x * gate
        # Apply normalization
        normalized = self.norm(residual_output)
        # Apply scale and shift
-        if isinstance(scale, torch.Tensor) and scale.dim() == 4:
-            num_frames = scale.shape[1]
-            frame_seqlen = normalized.shape[1] // num_frames
-            modulated = (normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2)
-        else:
-            modulated = normalized * (1.0 + scale) + shift
+        modulated = normalized * (1.0 + scale) + shift
        return modulated, residual_output


@@ -252,15 +219,7 @@ class LayerNormScaleShift(nn.Module):
                scale: torch.Tensor) -> torch.Tensor:
        """Apply ln followed by scale and shift in a single fused operation."""
        normalized = self.norm(x)
-        if scale.dim() == 4:
-            num_frames = scale.shape[1]
-            frame_seqlen = normalized.shape[1] // num_frames
-            if self.compute_dtype == torch.float32:
-                return (normalized.float().unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2).to(x.dtype)
-            else:
-                return (normalized.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1.0 + scale) + shift).flatten(1, 2)
+        if self.compute_dtype == torch.float32:
+            return (normalized.float() * (1.0 + scale) + shift).to(x.dtype)
        else:
-            if self.compute_dtype == torch.float32:
-                return (normalized.float() * (1.0 + scale) + shift).to(x.dtype)
-            else:
-                return normalized * (1.0 + scale) + shift
+            return normalized * (1.0 + scale) + shift
@@ -33,7 +33,6 @@ from fastvideo.logger import init_logger

 logger = init_logger(__name__)

-logger = init_logger(__name__)

 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., :x.shape[-1] // 2]
@@ -176,4 +176,4 @@ def unpatchify(x, t, h, w, patch_size, channels) -> torch.Tensor:
    x = torch.einsum("nthwcopq->nctohpwq", x)
    imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw))

-    return imgs
+    return imgs
@@ -147,10 +147,6 @@ class CausalWanSelfAttention(nn.Module):
                # 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"] = kv_cache["k"].clone()
-                kv_cache["v"] = kv_cache["v"].clone()
-
                kv_cache["k"][:, local_start_index:local_end_index] = roped_key
                kv_cache["v"][:, local_start_index:local_end_index] = v
            x = self.attn(
@@ -248,36 +244,19 @@ class CausalWanTransformerBlock(nn.Module):
        current_start: int = 0,
        cache_start: int | None = None,
    ) -> torch.Tensor:
-        logger.info("temb.shape: %s", temb.shape)
-        num_frames = temb.shape[1]
-        logger.info("first hidden_states.shape: %s", hidden_states.shape)
-        logger.info("num_frames: %s", num_frames)
        if hidden_states.dim() == 4:
            hidden_states = hidden_states.squeeze(1)
-        frame_seqlen = hidden_states.shape[1] // temb.shape[1]
-        logger.info("frame_seqlen: %s", frame_seqlen)
        bs, seq_length, _ = hidden_states.shape
        orig_dtype = hidden_states.dtype
        # assert orig_dtype != torch.float32
        e = self.scale_shift_table + temb.float()
-        logger.info("e.shape: %s", e.shape)
        shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
-            6, dim=2)
+            6, dim=1)
        assert shift_msa.dtype == torch.float32

        # 1. Self-attention
-        logger.info("hidden_states.shape: %s", hidden_states.shape)
-        logger.info("scale_msa.shape: %s", scale_msa.shape)
-        logger.info("shift_msa.shape: %s", shift_msa.shape)
-        
-        norm_hidden_states_unflattened = self.norm1(hidden_states.float()).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
-        # logger.info("norm_hidden_states_unflattened.shape: %s", norm_hidden_states_unflattened.shape)
-        
-        # norm_hidden_states = (self.norm1(hidden_states.float()) *
-        #                       (1 + scale_msa) + shift_msa).to(orig_dtype)
-        norm_hidden_states = (norm_hidden_states_unflattened *
-                              (1 + scale_msa) + shift_msa).flatten(1, 2).to(orig_dtype)
-        # logger.info("1 norm_hidden_states.shape: %s", norm_hidden_states.shape)
+        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)
@@ -299,8 +278,6 @@ class CausalWanTransformerBlock(nn.Module):
        null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
        norm_hidden_states, hidden_states = self.self_attn_residual_norm(
            hidden_states, attn_output, gate_msa, null_shift, null_scale)
-        # logger.info("after self_attn_residual_norm norm_hidden_states.shape: %s", norm_hidden_states.shape)
-        # logger.info("after self_attn_residual_norm hidden_states.shape: %s", hidden_states.shape)
        norm_hidden_states, hidden_states = norm_hidden_states.to(
            orig_dtype), hidden_states.to(orig_dtype)

@@ -311,16 +288,12 @@ class CausalWanTransformerBlock(nn.Module):
                                 crossattn_cache=crossattn_cache)
        norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
            hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
-        # logger.info("after cross_attn_residual_norm norm_hidden_states.shape: %s", norm_hidden_states.shape)
-        # logger.info("after cross_attn_residual_norm hidden_states.shape: %s", hidden_states.shape)
        norm_hidden_states, hidden_states = norm_hidden_states.to(
            orig_dtype), hidden_states.to(orig_dtype)

        # 3. Feed-forward
        ff_output = self.ffn(norm_hidden_states)
        hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
-        # logger.info("after mlp_residual norm_hidden_states.shape: %s", norm_hidden_states.shape)
-        logger.info("after mlp_residual hidden_states.shape: %s", hidden_states.shape)
        hidden_states = hidden_states.to(orig_dtype)

        return hidden_states
@@ -386,10 +359,8 @@ class CausalWanTransformer3DModel(BaseDiT):
                                            elementwise_affine=False,
                                            dtype=torch.float32,
                                            compute_dtype=torch.float32)
-        # Debug: Log configuration values
-        proj_out_dim = config.out_channels * math.prod(config.patch_size)
-        
-        self.proj_out = nn.Linear(inner_dim, proj_out_dim)
+        self.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)

@@ -478,7 +449,6 @@ class CausalWanTransformer3DModel(BaseDiT):
        This function will be run for num_frame times.
        Process the latent frames one by one (1560 tokens each)
        """
-        # logger.info("forward inference hidden_states.shape: %s", hidden_states.shape)

        orig_dtype = hidden_states.dtype
        if not isinstance(encoder_hidden_states, torch.Tensor):
@@ -515,13 +485,10 @@ class CausalWanTransformer3DModel(BaseDiT):

        hidden_states = self.patch_embedding(hidden_states)
        hidden_states = hidden_states.flatten(2).transpose(1, 2)
-        # logger.info("forward inference flattened and transposed hidden_states.shape: %s", hidden_states.shape)
-
-        logger.info("timestep shape: %s", timestep.shape)

        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
-        timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
+            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(
@@ -559,14 +526,8 @@ class CausalWanTransformer3DModel(BaseDiT):
                                        **causal_kwargs)

        # 5. Output norm, projection & unpatchify
-        # logger.info("===== INFERENCE 5. Output norm, projection & unpatchify")
-        # logger.info("hidden_states.shape: %s", hidden_states.shape)
-        # logger.info("temb.shape: %s", temb.shape)
-        temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
-        # logger.info("WTFWTF train temb.shape: %s", temb.shape)
-        # logger.info("WTFWTF train self.scale_shift_table.shape: %s", self.scale_shift_table.shape)
-        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2,
-                                                                          dim=2)
+        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)

@@ -588,8 +549,6 @@ class CausalWanTransformer3DModel(BaseDiT):
                start_frame: int = 0,
                **kwargs) -> torch.Tensor:

-        # logger.info("===== forward train hidden_states.shape: %s", hidden_states.shape)
-        # logger.info("===== forward train timestep.shape: %s", timestep.shape)
        orig_dtype = hidden_states.dtype
        if not isinstance(encoder_hidden_states, torch.Tensor):
            encoder_hidden_states = encoder_hidden_states[0]
@@ -635,14 +594,10 @@ class CausalWanTransformer3DModel(BaseDiT):

        hidden_states = self.patch_embedding(hidden_states)
        hidden_states = hidden_states.flatten(2).transpose(1, 2)
-        # logger.info("forward train flattened and transposed hidden_states.shape: %s", hidden_states.shape)

        temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
-            timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
-        # logger.info("forward train timestep_proj.shape: %s", timestep_proj.shape)
-        # logger.info("forward train timestep.shape: %s", timestep.shape)
-        # logger.info("forward train temb.shape: %s", temb.shape)
-        timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
+            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(
@@ -662,35 +617,16 @@ class CausalWanTransformer3DModel(BaseDiT):
                    timestep_proj, freqs_cis,
                    block_mask=self.block_mask)
        else:
-            for block_index, block in enumerate(self.blocks):
-                logger.info("===== TRAIN block %d", block_index)
-                logger.info("hidden_states.shape: %s", hidden_states.shape)
-                # logger.info("encoder_hidden_states.shape: %s", encoder_hidden_states.shape)
-                logger.info("timestep_proj.shape: %s", timestep_proj.shape)
-                # logger.info("freqs_cis.shape: %s", freqs_cis.shape)
-                # logger.info("block_mask.shape: %s", self.block_mask.shape)
+            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
-        # logger.info("===== TRAIN 5. Output norm, projection & unpatchify")
-        # logger.info("hidden_states.shape: %s", hidden_states.shape)
-        # logger.info("temb.shape: %s", temb.shape)
-        # shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
-        temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)
-        # logger.info("WTFWTF train temb.shape: %s", temb.shape)
-        # logger.info("WTFWTF train self.scale_shift_table.shape: %s", self.scale_shift_table.shape)
-        shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2,
-                                                                          dim=2)
-        # logger.info("DEBUG scale.shape: %s", scale.shape)
-        # logger.info("DEBUG shift.shape: %s", shift.shape)
+        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)
-        # logger.info("DEBUG after proj_out hidden_states.shape: %s", hidden_states.shape)
-        # logger.info(f"DEBUG reshape dimensions: batch_size={batch_size}, post_patch_num_frames={post_patch_num_frames}")
-        # logger.info(f"DEBUG reshape dimensions: post_patch_height={post_patch_height}, post_patch_width={post_patch_width}")
-        # logger.info(f"DEBUG patch dimensions: p_t={p_t}, p_h={p_h}, p_w={p_w}")

        hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
                                              post_patch_height,
@@ -709,4 +645,4 @@ class CausalWanTransformer3DModel(BaseDiT):
        if kwargs.get('kv_cache', None) is not None:
            return self._forward_inference(*args, **kwargs)
        else:
-            return self._forward_train(*args, **kwargs)
+            return self._forward_train(*args, **kwargs)
@@ -430,16 +430,6 @@ class TransformerLoader(ComponentLoader):
        if not safetensors_list:
            raise ValueError(f"No safetensors files found in {model_path}")

-        # Check if we should use custom initialization weights
-        custom_weights_path = getattr(fastvideo_args, 'init_weights_from_safetensors', None)
-        use_custom_weights = (custom_weights_path and os.path.exists(custom_weights_path) and 
-                            fastvideo_args.training_mode and 
-                            not hasattr(fastvideo_args, '_loading_teacher_critic_model'))
-        
-        if use_custom_weights:
-            logger.info("Using custom initialization weights from: %s", custom_weights_path)
-            safetensors_list = [custom_weights_path]
-
        logger.info("Loading model from %s safetensors files in %s",
                    len(safetensors_list), model_path)

@@ -248,13 +248,6 @@ def load_model_from_full_model_state_dict(
    sharded_sd = {}
    custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
        full_sd_iterator, param_names_mapping)  # type: ignore
-    print(custom_param_sd.keys())
-    print("--------------------------------")
-    print("--------------------------------")
-    print("--------------------------------")
-    print("--------------------------------")
-    print("--------------------------------")
-    print(meta_sd.keys())
    for target_param_name, full_tensor in custom_param_sd.items():
        meta_sharded_param = meta_sd.get(target_param_name)
        if meta_sharded_param is None:
@@ -636,14 +636,8 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
        timestep: torch.IntTensor,
    ) -> torch.Tensor:
        self.sigmas = self.sigmas.to(noise.device)
-        # TODO: hack
-        if timestep.ndim == 2 and timestep.shape[1] == 1:
-            timestep = timestep.expand(clean_latent.shape[0])
+        timestep = timestep.expand(clean_latent.shape[0])
        self.timesteps = self.timesteps.to(noise.device)
-        logger.info("self.timesteps shape: %s", self.timesteps.shape)
-        logger.info("timestep shape: %s", timestep.shape)
-        if timestep.ndim > 1:
-            timestep = timestep.squeeze(0)
        timestep_id = torch.argmin(
            (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
@@ -6,9 +6,6 @@ from typing import Any
 import torch


-from fastvideo.logger import init_logger
-logger = init_logger(__name__)
-
 # TODO(PY): move it elsewhere
 def auto_attributes(init_func):
    """
@@ -149,9 +146,6 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
    """
    Convert predicted noise to clean latent.
    """
-    logger.info(f"timestep: {timestep.shape}")
-    logger.info(f"noise_input_latent: {noise_input_latent.shape}")
-    logger.info(f"pred_noise: {pred_noise.shape}")
    timestep = timestep.expand(noise_input_latent.shape[0])
    dtype = pred_noise.dtype
    device = pred_noise.device
@@ -163,4 +157,4 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor,
        (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)
+    return pred_video.to(dtype)
@@ -27,7 +27,6 @@ from fastvideo.layers.activation import get_act_fn
 from fastvideo.models.vaes.common import (DiagonalGaussianDistribution,
                                          ParallelTiledVAE)
 from fastvideo.platforms import current_platform
-from fastvideo.logger import init_logger

 CACHE_T = 2

@@ -36,7 +35,6 @@ feat_cache = contextvars.ContextVar("feat_cache", default=None)
 feat_idx = contextvars.ContextVar("feat_idx", default=0)
 first_chunk = contextvars.ContextVar("first_chunk", default=None)

-logger = init_logger(__name__)

@contextmanager
 def forward_context(first_frame_arg=False,
@@ -1131,7 +1129,6 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE):
            self._conv_idx = 0
            self._feat_map = [None] * self._conv_num
        # cache encode
-        logger.info("self.config.load_encoder: %s", self.config.load_encoder)
        if self.config.load_encoder:
            self._enc_conv_num = _count_conv3d(self.encoder)
            self._enc_conv_idx = 0
@@ -63,6 +63,7 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase):
                           transformer=self.get_module("transformer"),
                           transformer_2=self.get_module("transformer_2", None),
                           scheduler=self.get_module("scheduler"),
+                           vae=self.get_module("vae"),
                           pipeline=self))

        self.add_stage(stage_name="decoding_stage",
@@ -136,11 +136,9 @@ class ComposedPipelineBase(ABC):

            kwargs['model_path'] = model_path
            fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
-            logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
        else:
            assert args is not None, "args must be provided for training mode"
            fastvideo_args = TrainingArgs.from_cli_args(args)
-            logger.info("training args in from_pretrained: %s", fastvideo_args)
            # TODO(will): fix this so that its not so ugly
            fastvideo_args.model_path = model_path
            for key, value in kwargs.items():
@@ -241,5 +241,5 @@ class TrainingBatch:

@dataclass
 class PreprocessBatch(ForwardBatch):
-    video_loader: list["VideoDecoder"] = field(default_factory=list)
-    video_file_name: list[str] = field(default_factory=list)
+    video_loader: list["VideoDecoder"] | list[str] = field(default_factory=list)
+    video_file_name: list[str] = field(default_factory=list)
@@ -0,0 +1,762 @@
+# 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
@@ -4,9 +4,11 @@ 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
@@ -61,7 +63,16 @@ class VideoTransformStage(PipelineStage):
                else:
                    frame_indices = frame_indices[:self.num_frames]

-            video = batch.video_loader[i].get_frames_at(frame_indices).data
+            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 = self.video_transform(video)
            video_pixel_batch.append(video)

@@ -9,8 +9,11 @@ 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__)
@@ -21,12 +24,22 @@ 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"
-    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)
+    
+    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)
+    
    fastvideo_args = FastVideoArgs(
        model_path=args.model_path,
        num_gpus=get_world_size(),
@@ -35,7 +48,25 @@ def main(args) -> None:
        text_encoder_cpu_offload=False,
        pipeline_config=pipeline_config,
    )
-    PreprocessPipeline = PreprocessPipeline_I2V if args.preprocess_task == "i2v" else PreprocessPipeline_T2V
+    
+    if args.preprocess_task == "t2v":
+        PreprocessPipeline = PreprocessPipeline_T2V
+    elif args.preprocess_task == "i2v":
+        PreprocessPipeline = PreprocessPipeline_I2V
+    elif args.preprocess_task == "ode_trajectory":
+        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__}"
+    )
+
    pipeline = PreprocessPipeline(args.model_path, fastvideo_args)
    pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)

@@ -74,7 +105,11 @@ 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")
+    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("--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)
@@ -98,4 +133,4 @@ if __name__ == "__main__":
    )

    args = parser.parse_args()
-    main(args)
+    main(args)
@@ -0,0 +1,216 @@
+# 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
@@ -228,10 +228,7 @@ class CausalDMDDenosingStage(DenoisingStage):
                                                       dim=2)

                    # Prepare inputs
-                    t_expand = t_cur.expand(latent_model_input.shape[0])
-                    # t_expand = t_cur * torch.ones((latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long)
-                    # t_expand = t_expand.repeat(1, self.sliding_window_num_frames)
-
+                    t_expand = t_cur.repeat(latent_model_input.shape[0])

                    # Attention metadata if needed
                    if (vsa_available and self.attn_backend
@@ -265,11 +262,10 @@ class CausalDMDDenosingStage(DenoisingStage):
                                            attn_metadata=attn_metadata,
                                            forward_batch=batch):
                        # Run transformer; follow DMD stage pattern
-                        t_expanded_noise= t_cur * torch.ones((latent_model_input.shape[0], 1), device=latent_model_input.device, dtype=torch.long)
                        pred_noise_btchw = self.transformer(
                            latent_model_input,
                            prompt_embeds,
-                            t_expanded_noise,
+                            t_expand,
                            kv_cache=self.kv_cache1,
                            crossattn_cache=self.crossattn_cache,
                            current_start=(pos_start_base + start_index) *
@@ -330,11 +326,10 @@ class CausalDMDDenosingStage(DenoisingStage):
                    set_forward_context(current_timestep=0,
                                        attn_metadata=attn_metadata,
                                        forward_batch=batch):
-                    t_expanded_context = t_context * torch.ones((context_bcthw.shape[0], 1), device=context_bcthw.device, dtype=torch.long)
                    _ = self.transformer(
                        context_bcthw,
                        prompt_embeds,
-                        t_expanded_context,
+                        t_context,
                        kv_cache=self.kv_cache1,
                        crossattn_cache=self.crossattn_cache,
                        current_start=(pos_start_base + start_index) *
@@ -411,4 +406,4 @@ class CausalDMDDenosingStage(DenoisingStage):
                "is_init":
                False,
            })
-        self.crossattn_cache = crossattn_cache
+        self.crossattn_cache = crossattn_cache
@@ -4,6 +4,7 @@ Denoising stage for diffusion pipelines.
 """

 import inspect
+import math
 import weakref
 from collections.abc import Iterable
 from typing import Any
@@ -24,12 +25,13 @@ from fastvideo.logger import init_logger
 from fastvideo.models.loader.component_loader import TransformerLoader
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
    FlowMatchEulerDiscreteScheduler)
+from fastvideo.models.utils import pred_noise_to_pred_video
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.platforms import AttentionBackendEnum
-from fastvideo.utils import dict_to_3d_list
+from fastvideo.utils import dict_to_3d_list, masks_like

 try:
    from fastvideo.attention.backends.sliding_tile_attn import (
@@ -60,11 +62,13 @@ class DenoisingStage(PipelineStage):
                 transformer,
                 scheduler,
                 pipeline=None,
-                 transformer_2=None) -> None:
+                 transformer_2=None,
+                 vae=None) -> None:
        super().__init__()
        self.transformer = transformer
        self.transformer_2 = transformer_2
        self.scheduler = scheduler
+        self.vae = vae
        self.pipeline = weakref.ref(pipeline) if pipeline else None
        attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
        self.attn_backend = get_attn_backend(
@@ -193,6 +197,44 @@ 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:
@@ -217,19 +259,32 @@ 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] *
@@ -329,6 +384,11 @@ 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
@@ -778,7 +838,6 @@ class DmdDenoisingStage(DenoisingStage):
                            **pos_cond_kwargs,
                        ).permute(0, 2, 1, 3, 4)

-                    from fastvideo.models.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),
@@ -92,24 +92,6 @@ 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

@@ -4,6 +4,8 @@ Input validation stage for diffusion pipelines.
 """

 import torch
+import torchvision.transforms.functional as TF
+from PIL import Image

 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
@@ -12,6 +14,7 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import (StageValidators,
                                                   VerificationResult)
+from fastvideo.utils import best_output_size

 logger = init_logger(__name__)

@@ -94,6 +97,7 @@ class InputValidationStage(PipelineStage):
            )

        # for i2v, get image from image_path
+        # @TODO(Wei) hard-coded for wan2.2 5b ti2v for now. Should put this in image_encoding stage
        if batch.image_path is not None:
            if batch.image_path.endswith(".mp4"):
                image = load_video(batch.image_path)[0]
@@ -101,6 +105,36 @@ class InputValidationStage(PipelineStage):
                image = load_image(batch.image_path)
            batch.pil_image = image

+        # further processing for ti2v task
+        if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+            img = batch.pil_image
+            ih, iw = img.height, img.width
+            patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
+            vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
+            dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride
+            max_area = 704 * 1280
+            ow, oh = best_output_size(iw, ih, dw, dh, max_area)
+
+            scale = max(ow / iw, oh / ih)
+            img = img.resize((round(iw * scale), round(ih * scale)),
+                             Image.LANCZOS)
+            logger.info("resized img height: %s, img width: %s", img.height,
+                        img.width)
+
+            # center-crop
+            x1 = (img.width - ow) // 2
+            y1 = (img.height - oh) // 2
+            img = img.crop((x1, y1, x1 + ow, y1 + oh))
+            assert img.width == ow and img.height == oh
+
+            # to tensor
+            img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(
+                self.device).unsqueeze(1)
+            img = img.unsqueeze(0)
+            batch.height = oh
+            batch.width = ow
+            batch.pil_image = img
+
        return batch

    def verify_input(self, batch: ForwardBatch,
@@ -59,58 +59,37 @@ class TextEncodingStage(PipelineStage):
        assert len(self.text_encoders) == len(
            fastvideo_args.pipeline_config.text_encoder_configs)

-        for tokenizer, text_encoder, encoder_config, preprocess_func, postprocess_func in zip(
-                self.tokenizers,
-                self.text_encoders,
-                fastvideo_args.pipeline_config.text_encoder_configs,
-                fastvideo_args.pipeline_config.preprocess_text_funcs,
-                fastvideo_args.pipeline_config.postprocess_text_funcs,
-                strict=True):
+        # Encode positive prompt with all available encoders
+        assert batch.prompt is not None
+        prompt_text: str | list[str] = batch.prompt
+        all_indices: list[int] = list(range(len(self.text_encoders)))
+        prompt_embeds_list, prompt_masks_list = self.encode_text(
+            prompt_text,
+            fastvideo_args,
+            encoder_index=all_indices,
+            return_attention_mask=True,
+        )
+        for pe in prompt_embeds_list:
+            batch.prompt_embeds.append(pe)
+        if batch.prompt_attention_mask is not None:
+            for am in prompt_masks_list:
+                batch.prompt_attention_mask.append(am)

-            assert isinstance(batch.prompt, str | list)
-            if isinstance(batch.prompt, str):
-                batch.prompt = [batch.prompt]
-            texts = []
-            for prompt_str in batch.prompt:
-                texts.append(preprocess_func(prompt_str))
-            text_inputs = tokenizer(texts,
-                                    **encoder_config.tokenizer_kwargs).to(
-                                        get_local_torch_device())
-            input_ids = text_inputs["input_ids"]
-            attention_mask = text_inputs["attention_mask"]
-            with set_forward_context(current_timestep=0, attn_metadata=None):
-                outputs = text_encoder(
-                    input_ids=input_ids,
-                    attention_mask=attention_mask,
-                    output_hidden_states=True,
-                )
-            prompt_embeds = postprocess_func(outputs)
-            batch.prompt_embeds.append(prompt_embeds)
-            if batch.prompt_attention_mask is not None:
-                batch.prompt_attention_mask.append(attention_mask)
-
-            if batch.do_classifier_free_guidance:
-                assert isinstance(batch.negative_prompt, str)
-                negative_text = preprocess_func(batch.negative_prompt)
-                negative_text_inputs = tokenizer(
-                    negative_text, **encoder_config.tokenizer_kwargs).to(
-                        get_local_torch_device())
-                negative_input_ids = negative_text_inputs["input_ids"]
-                negative_attention_mask = negative_text_inputs["attention_mask"]
-                with set_forward_context(current_timestep=0,
-                                         attn_metadata=None):
-                    negative_outputs = text_encoder(
-                        input_ids=negative_input_ids,
-                        attention_mask=negative_attention_mask,
-                        output_hidden_states=True,
-                    )
-                negative_prompt_embeds = postprocess_func(negative_outputs)
-
-                assert batch.negative_prompt_embeds is not None
-                batch.negative_prompt_embeds.append(negative_prompt_embeds)
-                if batch.negative_attention_mask is not None:
-                    batch.negative_attention_mask.append(
-                        negative_attention_mask)
+        # Encode negative prompt if CFG is enabled
+        if batch.do_classifier_free_guidance:
+            assert isinstance(batch.negative_prompt, str)
+            neg_embeds_list, neg_masks_list = self.encode_text(
+                batch.negative_prompt,
+                fastvideo_args,
+                encoder_index=all_indices,
+                return_attention_mask=True,
+            )
+            assert batch.negative_prompt_embeds is not None
+            for ne in neg_embeds_list:
+                batch.negative_prompt_embeds.append(ne)
+            if batch.negative_attention_mask is not None:
+                for nm in neg_masks_list:
+                    batch.negative_attention_mask.append(nm)

        return batch

@@ -129,6 +108,171 @@ class TextEncodingStage(PipelineStage):
                         V.none_or_list)
        return result

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

-all reference videos are were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
+reference videos were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
+causal videos were generated on commit b318063c0a4618f1d5d99ea82ca67a06aad0d19d

 ## Generation Details

@@ -76,4 +77,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"
@@ -0,0 +1,152 @@
+# 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}"
@@ -101,7 +101,7 @@ I2V_IMAGE_PATHS = [


@pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS)
-@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"])
+@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
@pytest.mark.parametrize("model_id", list(I2V_MODEL_TO_PARAMS.keys()))
 def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
    """
@@ -0,0 +1,128 @@
+import torch
+import types
+import pytest
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.configs.models.encoders.base import TextEncoderArchConfig, TextEncoderConfig, BaseEncoderOutput
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+class TensorDict(dict):
+    def to(self, device):
+        return TensorDict({k: v.to(device) for k, v in self.items()})
+
+class FakeTokenizer:
+    def __call__(self, texts, **kwargs):
+        B = len(texts)
+        seq_len = int(kwargs.get("max_length", 4))
+        return TensorDict({
+            "input_ids": torch.arange(B * seq_len).view(B, seq_len),
+            "attention_mask": torch.ones(B, seq_len, dtype=torch.long),
+        })
+
+class FakeTextEncoder(torch.nn.Module):
+    def __init__(self, hidden_size=8):
+        super().__init__()
+        self.hidden_size = hidden_size
+
+    def forward(self, input_ids, attention_mask, output_hidden_states=True):
+        B, T = input_ids.shape
+        last_hidden_state = torch.arange(B * T * self.hidden_size, dtype=torch.float32).view(B, T, self.hidden_size)
+        return types.SimpleNamespace(last_hidden_state=last_hidden_state)
+
+def id_preprocess(x: str) -> str:
+    return x
+
+def take_mean_postprocess(outputs: BaseEncoderOutput) -> torch.Tensor:
+    # [B, T, H] -> [B, H]
+    return outputs.last_hidden_state.mean(dim=1)
+
+def make_args(num_encoders=2, text_len=4, hidden_size=8):
+    enc_cfgs = []
+    preprocess_fns = []
+    postprocess_fns = []
+    for _ in range(num_encoders):
+        arch = TextEncoderArchConfig(text_len=text_len)
+        enc_cfgs.append(TextEncoderConfig(arch_config=arch))
+        preprocess_fns.append(id_preprocess)
+        postprocess_fns.append(take_mean_postprocess)
+    pipe_cfg = PipelineConfig(
+        text_encoder_configs=tuple(enc_cfgs),
+        text_encoder_precisions=tuple(["fp32"] * num_encoders),
+        preprocess_text_funcs=tuple(preprocess_fns),
+        postprocess_text_funcs=tuple(postprocess_fns),
+    )
+    return FastVideoArgs(model_path="", pipeline_config=pipe_cfg), hidden_size
+
+def make_stage(num_encoders=2, hidden_size=8):
+    tokenizers = [FakeTokenizer() for _ in range(num_encoders)]
+    encoders = [FakeTextEncoder(hidden_size=hidden_size) for _ in range(num_encoders)]
+    return TextEncodingStage(text_encoders=encoders, tokenizers=tokenizers)
+
+def test_encode_text_selection_and_shapes():
+    fastvideo_args, hidden = make_args(num_encoders=2, text_len=4, hidden_size=8)
+    stage = make_stage(num_encoders=2, hidden_size=hidden)
+
+    # list return, two encoders
+    embeds = stage.encode_text(["a", "b"], fastvideo_args, encoder_index=[0, 1])
+    assert isinstance(embeds, list) and len(embeds) == 2
+    for e in embeds:
+        assert e.shape == (2, hidden)
+
+    # with masks
+    embeds2, masks2 = stage.encode_text("a", fastvideo_args, encoder_index=[1], return_attention_mask=True)
+    assert len(embeds2) == 1 and len(masks2) == 1
+    assert embeds2[0].shape == (1, hidden)
+    assert masks2[0].shape == (1, 4)
+
+    # dict return
+    d = stage.encode_text(["a","b"], fastvideo_args, encoder_index=[0,1], return_type="dict")
+    assert set(d.keys()) == {"0", "1"}
+    assert d["0"].shape == (2, hidden)
+
+    # stack return
+    s = stage.encode_text(["a","b"], fastvideo_args, encoder_index=[0,1], return_type="stack")
+    assert s.shape == (2, 2, hidden)  # [encoders, batch, hidden]
+
+    # overrides: dtype + max_length
+    e3, m3 = stage.encode_text(["a"], fastvideo_args, encoder_index=[0], dtype=torch.float16, return_attention_mask=True, max_length=3)
+    assert e3[0].dtype == torch.float16
+    assert m3[0].shape[1] == 3
+
+def test_forward_integration_cfg_off_and_on():
+    fastvideo_args, hidden = make_args(num_encoders=2, text_len=4, hidden_size=8)
+    stage = make_stage(num_encoders=2, hidden_size=hidden)
+
+    # CFG off
+    batch = ForwardBatch(
+        data_type="video",
+        prompt="a cat",
+        negative_prompt="",
+        do_classifier_free_guidance=False,
+        prompt_embeds=[],
+        negative_prompt_embeds=None,
+        prompt_attention_mask=[],
+        negative_attention_mask=None,
+    )
+    out = stage.forward(batch, fastvideo_args)
+    assert len(out.prompt_embeds) == 2
+    for e in out.prompt_embeds:
+        assert e.shape[1] == hidden
+
+    # CFG on
+    batch2 = ForwardBatch(
+        data_type="video",
+        prompt=["a cat", "a dog"],
+        negative_prompt="bad picture",
+        do_classifier_free_guidance=True,
+        prompt_embeds=[],
+        negative_prompt_embeds=[],
+        prompt_attention_mask=[],
+        negative_attention_mask=[],
+    )
+    out2 = stage.forward(batch2, fastvideo_args)
+    assert len(out2.prompt_embeds) == 2
+    assert len(out2.negative_prompt_embeds) == 2
+    assert len(out2.prompt_attention_mask) == 2
+    assert len(out2.negative_attention_mask) == 2
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 import copy
 import gc
-import json
 import os
 import time
 from abc import abstractmethod
@@ -37,11 +36,9 @@ from fastvideo.training.activation_checkpoint import (
    apply_activation_checkpointing)
 from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.training_utils import (
-    EMA_FSDP, clip_grad_norm_while_handling_failing_dtensor_cases,
-    get_scheduler, load_distillation_checkpoint, save_distillation_checkpoint,
-    shift_timestep)
-from fastvideo.utils import (is_vsa_available, maybe_download_model,
-                             set_random_seed, verify_model_config_and_directory)
+    clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler,
+    load_distillation_checkpoint, save_distillation_checkpoint, shift_timestep)
+from fastvideo.utils import is_vsa_available, set_random_seed

 import wandb  # isort: skip

@@ -90,27 +87,9 @@ class DistillationPipeline(TrainingPipeline):
        self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
            shift=self.timestep_shift)

-        if training_args.real_score_model_path:
-            logger.info(
-                f"Loading real score transformer from: {training_args.real_score_model_path}"
-            )
-            self.real_score_transformer = self.load_module_from_path(
-                training_args.real_score_model_path, "transformer",
-                training_args)
-        else:
-            self.real_score_transformer = self.get_module(
-                "real_score_transformer")
-
-        if training_args.fake_score_model_path:
-            logger.info(
-                f"Loading fake score transformer from: {training_args.fake_score_model_path}"
-            )
-            self.fake_score_transformer = self.load_module_from_path(
-                training_args.fake_score_model_path, "transformer",
-                training_args)
-        else:
-            self.fake_score_transformer = self.get_module(
-                "fake_score_transformer")
+        # self.transformer is the generator model
+        self.real_score_transformer = self.get_module("real_score_transformer")
+        self.fake_score_transformer = self.get_module("fake_score_transformer")

        self.real_score_transformer.requires_grad_(False)
        self.real_score_transformer.eval()
@@ -137,13 +116,10 @@ class DistillationPipeline(TrainingPipeline):
        if fake_score_lr == 0.0:
            fake_score_lr = training_args.learning_rate

-        betas_str = training_args.fake_score_betas
-        betas = tuple(float(x.strip()) for x in betas_str.split(","))
-
        self.fake_score_optimizer = torch.optim.AdamW(
            fake_score_params,
            lr=fake_score_lr,
-            betas=betas,
+            betas=(0.9, 0.999),
            weight_decay=training_args.weight_decay,
            eps=1e-8,
        )
@@ -171,19 +147,8 @@ class DistillationPipeline(TrainingPipeline):
            self.training_args.pipeline_config.dmd_denoising_steps,
            dtype=torch.long,
            device=get_local_torch_device())
-
-        if training_args.warp_denoising_step:  # Warp the denoising step according to the scheduler time shift
-            timesteps = torch.cat((self.noise_scheduler.timesteps.cpu(),
-                                   torch.tensor([0],
-                                                dtype=torch.float32))).cuda()
-            self.denoising_step_list = timesteps[1000 -
-                                                 self.denoising_step_list]
-            logger.info("Warping denoising_step_list")
-
-        self.denoising_step_list = self.denoising_step_list.to(
-            get_local_torch_device())
-        logger.info("Distillation generator model to %s denoising steps: %s",
-                    len(self.denoising_step_list), self.denoising_step_list)
+        logger.info("Distillation generator model to %s denoising steps",
+                    len(self.denoising_step_list))
        self.num_train_timestep = self.noise_scheduler.num_train_timesteps

        self.min_timestep = int(self.training_args.min_timestep_ratio *
@@ -193,82 +158,6 @@ class DistillationPipeline(TrainingPipeline):

        self.real_score_guidance_scale = self.training_args.real_score_guidance_scale

-        self.generator_ema = None
-        if (self.training_args.ema_decay
-                is not None) and (self.training_args.ema_decay > 0.0):
-            self.generator_ema = EMA_FSDP(self.transformer,
-                                          decay=self.training_args.ema_decay)
-            logger.info(
-                f"Initialized generator EMA with decay={self.training_args.ema_decay}"
-            )
-        else:
-            logger.info("Generator EMA disabled (ema_decay <= 0.0)")
-
-    def load_module_from_path(self, model_path: str, module_type: str,
-                              training_args: "TrainingArgs"):
-        """
-        Load a module from a specific path using the same loading logic as the pipeline.
-        
-        Args:
-            model_path: Path to the model
-            module_type: Type of module to load (e.g., "transformer")
-            training_args: Training arguments
-            
-        Returns:
-            The loaded module
-        """
-        logger.info(f"Loading {module_type} from custom path: {model_path}")
-        # Set flag to prevent custom weight loading for teacher/critic models
-        training_args._loading_teacher_critic_model = True
-
-        try:
-            from fastvideo.models.loader.component_loader import (
-                PipelineComponentLoader)
-
-            # Download the model if it's a Hugging Face model ID
-            local_model_path = maybe_download_model(model_path)
-            logger.info(f"Model downloaded/found at: {local_model_path}")
-            config = verify_model_config_and_directory(local_model_path)
-
-            if module_type not in config:
-                if hasattr(self, '_extra_config_module_map'
-                           ) and module_type in self._extra_config_module_map:
-                    extra_module = self._extra_config_module_map[module_type]
-                    if extra_module in config:
-                        module_type = extra_module
-                        logger.info(f"Using {extra_module} for {module_type}")
-                    else:
-                        raise ValueError(
-                            f"Module {module_type} not found in config at {local_model_path}"
-                        )
-                else:
-                    raise ValueError(
-                        f"Module {module_type} not found in config at {local_model_path}"
-                    )
-
-            module_info = config[module_type]
-            if module_info is None:
-                raise ValueError(
-                    f"Module {module_type} has null value in config at {local_model_path}"
-                )
-
-            transformers_or_diffusers, architecture = module_info
-            component_path = os.path.join(local_model_path, module_type)
-            module = PipelineComponentLoader.load_module(
-                module_name=module_type,
-                component_model_path=component_path,
-                transformers_or_diffusers=transformers_or_diffusers,
-                fastvideo_args=training_args,
-            )
-
-            logger.info(
-                f"Successfully loaded {module_type} from {component_path}")
-            return module
-        finally:
-            # Always clean up the flag
-            if hasattr(training_args, '_loading_teacher_critic_model'):
-                delattr(training_args, '_loading_teacher_critic_model')
-
    @abstractmethod
    def initialize_validation_pipeline(self, training_args: TrainingArgs):
        """Initialize validation pipeline - must be implemented by subclasses."""
@@ -285,110 +174,6 @@ class DistillationPipeline(TrainingPipeline):

        return training_batch

-    def apply_ema_to_model(self, model):
-        """Apply EMA weights to the model for validation or inference."""
-        if self.generator_ema is not None:
-            with self.generator_ema.apply_to_model(model):
-                return model
-        return model
-
-    def get_ema_model_copy(self):
-        """Get a copy of the model with EMA weights applied."""
-        if self.generator_ema is not None:
-            ema_model = copy.deepcopy(self.transformer)
-            self.generator_ema.copy_to_unwrapped(ema_model)
-            return ema_model
-        return None
-
-    def is_ema_ready(self, current_step: int = None):
-        """Check if EMA is ready for use (after ema_start_step)."""
-        if current_step is None:
-            current_step = getattr(self, 'current_trainstep', 0)
-        return (self.generator_ema is not None
-                and current_step >= self.training_args.ema_start_step)
-
-    def save_ema_weights(self, output_dir: str, step: int):
-        """Save EMA weights separately for inference purposes."""
-        if self.generator_ema is None:
-            logger.warning("Cannot save EMA weights: EMA not initialized")
-            return
-
-        if not self.is_ema_ready():
-            logger.warning(
-                "Cannot save EMA weights: EMA not ready yet (step < ema_start_step)"
-            )
-            return
-
-        try:
-            ema_model = self.get_ema_model_copy()
-            if ema_model is None:
-                logger.warning("Failed to create EMA model copy")
-                return
-
-            ema_save_dir = os.path.join(output_dir, f"ema_checkpoint-{step}")
-            os.makedirs(ema_save_dir, exist_ok=True)
-
-            # save as diffusers format
-            from safetensors.torch import save_file
-
-            from fastvideo.training.training_utils import (
-                custom_to_hf_state_dict, gather_state_dict_on_cpu_rank0)
-            cpu_state = gather_state_dict_on_cpu_rank0(ema_model, device=None)
-
-            if self.global_rank == 0:
-                weight_path = os.path.join(
-                    ema_save_dir, "diffusion_pytorch_model.safetensors")
-                diffusers_state_dict = custom_to_hf_state_dict(
-                    cpu_state, ema_model.reverse_param_names_mapping)
-                save_file(diffusers_state_dict, weight_path)
-
-                config_dict = ema_model.hf_config
-                if "dtype" in config_dict:
-                    del config_dict["dtype"]
-                config_path = os.path.join(ema_save_dir, "config.json")
-                with open(config_path, "w") as f:
-                    json.dump(config_dict, f, indent=4)
-
-                logger.info(f"EMA weights saved to {weight_path}")
-
-            del ema_model
-
-        except Exception as e:
-            logger.error(f"Failed to save EMA weights: {str(e)}")
-
-    def get_ema_stats(self):
-        """Get EMA statistics for monitoring."""
-        if self.generator_ema is None:
-            return {
-                "ema_enabled": False,
-                "ema_decay": None,
-                "ema_start_step": self.training_args.ema_start_step,
-                "ema_ready": False,
-                "ema_step": self.current_trainstep,
-            }
-
-        return {
-            "ema_enabled": True,
-            "ema_decay": self.training_args.ema_decay,
-            "ema_start_step": self.training_args.ema_start_step,
-            "ema_ready": self.is_ema_ready(),
-            "ema_step": self.current_trainstep,
-        }
-
-    def reset_ema(self):
-        """Reset EMA to current model weights."""
-        if self.generator_ema is not None:
-            logger.info("Resetting EMA to current model weights")
-            self.generator_ema.update(self.transformer)
-            # Force update to current weights by setting decay to 0 temporarily
-            original_decay = self.generator_ema.decay
-            self.generator_ema.decay = 0.0
-            self.generator_ema.update(self.transformer)
-            self.generator_ema.decay = original_decay
-            logger.info("EMA reset completed")
-        else:
-            logger.warning("Cannot reset EMA: EMA not initialized")
-
    def _build_distill_input_kwargs(
            self, noise_input: torch.Tensor, timestep: torch.Tensor,
            text_dict: dict[str, torch.Tensor] | None,
@@ -802,10 +587,6 @@ class DistillationPipeline(TrainingPipeline):
            self._clip_model_grad_norm_(batch_gen, self.transformer)
            self.optimizer.step()
            self.optimizer.zero_grad(set_to_none=True)
-
-            if self.generator_ema is not None:
-                self.generator_ema.update(self.transformer)
-
            avg_dmd_loss = torch.tensor(total_dmd_loss /
                                        gradient_accumulation_steps,
                                        device=self.device)
@@ -856,8 +637,7 @@ class DistillationPipeline(TrainingPipeline):
            self.transformer, self.fake_score_transformer, self.global_rank,
            self.training_args.resume_from_checkpoint, self.optimizer,
            self.fake_score_optimizer, self.train_dataloader, self.lr_scheduler,
-            self.fake_score_lr_scheduler, self.noise_random_generator,
-            self.generator_ema)
+            self.fake_score_lr_scheduler, self.noise_random_generator)

        if resumed_step > 0:
            self.init_steps = resumed_step
@@ -888,14 +668,6 @@ class DistillationPipeline(TrainingPipeline):
            sum(p.numel()
                for p in self.fake_score_transformer.parameters()) / 1e9)

-        if self.generator_ema is not None:
-            logger.info("  Generator EMA enabled with decay: %s",
-                        self.training_args.ema_decay)
-            logger.info("  Generator EMA start step: %s",
-                        self.training_args.ema_start_step)
-        else:
-            logger.info("  Generator EMA disabled")
-
    @torch.no_grad()
    def _log_validation(self, transformer, training_args, global_step) -> None:
        training_args.inference_mode = True
@@ -927,18 +699,6 @@ class DistillationPipeline(TrainingPipeline):

        transformer.eval()

-        # Optionally use EMA model for validation if available and ready
-        use_ema_for_validation = (self.training_args.use_ema
-                                  and self.is_ema_ready(global_step))
-        if use_ema_for_validation:
-            logger.info("Using EMA model for validation")
-            validation_transformer = self.transformer
-            ema_context = self.generator_ema.apply_to_model(
-                validation_transformer)
-        else:
-            validation_transformer = transformer
-            ema_context = None
-
        validation_steps = training_args.validation_sampling_steps.split(",")
        validation_steps = [int(step) for step in validation_steps]
        validation_steps = [step for step in validation_steps if step > 0]
@@ -954,98 +714,50 @@ class DistillationPipeline(TrainingPipeline):
            step_videos: list[np.ndarray] = []
            step_captions: list[str] = []

-            if ema_context is not None:
-                with ema_context:
-                    for validation_batch in validation_dataloader:
-                        batch = self._prepare_validation_batch(
-                            sampling_param, training_args, validation_batch,
-                            num_inference_steps)
+            for validation_batch in validation_dataloader:
+                batch = self._prepare_validation_batch(sampling_param,
+                                                       training_args,
+                                                       validation_batch,
+                                                       num_inference_steps)

-                        negative_prompt = batch.negative_prompt
-                        batch_negative = ForwardBatch(
-                            data_type="video",
-                            prompt=negative_prompt,
-                            prompt_embeds=[],
-                            prompt_attention_mask=[],
-                        )
-                        result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
-                            batch_negative, training_args)
-                        self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
-                            0], result_batch.prompt_attention_mask[0]
+                negative_prompt = batch.negative_prompt
+                batch_negative = ForwardBatch(
+                    data_type="video",
+                    prompt=negative_prompt,
+                    prompt_embeds=[],
+                    prompt_attention_mask=[],
+                )
+                result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
+                    batch_negative, training_args)
+                self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
+                    0], result_batch.prompt_attention_mask[0]

-                        logger.info(
-                            "rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
+                logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
                            self.global_rank,
                            self.rank_in_sp_group,
                            batch.prompt,
                            local_main_process_only=False)

-                        assert batch.prompt is not None and isinstance(
-                            batch.prompt, str)
-                        step_captions.append(batch.prompt)
+                assert batch.prompt is not None and isinstance(
+                    batch.prompt, str)
+                step_captions.append(batch.prompt)

-                        # Run validation inference
-                        with torch.no_grad():
-                            output_batch = self.validation_pipeline.forward(
-                                batch, training_args)
-                        samples = output_batch.output
-                        if self.rank_in_sp_group != 0:
-                            continue
+                # Run validation inference
+                with torch.no_grad():
+                    output_batch = self.validation_pipeline.forward(
+                        batch, training_args)
+                samples = output_batch.output
+                if self.rank_in_sp_group != 0:
+                    continue

-                        # Process outputs
-                        video = rearrange(samples, "b c t h w -> t b c h w")
-                        frames = []
-                        for x in video:
-                            x = torchvision.utils.make_grid(x, nrow=6)
-                            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-                            frames.append((x * 255).numpy().astype(np.uint8))
-                        step_videos.append(frames)
-            else:
-                # Use original transformer without EMA
-                for validation_batch in validation_dataloader:
-                    batch = self._prepare_validation_batch(
-                        sampling_param, training_args, validation_batch,
-                        num_inference_steps)
-
-                    negative_prompt = batch.negative_prompt
-                    batch_negative = ForwardBatch(
-                        data_type="video",
-                        prompt=negative_prompt,
-                        prompt_embeds=[],
-                        prompt_attention_mask=[],
-                    )
-                    result_batch = self.validation_pipeline.prompt_encoding_stage(  # type: ignore
-                        batch_negative, training_args)
-                    self.negative_prompt_embeds, self.negative_prompt_attention_mask = result_batch.prompt_embeds[
-                        0], result_batch.prompt_attention_mask[0]
-
-                    logger.info(
-                        "rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
-                        self.global_rank,
-                        self.rank_in_sp_group,
-                        batch.prompt,
-                        local_main_process_only=False)
-
-                    assert batch.prompt is not None and isinstance(
-                        batch.prompt, str)
-                    step_captions.append(batch.prompt)
-
-                    # Run validation inference
-                    with torch.no_grad():
-                        output_batch = self.validation_pipeline.forward(
-                            batch, training_args)
-                    samples = output_batch.output
-                    if self.rank_in_sp_group != 0:
-                        continue
-
-                    # Process outputs
-                    video = rearrange(samples, "b c t h w -> t b c h w")
-                    frames = []
-                    for x in video:
-                        x = torchvision.utils.make_grid(x, nrow=6)
-                        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-                        frames.append((x * 255).numpy().astype(np.uint8))
-                    step_videos.append(frames)
+                # Process outputs
+                video = rearrange(samples, "b c t h w -> t b c h w")
+                frames = []
+                for x in video:
+                    x = torchvision.utils.make_grid(x, nrow=6)
+                    x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+                    frames.append((x * 255).numpy().astype(np.uint8))
+                step_videos.append(frames)

            # Log validation results for this step
            world_group = get_world_group()
@@ -1122,16 +834,16 @@ class DistillationPipeline(TrainingPipeline):
                                                        latents.dtype)
                else:
                    latents += self.vae.shift_factor
-                with torch.autocast("cuda", dtype=torch.bfloat16):
-                    video = self.vae.decode(latents)
-                video = (video / 2 + 0.5).clamp(0, 1)
-                video = video.cpu().float()
-                video = video.permute(0, 2, 1, 3, 4)
-                video = (video * 255).numpy().astype(np.uint8)
-                wandb_loss_dict[latent_key] = wandb.Video(
-                    video, fps=24, format="mp4")  # change to 16 for Wan2.1
-                # Clean up references
-                del video, latents
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                video = self.vae.decode(latents)
+            video = (video / 2 + 0.5).clamp(0, 1)
+            video = video.cpu().float()
+            video = video.permute(0, 2, 1, 3, 4)
+            video = (video * 255).numpy().astype(np.uint8)
+            wandb_loss_dict[latent_key] = wandb.Video(
+                video, fps=24, format="mp4")  # change to 16 for Wan2.1
+            # Clean up references
+            del video, latents

        # Process DMD training data if available - use decode_stage instead of self.vae.decode
        if 'generator_pred_video' in dmd_latents_vis_dict:
@@ -1192,10 +904,6 @@ class DistillationPipeline(TrainingPipeline):
            device="cpu").manual_seed(self.seed)
        logger.info("Initialized random seeds with seed: %s", seed)

-        # Initialize current_trainstep for EMA ready checks
-        #TODO: check if needed
-        self.current_trainstep = self.init_steps
-
        # Resume from checkpoint if specified (this will restore random states)
        if self.training_args.resume_from_checkpoint:
            self._resume_from_checkpoint()
@@ -1239,14 +947,6 @@ class DistillationPipeline(TrainingPipeline):
            self.current_trainstep = step
            training_batch.current_vsa_sparsity = current_vsa_sparsity

-            if (step >= self.training_args.ema_start_step) and \
-                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
-                self.generator_ema = EMA_FSDP(
-                    self.transformer, decay=self.training_args.ema_decay)
-                logger.info(
-                    f"Created generator EMA at step {step} with decay={self.training_args.ema_decay}"
-                )
-
            with torch.autocast("cuda", dtype=torch.bfloat16):
                training_batch = self.train_one_step(training_batch)

@@ -1260,19 +960,11 @@ class DistillationPipeline(TrainingPipeline):
            avg_step_time = sum(step_times) / len(step_times)

            progress_bar.set_postfix({
-                "total_loss":
-                f"{total_loss:.4f}",
-                "generator_loss":
-                f"{generator_loss:.4f}",
-                "fake_score_loss":
-                f"{fake_score_loss:.4f}",
-                "step_time":
-                f"{step_time:.2f}s",
-                "grad_norm":
-                grad_norm,
-                "ema":
-                "✓" if (self.generator_ema is not None and self.is_ema_ready())
-                else "✗",
+                "total_loss": f"{total_loss:.4f}",
+                "generator_loss": f"{generator_loss:.4f}",
+                "fake_score_loss": f"{fake_score_loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
            })
            progress_bar.update(1)

@@ -1300,15 +992,6 @@ class DistillationPipeline(TrainingPipeline):
                if use_vsa:
                    log_data["VSA_train_sparsity"] = current_vsa_sparsity

-                if self.generator_ema is not None:
-                    log_data["ema_enabled"] = True
-                    log_data["ema_decay"] = self.training_args.ema_decay
-                else:
-                    log_data["ema_enabled"] = False
-
-                ema_stats = self.get_ema_stats()
-                log_data.update(ema_stats)
-
                if training_batch.dmd_latent_vis_dict:
                    dmd_additional_logs = {
                        "generator_timestep":
@@ -1340,8 +1023,7 @@ class DistillationPipeline(TrainingPipeline):
                    self.global_rank, self.training_args.output_dir, step,
                    self.optimizer, self.fake_score_optimizer,
                    self.train_dataloader, self.lr_scheduler,
-                    self.fake_score_lr_scheduler, self.noise_random_generator,
-                    self.generator_ema)
+                    self.fake_score_lr_scheduler, self.noise_random_generator)

                if self.transformer:
                    self.transformer.train()
@@ -1358,11 +1040,7 @@ class DistillationPipeline(TrainingPipeline):
                                             self.global_rank,
                                             self.training_args.output_dir,
                                             f"{step}_weight_only",
-                                             only_save_generator_weight=True,
-                                             generator_ema=self.generator_ema)
-
-                if self.training_args.use_ema and self.is_ema_ready():
-                    self.save_ema_weights(self.training_args.output_dir, step)
+                                             only_save_generator_weight=True)

            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
                if self.training_args.log_visualization:
@@ -1382,11 +1060,7 @@ class DistillationPipeline(TrainingPipeline):
            self.training_args.output_dir, self.training_args.max_train_steps,
            self.optimizer, self.fake_score_optimizer, self.train_dataloader,
            self.lr_scheduler, self.fake_score_lr_scheduler,
-            self.noise_random_generator, self.generator_ema)
-
-        if self.training_args.use_ema and self.is_ema_ready():
-            self.save_ema_weights(self.training_args.output_dir,
-                                  self.training_args.max_train_steps)
+            self.noise_random_generator)

        if get_sp_group():
            cleanup_dist_env_and_memory()
@@ -1,996 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import copy
-import gc
-import logging
-from typing import Any
-from collections import deque
-import time
-
-import torch
-import torch.nn.functional as F
-import wandb
-from tqdm.auto import tqdm
-
-from fastvideo.fastvideo_args import TrainingArgs
-from fastvideo.forward_context import set_forward_context
-from fastvideo.logger import init_logger
-from fastvideo.pipelines import TrainingBatch
-from fastvideo.training.distillation_pipeline import DistillationPipeline
-from fastvideo.training.training_utils import (
-    clip_grad_norm_while_handling_failing_dtensor_cases,
-    EMA_FSDP,
-    save_distillation_checkpoint,
-)
-from fastvideo.models.utils import pred_noise_to_pred_video
-from fastvideo.distributed import get_world_group
-import torch.distributed as dist
-import numpy as np
-from fastvideo.utils import set_random_seed, is_vsa_available
-import fastvideo.envs as envs
-from einops import rearrange
-
-logger = init_logger(__name__)
-
-vsa_available = is_vsa_available()
-
-
-class SelfForcingDistillationPipeline(DistillationPipeline):
-    """
-    A self-forcing distillation pipeline that alternates between training
-    the generator and critic based on the self-forcing methodology.
-    
-    This implementation follows the self-forcing approach where:
-    1. Generator and critic are trained in alternating steps
-    2. Generator loss uses DMD-style loss with the critic as fake score
-    3. Critic loss trains the fake score model to distinguish real vs fake
-    """
-
-    def initialize_training_pipeline(self, training_args: TrainingArgs):
-        """Initialize the self-forcing training pipeline."""
-        logger.info("Initializing self-forcing distillation pipeline...")
-
-        super().initialize_training_pipeline(training_args)
-        self.dfake_gen_update_ratio = getattr(training_args, 'dfake_gen_update_ratio', 5)
-
-        self.num_frame_per_block = getattr(training_args, 'num_frame_per_block', 3)
-        self.independent_first_frame = getattr(training_args, 'independent_first_frame', False)
-        self.same_step_across_blocks = getattr(training_args, 'same_step_across_blocks', False)
-        self.last_step_only = getattr(training_args, 'last_step_only', False)
-        self.context_noise = getattr(training_args, 'context_noise', 0)
-        
-        self.frame_seq_length = 1560  # TODO: Calculate this dynamically based on patch size
-        
-        self.kv_cache1 = None
-        self.crossattn_cache = None
-
-        logger.info(f"Self-forcing generator update ratio: {self.dfake_gen_update_ratio}")
-
-    def generate_and_sync_list(self, num_blocks, num_denoising_steps, device):
-        """Generate and synchronize random exit flags across distributed processes."""
-        rank = dist.get_rank() if dist.is_initialized() else 0
-
-        if rank == 0:
-            # Generate random indices
-            indices = torch.randint(
-                low=0,
-                high=num_denoising_steps,
-                size=(num_blocks,),
-                device=device
-            )
-            if self.last_step_only:
-                indices = torch.ones_like(indices) * (num_denoising_steps - 1)
-        else:
-            indices = torch.empty(num_blocks, dtype=torch.long, device=device)
-
-        if dist.is_initialized():
-            dist.broadcast(indices, src=0)  # Broadcast the random indices to all ranks
-        return indices.tolist()
-
-    def generator_loss(self, training_batch: TrainingBatch) -> tuple[torch.Tensor, dict[str, Any]]:
-        """
-        Compute generator loss using DMD-style approach.
-        The generator tries to fool the critic (fake_score_transformer).
-        """
-        with set_forward_context(
-                current_timestep=training_batch.timesteps,
-                attn_metadata=training_batch.attn_metadata_vsa):
-            if self.training_args.simulate_generator_forward:
-                generator_pred_video = self._generator_multi_step_simulation_forward(training_batch)
-            else:
-                generator_pred_video = self._generator_forward(training_batch)
-
-        with set_forward_context(
-                current_timestep=training_batch.timesteps,
-                attn_metadata=training_batch.attn_metadata):
-            dmd_loss = self._dmd_forward(
-                generator_pred_video=generator_pred_video,
-                training_batch=training_batch
-            )
-
-        log_dict = {
-            "dmdtrain_gradient_norm": torch.tensor(0.0, device=self.device)
-        }
-        
-        return dmd_loss, log_dict
-
-    def critic_loss(self, training_batch: TrainingBatch) -> tuple[torch.Tensor, dict[str, Any]]:
-        """
-        Compute critic loss using flow matching between noise and generator output.
-        The critic learns to predict the flow from noise to the generator's output.
-        """
-        updated_batch, flow_matching_loss = self.faker_score_forward(training_batch)
-        training_batch.fake_score_latent_vis_dict = updated_batch.fake_score_latent_vis_dict
-        log_dict = {}
-        
-        return flow_matching_loss, log_dict
-
-    def _generator_forward(self, training_batch: TrainingBatch) -> torch.Tensor:
-        """Forward pass through generator with KV cache support for causal generation."""
-        latents = training_batch.latents
-        dtype = latents.dtype
-        batch_size = latents.shape[0]
-        
-        # Step 1: Sample a timestep from denoising_step_list
-        index = torch.randint(0,
-                              len(self.denoising_step_list), [1],
-                              device=self.device,
-                              dtype=torch.long)
-        timestep = self.denoising_step_list[index]
-        training_batch.dmd_latent_vis_dict["generator_timestep"] = timestep
-
-        # Step 2: Initialize KV cache and cross-attention cache for causal generation
-        kv_cache, crossattn_cache = self._initialize_simulation_caches(batch_size, dtype, self.device)
-        
-        if getattr(self.training_args, 'validate_cache_structure', False):
-            self._validate_cache_structure(kv_cache, crossattn_cache, batch_size)
-
-        # Step 3: Add noise to latents
-        noise = torch.randn(self.video_latent_shape,
-                            device=self.device,
-                            dtype=dtype)
-        if self.sp_world_size > 1:
-            noise = rearrange(noise,
-                              "b (n t) c h w -> b n t c h w",
-                              n=self.sp_world_size).contiguous()
-            noise = noise[:, self.rank_in_sp_group, :, :, :, :]
-        noisy_latent = self.noise_scheduler.add_noise(latents.flatten(0, 1),
-                                                      noise.flatten(0, 1),
-                                                      timestep).unflatten(
-                                                          0,
-                                                          (1, latents.shape[1]))
-
-        # Step 4: Build input kwargs with KV cache support
-        training_batch = self._build_distill_input_kwargs(
-            noisy_latent, timestep, training_batch.conditional_dict,
-            training_batch)
-        
-        # Step 5: Forward pass with KV cache if available
-        if hasattr(self.transformer, '_forward_inference'):
-            # Use causal inference forward with KV cache
-            pred_noise = self.transformer(
-                hidden_states=training_batch.input_kwargs['hidden_states'],
-                encoder_hidden_states=training_batch.input_kwargs['encoder_hidden_states'],
-                timestep=training_batch.input_kwargs['timestep'],
-                encoder_hidden_states_image=training_batch.input_kwargs.get('encoder_hidden_states_image'),
-                kv_cache=kv_cache,
-                crossattn_cache=crossattn_cache,
-                current_start=0,  # Start from beginning for single-step
-                cache_start=0
-            ).permute(0, 2, 1, 3, 4)
-        else:
-            # Fallback to regular forward
-            pred_noise = self.transformer(**training_batch.input_kwargs).permute(
-                0, 2, 1, 3, 4)
-        
-        # Step 6: Convert noise prediction to video prediction
-        pred_video = pred_noise_to_pred_video(
-            pred_noise=pred_noise.flatten(0, 1),
-            noise_input_latent=noisy_latent.flatten(0, 1),
-            timestep=timestep,
-            scheduler=self.noise_scheduler).unflatten(0, pred_noise.shape[:2])
-
-        self._reset_simulation_caches(kv_cache, crossattn_cache)
-
-        return pred_video
-
-    def _generator_multi_step_simulation_forward(
-            self, training_batch: TrainingBatch, return_sim_steps: bool = False) -> torch.Tensor:
-        """Forward pass through student transformer matching inference procedure with KV cache management.
-        
-        This function is adapted from the reference self-forcing implementation's inference_with_trajectory
-        and includes gradient masking logic for dynamic frame generation.
-        """
-        latents = training_batch.latents
-        dtype = latents.dtype
-        batch_size = latents.shape[0]
-        initial_latent = getattr(training_batch, 'image_latent', None)
-        
-        # Dynamic frame generation logic (adapted from _run_generator)
-        num_training_frames = getattr(self.training_args, 'num_frames', 21)
-        
-        # During training, the number of generated frames should be uniformly sampled from
-        # [21, self.num_training_frames], but still being a multiple of self.num_frame_per_block
-        min_num_frames = 20 if self.independent_first_frame else 21
-        max_num_frames = num_training_frames - 1 if self.independent_first_frame else num_training_frames
-        assert max_num_frames % self.num_frame_per_block == 0
-        assert min_num_frames % self.num_frame_per_block == 0
-        max_num_blocks = max_num_frames // self.num_frame_per_block
-        min_num_blocks = min_num_frames // self.num_frame_per_block
-        
-        # Sample number of blocks and sync across processes
-        num_generated_blocks = torch.randint(min_num_blocks, max_num_blocks + 1, (1,), device=self.device)
-        if dist.is_initialized():
-            dist.broadcast(num_generated_blocks, src=0)
-        num_generated_blocks = num_generated_blocks.item()
-        num_generated_frames = num_generated_blocks * self.num_frame_per_block
-        if self.independent_first_frame and initial_latent is None:
-            num_generated_frames += 1
-            min_num_frames += 1
-        
-        # Create noise with dynamic shape
-        if initial_latent is not None:
-            noise_shape = [batch_size, num_generated_frames - 1, *self.video_latent_shape[2:]]
-        else:
-            noise_shape = [batch_size, num_generated_frames, *self.video_latent_shape[2:]]
-        
-        noise = torch.randn(noise_shape, device=self.device, dtype=dtype)
-        if self.sp_world_size > 1:
-            noise = rearrange(noise, "b (n t) c h w -> b n t c h w", n=self.sp_world_size).contiguous()
-            noise = noise[:, self.rank_in_sp_group, :, :, :, :]
-        
-        batch_size, num_frames, num_channels, height, width = noise.shape
-        
-        # Block size calculation
-        if not self.independent_first_frame or (self.independent_first_frame and initial_latent is not None):
-            assert num_frames % self.num_frame_per_block == 0
-            num_blocks = num_frames // self.num_frame_per_block
-        else:
-            assert (num_frames - 1) % self.num_frame_per_block == 0
-            num_blocks = (num_frames - 1) // self.num_frame_per_block
-            
-        num_input_frames = initial_latent.shape[1] if initial_latent is not None else 0
-        num_output_frames = num_frames + num_input_frames
-        output = torch.zeros([batch_size, num_output_frames, num_channels, height, width],
-                           device=noise.device, dtype=noise.dtype)
-
-        # Step 1: Initialize KV cache to all zeros
-        self.kv_cache1, self.crossattn_cache = self._initialize_simulation_caches(batch_size, dtype, self.device)
-        
-        # Validate cache structure (can be disabled in production)
-        if getattr(self.training_args, 'validate_cache_structure', False):
-            self._validate_cache_structure(self.kv_cache1, self.crossattn_cache, batch_size)
-
-        # Step 2: Cache context feature
-        current_start_frame = 0
-        if initial_latent is not None:
-            timestep = torch.ones([batch_size, 1], device=noise.device, dtype=torch.int64) * 0
-            output[:, :1] = initial_latent
-            with torch.no_grad():
-                # Build input kwargs for initial latent
-                training_batch_temp = self._build_distill_input_kwargs(
-                    initial_latent, timestep * 0, training_batch.conditional_dict, training_batch)
-            
-                self.transformer(
-                    hidden_states=training_batch_temp.input_kwargs['hidden_states'],
-                    encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'],
-                    timestep=training_batch_temp.input_kwargs['timestep'],
-                    encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'),
-                    kv_cache=self.kv_cache1,
-                    crossattn_cache=self.crossattn_cache,
-                    current_start=current_start_frame * self.frame_seq_length
-                )
-            current_start_frame += 1
-
-        # Step 3: Temporal denoising loop
-        all_num_frames = [self.num_frame_per_block] * num_blocks
-        if self.independent_first_frame and initial_latent is None:
-            all_num_frames = [1] + all_num_frames
-        num_denoising_steps = len(self.denoising_step_list)
-        exit_flags = self.generate_and_sync_list(len(all_num_frames), num_denoising_steps, device=noise.device)
-        start_gradient_frame_index = num_output_frames - 21
-
-        for block_index, current_num_frames in enumerate(all_num_frames):
-            noisy_input = noise[:, current_start_frame - num_input_frames:current_start_frame + current_num_frames - num_input_frames]
-
-            # Step 3.1: Spatial denoising loop
-            for index, current_timestep in enumerate(self.denoising_step_list):
-                if self.same_step_across_blocks:
-                    exit_flag = (index == exit_flags[0])
-                else:
-                    exit_flag = (index == exit_flags[block_index])
-                    
-                timestep = torch.ones([batch_size, current_num_frames], device=noise.device, dtype=torch.int64) * current_timestep
-                logger.info("timestep shape at initalization: %s", timestep.shape)
-
-                if not exit_flag:
-                    with torch.no_grad():
-                        # Build input kwargs
-                        training_batch_temp = self._build_distill_input_kwargs(
-                            noisy_input, timestep, training_batch.conditional_dict, training_batch)
-                        
-                        logger.info("timestep shape in generator: %s", timestep.shape)
-                        pred_flow = self.transformer(
-                            hidden_states=training_batch_temp.input_kwargs['hidden_states'],
-                            encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'],
-                            timestep=training_batch_temp.input_kwargs['timestep'],
-                            encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'),
-                            kv_cache=self.kv_cache1,
-                            crossattn_cache=self.crossattn_cache,
-                            current_start=current_start_frame * self.frame_seq_length
-                        ).permute(0, 2, 1, 3, 4)
-                        
-                        denoised_pred = pred_noise_to_pred_video(
-                            pred_noise=pred_flow.flatten(0, 1),
-                            noise_input_latent=noisy_input.flatten(0, 1),
-                            timestep=timestep.flatten(),
-                            scheduler=self.noise_scheduler).unflatten(0, pred_flow.shape[:2])
-                        
-                        next_timestep = self.denoising_step_list[index + 1]
-                        logger.info("denoised_pred shape before: %s", denoised_pred.shape)
-                        noisy_input = self.noise_scheduler.add_noise(
-                            denoised_pred.flatten(0, 1),
-                            torch.randn_like(denoised_pred.flatten(0, 1)),
-                            next_timestep * torch.ones([batch_size * current_num_frames], device=noise.device, dtype=torch.long)
-                        ).unflatten(0, denoised_pred.shape[:2])
-                        logger.info("denoised_pred shape after: %s", denoised_pred.shape)
-                else:
-                    # Final prediction with gradient control
-                    if current_start_frame < start_gradient_frame_index:
-                        logger.info("timestep shape in generator: %s", timestep.shape)
-                        with torch.no_grad():
-                            training_batch_temp = self._build_distill_input_kwargs(
-                                noisy_input, timestep, training_batch.conditional_dict, training_batch)
-                            
-                            pred_flow = self.transformer(
-                                hidden_states=training_batch_temp.input_kwargs['hidden_states'],
-                                encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'],
-                                timestep=training_batch_temp.input_kwargs['timestep'],
-                                encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'),
-                                kv_cache=self.kv_cache1,
-                                crossattn_cache=self.crossattn_cache,
-                                current_start=current_start_frame * self.frame_seq_length
-                            ).permute(0, 2, 1, 3, 4)
-                    else:
-                        logger.info("timestep shape in generator: %s", timestep.shape)
-                        training_batch_temp = self._build_distill_input_kwargs(
-                            noisy_input, timestep, training_batch.conditional_dict, training_batch)
-    
-                        pred_flow = self.transformer(
-                            hidden_states=training_batch_temp.input_kwargs['hidden_states'],
-                            encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'],
-                            timestep=training_batch_temp.input_kwargs['timestep'],
-                            encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'),
-                            kv_cache=self.kv_cache1,
-                            crossattn_cache=self.crossattn_cache,
-                            current_start=current_start_frame * self.frame_seq_length
-                        ).permute(0, 2, 1, 3, 4)
-                    
-                    denoised_pred = pred_noise_to_pred_video(
-                        pred_noise=pred_flow.flatten(0, 1),
-                        noise_input_latent=noisy_input.flatten(0, 1),
-                        timestep=timestep.flatten(),
-                        scheduler=self.noise_scheduler).unflatten(0, pred_flow.shape[:2])
-                    break
-
-            # Step 3.2: record the model's output
-            output[:, current_start_frame:current_start_frame + current_num_frames] = denoised_pred
-
-            # Step 3.3: rerun with timestep zero to update the cache
-            context_timestep = torch.ones_like(timestep) * self.context_noise
-            logger.info("denoised_pred shape before: %s", denoised_pred.shape)
-            denoised_pred = self.noise_scheduler.add_noise(
-                denoised_pred.flatten(0, 1),
-                torch.randn_like(denoised_pred.flatten(0, 1)),
-                context_timestep * torch.ones([batch_size * current_num_frames], device=noise.device, dtype=torch.long)
-            ).unflatten(0, denoised_pred.shape[:2])
-            logger.info("denoised_pred shape after: %s", denoised_pred.shape)
-            
-            with torch.no_grad():
-                training_batch_temp = self._build_distill_input_kwargs(
-                    denoised_pred, context_timestep, training_batch.conditional_dict, training_batch)
-                
-                self.transformer(
-                    hidden_states=training_batch_temp.input_kwargs['hidden_states'],
-                    encoder_hidden_states=training_batch_temp.input_kwargs['encoder_hidden_states'],
-                    timestep=training_batch_temp.input_kwargs['timestep'],
-                    encoder_hidden_states_image=training_batch_temp.input_kwargs.get('encoder_hidden_states_image'),
-                    kv_cache=self.kv_cache1,
-                    crossattn_cache=self.crossattn_cache,
-                    current_start=current_start_frame * self.frame_seq_length
-                )
-
-            # Step 3.4: update the start and end frame indices
-            current_start_frame += current_num_frames
-
-        # Handle last 21 frames logic 
-        pred_image_or_video = output
-        if num_input_frames > 0:
-            pred_image_or_video = output[:, num_input_frames:]
-        
-        # Slice last 21 frames if we generated more
-        gradient_mask = None
-        if pred_image_or_video.shape[1] > 21:
-            with torch.no_grad():
-                # Reencode to get image latent
-                latent_to_decode = pred_image_or_video[:, :-20, ...]
-                # Decode to video
-                latent_to_decode = latent_to_decode.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
-                
-                # Apply VAE scaling and shift factors
-                if isinstance(self.vae.scaling_factor, torch.Tensor):
-                    latent_to_decode = latent_to_decode / self.vae.scaling_factor.to(latent_to_decode.device, latent_to_decode.dtype)
-                else:
-                    latent_to_decode = latent_to_decode / self.vae.scaling_factor
-                
-                if hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None:
-                    if isinstance(self.vae.shift_factor, torch.Tensor):
-                        latent_to_decode += self.vae.shift_factor.to(latent_to_decode.device, latent_to_decode.dtype)
-                    else:
-                        latent_to_decode += self.vae.shift_factor
-                
-                # Decode to pixels
-                pixels = self.vae.decode(latent_to_decode)
-                frame = pixels[:, :, -1:, :, :].to(dtype)  # Last frame [B, C, 1, H, W]
-                
-                # Encode frame back to get image latent
-                image_latent = self.vae.encode(frame).mean.to(dtype)
-                image_latent = image_latent.permute(0, 2, 1, 3, 4)  # [B, F, C, H, W]
-            
-            pred_image_or_video_last_21 = torch.cat([image_latent, pred_image_or_video[:, -20:, ...]], dim=1)
-        else:
-            pred_image_or_video_last_21 = pred_image_or_video
-        
-        # Set up gradient mask if we generated more than minimum frames
-        if num_generated_frames != min_num_frames:
-            # Currently, we do not use gradient for the first chunk, since it contains image latents
-            gradient_mask = torch.ones_like(pred_image_or_video_last_21, dtype=torch.bool)
-            if self.independent_first_frame:
-                gradient_mask[:, :1] = False
-            else:
-                gradient_mask[:, :self.num_frame_per_block] = False
-        
-        # Apply gradient masking if needed
-        final_output = pred_image_or_video_last_21.to(dtype)
-        if gradient_mask is not None:
-            # Apply gradient masking: detach frames that shouldn't contribute gradients
-            final_output = torch.where(
-                gradient_mask,
-                pred_image_or_video_last_21,  # Keep original values where gradient_mask is True
-                pred_image_or_video_last_21.detach()  # Detach where gradient_mask is False
-            )
-
-        # Store visualization data
-        training_batch.dmd_latent_vis_dict["generator_timestep"] = torch.tensor(
-            self.denoising_step_list[exit_flags[0]], dtype=torch.float32, device=self.device)
-        
-        # Store gradient mask information for debugging
-        if gradient_mask is not None:
-            training_batch.dmd_latent_vis_dict["gradient_mask"] = gradient_mask.float()
-            training_batch.dmd_latent_vis_dict["num_generated_frames"] = torch.tensor(
-                num_generated_frames, dtype=torch.float32, device=self.device)
-            training_batch.dmd_latent_vis_dict["min_num_frames"] = torch.tensor(
-                min_num_frames, dtype=torch.float32, device=self.device)
-        
-        self._reset_simulation_caches(self.kv_cache1, self.crossattn_cache)
-
-        return final_output if gradient_mask is not None else pred_image_or_video
-
-    def _initialize_simulation_caches(self, batch_size: int, dtype: torch.dtype, device: torch.device):
-        """Initialize KV cache and cross-attention cache for multi-step simulation."""
-        num_transformer_blocks = len(self.transformer.blocks)
-        
-        # Calculate frame sequence length based on input dimensions and patch size
-        # From the training batch, we can get the actual latent dimensions
-        latent_shape = self.video_latent_shape_sp  # This is set in _prepare_dit_inputs
-        batch_size_actual, num_frames, num_channels, height, width = latent_shape
-        
-        # Get patch size from transformer config
-        p_t, p_h, p_w = self.transformer.patch_size
-        post_patch_height = height // p_h
-        post_patch_width = width // p_w
-        
-        # Frame sequence length is the spatial sequence length per frame
-        frame_seq_length = post_patch_height * post_patch_width
-        
-        # Get local attention size from transformer config
-        local_attn_size = getattr(self.transformer, 'local_attn_size', -1)
-        
-        # Get model configuration parameters - handle FSDP wrapping
-        if hasattr(self.transformer, 'config'):
-            config = self.transformer.config
-            num_attention_heads = config.num_attention_heads
-            attention_head_dim = config.attention_head_dim
-            text_len = config.text_len
-        else:
-            # Fallback to direct attribute access for non-FSDP models
-            num_attention_heads = getattr(self.transformer, 'num_attention_heads', 40)
-            attention_head_dim = getattr(self.transformer, 'attention_head_dim', 128)
-            text_len = getattr(self.transformer, 'text_len', 512)
-        
-        num_max_frames = getattr(self.training_args, "num_frames", num_frames)
-        kv_cache_size = num_max_frames * frame_seq_length
-
-        kv_cache = []
-        for _ in range(num_transformer_blocks):
-            kv_cache.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)
-            })
-        
-        # Initialize cross-attention cache
-        crossattn_cache = []
-        for _ in range(num_transformer_blocks):
-            crossattn_cache.append({
-                "k": torch.zeros([batch_size, text_len, num_attention_heads, attention_head_dim], dtype=dtype, device=device),
-                "v": torch.zeros([batch_size, text_len, num_attention_heads, attention_head_dim], dtype=dtype, device=device),
-                "is_init": False
-            })
-        
-        return kv_cache, crossattn_cache
-
-    def _reset_simulation_caches(self, kv_cache, crossattn_cache):
-        """Reset KV cache and cross-attention cache to clean state."""
-        if kv_cache is not None:
-            for cache_dict in kv_cache:
-                cache_dict["global_end_index"].fill_(0)
-                cache_dict["local_end_index"].fill_(0)
-                cache_dict["k"].zero_()
-                cache_dict["v"].zero_()
-        
-        if crossattn_cache is not None:
-            for cache_dict in crossattn_cache:
-                cache_dict["is_init"] = False
-                cache_dict["k"].zero_()
-                cache_dict["v"].zero_()
-
-    def _validate_cache_structure(self, kv_cache, crossattn_cache, batch_size: int):
-        """Validate that cache structures are correctly initialized."""
-        num_transformer_blocks = len(self.transformer.blocks)
-        
-        # Get model configuration parameters - handle FSDP wrapping
-        if hasattr(self.transformer, 'config'):
-            config = self.transformer.config
-            num_attention_heads = config.num_attention_heads
-            attention_head_dim = config.attention_head_dim
-            text_len = config.text_len
-        else:
-            # Fallback to direct attribute access for non-FSDP models
-            num_attention_heads = getattr(self.transformer, 'num_attention_heads', 40)
-            attention_head_dim = getattr(self.transformer, 'attention_head_dim', 128)
-            text_len = getattr(self.transformer, 'text_len', 512)
-        
-        if kv_cache is not None:
-            assert len(kv_cache) == num_transformer_blocks, f"Expected {num_transformer_blocks} transformer blocks, got {len(kv_cache)}"
-            for i, cache_dict in enumerate(kv_cache):
-                assert "k" in cache_dict and "v" in cache_dict, f"Missing k/v in kv_cache block {i}"
-                assert "global_end_index" in cache_dict and "local_end_index" in cache_dict, f"Missing indices in kv_cache block {i}"
-                assert cache_dict["k"].shape[0] == batch_size, f"Batch size mismatch in kv_cache block {i}"
-                assert cache_dict["v"].shape[0] == batch_size, f"Batch size mismatch in kv_cache block {i}"
-                assert cache_dict["k"].shape[2] == num_attention_heads, f"Attention heads mismatch in kv_cache block {i}"
-                assert cache_dict["k"].shape[3] == attention_head_dim, f"Attention head dim mismatch in kv_cache block {i}"
-        
-        if crossattn_cache is not None:
-            assert len(crossattn_cache) == num_transformer_blocks, f"Expected {num_transformer_blocks} transformer blocks, got {len(crossattn_cache)}"
-            for i, cache_dict in enumerate(crossattn_cache):
-                assert "k" in cache_dict and "v" in cache_dict, f"Missing k/v in crossattn_cache block {i}"
-                assert "is_init" in cache_dict, f"Missing is_init in crossattn_cache block {i}"
-                assert cache_dict["k"].shape[0] == batch_size, f"Batch size mismatch in crossattn_cache block {i}"
-                assert cache_dict["v"].shape[0] == batch_size, f"Batch size mismatch in crossattn_cache block {i}"
-                assert cache_dict["k"].shape[1] == text_len, f"Text length mismatch in crossattn_cache block {i}"
-                assert cache_dict["k"].shape[2] == num_attention_heads, f"Attention heads mismatch in crossattn_cache block {i}"
-                assert cache_dict["k"].shape[3] == attention_head_dim, f"Attention head dim mismatch in crossattn_cache block {i}"
-
-    def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
-        """
-        Self-forcing training step that alternates between generator and critic training.
-        """
-        gradient_accumulation_steps = getattr(self.training_args, 'gradient_accumulation_steps', 1)
-        train_generator = (self.current_trainstep % self.dfake_gen_update_ratio == 0)
-        
-        batches = []
-        for _ in range(gradient_accumulation_steps):
-            batch = self._prepare_distillation(training_batch)
-            batch = self._get_next_batch(batch)
-            batch = self._normalize_dit_input(batch)
-            batch = self._prepare_dit_inputs(batch)
-            batch = self._build_attention_metadata(batch)
-            batch.attn_metadata_vsa = copy.deepcopy(batch.attn_metadata)
-            if batch.attn_metadata is not None:
-                batch.attn_metadata.VSA_sparsity = 0.0
-            batches.append(batch)
-
-        training_batch.dmd_latent_vis_dict = {}
-        training_batch.fake_score_latent_vis_dict = {}
-
-        if train_generator:
-            logger.debug(f"Training generator at step {self.current_trainstep}")
-            self.optimizer.zero_grad()
-            total_generator_loss = 0.0
-            generator_log_dict = {}
-            
-            for batch in batches:
-                # Create a new batch with detached tensors
-                batch_gen = TrainingBatch()
-                for key, value in batch.__dict__.items():
-                    if isinstance(value, torch.Tensor):
-                        setattr(batch_gen, key, value.detach().clone())
-                    elif isinstance(value, dict):
-                        setattr(batch_gen, key, {k: v.detach().clone() if isinstance(v, torch.Tensor) else copy.deepcopy(v) for k, v in value.items()})
-                    else:
-                        setattr(batch_gen, key, copy.deepcopy(value))
-                
-                generator_loss, gen_log_dict = self.generator_loss(batch_gen)
-                with set_forward_context(
-                        current_timestep=batch_gen.timesteps,
-                        attn_metadata=batch_gen.attn_metadata):
-                    (generator_loss / gradient_accumulation_steps).backward()
-                total_generator_loss += generator_loss.detach().item()
-                generator_log_dict.update(gen_log_dict)
-                # Store visualization data from generator training
-                if hasattr(batch_gen, 'dmd_latent_vis_dict'):
-                    training_batch.dmd_latent_vis_dict.update(batch_gen.dmd_latent_vis_dict)
-            
-            self._clip_model_grad_norm_(batch_gen, self.transformer)
-            self.optimizer.step()
-            self.lr_scheduler.step()
-            
-            if self.generator_ema is not None:
-                self.generator_ema.update(self.transformer)
-            
-            avg_generator_loss = torch.tensor(
-                total_generator_loss / gradient_accumulation_steps,
-                device=self.device
-            )
-            world_group = get_world_group()
-            world_group.all_reduce(avg_generator_loss, op=torch.distributed.ReduceOp.AVG)
-            training_batch.generator_loss = avg_generator_loss.item()
-        else:
-            training_batch.generator_loss = 0.0
-
-        logger.debug(f"Training critic at step {self.current_trainstep}")
-        self.fake_score_optimizer.zero_grad()
-        total_critic_loss = 0.0
-        critic_log_dict = {}
-        
-        for batch in batches:
-            # Create a new batch with detached tensors
-            batch_critic = TrainingBatch()
-            for key, value in batch.__dict__.items():
-                if isinstance(value, torch.Tensor):
-                    setattr(batch_critic, key, value.detach().clone())
-                elif isinstance(value, dict):
-                    setattr(batch_critic, key, {k: v.detach().clone() if isinstance(v, torch.Tensor) else copy.deepcopy(v) for k, v in value.items()})
-                else:
-                    setattr(batch_critic, key, copy.deepcopy(value))
-            
-            critic_loss, crit_log_dict = self.critic_loss(batch_critic)
-            with set_forward_context(
-                    current_timestep=batch_critic.timesteps,
-                    attn_metadata=batch_critic.attn_metadata):
-                (critic_loss / gradient_accumulation_steps).backward()
-            total_critic_loss += critic_loss.detach().item()
-            critic_log_dict.update(crit_log_dict)
-            # Store visualization data from critic training
-            if hasattr(batch_critic, 'fake_score_latent_vis_dict'):
-                training_batch.fake_score_latent_vis_dict.update(batch_critic.fake_score_latent_vis_dict)
-        
-        self._clip_model_grad_norm_(batch_critic, self.fake_score_transformer)
-        self.fake_score_optimizer.step()
-        self.fake_score_lr_scheduler.step()
-        
-        avg_critic_loss = torch.tensor(
-            total_critic_loss / gradient_accumulation_steps,
-            device=self.device
-        )
-        world_group = get_world_group()
-        world_group.all_reduce(avg_critic_loss, op=torch.distributed.ReduceOp.AVG)
-        training_batch.fake_score_loss = avg_critic_loss.item()
-
-        training_batch.total_loss = training_batch.generator_loss + training_batch.fake_score_loss
-        
-        return training_batch
-
-    def _log_training_info(self) -> None:
-        """Log self-forcing specific training information."""
-        super()._log_training_info()
-        logger.info("Self-forcing specific settings:")
-        logger.info("  Generator update ratio: %s", self.dfake_gen_update_ratio) 
-
-    def visualize_intermediate_latents(self, training_batch: TrainingBatch,
-                                     training_args: TrainingArgs, step: int):
-        """Add visualization data to wandb logging and save frames to disk."""
-        wandb_loss_dict = {}
-        
-        # Debug logging
-        logger.info(f"Step {step}: Starting visualization")
-        if hasattr(training_batch, 'dmd_latent_vis_dict'):
-            logger.info(f"DMD latent keys: {list(training_batch.dmd_latent_vis_dict.keys())}")
-        if hasattr(training_batch, 'fake_score_latent_vis_dict'):
-            logger.info(f"Fake score latent keys: {list(training_batch.fake_score_latent_vis_dict.keys())}")
-        
-        # Process generator predictions if available
-        if hasattr(training_batch, 'dmd_latent_vis_dict') and training_batch.dmd_latent_vis_dict:
-            dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict
-            dmd_log_keys = ['generator_pred_video', 'real_score_pred_video', 'faker_score_pred_video']
-            
-            for latent_key in dmd_log_keys:
-                if latent_key in dmd_latents_vis_dict:
-                    logger.info(f"Processing DMD latent: {latent_key}")
-                    latents = dmd_latents_vis_dict[latent_key]
-                    if not isinstance(latents, torch.Tensor):
-                        logger.warning(f"Expected tensor for {latent_key}, got {type(latents)}")
-                        continue
-                        
-                    latents = latents.detach()
-                    latents = latents.permute(0, 2, 1, 3, 4)
-
-                    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
-
-                    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
-
-                    try:
-                        with torch.autocast("cuda", dtype=torch.bfloat16):
-                            video = self.vae.decode(latents)
-                        video = (video / 2 + 0.5).clamp(0, 1)
-                        video = video.cpu().float()
-                        video = video.permute(0, 2, 1, 3, 4)
-                        video = (video * 255).numpy().astype(np.uint8)
-                        wandb_loss_dict[f"dmd_{latent_key}"] = wandb.Video(video, fps=24, format="mp4")
-                        logger.info(f"Successfully processed DMD latent: {latent_key}")
-                    except Exception as e:
-                        logger.error(f"Error processing DMD latent {latent_key}: {str(e)}")
-                    del video, latents
-
-        # Process critic predictions
-        if hasattr(training_batch, 'fake_score_latent_vis_dict') and training_batch.fake_score_latent_vis_dict:
-            fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict
-            fake_score_log_keys = ['generator_pred_video']
-            
-            for latent_key in fake_score_log_keys:
-                if latent_key in fake_score_latents_vis_dict:
-                    logger.info(f"Processing critic latent: {latent_key}")
-                    latents = fake_score_latents_vis_dict[latent_key]
-                    if not isinstance(latents, torch.Tensor):
-                        logger.warning(f"Expected tensor for {latent_key}, got {type(latents)}")
-                        continue
-                        
-                    latents = latents.detach()
-                    latents = latents.permute(0, 2, 1, 3, 4)
-
-                    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
-
-                    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
-
-                    try:
-                        with torch.autocast("cuda", dtype=torch.bfloat16):
-                            video = self.vae.decode(latents)
-                        video = (video / 2 + 0.5).clamp(0, 1)
-                        video = video.cpu().float()
-                        video = video.permute(0, 2, 1, 3, 4)
-                        video = (video * 255).numpy().astype(np.uint8)
-                        wandb_loss_dict[f"critic_{latent_key}"] = wandb.Video(video, fps=24, format="mp4")
-                        logger.info(f"Successfully processed critic latent: {latent_key}")
-                    except Exception as e:
-                        logger.error(f"Error processing critic latent {latent_key}: {str(e)}")
-                    del video, latents
-
-        # Log metadata
-        if hasattr(training_batch, 'dmd_latent_vis_dict') and training_batch.dmd_latent_vis_dict:
-            if "generator_timestep" in training_batch.dmd_latent_vis_dict:
-                wandb_loss_dict["generator_timestep"] = training_batch.dmd_latent_vis_dict["generator_timestep"].item()
-            if "dmd_timestep" in training_batch.dmd_latent_vis_dict:
-                wandb_loss_dict["dmd_timestep"] = training_batch.dmd_latent_vis_dict["dmd_timestep"].item()
-
-        if hasattr(training_batch, 'fake_score_latent_vis_dict') and training_batch.fake_score_latent_vis_dict:
-            if "fake_score_timestep" in training_batch.fake_score_latent_vis_dict:
-                wandb_loss_dict["fake_score_timestep"] = training_batch.fake_score_latent_vis_dict["fake_score_timestep"].item()
-
-        # Log final dict contents
-        logger.info(f"Final wandb_loss_dict keys: {list(wandb_loss_dict.keys())}")
-
-        if self.global_rank == 0:
-            wandb.log(wandb_loss_dict, step=step)
-
-    def train(self) -> None:
-        """Main training loop with self-forcing specific logging."""
-        assert self.training_args.seed is not None, "seed must be set"
-        seed = self.training_args.seed
-
-        # Set the same seed within each SP group to ensure reproducibility
-        if self.sp_world_size > 1:
-            # Use the same seed for all processes within the same SP group
-            sp_group_seed = seed + (self.global_rank // self.sp_world_size)
-            set_random_seed(sp_group_seed)
-            logger.info("Rank %s: Using SP group seed %s", self.global_rank,
-                        sp_group_seed)
-        else:
-            set_random_seed(seed + self.global_rank)
-
-        self.noise_random_generator = torch.Generator(device="cpu").manual_seed(
-            self.seed)
-        self.noise_gen_cuda = torch.Generator(device="cuda").manual_seed(
-            self.seed)
-        self.validation_random_generator = torch.Generator(
-            device="cpu").manual_seed(self.seed)
-        logger.info("Initialized random seeds with seed: %s", seed)
-        
-        self.current_trainstep = self.init_steps
-
-        if self.training_args.resume_from_checkpoint:
-            self._resume_from_checkpoint()
-            logger.info("Resumed from checkpoint, random states restored")
-        else:
-            logger.info("Starting training from scratch")
-
-        self.train_loader_iter = iter(self.train_dataloader)
-
-        step_times = deque(maxlen=100)
-
-        self._log_training_info()
-        self._log_validation(self.transformer, self.training_args,
-                             self.init_steps)
-
-        progress_bar = tqdm(
-            range(0, self.training_args.max_train_steps),
-            initial=self.init_steps,
-            desc="Steps",
-            disable=self.local_rank > 0,
-        )
-
-        use_vsa = vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN"
-        for step in range(self.init_steps + 1,
-                          self.training_args.max_train_steps + 1):
-            start_time = time.perf_counter()
-            if use_vsa:
-                vsa_sparsity = self.training_args.VSA_sparsity
-                vsa_decay_rate = self.training_args.VSA_decay_rate
-                vsa_decay_interval_steps = self.training_args.VSA_decay_interval_steps
-                if vsa_decay_interval_steps > 1:
-                    current_decay_times = min(step // vsa_decay_interval_steps,
-                                              vsa_sparsity // vsa_decay_rate)
-                    current_vsa_sparsity = current_decay_times * vsa_decay_rate
-                else:
-                    current_vsa_sparsity = vsa_sparsity
-            else:
-                current_vsa_sparsity = 0.0
-
-            training_batch = TrainingBatch()
-            self.current_trainstep = step
-            training_batch.current_vsa_sparsity = current_vsa_sparsity
-
-            if (step >= self.training_args.ema_start_step) and \
-                    (self.generator_ema is None) and (self.training_args.ema_decay > 0):
-                self.generator_ema = EMA_FSDP(self.transformer, decay=self.training_args.ema_decay)
-                logger.info(f"Created generator EMA at step {step} with decay={self.training_args.ema_decay}")
-
-            with torch.autocast("cuda", dtype=torch.bfloat16):
-                training_batch = self.train_one_step(training_batch)
-
-            total_loss = training_batch.total_loss
-            generator_loss = training_batch.generator_loss
-            fake_score_loss = training_batch.fake_score_loss
-            grad_norm = training_batch.grad_norm
-
-            step_time = time.perf_counter() - start_time
-            step_times.append(step_time)
-            avg_step_time = sum(step_times) / len(step_times)
-
-            progress_bar.set_postfix({
-                "total_loss": f"{total_loss:.4f}",
-                "generator_loss": f"{generator_loss:.4f}",
-                "fake_score_loss": f"{fake_score_loss:.4f}",
-                "step_time": f"{step_time:.2f}s",
-                "grad_norm": grad_norm,
-                "ema": "✓" if (self.generator_ema is not None and self.is_ema_ready()) else "✗",
-            })
-            progress_bar.update(1)
-
-            if self.global_rank == 0:
-                log_data = {
-                    "train_total_loss": total_loss,
-                    "train_fake_score_loss": fake_score_loss,
-                    "learning_rate": self.lr_scheduler.get_last_lr()[0],
-                    "fake_score_learning_rate": self.fake_score_lr_scheduler.get_last_lr()[0],
-                    "step_time": step_time,
-                    "avg_step_time": avg_step_time,
-                    "grad_norm": grad_norm,
-                }
-                if (step % self.dfake_gen_update_ratio == 0):
-                    log_data["train_generator_loss"] = generator_loss
-                if use_vsa:
-                    log_data["VSA_train_sparsity"] = current_vsa_sparsity
-
-                if self.generator_ema is not None:
-                    log_data["ema_enabled"] = True
-                    log_data["ema_decay"] = self.training_args.ema_decay
-                else:
-                    log_data["ema_enabled"] = False
-
-                ema_stats = self.get_ema_stats()
-                log_data.update(ema_stats)
-
-                if training_batch.dmd_latent_vis_dict:
-                    dmd_additional_logs = {
-                        "generator_timestep": training_batch.dmd_latent_vis_dict["generator_timestep"].item(),
-                        "dmd_timestep": training_batch.dmd_latent_vis_dict["dmd_timestep"].item(),
-                    }
-                    log_data.update(dmd_additional_logs)
-
-                faker_score_additional_logs = {
-                    "fake_score_timestep": training_batch.fake_score_latent_vis_dict["fake_score_timestep"].item(),
-                }
-                log_data.update(faker_score_additional_logs)
-
-                wandb.log(log_data, step=step)
-
-                if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
-                    if self.training_args.log_visualization:
-                        self.visualize_intermediate_latents(training_batch, self.training_args, step)
-
-            if (self.training_args.training_state_checkpointing_steps > 0
-                    and step % self.training_args.training_state_checkpointing_steps == 0):
-                print("rank", self.global_rank,
-                      "save training state checkpoint at step", step)
-                save_distillation_checkpoint(
-                    self.transformer, self.fake_score_transformer,
-                    self.global_rank, self.training_args.output_dir, step,
-                    self.optimizer, self.fake_score_optimizer,
-                    self.train_dataloader, self.lr_scheduler,
-                    self.fake_score_lr_scheduler, self.noise_random_generator,
-                    self.generator_ema)
-
-                if self.transformer:
-                    self.transformer.train()
-                self.sp_group.barrier()
-
-            if (self.training_args.weight_only_checkpointing_steps > 0
-                    and step % self.training_args.weight_only_checkpointing_steps == 0):
-                print("rank", self.global_rank,
-                      "save weight-only checkpoint at step", step)
-                save_distillation_checkpoint(self.transformer,
-                                             self.fake_score_transformer,
-                                             self.global_rank,
-                                             self.training_args.output_dir,
-                                             f"{step}_weight_only",
-                                             only_save_generator_weight=True,
-                                             generator_ema=self.generator_ema)
-                
-                if self.training_args.use_ema and self.is_ema_ready():
-                    self.save_ema_weights(self.training_args.output_dir, step)
-
-            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
-                self._log_validation(self.transformer, self.training_args, step)
-
-        wandb.finish()
-
-        print("rank", self.global_rank,
-              "save final training state checkpoint at step",
-              self.training_args.max_train_steps)
-        save_distillation_checkpoint(
-            self.transformer, self.fake_score_transformer, self.global_rank,
-            self.training_args.output_dir, self.training_args.max_train_steps,
-            self.optimizer, self.fake_score_optimizer, self.train_dataloader,
-            self.lr_scheduler, self.fake_score_lr_scheduler,
-            self.noise_random_generator, self.generator_ema)
-
-        if self.training_args.use_ema and self.is_ema_ready():
-            self.save_ema_weights(self.training_args.output_dir, self.training_args.max_train_steps)
-
-        if get_sp_group():
-            cleanup_dist_env_and_memory() 
@@ -117,14 +117,10 @@ class TrainingPipeline(LoRAPipeline, ABC):
        params_to_optimize = self.transformer.parameters()
        params_to_optimize = list(
            filter(lambda p: p.requires_grad, params_to_optimize))
-        # Parse betas from string format "beta1,beta2"
-        betas_str = training_args.betas
-        betas = tuple(float(x.strip()) for x in betas_str.split(","))
-        
        self.optimizer = torch.optim.AdamW(
            params_to_optimize,
            lr=training_args.learning_rate,
-            betas=betas,
+            betas=(0.9, 0.999),
            weight_decay=training_args.weight_decay,
            eps=1e-8,
        )
@@ -202,7 +202,6 @@ def save_distillation_checkpoint(generator_transformer,
                                 generator_scheduler=None,
                                 fake_score_scheduler=None,
                                 noise_generator=None,
-                                 generator_ema=None,
                                 only_save_generator_weight=False) -> None:
    """
    Save distillation checkpoint with both generator and fake_score models.
@@ -234,8 +233,6 @@ def save_distillation_checkpoint(generator_transformer,
        if generator_scheduler is not None:
            generator_states["scheduler"] = SchedulerWrapper(
                generator_scheduler)
-        if generator_ema is not None:
-            generator_states["ema"] = generator_ema.state_dict()

        generator_dcp_dir = os.path.join(save_dir, "distributed_checkpoint",
                                         "generator")
@@ -405,8 +402,7 @@ def load_distillation_checkpoint(generator_transformer,
                                 dataloader=None,
                                 generator_scheduler=None,
                                 fake_score_scheduler=None,
-                                 noise_generator=None,
-                                 generator_ema=None) -> int:
+                                 noise_generator=None) -> int:
    """
    Load distillation checkpoint with both generator and fake_score models.
    Returns the step number from which training should resume.
@@ -460,18 +456,6 @@ def load_distillation_checkpoint(generator_transformer,
        end_time - begin_time,
        local_main_process_only=False)

-    # Load EMA state if available and generator_ema is provided
-    if generator_ema is not None:
-        try:
-            ema_state = generator_states.get("ema")
-            if ema_state is not None:
-                generator_ema.load_state_dict(ema_state)
-                logger.info("rank: %s, generator EMA state loaded successfully", rank)
-            else:
-                logger.info("rank: %s, no EMA state found in checkpoint", rank)
-        except Exception as e:
-            logger.warning("rank: %s, failed to load EMA state: %s", rank, str(e))
-
    # Load critic distributed checkpoint
    critic_dcp_dir = os.path.join(checkpoint_path, "distributed_checkpoint",
                                  "critic")
@@ -1294,156 +1278,3 @@ def get_scheduler(
                         num_warmup_steps=num_warmup_steps,
                         num_training_steps=num_training_steps,
                         last_epoch=last_epoch)
-
-
-class EMA_FSDP:
-    """
-    FSDP2-friendly EMA with two modes:
-      - mode="local_shard" (default): maintain float32 CPU EMA of local parameter shards on every rank.
-        Provides a context manager to temporarily swap EMA weights into the live model for teacher forward.
-      - mode="rank0_full": maintain a consolidated float32 CPU EMA of full parameters on rank 0 only
-        using gather_state_dict_on_cpu_rank0(). Useful for checkpoint export; not for teacher forward.
-
-    Usage (local_shard for CM teacher):
-      ema = EMA_FSDP(model, decay=0.999, mode="local_shard")
-      for step in ...:
-          ema.update(model)
-      with ema.apply_to_model(model):
-          with torch.no_grad():
-              y_teacher = model(...)
-
-    Usage (rank0_full for export):
-      ema = EMA_FSDP(model, decay=0.999, mode="rank0_full")
-      ema.update(model)
-      ema.state_dict()  # on rank 0
-    """
-    def __init__(self, module, decay: float = 0.999, mode: str = "local_shard"):
-        self.decay = float(decay)
-        self.mode = mode
-        self.shadow: dict[str, torch.Tensor] = {}
-        self.rank = dist.get_rank() if dist.is_initialized() else 0
-        if self.mode not in {"local_shard", "rank0_full"}:
-            raise ValueError(f"Unsupported EMA_FSDP mode: {self.mode}")
-        self._init_shadow(module)
-
-    @staticmethod
-    def _to_local_tensor(t: torch.Tensor) -> torch.Tensor:
-        # DTensor-aware to_local fetch; fall back to raw tensor
-        try:
-            from torch.distributed.tensor import DTensor  # type: ignore
-            if isinstance(t, DTensor):
-                return t.to_local()
-        except Exception:
-            pass
-        return t
-
-    @torch.no_grad()
-    def _init_shadow(self, module):
-        if self.mode == "rank0_full":
-            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
-            if self.rank == 0:
-                self.shadow = {k: v.detach().clone().float().cpu() for k, v in cpu_state.items()}
-            else:
-                self.shadow = {}
-            return
-
-        # local_shard: maintain EMA of local shards for requires_grad params
-        self.shadow = {}
-        for name, p in module.named_parameters():
-            if not p.requires_grad:
-                continue
-            local = self._to_local_tensor(p.detach())
-            self.shadow[name] = local.clone().float().cpu()
-
-    @torch.no_grad()
-    def update(self, module):
-        d = self.decay
-        if self.mode == "rank0_full":
-            if self.rank != 0:
-                return
-            cpu_state = gather_state_dict_on_cpu_rank0(module, device=None)
-            for n, v in cpu_state.items():
-                v_cpu = v.detach().float().cpu()
-                if n not in self.shadow:
-                    self.shadow[n] = v_cpu.clone()
-                else:
-                    self.shadow[n].mul_(d).add_(v_cpu, alpha=1.0 - d)
-            return
-
-        # local_shard: update local shard EMA on every rank
-        for name, p in module.named_parameters():
-            if not p.requires_grad:
-                continue
-            local = self._to_local_tensor(p.detach())
-            v_cpu = local.float().cpu()
-            if name not in self.shadow:
-                self.shadow[name] = v_cpu.clone()
-            else:
-                self.shadow[name].mul_(d).add_(v_cpu, alpha=1.0 - d)
-
-    def state_dict(self) -> dict[str, torch.Tensor]:
-        if self.mode == "rank0_full":
-            return {k: v.clone() for k, v in self.shadow.items()} if self.rank == 0 else {}
-        return {k: v.clone() for k, v in self.shadow.items()}
-
-    def load_state_dict(self, sd: dict[str, torch.Tensor]):
-        self.shadow = {k: v.clone() for k, v in sd.items()}
-
-    @torch.no_grad()
-    def copy_to_unwrapped(self, module) -> None:
-        """
-        Copy EMA weights into a non-sharded (unwrapped) module. Intended for export/eval.
-        For mode="rank0_full", only rank 0 has the full EMA state.
-        """
-        if self.mode == "rank0_full" and self.rank != 0:
-            return
-        name_to_param = dict(module.named_parameters())
-        for n, w in self.shadow.items():
-            if n in name_to_param:
-                p = name_to_param[n]
-                p.data.copy_(w.to(dtype=p.dtype, device=p.device))
-
-    class _ApplyEMACtx:
-        def __init__(self, ema: "EMA_FSDP", module):
-            self.ema = ema
-            self.module = module
-            self.saved: dict[str, torch.Tensor] = {}
-
-        def __enter__(self):
-            if self.ema.mode != "local_shard":
-                raise RuntimeError("EMA apply_to_model is only supported for mode='local_shard'")
-            with torch.no_grad():
-                for name, p in self.module.named_parameters():
-                    if not p.requires_grad:
-                        continue
-                    # Save local shard
-                    p_local = EMA_FSDP._to_local_tensor(p.detach())
-                    if p_local.numel() == 0:
-                        # Nothing to swap on this rank for this param
-                        continue
-                    self.saved[name] = p_local.clone().to(device=p_local.device, dtype=p_local.dtype)
-                    if name in self.ema.shadow:
-                        ema_cpu = self.ema.shadow[name]
-                        if ema_cpu.numel() != p_local.numel():
-                            # Shard shape mismatch (e.g., empty shard here), skip
-                            continue
-                        # Copy EMA shard into local param shard
-                        p_local.copy_(ema_cpu.to(dtype=p_local.dtype, device=p_local.device))
-            return self.module
-
-        def __exit__(self, exc_type, exc, tb):
-            with torch.no_grad():
-                for name, p in self.module.named_parameters():
-                    if name in self.saved:
-                        p_local = EMA_FSDP._to_local_tensor(p.detach())
-                        if p_local.numel() == 0:
-                            continue
-                        saved_local = self.saved[name]
-                        if saved_local.numel() != p_local.numel():
-                            continue
-                        p_local.copy_(saved_local)
-            self.saved.clear()
-            return False
-
-    def apply_to_model(self, module):
-        return EMA_FSDP._ApplyEMACtx(self, module)
@@ -1,77 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import sys
-from copy import deepcopy
-
-from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
-from fastvideo.logger import init_logger
-from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler)
-from fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline import WanCausalDMDPipeline
-from fastvideo.training.self_forcing_distillation_pipeline import SelfForcingDistillationPipeline
-from fastvideo.utils import is_vsa_available
-
-vsa_available = is_vsa_available()
-
-logger = init_logger(__name__)
-
-
-class WanSelfForcingDistillationPipeline(SelfForcingDistillationPipeline):
-    """
-    A self-forcing distillation pipeline for Wan that uses the self-forcing methodology
-    with DMD for video generation.
-    """
-    _required_config_modules = [
-        "scheduler", "transformer", "vae", "real_score_transformer",
-        "fake_score_transformer"
-    ]
-
-    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
-        """Initialize Wan-specific scheduler."""
-        self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(
-            shift=fastvideo_args.pipeline_config.flow_shift)
-
-    def create_training_stages(self, training_args: TrainingArgs):
-        """
-        May be used in future refactors.
-        """
-        pass
-
-    def initialize_validation_pipeline(self, training_args: TrainingArgs):
-        logger.info("Initializing validation pipeline...")
-        args_copy = deepcopy(training_args)
-
-        args_copy.inference_mode = True
-        validation_pipeline = WanCausalDMDPipeline.from_pretrained(
-            training_args.model_path,
-            args=args_copy,  # type: ignore
-            inference_mode=True,
-            loaded_modules={"transformer": self.get_module("transformer")},
-            tp_size=training_args.tp_size,
-            sp_size=training_args.sp_size,
-            num_gpus=training_args.num_gpus,
-            pin_cpu_memory=training_args.pin_cpu_memory,
-            dit_cpu_offload=True)
-
-        self.validation_pipeline = validation_pipeline
-
-
-def main(args) -> None:
-    logger.info("Starting Wan self-forcing distillation pipeline...")
-
-    pipeline = WanSelfForcingDistillationPipeline.from_pretrained(
-        args.pretrained_model_name_or_path, args=args)
-
-    args = pipeline.training_args
-    pipeline.train()
-    logger.info("Wan self-forcing distillation pipeline completed")
-
-
-if __name__ == "__main__":
-    argv = sys.argv
-    from fastvideo.fastvideo_args import TrainingArgs
-    from fastvideo.utils import FlexibleArgumentParser
-    parser = FlexibleArgumentParser()
-    parser = TrainingArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    main(args) 
@@ -34,7 +34,6 @@ from torch.distributed.fsdp import MixedPrecisionPolicy
 import fastvideo.envs as envs
 from fastvideo.logger import init_logger

-
 logger = init_logger(__name__)

 T = TypeVar("T")
@@ -615,6 +614,7 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]:
            f"Failed to download or parse model_index.json for {model_name_or_path}: {e}"
        ) from e

+
 def update_environment_variables(envs: dict[str, str]):
    for k, v in envs.items():
        if k in os.environ and os.environ[k] != v:
@@ -812,3 +812,66 @@ 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
@@ -1 +1 @@
-__version__ = "0.1.5"
+__version__ = "0.1.6"
@@ -13,7 +13,8 @@ import pyarrow.parquet as pq
 import torch
 from datasets import Dataset, Video, load_dataset

-from fastvideo.configs.configs import DatasetType, PreprocessConfig
+from fastvideo.configs.configs import (DatasetType, PreprocessConfig,
+                                       VideoLoaderType)
 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
@@ -431,7 +432,8 @@ def build_dataset(preprocess_config: PreprocessConfig, split: str,
            return item

        dataset = dataset.map(add_video_column)
-        dataset = dataset.cast_column("video", Video())
+        if preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC:
+            dataset = dataset.cast_column("video", Video())
    else:
        raise ValueError(
            f"Invalid dataset type: {preprocess_config.dataset_type}")
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"

 [project]
 name = "fastvideo"
-version = "0.1.5"
+version = "0.1.6"
 description = "FastVideo"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -3,7 +3,7 @@ from huggingface_hub import HfApi
 api = HfApi()

 api.upload_folder(
-    folder_path="wow",
-    repo_id="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",
+    folder_path="Wan2.2-TI2V-5B-Diffusers",
+    repo_id="FastVideo/FastWan2.2-TI2V-5B-Diffusers",
    repo_type="model",
 )
@@ -0,0 +1,21 @@
+#!/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!"
@@ -0,0 +1,21 @@
+#!/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!" 
@@ -0,0 +1,88 @@
+#!/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!"
@@ -0,0 +1,89 @@
+#!/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!"
Author	SHA1	Message	Date
RandNMR73	da5ca94091	preprocessing text	2025-09-06 05:56:13 +00:00
William Lin	d3ceb67e66	[misc] Update Slack invite link (#786 )	2025-09-05 12:16:18 -07:00
Zhang Peiyuan	7ac153a5ca	Update WeChat Link	2025-09-05 11:40:47 -07:00
William Lin	d1e7aa0abd	[CI] Add ssim test for causal inference (#784 )	2025-09-05 01:23:01 -07:00
William Lin	2d846c55a1	[misc] Improve text encoding stage (#774 )	2025-09-04 17:51:27 -07:00
Jinzhe Pan	b318063c0a	[Preprocess][Fix] video quality issue (#773 )	2025-09-03 20:47:33 -07:00
Jinzhe Pan	4aa307be55	[Preprocess][Feat] support torchvision to load video in new preprocessing (#761 )	2025-09-01 23:37:01 -07:00
William Lin	055e52e5ea	[misc] [VSA] [STA] fix tk_root in setup.py for VSA and STA (#772 )	2025-08-29 01:13:37 -07:00
William Lin	7d2069596b	[bugfix] [VSA] [STA] Fix MANIFEST.in for VSA and STA; Move tk into both directories (#771 )	2025-08-29 00:51:05 -07:00
William Lin	c45009c9a4	[bugfix] fix STA install setup.py import (#770 )	2025-08-28 23:02:53 -07:00
William LinandPeiyuan Zhang	b91020b407	[VSA] [STA] Fix directory structure for pypi publishing (#769 ) Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>	2025-08-28 22:34:03 -07:00
William Lin	2dcc5ea4f6	[chore] Release 0.1.6 (#768 )	2025-08-28 20:56:21 -07:00
Wei ZhouandSolitaryThinker	359151d9a0	[Feature] Add wan2.2 5b i2v (#760 ) Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>	2025-08-28 18:15:59 -07:00
Wei ZhouandSolitaryThinker	ce67cd3729	[Feat] Support Self-Forcing's Causal Inference for Wan2.1 T2V 1.3B (#766 ) Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>	2025-08-28 16:47:49 -07:00