update

Co-authored-by: RandNMR73 <notomatthew31@gmail.com>

.buildkite/pipeline.yml

@@ -163,6 +163,7 @@ steps:
             - path:
                 - "csrc/attn/vsa/**"
                 - "csrc/attn/tk/**"
+                - "csrc/attn/tests/test_vsa.py"
                 - "csrc/attn/setup_vsa.py"
                 - "csrc/attn/config_vsa.py"
                 - "csrc/attn/vsa.cpp"

.github/workflows/vsa-publish.yml

@@ -0,0 +1,257 @@
+name: Publish Video Sparse Attention Kernel to PyPI on Version Change
+
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - "csrc/attn/setup_vsa.py"
+  workflow_dispatch:
+
+jobs:
+  check-version-change:
+    runs-on: ubuntu-latest
+    outputs:
+      version-changed: ${{ steps.check-version.outputs.changed }}
+      new-version: ${{ steps.check-version.outputs.new-version }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 2
+
+      - name: Check if version changed
+        id: check-version
+        run: |
+          cd csrc/attn
+          # Get current commit's version
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_vsa.py)
+          echo "New version: $NEW_VERSION"
+          
+          # Get previous version from git history
+          OLD_VERSION=$(git show HEAD~1:./setup_vsa.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          echo "Old version: $OLD_VERSION"
+          
+          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
+            echo "Version changed from $OLD_VERSION to $NEW_VERSION"
+            echo "changed=true" >> $GITHUB_OUTPUT
+            echo "new-version=$NEW_VERSION" >> $GITHUB_OUTPUT
+          else
+            echo "Version did not change"
+            echo "changed=false" >> $GITHUB_OUTPUT
+          fi
+
+  build_wheels:
+    name: Build Wheel
+    needs: check-version-change
+    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+          # Using ubuntu-20.04 instead of 22.04 for more compatibility (glibc). Ideally we'd use the
+          # manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
+          os: [ubuntu-22.04]
+          python-version: ['3.10', '3.11', '3.12', '3.13']
+          # For version reference https://pytorch.org/get-started/previous-versions/
+          torch-cuda:
+            - torch-version: '2.5.1'
+              cuda-version: '12.4.1'
+              torch-cuda-short: 'cu124'
+            - torch-version: '2.6.0'
+              cuda-version: '12.6.3'
+              torch-cuda-short: 'cu126'
+            - torch-version: '2.7.1'
+              cuda-version: '12.8.0'
+              torch-cuda-short: 'cu128'
+
+    steps:
+      - name: Free up disk space
+        run: |
+          echo "Initial disk space:"
+          df -h
+          
+          # Remove large directories
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          sudo rm -rf /usr/share/swift
+          sudo rm -rf /usr/local/lib/node_modules
+          sudo rm -rf /usr/local/share/powershell
+          sudo rm -rf /usr/share/rust
+          sudo rm -rf /usr/local/.ghcup
+          
+          # Remove cached files
+          sudo rm -rf /var/lib/apt/lists/*
+          sudo rm -rf /var/cache/apt/archives/*
+          
+          echo "Disk space after cleanup:"
+          df -h
+
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install CUDA ${{ matrix.torch-cuda.cuda-version }}
+        uses: Jimver/cuda-toolkit@v0.2.21
+        id: cuda-toolkit
+        with:
+          cuda: ${{ matrix.torch-cuda.cuda-version }}
+          linux-local-args: '["--toolkit"]'
+          method: 'network'
+
+      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
+        run: |
+          sudo apt update
+          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
+          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+          # Allow Git to Access Safe Directory
+          git config --global --add safe.directory /__w/FastVideo/FastVideo
+
+          # Set CUDA environment variables
+          export CUDA_HOME=/usr/local/cuda-${{ matrix.torch-cuda.cuda-version }}
+          export PATH=${CUDA_HOME}/bin:${PATH}
+          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+          
+          # Verify installation
+          gcc --version
+          g++ --version
+          clang-11 --version
+          nvcc --version
+
+      - name: Install PyTorch ${{ matrix.torch-cuda.torch-version }}+cu${{ matrix.torch-cuda.cuda-version }}
+        run: |
+          pip install --upgrade pip
+          # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
+          # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
+          pip install typing-extensions==4.12.2
+          # We want to figure out the CUDA version to download pytorch
+          # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
+          # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
+          pip install --no-cache-dir torch==${{ matrix.torch-cuda.torch-version }} --index-url https://download.pytorch.org/whl/${{matrix.torch-cuda.torch-cuda-short}}
+          nvcc --version
+          python --version
+          python -c "import torch; print('PyTorch:', torch.__version__)"
+          python -c "import torch; print('CUDA:', torch.version.cuda)"
+          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
+
+      - name: Build wheel
+        run: |
+          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
+          
+          # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
+          # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
+          # However this still fails so I'm using a newer version of setuptools
+          pip install setuptools
+          pip install ninja packaging wheel
+
+          cd csrc/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
+
+      - name: Rename wheel file
+        run: |
+          cd csrc/attn
+
+          CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
+          TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2)
+          # Get the correct version format
+          tmpname=cu${CUDA_SHORT_VERSION}torch${TORCH_SHORT_VERSION}
+          wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
+          # Rename with version information
+          ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
+          echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
+
+      - name: Upload wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: ${{ env.wheel_name }}
+          path: csrc/attn/dist/*.whl
+          retention-days: 90
+
+  publish_package:
+    name: Publish package
+    needs: [build_wheels, check-version-change]
+    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
+    runs-on: ubuntu-22.04
+    permissions:
+      id-token: write  # Needed for OIDC Trusted Publishing
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.10'
+
+      - name: Install CUDA 12.4.1
+        uses: Jimver/cuda-toolkit@v0.2.21
+        id: cuda-toolkit
+        with:
+          cuda: 12.4.1
+          linux-local-args: '["--toolkit"]'
+          method: 'network'
+          sub-packages: '["nvcc"]'
+
+      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
+        run: |
+          sudo apt update
+          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
+          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+
+          # Allow Git to Access Safe Directory
+          git config --global --add safe.directory /__w/FastVideo/FastVideo
+
+          # Set CUDA environment variables
+          export CUDA_HOME=/usr/local/cuda-12.4.1
+          export PATH=${CUDA_HOME}/bin:${PATH}
+          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+          
+          # Verify installation
+          gcc --version
+          g++ --version
+          clang-11 --version
+          nvcc --version
+
+      - name: Install PyTorch 2.5.1+cu12.4.1
+        run: |
+          pip install --upgrade pip
+          # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
+          # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
+          pip install typing-extensions==4.12.2
+          # We want to figure out the CUDA version to download pytorch
+          # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
+          # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
+          export TORCH_CUDA_VERSION=124
+          pip install --no-cache-dir torch==2.5.1 --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
+          nvcc --version
+          python --version
+          python -c "import torch; print('PyTorch:', torch.__version__)"
+          python -c "import torch; print('CUDA:', torch.version.cuda)"
+          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
+
+      - name: Build source distribution
+        run: |
+          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
+          
+          # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
+          # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
+          # However this still fails so I'm using a newer version of setuptools
+          pip install setuptools
+          pip install ninja packaging wheel
+
+          cd csrc/attn  # Move into the correct folder
+          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
+          python setup_vsa.py sdist --dist-dir=dist
+
+      - name: Publish release distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: csrc/attn/dist/

README.md

@@ -1,21 +1,21 @@
 <div align="center">
-<img src=assets/logo.jpg width="30%"/>
+<img src=assets/logo.png width="30%"/>
 </div>
 
-**FastVideo is a unified framework for accelerated video generation.**
+**FastVideo is a unified post-training and inference framework for accelerated video generation.**
 
-It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
-With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
+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://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/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |
+    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/qqPzbrw" target="_blank"> <b> WeChat </b> </a> |
 </p>
 
 <div align="center">
-<img src=assets/perf.png width="90%"/>
+<img src=assets/fastwan.png width="90%"/>
 </div>
 
 ## NEWS
+- ```2025/08/04```: Release [FastWan](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html) models and [Sparse-Distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
 - ```2025/06/14```: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
 - ```2025/04/24```: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
 - ```2025/02/18```: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
@@ -23,20 +23,19 @@ With FastVideo's optimizations, you can achieve more than 3x inference improveme
 ## Key Features
 
 FastVideo has the following features:
+- End-to-end post-training support:
+  - [Sparse distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/) for Wan2.1 and Wan2.2 to achineve >50x denoising speedup
+  - Data preprocessing pipeline for video data
+  - Support full finetuning and LoRA finetuning for state-of-the-art open video DiTs
+  - Scalable training with FSDP2, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs
 - State-of-the-art performance optimizations for inference
+  - [Video Sparse Attention](https://arxiv.org/pdf/2505.13389)
   - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
   - [TeaCache](https://arxiv.org/pdf/2411.19108)
   - [Sage Attention](https://arxiv.org/abs/2410.02367)
-- Cutting edge models
-  - Wan2.1 T2V, I2V
-  - HunyuanVideo
-  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
-  - StepVideo T2V
-- Distillation support
-  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
-  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
-- Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
-- Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
+- Diverse hardware and OS support
+  - Support H100, A100, 4090
+  - Support Linux, Windows, MacOS
 
 ## Getting Started
 We recommend using an environment manager such as `Conda` to create a clean environment:
@@ -52,17 +51,31 @@ pip install fastvideo
 
 Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) for more detailed installation instructions.
 
+## Sparse Distillation
+For our sparse distillation techniques, please see our [distillation docs](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html) and check out our [blog](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).
+
+See below for recipes and datasets:
+
+|                                            Model                                              |                                               Sparse Distillation                                                 |                                                  Dataset                                                  |
+|:-------------------------------------------------------------------------------------------:  |:---------------------------------------------------------------------------------------------------------------:  |:--------------------------------------------------------------------------------------------------------: |
+| [FastWan2.1-T2V-1.3B](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers)         |    [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P)      | [FastVideo Synthetic Wan2.1 480P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)      |
+| [FastWan2.1-T2V-14B-Preview](https://huggingface.co/FastVideo/FastWan2.1-T2V-14B-Diffusers)   |                                                   Coming soon!                                                    |   [FastVideo Synthetic Wan2.1 720P](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x768x1280_250k)   |
+| [FastWan2.2-TI2V-5B](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-Diffusers)           | [Recipe](https://github.com/hao-ai-lab/FastVideo/tree/main/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free)   | [FastVideo Synthetic Wan2.2 720P](https://huggingface.co/datasets/FastVideo/Wan2.2-Syn-121x704x1280_32k)  |
+
 ## Inference
 ### Generating Your First Video
-Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
+Here's a minimal example to generate a video using the default settings. Make sure VSA kernels are [installed](https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation.html). Create a file called `example.py` with the following code:
 
 ```python
+import os
 from fastvideo import VideoGenerator
 
 def main():
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+
     # Create a video generator with a pre-trained model
     generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
         num_gpus=1,  # Adjust based on your hardware
     )
 
@@ -96,41 +109,41 @@ For a more detailed guide, please see our [inference quick start](https://hao-ai
 
 ## Distillation and Finetuning
 - [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html)
-- [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html)
+<!-- - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html) -->
 
 ## 📑 Development Plan
-
 <!-- - More distillation methods -->
   <!-- - [ ] Add Distribution Matching Distillation -->
-- More models support
-  <!-- - [ ] Add CogvideoX model -->
-  - [x] Add StepVideo to V1
-- Optimization features
-  - [x] Teacache in V1
-  - [x] SageAttention in V1
-- Code updates
-  - [x] V1 Configuration API
-  - [ ] Support Training in V1
+More FastWan Models Coming Soon!
+- [ ] Add FastWan2.1-T2V-14B
+- [ ] Add FastWan2.2-T2V-14B
+- [ ] Add FastWan2.2-I2V-14B
+<!-- - Optimization features
+- Code updates -->
   <!-- - [ ] fp8 support -->
   <!-- - [ ] faster load model and save model support -->
 
+See details in [development roadmap](https://github.com/hao-ai-lab/FastVideo/issues/468).
+
 ## 🤝 Contributing
 
 We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview.html)
 
 ## Acknowledgement
 We learned and reused code from the following projects:
-- [PCM](https://github.com/G-U-N/Phased-Consistency-Model)
+- [Wan-Video](https://github.com/Wan-Video)
+- [ThunderKittens](https://github.com/HazyResearch/ThunderKittens)
+- [Triton](https://github.com/triton-lang/triton)
+- [DMD2](https://github.com/tianweiy/DMD2)
 - [diffusers](https://github.com/huggingface/diffusers)
-- [OpenSoraPlan](https://github.com/PKU-YuanGroup/Open-Sora-Plan)
 - [xDiT](https://github.com/xdit-project/xDiT)
 - [vLLM](https://github.com/vllm-project/vllm)
 - [SGLang](https://github.com/sgl-project/sglang)
 
-We thank MBZUAI and [Anyscale](https://www.anyscale.com/) for their support throughout this project.
+We thank [MBZUAI](https://ifm.mbzuai.ac.ae/), [Anyscale](https://www.anyscale.com/), and [GMI Cloud](https://www.gmicloud.ai/) for their support throughout this project.
 
 ## Citation
-If you use FastVideo for your research, please cite our work:
+If you find FastVideo useful, please considering citing our work:
 
 ```bibtex
 @software{fastvideo2024,
@@ -141,31 +154,17 @@ If you use FastVideo for your research, please cite our work:
   year         = {2024},
 }
 
-@misc{zhang2025vsafastervideodiffusion,
-      title={VSA: Faster Video Diffusion with Trainable Sparse Attention}, 
-      author={Peiyuan Zhang and Haofeng Huang and Yongqi Chen and Will Lin and Zhengzhong Liu and Ion Stoica and Eric Xing and Hao Zhang},
-      year={2025},
-      eprint={2505.13389},
-      archivePrefix={arXiv},
-      primaryClass={cs.CV},
-      url={https://arxiv.org/abs/2505.13389}, 
+@article{zhang2025vsa,
+  title={VSA: Faster Video Diffusion with Trainable Sparse Attention},
+  author={Zhang, Peiyuan and Huang, Haofeng and Chen, Yongqi and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
+  journal={arXiv preprint arXiv:2505.13389},
+  year={2025}
 }
-@misc{zhang2025fastvideogenerationsliding,
-      title={Fast Video Generation with Sliding Tile Attention},
-      author={Peiyuan Zhang and Yongqi Chen and Runlong Su and Hangliang Ding and Ion Stoica and Zhenghong Liu and Hao Zhang},
-      year={2025},
-      eprint={2502.04507},
-      archivePrefix={arXiv},
-      primaryClass={cs.CV},
-      url={https://arxiv.org/abs/2502.04507},
-}
-@misc{ding2025efficientvditefficientvideodiffusion,
-      title={Efficient-vDiT: Efficient Video Diffusion Transformers With Attention Tile},
-      author={Hangliang Ding and Dacheng Li and Runlong Su and Peiyuan Zhang and Zhijie Deng and Ion Stoica and Hao Zhang},
-      year={2025},
-      eprint={2502.06155},
-      archivePrefix={arXiv},
-      primaryClass={cs.CV},
-      url={https://arxiv.org/abs/2502.06155},
+
+@article{zhang2025fast,
+  title={Fast video generation with sliding tile attention},
+  author={Zhang, Peiyuan and Chen, Yongqi and Su, Runlong and Ding, Hangliang and Stoica, Ion and Liu, Zhengzhong and Zhang, Hao},
+  journal={arXiv preprint arXiv:2502.04507},
+  year={2025}
 }
 ```

csrc/attn/README.md

@@ -84,7 +84,7 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False)
 ### Test
 ```bash
 python tests/test_sta.py # test STA
-python tests/test_block_sparse.py # test VSA
+python tests/test_vsa.py # test VSA
 ```
 ### Benchmark
 ```bash

csrc/attn/benchmarks/bench_sta.py

@@ -86,8 +86,6 @@ def benchmark_attention(configurations):
         # print(f"Average TFLOPS: {tflops_bwd}")
         # print("=" * 60)
 
-        torch.cuda.empty_cache()
-
     return results
 
 

csrc/attn/setup_vsa.py

@@ -51,19 +51,16 @@ for k in kernels:
         source_files.append(sources[k]['source_files'][target])
     cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}')
 
-ext_modules = []
-import torch
-major, minor = torch.cuda.get_device_capability(0)
-if major == 9 and minor == 0:# check if H100
-    ext_modules = [
-        CUDAExtension('vsa_cuda',
-            sources=source_files,
-            extra_compile_args={
-                'cxx': cpp_flags,
-                'nvcc': cuda_flags
-            },
-            libraries=['cuda'])
-    ]
+
+ext_modules = [
+    CUDAExtension('vsa_cuda',
+        sources=source_files,
+        extra_compile_args={
+            'cxx': cpp_flags,
+            'nvcc': cuda_flags
+        },
+        libraries=['cuda'])
+]
 
 
 

csrc/attn/vsa/__init__.py

@@ -1,12 +1,13 @@
 import torch
 from typing import Tuple
 block_sparse_attn=None
-
-try:
+import torch
+major, minor = torch.cuda.get_device_capability(0)
+if major == 9 and minor == 0:# check if H100
     from vsa_cuda import block_sparse_fwd, block_sparse_bwd
     from vsa.block_sparse_wrapper import block_sparse_attn_SM90
     block_sparse_attn = block_sparse_attn_SM90
-except ImportError:
+else:
     from vsa.block_sparse_wrapper import block_sparse_attn_triton
     block_sparse_fwd = None
     block_sparse_bwd = None

csrc/attn/vsa/block_sparse_h100.cu

@@ -578,6 +578,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
         if (threadIdx.x / 32 == 0) {
             coord<qg_tile> tile_idx = {blockIdx.z, blockIdx.y, store_qg_block_index, 0};
             tma::store_add_async(g.qg, qg_smem, tile_idx);
+            tma::store_async_wait();
         }
     }
 
@@ -634,13 +635,14 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
         if (threadIdx.x / 32 == 0) {
             coord<qg_tile> tile_idx = {blockIdx.z, blockIdx.y, store_qg_block_index, 0};
             tma::store_add_async(g.qg, qg_smem, tile_idx);
+            tma::store_async_wait();
         }
     }
 
     // store kq and vq
 
     // ! the following two line seems unnecessary.
-    tma::store_async_wait(); // ensure qg is finished
+    // tma::store_async_wait(); // ensure qg is finished
     __syncthreads();
 
     warpgroup::store(kg_smem[0], kg_reg);
@@ -1174,4 +1176,4 @@ block_sparse_attention_backward(torch::Tensor q,
 
     return {qg, kg, vg};
     //cudadevicesynchronize();
-}
+}

csrc/attn/vsa/block_sparse_wrapper.py

@@ -11,95 +11,6 @@ from typing import Tuple, Optional
 
 
 
-@torch.library.custom_op("vsa::block_sparse_attn_SM90", mutates_args=(), device_types="cuda")
-def block_sparse_attn_SM90(
-    q_padded: torch.Tensor,
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor, 
-    )-> Tuple[torch.Tensor, torch.Tensor]:
-    q_padded = q_padded.contiguous()
-    k_padded = k_padded.contiguous()
-    v_padded = v_padded.contiguous()
-    q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
-    variable_block_sizes = variable_block_sizes.int()
-    o_padded, lse_padded = block_sparse_fwd(q_padded, k_padded, v_padded, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
-    return o_padded, lse_padded
-
-
-
-
-@torch.library.register_fake("vsa::block_sparse_attn_SM90")
-def _block_sparse_attn_SM90_fake(
-    q_padded: torch.Tensor,
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor, 
-) -> Tuple[torch.Tensor,  torch.Tensor]:
-    q_padded, k_padded, v_padded = [x.contiguous() for x in (q_padded, k_padded, v_padded)]
-    B, H, S, D = q_padded.shape
-    o_padded = torch.empty_like(q_padded)
-    lse_padded = torch.empty((B, H, S, 1), device=q_padded.device, dtype=torch.float32)
-    return o_padded, lse_padded
-
-
-@torch.library.custom_op("vsa::block_sparse_attn_backward_SM90", mutates_args=(), device_types="cuda")
-def block_sparse_attn_backward_SM90(
-    grad_output_padded: torch.Tensor, 
-    q_padded: torch.Tensor, 
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    o_padded: torch.Tensor, 
-    lse_padded: torch.Tensor, 
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor, 
-    )-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    grad_output_padded = grad_output_padded.contiguous()
-    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2))
-    grad_q_padded, grad_k_padded, grad_v_padded = block_sparse_bwd(
-        q_padded, k_padded, v_padded, o_padded, lse_padded, grad_output_padded, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes
-    )
-    grad_q_padded = grad_q_padded.to(grad_output_padded.dtype)
-    grad_k_padded = grad_k_padded.to(grad_output_padded.dtype)
-    grad_v_padded = grad_v_padded.to(grad_output_padded.dtype)
-    return grad_q_padded, grad_k_padded, grad_v_padded
-
-@torch.library.register_fake("vsa::block_sparse_attn_backward_SM90")
-def _block_sparse_attn_backward_SM90_fake(
-    grad_output_padded: torch.Tensor, 
-    q_padded: torch.Tensor, 
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    o_padded: torch.Tensor, 
-    lse_padded: torch.Tensor, 
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor, 
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    torch._check(grad_output_padded.dtype == torch.bfloat16)
-    torch._check(lse_padded.dtype == torch.float32)
-    grad_output_padded = grad_output_padded.contiguous()
-    dq = torch.empty_like(grad_output_padded)
-    dk = torch.empty_like(grad_output_padded)
-    dv = torch.empty_like(grad_output_padded)
-    return dq, dk, dv
-
-
-def backward_SM90(ctx, grad_output1, grad_output2):
-    q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes= ctx.saved_tensors
-    dq, dk, dv = block_sparse_attn_backward_SM90(grad_output1, q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
-    return dq, dk, dv, None, None
-
-def setup_context_SM90(ctx, inputs, output):
-    q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs
-    o_padded, lse_padded = output
-    ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
-
-
-block_sparse_attn_SM90.register_autograd(backward_SM90, setup_context=setup_context_SM90)
-
-
 @torch.library.custom_op("vsa::block_sparse_attn_triton", mutates_args=(), device_types="cuda")
 def block_sparse_attn_triton(
     q: torch.Tensor,
@@ -179,4 +90,96 @@ def setup_context_triton(ctx, inputs, output):
     o_padded, M = output
     ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes)
     
-block_sparse_attn_triton.register_autograd(backward_triton, setup_context=setup_context_triton)
+block_sparse_attn_triton.register_autograd(backward_triton, setup_context=setup_context_triton)
+
+
+major, minor = torch.cuda.get_device_capability(0)
+
+if major == 9 and minor == 0:# check if H100
+    @torch.library.custom_op("vsa::block_sparse_attn_SM90", mutates_args=(), device_types="cuda")
+    def block_sparse_attn_SM90(
+        q_padded: torch.Tensor,
+        k_padded: torch.Tensor, 
+        v_padded: torch.Tensor, 
+        block_map: torch.Tensor,
+        variable_block_sizes: torch.Tensor, 
+        )-> Tuple[torch.Tensor, torch.Tensor]:
+        q_padded = q_padded.contiguous()
+        k_padded = k_padded.contiguous()
+        v_padded = v_padded.contiguous()
+        q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
+        variable_block_sizes = variable_block_sizes.int()
+        o_padded, lse_padded = block_sparse_fwd(q_padded, k_padded, v_padded, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
+        return o_padded, lse_padded
+
+
+
+
+    @torch.library.register_fake("vsa::block_sparse_attn_SM90")
+    def _block_sparse_attn_SM90_fake(
+        q_padded: torch.Tensor,
+        k_padded: torch.Tensor, 
+        v_padded: torch.Tensor, 
+        block_map: torch.Tensor,
+        variable_block_sizes: torch.Tensor, 
+    ) -> Tuple[torch.Tensor,  torch.Tensor]:
+        q_padded, k_padded, v_padded = [x.contiguous() for x in (q_padded, k_padded, v_padded)]
+        B, H, S, D = q_padded.shape
+        o_padded = torch.empty_like(q_padded)
+        lse_padded = torch.empty((B, H, S, 1), device=q_padded.device, dtype=torch.float32)
+        return o_padded, lse_padded
+
+
+    @torch.library.custom_op("vsa::block_sparse_attn_backward_SM90", mutates_args=(), device_types="cuda")
+    def block_sparse_attn_backward_SM90(
+        grad_output_padded: torch.Tensor, 
+        q_padded: torch.Tensor, 
+        k_padded: torch.Tensor, 
+        v_padded: torch.Tensor, 
+        o_padded: torch.Tensor, 
+        lse_padded: torch.Tensor, 
+        block_map: torch.Tensor,
+        variable_block_sizes: torch.Tensor, 
+        )-> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        grad_output_padded = grad_output_padded.contiguous()
+        k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2))
+        grad_q_padded, grad_k_padded, grad_v_padded = block_sparse_bwd(
+            q_padded, k_padded, v_padded, o_padded, lse_padded, grad_output_padded, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes
+        )
+        grad_q_padded = grad_q_padded.to(grad_output_padded.dtype)
+        grad_k_padded = grad_k_padded.to(grad_output_padded.dtype)
+        grad_v_padded = grad_v_padded.to(grad_output_padded.dtype)
+        return grad_q_padded, grad_k_padded, grad_v_padded
+
+    @torch.library.register_fake("vsa::block_sparse_attn_backward_SM90")
+    def _block_sparse_attn_backward_SM90_fake(
+        grad_output_padded: torch.Tensor, 
+        q_padded: torch.Tensor, 
+        k_padded: torch.Tensor, 
+        v_padded: torch.Tensor, 
+        o_padded: torch.Tensor, 
+        lse_padded: torch.Tensor, 
+        block_map: torch.Tensor,
+        variable_block_sizes: torch.Tensor, 
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        torch._check(grad_output_padded.dtype == torch.bfloat16)
+        torch._check(lse_padded.dtype == torch.float32)
+        grad_output_padded = grad_output_padded.contiguous()
+        dq = torch.empty_like(grad_output_padded)
+        dk = torch.empty_like(grad_output_padded)
+        dv = torch.empty_like(grad_output_padded)
+        return dq, dk, dv
+
+
+    def backward_SM90(ctx, grad_output1, grad_output2):
+        q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes= ctx.saved_tensors
+        dq, dk, dv = block_sparse_attn_backward_SM90(grad_output1, q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
+        return dq, dk, dv, None, None
+
+    def setup_context_SM90(ctx, inputs, output):
+        q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs
+        o_padded, lse_padded = output
+        ctx.save_for_backward(q_padded, k_padded, v_padded, o_padded, lse_padded, block_map, variable_block_sizes)
+
+
+    block_sparse_attn_SM90.register_autograd(backward_SM90, setup_context=setup_context_SM90)

docs/source/distillation/dmd.md

@@ -6,11 +6,19 @@ We introduce a new finetuning strategy - **Sparse-distill**, which jointly integ
 
 We provide two distilled models:
 
-- **[FastWan2.1-T2V-1.3B-Diffusers](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers)**: 3-step inference, up to **20 FPS** on H100 GPU
-- **[FastWan2.1-T2V-14B-480P-Diffusers](https://huggingface.co/FastVideo/FastWan2.1-T2V-14B-480P-Diffusers)**: 3-step inference, up to **50x speed up** at 480P, **70x speed up** at 720P for denoising loop
+- **[FastWan2.1-T2V-1.3B-Diffusers](https://huggingface.co/FastVideo/FastWan2.1-T2V-1.3B-Diffusers)**: 3-step inference, up to **16 FPS** on H100 GPU
+- **[FastWan2.1-T2V-14B-480P-Diffusers](https://huggingface.co/FastVideo/FastWan2.1-T2V-14B-480P-Diffusers)**: 3-step inference, up to **60x speed up** at 480P, **90x speed up** at 720P for denoising loop
+- **[FastWan2.2-TI2V-5B-FullAttn-Diffusers](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers)**: 3-step inference, up to **50x speed up** at 720P for denoising loop
 
 Both models are trained on **61×448×832** resolution but support generating videos with **any resolution** (1.3B  model mainly support 480P, 14B model support 480P and 720P, quality may degrade for different resolutions).
 
+## ⚙️ Inference
+First install [VSA](https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation.html). Set `MODEL_BASE` to your own model path and run:
+
+```bash
+bash scripts/inference/v1_inference_wan_dmd.sh
+```
+
 ## 🗂️ Dataset
 
 We use the **FastVideo 480P Synthetic Wan dataset** ([FastVideo/Wan-Syn_77x448x832_600k](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)) for distillation, which contains 600k synthetic latents.
@@ -27,7 +35,7 @@ python scripts/huggingface/download_hf.py \
 
 ## 🚀 Training Scripts
 
-### 1.3B Model Sparse-Distill
+### Wan2.1 1.3B Model Sparse-Distill
 
 For the 1.3B model, we use **4 nodes with 32 H200 GPUs** (8 GPUs per node):
 
@@ -43,7 +51,7 @@ sbatch examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.sl
 - VSA attention sparsity: 0.8
 - Training steps: 4000 (~12 hours)
 
-### 14B Model Sparse-Distill
+### Wan2.1 14B Model Sparse-Distill
 
 For the 14B model, we use **8 nodes with 64 H200 GPUs** (8 GPUs per node):
 
@@ -61,9 +69,19 @@ sbatch examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slu
 - Training steps: 3000 (~52 hours)
 - HSDP shard dim: 8
 
-## ⚙️ Inference
-Set `MODEL_BASE` to your own model path and run:
+### Wan2.2 5B Model Sparse-Distill
+
+For the 5B model, we use **8 nodes with 64 H200 GPUs** (8 GPUs per node):
 
 ```bash
-bash scripts/inference/v1_inference_wan_dmd.sh
+# Multi-node training (8 nodes, 64 GPUs total)
+sbatch examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh 
 ```
+
+**Key Configuration:**
+- Global batch size: 64
+- Sequence parallel size: 1
+- Gradient accumulation steps: 1
+- Learning rate: 2e-5
+- Training steps: 3000 (~12 hours)
+- HSDP shard dim: 1

docs/source/index.md

@@ -1,6 +1,6 @@
 # Welcome to FastVideo
 
-:::{figure} ../../assets/logo.jpg
+:::{figure} ../../assets/logo.png
 :align: center
 :alt: FastVideo
 :class: no-scaled-link
@@ -9,7 +9,7 @@
 
 :::{raw} html
 <p style="text-align:center">
-<strong>FastVideo is a unified framework for accelerated video generation.
+<strong>FastVideo is a unified inference and post-training framework for accelerated video generation.
 </strong>
 </p>
 
@@ -21,11 +21,10 @@
 </p>
 :::
 
-It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
-With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
+FastVideo is an inference and post-training framework for diffusion models. It 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.
 
 <div style="text-align: center;">
-  <img src=_static/images/perf.png width="100%"/>
+  <img src=_static/images/fastwan.png width="100%"/>
 </div>
 
 ## Key Features
@@ -33,19 +32,13 @@ With FastVideo's optimizations, you can achieve more than 3x inference improveme
 FastVideo has the following features:
 - State-of-the-art performance optimizations for inference
   - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
-  - [Video Sparse Attention](https://arxiv.org/pdf/2505.13389)
   - [TeaCache](https://arxiv.org/pdf/2411.19108)
   - [Sage Attention](https://arxiv.org/abs/2410.02367)
-- Cutting edge models
-  - Wan2.1 T2V, I2V
-  - HunyuanVideo
-  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
-  - StepVideo T2V
-- Distillation support
-  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
-  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
-- Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
-- Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
+- E2E post-training support
+  - Data preprocessing pipeline for video data.
+  - [Sparse distillation](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/) for Wan2.1 and Wan2.2 using [Video Sparse Attention](https://arxiv.org/pdf/2505.13389) and [Distribution Matching Distillation](https://tianweiy.github.io/dmd2/)
+  - Support full finetuning and LoRA finetuning for state-of-the-art open video DiTs.
+  - Scalable training with FSDP2, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
 
 ## Documentation
 

docs/source/sliding_tile_attention/installation.md

@@ -53,3 +53,9 @@ out = sliding_tile_attention(q, k, v, window_size, text_length)
 out = sliding_tile_attention(q, k, v, window_size, 0, False)
 
 ```
+
+# 🚀Inference
+
+```bash
+bash scripts/inference/v1_inference_wan_STA.sh
+```

docs/source/video_sparse_attention/installation.md

@@ -59,3 +59,9 @@ from vsa import video_sparse_attn
 output = video_sparse_attn(q, k, v, variable_block_sizes, topk, block_size, compress_attn_weight)
 
 ``` 
+
+# 🚀Inference
+
+```bash
+bash scripts/inference/v1_inference_wan_VSA.sh
+```

examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm

@@ -86,7 +86,7 @@ validation_args=(
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
   --validation_steps 200
   --validation_sampling_steps "3"
-  --validation_guidance_scale "1.0" # not used for dmd inference
+  --validation_guidance_scale "6.0" # not used for dmd inference
 )
 
 # Optimizer arguments
@@ -102,7 +102,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.0
   --dit_precision "fp32"

examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm

@@ -86,7 +86,7 @@ validation_args=(
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
   --validation_steps 200
   --validation_sampling_steps "3"
-  --validation_guidance_scale "1.0" # not used for dmd inference
+  --validation_guidance_scale "6.0" # not used for dmd inference
 )
 
 # Optimizer arguments
@@ -102,7 +102,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.0
   --dit_precision "fp32"

examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm

@@ -86,7 +86,7 @@ validation_args=(
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
   --validation_steps 200
   --validation_sampling_steps "3"
-  --validation_guidance_scale "1.0" # not used for dmd inference
+  --validation_guidance_scale "6.0" # not used for dmd inference
 )
 
 # Optimizer arguments
@@ -102,7 +102,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.0
   --dit_precision "fp32"

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

@@ -0,0 +1,13 @@
+# Wan2.2-5B Distill Example
+These are end-to-end example scripts for distilling Wan2.2 TI2V 5B model DMD+VSA methods.
+
+### 0. Make sure you have installed VSA
+
+```bash
+cd csrc/attn
+git submodule update --init --recursive
+python setup_vsa.py install
+```
+
+### Data-free Distillation
+When `--simulate_generator_forward` is enabled, distillation becomes data-free by simulating intermediate steps through forward inference of the generator. This helps avoid training–inference mismatch. See Section 4.5 of [DMD2](https://arxiv.org/pdf/2405.14867) for details.

examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh

@@ -0,0 +1,144 @@
+#!/bin/bash
+#SBATCH --job-name=t2v
+#SBATCH --partition=main
+#SBATCH --nodes=8
+#SBATCH --ntasks=8
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=dmd_Wan2.2/t2v_g2e5_f1e5_%j.out
+#SBATCH --error=dmd_Wan2.2/t2v_g2e5_f1e5_%j.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate your_env
+
+# 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 WANDB_API_KEY=your_wandb_api_key
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+# Configs
+NUM_GPUS=8
+MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+DATA_DIR=your_data_dir
+VALIDATION_DIR=your_validation_path  #(example:validation_64.json)
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name Wan_distillation
+  --output_dir "your_output_dir"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 31
+  --num_height 704
+  --num_width 1280
+  --num_frames 121
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 64
+  --hsdp_shard_dim 1
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DIR"
+  --validation_steps 200
+  --validation_sampling_steps "3"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-5
+  --lr_scheduler "cosine_with_min_lr"
+  --min_lr_ratio 0.5
+  --lr_warmup_steps 100
+  --fake_score_learning_rate 1e-5
+  --fake_score_lr_scheduler "cosine_with_min_lr"
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 200
+  --weight_decay 0.01
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --ema_start_step 0
+  --flow_shift 5
+  --seed 1000
+)
+
+# DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '1000,757,522'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --generator_update_interval 5
+  --real_score_guidance_scale 3
+  --simulate_generator_forward 
+  --log_visualization # disable if oom
+)
+
+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_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}"

examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh

@@ -0,0 +1,145 @@
+#!/bin/bash
+#SBATCH --job-name=t2v
+#SBATCH --partition=main
+#SBATCH --nodes=8
+#SBATCH --ntasks=8
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=dmd_Wan2.2/t2v_g2e5_f1e5_%j.out
+#SBATCH --error=dmd_Wan2.2/t2v_g2e5_f1e5_%j.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate your_env
+
+# 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=VIDEO_SPARSE_ATTN
+export WANDB_API_KEY=your_wandb_api_key
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+# Configs
+NUM_GPUS=8
+MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+DATA_DIR=your_data_dir
+VALIDATION_DIR=your_validation_path  #(example:validation_64.json)
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name Wan_distillation
+  --output_dir "your_output_dir"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 31
+  --num_height 704
+  --num_width 1280
+  --num_frames 121
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 64
+  --hsdp_shard_dim 1
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DIR"
+  --validation_steps 200
+  --validation_sampling_steps "3"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-5
+  --lr_scheduler "cosine_with_min_lr"
+  --min_lr_ratio 0.5
+  --lr_warmup_steps 100
+  --fake_score_learning_rate 1e-5
+  --fake_score_lr_scheduler "cosine_with_min_lr"
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 200
+  --weight_decay 0.01
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --ema_start_step 0
+  --flow_shift 5
+  --seed 1000
+)
+
+# DMD arguments
+dmd_args=(
+  --dmd_denoising_steps '1000,757,522'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --generator_update_interval 5
+  --real_score_guidance_scale 3
+  --simulate_generator_forward 
+  --log_visualization # disable if oom
+  --VSA_sparsity 0.8
+)
+
+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_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}"

examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json

@@ -0,0 +1,516 @@
+{
+  "data": [
+    {
+      "caption": "In the video, a woman is elegantly showcasing her earrings, bringing attention to their intricate design with a gentle touch of her fingers. She is bathed in ambient purple and pink lighting, which casts a soft glow on her delicate features and enhances the vivid tones of her lipstick and eye makeup. Her hair is styled to frame her face smoothly, emphasizing the contours of her jawline and cheekbones. The background features a blurred neon light, adding an artistic and modern touch to the overall aesthetic.",
+      "video_path": "Fashion/mixkit-face-of-an-elegant-and-captivating-woman-41914_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.",
+      "video_path": "Man/mixkit-a-rancher-riding-a-horse-at-sunset-1143_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a dimly lit, eerie setting, a mysterious pink bottle labeled \"Authentic 100% organic POISON\" sits prominently in the foreground, casting a menacing aura. The bottle is accentuated by green fog, which swirls lightly around it, enhancing its sinister allure. Behind it, a shadowy golden bottle adorned with a spider emblem subtly emerges, adding an extra layer of mystery to the scene. Dim candles provide faint, flickering light, which complements the dark atmosphere, making the setting ideal for an illusion of hidden dangers.",
+      "video_path": "smoke/mixkit-poison-in-halloween-ritual-33879_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The video opens with a tranquil scene in the heart of a dense forest, emphasizing two large, textured tree trunks in the foreground framing the view. Sunlight filters through the canopy above, casting intricate patterns of light and shadow on the trees and the ground. Between the tree trunks, a clear view of a calm, muddy river unfolds, its surface shimmering under the gentle sunlight. The riverbank is decorated with a variety of small bushes and vibrant foliage, subtly transitioning into the deep greens of tall, leafy plants. In the background, the dense forest looms, filled with dark, towering trees, their branches intertwining to form an intricate canopy. The scene is bathed in the soft glow of the sun, creating a serene and picturesque setting. Occasional sunbeams pierce through the foliage, adding a magical aura to the landscape. The vibrant reds and oranges of the smaller plants add contrast, bringing warmth to the earthy tones of the scenery. Overall, this harmonious blend of natural elements creates a peaceful and idyllic forest setting.",
+      "video_path": "forest/mixkit-view-of-a-river-between-two-old-trees-560_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a martial artist dressed in a traditional white uniform with a black belt demonstrates a series of precise movements against a stark black background. The individual gracefully transitions between stances, embodying a sense of focused discipline and control. Each motion is executed with a deliberate pace, showcasing the fluidity of martial arts techniques. The soft lighting creates subtle highlights on the uniform, adding depth to the figure as it moves. The practitioner begins with an open-hand pose, feet firmly grounded, gradually shifting to a powerful forward punch. The fluidity of the sequence displays a mastery of balance and poise. Every trajectory of the limbs is precise and deliberate, capturing the elegance and strength of martial arts. The serene, isolated setting enhances the intensity and concentration of the practitioner. This visual presentation is an elegant interplay of motion and stillness, displaying the art form's discipline and grace.",
+      "video_path": "Man/mixkit-a-young-man-practicing-his-karate-moves-49635_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A tranquil coastal scene unfolds with a drone's aerial view capturing a serene beach landscape. The camera glides over a quiet stretch of sandy shoreline, where gentle waves kiss the shore under a clear blue sky. Nestled amidst lush palm trees are a series of traditional thatched-roof huts, their earthy tones blending harmoniously with the natural surroundings. The sandy beach stretches endlessly, bordered by the rhythmic dance of ocean waves on one side and verdant greenery on the other. A pair of white umbrellas is set up on the sand, suggesting a place to relax and enjoy the sun. In the distance, two small human figures can be seen walking leisurely along the water's edge, leaving faint footprints behind them. The scene exudes a calm and inviting atmosphere, with the soft rustle of palm leaves and the whisper of the ocean breeze almost audible. The overall composition is a captivating blend of nature's tranquility and architectural simplicity. This picturesque setting invites viewers to imagine themselves steps away from this idyllic coastal escape.",
+      "video_path": "beach/mixkit-sunny-beach-in-a-dynamic-shot-from-a-drone-44383_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A lone figure stands on a large, moss-covered rock, surrounded by the soft rush of a nearby stream. The figure is wearing white sneakers and shorts, with a plaid shirt that hangs loosely in the breeze. The lighting creates dramatic shadows, enhancing the textures of the rock and the subtle movement of the water below. In the background, a waterfall cascades into the stream, completing this tranquil and serene nature scene.",
+      "video_path": "forest/mixkit-woman-standing-in-front-of-waterfall-559_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In an industrial setting, a person leans casually against a railing, exuding a sense of confidence and composure. They are wearing a striking outfit, consisting of a vibrant, patterned jacket over a simple white crop top, creating a bold contrast. The atmosphere is infused with warm, ambient lighting that casts soft shadows on the concrete walls and metallic surfaces. Intricate wiring and pipes form an intricate backdrop, enhancing the urban aesthetic. Their relaxed posture and direct, engaging gaze suggest a sense of ease in this industrial environment. This scene encapsulates a blend of modern fashion and gritty, urban architecture, creating a visually compelling narrative.",
+      "video_path": "Fashion/mixkit-portrait-of-a-hipster-woman-walking-down-a-stairs-1297_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A man is energetically stretching in an open-air setting, surrounded by rows of vibrant red seats that suggest an amphitheater or outdoor venue. He wears a sleeveless black shirt layered with a hooded vest, emphasizing his athletic build as he engages in a warm-up routine. Behind him, the striking modern architecture of the building features geometric panels, with large sections of glass and overlapping metallic beams creating a dynamic backdrop. The scene captures the contrast between his focused movements and the static, bold design of the structure, while the surrounding greenery adds a touch of nature to the environment. The overall atmosphere is one of preparation and anticipation, with the man appearing determined and ready for an upcoming event or performance.",
+      "video_path": "Sport/mixkit-man-doing-arm-stretches-595_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A young woman is seated on the floor in front of a plush, beige tufted couch, fully engrossed in sorting through a stack of papers. Her dark hair falls loosely past her shoulders, and she wears a green plaid shirt, contributing to the casual yet focused atmosphere. She gently places the papers onto a small round white table, occasionally lifting individual sheets to examine them more closely. Her expression shifts subtly, reflecting concentration and contemplation as she processes the information on the pages. Two small, round nested tables hold her documents, along with a small plant in a gray pot, adding a touch of greenery to the scene. The background features a dark paneled wall, creating a contrasting backdrop for the light-colored furniture. The setting is tranquil and organized, the couch and tables arranged symmetrically, conveying a sense of harmony. A calculator rests on the smaller table, hinting at a task involving calculations or budgeting.",
+      "video_path": "Woman/mixkit-frustrated-woman-throws-paperwork-on-the-floor-4526_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A heavily rusted metal gate stands firmly locked, with two vertical bars joined by a thick, old chain that loops elegantly around them. The chain's texture is coarse and rugged, its surface reflecting varying shades of orange and brown, indicative of years exposed to the elements. At the heart of the chain, a black iron padlock, slightly worn yet imposing, secures the gate, its curves and edges smooth against the aged links. The gate's metalwork is outlined by a backdrop of soft, blurred greenery, suggesting a serene and isolated location beyond the barrier. Tall trees rise in the distance, their trunks and leaves creating a lush, forest-like setting that contrasts with the gate's severe rust. A pathway leads away from the gate, its surface uneven with patches of moss and weathered stone visible in the soft focus, inviting yet inaccessible. The ambiance is quiet and mysterious, with a sense of abandonment hanging subtly in the air, evoking curiosity about what lies beyond. Shadows play across the gate, cast by branches swaying gently in the breeze, adding to the dynamic interaction of light and texture. This scene, rich in detail and atmosphere, captures the viewer's imagination, evoking both the allure of the forbidden and the beauty of decay.",
+      "video_path": "forest/mixkit-rusty-fence-with-a-chain-of-a-property-in-nature-5294_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a serene and softly lit yoga studio, three individuals engage in a yoga session, each performing an upward-facing stretch. The central figure is a woman with shoulder-length brown hair, dressed in a light cropped top and green leggings, her posture reflecting grace and concentration. To her right, another participant, a woman in a purple outfit, mirrors the pose with equal poise. On her left, a person with a bun focuses intently, supported slightly by yoga blocks beneath their hands. The warm-colored wooden floor contrasts soothingly with the soft pastel mural on the back wall, featuring an abstract design and partial visage of a serene face. Natural light floods the space from a large window on the right, where lush greens peek through, adding an element of tranquility. In the corner of the room, a collection of meditation instruments, including a gong and a Buddha statue, subtly frame the peaceful setting. The mood is calm yet focused, as all three participants are deeply engaged in their practice. The scene combines elements of balance, harmony, and a shared journey towards mindfulness. This depiction captures the essence of a yoga session that blends personal growth with collective experience.",
+      "video_path": "People/mixkit-small-group-of-people-doing-yoga-together-43730_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the deep blue expanse of the ocean, two dolphins glide effortlessly, their sleek bodies reflecting the sunlight filtering through the water. The prominent shadows and caustics create a shimmering effect on their skin, capturing the beauty of their natural habitat. Each dolphin moves with a fluid grace, occasionally interacting with gentle nudges, showcasing their playful and social nature. The scene is vibrant and dynamic, with the clear blue background accentuating the dolphins' movements, making it an ideal subject for AI recreation.",
+      "video_path": "sea/mixkit-dolphins-underwater-4133_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a young woman stands against a vibrant graffiti-covered wall, deeply engrossed in her smartphone. Her expression reflects a mix of focus and subtle satisfaction as she interacts with the screen. She wears a black floral-patterned top, which contrasts with the bright, abstract shapes and bold colors of the mural behind her. As she continues to engage with her phone, a series of like count notifications appear on the screen, indicating a growing online appreciation. The wall behind her features a striking mix of geometric and organic shapes, including swirls of teal, orange, and black, with large humanoid figures in a pop-art style. Her long, light-brown hair frames her face, adding a calm, composed aura amidst the lively backdrop. The video captures a blend of contemporary digital interaction and expressive urban art, creating a dynamic yet harmonious scene.",
+      "video_path": "Girl/mixkit-girl-looking-at-the-likes-in-her-post-4914_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A young mother and her baby sit comfortably on a bed, surrounded by an inviting, cozy atmosphere. The woman, wearing a sleeveless top and jeans, is gently engaging with the baby, who is dressed in an adorable animal-print onesie. The child is seated on the bed with colorful toys scattered around, including a plush toy and a board book. The warm glow from a hanging lamp casts a soft light on them, enhancing the serene environment. Pillows are propped up against the headboard, providing a cushioned backdrop as the mother leans slightly over to interact with the baby. A small bottle is visible beside her, suggesting a nurturing setting. Her hand gestures animatedly as she holds up a soft, white cushion with red and blue accents, likely stimulating the baby\u2019s curiosity. Their shared moment is filled with affection and joy, a perfect snapshot of familial bonding.",
+      "video_path": "Baby/mixkit-loving-mother-and-her-baby-playing-with-soft-toys-49966_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A young girl with long brown hair sits at a round wooden table, engrossed in working on her laptop. The laptop screen is a vivid green, suggesting a green screen effect is in use. To her left, a doll dressed in a yellow and white outfit is casually laid on top of some books, adding a playful and innocent touch to the scene. The setting is cozy, with sheer curtains in the background allowing soft natural light to spill into the room. The girl's posture and focused attention on the laptop suggest she is either playing a game or learning something new. This serene and domestic atmosphere is complemented by the slight blur of a dark couch in the foreground, framing the focused activity of the child.",
+      "video_path": "Girl/mixkit-little-girl-doing-homework-on-a-laptop-4757_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "An expansive view of a calm bay reveals a fleet of sailboats, each anchored in a regimented line stretching toward the horizon. The water is a serene blue, reflecting the soft hues of the early morning sky. A gentle breeze is indicated by the subtle ripples trailing behind the boats, while a single, larger vessel cuts a distinct path, leaving a graceful wake in its journey to the open sea. On one side, a cluster of modern high-rise buildings stands, contrasting against the natural simplicity of the water, suggesting a blend of urban and marine life. The distant shoreline is barely visible, softened by the atmospheric perspective, giving a sense of endless waters meeting the sky. The overall mood is peaceful and orderly, with the boats appearing almost as sentinels guarding the expanse of the tranquil bay.",
+      "video_path": "beach/mixkit-flying-backwards-over-the-sea-near-a-coast-50187_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a person is standing in the center of a dark, featureless space, illuminated by a spotlight that emphasizes their presence. The individual is dressed in a traditional martial arts uniform, known as a gi, which is predominantly white with a black belt tied around the waist, indicating a high level of expertise. The background remains pitch black, creating a stark contrast with the brightly lit figure, ensuring complete focus on them. The person's expression is serious and focused, reflecting a deep sense of discipline and concentration. Their hands move gracefully, transitioning through various martial arts stances, demonstrating practiced skill and fluidity. The uniform's crisp fabric folds and subtly reflects the light, further highlighting each precise movement. Despite the simplicity of the environment, the scene is dynamic, with each motion capturing the essence of martial arts practice. The video effectively conveys a sense of calm strength and mastery, making it ideal for an AI to recreate with attention to posture, lighting, and attire.",
+      "video_path": "Sport/mixkit-karate-fighter-bowing-to-the-front-49706_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a dimly lit room bathed in a mix of neon purple and blue lights, a focused individual is seated in a gaming chair. She wears a white hoodie and large headphones with cat ears that glow softly, creating a striking silhouette. Her hands rest on a keyboard, typing swiftly as she concentrates intently on the screen in front of her. The atmosphere exudes a sense of intensity and immersion, with the soft-colored lighting enhancing the futuristic vibe. Her long hair cascades down her shoulders, adding a touch of elegance to the otherwise tech-centric setting. The overall scene captures the essence of a dedicated gamer deeply engaged in her virtual world.",
+      "video_path": "earth/mixkit-a-young-woman-wearing-headphones-with-rgb-lights-suddenly-gets-51621_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "Inside a dimly-lit bus, five individuals are seated along the rows of worn seats, each subtly illuminated by the colorful lights emanating from overhead. On the left, a woman sits with a relaxed posture, her curly hair accented by a patterned scarf, wearing a plaid outfit paired with bright neon socks. Next to her, a person clad in a denim jacket appears deep in thought, resting their head on a hand. Further back, another figure in a bucket hat and oversized yellow attire gazes across the aisle, evoking a sense of introspection. The atmosphere is enriched by the soft glow of red and green lights, bathing the bus interior in an almost surreal ambiance, creating a compelling tableau of urban life.",
+      "video_path": "Music/mixkit-conceptual-urban-fashion-42581_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "An aerial view captures two tennis players on a court, with one dressed in white on the left and another in red on the right. They are mid-game, each poised for action with rackets in hand, accentuated by their strategic positioning at opposite baselines. The court itself is a stark, deep blue, bordered by the vibrant green of the surrounding area, with a dark central net dividing the space. Long shadows stretch dramatically across the ground, suggesting a late afternoon setting. The subtly textured surface of the court contrasts with the crisp, white lines marking its boundaries and sections. This scene creates a vivid, balanced composition, highlighting both the competitive tension and serene atmosphere of the game.",
+      "video_path": "People/mixkit-two-people-playing-tennis-aerial-view-880_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a vibrant, dreamlike setting, a lone figure moves energetically against a backdrop of deep blue and purple hues, casting emotive shadows that ripple with dynamic motion. The figure, almost obscured by a smeared effect, suggests a rhythmic dance or a passionate performance, arms blurred as they sweep through colorful, streaked lighting. A neon glow accentuates their form, particularly highlighting the face which is abstractly illuminated in bursts of orange and red, suggesting intense emotional expression. The scene is dominated by two primary elements \u2013 the figure\u2019s motion and the dramatic lighting, creating a synergy of human emotion and visual spectacle. Swirling trails of light seem to intertwine with the figure, like a visual symphony of movement and color that floods the space. The lighting changes, casting intricate patterns on the figure and the surrounding space, giving the impression of a kaleidoscope in motion. Despite the blurred and abstract portrayal, there is a sense of focus conveyed through the figure\u2019s intent movements, akin to a conductor orchestrating a visual and auditory performance. The environment resonates with an electric energy, suggesting a seamless fusion of art and technology. As the visual drama unfolds, the scene invites viewers to lose themselves in the abstract dance and the play of vivid luminance.",
+      "video_path": "Music/mixkit-dancer-dancing-with-a-light-bar-in-his-hands-42221_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a brightly lit studio, a photographer wearing a denim jacket focuses intently, capturing shots with a professional camera. Facing him, a model stands gracefully, adjusting her long, flowing hair with delicate movements. The scene is characterized by strong contrasts; the model's soft pink attire and gentle gestures complement the rugged, precise demeanor of the photographer. Positioned against a minimalist backdrop, the pair work seamlessly, with the camera\u2019s lens pointed directly at the model, capturing her elegance. The soft, diffused lighting casts a gentle glow on both subjects, creating an airy and ethereal atmosphere perfect for a high-fashion photo shoot.",
+      "video_path": "Fashion/mixkit-professional-photo-session-with-a-young-female-model-41621_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The video showcases a serene, expansive landscape covered with a variety of trees dotting the hills. The hills gently slope across the frame, with patches of dry grass contrasting against the lush green foliage. Tall trees with dense canopies stand elegantly, casting soft shadows on the ground below. The sunlight bathes the entire scene, highlighting the varied textures of the leaves and terrain. Gaps between the trees reveal a narrow dirt path meandering through the hills, suggesting a sense of quiet solitude. The undulating hills extend into the distance, creating depth and a calming sense of vast space. The verdant hues of the leaves contrast with the earthy tones of the hills, enhancing the visual richness. In the background, a faint outline of distant hills can be seen, blurred softly by the atmospheric perspective. This tranquil setting could be efficiently recreated in a virtual environment by focusing on its layered composition, color palette, and natural textures.",
+      "video_path": "forest/mixkit-aerial-panorama-of-a-sunny-mountain-landscape-40846_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A bustling ski slope comes alive with skiers descending a pristine, snow-covered hill, surrounded by towering, snow-draped evergreens. Several figures stand atop the slope, silhouetted against a clear blue sky, preparing to embark on their ski run. The chair lift on the right continuously drops off eager adventurers, adding to the excitement at the hilltop. Each skier, clad in colorful winter gear, carves distinct paths into the textured snow as they weave their way down. The interplay of sunlight and shadows accentuates the myriad tracks etched into the slope, creating a dynamic visual rhythm. The scene captures a vibrant winter wonderland, full of action and the thrill of a perfect ski day.",
+      "video_path": "Car/mixkit-skiers-on-a-snowy-slope-3327_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The scene unfolds within a dimly lit bus, where three young individuals are seated, each absorbed in their unique world. To the left, a person with tied-back hair rests their head on their hand, dressed casually in a jacket and jeans, projecting a relaxed demeanor. Central to the frame is another individual, sitting upright with intense focus, donning a plaid blazer and oversize hoops, enhancing their confident presence. The muted green and red lighting casts an atmospheric glow, adding depth and intrigue to the setting. On the right, a person in a bucket hat and striped shirt leans back, appearing contemplative as they adjust their hat with a nonchalant gesture. The interplay of light and shadow highlights their expressions, creating an intimate and cinematic ambiance. Together, these figures form a cohesive tableau, capturing a moment of introspection amid a bustling yet serene urban environment.",
+      "video_path": "City/mixkit-three-models-posing-to-the-lens-while-on-board-a-42575_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A determined climber is scaling a massive rock face, showcasing exceptional strength and skill. The person, clad in a teal shirt and dark pants, climbs with precision, their movements measured and deliberate. They are secured by climbing gear, which includes ropes and a harness, emphasizing their commitment to safety. The rugged texture of the sandy-colored rock provides an imposing backdrop, adding drama and scale to the climb. In the distance, other large rock formations and sparse vegetation can be seen under a bright, overcast sky, contributing to the natural and adventurous atmosphere. The scene captures a moment of focus and challenge, highlighting the climber's tenacity and the breathtaking environment.",
+      "video_path": "Sport/mixkit-alpinist-climbing-a-huge-rock-in-a-desert-43306_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A woman stands confidently in front of a large array of solar panels, her navy blue jumpsuit contrasting against the lush green grass beneath her feet. Her expression is calm and focused, eyes facing directly ahead, suggesting a deep connection to the subject matter\u2014renewable energy. The sunlight bathes the scene in warm hues, casting gentle shadows and highlighting the geometric precision of the solar panels' grid-like structure. The background reveals a blend of nature and technology, as the panels are anchored on a grassy slope with foliage on the left side of the frame. This composition captures a harmonious blend of human innovation and environmental consciousness, accentuated by the serene outdoor setting.",
+      "video_path": "Business/mixkit-woman-standing-in-front-of-a-solar-panel-4880_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, two people are working at a wooden desk, using an iMac computer. One person, wearing a white knit sweater, is using the apple wireless mouse with their right hand, while their left hand rests on the sleek white keyboard. Their movements are smooth yet intentional, suggesting they are focused on a task on the computer screen. The monitor displays a well-organized array of files and folders, hinting at a task that involves detailed organization or detailed data navigation. The second person, only subtly visible, sits closely by and appears to observe or assist, creating a collaborative atmosphere. Their presence adds a quiet dynamic to the scene, as if they are ready to provide input or guidance. Sticky notes with handwritten notes are attached to the monitor\u2019s stand, adding a touch of personal organization amidst the digital workspace. The focus on the keyboard and mouse emphasizes a streamlined workflow, indicative of a productive work environment. The overall ambiance is calm and focuses on teamwork, technology, and efficient workspace management.",
+      "video_path": "People/mixkit-person-with-glasses-working-on-a-desktop-computer-3248_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A man stands in front of a modern glass facade, taking off a dark hoodie to reveal his gray tank top underneath. His arms are lifted high as he maneuvers the hoodie over his head, showcasing a fluid motion that conveys a sense of calm and routine. The lighting highlights the contours of his muscles, emphasizing a combination of strength and quiet determination. Behind him, the reflective surface of the glass panels provides a subtle backdrop, enhancing the focus on his focused and serene demeanor.",
+      "video_path": "Sport/mixkit-man-puts-on-sleeveless-hoodie-603_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The video displays a captivating dance of fiery orange flames against a stark black background, creating an intense visual contrast. The flames twist and intertwine, forming symmetrical, swirling patterns that expand and contract rhythmically across the frame. Each fiery tendril seems to be alive, moving with an almost hypnotic fluidity that captures the viewer's attention. The illumination from the flames casts subtle shadows, enhancing the depth and texture of the scene. Overall, the dynamic movement and vibrant color palette create an atmosphere of both beauty and power.",
+      "video_path": "fire/mixkit-two-orange-flames-on-black-background-685_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In this scene, a person is seated in a dimly lit room, possibly a recording studio, holding several drumsticks in their hands. The individual's face is partially obscured by sunglasses, adding a touch of mystery to their demeanor. They are wearing a colorful, patterned shirt with a mix of orange and blue tones that stands out against the darker background. The person appears focused and engaged with the drumsticks, their hands prominently displayed. The ambient light casts warm, soft shadows, emphasizing the texture and colors of their shirt and the wooden drumsticks. The room features wooden paneling, which complements the overall cozy, music-centric setting of the scene. The use of perspective centers on the drumsticks, highlighting the importance of rhythm and music in the captured moment.",
+      "video_path": "Music/mixkit-drummer-stretching-before-playing-42783_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A man is casually sitting on a sofa, engrossed in his meal and entertainment. He is holding a TV remote in one hand while reaching for food with the other, indicating a laid-back, comfortable evening. The table before him is filled with takeout containers, revealing a variety of appetizers and dishes, suggestive of a casual dining experience at home. The background is defined by colorful patterned cushions, adding a cozy, homey feel to the scene. Warm, ambient lighting highlights the relaxed atmosphere, casting soft shadows that contribute to the intimate setting. In this moment, he takes a bite of a sandwich, comfortably balancing his attention between food and whatever is playing on the screen.",
+      "video_path": "Man/mixkit-man-watching-tv-and-eating-fast-food-26089_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The scene opens to a breathtaking view of a tranquil ocean horizon at dusk, displaying a vibrant tapestry of oranges, pinks, and purples as the sun sets. In the foreground, tall, swaying palm trees frame the scene, their silhouettes stark against the colorful sky. The ocean itself shimmers with reflections of the sunset, creating a peaceful, almost ethereal atmosphere. A small boat can be seen in the distance, centered on the horizon, adding a sense of scale and solitude to the scene. The waves gently lap the shore, creating faint patterns on the sandy beach, which stretches across the foreground. Above, the sky is dotted with scattered clouds that catch the last light of the day, enhancing the drama and beauty of the scene. The overall mood is serene and contemplative, capturing a perfect moment of nature\u2019s grandeur.",
+      "video_path": "beach/mixkit-sunset-with-sailing-boats-2166_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A man sits hunched on a couch, the weight of emotions clearly visible on his posture. He wears a simple, gray t-shirt, and his head is bowed, resting in his hands, which cover most of his face, obscuring his features. The gentle light filtering through sheer curtains in the background casts a soft glow upon him, emphasizing the contrast between his static form and the hazy brightness behind. His elbows rest upon his knees, suggesting a posture of deep contemplation or distress. The simplicity of the room, with its muted colors, highlights the focus on the man's internal struggle. Delicate detailing on the fabric of his shirt adds texture, enhancing the scene's realism. Subtle changes in the natural light indicate the passage of time, as the man remains unmoving, absorbed in thought. This intimate moment captures a profound vulnerability, making the scene universally relatable and poignant.",
+      "video_path": "Man/mixkit-worried-and-sad-man-with-his-head-down-4701_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A pair of hands, belonging to an unseen figure, carefully unrolls a large sheet of crisp, white paper on a dark wooden table. The lighting is warm, casting a gentle glow that highlights the textures of the paper and the wood grain of the table. As the paper unfurls, the edges reveal the faint beginnings of a colorful map printed on its surface. The arms, clad in a casual gray T-shirt, suggest a relaxed and focused task at hand. Each motion is deliberate, with fingers deftly guiding the paper, ensuring it lays flat without creases. In the background, a hint of a red curtain can be seen, adding a touch of color and depth to the setting. The composition of the scene emphasizes the contrast between the bright paper and the rich tones of the surroundings. This serene and methodical action evokes a sense of exploration and preparation.",
+      "video_path": "Man/mixkit-unrolling-a-world-map-on-a-table-21626_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A young woman sits on a vibrant green seat inside a bus, illuminated by the soft glow of pink and blue lights. Her outfit is a striking mix of colors: a neon pink top paired with a jacket featuring dark sleeves, and jeans that provide a neutral contrast. She wears large, hoop earrings that catch the light as she moves slightly, exuding an air of cool confidence. Her gaze is directed thoughtfully to the side, suggesting contemplation or daydreaming during her commute. The metallic pole beside her adds a geometric element to the composition, reflecting the kaleidoscope of neon hues. The background is a clean, futuristic white, serving as a blank canvas that amplifies the neon atmosphere. Her relaxed posture and the modern bus setting create a scene that captures a blend of urban life and personal introspection.",
+      "video_path": "City/mixkit-fashion-model-posing-on-a-bus-42578_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A silver SUV drives along a winding, snow-covered mountain road, with dense pine trees blanketed in snow lining both sides. The scene is serene, with the vehicle moving smoothly, possibly on a winter journey or vacation. As the SUV disappears around the bend, another, darker SUV follows, creating a sense of motion and perspective on the snow-dusted asphalt. The towering, snow-laden rock formation to the right contrasts with the dark green of the pines, highlighting the peacefulness of the wintry landscape.",
+      "video_path": "Car/mixkit-curve-on-a-snowy-forest-road-3317_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The video showcases a vibrant urban skyline during twilight, with towering buildings reflecting the warm hues of the setting sun. A series of tall, cylindrical structures dominate the foreground, adjacent to a complex of industrial equipment and grids. The scene includes modern high-rise buildings with glass exteriors, capturing the evolving architecture of a bustling cityscape. A prominent structure labeled \"CITY OF AUSTIN POWER PLANT\" stands out, highlighting the industrial theme amidst the urban backdrop. The soft glow of city lights begins to pierce the approaching dusk, creating an inviting yet dynamic atmosphere. Shadows cast by the buildings add depth and contrast, emphasizing their massive scale and intricate designs. The overall composition is balanced between the natural light of the sunset and the artificial illumination of the city, offering a compelling visual narrative.",
+      "video_path": "Car/mixkit-slow-air-travel-in-reverse-over-a-big-city-49841_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the scene, a striking architectural structure dominates the view, bathed in a soft, ambient light. The enormous yellow arches serve as the centerpiece, drawing the eye upwards with their majestic curves and towering presence. The smooth, clean surfaces of the structure reflect the light, highlighting the texture and depth of the architecture. In the foreground, blurred streaks of headlights and taillights suggest the motion of vehicles passing by, adding dynamic energy to the otherwise still scene. The contrast between the fast-moving lights and the static arches creates a balanced composition. To the left, a lone streetlamp and a small tree provide a touch of nature and urban elements against the monumental backdrop. The night sky subtly peeks through the gaps in the structure, hinting at a clear, calm evening. Shadows from the arches create patterns on the ground, adding an intricate detail to the scene. Overall, the combination of light, shadow, and movement makes for a dramatic and visually captivating moment.",
+      "video_path": "Car/mixkit-a-fast-timelapse-of-the-street-with-a-monumental-yellow-50993_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A tranquil marina comes into full view under the golden hues of a setting sun. A collection of gleaming yachts and boats are neatly moored, their reflections shimmering softly on the gentle water. The sun's low position casts elongated shadows over the bustling harbor scene, while rolling hillsides surround the distant cityscape. The skyline is interspersed with modern buildings and clusters of residences, adding layers to the vibrant community. At the center, a broad wooden pier juts confidently into the harbor, extending an invitation for leisurely strolls. To the left, various shops and colorful structures line the waterfront, indicating a vibrant coastal economy. The entire atmosphere exudes a serene yet lively charm, balancing the hustle of maritime activity with the peacefulness of the encroaching dusk. It's a scene of calm anticipation, as if the whole place holds its breath before the night's events unfold.",
+      "video_path": "beach/mixkit-harbor-on-a-tourist-coast-with-many-boats-and-yachts-40077_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The video features a confident individual standing atop a structure against a clear blue sky, exuding a sense of freedom and style. The person is clad in a striking yellow button-up shirt tied at the waist, and beneath it, they wear a simple white top that adds to their relaxed yet stylish appearance. Completing the ensemble are high-waisted white jeans paired with a black belt, adding a touch of contrast. Around their neck is a bold red scarf, providing a splash of color and an air of vintage flair. The person's sunglasses, tinted in yellow, reflect the sunlight and contribute to the overall cool and composed demeanor. Their hair is styled elegantly, pulled back with headphones resting over the ears, suggesting they are immersed in music. One hand casually grazes the headphones, while the other rests gently on the railing, grounding the individual in the moment. The scene is an effortless blend of fashion and tranquility, capturing the spirit of sunny, carefree days.",
+      "video_path": "Music/mixkit-standing-woman-listening-to-music-460_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A ballerina gracefully spins and moves across a pink-hued studio, her poised figure accentuated by a shimmering white tutu and bodice. The background, a continuous wash of soft pink, provides a serene and ethereal atmosphere, emphasizing her fluid movements. Her arms extend with elegance, highlighting the delicacy and precision of her ballet pose, while her focused expression adds intensity to the scene. The subtle details of her costume, combined with the pink monochromatic ambiance, create a dreamlike spectacle, ideal for an AI to envision a oneiric dance setting.",
+      "video_path": "Dance/mixkit-portrait-of-a-ballerina-spinning-with-pink-background-40163_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "The scene unfolds with two human figures in the distance, making their way through a serene meadow, thick with tall golden grass swaying gently in the breeze. The sun hangs low in the sky, casting a soft, diffused glow that illuminates the landscape with a warm, ethereal light. These figures, clad in hiking gear, move deliberately, suggesting they're either embarking on or concluding a journey. Their silhouettes contrast against the lush greenery of the surrounding trees, whose branches reach out, framing the horizon. The play of light and shadow among the trees creates a quilt of textures, with each leaf catching a hint of the sun's dying rays. This tranquil setting evokes a sense of calm and adventure, capturing the quintessential beauty of nature\u2019s landscape.",
+      "video_path": "People/mixkit-landscape-in-nature-while-two-people-are-jogging-44348_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A large cargo ship is docked at an industrial port, its white superstructure contrasting with the deep green and yellow of its deck. The foreground is dominated by the calm, deep blue waters of the harbor, which reflect the vessel\u2019s imposing presence. Surrounding the ship, a series of industrial buildings and storage facilities are visible, hinting at the bustling activity of the port. The deck is intricately detailed, featuring an array of pipes, equipment, and railings, showcasing the ship's functionality and purpose. In the background, a paved area with green patches and a few parked vehicles adds to the busy, industrious atmosphere of the scene.",
+      "video_path": "sea/mixkit-empty-cargo-ship-waiting-at-the-port-4209_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A lone climber ascends a towering rock face, clad in a pink shirt and gray pants, displaying a determined and focused expression. The climber navigates the rugged surface, where the texture of the rock is peppered with natural pockets and crevices that offer handholds and footholds. Sunlight casts soft shadows across the cliff, highlighting the intricate patterns and the climber\u2019s strategic movements. The cliff looms high, with sparse vegetation breaking the monotony of the stone, while distant rocky formations form a dramatic backdrop against the clear blue sky. The climber\u2019s gear, including a harness and chalk bag, underscores the adventure and challenge woven into this majestic, vertical journey.",
+      "video_path": "Sport/mixkit-mountaineer-girl-climbing-a-steep-rocky-mountain-41089_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A person is seen in a close-up shot, skillfully adjusting the tuning pegs of a guitar, showcasing a focused and practiced hand. The image is in black and white, highlighting the contrast between the textures of the instrument and the clothing. The individual's shirt, visible in the background, adds a soft, subtle texture, while the dark tones of the guitar neck create depth in the scene. This composition captures a moment of concentration and finesse, perfect for recreating an intimate musical setting.",
+      "video_path": "Music/mixkit-guitarist-playing-so-inspired-black-and-white-shot-44178_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A musician is playing a large brass instrument with the words \"Brass Band\" clearly visible on its bell. The scene is set against a vibrant yellow backdrop, casting a warm glow on the subject. The musician wears a dark cap and a matching suit, adding a formal touch to his attire. He is deeply focused on his performance, with the instrument's intricate tubing adding complexity to the visual composition. The lighting creates dramatic shadows and highlights, emphasizing the musician's expression and the instrument's metallic sheen. This harmonious blend of color and form captures the essence of a live brass band performance.",
+      "video_path": "Music/mixkit-musician-playing-the-trombone-while-dancing-43752_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a lone musician stands gracefully in front of a grand cathedral, playing an accordion while surrounded by the lively water display of a central fountain. Dressed in a casual ensemble, he wears a light-colored shirt, dark pants, and a flat cap that gives him a vintage charm. His posture is relaxed, yet engaged, as he sways gently in rhythm with the music, casting soft shadows on the cobblestone steps beneath him. The backdrop features the cathedral's towering twin spires, with intricate stonework that casts a rich, historical aura around the scene. Sunlight bathes the entire setting, enhancing the golden hues of the cathedral facade and creating a halo-like effect around the musician. The fountain's water jets splash playfully, catching glimmers of light and adding a dynamic element to the tranquil atmosphere. The scene captures a harmonious blend of architectural majesty and human creativity, framed by the clear, azure sky that extends infinitely above. It's a vivid depiction of solitude and artistry, set against a timeless urban landscape.",
+      "video_path": "Music/mixkit-man-plays-an-accordion-in-front-of-a-fountain-630_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the tranquil video, a person sits in a meditative pose on a gentle hillside, silhouetted against the dawning sky. The person is facing the breathtaking sunrise, with their back slightly turned to the viewer, wearing a simple, light-colored shirt. Their right hand rests on their knee, fingers relaxed in a common meditation mudra, symbolizing calmness and peace. The sky, a stunning blend of soft oranges and deep purples, gradually brightens, casting a warm glow over the lush, green landscape. To the left, the outlines of distant urban buildings can be seen against the horizon, adding a contrast between nature and city life. A river reflecting the sky's colors meanders through the scene, lending a serene, flowing dynamic to the landscape. Trees rise and fall gently across the terrain, their leaves rustling only faintly in the morning breeze. The person remains still and focused, embodying a moment of mindfulness and connection with nature. This visual captures a harmonious balance, evoking a sense of tranquility and introspection.",
+      "video_path": "City/mixkit-girl-meditating-in-yoga-pose-at-sunset-4803_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A serene landscape video captures a breathtaking panoramic view of a vast valley covered in a gentle mist. The undulating hills are lush with dense greenery, their rich foliage creating a vibrant border on the left side of the frame. The mist weaves through the landscape like a soft, ethereal blanket, lending a dream-like quality to the scene. In the distance, several mountain peaks emerge, their dark outlines contrasting against the pale blue sky. A few faint, wispy clouds drift lazily across the horizon, complementing the tranquil atmosphere. The sunlight filters through the haze, casting a warm glow and highlighting different textures of the flora. The overall mood is calm and contemplative, inviting the viewer to pause and appreciate nature's untouched beauty. The composition emphasizes depth and expansiveness, drawing attention to the harmony between earth and sky. This captivating scene embodies tranquility, offering a perfect backdrop for meditation or relaxation.",
+      "video_path": "forest/mixkit-flying-over-a-hill-with-a-view-of-the-surrounding-49743_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In this scene, a bearded individual is intently focused on their smartphone, with the sun setting in the background, casting a warm glow across the cityscape. The person, partially visible, is wearing a dark, buttoned shirt that contrasts with the golden hue of the sunset. Their hands are holding the smartphone delicately but purposefully, reflecting a sense of engagement and focus on the screen. The sunlight creates a striking lens flare effect, enhancing the dramatic atmosphere of the moment as it glimmers off the phone\u2019s surface. The surrounding environment hints at an elevated vantage point, providing a panoramic view of the urban landscape below.",
+      "video_path": "City/mixkit-guy-texting-at-sunset-265_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In an expansive, industrial space defined by towering columns and high ceilings, a solitary figure takes center stage. The person, dressed in dark, fitted clothing, assumes a powerful, dynamic stance with one leg bent forward and both arms outstretched in a horizontal arc. Framing this pose are intense flames that engulf their arms, creating a striking visual contrast against the muted tones of the room. The fire forms a brilliant halo of orange and yellow, casting flickering shadows on the weathered walls and worn, tiled floor. This interplay between light and dark showcases the dancer's poise and agility, as they maintain balance amidst the intense heat. Windows line the background, their panes dimly illuminated by the daylight filtering in, adding depth and perspective to the scene. The entire performance evokes a sense of raw energy and elemental mastery, as the figure continues to manipulate the fire in a seamless, mesmerizing display.",
+      "video_path": "fire/mixkit-expert-juggler-doing-tricks-with-a-stick-with-fire-43663_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A man is playing the violin, focused intently on his music. His fingers gracefully dance along the strings, flawlessly executing each note. He holds the violin close to his chin with a sense of familiarity and expertise. The rich, warm tones of the violin reflect in the soft lighting of the room. He wears a dark shirt, and a subtle necklace rests against his chest, adding a personal touch to his attire. The bow moves smoothly across the strings, producing a melody that seems to fill the space with emotion. His expression is one of concentration and passion, immersing himself fully in the performance. The background is softly blurred, bringing the violin's intricate craftsmanship and his precise movements into sharp focus. This serene and intimate moment captures the essence of his musical artistry.",
+      "video_path": "Music/mixkit-fiddler-playing-a-song-639_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the dimly lit parking garage, two figures engage in an impromptu game of soccer. The first person, wearing a light grey shirt and black pants with three white stripes, skillfully maneuvers the ball with precise footwork. The ground is slick with patches of water, reflecting the vibrant neon lights above. A second figure, clad in dark clothing, stands poised in the background, ready to intercept. The space is defined by stark yellow lines and orange safety bollards, adding structure to the chaotic energy of the scene. The soccer ball glides smoothly across the wet floor, kicking up droplets as it passes. Despite the muted colors of the environment, the players' movements are dynamic and full of life. Their shadowy silhouettes dance with the reflecting light, creating a mesmerizing visual interplay. The atmosphere is charged with focus and camaraderie, encapsulating the essence of a late-night urban soccer experience.",
+      "video_path": "Sport/mixkit-player-making-skillful-play-in-a-street-soccer-game-43504_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A lone climber is seen scaling a towering vertical rock face, demonstrating remarkable strength and focus. Dressed in a light-colored shirt and jeans, the climber grips the stone tightly, navigating the rough textures and crevices with precision. The sheer cliff is massive, exhibiting a range of natural hues from light tan to deep gray, accentuating the climber's figure against the vast rocky backdrop. Surrounding the cliff, scattered greenery and rugged terrain provide a sense of wilderness and isolation. The scene portrays a daring ascension requiring concentration and skill, capturing the essence of human endeavor against nature's formidable beauty.",
+      "video_path": "Sport/mixkit-skilled-mountaineer-climbing-a-gigantic-mountain-41083_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In this serene landscape, a lush meadow stretches across the foreground, dotted with vibrant yellow wildflowers swaying gently in the breeze. A towering tree stands majestically on the right side, its branches reaching wide under the bright blue sky filled with fluffy white clouds. On the left, dense trees form a natural corridor leading to the horizon, suggesting a sense of journey and possibility. The richness of the green grass contrasts beautifully with the golden hue of the distant fields, creating a harmonious palette of nature\u2019s colors. The play of light and shadow adds depth and dimension, evoking a tranquil, inviting atmosphere. It's a scene where nature\u2019s beauty simply commands attention, offering a perfect escape into tranquility.",
+      "video_path": "sky/mixkit-countryside-meadow-4075_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A solitary boat glides across the expansive, tranquil expanse of a serene lake. The vessel leaves a gentle wake behind, creating delicate ripples across the mirror-like surface. The water appears a rich shade of teal, seamlessly blending with the sky at the horizon. Silhouettes of distant trees are faintly visible, creating a picturesque backdrop that enhances the solitary journey of the boat. The sky is a calm gradient, shifting from soft oranges near the shore to the pale blues above. In the distance, a few slender poles emerge from the water, remnants of an old structure or natural formation. The mood of the scene is one of peace and solitude, with the boat journeying steadily through the quiet landscape. There is a sense of endless possibilities as the boat moves toward the unseen beyond the frame. The simplicity and stillness of the scene invite contemplation and reflection, encapsulating a perfect moment of quietude on the water.",
+      "video_path": "mountain/mixkit-motorboat-on-a-large-lake-with-turquoise-blue-waters-4996_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a cozy, dimly lit caf\u00e9, a woman sits alone at a rustic wooden table, fully engrossed in her reading. Her dark, wavy hair frames her face as she leans forward over an open book, suggesting deep focus and contemplation. The caf\u00e9\u2019s ambiance is warm, with hanging pendant lights casting a soft glow over the wooden shelves lined with jars and coffee paraphernalia in the background. A small cup of coffee rests just within her reach, alongside a glass dome encasing a solitary pastry, adding a touch of tranquility to the scene. Her casual attire, a denim jacket over a simple shirt, complements the laid-back, comfortable setting of the caf\u00e9. The contrast between her concentrated expression and the bustling, yet subdued caf\u00e9 atmosphere creates a harmonious, serene visual. The overall composition captures a quiet moment of introspection amidst the gentle hum of caf\u00e9 life.",
+      "video_path": "Woman/mixkit-woman-drinking-coffee-in-a-cafe-223_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a vast, deserted landscape under the night sky, a solitary figure stands at a small music setup, illuminated by strategically placed lights. The person is engrossed in playing a keyboard, with various electronic equipment surrounding them, casting soft glows of orange and blue hues across the scene. To the left, a large circular light adds a dramatic focal point, highlighting the intense contrast between the darkness and the lit performance area. This setup, with its minimalistic design and strategic lighting, creates a captivating and easily recognizable scene that merges the serene, expansive backdrop with an intimate, focused music performance.",
+      "video_path": "Music/mixkit-talented-dj-playing-in-a-lonely-desert-42414_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In a bustling urban scene, cars zoom past a weathered building, their blurred motion a testament to the city\u2019s lively pace. The building, with its faded yellow and brown facade, boasts graffiti that speaks of both art and decay, framing the scene with an air of urban grit. A solitary figure stands slightly to the side, clad casually in a gray top and mustard trousers, gazing into the street, seemingly detached from the surrounding flurry. The motion of the traffic creates a dynamic contrast against the static backdrop, emphasizing the relentless movement of the city. As the video progresses, a bright yellow taxi appears, slowing down as it approaches the figure, adding a pop of color to the desaturated hues of the environment. The interaction suggests a routine, a possibly daily exchange between the driver and the pedestrian, hinting at the rhythms of city life. Overhead, a soft, overcast sky casts a diffused light, lending the scene a subdued, timeless quality. Small elements, like the vertical pole cutting through the frame and the distant chatter of urban sounds, complete this vivid tableau of urban existence.",
+      "video_path": "Car/mixkit-morning-in-the-street-time-lapse-1648_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "A young woman sits on a curb in a tranquil park, basking in the golden hue of the setting sun. Beside her, a collie dog rests calmly, its fur illuminated by the warm sunlight, creating a serene glow. The woman's hand gently strokes the dog's back, highlighting the bond and affection between them. Tall trees surround the pair, casting elongated shadows on the leaf-laden ground, adding to the peaceful and intimate ambiance of the scene.",
+      "video_path": "Pets/mixkit-a-woman-pets-a-dog-in-a-park-1562_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a grand, majestic elephant stands in an open, sunlit field, its massive form dominating the scene. The elephant's skin is a tapestry of earthy tones, with rough, textured wrinkles that add character to its already imposing presence. Its trunk, a powerful and flexible appendage, moves gently, swaying as the elephant possibly enjoys the warmth of the day. The background is a blur of greenery, suggesting a lively environment filled with trees and shrubs that provide a natural habitat. Light plays on the elephant's skin, highlighting patches of dust and dirt that give it an authentic wilderness look. The scene captures the tranquility and majesty of this gentle giant in its natural surroundings.",
+      "video_path": "Zoo/mixkit-wet-elephant-in-the-savanna-3663_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    },
+    {
+      "caption": "In the video, a fluffy dog with brown patches is intently engaged with a bright red toy shaped like a fire hydrant, which has a yellow and orange rope attached. The dog's body is relaxed as it lies on a plain white background, concentrating on nudging and playfully biting the toy. Its ears perk up slightly with curiosity, and its eyes are fixated on the toy, suggesting a scene of focused playfulness. The neutral tones of the dog's fur contrast starkly against the vivid red of the toy, creating a visually striking moment.",
+      "video_path": "Pets/mixkit-a-cute-border-collie-dog-play-with-a-fire-street-50662_clip_1.mp4",
+      "num_inference_steps":  3,
+      "height": 448,
+      "width": 832,
+      "num_frames": 61
+    }
+  ]
+}

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

@@ -9,4 +9,14 @@ git submodule update --init --recursive
 python setup_vsa.py install
 ```
 
-### TODO
+### 1. Download dataset:
+```bash
+bash examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/download_dataset.sh
+```
+
+### 2. Configure and run distillation:
+
+#### For DMD-only distillation:
+```bash
+bash examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/
+```

examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh

@@ -63,7 +63,7 @@ validation_args=(
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
   --validation_steps 200
   --validation_sampling_steps "3"
-  --validation_guidance_scale "1.0" # not used for dmd inference
+  --validation_guidance_scale "6.0" # not used for dmd inference
 )
 
 # Optimizer arguments
@@ -77,7 +77,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.0
   --dit_precision "fp32"

examples/inference/basic/basic.py

@@ -16,8 +16,7 @@ def main():
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,
-        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
-        pin_cpu_memory=False,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
         # image_encoder_cpu_offload=False,
     )
 

examples/inference/basic/basic_dmd.py

@@ -17,7 +17,7 @@ def main():
         use_fsdp_inference=True,
         # Adjust these offload parameters if you have < 32GB of VRAM
         text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         VSA_sparsity=0.8,

examples/inference/basic/basic_wan2_2.py

@@ -0,0 +1,48 @@
+from fastvideo import VideoGenerator
+
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_wan2_2_14B_t2v"
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=True, # DiT need to be offloaded for MoE
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
+        pin_cpu_memory=True,
+        # image_encoder_cpu_offload=False,
+    )
+
+    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    # sampling_param.num_frames = 45
+    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    # Generate videos with the same simple API, regardless of GPU count
+    prompt = (
+        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+        "wide with interest. The playful yet serene atmosphere is complemented by soft "
+        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    )
+    _ = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, height=720, width=1280, num_frames=81)
+    # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+    prompt2 = (
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+    _ = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True, height=720, width=1280, num_frames=81)
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_wan2_2_i2v.py

@@ -0,0 +1,30 @@
+from fastvideo import VideoGenerator
+
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_wan2_2_14B_i2v"
+def main():
+    # FastVideo will automatically use the optimal default arguments for the
+    # model.
+    # If a local path is provided, FastVideo will make a best effort
+    # attempt to identify the optimal arguments.
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+        # FastVideo will automatically handle distributed setup
+        num_gpus=1,
+        use_fsdp_inference=True,
+        dit_cpu_offload=True, # DiT need to be offloaded for MoE
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
+        pin_cpu_memory=True,
+        # image_encoder_cpu_offload=False,
+    )
+
+    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, image_path=image_path, output_path=OUTPUT_PATH, save_video=True, height=832, width=480, num_frames=81)
+
+if __name__ == "__main__":
+    main()

examples/inference/gradio/README.md

@@ -1,59 +0,0 @@
-# FastVideo Gradio Demo
-
-This is a Gradio-based web interface for generating videos using the FastVideo framework. The demo allows users to create videos from text prompts with various customization options.
-
-## Overview
-
-The demo uses the FastVideo framework to generate videos based on text prompts. It provides a simple web interface built with Gradio that allows users to:
-
-- Enter text prompts to generate videos
-- Customize video parameters (dimensions, number of frames, etc.)
-- Use negative prompts to guide the generation process
-- Set or randomize seeds for reproducibility
-
----
-
-## Usage
-
-Run the demo with:
-
-```bash
-python examples/inference/gradio/gradio_demo.py
-```
-
-This will start a web server at `http://0.0.0.0:7860` where you can access the interface.
-
----
-
-## Model Initialization
-
-This demo initializes a `VideoGenerator` with the minimum required arguments for inference. Users can seamlessly adjust inference options between generations, including prompts, resolution, video length, or even the number of inference steps, *without ever needing to reload the model*.
-
-## Video Generation
-
-The core functionality is in the `generate_video` function, which:
-1. Processes user inputs
-2. Uses the FastVideo VideoGenerator from earlier to run inference (`generator.generate_video()`)
-3. Returns an output path that Gradio uses to display the generated video
-
-## Gradio Interface
-
-The interface is built with several components:
-- A text input for the prompt
-- A video display for the result
-- Inference options in a collapsible accordion:
-  - Height and width sliders
-  - Number of frames slider
-  - Guidance scale slider
-  - Inference steps slider
-  - Negative prompt options
-  - Seed controls
-
-### Inference Options
-
-- **Height/Width**: Control the resolution of the generated video
-- **Number of Frames**: Set how many frames to generate
-- **Guidance Scale**: Control how closely the generation follows the prompt
-- **Inference Steps**: More steps can improve quality but take longer
-- **Negative Prompt**: Specify what you don't want to see in the video
-- **Seed**: Control randomness for reproducible results

examples/inference/gradio/gradio_demo.py

@@ -1,169 +0,0 @@
-import argparse
-import os
-from copy import deepcopy
-
-import gradio as gr
-import torch
-
-from fastvideo import VideoGenerator
-from fastvideo.configs.sample.base import SamplingParam
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="FastVideo Gradio Demo")
-    parser.add_argument("--model_path",
-                        type=str,
-                        default="FastVideo/FastHunyuan-diffusers",
-                        help="Path to the model")
-    parser.add_argument("--num_gpus",
-                        type=int,
-                        default=1,
-                        help="Number of GPUs to use")
-    parser.add_argument("--output_path",
-                        type=str,
-                        default="outputs",
-                        help="Path to save generated videos")
-    parsed_args = parser.parse_args()
-
-    # args = FastVideoArgs(model_path="FastVideo/FastHunyuan-Diffusers", num_gpus=2)
-
-    generator = VideoGenerator.from_pretrained(
-        model_path=parsed_args.model_path, num_gpus=parsed_args.num_gpus)
-
-    default_params = SamplingParam.from_pretrained(parsed_args.model_path)
-
-    def generate_video(
-        prompt,
-        negative_prompt,
-        use_negative_prompt,
-        seed,
-        guidance_scale,
-        num_frames,
-        height,
-        width,
-        num_inference_steps,
-        randomize_seed=False,
-    ):
-        params = deepcopy(default_params)
-        params.prompt = prompt
-        params.negative_prompt = negative_prompt
-        params.seed = seed
-        params.guidance_scale = guidance_scale
-        params.num_frames = num_frames
-        params.height = height
-        params.width = width
-        params.num_inference_steps = num_inference_steps
-
-        if randomize_seed:
-            params.seed = torch.randint(0, 1000000, (1, )).item()
-
-        if not use_negative_prompt:
-            params.negative_prompt = None
-
-        generator.generate_video(prompt=prompt, sampling_param=params)
-
-        output_path = os.path.join(parsed_args.output_path,
-                                   f"{params.prompt[:100]}.mp4")
-
-        return output_path, params.seed
-
-    examples = [
-        "A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface. The plastic wrap is stretched to cover the dough more securely. The hand adjusts the wrap, ensuring that it is tight and smooth over the dough. The scene focuses on the hand’s movements as it secures the edges of the plastic wrap. No new objects appear, and the camera remains stationary, focusing on the action of covering the dough.",
-        "A vintage train snakes through the mountains, its plume of white steam rising dramatically against the jagged peaks. The cars glint in the late afternoon sun, their deep crimson and gold accents lending a touch of elegance. The tracks carve a precarious path along the cliffside, revealing glimpses of a roaring river far below. Inside, passengers peer out the large windows, their faces lit with awe as the landscape unfolds.",
-        "A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background. Strings of fairy lights hang above, casting a warm, golden glow over the scene. Groups of people gather around high tables, their laughter blending with the soft rhythm of live jazz. The aroma of freshly mixed cocktails and charred appetizers wafts through the air, mingling with the cool night breeze.",
-    ]
-
-    with gr.Blocks() as demo:
-        gr.Markdown("# FastVideo Inference Demo")
-
-        with gr.Group():
-            with gr.Row():
-                prompt = gr.Text(
-                    label="Prompt",
-                    show_label=False,
-                    max_lines=1,
-                    placeholder="Enter your prompt",
-                    container=False,
-                )
-                run_button = gr.Button("Run", scale=0)
-            result = gr.Video(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced options", open=False):
-            with gr.Group():
-                with gr.Row():
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=256,
-                        maximum=1024,
-                        step=32,
-                        value=default_params.height,
-                    )
-                    width = gr.Slider(label="Width",
-                                      minimum=256,
-                                      maximum=1024,
-                                      step=32,
-                                      value=default_params.width)
-
-                with gr.Row():
-                    num_frames = gr.Slider(
-                        label="Number of Frames",
-                        minimum=21,
-                        maximum=163,
-                        value=default_params.num_frames,
-                    )
-                    guidance_scale = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=1,
-                        maximum=12,
-                        value=default_params.guidance_scale,
-                    )
-                    num_inference_steps = gr.Slider(
-                        label="Inference Steps",
-                        minimum=4,
-                        maximum=100,
-                        value=default_params.num_inference_steps,
-                    )
-
-                with gr.Row():
-                    use_negative_prompt = gr.Checkbox(
-                        label="Use negative prompt", value=False)
-                    negative_prompt = gr.Text(
-                        label="Negative prompt",
-                        max_lines=1,
-                        placeholder="Enter a negative prompt",
-                        visible=False,
-                    )
-
-                seed = gr.Slider(label="Seed",
-                                 minimum=0,
-                                 maximum=1000000,
-                                 step=1,
-                                 value=default_params.seed)
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                seed_output = gr.Number(label="Used Seed")
-
-        gr.Examples(examples=examples, inputs=prompt)
-
-        use_negative_prompt.change(
-            fn=lambda x: gr.update(visible=x),
-            inputs=use_negative_prompt,
-            outputs=default_params.negative_prompt,
-        )
-
-        run_button.click(
-            fn=generate_video,
-            inputs=[
-                prompt,
-                negative_prompt,
-                use_negative_prompt,
-                seed,
-                guidance_scale,
-                num_frames,
-                height,
-                width,
-                num_inference_steps,
-                randomize_seed,
-            ],
-            outputs=[result, seed_output],
-        )
-
-    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860)

examples/inference/gradio/gradio_frontend.py

@@ -0,0 +1,730 @@
+import argparse
+import os
+import requests
+import base64
+import time
+
+import gradio as gr
+
+from fastvideo.configs.sample.base import SamplingParam
+
+
+MODEL_PATH_MAPPING = {
+    "FastWan2.1-T2V-1.3B": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
+    "FastWan2.2-TI2V-5B-FullAttn": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+}
+
+
+class RayServeClient:
+    def __init__(self, backend_url: str):
+        self.backend_url = backend_url
+        self.session = requests.Session()
+    
+    def check_health(self) -> bool:
+        try:
+            response = self.session.get(f"{self.backend_url}/health", timeout=5)
+            return response.status_code == 200
+        except requests.exceptions.RequestException:
+            return False
+    
+    def generate_video(self, request_data: dict) -> dict:
+        start_time = time.time()
+        
+        try:
+            headers = {"Content-Type": "application/json"}
+            
+            response = self.session.post(
+                f"{self.backend_url}/generate_video",
+                json=request_data,
+                headers=headers,
+                timeout=300
+            )
+            
+            round_trip_time = time.time() - start_time
+            
+            if response.status_code == 200:
+                result = response.json()
+                backend_total = result.get("total_time", 0)
+                network_time = round_trip_time - backend_total
+                result["network_time"] = network_time
+                return result
+            else:
+                return {"success": False, "error_message": f"HTTP {response.status_code}: {response.text}"}
+                
+        except requests.exceptions.RequestException as e:
+            return {"success": False, "error_message": f"Request failed: {str(e)}"}
+
+
+def save_video_from_base64(video_data: str, output_dir: str, prompt: str) -> str:
+    if not video_data:
+        return None
+    
+    try:
+        if video_data.startswith('data:video/'):
+            video_data = video_data.split(',')[1]
+        
+        video_bytes = base64.b64decode(video_data)
+        
+        safe_prompt = prompt[:50].replace(' ', '_').replace('/', '_').replace('\\', '_')
+        video_filename = f"{safe_prompt}.mp4"
+        video_path = os.path.join(output_dir, video_filename)
+        
+        os.makedirs(output_dir, exist_ok=True)
+        
+        with open(video_path, 'wb') as f:
+            f.write(video_bytes)
+        
+        return video_path
+        
+    except Exception as e:
+        print(f"Failed to save video: {e}")
+        return None
+
+
+def create_timing_display(inference_time, encoding_time, network_time, total_time, stage_execution_times, num_frames):
+    dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A"
+    
+    timing_html = f"""
+    <div style="margin: 10px 0;">
+        <h3 style="text-align: center; margin-bottom: 10px;">⏱️ Timing Breakdown</h3>
+        <div style="display: grid; grid-template-columns: repeat(5, 1fr); gap: 10px; margin-bottom: 10px;">
+            <div class="timing-card timing-card-highlight">
+                <div style="font-size: 20px;">🚀</div>
+                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">DiT Denoising</div>
+                <div style="font-size: 16px; color: #ffa200; font-weight: bold;">{dit_denoising_time}</div>
+            </div>
+            <div class="timing-card">
+                <div style="font-size: 20px;">🧠</div>
+                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">E2E (w. vae/text encoder)</div>
+                <div style="font-size: 16px; color: #2563eb;">{inference_time:.2f}s</div>
+            </div>
+            <div class="timing-card">
+                <div style="font-size: 20px;">🎬</div>
+                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Video Encoding</div>
+                <div style="font-size: 16px; color: #dc2626;">{encoding_time:.2f}s</div>
+            </div>
+            <div class="timing-card">
+                <div style="font-size: 20px;">🌐</div>
+                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Network Transfer</div>
+                <div style="font-size: 16px; color: #059669;">{network_time:.2f}s</div>
+            </div>
+            <div class="timing-card">
+                <div style="font-size: 20px;">📊</div>
+                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Total Processing</div>
+                <div style="font-size: 18px; color: #0277bd;">{total_time:.2f}s</div>
+            </div>
+        </div>"""
+    
+    if inference_time > 0:
+        fps = num_frames / inference_time
+        timing_html += f"""
+        <div class="performance-card" style="margin-top: 15px;">
+            <span style="font-weight: bold;">Generation Speed: </span>
+            <span style="font-size: 18px; color: #6366f1; font-weight: bold;">{fps:.1f} frames/second</span>
+        </div>"""
+    
+    return timing_html + "</div>"
+
+
+def load_example_prompts():
+    def contains_chinese(text):
+        return any('\u4e00' <= char <= '\u9fff' for char in text)
+    
+    def load_from_file(filepath):
+        prompts, labels = [], []
+        try:
+            with open(filepath, "r", encoding='utf-8') as f:
+                for line in f:
+                    line = line.strip()
+                    if line and not contains_chinese(line):
+                        label = line[:100] + "..." if len(line) > 100 else line
+                        labels.append(label)
+                        prompts.append(line)
+        except Exception as e:
+            print(f"Warning: Could not read {filepath}: {e}")
+        return prompts, labels
+    
+    examples, example_labels = load_from_file("prompts/prompts_final.txt")
+    
+    if not examples:
+        examples = ["A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background."]
+        example_labels = ["Crowded rooftop bar at night"]
+    
+    return examples, example_labels
+
+
+def create_gradio_interface(backend_url: str, default_params: dict[str, SamplingParam]):
+    
+    client = RayServeClient(backend_url)
+    
+    def generate_video(
+        prompt, negative_prompt, use_negative_prompt, seed, guidance_scale,
+        num_frames, height, width, randomize_seed, model_selection, progress
+    ):
+        if not client.check_health():
+            return None, f"Backend is not available. Please check if Ray Serve is running at {backend_url}", ""
+        
+        # Validate dimensions
+        max_pixels = 720 * 1280
+        if height * width > max_pixels:
+            return None, f"Video dimensions too large. Maximum: 720x1280 pixels", ""
+        
+        if progress:
+            progress(0.1, desc="Checking backend health...")
+        
+        request_data = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "use_negative_prompt": use_negative_prompt,
+            "seed": seed,
+            "guidance_scale": guidance_scale,
+            "num_frames": num_frames,
+            "height": height,
+            "width": width,
+            "randomize_seed": randomize_seed,
+            "return_frames": False,
+            "image_path": None,
+            "model_path": MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers")
+        }
+        
+        if progress:
+            progress(0.4, desc="Generating video...")
+        
+        response = client.generate_video(request_data)
+        
+        if progress:
+            progress(0.8, desc="Processing response...")
+        
+        if response.get("success", False):
+            video_data = response.get("video_data", "")
+            used_seed = response.get("seed", seed)
+            inference_time = response.get("inference_time", 0.0)
+            encoding_time = response.get("encoding_time", 0.0)
+            total_time = response.get("total_time", 0.0)
+            network_time = response.get("network_time", 0.0)
+            stage_execution_times = response.get("stage_execution_times", [])
+            
+            timing_details = create_timing_display(
+                inference_time, encoding_time, network_time, total_time, 
+                stage_execution_times, num_frames
+            )
+            
+            if video_data:
+                if progress:
+                    progress(0.9, desc="Saving video...")
+                
+                video_path = save_video_from_base64(video_data, "outputs", prompt)
+                
+                if progress:
+                    progress(1.0, desc="Generation complete!")
+                
+                if video_path and os.path.exists(video_path):
+                    return video_path, used_seed, timing_details
+                else:
+                    return None, "Failed to save video", ""
+            else:
+                return None, "No video data received from backend", ""
+        else:
+            error_msg = response.get("error_message", "Unknown error occurred")
+            return None, f"Generation failed: {error_msg}", ""
+    
+    examples, example_labels = load_example_prompts()
+    
+    theme = gr.themes.Base().set(
+        button_primary_background_fill="#2563eb",
+        button_primary_background_fill_hover="#1d4ed8",
+        button_primary_text_color="white",
+        slider_color="#2563eb",
+        checkbox_background_color_selected="#2563eb",
+    )
+    
+    def get_default_values(model_name):
+        model_path = MODEL_PATH_MAPPING.get(model_name)
+        if model_path and model_path in default_params:
+            params = default_params[model_path]
+            return {
+                'height': params.height,
+                'width': params.width,
+                'num_frames': params.num_frames,
+                'guidance_scale': params.guidance_scale,
+                'seed': params.seed,
+            }
+        
+        return {
+            'height': 448,
+            'width': 832,
+            'num_frames': 61,
+            'guidance_scale': 3.0,
+            'seed': 1024,
+        }
+    
+    initial_values = get_default_values("FastWan2.1-T2V-1.3B")
+    
+    with gr.Blocks(title="FastWan", theme=theme) as demo:
+        gr.Image("fastvideo-logos/main/svg/full.svg", show_label=False, container=False, height=80)
+        gr.HTML("""
+        <div style="text-align: center; margin-bottom: 10px;">
+            <p style="font-size: 18px;"> Make Video Generation Go Blurrrrrrr </p>
+            <p style="font-size: 18px;"> <a href="https://github.com/hao-ai-lab/FastVideo/tree/main" target="_blank">Code</a> | <a href="https://hao-ai-lab.github.io/blogs/fastvideo_post_training/" target="_blank">Blog</a> | <a href="https://hao-ai-lab.github.io/FastVideo/" target="_blank">Docs</a>  </p>
+        </div>
+        """)
+        
+        with gr.Accordion("🎥 What Is FastVideo?", open=False):
+            gr.HTML("""
+            <div style="padding: 20px; line-height: 1.6;">
+                <p style="font-size: 16px; margin-bottom: 15px;">
+                    FastVideo is an inference and post-training framework for diffusion models. It 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>
+            </div>
+            """)
+        
+        with gr.Row():
+            model_selection = gr.Dropdown(
+                choices=list(MODEL_PATH_MAPPING.keys()),
+                value="FastWan2.1-T2V-1.3B",
+                label="Select Model",
+                interactive=True
+            )
+
+        with gr.Row():
+            example_dropdown = gr.Dropdown(
+                choices=example_labels,
+                label="Example Prompts",
+                value=None,
+                interactive=True,
+                allow_custom_value=False
+            )
+        
+        with gr.Row():
+            with gr.Column(scale=6):
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=3,
+                    placeholder="Describe your scene...",
+                    container=False,
+                    lines=3,
+                    autofocus=True,
+                )
+            with gr.Column(scale=1, min_width=120, elem_classes="center-button"):
+                run_button = gr.Button("Run", variant="primary", size="lg")
+        
+        with gr.Row():
+            with gr.Column():
+                error_output = gr.Text(label="Error", visible=False)
+                timing_display = gr.Markdown(label="Timing Breakdown", visible=False)
+
+        with gr.Row(equal_height=True, elem_classes="main-content-row"):
+            with gr.Column(scale=1, elem_classes="advanced-options-column"):
+                with gr.Group():
+                    gr.HTML("<div style='margin: 0 0 15px 0; text-align: center; font-size: 16px;'>Advanced Options</div>")
+                    with gr.Row():
+                        height = gr.Number(
+                            label="Height",
+                            value=initial_values['height'],
+                            interactive=False,
+                            container=True
+                        )
+                        width = gr.Number(
+                            label="Width",
+                            value=initial_values['width'],
+                            interactive=False,
+                            container=True
+                        )
+                    
+                    with gr.Row():
+                        num_frames = gr.Number(
+                            label="Number of Frames",
+                            value=initial_values['num_frames'],
+                            interactive=False,
+                            container=True
+                        )
+                        guidance_scale = gr.Slider(
+                            label="Guidance Scale",
+                            minimum=1,
+                            maximum=12,
+                            value=initial_values['guidance_scale'],
+                        )
+                    
+                    with gr.Row():
+                        use_negative_prompt = gr.Checkbox(
+                            label="Use negative prompt", value=False)
+                        negative_prompt = gr.Text(
+                            label="Negative prompt",
+                            max_lines=3,
+                            lines=3,
+                            placeholder="Enter a negative prompt",
+                            visible=False,
+                        )
+
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=1000000,
+                        step=1,
+                        value=initial_values['seed'],
+                    )
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+                    seed_output = gr.Number(label="Used Seed")
+        
+            with gr.Column(scale=1, elem_classes="video-column"):
+                result = gr.Video(
+                    label="Generated Video", 
+                    show_label=True,
+                    height=466,
+                    width=600,
+                    container=True,
+                    elem_classes="video-component"
+                )
+        
+        gr.HTML("""
+        <style>
+        .center-button {
+            display: flex !important;
+            justify-content: center !important;
+            height: 100% !important;
+            padding-top: 1.4em !important;
+        }
+        
+        .gradio-container {
+            max-width: 1200px !important;
+            margin: 0 auto !important;
+        }
+        
+        .main {
+            max-width: 1200px !important;
+            margin: 0 auto !important;
+        }
+        
+        .gr-form, .gr-box, .gr-group {
+            max-width: 1200px !important;
+        }
+        
+        .gr-video {
+            max-width: 500px !important;
+            margin: 0 auto !important;
+        }
+        
+        .main-content-row {
+            display: flex !important;
+            align-items: flex-start !important;
+            min-height: 500px !important;
+            gap: 20px !important;
+        }
+        
+        .advanced-options-column,
+        .video-column {
+            display: flex !important;
+            flex-direction: column !important;
+            flex: 1 !important;
+            min-height: 400px !important;
+            align-items: stretch !important;
+        }
+        
+        .video-column > * {
+            margin-top: 0 !important;
+        }
+        
+        .video-column .gr-video,
+        .video-component {
+            margin-top: 0 !important;
+            padding-top: 0 !important;
+        }
+        
+        .video-column .gr-video .gr-form {
+            margin-top: 0 !important;
+        }
+        
+        .advanced-options-column .gr-group,
+        .video-column .gr-video {
+            margin-top: 0 !important;
+            vertical-align: top !important;
+        }
+        
+        .advanced-options-column > *:last-child,
+        .video-column > *:last-child {
+            flex-grow: 0 !important;
+        }
+        
+        @media (max-width: 1400px) {
+            .main-content-row {
+                min-height: 600px !important;
+            }
+            
+            .advanced-options-column,
+            .video-column {
+                min-height: 600px !important;
+            }
+        }
+        
+        @media (max-width: 1200px) {
+            .main-content-row {
+                flex-direction: column !important;
+                align-items: stretch !important;
+            }
+            
+            .advanced-options-column,
+            .video-column {
+                min-height: auto !important;
+                width: 100% !important;
+            }
+        }
+        
+        .timing-card {
+            background: var(--background-fill-secondary) !important;
+            border: 1px solid var(--border-color-primary) !important;
+            color: var(--body-text-color) !important;
+            padding: 10px;
+            border-radius: 8px;
+            text-align: center;
+            min-height: 80px;
+            display: flex;
+            flex-direction: column;
+            justify-content: center;
+        }
+        
+        .timing-card-highlight {
+            background: var(--background-fill-primary) !important;
+            border: 2px solid var(--color-accent) !important;
+        }
+        
+        .performance-card {
+            background: var(--background-fill-secondary) !important;
+            border: 1px solid var(--border-color-primary) !important;
+            color: var(--body-text-color) !important;
+            padding: 10px;
+            border-radius: 6px;
+            text-align: center;
+        }
+        
+        .gr-number input[readonly] {
+            background-color: var(--background-fill-secondary) !important;
+            border: 1px solid var(--border-color-primary) !important;
+            color: var(--body-text-color-subdued) !important;
+            cursor: default !important;
+            text-align: center !important;
+            font-weight: 500 !important;
+        }
+        </style>
+        """)
+        
+        def on_example_select(example_label):
+            if example_label and example_label in example_labels:
+                index = example_labels.index(example_label)
+                return examples[index]
+            return ""
+        
+        example_dropdown.change(
+            fn=on_example_select,
+            inputs=example_dropdown,
+            outputs=prompt,
+        )
+        
+        gr.HTML("""
+        <div style="text-align: center; margin-top: 10px; margin-bottom: 15px;">
+            <p style="font-size: 16px; margin: 0;">The compute for this demo is generously provided by <a href="https://www.gmicloud.ai/" target="_blank">GMI Cloud</a>. Note that this demo is meant to showcase FastWan's quality and that under a large number of requests, generation speed may be affected. We are also rate-limiting users to 3 requests per minute.</p>
+        </div>
+        """)
+        
+        use_negative_prompt.change(
+            fn=lambda x: gr.update(visible=x),
+            inputs=use_negative_prompt,
+            outputs=negative_prompt,
+        )
+        
+        def on_model_selection_change(selected_model):
+            if not selected_model:
+                selected_model = "FastWan2.1-T2V-1.3B"
+            
+            model_path = MODEL_PATH_MAPPING.get(selected_model)
+            
+            if model_path and model_path in default_params:
+                params = default_params[model_path]
+                return (
+                    gr.update(value=params.height),
+                    gr.update(value=params.width),
+                    gr.update(value=params.num_frames),
+                    gr.update(value=params.guidance_scale),
+                    gr.update(value=params.seed),
+                )
+            
+            return (
+                gr.update(value=448),
+                gr.update(value=832),
+                gr.update(value=61),
+                gr.update(value=3.0),
+                gr.update(value=1024),
+            )
+        
+        model_selection.change(
+            fn=on_model_selection_change,
+            inputs=model_selection,
+            outputs=[height, width, num_frames, guidance_scale, seed],
+        )
+        
+        def handle_generation(*args, progress=None, request: gr.Request = None):
+            model_selection, prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, num_frames, height, width, randomize_seed = args
+            
+            result_path, seed_or_error, timing_details = generate_video(
+                prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, 
+                num_frames, height, width, randomize_seed, model_selection, progress
+            )
+            
+            if result_path and os.path.exists(result_path):
+                return (
+                    result_path, 
+                    seed_or_error, 
+                    gr.update(visible=False),
+                    gr.update(visible=True, value=timing_details),
+                )
+            else:
+                return (
+                    None, 
+                    seed_or_error, 
+                    gr.update(visible=True, value=seed_or_error),
+                    gr.update(visible=False),
+                )
+        
+        run_button.click(
+            fn=handle_generation,
+            inputs=[
+                model_selection,
+                prompt,
+                negative_prompt,
+                use_negative_prompt,
+                seed,
+                guidance_scale,
+                num_frames,
+                height,
+                width,
+                randomize_seed,
+            ],
+            outputs=[result, seed_output, error_output, timing_display],
+            concurrency_limit=20,
+        )
+    
+    return demo
+
+
+def main():
+    parser = argparse.ArgumentParser(description="FastVideo Gradio Frontend")
+    parser.add_argument("--backend_url", type=str, default="http://localhost:8000",
+                        help="URL of the Ray Serve backend")
+    parser.add_argument("--t2v_model_paths", type=str,
+                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.1-T2V-14B-Diffusers",
+                        help="Comma separated list of paths to the T2V model(s)")
+    parser.add_argument("--host", type=str, default="0.0.0.0",
+                        help="Host to bind to")
+    parser.add_argument("--port", type=int, default=7860,
+                        help="Port to bind to")
+    
+    args = parser.parse_args()
+    
+    default_params = {}
+    model_paths = args.t2v_model_paths.split(",")
+    for model_path in model_paths:
+        default_params[model_path] = SamplingParam.from_pretrained(model_path)
+    
+    demo = create_gradio_interface(args.backend_url, default_params)
+    
+    print(f"Starting Gradio frontend at http://{args.host}:{args.port}")
+    print(f"Backend URL: {args.backend_url}")
+    print(f"T2V Models: {args.t2v_model_paths}")
+    
+    from fastapi import FastAPI, Request, HTTPException
+    from fastapi.responses import HTMLResponse, FileResponse
+    import uvicorn
+    
+    app = FastAPI()
+    
+    @app.get("/logo.png")
+    def get_logo():
+        return FileResponse(
+            "fastvideo-logos/main/svg/full.svg", 
+            media_type="image/svg+xml",
+            headers={
+                "Cache-Control": "public, max-age=3600",
+                "Access-Control-Allow-Origin": "*"
+            }
+        )
+    
+    @app.get("/favicon.ico")
+    def get_favicon():
+        favicon_path = "fastvideo-logos/main/svg/icon-simple.svg"
+        
+        if os.path.exists(favicon_path):
+            return FileResponse(
+                favicon_path, 
+                media_type="image/svg+xml",
+                headers={
+                    "Cache-Control": "public, max-age=3600",
+                    "Access-Control-Allow-Origin": "*"
+                }
+            )
+        else:
+            raise HTTPException(status_code=404, detail="Favicon not found")
+    
+    @app.get("/", response_class=HTMLResponse)
+    def index(request: Request):
+        base_url = str(request.base_url).rstrip('/')
+        return f"""
+        <!DOCTYPE html>
+        <html lang="en">
+        <head>
+            <meta charset="UTF-8" />
+            <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+            
+            <title>FastWan</title>
+            <meta name="title" content="FastWan">
+            <meta name="description" content="Make video generation go blurrrrrrr">
+            <meta name="keywords" content="FastVideo, video generation, AI, machine learning, FastWan">
+            
+            <meta property="og:type" content="website">
+            <meta property="og:url" content="{base_url}/">
+            <meta property="og:title" content="FastWan">
+            <meta property="og:description" content="Make video generation go blurrrrrrr">
+            <meta property="og:image" content="{base_url}/logo.png">
+            <meta property="og:image:width" content="1200">
+            <meta property="og:image:height" content="630">
+            <meta property="og:site_name" content="FastWan">
+            
+            <meta property="twitter:card" content="summary_large_image">
+            <meta property="twitter:url" content="{base_url}/">
+            <meta property="twitter:title" content="FastWan">
+            <meta property="twitter:description" content="Make video generation go blurrrrrrr">
+            <meta property="twitter:image" content="{base_url}/logo.png">
+            <link rel="icon" type="image/png" sizes="32x32" href="/favicon.ico">
+            <link rel="icon" type="image/png" sizes="16x16" href="/favicon.ico">
+            <link rel="apple-touch-icon" href="/favicon.ico">
+            <style>
+                body, html {{
+                    margin: 0;
+                    padding: 0;
+                    height: 100%;
+                    overflow: hidden;
+                }}
+                iframe {{
+                    width: 100%;
+                    height: 100vh;
+                    border: none;
+                }}
+            </style>
+        </head>
+        <body>
+            <iframe src="/gradio" width="100%" height="100%" style="border: none;"></iframe>
+        </body>
+        </html>
+        """
+    
+    app = gr.mount_gradio_app(
+        app, 
+        demo, 
+        path="/gradio",
+        allowed_paths=[os.path.abspath("outputs"), os.path.abspath("fastvideo-logos")]
+    )
+    
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main() 

examples/inference/gradio/ray_serve_backend.py

@@ -0,0 +1,379 @@
+import time
+import os
+import torch
+import base64
+import io
+from copy import deepcopy
+from typing import Dict, Any, Optional, List
+import signal
+import sys
+
+import ray
+from ray import serve
+from fastapi import FastAPI, Request, Response
+from pydantic import BaseModel
+import numpy as np
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+import imageio
+from ray.serve.handle import DeploymentHandle
+from prometheus_client import Counter, Histogram, generate_latest
+
+NUM_GPUS = 16
+DEFAULT_FPS = 16
+SEED_RANGE_MAX = 1_000_000
+SUPPORTED_MODELS = [
+    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", 
+    "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+]
+
+MODEL_CONFIGS = {
+    "1.3B": {
+        "num_cpus": 2,
+        "text_encoder_cpu_offload": False,
+        "dit_cpu_offload": False,
+        "vae_cpu_offload": False,
+        "VSA_sparsity": 0.8,
+    },
+    "14B": {
+        "num_cpus": 16,
+        "text_encoder_cpu_offload": True,
+        "dit_cpu_offload": True,
+        "vae_cpu_offload": False,
+        "VSA_sparsity": 0.9,
+    }
+}
+
+
+class VideoGenerationRequest(BaseModel):
+    prompt: str
+    negative_prompt: Optional[str] = None
+    use_negative_prompt: bool = False
+    seed: int = 42
+    guidance_scale: float = 7.5
+    num_frames: int = 21
+    height: int = 448
+    width: int = 832
+    randomize_seed: bool = False
+    return_frames: bool = False
+    model_path: Optional[str] = None
+
+
+class VideoGenerationResponse(BaseModel):
+    video_data: Optional[str] = None
+    seed: int
+    success: bool
+    error_message: Optional[str] = None
+    generation_time: Optional[float] = None
+    model_load_time: Optional[float] = None
+    inference_time: Optional[float] = None
+    encoding_time: Optional[float] = None
+    total_time: Optional[float] = None
+    stage_names: Optional[List[str]] = None
+    stage_execution_times: Optional[List[float]] = None
+
+
+def encode_video_to_base64(frames: List[np.ndarray], fps: int = DEFAULT_FPS) -> str:
+    if not frames:
+        return ""
+    
+    try:
+        buffer = io.BytesIO()
+        imageio.mimsave(buffer, frames, fps=fps, format="mp4")
+        buffer.seek(0)
+        
+        video_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+        return f"data:video/mp4;base64,{video_base64}"
+        
+    except Exception as e:
+        print(f"Warning: Failed to encode video: {e}")
+        return ""
+
+
+def setup_model_environment(model_path: str) -> None:
+    if "fullattn" in model_path.lower():
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
+    else:
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+    os.environ["FASTVIDEO_STAGE_LOGGING"] = "1"
+
+
+def process_generation_result(result: Any) -> tuple[List[np.ndarray], float, List[str], List[float]]:
+    frames = result if isinstance(result, list) else result.get("frames", [])
+    generation_time = result.get("generation_time", 0.0) if isinstance(result, dict) else 0.0
+    
+    logging_info = result.get("logging_info", None)
+    if logging_info:
+        stage_names = logging_info.get_execution_order()
+        stage_execution_times = [
+            logging_info.get_stage_info(stage_name).get("execution_time", 0.0) 
+            for stage_name in stage_names
+        ]
+    else:
+        stage_names = []
+        stage_execution_times = []
+    
+    return frames, generation_time, stage_names, stage_execution_times
+
+
+def prepare_sampling_params(video_request: VideoGenerationRequest, default_params: Any) -> Any:
+    params = deepcopy(default_params)
+    params.prompt = video_request.prompt
+    
+    if video_request.use_negative_prompt:
+        params.negative_prompt = video_request.negative_prompt
+
+    params.seed = (video_request.seed if not video_request.randomize_seed 
+                  else torch.randint(0, SEED_RANGE_MAX, (1,)).item())
+    params.randomize_seed = video_request.randomize_seed
+    params.guidance_scale = video_request.guidance_scale
+    params.num_frames = video_request.num_frames
+    params.height = video_request.height
+    params.width = video_request.width
+    params.save_video = False
+    params.return_frames = False
+    
+    return params
+
+
+class BaseModelDeployment:
+    def __init__(self, model_path: str, output_path: str = "outputs"):
+        self.model_path = model_path
+        self.output_path = output_path
+        self.generator = None
+        self.default_params = None
+
+        os.makedirs(self.output_path, exist_ok=True)
+        setup_model_environment(self.model_path)
+
+    def _initialize_generator(self, config: Dict[str, Any]) -> None:
+        from fastvideo.entrypoints.video_generator import VideoGenerator
+        from fastvideo.configs.sample.base import SamplingParam
+
+        print(f"Initializing model: {self.model_path}")
+        self.generator = VideoGenerator.from_pretrained(
+            model_path=self.model_path,
+            num_gpus=1,
+            use_fsdp_inference=True,
+            text_encoder_cpu_offload=config["text_encoder_cpu_offload"],
+            dit_cpu_offload=config["dit_cpu_offload"],
+            vae_cpu_offload=config["vae_cpu_offload"],
+            VSA_sparsity=config["VSA_sparsity"],
+            enable_stage_verification=False,
+        )
+        self.default_params = SamplingParam.from_pretrained(self.model_path)
+
+    def generate_video(self, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
+        total_start_time = time.time()
+        
+        params = prepare_sampling_params(video_request, self.default_params)
+
+        inference_start_time = time.time()
+        result = self.generator.generate_video(
+            prompt=video_request.prompt,
+            sampling_param=params,
+            save_video=False,
+            return_frames=False,
+        )
+        inference_time = time.time() - inference_start_time
+
+        frames, generation_time, stage_names, stage_execution_times = process_generation_result(result)
+
+        encoding_start_time = time.time()
+        video_data = encode_video_to_base64(frames, fps=DEFAULT_FPS)
+        encoding_time = time.time() - encoding_start_time
+        
+        total_time = time.time() - total_start_time
+
+        return VideoGenerationResponse(
+            video_data=video_data,
+            seed=params.seed,
+            success=True,
+            generation_time=generation_time,
+            inference_time=inference_time,
+            encoding_time=encoding_time,
+            total_time=total_time,
+            stage_names=stage_names,
+            stage_execution_times=stage_execution_times,
+        )
+
+
+@serve.deployment(
+    ray_actor_options={"num_cpus": 2, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
+)
+class T2VModelDeployment(BaseModelDeployment):
+    def __init__(self, t2v_model_path: str, output_path: str = "outputs"):
+        super().__init__(t2v_model_path, output_path)
+        self._initialize_generator(MODEL_CONFIGS["1.3B"])
+        print("✅ T2V model initialized successfully")
+
+
+@serve.deployment(
+    ray_actor_options={"num_cpus": 16, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
+)
+class T2V14BModelDeployment(BaseModelDeployment):
+    def __init__(self, t2v_14b_model_path: str, output_path: str = "outputs"):
+        super().__init__(t2v_14b_model_path, output_path)
+        # Override environment for 14B model
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+        self._initialize_generator(MODEL_CONFIGS["14B"])
+        print("✅ T2V 14B model initialized successfully")
+
+
+app = FastAPI()
+limiter = Limiter(key_func=get_remote_address)
+app.state.limiter = limiter
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+
+
+@serve.deployment(num_replicas=50, ray_actor_options={"num_cpus": 2})
+@serve.ingress(app)
+class FastVideoAPI:
+
+    def __init__(self, t2v_deployments: Dict[str, DeploymentHandle]):
+        self.t2v_deployments = t2v_deployments
+        
+        # Initialize Prometheus metrics
+        self.request_count = Counter('fastvideo_requests_total', 'Total FastVideo requests', ['model_type', 'status'])
+        self.request_duration = Histogram('fastvideo_request_duration_seconds', 'FastVideo request duration', ['model_type'])
+        self.video_generation_time = Histogram('fastvideo_video_generation_seconds', 'Video generation time', ['model_type'])
+    
+    def _get_model_name(self, model_path: Optional[str]) -> str:
+        return model_path.split('/')[-1] if model_path else "unknown"
+    
+    def _record_metrics(self, model_name: str, status: str, duration: float, response: Optional[VideoGenerationResponse] = None) -> None:
+        self.request_count.labels(model_type=model_name, status=status).inc()
+        self.request_duration.labels(model_type=model_name).observe(duration)
+        
+        if response and hasattr(response, 'generation_time') and response.generation_time:
+            self.video_generation_time.labels(model_type=model_name).observe(response.generation_time)
+    
+    @app.post("/generate_video", response_model=VideoGenerationResponse)
+    @limiter.limit("10/minute")
+    async def generate_video(self, request: Request, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
+        """Route the request to the appropriate model deployment based on model_path."""
+        start_time = time.time()
+        model_name = self._get_model_name(video_request.model_path)
+        
+        try:
+            if video_request.model_path not in self.t2v_deployments:
+                raise ValueError(f"Model {video_request.model_path} not found")
+            
+            response_ref = self.t2v_deployments[video_request.model_path].generate_video.remote(video_request)
+            response = await response_ref
+            
+            self._record_metrics(model_name, "success", time.time() - start_time, response)
+            return response
+
+        except Exception as e:
+            self._record_metrics(model_name, "error", time.time() - start_time)
+            
+            return VideoGenerationResponse(
+                video_data=None,
+                seed=video_request.seed,
+                success=False,
+                error_message=str(e),
+                generation_time=0,
+                inference_time=0,
+                encoding_time=0,
+                total_time=0,
+            )
+    
+    @app.get("/health")
+    @limiter.limit("10/minute")
+    async def health_check(self, request: Request) -> Dict[str, str]:
+        return {"status": "healthy"}
+    
+    @app.get("/metrics")
+    async def metrics(self) -> Response:
+        return Response(generate_latest(), media_type="text/plain")
+
+
+def validate_configuration(model_paths: List[str], replicas: List[int]) -> None:
+    assert len(model_paths) == len(replicas), "Number of models and replicas must match"
+    assert sum(replicas) <= NUM_GPUS, f"Total replicas ({sum(replicas)}) must be <= {NUM_GPUS}"
+    
+    for model, replica_count in zip(model_paths, replicas):
+        assert model in SUPPORTED_MODELS, f"Model {model} not supported"
+        assert replica_count > 0, f"Replicas must be greater than 0"
+
+
+def start_ray_serve(
+    *,
+    t2v_model_paths: str,
+    t2v_model_replicas: str,
+    output_path: str = "outputs",
+    host: str = "0.0.0.0",
+    port: int = 8000,
+) -> None:
+    if not ray.is_initialized():
+        ray.init()
+
+    model_paths = t2v_model_paths.split(",")
+    replicas = [int(r) for r in t2v_model_replicas.split(",")]
+    validate_configuration(model_paths, replicas)
+
+    t2v_deps = {}
+    for model_path, replica_count in zip(model_paths, replicas):
+        t2v_dep = T2VModelDeployment.options(num_replicas=replica_count).bind(model_path, output_path)
+        t2v_deps[model_path] = t2v_dep
+
+    api = FastVideoAPI.bind(t2v_deps)
+    serve.run(api, route_prefix="/", name="fast_video")
+
+    print(f"Ray Serve backend started at http://{host}:{port}")
+    for model_path, replica_count in zip(model_paths, replicas):
+        print(f"T2V Model: {model_path} | Replicas: {replica_count}")
+    print(f"Health check: http://{host}:{port}/health")
+    print(f"Video generation endpoint: http://{host}:{port}/generate_video")
+
+
+def setup_signal_handlers() -> None:
+    signal.signal(signal.SIGINT, lambda *_: sys.exit(0))
+    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
+
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="FastVideo Ray Serve Backend")
+    parser.add_argument("--t2v_model_paths",
+                        type=str,
+                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+                        help="Comma separated list of paths to the T2V model(s)")
+    parser.add_argument("--t2v_model_replicas",
+                        type=str,
+                        default="4,4",
+                        help="Comma separated list of number of replicas for the T2V model(s)")
+    parser.add_argument("--output_path",
+                        type=str,
+                        default="outputs",
+                        help="Path to save generated videos")
+    parser.add_argument("--host",
+                        type=str,
+                        default="0.0.0.0",
+                        help="Host to bind to")
+    parser.add_argument("--port",
+                        type=int,
+                        default=8000,
+                        help="Port to bind to")
+    
+    args = parser.parse_args()
+
+    model_paths = args.t2v_model_paths.split(",")
+    replicas = [int(r) for r in args.t2v_model_replicas.split(",")]
+    validate_configuration(model_paths, replicas)
+    
+    start_ray_serve(
+        t2v_model_paths=args.t2v_model_paths,
+        t2v_model_replicas=args.t2v_model_replicas,
+        output_path=args.output_path,
+        host=args.host,
+        port=args.port,
+    )
+
+    setup_signal_handlers()
+    print("✅ FastVideo backend is running. Press Ctrl-C to stop.")
+    while True:
+        time.sleep(3600) 

examples/inference/gradio/start.sh

@@ -0,0 +1,3 @@
+python examples/inference/gradio/start_ray_serve_app.py \
+--t2v_model_paths "FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers" \
+--t2v_model_replicas "4,4"

examples/inference/gradio/start_ray_serve_app.py

@@ -0,0 +1,257 @@
+"""
+Startup script for FastVideo with Ray Serve backend and Gradio frontend.
+This script starts both the backend and frontend services.
+"""
+
+import argparse
+import os
+import subprocess
+import sys
+import time
+import threading
+import signal
+import requests
+from pathlib import Path
+from typing import Dict, Any, Optional
+
+
+DEFAULT_BACKEND_HOST = "0.0.0.0"
+DEFAULT_BACKEND_PORT = 8000
+DEFAULT_FRONTEND_HOST = "0.0.0.0"
+DEFAULT_FRONTEND_PORT = 7860
+DEFAULT_OUTPUT_PATH = "outputs"
+DEFAULT_T2V_MODELS = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
+DEFAULT_T2V_REPLICAS = "4,4"
+
+HEALTH_CHECK_TIMEOUT = 5
+HEALTH_CHECK_MAX_RETRIES = 100
+HEALTH_CHECK_INTERVAL = 2
+PROCESS_SHUTDOWN_TIMEOUT = 5
+PROCESS_MONITOR_INTERVAL = 1
+
+PROJECT_ROOT = Path(__file__).parent.parent.parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+
+class ServiceManager:
+    
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.backend_process: Optional[subprocess.Popen] = None
+        self.frontend_process: Optional[subprocess.Popen] = None
+        self.backend_url = f"http://{args.backend_host}:{args.backend_port}"
+    
+    def check_backend_health(self, max_retries: int = HEALTH_CHECK_MAX_RETRIES) -> bool:
+        health_url = f"{self.backend_url}/health"
+        
+        for attempt in range(max_retries):
+            try:
+                response = requests.get(health_url, timeout=HEALTH_CHECK_TIMEOUT)
+                if response.status_code == 200:
+                    print(f"✅ Backend is healthy at {self.backend_url}")
+                    return True
+            except requests.exceptions.RequestException:
+                pass
+            
+            if attempt < max_retries - 1:
+                print(f"⏳ Waiting for backend to start... ({attempt + 1}/{max_retries})")
+                time.sleep(HEALTH_CHECK_INTERVAL)
+        
+        print(f"❌ Backend failed to start within {max_retries * HEALTH_CHECK_INTERVAL} seconds")
+        return False
+    
+    def _create_monitor_thread(self, process: subprocess.Popen, service_name: str) -> threading.Thread:
+        def monitor():
+            if process.stdout:
+                for line in process.stdout:
+                    print(f"[{service_name}] {line.rstrip()}")
+        
+        thread = threading.Thread(target=monitor, daemon=True)
+        thread.start()
+        return thread
+    
+    def _start_service(self, script_name: str, args_dict: Dict[str, Any], service_name: str) -> subprocess.Popen:
+        script_path = Path(__file__).parent / script_name
+        
+        cmd = [sys.executable, str(script_path)]
+        for key, value in args_dict.items():
+            cmd.extend([f"--{key}", str(value)])
+        
+        print(f"🚀 Starting {service_name}...")
+        print(f"Command: {' '.join(cmd)}")
+        
+        process = subprocess.Popen(
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            universal_newlines=True,
+            bufsize=1
+        )
+        
+        self._create_monitor_thread(process, service_name.upper())
+        return process
+    
+    def start_backend(self) -> subprocess.Popen:
+        backend_args = {
+            "t2v_model_paths": self.args.t2v_model_paths,
+            "t2v_model_replicas": self.args.t2v_model_replicas,
+            "output_path": self.args.output_path,
+            "host": self.args.backend_host,
+            "port": self.args.backend_port
+        }
+        
+        self.backend_process = self._start_service("ray_serve_backend.py", backend_args, "backend")
+        return self.backend_process
+    
+    def start_frontend(self) -> subprocess.Popen:
+        frontend_args = {
+            "backend_url": self.backend_url,
+            "t2v_model_paths": self.args.t2v_model_paths,
+            "host": self.args.frontend_host,
+            "port": self.args.frontend_port
+        }
+        
+        self.frontend_process = self._start_service("gradio_frontend.py", frontend_args, "frontend")
+        return self.frontend_process
+    
+    def shutdown_services(self) -> None:
+        print("\n🛑 Shutting down services...")
+        
+        processes = []
+        if self.frontend_process:
+            self.frontend_process.terminate()
+            processes.append(("frontend", self.frontend_process))
+        
+        if self.backend_process:
+            self.backend_process.terminate()
+            processes.append(("backend", self.backend_process))
+        
+        for name, process in processes:
+            try:
+                process.wait(timeout=PROCESS_SHUTDOWN_TIMEOUT)
+                print(f"✅ {name.capitalize()} stopped gracefully")
+            except subprocess.TimeoutExpired:
+                print(f"⚠️  Force killing {name} process...")
+                process.kill()
+        
+        print("✅ All services stopped")
+    
+    def monitor_processes(self) -> None:
+        if not self.backend_process or not self.frontend_process:
+            print("❌ Processes not properly initialized")
+            return
+        
+        try:
+            while True:
+                if self.frontend_process.poll() is not None:
+                    print("❌ Frontend process died unexpectedly")
+                    break
+                
+                if self.backend_process.poll() is not None:
+                    print("❌ Backend process died unexpectedly")
+                    break
+                
+                time.sleep(PROCESS_MONITOR_INTERVAL)
+        
+        except KeyboardInterrupt:
+            pass
+        
+        self.shutdown_services()
+
+
+def setup_signal_handlers(service_manager: ServiceManager) -> None:
+    def signal_handler(signum: int, frame: Any) -> None:
+        service_manager.shutdown_services()
+        sys.exit(0)
+    
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+
+
+def print_startup_info(args: argparse.Namespace) -> None:
+    print("🎬 FastVideo Ray Serve App")
+    print("=" * 50)
+    print(f"T2V Models: {args.t2v_model_paths}")
+    print(f"T2V Model Replicas: {args.t2v_model_replicas}")
+    print(f"Output: {args.output_path}")
+    print(f"Backend: http://{args.backend_host}:{args.backend_port}")
+    print(f"Frontend: http://{args.frontend_host}:{args.frontend_port}")
+    print("=" * 50)
+
+
+def parse_arguments() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="FastVideo Ray Serve App")
+    
+    parser.add_argument("--t2v_model_paths", 
+                        type=str,
+                        default=DEFAULT_T2V_MODELS,
+                        help="Comma separated list of paths to the T2V model(s)")
+    parser.add_argument("--t2v_model_replicas", 
+                        type=str, 
+                        default=DEFAULT_T2V_REPLICAS,
+                        help="Comma separated list of number of replicas for the T2V model(s)")
+    parser.add_argument("--output_path", 
+                        type=str, 
+                        default=DEFAULT_OUTPUT_PATH,
+                        help="Path to save generated videos")
+    
+    parser.add_argument("--backend_host", 
+                        type=str, 
+                        default=DEFAULT_BACKEND_HOST,
+                        help="Backend host to bind to")
+    parser.add_argument("--backend_port", 
+                        type=int, 
+                        default=DEFAULT_BACKEND_PORT,
+                        help="Backend port to bind to")
+    
+    parser.add_argument("--frontend_host", 
+                        type=str, 
+                        default=DEFAULT_FRONTEND_HOST,
+                        help="Frontend host to bind to")
+    parser.add_argument("--frontend_port", 
+                        type=int, 
+                        default=DEFAULT_FRONTEND_PORT,
+                        help="Frontend port to bind to")
+    
+    parser.add_argument("--skip_backend_check", 
+                        action="store_true",
+                        help="Skip backend health check")
+    
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_arguments()
+    
+    os.makedirs(args.output_path, exist_ok=True)
+    print_startup_info(args)
+    
+    service_manager = ServiceManager(args)
+    setup_signal_handlers(service_manager)
+    
+    try:
+        service_manager.start_backend()
+        
+        if not args.skip_backend_check:
+            if not service_manager.check_backend_health():
+                print("❌ Backend failed to start. Terminating...")
+                service_manager.shutdown_services()
+                sys.exit(1)
+        
+        service_manager.start_frontend()
+        
+        print("\n🎉 Both services are starting up!")
+        print(f"📺 Frontend will be available at: http://{args.frontend_host}:{args.frontend_port}")
+        print(f"🔧 Backend API will be available at: http://{args.backend_host}:{args.backend_port}")
+        print("\nPress Ctrl+C to stop both services...")
+        
+        service_manager.monitor_processes()
+    
+    except Exception as e:
+        print(f"❌ Unexpected error: {e}")
+        service_manager.shutdown_services()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

examples/inference/lora/wan_lora_inference.py

@@ -10,7 +10,7 @@ def main():
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
         lora_path="benjamin-paine/steamboat-willie-1.3b",
         lora_nickname="steamboat"
     )
@@ -18,7 +18,7 @@ def main():
         "height": 480,
         "width": 832,
         "num_frames": 81,
-        "guidance_scale": 5.0,
+        "guidance_scale": 6.0,
         "num_inference_steps": 32,
         "seed": 42,
     }

examples/inference/lora/wan_lora_inference_from_ckpt.py

@@ -11,22 +11,22 @@ def main():
         "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         num_gpus=1,
         dit_cpu_offload=False,
-        vae_cpu_offload=False,
+        vae_cpu_offload=True,
         text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
-        lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-1250/transformer",
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument" 
+        lora_path="checkpoints/wan_t2v_finetune_lora/checkpoint-1000/transformer",
         lora_nickname="crush_smol"
     )
     kwargs = {
         "height": 480,
         "width": 832,
         "num_frames": 77,
-        "guidance_scale": 5.0,
+        "guidance_scale": 6.0,
         "num_inference_steps": 50,
         "seed": 42,
     }
     # Generate video with LoRA style
-    prompt = "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."
+    prompt = "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press."
 
     video = generator.generate_video(
         prompt,
@@ -34,6 +34,12 @@ def main():
         save_video=True,
         **kwargs
     )    
-
+    prompt = "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press."
+    video = generator.generate_video(
+        prompt,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+        **kwargs
+    )
 if __name__ == "__main__":
     main()

examples/inference/optimizations/attention_example.py

@@ -14,7 +14,7 @@ def main():
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
     )
     load_time = time.perf_counter() - start_time
     print(f"Model loading time: {load_time:.2f} seconds")

examples/inference/optimizations/teacache_example.py

@@ -12,7 +12,7 @@ def main():
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,
-        pin_cpu_memory=False,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
     )
     load_time = time.perf_counter() - start_time
     print(f"Model loading time: {load_time:.2f} seconds")

examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh

@@ -54,7 +54,7 @@ validation_args=(
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
   --validation_steps 100
   --validation_sampling_steps "40"
-  --validation_guidance_scale "1.0"
+  --validation_guidance_scale "6.0"
 )
 
 # Optimizer arguments
@@ -69,7 +69,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"

examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm

@@ -88,7 +88,7 @@ validation_args=(
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
   --validation_steps 100
   --validation_sampling_steps "40"
-  --validation_guidance_scale "1.0"
+  --validation_guidance_scale "6.0"
 )
 
 # Optimizer arguments
@@ -103,7 +103,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/I2V-14B-VSA.slurm

@@ -101,7 +101,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.1
   --dit_precision "fp32"

examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm

@@ -101,7 +101,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.1
   --dit_precision "fp32"

examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh

@@ -54,7 +54,7 @@ validation_args=(
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
   --validation_steps 100
   --validation_sampling_steps "40"
-  --validation_guidance_scale "1.0"
+  --validation_guidance_scale "6.0"
 )
 
 # Optimizer arguments
@@ -69,7 +69,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"

examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm

@@ -88,7 +88,7 @@ validation_args=(
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
   --validation_steps 100
   --validation_sampling_steps "40"
-  --validation_guidance_scale "1.0"
+  --validation_guidance_scale "6.0"
 )
 
 # Optimizer arguments
@@ -103,7 +103,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"

examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh

@@ -54,7 +54,7 @@ validation_args=(
   --validation_preprocessed_path "$VALIDATION_DIR"
   --validation_steps 100
   --validation_sampling_steps "40"
-  --validation_guidance_scale "1.0"
+  --validation_guidance_scale "6.0"
 )
 
 # Optimizer arguments
@@ -69,7 +69,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"

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

@@ -52,9 +52,9 @@ dataset_args=(
 validation_args=(
   --log_validation 
   --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 50
+  --validation_steps 200
   --validation_sampling_steps "50" 
-  --validation_guidance_scale "1.0"
+  --validation_guidance_scale "6.0"
 )
 
 # Optimizer arguments
@@ -69,7 +69,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"
@@ -78,6 +77,7 @@ miscellaneous_args=(
   --num_euler_timesteps 50
   --ema_start_step 0
   --enable_gradient_checkpointing_type "full"
+  # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500"
 )
 
 torchrun \

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

@@ -85,14 +85,14 @@ validation_args=(
   --validation_dataset_file "$VALIDATION_DATASET_FILE"
   --validation_steps 100
   --validation_sampling_steps "50"
-  --validation_guidance_scale "1.0"
+  --validation_guidance_scale "6.0"
 )
 
 # Optimizer arguments
 optimizer_args=(
   --learning_rate 5e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 500
+  --checkpointing_steps 400
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )
@@ -100,7 +100,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"

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

@@ -52,16 +52,16 @@ dataset_args=(
 validation_args=(
   --log_validation 
   --validation_dataset_file $VALIDATION_DATASET_FILE
-  --validation_steps 50
+  --validation_steps 200
   --validation_sampling_steps "50" 
-  --validation_guidance_scale "1.0"
+  --validation_guidance_scale "6.0"
 )
 
 # Optimizer arguments
 optimizer_args=(
   --learning_rate 5e-5
   --mixed_precision "bf16"
-  --checkpointing_steps 500
+  --checkpointing_steps 400
   --weight_decay 1e-4
   --max_grad_norm 1.0
 )
@@ -69,7 +69,6 @@ optimizer_args=(
 # Miscellaneous arguments
 miscellaneous_args=(
   --inference_mode False
-  --allow_tf32
   --checkpoints_total_limit 3
   --training_cfg_rate 0.1
   --multi_phased_distill_schedule "4000-1"

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

@@ -0,0 +1,24 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_TYPE="wan"
+DATASET_PATH="data/crush-smol-test/"
+OUTPUT_DIR="data/crush-smol_processed_t2v/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/pipelines/preprocess/v1_preprocessing_new.py \
+    --model_path $MODEL_PATH \
+    --mode preprocess \
+    --workload_type t2v \
+    --preprocess.dataset_path $DATASET_PATH \
+    --preprocess.dataset_output_dir $OUTPUT_DIR \
+    --preprocess.preprocess_video_batch_size 4 \
+    --preprocess.dataloader_num_workers 0 \
+    --preprocess.max_height 480 \
+    --preprocess.max_width 832 \
+    --preprocess.num_frames 77 \
+    --preprocess.train_fps 16 \
+    --preprocess.samples_per_file 8 \
+    --preprocess.flush_frequency 8 \
+    --preprocess.video_length_tolerance_range 5

fastvideo/configs/models/dits/wanvideo.py

@@ -96,7 +96,7 @@ class WanVideoArchConfig(DiTArchConfig):
         super().__post_init__()
         self.out_channels = self.out_channels or self.in_channels
         self.hidden_size = self.num_attention_heads * self.attention_head_dim
-        self.num_channels_latents = self.in_channels if self.added_kv_proj_dim is None else self.out_channels
+        self.num_channels_latents = self.out_channels
 
 
 @dataclass

fastvideo/configs/pipelines/registry.py

@@ -31,7 +31,9 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = {
     "FastVideo/FastWan2.2-TI2V-5B-Diffusers": FastWan2_2_TI2V_5B_Config,
     "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
     "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers": WanT2V720PConfig,
-    "Wan-AI/Wan2.2-TI2V-5B-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,
     # Add other specific weight variants
 }
 

fastvideo/configs/sample/base.py

@@ -20,7 +20,7 @@ class SamplingParam:
 
     # Text inputs
     prompt: str | list[str] | None = None
-    negative_prompt: str | None = None
+    negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
     prompt_path: str | None = None
     output_path: str = "outputs/"
     output_video_name: str | None = None

fastvideo/configs/sample/registry.py

@@ -6,12 +6,19 @@ from typing import Any
 from fastvideo.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                               HunyuanSamplingParam)
 from fastvideo.configs.sample.stepvideo import StepVideoT2VSamplingParam
-from fastvideo.configs.sample.wan import (FastWanT2V480PConfig,
-                                          Wan2_2_TI2V_5B_SamplingParam,
-                                          WanI2V_14B_480P_SamplingParam,
-                                          WanI2V_14B_720P_SamplingParam,
-                                          WanT2V_1_3B_SamplingParam,
-                                          WanT2V_14B_SamplingParam)
+
+# isort: off
+from fastvideo.configs.sample.wan import (
+    FastWanT2V480PConfig,
+    Wan2_2_I2V_A14B_SamplingParam,
+    Wan2_2_T2V_A14B_SamplingParam,
+    Wan2_2_TI2V_5B_SamplingParam,
+    WanI2V_14B_480P_SamplingParam,
+    WanI2V_14B_720P_SamplingParam,
+    WanT2V_1_3B_SamplingParam,
+    WanT2V_14B_SamplingParam,
+)
+# isort: on
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index,
                              verify_model_config_and_directory)
@@ -29,10 +36,8 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
     "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig,
     "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
     "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam,
-    # "Wan-AI/Wan2.2-T2V-A14B-Diffusers":
-    # Wan2_2_T2V_A14B_SamplingParam,
-    # "Wan-AI/Wan2.2-I2V-A14B-Diffusers":
-    # Wan2_2_I2V_A14B_SamplingParam,
+    "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam,
+    "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam,
     # Add other specific weight variants
 }
 

fastvideo/configs/sample/wan.py

@@ -129,9 +129,15 @@ class Wan2_2_TI2V_5B_SamplingParam(Wan2_2_Base_SamplingParam):
 
 @dataclass
 class Wan2_2_T2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
-    pass
+    guidance_scale: float = 4.0
+    guidance_scale_2: float = 3.0
+    num_inference_steps: int = 40
+    fps: int = 16
 
 
 @dataclass
 class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam):
-    pass
+    guidance_scale: float = 3.5
+    guidance_scale_2: float = 3.5
+    num_inference_steps: int = 40
+    fps: int = 16

fastvideo/entrypoints/video_generator.py

@@ -9,6 +9,7 @@ diffusion models.
 import math
 import os
 import time
+from copy import deepcopy
 from typing import Any
 
 import imageio
@@ -202,6 +203,8 @@ class VideoGenerator:
         if sampling_param is None:
             sampling_param = SamplingParam.from_pretrained(
                 fastvideo_args.model_path)
+        else:
+            sampling_param = deepcopy(sampling_param)
 
         kwargs["prompt"] = prompt
         sampling_param.update(kwargs)
@@ -275,6 +278,7 @@ class VideoGenerator:
                        width: {target_width}
                 video_length: {sampling_param.num_frames}
                       prompt: {prompt}
+                      image_path: {sampling_param.image_path}
                   neg_prompt: {sampling_param.negative_prompt}
                         seed: {sampling_param.seed}
                  infer_steps: {sampling_param.num_inference_steps}
@@ -304,7 +308,8 @@ class VideoGenerator:
         # Run inference
         start_time = time.perf_counter()
         output_batch = self.executor.execute_forward(batch, fastvideo_args)
-        samples = output_batch
+        samples = output_batch.output
+        logging_info = output_batch.logging_info
 
         gen_time = time.perf_counter() - start_time
         logger.info("Generated successfully in %.2f seconds", gen_time)
@@ -334,9 +339,11 @@ class VideoGenerator:
         else:
             return {
                 "samples": samples,
+                "frames": frames,
                 "prompts": prompt,
                 "size": (target_height, target_width, batch.num_frames),
-                "generation_time": gen_time
+                "generation_time": gen_time,
+                "logging_info": logging_info,
             }
 
     def set_lora_adapter(self,

fastvideo/fastvideo_args.py

@@ -119,7 +119,7 @@ class FastVideoArgs:
     output_type: str = "pil"
 
     # CPU offload parameters
-    dit_cpu_offload: bool = True
+    dit_cpu_offload: bool = False
     use_fsdp_inference: bool = True
     text_encoder_cpu_offload: bool = True
     image_encoder_cpu_offload: bool = True
@@ -158,6 +158,9 @@ class FastVideoArgs:
     # # DMD parameters
     # dmd_denoising_steps: List[int] | None = field(default=None)
 
+    # MoE parameters used by Wan2.2
+    boundary_ratio: float | None = None
+
     @property
     def training_mode(self) -> bool:
         return not self.inference_mode
@@ -622,7 +625,6 @@ class TrainingArgs(FastVideoArgs):
     max_grad_norm: float = 0.0
     enable_gradient_checkpointing_type: str | None = None
     selective_checkpointing: float = 0.0
-    allow_tf32: bool = False
     mixed_precision: str = ""
     train_sp_batch_size: int = 0
     fsdp_sharding_startegy: str = ""
@@ -670,8 +672,8 @@ class TrainingArgs(FastVideoArgs):
     training_state_checkpointing_steps: int = 0  # for resuming training
     weight_only_checkpointing_steps: int = 0  # for inference
     log_visualization: bool = False
-    # simulate student forward to match inference
-    simulate_student_forward: bool = False
+    # simulate generator forward to match inference
+    simulate_generator_forward: bool = False
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -883,9 +885,6 @@ class TrainingArgs(FastVideoArgs):
         parser.add_argument("--selective-checkpointing",
                             type=float,
                             help="Selective checkpointing threshold")
-        parser.add_argument("--allow-tf32",
-                            action=StoreBoolean,
-                            help="Whether to allow TF32")
         parser.add_argument("--mixed-precision",
                             type=str,
                             help="Mixed precision training type")
@@ -1016,9 +1015,9 @@ class TrainingArgs(FastVideoArgs):
                             action=StoreBoolean,
                             help="Whether to log visualization")
         parser.add_argument(
-            "--simulate-student-forward",
+            "--simulate-generator-forward",
             action=StoreBoolean,
-            help="Whether to simulate student forward to match inference")
+            help="Whether to simulate generator forward to match inference")
 
         return parser
 

fastvideo/layers/lora/linear.py

@@ -101,15 +101,15 @@ class BaseLayerWithLoRA(nn.Module):
                          B: torch.Tensor,
                          training_mode: bool = False,
                          lora_path: str | None = None) -> None:
-        self.lora_A = A  # share storage with weights in the pipeline
-        self.lora_B = B
+        self.lora_A = torch.nn.Parameter(
+            A)  # share storage with weights in the pipeline
+        self.lora_B = torch.nn.Parameter(B)
         self.disable_lora = False
         if not training_mode:
             self.merge_lora_weights()
         self.lora_path = lora_path
 
     @torch.no_grad()
-    # @torch.compile()
     def merge_lora_weights(self) -> None:
         if self.disable_lora:
             return

fastvideo/models/dits/wanvideo.py

@@ -672,11 +672,9 @@ class WanTransformer3DModel(CachableDiT):
                 for block in self.blocks:
                     hidden_states = block(hidden_states, encoder_hidden_states,
                                           timestep_proj, freqs_cis)
-
             # if teacache is enabled, we need to cache the original hidden states
             if enable_teacache:
                 self.maybe_cache_states(hidden_states, original_hidden_states)
-
         # 5. Output norm, projection & unpatchify
         shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2,
                                                                           dim=1)

fastvideo/models/loader/component_loader.py

@@ -72,8 +72,7 @@ class ComponentLoader(ABC):
         module_loaders = {
             "scheduler": (SchedulerLoader, "diffusers"),
             "transformer": (TransformerLoader, "diffusers"),
-            "real_score_transformer": (TransformerLoader, "diffusers"),
-            "fake_score_transformer": (TransformerLoader, "diffusers"),
+            "transformer_2": (TransformerLoader, "diffusers"),
             "vae": (VAELoader, "diffusers"),
             "text_encoder": (TextEncoderLoader, "transformers"),
             "text_encoder_2": (TextEncoderLoader, "transformers"),
@@ -452,6 +451,7 @@ class TransformerLoader(ComponentLoader):
             hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
             hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
             cpu_offload=fastvideo_args.dit_cpu_offload,
+            pin_cpu_memory=fastvideo_args.pin_cpu_memory,
             fsdp_inference=fastvideo_args.use_fsdp_inference,
             # TODO(will): make these configurable
             param_dtype=torch.bfloat16,

fastvideo/models/schedulers/base.py

@@ -8,10 +8,11 @@ import torch
 class BaseScheduler(ABC):
     timesteps: torch.Tensor
     order: int
+    num_train_timesteps: int
 
     def __init__(self, *args, **kwargs) -> None:
         # Check if subclass has defined all required properties
-        required_attributes = ['timesteps', 'order']
+        required_attributes = ['timesteps', 'order', 'num_train_timesteps']
 
         for attr in required_attributes:
             if not hasattr(self, attr):

fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -134,6 +134,7 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
             sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
 
         self.timesteps = sigmas * num_train_timesteps
+        self.num_train_timesteps = num_train_timesteps
 
         self._step_index: int | None = None
         self._begin_index: int | None = None
@@ -145,6 +146,8 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin,
         self.sigma_min = self.sigmas[-1].item()
         self.sigma_max = self.sigmas[0].item()
 
+        BaseScheduler.__init__(self)
+
     @property
     def shift(self) -> float:
         """

fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py

@@ -117,6 +117,7 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
 
         self.sigmas = sigmas
         self.timesteps = sigmas * num_train_timesteps
+        self.num_train_timesteps = num_train_timesteps
 
         self.model_outputs = [None] * solver_order
         self.timestep_list: list[Any | None] = [None] * solver_order
@@ -132,6 +133,8 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self.sigma_min = self.sigmas[-1].item()
         self.sigma_max = self.sigmas[0].item()
 
+        BaseScheduler.__init__(self)
+
     @property
     def step_index(self):
         """

fastvideo/models/schedulers/scheduling_unipc_multistep.py

@@ -273,6 +273,7 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         self.num_inference_steps = None
         timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy()
         self.timesteps = torch.from_numpy(timesteps)
+        self.num_train_timesteps = num_train_timesteps
         self.model_outputs = [None] * solver_order
         self.timestep_list = [None] * solver_order
         self.lower_order_nums = 0

fastvideo/pipelines/basic/wan/wan_i2v_dmd_pipeline.py

@@ -12,12 +12,10 @@ from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
 from fastvideo.pipelines.lora_pipeline import LoRAPipeline
 
 # isort: off
-from fastvideo.pipelines.stages import (ImageEncodingStage, ConditioningStage,
-                                        DecodingStage, DmdDenoisingStage,
-                                        EncodingStage, InputValidationStage,
-                                        LatentPreparationStage,
-                                        TextEncodingStage,
-                                        TimestepPreparationStage)
+from fastvideo.pipelines.stages import (
+    ImageEncodingStage, ConditioningStage, DecodingStage, DmdDenoisingStage,
+    ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage,
+    TextEncodingStage, TimestepPreparationStage)
 # isort: on
 from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
     FlowMatchEulerDiscreteScheduler)
@@ -67,7 +65,7 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):
                            transformer=self.get_module("transformer")))
 
         self.add_stage(stage_name="image_latent_preparation_stage",
-                       stage=EncodingStage(vae=self.get_module("vae")))
+                       stage=ImageVAEEncodingStage(vae=self.get_module("vae")))
 
         self.add_stage(stage_name="denoising_stage",
                        stage=DmdDenoisingStage(

fastvideo/pipelines/basic/wan/wan_i2v_pipeline.py

@@ -12,12 +12,10 @@ from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
 from fastvideo.pipelines.lora_pipeline import LoRAPipeline
 
 # isort: off
-from fastvideo.pipelines.stages import (ImageEncodingStage, ConditioningStage,
-                                        DecodingStage, DenoisingStage,
-                                        EncodingStage, InputValidationStage,
-                                        LatentPreparationStage,
-                                        TextEncodingStage,
-                                        TimestepPreparationStage)
+from fastvideo.pipelines.stages import (
+    ImageEncodingStage, ConditioningStage, DecodingStage, DenoisingStage,
+    ImageVAEEncodingStage, InputValidationStage, LatentPreparationStage,
+    TextEncodingStage, TimestepPreparationStage)
 # isort: on
 from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import (
     FlowUniPCMultistepScheduler)
@@ -48,11 +46,14 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
                            tokenizers=[self.get_module("tokenizer")],
                        ))
 
-        self.add_stage(stage_name="image_encoding_stage",
-                       stage=ImageEncodingStage(
-                           image_encoder=self.get_module("image_encoder"),
-                           image_processor=self.get_module("image_processor"),
-                       ))
+        if (self.get_module("image_encoder") is not None
+                and self.get_module("image_processor") is not None):
+            self.add_stage(
+                stage_name="image_encoding_stage",
+                stage=ImageEncodingStage(
+                    image_encoder=self.get_module("image_encoder"),
+                    image_processor=self.get_module("image_processor"),
+                ))
 
         self.add_stage(stage_name="conditioning_stage",
                        stage=ConditioningStage())
@@ -67,11 +68,12 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
                            transformer=self.get_module("transformer")))
 
         self.add_stage(stage_name="image_latent_preparation_stage",
-                       stage=EncodingStage(vae=self.get_module("vae")))
+                       stage=ImageVAEEncodingStage(vae=self.get_module("vae")))
 
         self.add_stage(stage_name="denoising_stage",
                        stage=DenoisingStage(
                            transformer=self.get_module("transformer"),
+                           transformer_2=self.get_module("transformer_2"),
                            scheduler=self.get_module("scheduler")))
 
         self.add_stage(stage_name="decoding_stage",

fastvideo/pipelines/basic/wan/wan_pipeline.py

@@ -61,6 +61,7 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase):
         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))
 

fastvideo/pipelines/composed_pipeline_base.py

@@ -37,6 +37,7 @@ class ComposedPipelineBase(ABC):
 
     is_video_pipeline: bool = False  # To be overridden by video pipelines
     _required_config_modules: list[str] = []
+    _extra_config_module_map: dict[str, str] = {}
     training_args: TrainingArgs | None = None
     fastvideo_args: FastVideoArgs | TrainingArgs | None = None
     modules: dict[str, torch.nn.Module] = {}
@@ -148,9 +149,7 @@ class ComposedPipelineBase(ABC):
             # model is loaded with the correct precision. Subsequently we will
             # use FSDP2's MixedPrecisionPolicy to set the precision for the
             # fwd, bwd, and other operations' precision.
-            # fastvideo_args.precision = fastvideo_args.master_weight_type
             assert fastvideo_args.pipeline_config.dit_precision == 'fp32', 'only fp32 is supported for training'
-            # assert fastvideo_args.precision == 'fp32', 'only fp32 is supported for training'
 
         logger.info("fastvideo_args in from_pretrained: %s", fastvideo_args)
 
@@ -239,6 +238,18 @@ class ComposedPipelineBase(ABC):
         model_index.pop("_class_name")
         model_index.pop("_diffusers_version")
         # @TODO(Wei): Temporary hack
+        if "boundary_ratio" in model_index and model_index[
+                "boundary_ratio"] is not None:
+            logger.info(
+                "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..."
+            )
+            self.required_config_modules.append("transformer_2")
+            if fastvideo_args.boundary_ratio is None:
+                logger.info(
+                    "MoE pipeline detected. Setting boundary ratio to %s",
+                    model_index["boundary_ratio"])
+                fastvideo_args.boundary_ratio = model_index["boundary_ratio"]
+
         model_index.pop("boundary_ratio", None)
         model_index.pop("expand_timesteps", None)
 
@@ -248,12 +259,20 @@ class ComposedPipelineBase(ABC):
         ) > 1, "model_index.json must contain at least one pipeline module"
 
         for module_name in self.required_config_modules:
-            if module_name not in model_index:
+            if module_name not in model_index and module_name in self._extra_config_module_map:
+                extra_module_value = self._extra_config_module_map[module_name]
                 logger.warning(
-                    "model_index.json does not contain a %s module, adding %s to model_index",
-                    module_name, module_name)
-                if 'transformer' in module_name:
-                    model_index[module_name] = model_index['transformer']
+                    "model_index.json does not contain a %s module, but found {%s: %s} in _extra_config_module_map, adding to model_index.",
+                    module_name, module_name, extra_module_value)
+                if extra_module_value in model_index:
+                    logger.info("Using module %s for %s", extra_module_value,
+                                module_name)
+                    model_index[module_name] = model_index[extra_module_value]
+                    continue
+                else:
+                    raise ValueError(
+                        f"Required module key: {module_name} value: {model_index.get(module_name)} was not found in loaded modules {model_index.keys()}"
+                    )
 
         # all the component models used by the pipeline
         required_modules = self.required_config_modules
@@ -263,6 +282,7 @@ class ComposedPipelineBase(ABC):
         for module_name, (transformers_or_diffusers,
                           architecture) in model_index.items():
             if transformers_or_diffusers is None:
+                self.required_config_modules.remove(module_name)
                 continue
             if module_name not in required_modules:
                 logger.info("Skipping module %s", module_name)
@@ -271,14 +291,17 @@ class ComposedPipelineBase(ABC):
                 logger.info("Using module %s already provided", module_name)
                 modules[module_name] = loaded_modules[module_name]
                 continue
-            if 'transformer' in module_name:
-                loading_module_name = module_name.split("_")[-1]
+
+            # we load the module from the extra config module map if it exists
+            if module_name in self._extra_config_module_map:
+                load_module_name = self._extra_config_module_map[module_name]
             else:
-                loading_module_name = module_name
+                load_module_name = module_name
+
             component_model_path = os.path.join(self.model_path,
-                                                loading_module_name)
+                                                load_module_name)
             module = PipelineComponentLoader.load_module(
-                module_name=module_name,
+                module_name=load_module_name,
                 component_model_path=component_model_path,
                 transformers_or_diffusers=transformers_or_diffusers,
                 fastvideo_args=fastvideo_args,

fastvideo/pipelines/pipeline_batch_info.py

@@ -9,15 +9,55 @@ in a functional manner, reducing the need for explicit parameter passing.
 
 import pprint
 from dataclasses import asdict, dataclass, field
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import PIL.Image
 import torch
 
+if TYPE_CHECKING:
+    from torchcodec.decoders import VideoDecoder
+
+import time
+from collections import OrderedDict
+
 from fastvideo.attention import AttentionMetadata
 from fastvideo.configs.sample.teacache import TeaCacheParams, WanTeaCacheParams
 
 
+class PipelineLoggingInfo:
+    """Simple approach using OrderedDict to track stage metrics."""
+
+    def __init__(self):
+        # OrderedDict preserves insertion order and allows easy access
+        self.stages: OrderedDict[str, dict[str, Any]] = OrderedDict()
+
+    def add_stage_execution_time(self, stage_name: str, execution_time: float):
+        """Add execution time for a stage."""
+        if stage_name not in self.stages:
+            self.stages[stage_name] = {}
+        self.stages[stage_name]['execution_time'] = execution_time
+        self.stages[stage_name]['timestamp'] = time.time()
+
+    def add_stage_metric(self, stage_name: str, metric_name: str, value: Any):
+        """Add any metric for a stage."""
+        if stage_name not in self.stages:
+            self.stages[stage_name] = {}
+        self.stages[stage_name][metric_name] = value
+
+    def get_stage_info(self, stage_name: str) -> dict[str, Any]:
+        """Get all info for a specific stage."""
+        return self.stages.get(stage_name, {})
+
+    def get_execution_order(self) -> list[str]:
+        """Get stages in execution order."""
+        return list(self.stages.keys())
+
+    def get_total_execution_time(self) -> float:
+        """Get total pipeline execution time."""
+        return sum(
+            stage.get('execution_time', 0) for stage in self.stages.values())
+
+
 @dataclass
 class ForwardBatch:
     """
@@ -75,15 +115,15 @@ class ForwardBatch:
     image_latent: torch.Tensor | None = None
 
     # Latent dimensions
-    height_latents: int | None = None
-    width_latents: int | None = None
-    num_frames: int = 1  # Default for image models
+    height_latents: list[int] | int | None = None
+    width_latents: list[int] | int | None = None
+    num_frames: list[int] | int = 1  # Default for image models
     num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
 
     # Original dimensions (before VAE scaling)
-    height: int | None = None
-    width: int | None = None
-    fps: int | None = None
+    height: list[int] | int | None = None
+    width: list[int] | int | None = None
+    fps: list[int] | int | None = None
 
     # Timesteps
     timesteps: torch.Tensor | None = None
@@ -93,6 +133,7 @@ class ForwardBatch:
     # Scheduler parameters
     num_inference_steps: int = 50
     guidance_scale: float = 1.0
+    guidance_scale_2: float | None = None
     guidance_rescale: float = 0.0
     eta: float = 0.0
     sigmas: list[float] | None = None
@@ -128,6 +169,10 @@ class ForwardBatch:
     # VSA parameters
     VSA_sparsity: float = 0.0
 
+    # Logging info
+    logging_info: PipelineLoggingInfo = field(
+        default_factory=PipelineLoggingInfo)
+
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""
 
@@ -136,6 +181,8 @@ class ForwardBatch:
             self.do_classifier_free_guidance = True
         if self.negative_prompt_embeds is None:
             self.negative_prompt_embeds = []
+        if self.guidance_scale_2 is None:
+            self.guidance_scale_2 = self.guidance_scale
 
     def __str__(self):
         return pprint.pformat(asdict(self), indent=2, width=120)
@@ -189,4 +236,10 @@ class TrainingBatch:
     fake_score_loss: float = 0.0
 
     dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
-    fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
+    fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class PreprocessBatch(ForwardBatch):
+    video_loader: list["VideoDecoder"] = field(default_factory=list)
+    video_file_name: list[str] = field(default_factory=list)

fastvideo/pipelines/pipeline_registry.py

@@ -25,6 +25,11 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = {
     "HunyuanVideoPipeline": "hunyuan",
 }
 
+_PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = {
+    WorkloadType.I2V: "PreprocessPipelineI2V",
+    WorkloadType.T2V: "PreprocessPipelineT2V",
+}
+
 
 class PipelineType(str, Enum):
     """
@@ -65,15 +70,11 @@ class _PipelineRegistry:
         arch = _PIPELINE_NAME_TO_ARCHITECTURE_NAME[pipeline_name_in_config]
         return set(self.pipelines[pipeline_type.value][arch].keys())
 
-    def _load_preprocessing_pipeline_cls(
+    def _load_preprocess_pipeline_cls(
             self, workload_type: WorkloadType,
             arch: str) -> type[ComposedPipelineBase] | None:
-        if workload_type == WorkloadType.I2V:
-            pipeline_name = "I2VPreprocessPipeline"
-        elif workload_type == WorkloadType.T2V:
-            pipeline_name = "T2VPreprocessPipeline"
-        else:
-            raise ValueError(f"Invalid workload type: {workload_type.value}")
+        pipeline_name = _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME[
+            workload_type]
 
         return self.pipelines[
             PipelineType.PREPROCESS.value][arch][pipeline_name]
@@ -90,7 +91,7 @@ class _PipelineRegistry:
             return None
 
         if pipeline_type == PipelineType.PREPROCESS:
-            return self._load_preprocessing_pipeline_cls(workload_type, arch)
+            return self._load_preprocess_pipeline_cls(workload_type, arch)
         elif pipeline_type == PipelineType.BASIC:
             return self.pipelines[
                 pipeline_type.value][arch][pipeline_name_in_config]
@@ -131,7 +132,7 @@ def import_pipeline_classes(
     Import pipeline classes based on the pipeline type and workload type.
     
     Args:
-        pipeline_types: The pipeline types to load (basic, preprocessing, training). 
+        pipeline_types: The pipeline types to load (basic, preprocess, training). 
                       If None, loads all types.
     
     Returns:

fastvideo/pipelines/preprocess/preprocess_stages.py

@@ -0,0 +1,73 @@
+from collections.abc import Callable
+from typing import cast
+
+import numpy as np
+import torch
+from einops import rearrange
+from torchvision import transforms
+
+from fastvideo.dataset.transform import (CenterCropResizeVideo,
+                                         TemporalRandomCrop)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import (ForwardBatch,
+                                                     PreprocessBatch)
+from fastvideo.pipelines.stages.base import PipelineStage
+
+
+class VideoTransformStage(PipelineStage):
+    """
+    Crop a video in temporal dimension.
+    """
+
+    def __init__(self, train_fps: int, num_frames: int, max_height: int,
+                 max_width: int, do_temporal_sample: bool) -> None:
+        self.train_fps = train_fps
+        self.num_frames = num_frames
+        if do_temporal_sample:
+            self.temporal_sample_fn: Callable | None = TemporalRandomCrop(
+                num_frames)
+        else:
+            self.temporal_sample_fn = None
+
+        self.video_transform = transforms.Compose([
+            CenterCropResizeVideo((max_height, max_width)),
+        ])
+
+    def forward(self, batch: ForwardBatch,
+                fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        batch = cast(PreprocessBatch, batch)
+        assert isinstance(batch.fps, list)
+        assert isinstance(batch.num_frames, list)
+
+        if batch.data_type != "video":
+            return batch
+
+        if len(batch.video_loader) == 0:
+            raise ValueError("Video loader is not set")
+
+        video_pixel_batch = []
+
+        for i in range(len(batch.video_loader)):
+            frame_interval = batch.fps[i] / self.train_fps
+            start_frame_idx = 0
+            frame_indices = np.arange(start_frame_idx, batch.num_frames[i],
+                                      frame_interval).astype(int)
+            if len(frame_indices) > self.num_frames:
+                if self.temporal_sample_fn is not None:
+                    begin_index, end_index = self.temporal_sample_fn(
+                        len(frame_indices))
+                    frame_indices = frame_indices[begin_index:end_index]
+                else:
+                    frame_indices = frame_indices[:self.num_frames]
+
+            video = batch.video_loader[i].get_frames_at(frame_indices).data
+            video = self.video_transform(video)
+            video_pixel_batch.append(video)
+
+        video_pixel_values = torch.stack(video_pixel_batch)
+        video_pixel_values = rearrange(video_pixel_values,
+                                       "b t c h w -> b c t h w")
+        video_pixel_values = video_pixel_values.to(torch.uint8)
+        video_pixel_values = video_pixel_values.float() / 255.0
+        batch.latents = video_pixel_values
+        return cast(ForwardBatch, batch)

fastvideo/pipelines/preprocess/v1_preprocessing_new.py

@@ -0,0 +1,27 @@
+import os
+
+from fastvideo.distributed import (
+    maybe_init_distributed_environment_and_model_parallel)
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.utils import FlexibleArgumentParser
+from fastvideo.workflow.preprocess.preprocess_workflow_t2v import (
+    PreprocessWorkflowT2V)
+
+logger = init_logger(__name__)
+
+
+def main(fastvideo_args: FastVideoArgs) -> 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"
+    preprocess_workflow = PreprocessWorkflowT2V(fastvideo_args)
+    preprocess_workflow.run()
+
+
+if __name__ == "__main__":
+    parser = FlexibleArgumentParser()
+    parser = FastVideoArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    fastvideo_args = FastVideoArgs.from_cli_args(args)
+    main(fastvideo_args)

fastvideo/pipelines/preprocess/wan/wan_preprocess_pipelines.py

@@ -0,0 +1,53 @@
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.preprocess.preprocess_stages import VideoTransformStage
+from fastvideo.pipelines.stages import (EncodingStage, ImageEncodingStage,
+                                        TextEncodingStage)
+
+
+class PreprocessPipelineI2V(ComposedPipelineBase):
+    _required_config_modules = [
+        "image_encoder", "image_processor", "text_encoder", "tokenizer", "vae"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        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="image_encoding_stage",
+                       stage=ImageEncodingStage(
+                           image_encoder=self.get_module("image_encoder"),
+                           image_processor=self.get_module("image_processor"),
+                       ))
+        self.add_stage(stage_name="video_encoding_stage",
+                       stage=EncodingStage(vae=self.get_module("vae"), ))
+
+
+class PreprocessPipelineT2V(ComposedPipelineBase):
+    _required_config_modules = ["text_encoder", "tokenizer", "vae"]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        assert fastvideo_args.preprocess_config is not None
+
+        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="video_transform_stage",
+            stage=VideoTransformStage(
+                train_fps=fastvideo_args.preprocess_config.train_fps,
+                num_frames=fastvideo_args.preprocess_config.num_frames,
+                max_height=fastvideo_args.preprocess_config.max_height,
+                max_width=fastvideo_args.preprocess_config.max_width,
+                do_temporal_sample=fastvideo_args.preprocess_config.
+                do_temporal_sample,
+            ))
+        self.add_stage(stage_name="video_encoding_stage",
+                       stage=EncodingStage(vae=self.get_module("vae"), ))
+
+
+EntryClass = [PreprocessPipelineI2V, PreprocessPipelineT2V]

fastvideo/pipelines/stages/__init__.py

@@ -12,7 +12,8 @@ from fastvideo.pipelines.stages.decoding import DecodingStage
 from fastvideo.pipelines.stages.denoising import (DenoisingStage,
                                                   DmdDenoisingStage)
 from fastvideo.pipelines.stages.encoding import EncodingStage
-from fastvideo.pipelines.stages.image_encoding import ImageEncodingStage
+from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage,
+                                                       ImageVAEEncodingStage)
 from fastvideo.pipelines.stages.input_validation import InputValidationStage
 from fastvideo.pipelines.stages.latent_preparation import LatentPreparationStage
 from fastvideo.pipelines.stages.stepvideo_encoding import (
@@ -32,6 +33,7 @@ __all__ = [
     "EncodingStage",
     "DecodingStage",
     "ImageEncodingStage",
+    "ImageVAEEncodingStage",
     "TextEncodingStage",
     "StepvideoPromptEncodingStage",
 ]

fastvideo/pipelines/stages/base.py

@@ -155,6 +155,8 @@ class PipelineStage(ABC):
                 execution_time = time.perf_counter() - start_time
                 logger.info("[%s] Execution completed in %s ms", stage_name,
                             execution_time * 1000)
+                batch.logging_info.add_stage_execution_time(
+                    stage_name, execution_time)
             except Exception as e:
                 execution_time = time.perf_counter() - start_time
                 logger.error("[%s] Error during execution after %s ms: %s",

fastvideo/pipelines/stages/decoding.py

@@ -3,7 +3,6 @@
 Decoding stage for diffusion pipelines.
 """
 
-import gc
 import weakref
 
 import torch
@@ -140,8 +139,6 @@ class DecodingStage(PipelineStage):
             del self.vae
             if pipeline is not None and "vae" in pipeline.modules:
                 del pipeline.modules["vae"]
-            gc.collect()
-            torch.mps.empty_cache()
             fastvideo_args.model_loaded["vae"] = False
 
         return batch

fastvideo/pipelines/stages/denoising.py

@@ -3,7 +3,6 @@
 Denoising stage for diffusion pipelines.
 """
 
-import gc
 import inspect
 import weakref
 from collections.abc import Iterable
@@ -57,9 +56,14 @@ class DenoisingStage(PipelineStage):
     the initial noise into the final output.
     """
 
-    def __init__(self, transformer, scheduler, pipeline=None) -> None:
+    def __init__(self,
+                 transformer,
+                 scheduler,
+                 pipeline=None,
+                 transformer_2=None) -> None:
         super().__init__()
         self.transformer = transformer
+        self.transformer_2 = transformer_2
         self.scheduler = scheduler
         self.pipeline = weakref.ref(pipeline) if pipeline else None
         attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads
@@ -184,6 +188,12 @@ class DenoisingStage(PipelineStage):
             assert neg_prompt_embeds is not None
             assert torch.isnan(neg_prompt_embeds[0]).sum() == 0
 
+        # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
+        if fastvideo_args.boundary_ratio is not None:
+            boundary_timestep = fastvideo_args.boundary_ratio * self.scheduler.num_train_timesteps
+        else:
+            boundary_timestep = None
+
         # Run denoising loop
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -191,6 +201,23 @@ class DenoisingStage(PipelineStage):
                 if hasattr(self, 'interrupt') and self.interrupt:
                     continue
 
+                if boundary_timestep is None or t >= boundary_timestep:
+                    if (fastvideo_args.dit_cpu_offload
+                            and self.transformer_2 is not None and next(
+                                self.transformer_2.parameters()).device.type
+                            == 'cuda'):
+                        self.transformer_2.to('cpu')
+                    current_model = self.transformer
+                    current_guidance_scale = batch.guidance_scale
+                else:
+                    # low-noise stage in wan2.2
+                    if fastvideo_args.dit_cpu_offload and next(
+                            self.transformer.parameters(
+                            )).device.type == 'cuda':
+                        self.transformer.to('cpu')
+                    current_model = self.transformer_2
+                    current_guidance_scale = batch.guidance_scale_2
+
                 # Expand latents for I2V
                 latent_model_input = latents.to(target_dtype)
                 if batch.image_latent is not None:
@@ -257,7 +284,7 @@ class DenoisingStage(PipelineStage):
                             # fastvideo_args=fastvideo_args
                     ):
                         # Run transformer
-                        noise_pred = self.transformer(
+                        noise_pred = current_model(
                             latent_model_input,
                             prompt_embeds,
                             t_expand,
@@ -276,7 +303,7 @@ class DenoisingStage(PipelineStage):
                                 # fastvideo_args=fastvideo_args
                         ):
                             # Run transformer
-                            noise_pred_uncond = self.transformer(
+                            noise_pred_uncond = current_model(
                                 latent_model_input,
                                 neg_prompt_embeds,
                                 t_expand,
@@ -285,7 +312,7 @@ class DenoisingStage(PipelineStage):
                                 **neg_cond_kwargs,
                             )
                         noise_pred_text = noise_pred
-                        noise_pred = noise_pred_uncond + batch.guidance_scale * (
+                        noise_pred = noise_pred_uncond + current_guidance_scale * (
                             noise_pred_text - noise_pred_uncond)
 
                         # Apply guidance rescale if needed
@@ -296,14 +323,12 @@ class DenoisingStage(PipelineStage):
                                 noise_pred_text,
                                 guidance_rescale=batch.guidance_rescale,
                             )
-
                     # Compute the previous noisy sample
                     latents = self.scheduler.step(noise_pred,
                                                   t,
                                                   latents,
                                                   **extra_step_kwargs,
                                                   return_dict=False)[0]
-
                 # Update progress bar
                 if i == len(timesteps) - 1 or (
                     (i + 1) > num_warmup_steps and
@@ -329,8 +354,6 @@ class DenoisingStage(PipelineStage):
             del self.transformer
             if pipeline is not None and "transformer" in pipeline.modules:
                 del pipeline.modules["transformer"]
-            gc.collect()
-            torch.mps.empty_cache()
             fastvideo_args.model_loaded["transformer"] = False
             logger.info("Memory after deallocating transformer: %s",
                         torch.mps.current_allocated_memory())

fastvideo/pipelines/stages/encoding.py

@@ -3,15 +3,12 @@
 Encoding stage for diffusion pipelines.
 """
 
-import PIL.Image
 import torch
 
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.logger import init_logger
 from fastvideo.models.vaes.common import ParallelTiledVAE
-from fastvideo.models.vision_utils import (get_default_height_width, normalize,
-                                           numpy_to_pt, pil_to_numpy, resize)
 from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.pipelines.stages.base import PipelineStage
 from fastvideo.pipelines.stages.validators import V  # Import validators
@@ -23,64 +20,67 @@ logger = init_logger(__name__)
 
 class EncodingStage(PipelineStage):
     """
-    Stage for encoding pixel representations into latent space.
+    Stage for encoding pixel space representations into latent space.
     
-    This stage handles the encoding of pixel representations into the final
-    input format (e.g., latents).
+    This stage handles the encoding of pixel-space video/images into latent
+    representations for further processing in the diffusion pipeline.
     """
 
     def __init__(self, vae: ParallelTiledVAE) -> None:
         self.vae: ParallelTiledVAE = vae
 
     @torch.no_grad()
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify encoding stage inputs."""
+        result = VerificationResult()
+        # Input video/images for VAE encoding: [batch_size, channels, frames, height, width]
+        result.add_check("latents", batch.latents,
+                         [V.is_tensor, V.with_dims(5)])
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify encoding stage outputs."""
+        result = VerificationResult()
+        # Encoded latents: [batch_size, channels, frames, height_latents, width_latents]
+        result.add_check("latents", batch.latents,
+                         [V.is_tensor, V.with_dims(5)])
+        return result
+
     def forward(
         self,
         batch: ForwardBatch,
         fastvideo_args: FastVideoArgs,
     ) -> ForwardBatch:
         """
-        Encode pixel representations into latent space.
+        Encode pixel space representations into latent space.
         
         Args:
             batch: The current batch information.
             fastvideo_args: The inference arguments.
             
         Returns:
-            The batch with encoded outputs.
+            The batch with encoded latents.
         """
+        assert batch.latents is not None and isinstance(batch.latents,
+                                                        torch.Tensor)
+
         self.vae = self.vae.to(get_local_torch_device())
 
-        assert batch.height is not None
-        assert batch.width is not None
-        latent_height = batch.height // self.vae.spatial_compression_ratio
-        latent_width = batch.width // self.vae.spatial_compression_ratio
-
-        assert batch.pil_image is not None
-        image = batch.pil_image
-        image = self.preprocess(
-            image,
-            vae_scale_factor=self.vae.spatial_compression_ratio,
-            height=batch.height,
-            width=batch.width).to(get_local_torch_device(), dtype=torch.float32)
-
-        image = image.unsqueeze(2)
-
-        video_condition = torch.cat([
-            image,
-            image.new_zeros(image.shape[0], image.shape[1],
-                            batch.num_frames - 1, batch.height, batch.width)
-        ],
-                                    dim=2)
-        video_condition = video_condition.to(device=get_local_torch_device(),
-                                             dtype=torch.float32)
-
         # Setup VAE precision
         vae_dtype = PRECISION_TO_TYPE[
             fastvideo_args.pipeline_config.vae_precision]
         vae_autocast_enabled = (
             vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
 
-        # Encode Image
+        # Normalize input to [-1, 1] range (reverse of decoding normalization)
+        latents = (batch.latents * 2.0 - 1.0).clamp(-1, 1)
+
+        # Move to appropriate device and dtype
+        latents = latents.to(get_local_torch_device())
+
+        # Encode image to latents
         with torch.autocast(device_type="cuda",
                             dtype=vae_dtype,
                             enabled=vae_autocast_enabled):
@@ -89,47 +89,29 @@ class EncodingStage(PipelineStage):
             # if fastvideo_args.vae_sp:
             #     self.vae.enable_parallel()
             if not vae_autocast_enabled:
-                video_condition = video_condition.to(vae_dtype)
-            encoder_output = self.vae.encode(video_condition)
+                latents = latents.to(vae_dtype)
+            latents = self.vae.encode(latents).mean
 
-        generator = batch.generator
-        if generator is None:
-            raise ValueError("Generator must be provided")
-        latent_condition = self.retrieve_latents(encoder_output, generator)
-
-        # Apply shifting if needed
+        # 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):
-                latent_condition -= self.vae.shift_factor.to(
-                    latent_condition.device, latent_condition.dtype)
+                latents -= self.vae.shift_factor.to(latents.device,
+                                                    latents.dtype)
             else:
-                latent_condition -= self.vae.shift_factor
+                latents -= self.vae.shift_factor
 
-        if isinstance(self.vae.scaling_factor, torch.Tensor):
-            latent_condition = latent_condition * self.vae.scaling_factor.to(
-                latent_condition.device, latent_condition.dtype)
-        else:
-            latent_condition = latent_condition * self.vae.scaling_factor
+        # 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
 
-        mask_lat_size = torch.ones(1, 1, batch.num_frames, latent_height,
-                                   latent_width)
-        mask_lat_size[:, :, list(range(1, batch.num_frames))] = 0
-        first_frame_mask = mask_lat_size[:, :, 0:1]
-        first_frame_mask = torch.repeat_interleave(
-            first_frame_mask,
-            dim=2,
-            repeats=self.vae.temporal_compression_ratio)
-        mask_lat_size = torch.concat(
-            [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
-        mask_lat_size = mask_lat_size.view(1, -1,
-                                           self.vae.temporal_compression_ratio,
-                                           latent_height, latent_width)
-        mask_lat_size = mask_lat_size.transpose(1, 2)
-        mask_lat_size = mask_lat_size.to(latent_condition.device)
-
-        batch.image_latent = torch.concat([mask_lat_size, latent_condition],
-                                          dim=1)
+        # Update batch with encoded latents
+        batch.latents = latents
 
         # Offload models if needed
         if hasattr(self, 'maybe_free_model_hooks'):
@@ -139,61 +121,3 @@ class EncodingStage(PipelineStage):
             self.vae.to("cpu")
 
         return batch
-
-    def retrieve_latents(self,
-                         encoder_output: torch.Tensor,
-                         generator: torch.Generator | None = None,
-                         sample_mode: str = "sample"):
-        if sample_mode == "sample":
-            return encoder_output.sample(generator)
-        elif sample_mode == "argmax":
-            return encoder_output.mode()
-        else:
-            raise AttributeError(
-                "Could not access latents of provided encoder_output")
-
-    def preprocess(
-            self,
-            image: PIL.Image.Image,
-            vae_scale_factor: int,
-            height: int | None = None,
-            width: int | None = None,
-            resize_mode: str = "default",  # "default", "fill", "crop"
-    ) -> torch.Tensor:
-        image = [image]
-
-        height, width = get_default_height_width(image[0], vae_scale_factor,
-                                                 height, width)
-        image = [
-            resize(i, height, width, resize_mode=resize_mode) for i in image
-        ]
-        image = pil_to_numpy(image)  # to np
-        image = numpy_to_pt(image)  # to pt
-
-        do_normalize = True
-        if image.min() < 0:
-            do_normalize = False
-        if do_normalize:
-            image = normalize(image)
-
-        return image
-
-    def verify_input(self, batch: ForwardBatch,
-                     fastvideo_args: FastVideoArgs) -> VerificationResult:
-        """Verify encoding stage inputs."""
-        result = VerificationResult()
-        # result.add_check("pil_image", batch.pil_image)
-        result.add_check("height", batch.height, V.positive_int)
-        result.add_check("width", batch.width, V.positive_int)
-        result.add_check("generator", batch.generator,
-                         V.generator_or_list_generators)
-        result.add_check("num_frames", batch.num_frames, V.positive_int)
-        return result
-
-    def verify_output(self, batch: ForwardBatch,
-                      fastvideo_args: FastVideoArgs) -> VerificationResult:
-        """Verify encoding stage outputs."""
-        result = VerificationResult()
-        result.add_check("image_latent", batch.image_latent,
-                         [V.is_tensor, V.with_dims(5)])
-        return result

fastvideo/pipelines/stages/image_encoding.py

@@ -5,16 +5,21 @@ Image encoding stages for I2V diffusion pipelines.
 This module contains implementations of image encoding stages for diffusion pipelines.
 """
 
+import PIL
 import torch
 
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs
 from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
+from fastvideo.models.vaes.common import ParallelTiledVAE
+from fastvideo.models.vision_utils import (get_default_height_width, normalize,
+                                           numpy_to_pt, pil_to_numpy, resize)
 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.utils import PRECISION_TO_TYPE
 
 logger = init_logger(__name__)
 
@@ -87,3 +92,182 @@ class ImageEncodingStage(PipelineStage):
         result.add_check("image_embeds", batch.image_embeds,
                          V.list_of_tensors_dims(3))
         return result
+
+
+class ImageVAEEncodingStage(PipelineStage):
+    """
+    Stage for encoding pixel representations into latent space.
+    
+    This stage handles the encoding of pixel representations into the final
+    input format (e.g., latents).
+    """
+
+    def __init__(self, vae: ParallelTiledVAE) -> None:
+        self.vae: ParallelTiledVAE = vae
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """
+        Encode pixel representations into latent space.
+        
+        Args:
+            batch: The current batch information.
+            fastvideo_args: The inference arguments.
+            
+        Returns:
+            The batch with encoded outputs.
+        """
+        assert batch.height is not None and isinstance(batch.height, int)
+        assert batch.width is not None and isinstance(batch.width, int)
+        assert batch.num_frames is not None and isinstance(
+            batch.num_frames, int)
+
+        self.vae = self.vae.to(get_local_torch_device())
+
+        latent_height = batch.height // self.vae.spatial_compression_ratio
+        latent_width = batch.width // self.vae.spatial_compression_ratio
+
+        assert batch.pil_image is not None
+        image = batch.pil_image
+        image = self.preprocess(
+            image,
+            vae_scale_factor=self.vae.spatial_compression_ratio,
+            height=batch.height,
+            width=batch.width).to(get_local_torch_device(), dtype=torch.float32)
+
+        image = image.unsqueeze(2)
+
+        video_condition = torch.cat([
+            image,
+            image.new_zeros(image.shape[0], image.shape[1],
+                            batch.num_frames - 1, batch.height, batch.width)
+        ],
+                                    dim=2)
+        video_condition = video_condition.to(device=get_local_torch_device(),
+                                             dtype=torch.float32)
+
+        # Setup VAE precision
+        vae_dtype = PRECISION_TO_TYPE[
+            fastvideo_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (
+            vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
+
+        # Encode Image
+        with torch.autocast(device_type="cuda",
+                            dtype=vae_dtype,
+                            enabled=vae_autocast_enabled):
+            if fastvideo_args.pipeline_config.vae_tiling:
+                self.vae.enable_tiling()
+            # if fastvideo_args.vae_sp:
+            #     self.vae.enable_parallel()
+            if not vae_autocast_enabled:
+                video_condition = video_condition.to(vae_dtype)
+            encoder_output = self.vae.encode(video_condition)
+
+        generator = batch.generator
+        if generator is None:
+            raise ValueError("Generator must be provided")
+        latent_condition = self.retrieve_latents(encoder_output, generator)
+
+        # Apply shifting if needed
+        if (hasattr(self.vae, "shift_factor")
+                and self.vae.shift_factor is not None):
+            if isinstance(self.vae.shift_factor, torch.Tensor):
+                latent_condition -= self.vae.shift_factor.to(
+                    latent_condition.device, latent_condition.dtype)
+            else:
+                latent_condition -= self.vae.shift_factor
+
+        if isinstance(self.vae.scaling_factor, torch.Tensor):
+            latent_condition = latent_condition * self.vae.scaling_factor.to(
+                latent_condition.device, latent_condition.dtype)
+        else:
+            latent_condition = latent_condition * self.vae.scaling_factor
+
+        mask_lat_size = torch.ones(1, 1, batch.num_frames, latent_height,
+                                   latent_width)
+        mask_lat_size[:, :, list(range(1, batch.num_frames))] = 0
+        first_frame_mask = mask_lat_size[:, :, 0:1]
+        first_frame_mask = torch.repeat_interleave(
+            first_frame_mask,
+            dim=2,
+            repeats=self.vae.temporal_compression_ratio)
+        mask_lat_size = torch.concat(
+            [first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
+        mask_lat_size = mask_lat_size.view(1, -1,
+                                           self.vae.temporal_compression_ratio,
+                                           latent_height, latent_width)
+        mask_lat_size = mask_lat_size.transpose(1, 2)
+        mask_lat_size = mask_lat_size.to(latent_condition.device)
+
+        batch.image_latent = torch.concat([mask_lat_size, latent_condition],
+                                          dim=1)
+
+        # Offload models if needed
+        if hasattr(self, 'maybe_free_model_hooks'):
+            self.maybe_free_model_hooks()
+
+        self.vae.to("cpu")
+
+        return batch
+
+    def retrieve_latents(self,
+                         encoder_output: torch.Tensor,
+                         generator: torch.Generator | None = None,
+                         sample_mode: str = "sample"):
+        if sample_mode == "sample":
+            return encoder_output.sample(generator)
+        elif sample_mode == "argmax":
+            return encoder_output.mode()
+        else:
+            raise AttributeError(
+                "Could not access latents of provided encoder_output")
+
+    def preprocess(
+            self,
+            image: PIL.Image.Image,
+            vae_scale_factor: int,
+            height: int | None = None,
+            width: int | None = None,
+            resize_mode: str = "default",  # "default", "fill", "crop"
+    ) -> torch.Tensor:
+        image = [image]
+
+        height, width = get_default_height_width(image[0], vae_scale_factor,
+                                                 height, width)
+        image = [
+            resize(i, height, width, resize_mode=resize_mode) for i in image
+        ]
+        image = pil_to_numpy(image)  # to np
+        image = numpy_to_pt(image)  # to pt
+
+        do_normalize = True
+        if image.min() < 0:
+            do_normalize = False
+        if do_normalize:
+            image = normalize(image)
+
+        return image
+
+    def verify_input(self, batch: ForwardBatch,
+                     fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify encoding stage inputs."""
+        result = VerificationResult()
+        # result.add_check("pil_image", batch.pil_image)
+        result.add_check("height", batch.height, V.positive_int)
+        result.add_check("width", batch.width, V.positive_int)
+        result.add_check("generator", batch.generator,
+                         V.generator_or_list_generators)
+        result.add_check("num_frames", batch.num_frames, V.positive_int)
+        return result
+
+    def verify_output(self, batch: ForwardBatch,
+                      fastvideo_args: FastVideoArgs) -> VerificationResult:
+        """Verify encoding stage outputs."""
+        result = VerificationResult()
+        result.add_check("image_latent", batch.image_latent,
+                         [V.is_tensor, V.with_dims(5)])
+        return result

fastvideo/pipelines/stages/input_validation.py

@@ -56,6 +56,7 @@ class InputValidationStage(PipelineStage):
         Returns:
             The validated batch information.
         """
+
         self._generate_seeds(batch, fastvideo_args)
 
         # Ensure prompt is properly formatted

fastvideo/pipelines/stages/latent_preparation.py

@@ -90,7 +90,6 @@ class LatentPreparationStage(PipelineStage):
                 f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                 f" size of {batch_size}. Make sure the batch size matches the length of the generators."
             )
-
         # Generate or use provided latents
         if latents is None:
             latents = randn_tensor(shape,
@@ -103,7 +102,6 @@ class LatentPreparationStage(PipelineStage):
         # Scale the initial noise if needed
         if hasattr(self.scheduler, "init_noise_sigma"):
             latents = latents * self.scheduler.init_noise_sigma
-
         # Update batch with prepared latents
         batch.latents = latents
         batch.raw_latent_shape = latents.shape
@@ -122,6 +120,7 @@ class LatentPreparationStage(PipelineStage):
         Returns:
             The batch with adjusted video length.
         """
+
         video_length = batch.num_frames
         use_temporal_scaling_frames = fastvideo_args.pipeline_config.vae_config.use_temporal_scaling_frames
         if use_temporal_scaling_frames:

fastvideo/tests/modal/pr_test.py

@@ -11,7 +11,7 @@ print(f"Using image: {image_tag}")
 image = (
     modal.Image.from_registry(image_tag, add_python="3.12")
     .run_commands("rm -rf /FastVideo")
-    .apt_install("cmake", "pkg-config", "build-essential", "curl", "libssl-dev")
+    .apt_install("cmake", "pkg-config", "build-essential", "curl", "libssl-dev", "ffmpeg")
     .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable")
     .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
     .env({
@@ -100,7 +100,7 @@ def run_precision_tests_STA():
 
 @app.function(gpu="H100:1", image=image, timeout=900)
 def run_precision_tests_VSA():
-    run_test("python csrc/attn/tests/test_block_sparse.py")
+    run_test("python csrc/attn/tests/test_vsa.py")
 
 @app.function(gpu="L40S:1", image=image, timeout=3600)
 def run_inference_lora_tests():

fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py

@@ -116,7 +116,6 @@ def run_training():
         "--validation_sampling_steps", "50",
         "--log_validation",
         "--checkpoints_total_limit", "3",
-        "--allow_tf32",
         "--ema_start_step", "0",
         "--training_cfg_rate", "0.1",
         "--output_dir", LOCAL_OUTPUT_DIR,

fastvideo/tests/nightly/test_e2e_overfit_single_sample.py

@@ -116,7 +116,6 @@ def run_training():
         "--validation_sampling_steps", "50",
         "--log_validation",
         "--checkpoints_total_limit", "3",
-        "--allow_tf32",
         "--ema_start_step", "0",
         "--training_cfg_rate", "0.0",
         "--output_dir", LOCAL_OUTPUT_DIR,

fastvideo/tests/training/VSA/test_training_loss_VSA.py

@@ -51,7 +51,6 @@ def run_worker():
         "--validation_sampling_steps", "50",
         "--log_validation",
         "--checkpoints_total_limit", "3",
-        "--allow_tf32",
         "--ema_start_step", "0",
         "--training_cfg_rate", "0.0",
         "--output_dir", "data/wan_finetune_test_VSA",

fastvideo/tests/training/Vanilla/test_training_loss.py

@@ -55,7 +55,6 @@ def run_worker():
         "--validation_sampling_steps", "8",
         "--log_validation",
         "--checkpoints_total_limit", "3",
-        "--allow_tf32",
         "--ema_start_step", "0",
         "--training_cfg_rate", "0.0",
         "--output_dir", "data/wan_finetune_test",

fastvideo/tests/training/distill/test_distill_dmd.py

@@ -54,7 +54,6 @@ def run_worker():
         "--validation_sampling_steps", "3",
         "--log_validation",
         "--checkpoints_total_limit", "3",
-        "--allow_tf32",
         "--ema_start_step", "0",
         "--training_cfg_rate", "0.0",
         "--output_dir", "data/wan_finetune_test",

fastvideo/tests/training/lora/test_lora_training.py

@@ -60,7 +60,6 @@ def test_lora_training():
         "--validation_sampling_steps", "50",
         "--log_validation",
         "--checkpoints_total_limit", "3",
-        "--allow_tf32",
         "--ema_start_step", "0",
         "--training_cfg_rate", "0.1",
         "--output_dir", "/workspace",

fastvideo/training/distillation_pipeline.py

@@ -56,6 +56,10 @@ class DistillationPipeline(TrainingPipeline):
         "scheduler", "transformer", "vae", "real_score_transformer",
         "fake_score_transformer"
     ]
+    _extra_config_module_map = {
+        "real_score_transformer": "transformer",
+        "fake_score_transformer": "transformer"
+    }
     validation_pipeline: ComposedPipelineBase
     train_dataloader: StatefulDataLoader
     train_loader_iter: Iterator[dict[str, Any]]
@@ -237,6 +241,7 @@ class DistillationPipeline(TrainingPipeline):
                                             len(self.denoising_step_list), [1],
                                             device=self.device,
                                             dtype=torch.long)
+        target_timestep_idx_int = target_timestep_idx.item()
         target_timestep = self.denoising_step_list[target_timestep_idx]
 
         # Step 2: Simulate the multi-step inference process up to the target timestep
@@ -254,6 +259,8 @@ class DistillationPipeline(TrainingPipeline):
 
         # Only run intermediate steps if target_timestep_idx > 0
         max_target_idx = len(self.denoising_step_list) - 1
+        noise_latents = []
+        noise_latent_index = target_timestep_idx_int - 1
         if max_target_idx > 0:
             # Run student model for all steps before the target timestep
             with torch.no_grad():
@@ -290,11 +297,19 @@ class DistillationPipeline(TrainingPipeline):
                     current_noise_latents = self.noise_scheduler.add_noise(
                         pred_clean.flatten(0, 1), noise.flatten(0, 1),
                         next_timestep_tensor).unflatten(0, pred_clean.shape[:2])
+                    latent_copy = current_noise_latents.clone()
+                    noise_latents.append(latent_copy)
 
         # Step 3: Use the simulated noisy input for the final training step
         # For timestep index 0, this is pure noise
         # For timestep index k > 0, this is the result after k denoising steps + noise at target level
-        noisy_input = current_noise_latents
+        if noise_latent_index >= 0:
+            assert noise_latent_index < len(
+                self.denoising_step_list
+            ) - 1, "noise_latent_index is out of bounds"
+            noisy_input = noise_latents[noise_latent_index]
+        else:
+            noisy_input = current_noise_latents
 
         # Step 4: Final student prediction (this is what we train on)
         training_batch = self._build_distill_input_kwargs(
@@ -415,7 +430,7 @@ class DistillationPipeline(TrainingPipeline):
         with torch.no_grad(), set_forward_context(
                 current_timestep=training_batch.timesteps,
                 attn_metadata=training_batch.attn_metadata_vsa):
-            if self.training_args.simulate_student_forward:
+            if self.training_args.simulate_generator_forward:
                 generator_pred_video = self._generator_multi_step_simulation_forward(
                     training_batch)
             else:
@@ -550,7 +565,7 @@ class DistillationPipeline(TrainingPipeline):
                 with set_forward_context(
                         current_timestep=batch_gen.timesteps,
                         attn_metadata=batch_gen.attn_metadata_vsa):
-                    if self.training_args.simulate_student_forward:
+                    if self.training_args.simulate_generator_forward:
                         generator_pred_video = self._generator_multi_step_simulation_forward(
                             batch_gen)
                     else:
@@ -790,17 +805,15 @@ class DistillationPipeline(TrainingPipeline):
         # Re-enable gradients for training
         transformer.train()
         gc.collect()
-        torch.cuda.empty_cache()
 
     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 = {}
-        torch.cuda.empty_cache()
         dmd_latents_vis_dict = training_batch.dmd_latent_vis_dict
         fake_score_latents_vis_dict = training_batch.fake_score_latent_vis_dict
         fake_score_log_keys = ['generator_pred_video']
-        dmd_log_keys = ['faker_score_pred_video']
+        dmd_log_keys = ['faker_score_pred_video', 'real_score_pred_video']
 
         for latent_key in fake_score_log_keys:
             latents = fake_score_latents_vis_dict[latent_key]
@@ -830,7 +843,6 @@ class DistillationPipeline(TrainingPipeline):
                 video, fps=24, format="mp4")  # change to 16 for Wan2.1
             # Clean up references
             del video, latents
-            torch.cuda.empty_cache()
 
         # Process DMD training data if available - use decode_stage instead of self.vae.decode
         if 'generator_pred_video' in dmd_latents_vis_dict:
@@ -862,7 +874,6 @@ class DistillationPipeline(TrainingPipeline):
                     video, fps=24, format="mp4")  # change to 16 for Wan2.1
                 # Clean up references
                 del video, latents
-                torch.cuda.empty_cache()
 
         # Log to wandb
         if self.global_rank == 0:

fastvideo/training/training_pipeline.py

@@ -64,6 +64,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
     train_dataloader: StatefulDataLoader
     train_loader_iter: Iterator[dict[str, Any]]
     current_epoch: int = 0
+    train_transformer_2: bool = False
 
     def __init__(
             self,
@@ -76,6 +77,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         if self.lora_training and fastvideo_args.lora_rank is None:
             raise ValueError("lora rank must be set when using lora training")
 
+        set_random_seed(fastvideo_args.seed)  # for lora param init
         super().__init__(model_path, fastvideo_args, required_config_modules,
                          loaded_modules)  # type: ignore
 
@@ -98,6 +100,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         self.sp_world_size = self.sp_group.world_size
         self.local_rank = world_group.local_rank
         self.transformer = self.get_module("transformer")
+        self.transformer_2 = self.get_module("transformer_2", None)
         self.seed = training_args.seed
         self.set_schemas()
 
@@ -110,10 +113,16 @@ class TrainingPipeline(LoRAPipeline, ABC):
                 self.transformer,
                 checkpointing_type=training_args.
                 enable_gradient_checkpointing_type)
+            if self.transformer_2 is not None:
+                self.transformer_2 = apply_activation_checkpointing(
+                    self.transformer_2,
+                    checkpointing_type=training_args.
+                    enable_gradient_checkpointing_type)
 
         noise_scheduler = self.modules["scheduler"]
         self.set_trainable()
         params_to_optimize = self.transformer.parameters()
+
         params_to_optimize = list(
             filter(lambda p: p.requires_grad, params_to_optimize))
         self.optimizer = torch.optim.AdamW(
@@ -137,6 +146,30 @@ class TrainingPipeline(LoRAPipeline, ABC):
             min_lr_ratio=training_args.min_lr_ratio,
             last_epoch=self.init_steps - 1,
         )
+        if self.transformer_2 is not None:
+            # Ensure transformer_2 has trainable parameters before creating optimizer
+            self.transformer_2.train()
+            self.transformer_2.requires_grad_(True)
+            params_to_optimize_2 = self.transformer_2.parameters()
+            params_to_optimize_2 = list(
+                filter(lambda p: p.requires_grad, params_to_optimize_2))
+            self.optimizer_2 = torch.optim.AdamW(
+                params_to_optimize_2,
+                lr=training_args.learning_rate,
+                betas=(0.9, 0.999),
+                weight_decay=training_args.weight_decay,
+                eps=1e-8,
+            )
+            self.lr_scheduler_2 = get_scheduler(
+                training_args.lr_scheduler,
+                optimizer=self.optimizer_2,
+                num_warmup_steps=training_args.lr_warmup_steps,
+                num_training_steps=training_args.max_train_steps,
+                num_cycles=training_args.lr_num_cycles,
+                power=training_args.lr_power,
+                min_lr_ratio=training_args.min_lr_ratio,
+                last_epoch=self.init_steps - 1,
+            )
 
         self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader(
             training_args.data_path,
@@ -151,7 +184,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
             seed=self.seed)
 
         self.noise_scheduler = noise_scheduler
-
+        self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
         self.num_update_steps_per_epoch = math.ceil(
             len(self.train_dataloader) /
             training_args.gradient_accumulation_steps * training_args.sp_size /
@@ -177,6 +210,9 @@ class TrainingPipeline(LoRAPipeline, ABC):
     def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch:
         self.transformer.train()
         self.optimizer.zero_grad()
+        if self.transformer_2 is not None:
+            self.transformer_2.train()
+            self.optimizer_2.zero_grad()
         training_batch.total_loss = 0.0
         return training_batch
 
@@ -223,17 +259,8 @@ class TrainingPipeline(LoRAPipeline, ABC):
                             generator=self.noise_gen_cuda,
                             device=latents.device,
                             dtype=latents.dtype)
-        u = compute_density_for_timestep_sampling(
-            weighting_scheme=self.training_args.weighting_scheme,
-            batch_size=batch_size,
-            generator=self.noise_random_generator,
-            logit_mean=self.training_args.logit_mean,
-            logit_std=self.training_args.logit_std,
-            mode_scale=self.training_args.mode_scale,
-        )
-        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
-        timesteps = self.noise_scheduler.timesteps[indices].to(
-            device=latents.device)
+        timesteps = self._sample_timesteps(batch_size, latents.device)
+
         if self.training_args.sp_size > 1:
             # Make sure that the timesteps are the same across all sp processes.
             sp_group = get_sp_group()
@@ -256,6 +283,38 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         return training_batch
 
+    def _sample_timesteps(self, batch_size, device):
+        # Determine which model to train based on the boundary timestep
+        if (self.transformer_2 is not None and self.boundary_timestep is not None and
+                torch.rand(1, generator=self.noise_random_generator).item() > self.training_args.boundary_ratio):
+            self.train_transformer_2 = True
+        else:
+            self.train_transformer_2 = False
+
+        # Broadcast the decision to all processes
+        decision = torch.tensor(1.0 if self.train_transformer_2 else 0.0, device=self.device)
+        dist.broadcast(decision, src=0)
+        self.train_transformer_2 = decision.item() == 1.0
+
+        # Sample u from the appropriate range
+        u = compute_density_for_timestep_sampling(
+            weighting_scheme=self.training_args.weighting_scheme,
+            batch_size=batch_size,
+            generator=self.noise_random_generator,
+            logit_mean=self.training_args.logit_mean,
+            logit_std=self.training_args.logit_std,
+            mode_scale=self.training_args.mode_scale,
+        )
+
+        boundary_ratio = self.training_args.boundary_ratio
+        if self.train_transformer_2:
+            u = boundary_ratio + u * (1.0 - boundary_ratio)
+        else:
+            u = u * boundary_ratio
+
+        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
+        return self.noise_scheduler.timesteps[indices].to(device=device)
+        
     def _build_attention_metadata(
             self, training_batch: TrainingBatch) -> TrainingBatch:
         latents_shape = training_batch.raw_latent_shape
@@ -306,11 +365,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
         #         [1000.0],
         #         device=training_batch.noisy_model_input.device,
         #         dtype=torch.bfloat16)
+        current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
 
         with set_forward_context(
                 current_timestep=training_batch.current_timestep,
                 attn_metadata=training_batch.attn_metadata):
-            model_pred = self.transformer(**input_kwargs)
+            model_pred = current_model(**input_kwargs)
             if self.training_args.precondition_outputs:
                 assert training_batch.sigmas is not None
                 model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
@@ -341,7 +401,12 @@ class TrainingPipeline(LoRAPipeline, ABC):
         # the following:
         # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
         if max_grad_norm is not None:
-            model_parts = [self.transformer]
+            # Only clip gradients for the model that is currently training
+            if self.train_transformer_2 and self.transformer_2 is not None:
+                model_parts = [self.transformer_2]
+            else:
+                model_parts = [self.transformer]
+            
             grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
                 [p for m in model_parts for p in m.parameters()],
                 max_grad_norm,
@@ -386,9 +451,14 @@ class TrainingPipeline(LoRAPipeline, ABC):
 
         training_batch = self._clip_grad_norm(training_batch)
 
-        self.optimizer.step()
-        self.lr_scheduler.step()
-
+        # Only step the optimizer and scheduler for the model that is currently training
+        if self.train_transformer_2 and self.transformer_2 is not None:
+            self.optimizer_2.step()
+            self.lr_scheduler_2.step()
+        else:
+            self.optimizer.step()
+            self.lr_scheduler.step()
+            
         training_batch.total_loss = training_batch.total_loss
         training_batch.grad_norm = training_batch.grad_norm
         return training_batch
@@ -439,6 +509,7 @@ class TrainingPipeline(LoRAPipeline, ABC):
         step_times: deque[float] = deque(maxlen=100)
 
         self._log_training_info()
+
         self._log_validation(self.transformer, self.training_args,
                              self.init_steps)
 

fastvideo/training/training_pipeline_switch.py

@@ -0,0 +1,797 @@
+# SPDX-License-Identifier: Apache-2.0
+import dataclasses
+import math
+import os
+import time
+from abc import ABC, abstractmethod
+from collections import deque
+from collections.abc import Iterator
+from typing import Any
+
+import imageio
+import numpy as np
+import torch
+import torch.distributed as dist
+import torchvision
+from diffusers import FlowMatchEulerDiscreteScheduler
+from einops import rearrange
+from torch.utils.data import DataLoader
+from torchdata.stateful_dataloader import StatefulDataLoader
+from tqdm.auto import tqdm
+
+import fastvideo.envs as envs
+from fastvideo.attention.backends.video_sparse_attn import (
+    VideoSparseAttentionMetadataBuilder)
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset import build_parquet_map_style_dataloader
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
+from fastvideo.dataset.validation_dataset import ValidationDataset
+from fastvideo.distributed import (cleanup_dist_env_and_memory,
+                                   get_local_torch_device, get_sp_group,
+                                   get_world_group)
+from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.pipelines import (ComposedPipelineBase, ForwardBatch,
+                                 LoRAPipeline, TrainingBatch)
+from fastvideo.training.activation_checkpoint import (
+    apply_activation_checkpointing)
+from fastvideo.training.training_utils import (
+    clip_grad_norm_while_handling_failing_dtensor_cases,
+    compute_density_for_timestep_sampling, get_scheduler, get_sigmas,
+    load_checkpoint, normalize_dit_input, save_checkpoint,
+    shard_latents_across_sp)
+from fastvideo.utils import is_vsa_available, set_random_seed, shallow_asdict
+
+import wandb  # isort: skip
+
+vsa_available = is_vsa_available()
+
+logger = init_logger(__name__)
+
+
+def _get_trainable_params(model: torch.nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+
+class TrainingPipeline(LoRAPipeline, ABC):
+    """
+    A pipeline for training a model. All training pipelines should inherit from this class.
+    All reusable components and code should be implemented in this class.
+    """
+    _required_config_modules = ["scheduler", "transformer"]
+    validation_pipeline: ComposedPipelineBase
+    train_dataloader: StatefulDataLoader
+    train_loader_iter: Iterator[dict[str, Any]]
+    current_epoch: int = 0
+    train_transformer_2: bool = False
+
+    def __init__(
+            self,
+            model_path: str,
+            fastvideo_args: TrainingArgs,
+            required_config_modules: list[str] | None = None,
+            loaded_modules: dict[str, torch.nn.Module] | None = None) -> None:
+        fastvideo_args.inference_mode = False
+        self.lora_training = fastvideo_args.lora_training
+        if self.lora_training and fastvideo_args.lora_rank is None:
+            raise ValueError("lora rank must be set when using lora training")
+
+        set_random_seed(fastvideo_args.seed)  # for lora param init
+        super().__init__(model_path, fastvideo_args, required_config_modules,
+                         loaded_modules)  # type: ignore
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        raise RuntimeError(
+            "create_pipeline_stages should not be called for training pipeline")
+
+    def set_schemas(self) -> None:
+        self.train_dataset_schema = pyarrow_schema_t2v
+
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
+        logger.info("Initializing training pipeline...")
+        self.device = get_local_torch_device()
+        self.training_args = training_args
+        world_group = get_world_group()
+        self.world_size = world_group.world_size
+        self.global_rank = world_group.rank
+        self.sp_group = get_sp_group()
+        self.rank_in_sp_group = self.sp_group.rank_in_group
+        self.sp_world_size = self.sp_group.world_size
+        self.local_rank = world_group.local_rank
+        self.transformer = self.get_module("transformer")
+        self.transformer_2 = self.get_module("transformer_2", None)
+        self.seed = training_args.seed
+        self.set_schemas()
+
+        # Set random seeds for deterministic training
+        assert self.seed is not None, "seed must be set"
+        set_random_seed(self.seed)
+        self.transformer.train()
+        if training_args.enable_gradient_checkpointing_type is not None:
+            self.transformer = apply_activation_checkpointing(
+                self.transformer,
+                checkpointing_type=training_args.
+                enable_gradient_checkpointing_type)
+            if self.transformer_2 is not None:
+                self.transformer_2 = apply_activation_checkpointing(
+                    self.transformer_2,
+                    checkpointing_type=training_args.
+                    enable_gradient_checkpointing_type)
+
+        noise_scheduler = self.modules["scheduler"]
+        self.set_trainable()
+        params_to_optimize = self.transformer.parameters()
+
+        params_to_optimize = list(
+            filter(lambda p: p.requires_grad, params_to_optimize))
+        self.optimizer = torch.optim.AdamW(
+            params_to_optimize,
+            lr=training_args.learning_rate,
+            betas=(0.9, 0.999),
+            weight_decay=training_args.weight_decay,
+            eps=1e-8,
+        )
+
+        self.init_steps = 0
+        logger.info("optimizer: %s", self.optimizer)
+
+        self.lr_scheduler = get_scheduler(
+            training_args.lr_scheduler,
+            optimizer=self.optimizer,
+            num_warmup_steps=training_args.lr_warmup_steps,
+            num_training_steps=training_args.max_train_steps,
+            num_cycles=training_args.lr_num_cycles,
+            power=training_args.lr_power,
+            min_lr_ratio=training_args.min_lr_ratio,
+            last_epoch=self.init_steps - 1,
+        )
+        if self.transformer_2 is not None:
+            # Ensure transformer_2 has trainable parameters before creating optimizer
+            self.transformer_2.train()
+            self.transformer_2.requires_grad_(True)
+            params_to_optimize_2 = self.transformer_2.parameters()
+            params_to_optimize_2 = list(
+                filter(lambda p: p.requires_grad, params_to_optimize_2))
+            self.optimizer_2 = torch.optim.AdamW(
+                params_to_optimize_2,
+                lr=training_args.learning_rate,
+                betas=(0.9, 0.999),
+                weight_decay=training_args.weight_decay,
+                eps=1e-8,
+            )
+            self.lr_scheduler_2 = get_scheduler(
+                training_args.lr_scheduler,
+                optimizer=self.optimizer_2,
+                num_warmup_steps=training_args.lr_warmup_steps,
+                num_training_steps=training_args.max_train_steps,
+                num_cycles=training_args.lr_num_cycles,
+                power=training_args.lr_power,
+                min_lr_ratio=training_args.min_lr_ratio,
+                last_epoch=self.init_steps - 1,
+            )
+
+        self.train_dataset, self.train_dataloader = build_parquet_map_style_dataloader(
+            training_args.data_path,
+            training_args.train_batch_size,
+            parquet_schema=self.train_dataset_schema,
+            num_data_workers=training_args.dataloader_num_workers,
+            cfg_rate=training_args.training_cfg_rate,
+            drop_last=True,
+            text_padding_length=training_args.pipeline_config.
+            text_encoder_configs[0].arch_config.
+            text_len,  # type: ignore[attr-defined]
+            seed=self.seed)
+
+        self.noise_scheduler = noise_scheduler
+        self.boundary_timestep = self.training_args.boundary_ratio * self.noise_scheduler.num_train_timesteps
+        self.num_update_steps_per_epoch = math.ceil(
+            len(self.train_dataloader) /
+            training_args.gradient_accumulation_steps * training_args.sp_size /
+            training_args.train_sp_batch_size)
+        self.num_train_epochs = math.ceil(training_args.max_train_steps /
+                                          self.num_update_steps_per_epoch)
+
+        # TODO(will): is there a cleaner way to track epochs?
+        self.current_epoch = 0
+
+        if self.global_rank == 0:
+            project = training_args.tracker_project_name or "fastvideo"
+            wandb_config = dataclasses.asdict(training_args)
+            wandb.init(project=project,
+                       config=wandb_config,
+                       name=training_args.wandb_run_name)
+
+    @abstractmethod
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
+        raise NotImplementedError(
+            "Training pipelines must implement this method")
+
+    def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch:
+        # At the beginning, disable training for all models
+        self._disable_training(self.transformer, self.optimizer)
+        if self.transformer_2 is not None:
+            self._disable_training(self.transformer_2, self.optimizer_2)
+        
+        training_batch.total_loss = 0.0
+        return training_batch
+
+    def _enable_training(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        """Enable training mode and gradients for the specified model."""
+        for param in model.parameters():
+            param.requires_grad = True
+        model.train()
+        optimizer.zero_grad()
+
+    def _disable_training(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer) -> None:
+        """Disable training mode and gradients for the specified model."""
+        for param in model.parameters():
+            param.requires_grad = False
+        optimizer.zero_grad(set_to_none=True)
+
+    def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
+        batch = next(self.train_loader_iter, None)  # type: ignore
+        if batch is None:
+            self.current_epoch += 1
+            logger.info("Starting epoch %s", self.current_epoch)
+            # Reset iterator for next epoch
+            self.train_loader_iter = iter(self.train_dataloader)
+            # Get first batch of new epoch
+            batch = next(self.train_loader_iter)
+
+        # latents, encoder_hidden_states, encoder_attention_mask, infos = batch
+        latents = batch['vae_latent']
+        latents = latents[:, :, :self.training_args.num_latent_t]
+        encoder_hidden_states = batch['text_embedding']
+        encoder_attention_mask = batch['text_attention_mask']
+        infos = batch['info_list']
+
+        training_batch.latents = latents.to(get_local_torch_device(),
+                                            dtype=torch.bfloat16)
+        training_batch.encoder_hidden_states = encoder_hidden_states.to(
+            get_local_torch_device(), dtype=torch.bfloat16)
+        training_batch.encoder_attention_mask = encoder_attention_mask.to(
+            get_local_torch_device(), dtype=torch.bfloat16)
+        training_batch.infos = infos
+
+        return training_batch
+
+    def _normalize_dit_input(self,
+                             training_batch: TrainingBatch) -> TrainingBatch:
+        # TODO(will): support other models
+        training_batch.latents = normalize_dit_input('wan',
+                                                     training_batch.latents,
+                                                     self.get_module("vae"))
+        return training_batch
+
+    def _prepare_dit_inputs(self,
+                            training_batch: TrainingBatch) -> TrainingBatch:
+        latents = training_batch.latents
+        batch_size = latents.shape[0]
+        noise = torch.randn(latents.shape,
+                            generator=self.noise_gen_cuda,
+                            device=latents.device,
+                            dtype=latents.dtype)
+        timesteps = self._sample_timesteps(batch_size, latents.device)
+
+        # Enable training for the model that will be trained next and disable the other
+        if self.train_transformer_2:
+            self._enable_training(self.transformer_2, self.optimizer_2)
+            self._disable_training(self.transformer, self.optimizer)
+        else:
+            self._enable_training(self.transformer, self.optimizer)
+            if self.transformer_2 is not None:
+                self._disable_training(self.transformer_2, self.optimizer_2)
+
+        if self.training_args.sp_size > 1:
+            # Make sure that the timesteps are the same across all sp processes.
+            sp_group = get_sp_group()
+            sp_group.broadcast(timesteps, src=0)
+        sigmas = get_sigmas(
+            self.noise_scheduler,
+            latents.device,
+            timesteps,
+            n_dim=latents.ndim,
+            dtype=latents.dtype,
+        )
+        noisy_model_input = (1.0 -
+                             sigmas) * training_batch.latents + sigmas * noise
+
+        training_batch.noisy_model_input = noisy_model_input
+        training_batch.timesteps = timesteps
+        training_batch.sigmas = sigmas
+        training_batch.noise = noise
+        training_batch.raw_latent_shape = training_batch.latents.shape
+
+        return training_batch
+
+    def _sample_timesteps(self, batch_size, device):
+        # Determine which model to train based on the boundary timestep
+        if (self.transformer_2 is not None and self.boundary_timestep is not None and
+                torch.rand(1, generator=self.noise_random_generator).item() > self.training_args.boundary_ratio):
+            self.train_transformer_2 = True
+        else:
+            self.train_transformer_2 = False
+
+        # Broadcast the decision to all processes
+        decision = torch.tensor(1.0 if self.train_transformer_2 else 0.0, device=self.device)
+        dist.broadcast(decision, src=0)
+        self.train_transformer_2 = decision.item() == 1.0
+
+        # Sample u from the appropriate range
+        u = compute_density_for_timestep_sampling(
+            weighting_scheme=self.training_args.weighting_scheme,
+            batch_size=batch_size,
+            generator=self.noise_random_generator,
+            logit_mean=self.training_args.logit_mean,
+            logit_std=self.training_args.logit_std,
+            mode_scale=self.training_args.mode_scale,
+        )
+
+        boundary_ratio = self.training_args.boundary_ratio
+        if self.train_transformer_2:
+            u = boundary_ratio + u * (1.0 - boundary_ratio)
+        else:
+            u = u * boundary_ratio
+
+        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
+        return self.noise_scheduler.timesteps[indices].to(device=device)
+        
+    def _build_attention_metadata(
+            self, training_batch: TrainingBatch) -> TrainingBatch:
+        latents_shape = training_batch.raw_latent_shape
+        patch_size = self.training_args.pipeline_config.dit_config.patch_size
+        current_vsa_sparsity = training_batch.current_vsa_sparsity
+        assert latents_shape is not None
+        assert training_batch.timesteps is not None
+        if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
+            training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder(  # type: ignore
+            ).build(  # type: ignore
+                raw_latent_shape=latents_shape[2:5],
+                current_timestep=training_batch.timesteps,
+                patch_size=patch_size,
+                VSA_sparsity=current_vsa_sparsity,
+                device=get_local_torch_device())
+        else:
+            training_batch.attn_metadata = None
+
+        return training_batch
+
+    def _build_input_kwargs(self,
+                            training_batch: TrainingBatch) -> TrainingBatch:
+        training_batch.input_kwargs = {
+            "hidden_states":
+            training_batch.noisy_model_input,
+            "encoder_hidden_states":
+            training_batch.encoder_hidden_states,
+            "timestep":
+            training_batch.timesteps.to(get_local_torch_device(),
+                                        dtype=torch.bfloat16),
+            "encoder_attention_mask":
+            training_batch.encoder_attention_mask,
+            "return_dict":
+            False,
+        }
+        return training_batch
+
+    def _transformer_forward_and_compute_loss(
+            self, training_batch: TrainingBatch) -> TrainingBatch:
+        if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
+            assert training_batch.attn_metadata is not None
+        else:
+            assert training_batch.attn_metadata is None
+        input_kwargs = training_batch.input_kwargs
+
+        # if 'hunyuan' in self.training_args.model_type:
+        #     input_kwargs["guidance"] = torch.tensor(
+        #         [1000.0],
+        #         device=training_batch.noisy_model_input.device,
+        #         dtype=torch.bfloat16)
+        current_model = self.transformer_2 if self.train_transformer_2 else self.transformer
+
+        with set_forward_context(
+                current_timestep=training_batch.current_timestep,
+                attn_metadata=training_batch.attn_metadata):
+            model_pred = current_model(**input_kwargs)
+            if self.training_args.precondition_outputs:
+                assert training_batch.sigmas is not None
+                model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
+            assert training_batch.latents is not None
+            assert training_batch.noise is not None
+            target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents
+
+            # make sure no implicit broadcasting happens
+            assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
+            loss = (torch.mean((model_pred.float() - target.float())**2) /
+                    self.training_args.gradient_accumulation_steps)
+
+            loss.backward()
+            avg_loss = loss.detach().clone()
+
+        # logger.info(f"rank: {self.rank}, avg_loss: {avg_loss.item()}",
+        #             local_main_process_only=False)
+        world_group = get_world_group()
+        world_group.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
+        training_batch.total_loss += avg_loss.item()
+
+        return training_batch
+
+    def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch:
+        max_grad_norm = self.training_args.max_grad_norm
+
+        # TODO(will): perhaps move this into transformer api so that we can do
+        # the following:
+        # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
+        if max_grad_norm is not None:
+            # Only clip gradients for the model that is currently training
+            if self.train_transformer_2 and self.transformer_2 is not None:
+                model_parts = [self.transformer_2]
+            else:
+                model_parts = [self.transformer]
+            
+            grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases(
+                [p for m in model_parts for p in m.parameters()],
+                max_grad_norm,
+                foreach=None,
+            )
+            assert grad_norm is not float('nan') or grad_norm is not float(
+                'inf')
+            grad_norm = grad_norm.item() if grad_norm is not None else 0.0
+        else:
+            grad_norm = 0.0
+        training_batch.grad_norm = grad_norm
+        return training_batch
+
+    def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
+        training_batch = self._prepare_training(training_batch)
+
+        for _ in range(self.training_args.gradient_accumulation_steps):
+            training_batch = self._get_next_batch(training_batch)
+
+            # Normalize DIT input
+            training_batch = self._normalize_dit_input(training_batch)
+            # Create noisy model input
+            training_batch = self._prepare_dit_inputs(training_batch)
+
+            # Shard latents across sp groups
+            training_batch.latents = shard_latents_across_sp(
+                training_batch.latents,
+                num_latent_t=self.training_args.num_latent_t)
+            # shard noisy_model_input to match
+            training_batch.noisy_model_input = shard_latents_across_sp(
+                training_batch.noisy_model_input,
+                num_latent_t=self.training_args.num_latent_t)
+            # shard noise to match latents
+            training_batch.noise = shard_latents_across_sp(
+                training_batch.noise,
+                num_latent_t=self.training_args.num_latent_t)
+
+            training_batch = self._build_attention_metadata(training_batch)
+            training_batch = self._build_input_kwargs(training_batch)
+            training_batch = self._transformer_forward_and_compute_loss(
+                training_batch)
+
+        training_batch = self._clip_grad_norm(training_batch)
+
+        # Only step the optimizer and scheduler for the model that is currently training
+        if self.train_transformer_2 and self.transformer_2 is not None:
+            self.optimizer_2.step()
+            self.lr_scheduler_2.step()
+        else:
+            self.optimizer.step()
+            self.lr_scheduler.step()
+            
+        training_batch.total_loss = training_batch.total_loss
+        training_batch.grad_norm = training_batch.grad_norm
+        return training_batch
+
+    def _resume_from_checkpoint(self) -> None:
+        logger.info("Loading checkpoint from %s",
+                    self.training_args.resume_from_checkpoint)
+        resumed_step = load_checkpoint(
+            self.transformer, self.global_rank,
+            self.training_args.resume_from_checkpoint, self.optimizer,
+            self.train_dataloader, self.lr_scheduler,
+            self.noise_random_generator)
+        if resumed_step > 0:
+            self.init_steps = resumed_step
+            logger.info("Successfully resumed from step %s", resumed_step)
+        else:
+            logger.warning("Failed to load checkpoint, starting from step 0")
+            self.init_steps = 0
+
+    def train(self) -> None:
+        assert self.seed is not None, "seed must be set"
+        set_random_seed(self.seed + self.global_rank)
+        logger.info('rank: %s: start training',
+                    self.global_rank,
+                    local_main_process_only=False)
+        if not self.post_init_called:
+            self.post_init()
+        num_trainable_params = _get_trainable_params(self.transformer)
+        logger.info("Starting training with %s B trainable parameters",
+                    round(num_trainable_params / 1e9, 3))
+
+        # Set random seeds for deterministic training
+        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", self.seed)
+
+        self.noise_scheduler = FlowMatchEulerDiscreteScheduler()
+
+        if self.training_args.resume_from_checkpoint:
+            self._resume_from_checkpoint()
+
+        self.train_loader_iter = iter(self.train_dataloader)
+
+        step_times: deque[float] = deque(maxlen=100)
+
+        self._log_training_info()
+
+        self._log_validation(self.transformer, self.training_args,
+                             self.init_steps)
+
+        # Train!
+        progress_bar = tqdm(
+            range(0, self.training_args.max_train_steps),
+            initial=self.init_steps,
+            desc="Steps",
+            # Only show the progress bar once on each machine.
+            disable=self.local_rank > 0,
+        )
+        for step in range(self.init_steps + 1,
+                          self.training_args.max_train_steps + 1):
+            start_time = time.perf_counter()
+            if vsa_available:
+                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
+                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 = 0.0
+
+            training_batch = TrainingBatch()
+            training_batch.current_timestep = step
+            training_batch.current_vsa_sparsity = current_vsa_sparsity
+            training_batch = self.train_one_step(training_batch)
+
+            loss = training_batch.total_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({
+                "loss": f"{loss:.4f}",
+                "step_time": f"{step_time:.2f}s",
+                "grad_norm": grad_norm,
+            })
+            progress_bar.update(1)
+            if self.global_rank == 0:
+                wandb.log(
+                    {
+                        "train_loss": loss,
+                        "learning_rate": self.lr_scheduler.get_last_lr()[0],
+                        "step_time": step_time,
+                        "avg_step_time": avg_step_time,
+                        "grad_norm": grad_norm,
+                        "vsa_sparsity": current_vsa_sparsity,
+                    },
+                    step=step,
+                )
+            if step % self.training_args.checkpointing_steps == 0:
+                save_checkpoint(self.transformer, self.global_rank,
+                                self.training_args.output_dir, step,
+                                self.optimizer, self.train_dataloader,
+                                self.lr_scheduler, self.noise_random_generator)
+                self.transformer.train()
+                self.sp_group.barrier()
+            if self.training_args.log_validation and step % self.training_args.validation_steps == 0:
+                self._log_validation(self.transformer, self.training_args, step)
+                gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+                trainable_params = round(
+                    _get_trainable_params(self.transformer) / 1e9, 3)
+                logger.info(
+                    "GPU memory usage after validation: %s MB, trainable params: %sB",
+                    gpu_memory_usage, trainable_params)
+
+        wandb.finish()
+        save_checkpoint(self.transformer, self.global_rank,
+                        self.training_args.output_dir,
+                        self.training_args.max_train_steps, self.optimizer,
+                        self.train_dataloader, self.lr_scheduler,
+                        self.noise_random_generator)
+
+        if get_sp_group():
+            cleanup_dist_env_and_memory()
+
+    def _log_training_info(self) -> None:
+        total_batch_size = (self.world_size *
+                            self.training_args.gradient_accumulation_steps /
+                            self.training_args.sp_size *
+                            self.training_args.train_sp_batch_size)
+        logger.info("***** Running training *****")
+        logger.info("  Num examples = %s", len(self.train_dataset))
+        logger.info("  Dataloader size = %s", len(self.train_dataloader))
+        logger.info("  Num Epochs = %s", self.num_train_epochs)
+        logger.info("  Resume training from step %s",
+                    self.init_steps)  # type: ignore
+        logger.info("  Instantaneous batch size per device = %s",
+                    self.training_args.train_batch_size)
+        logger.info(
+            "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
+            total_batch_size)
+        logger.info("  Gradient Accumulation steps = %s",
+                    self.training_args.gradient_accumulation_steps)
+        logger.info("  Total optimization steps = %s",
+                    self.training_args.max_train_steps)
+        logger.info("  Total training parameters per FSDP shard = %s B",
+                    round(_get_trainable_params(self.transformer) / 1e9, 3))
+        # print dtype
+        logger.info("  Master weight dtype: %s",
+                    self.transformer.parameters().__next__().dtype)
+
+        gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
+        logger.info("GPU memory usage before train_one_step: %s MB",
+                    gpu_memory_usage)
+        logger.info("VSA validation sparsity: %s",
+                    self.training_args.VSA_sparsity)
+
+    def _prepare_validation_batch(self, sampling_param: SamplingParam,
+                                  training_args: TrainingArgs,
+                                  validation_batch: dict[str, Any],
+                                  num_inference_steps: int) -> ForwardBatch:
+        sampling_param.prompt = validation_batch['prompt']
+        sampling_param.height = training_args.num_height
+        sampling_param.width = training_args.num_width
+        sampling_param.num_inference_steps = num_inference_steps
+        sampling_param.data_type = "video"
+        assert self.seed is not None
+        sampling_param.seed = self.seed
+
+        latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
+                        sampling_param.height // 8, sampling_param.width // 8]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+        temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+        num_frames = (training_args.num_latent_t -
+                      1) * temporal_compression_factor + 1
+        sampling_param.num_frames = num_frames
+        batch = ForwardBatch(
+            **shallow_asdict(sampling_param),
+            latents=None,
+            generator=self.validation_random_generator,
+            n_tokens=n_tokens,
+            eta=0.0,
+            VSA_sparsity=training_args.VSA_sparsity,
+        )
+
+        return batch
+
+    @torch.no_grad()
+    def _log_validation(self, transformer, training_args, global_step) -> None:
+        """
+        Generate a validation video and log it to wandb to check the quality during training.
+        """
+        training_args.inference_mode = True
+        training_args.dit_cpu_offload = True
+        if not training_args.log_validation:
+            return
+        if self.validation_pipeline is None:
+            raise ValueError("Validation pipeline is not set")
+
+        logger.info("Starting validation")
+
+        # Create sampling parameters if not provided
+        sampling_param = SamplingParam.from_pretrained(training_args.model_path)
+
+        # Prepare validation prompts
+        logger.info('rank: %s: fastvideo_args.validation_dataset_file: %s',
+                    self.global_rank,
+                    training_args.validation_dataset_file,
+                    local_main_process_only=False)
+        validation_dataset = ValidationDataset(
+            training_args.validation_dataset_file)
+        validation_dataloader = DataLoader(validation_dataset,
+                                           batch_size=None,
+                                           num_workers=0)
+
+        transformer.eval()
+
+        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]
+        # Log validation results for this step
+        world_group = get_world_group()
+        num_sp_groups = world_group.world_size // self.sp_group.world_size
+
+        # Process each validation prompt for each validation step
+        for num_inference_steps in validation_steps:
+            logger.info("rank: %s: num_inference_steps: %s",
+                        self.global_rank,
+                        num_inference_steps,
+                        local_main_process_only=False)
+            step_videos: list[np.ndarray] = []
+            step_captions: list[str] = []
+
+            for validation_batch in validation_dataloader:
+                batch = self._prepare_validation_batch(sampling_param,
+                                                       training_args,
+                                                       validation_batch,
+                                                       num_inference_steps)
+                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
+                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)
+
+            # Only sp_group leaders (rank_in_sp_group == 0) need to send their
+            # results to global rank 0
+            if self.rank_in_sp_group == 0:
+                if self.global_rank == 0:
+                    # Global rank 0 collects results from all sp_group leaders
+                    all_videos = step_videos  # Start with own results
+                    all_captions = step_captions
+
+                    # Receive from other sp_group leaders
+                    for sp_group_idx in range(1, num_sp_groups):
+                        src_rank = sp_group_idx * self.sp_world_size  # Global rank of other sp_group leaders
+                        recv_videos = world_group.recv_object(src=src_rank)
+                        recv_captions = world_group.recv_object(src=src_rank)
+                        all_videos.extend(recv_videos)
+                        all_captions.extend(recv_captions)
+
+                    video_filenames = []
+                    for i, (video, caption) in enumerate(
+                            zip(all_videos, all_captions, strict=True)):
+                        os.makedirs(training_args.output_dir, exist_ok=True)
+                        filename = os.path.join(
+                            training_args.output_dir,
+                            f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4"
+                        )
+                        imageio.mimsave(filename, video, fps=sampling_param.fps)
+                        video_filenames.append(filename)
+
+                    logs = {
+                        f"validation_videos_{num_inference_steps}_steps": [
+                            wandb.Video(filename, caption=caption)
+                            for filename, caption in zip(
+                                video_filenames, all_captions, strict=True)
+                        ]
+                    }
+                    wandb.log(logs, step=global_step)
+                else:
+                    # Other sp_group leaders send their results to global rank 0
+                    world_group.send_object(step_videos, dst=0)
+                    world_group.send_object(step_captions, dst=0)
+
+        # Re-enable gradients for training
+        training_args.inference_mode = False
+        transformer.train()

fastvideo/training/training_utils.py

@@ -159,7 +159,6 @@ def save_checkpoint(transformer,
                 local_main_process_only=False)
 
     cpu_state = gather_state_dict_on_cpu_rank0(transformer, device=None)
-
     if rank == 0:
         # Save model weights (consolidated)
         transformer_save_dir = os.path.join(save_dir, "transformer")
@@ -1084,11 +1083,11 @@ def get_cosine_schedule_with_min_lr(optimizer: Optimizer,
         progress = float(current_step - num_warmup_steps) / float(
             max(1, num_training_steps - num_warmup_steps))
         # Cosine decay from 1.0 to min_lr_ratio over num_cycles periods
-        # Use the same formula as standard cosine but scale the output to [min_lr_ratio, 1.0]
+        # Use the same formula as standard cosine but ensure minimum is min_lr_ratio instead of 0
         cosine_value = 0.5 * (
             1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))
-        # Scale from [0, 1] to [min_lr_ratio, 1.0]
-        return min_lr_ratio + (1.0 - min_lr_ratio) * cosine_value
+        # Ensure the value doesn't go below min_lr_ratio
+        return max(min_lr_ratio, cosine_value)
 
     return LambdaLR(optimizer, lr_lambda, last_epoch)
 

fastvideo/training/wan_distillation_pipeline.py

@@ -41,8 +41,6 @@ class WanDistillationPipeline(DistillationPipeline):
         args_copy = deepcopy(training_args)
 
         args_copy.inference_mode = True
-        args_copy.dit_cpu_offload = True
-        args_copy.pipeline_config.vae_config.load_encoder = False
         validation_pipeline = WanDMDPipeline.from_pretrained(
             training_args.model_path,
             args=args_copy,  # type: ignore
@@ -50,7 +48,9 @@ class WanDistillationPipeline(DistillationPipeline):
             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)
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=True)
 
         self.validation_pipeline = validation_pipeline
 

fastvideo/training/wan_i2v_distillation_pipeline.py

@@ -7,8 +7,7 @@ import torch
 from einops import rearrange
 
 from fastvideo.configs.sample import SamplingParam
-from fastvideo.dataset.dataloader.schema import (pyarrow_schema_i2v,
-                                                 pyarrow_schema_i2v_validation)
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_i2v
 from fastvideo.distributed import get_local_torch_device
 from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.logger import init_logger
@@ -48,7 +47,6 @@ class WanI2VDistillationPipeline(DistillationPipeline):
 
     def set_schemas(self):
         self.train_dataset_schema = pyarrow_schema_i2v
-        self.validation_dataset_schema = pyarrow_schema_i2v_validation
 
     def initialize_validation_pipeline(self, training_args: TrainingArgs):
         logger.info("Initializing validation pipeline...")

fastvideo/training/wan_training_pipeline.py

@@ -34,10 +34,8 @@ class WanTrainingPipeline(TrainingPipeline):
     def initialize_validation_pipeline(self, training_args: TrainingArgs):
         logger.info("Initializing validation pipeline...")
         args_copy = deepcopy(training_args)
-
+        assert training_args.dit_cpu_offload == False
         args_copy.inference_mode = True
-        args_copy.dit_cpu_offload = True
-        args_copy.pipeline_config.vae_config.load_encoder = False
         validation_pipeline = WanPipeline.from_pretrained(
             training_args.model_path,
             args=args_copy,  # type: ignore
@@ -47,7 +45,9 @@ class WanTrainingPipeline(TrainingPipeline):
             },
             tp_size=training_args.tp_size,
             sp_size=training_args.sp_size,
-            num_gpus=training_args.num_gpus)
+            num_gpus=training_args.num_gpus,
+            pin_cpu_memory=training_args.pin_cpu_memory,
+            dit_cpu_offload=training_args.dit_cpu_offload)
 
         self.validation_pipeline = validation_pipeline
 

fastvideo/version.py

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

fastvideo/worker/gpu_worker.py

@@ -11,6 +11,7 @@ from typing import Any, TextIO, cast
 import psutil
 import torch
 
+import fastvideo.envs as envs
 from fastvideo.distributed import (
     cleanup_dist_env_and_memory,
     maybe_init_distributed_environment_and_model_parallel)
@@ -71,7 +72,6 @@ class Worker:
 
         # _check_if_gpu_supports_dtype(self.model_config.dtype)
         if current_platform.is_cuda_alike():
-            torch.cuda.empty_cache()
             self.init_gpu_memory = torch.cuda.mem_get_info()[0]
         else:
             # For MPS, we can't get memory info the same way
@@ -140,7 +140,13 @@ class Worker:
                     fastvideo_args = recv_rpc['kwargs']['fastvideo_args']
                     output_batch = self.execute_forward(forward_batch,
                                                         fastvideo_args)
-                    self.pipe.send({"output_batch": output_batch.output.cpu()})
+                    logging_info = None
+                    if envs.FASTVIDEO_STAGE_LOGGING:
+                        logging_info = output_batch.logging_info
+                    self.pipe.send({
+                        "output_batch": output_batch.output.cpu(),
+                        "logging_info": logging_info
+                    })
                 elif method_name == 'set_lora_adapter':
                     lora_nickname = recv_rpc['kwargs']['lora_nickname']
                     lora_path = recv_rpc['kwargs']['lora_path']