Fix OOM issue

.buildkite/pipeline.yml

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

.github/workflows/build-image.yml

@@ -18,12 +18,6 @@ on:
         required: false
         default: false
         type: boolean
-      python_3_12_cuda_12_9:
-        description: 'Build Python 3.12 image Cuda 12.9'
-        required: false
-        default: false
-        type: boolean
-
 
 permissions:
   contents: read
@@ -55,13 +49,4 @@ jobs:
       python_version: '3.12'
       dockerfile_path: docker/Dockerfile.python3.12
       tag_suffix: py3.12
-    secrets: inherit
-
-  build-python-3-12-cuda-12-9:
-    if: ${{ github.event.inputs.python_3_12_cuda_12_9 == 'true' }}
-    uses: ./.github/workflows/build-image-template.yml
-    with:
-      python_version: '3.12'
-      dockerfile_path: docker/Dockerfile.python3.12.cuda12.9.1
-      tag_suffix: py3.12-cuda12.9.1
     secrets: inherit

.github/workflows/pr-test.yml

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

.github/workflows/sta-publish.yml

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

.github/workflows/vsa-publish.yml

@@ -1,257 +0,0 @@
-name: Publish Video Sparse Attention Kernel to PyPI on Version Change
-
-on:
-  push:
-    branches:
-      - main
-    paths:
-      - "csrc/attn/video_sparse_attn/setup.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/video_sparse_attn
-          # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
-          echo "New version: $NEW_VERSION"
-          
-          # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup.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/video_sparse_attn  # Move into the correct folder
-          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py bdist_wheel --dist-dir=dist
-
-      - name: Rename wheel file
-        run: |
-          cd csrc/attn/video_sparse_attn
-
-          CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
-          TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2)
-          # 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/video_sparse_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/video_sparse_attn  # Move into the correct folder
-          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py sdist --dist-dir=dist
-
-      - name: Publish release distributions to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          packages-dir: csrc/attn/video_sparse_attn/dist/

.gitignore

@@ -42,7 +42,6 @@ docs/_build/
 docs/source/getting_started/examples/
 docs/source/inference/examples/
 docs/source/training/examples/
-docs/source/distillation/examples/
 
 # VSCode
 .vscode/

.gitmodules

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

.pre-commit-config.yaml

@@ -22,7 +22,6 @@ exclude: |
     examples/.*|
     .github/workflows/fastvideo-publish.yml|
     .github/workflows/sta-publish.yml|
-    .github/workflows/vsa-publish.yml|
     .github/workflows/build-image-template.yml|
     docs/source/inference/support_matrix.md
   )

README.md

@@ -1,21 +1,21 @@
 <div align="center">
-<img src=assets/logos/logo.svg width="30%"/>
+<img src=assets/logo.jpg width="30%"/>
 </div>
 
-**FastVideo is a unified post-training and inference framework for accelerated video generation.**
+**FastVideo is a unified framework for accelerated video generation.**
 
-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.
+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.
 
 <p align="center">
-    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/rG0QpZdw" target="_blank"> <b> WeChat </b> </a> |
+    | <a href="https://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> |
 </p>
 
 <div align="center">
-<img src=assets/fastwan.png width="90%"/>
+<img src=assets/perf.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,19 +23,20 @@ FastVideo features an end-to-end unified pipeline for accelerating diffusion mod
 ## 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)
-- Diverse hardware and OS support
-  - Support H100, A100, 4090
-  - Support Linux, Windows, MacOS
+- 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.
 
 ## Getting Started
 We recommend using an environment manager such as `Conda` to create a clean environment:
@@ -51,31 +52,17 @@ 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. 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:
+Here's a minimal example to generate a video using the default settings. 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(
-        "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         num_gpus=1,  # Adjust based on your hardware
     )
 
@@ -108,42 +95,42 @@ For a more detailed guide, please see our [inference quick start](https://hao-ai
 - [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
 
 ## 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) -->
+- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/training/distillation.html)
+- [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html)
 
 ## 📑 Development Plan
+
 <!-- - More distillation methods -->
   <!-- - [ ] Add Distribution Matching Distillation -->
-More FastWan Models Coming Soon!
-- [ ] Add FastWan2.1-T2V-14B
-- [ ] Add FastWan2.2-T2V-14B
-- [ ] Add FastWan2.2-I2V-14B
-<!-- - Optimization features
-- Code updates -->
+- 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
   <!-- - [ ] 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:
-- [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)
+- [PCM](https://github.com/G-U-N/Phased-Consistency-Model)
 - [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](https://ifm.mbzuai.ac.ae/), [Anyscale](https://www.anyscale.com/), and [GMI Cloud](https://www.gmicloud.ai/) for their support throughout this project.
+We thank MBZUAI and [Anyscale](https://www.anyscale.com/) for their support throughout this project.
 
 ## Citation
-If you find FastVideo useful, please considering citing our work:
+If you use FastVideo for your research, please cite our work:
 
 ```bibtex
 @software{fastvideo2024,
@@ -154,17 +141,31 @@ If you find FastVideo useful, please considering citing our work:
   year         = {2024},
 }
 
-@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{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{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}
+@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},
 }
 ```

assets/logos/icon_simple.svg

@@ -1,6 +0,0 @@
-<svg width="160" height="93" viewBox="0 0 160 93" fill="none" xmlns="http://www.w3.org/2000/svg">
-<path d="M28.8511 91.66L57.6319 1.86368H64.5394L35.7585 91.66H28.8511Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
-<path d="M15.0376 91.66L43.8185 1.86368H46.1209L17.3401 91.66H15.0376Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
-<path d="M1.22217 91.66L30.003 1.86366H31.1543L2.3734 91.66H1.22217Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.15122"/>
-<path d="M71.4465 1.86483L42.666 91.6599H69.144L78.3538 58.2746H123.251L129.007 39.855H84.1099L89.866 22.5868H152.032L157.788 1.86483H71.4465Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
-</svg>

assets/logos/logo.svg

@@ -1,18 +0,0 @@
-<svg width="252" height="105" viewBox="0 0 252 105" fill="none" xmlns="http://www.w3.org/2000/svg">
-<path d="M89.4843 55.5457H101.361L87.7028 101H74.638L89.4843 55.5457Z" fill="#356CFF"/>
-<path fill-rule="evenodd" clip-rule="evenodd" d="M96.0167 1.00057H112.645L118.583 48.273H104.924L103.737 39.7882H79.9827L67.5117 55.5457H85.3273L43.1638 101H28.3174L22.3789 55.5457H33.6621L38.4129 91.3031L58.604 68.2729H44.3515L96.0167 1.00057ZM100.768 13.1217L87.7028 29.4852H103.143L100.768 13.1217Z" fill="#356CFF"/>
-<path d="M37.2252 1.00057L22.3789 48.273H36.0375L40.7884 30.6974L62.6727 30.6974L69.6727 21.0005L43.7576 21.0004L46.7269 11.9096L77.6727 11.9096L86 1.00057L37.2252 1.00057Z" fill="#356CFF"/>
-<path fill-rule="evenodd" clip-rule="evenodd" d="M108.488 55.5457L94.2351 101C94.2351 101 105.518 101 120.959 101C136.399 101 144.078 93.0133 148.276 79.788C152.432 68.0157 153.027 55.5457 136.399 55.5457C119.771 55.5457 108.488 55.5457 108.488 55.5457ZM109.081 90.697L116.802 65.8487C116.802 65.8487 120.959 65.8487 132.242 65.8487C143.525 65.8487 137.586 78.5759 135.211 84.0304C133.307 88.4021 127.491 90.697 122.74 90.697C117.989 90.697 109.081 90.697 109.081 90.697Z" fill="#356CFF"/>
-<path d="M173.188 1.00056L168.625 11.9096C168.625 11.9096 149.386 11.9092 142.525 11.9095C135.664 11.9098 136.586 20.3944 141.337 20.3944H159.747C168.654 20.3944 166.961 33.6899 163.904 38.5761C160.467 44.0675 157.371 48.273 148.463 48.273L125.188 48.273L124 37.97L147.87 37.97C153.808 37.97 156.184 29.4852 151.433 29.4852H131.836C120.142 29.4852 125.897 1.00043 141.337 1.00043L173.188 1.00056Z" fill="#356CFF"/>
-<path d="M179.938 1.00056L175.688 11.9096L191.221 11.9096L179.938 48.273H192.409L203.692 11.9096L219.132 11.9095L223.289 1.00043L179.938 1.00056Z" fill="#356CFF"/>
-<path d="M161.341 55.5457H202.845L198.5 65.8487H169.654L167.279 73.7268H188.5L184.749 82.8177H164.31L161.934 90.697H190.251L186.624 101H146.494L161.341 55.5457Z" fill="#356CFF"/>
-<path fill-rule="evenodd" clip-rule="evenodd" d="M230.821 54.9391C255.169 54.9391 251.776 67.0602 249.231 77.9692C246.686 88.8783 240.917 101 217.757 101C194.596 101 195.606 88.8783 199.347 77.9692C203.089 67.0602 206.473 54.9391 230.821 54.9391ZM237.948 77.9692C239.984 70.6965 240.917 65.242 228.446 65.242C215.975 65.242 211.818 71.9087 210.037 77.9692C208.255 84.0298 208.255 91.3025 219.538 91.3025C230.821 91.3025 235.911 85.2419 237.948 77.9692Z" fill="#356CFF"/>
-<path d="M173.188 1.00056L168.625 11.9096C168.625 11.9096 149.386 11.9092 142.525 11.9095C135.664 11.9098 136.586 20.3944 141.337 20.3944M173.188 1.00056C173.188 1.00056 156.777 1.00043 141.337 1.00043M173.188 1.00056L141.337 1.00043M141.337 20.3944C146.088 20.3944 150.839 20.3944 159.747 20.3944M141.337 20.3944H159.747M159.747 20.3944C168.654 20.3944 166.961 33.6899 163.904 38.5761C160.467 44.0675 157.371 48.273 148.463 48.273M148.463 48.273C139.556 48.273 125.188 48.273 125.188 48.273M148.463 48.273L125.188 48.273M125.188 48.273L124 37.97M124 37.97C124 37.97 141.931 37.97 147.87 37.97M124 37.97L147.87 37.97M147.87 37.97C153.808 37.97 156.184 29.4852 151.433 29.4852M151.433 29.4852C146.682 29.4852 138.962 29.4852 131.836 29.4852M151.433 29.4852H131.836M131.836 29.4852C120.142 29.4852 125.897 1.00043 141.337 1.00043M37.2252 1.00057L22.3789 48.273H36.0375L40.7884 30.6974L62.6727 30.6974L69.6727 21.0005L43.7576 21.0004L46.7269 11.9096L77.6727 11.9096L86 1.00057L37.2252 1.00057ZM96.0167 1.00057H112.645L118.583 48.273H104.924L103.737 39.7882H79.9827L67.5117 55.5457H85.3273L43.1638 101H28.3174L22.3789 55.5457H33.6621L38.4129 91.3031L58.604 68.2729H44.3515L96.0167 1.00057ZM87.7028 29.4852L100.768 13.1217L103.143 29.4852H87.7028ZM89.4843 55.5457H101.361L87.7028 101H74.638L89.4843 55.5457ZM108.488 55.5457L94.2351 101C94.2351 101 105.518 101 120.959 101C136.399 101 144.078 93.0133 148.276 79.788C152.432 68.0157 153.027 55.5457 136.399 55.5457C119.771 55.5457 108.488 55.5457 108.488 55.5457ZM116.802 65.8487L109.081 90.697C109.081 90.697 117.989 90.697 122.74 90.697C127.491 90.697 133.307 88.4021 135.211 84.0304C137.586 78.5759 143.525 65.8487 132.242 65.8487C120.959 65.8487 116.802 65.8487 116.802 65.8487ZM179.938 1.00056L175.688 11.9096L191.221 11.9096L179.938 48.273H192.409L203.692 11.9096L219.132 11.9095L223.289 1.00043L179.938 1.00056ZM161.341 55.5457H202.845L198.5 65.8487H169.654L167.279 73.7268H188.5L184.749 82.8177H164.31L161.934 90.697H190.251L186.624 101H146.494L161.341 55.5457ZM230.821 54.9391C255.169 54.9391 251.776 67.0602 249.231 77.9692C246.686 88.8783 240.917 101 217.757 101C194.596 101 195.606 88.8783 199.347 77.9692C203.089 67.0602 206.473 54.9391 230.821 54.9391ZM228.446 65.242C240.917 65.242 239.984 70.6965 237.948 77.9692C235.911 85.2419 230.821 91.3025 219.538 91.3025C208.255 91.3025 208.255 84.0298 210.037 77.9692C211.818 71.9087 215.975 65.242 228.446 65.242Z" stroke="#356CFF" stroke-width="1.18771"/>
-<path d="M15.2524 55.5451L21.191 100.999L24.7541 100.999L18.8156 55.5451L15.2524 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
-<path d="M8.12646 55.5451L14.065 100.999L15.2527 100.999L9.31417 55.5451L8.12646 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
-<path d="M1 55.5451L6.93853 100.999L7.53239 100.999L1.59385 55.5451L1 55.5451Z" fill="#356CFF" stroke="#356CFF" stroke-width="0.593853"/>
-<path d="M15.2524 48.2724L30.0988 1H33.6619L18.8156 48.2724H15.2524Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
-<path d="M8.12646 48.2724L22.9728 1H24.1605L9.31417 48.2724H8.12646Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.18771"/>
-<path d="M1 48.2724L15.8463 1H16.4402L1.59385 48.2724H1Z" fill="#356CFF" stroke="#356CFF" stroke-width="0.593853"/>
-<path d="M85.3271 55.5457H67.5116L87 12.7363L44.3513 68.2729H58.6038L43.1636 101L85.3271 55.5457Z" fill="#FDC717" stroke="#FDC717" stroke-width="1.18771" stroke-miterlimit="16"/>
-</svg>

comfyui/assets/icon_simple.svg

@@ -1,6 +0,0 @@
-<svg width="160" height="93" viewBox="0 0 160 93" fill="none" xmlns="http://www.w3.org/2000/svg">
-<path d="M28.8511 91.66L57.6319 1.86368H64.5394L35.7585 91.66H28.8511Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
-<path d="M15.0376 91.66L43.8185 1.86368H46.1209L17.3401 91.66H15.0376Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
-<path d="M1.22217 91.66L30.003 1.86366H31.1543L2.3734 91.66H1.22217Z" fill="#356CFF" stroke="#356CFF" stroke-width="1.15122"/>
-<path d="M71.4465 1.86483L42.666 91.6599H69.144L78.3538 58.2746H123.251L129.007 39.855H84.1099L89.866 22.5868H152.032L157.788 1.86483H71.4465Z" fill="#356CFF" stroke="#356CFF" stroke-width="2.30244"/>
-</svg>

csrc/attn/MANIFEST.in

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

csrc/attn/README.md

@@ -25,9 +25,9 @@ sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave
 sudo apt update
 sudo apt install clang-11
 ```
-(If you use CUDA12.8)
+(If you use CUDA12.4)
 ```bash
-export CUDA_HOME=/usr/local/cuda-12.8
+export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 ```
@@ -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_vsa.py # test VSA
+python tests/test_block_sparse.py # test VSA
 ```
 ### Benchmark
 ```bash

csrc/attn/benchmarks/bench_sta.py

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

csrc/attn/scratch.sh

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

csrc/attn/setup_sta.py

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

csrc/attn/setup_vsa.py

@@ -9,10 +9,10 @@ target = target.lower()
 
 # Package metadata
 PACKAGE_NAME = "vsa"
-VERSION = "0.0.3"
+VERSION = "0.0.1"
 AUTHOR = "Hao AI Lab"
 DESCRIPTION = "Video Sparse Attention Kernel Used in FastVideo"
-URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn/video_sparse_attn"
+URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn"
 
 # Set environment variables
 tk_root = os.getenv('THUNDERKITTENS_ROOT', os.path.abspath(os.path.join(os.getcwd(), 'tk/')))
@@ -51,16 +51,19 @@ for k in kernels:
         source_files.append(sources[k]['source_files'][target])
     cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}')
 
-
-ext_modules = [
-    CUDAExtension('vsa_cuda',
-        sources=source_files,
-        extra_compile_args={
-            'cxx': cpp_flags,
-            'nvcc': cuda_flags
-        },
-        libraries=['cuda'])
-]
+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'])
+    ]
 
 
 

csrc/attn/sliding_tile_attn/README.md

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

csrc/attn/tests/test_vsa.py

@@ -13,9 +13,9 @@ BLOCK_M = 64
 BLOCK_N = 64
 
 def pytorch_test(Q, K, V, block_sparse_mask, dO):
-    q_ = Q.clone().float().requires_grad_()
-    k_ = K.clone().float().requires_grad_()
-    v_ = V.clone().float().requires_grad_()
+    q_ = Q.clone().requires_grad_()
+    k_ = K.clone().requires_grad_()
+    v_ = V.clone().requires_grad_()
 
     QK = torch.matmul(q_, k_.transpose(-2, -1))
     QK /= (q_.size(-1) ** 0.5)
@@ -35,9 +35,9 @@ def pytorch_test(Q, K, V, block_sparse_mask, dO):
 
 
 def block_sparse_kernel_test(Q, K, V, block_sparse_mask, variable_block_sizes, non_pad_index, dO):
-    Q = Q.detach().requires_grad_()
-    K = K.detach().requires_grad_()
-    V = V.detach().requires_grad_()
+    Q = Q.clone().requires_grad_()
+    K = K.clone().requires_grad_()
+    V = V.clone().requires_grad_()
     
     q_padded = vsa_pad(Q, non_pad_index, variable_block_sizes.shape[0], BLOCK_M)
     k_padded = vsa_pad(K, non_pad_index, variable_block_sizes.shape[0], BLOCK_M)
@@ -60,9 +60,11 @@ def get_non_pad_index(
 
     return index_pad[index_mask]
 
-def generate_tensor(shape, dtype, device):
+def generate_tensor(shape, mean, std, dtype, device):
     tensor = torch.randn(shape, dtype=dtype, device=device)
-    return tensor
+    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
+    scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
+    return scaled_tensor.contiguous()
 
 def generate_variable_block_sizes(num_blocks, min_size=32, max_size=64, device="cuda"):
     return torch.randint(min_size, max_size + 1, (num_blocks,), device=device, dtype=torch.int32)
@@ -73,7 +75,7 @@ def vsa_pad(x, non_pad_index, num_blocks, block_size):
     padded_x[:, :, non_pad_index, :] = x
     return padded_x
 
-def check_correctness(h, d, num_blocks, k,  num_iterations=20, error_mode='all'):
+def check_correctness(h, d, num_blocks, k, mean, std, num_iterations=20, error_mode='all'):
     results = {
         'gO': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
         'gQ': {'sum_diff': 0.0, 'sum_abs': 0.0, 'max_diff': 0.0},
@@ -89,10 +91,10 @@ def check_correctness(h, d, num_blocks, k,  num_iterations=20, error_mode='all')
     block_mask = generate_block_sparse_mask_for_function(h, num_blocks, k, device)
     full_mask = create_full_mask_from_block_mask(block_mask, variable_block_sizes, device)
     for _ in range(num_iterations):
-        Q = generate_tensor((1, h, S, d),  torch.bfloat16, device)
-        K = generate_tensor((1, h, S, d), torch.bfloat16, device)
-        V = generate_tensor((1, h, S, d),  torch.bfloat16, device)
-        dO = generate_tensor((1, h, S, d), torch.bfloat16, device)
+        Q = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
+        K = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
+        V = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
+        dO = generate_tensor((1, h, S, d), mean, std, torch.bfloat16, device)
 
         # dO_padded = torch.zeros_like(dO_padded)
         # dO_padded[:, :, non_pad_index, :] = dO
@@ -105,8 +107,7 @@ def check_correctness(h, d, num_blocks, k,  num_iterations=20, error_mode='all')
                 abs_diff = torch.abs(diff)
                 results[name]['sum_diff'] += torch.sum(abs_diff).item()
                 results[name]['sum_abs'] += torch.sum(torch.abs(pt)).item()
-                rel_max_diff = torch.max(abs_diff) / torch.mean(torch.abs(pt))
-                results[name]['max_diff'] = max(results[name]['max_diff'], rel_max_diff.item())
+                results[name]['max_diff'] = max(results[name]['max_diff'], torch.max(abs_diff).item())
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
@@ -118,27 +119,27 @@ def check_correctness(h, d, num_blocks, k,  num_iterations=20, error_mode='all')
 
     return results
 
-def generate_error_graphs(h, d, error_mode='all'):
+def generate_error_graphs(h, d, mean, std, error_mode='all'):
     test_configs = [
         {"num_blocks": 16, "k": 2, "description": "Small sequence"},
         {"num_blocks": 32, "k": 4, "description": "Medium sequence"},
         {"num_blocks": 53, "k": 6, "description": "Large sequence"},
     ]
     
-    print(f"\nError Analysis for h={h}, d={d}, mode={error_mode}")
+    print(f"\nError Analysis for h={h}, d={d}, mean={mean}, std={std}, mode={error_mode}")
     print("=" * 150)
     print(f"{'Config':<20} {'Blocks':<8} {'K':<4} "
-          f"{'gQ Avg':<12} {'Rel gQ Max':<12} "
-          f"{'gK Avg':<12} {'Rel gK Max':<12} "
-          f"{'gV Avg':<12} {'Rel gV Max':<12} "
-          f"{'gO Avg':<12} {'Rel gO Max':<12}")
+          f"{'gQ Avg':<12} {'gQ Max':<12} "
+          f"{'gK Avg':<12} {'gK Max':<12} "
+          f"{'gV Avg':<12} {'gV Max':<12} "
+          f"{'gO Avg':<12} {'gO Max':<12}")
     print("-" * 150)
 
     for config in test_configs:
         num_blocks = config["num_blocks"]
         k = config["k"]
         description = config["description"]
-        results = check_correctness(h, d, num_blocks, k, error_mode=error_mode)
+        results = check_correctness(h, d, num_blocks, k, mean, std, error_mode=error_mode)
         print(f"{description:<20} {num_blocks:<8} {k:<4} "
               f"{results['gQ']['avg_diff']:<12.6e} {results['gQ']['max_diff']:<12.6e} "
               f"{results['gK']['avg_diff']:<12.6e} {results['gK']['max_diff']:<12.6e} "
@@ -149,8 +150,10 @@ def generate_error_graphs(h, d, error_mode='all'):
 
 if __name__ == "__main__":
     h, d = 16, 128
+    mean = 0.0
+    std = 1
     print("Block Sparse Attention with Variable Block Sizes Analysis")
     print("=" * 60)
     for mode in ['backward']:
-        generate_error_graphs(h, d, error_mode=mode)
+        generate_error_graphs(h, d, mean, std, error_mode=mode)
     print("\nAnalysis completed for all modes.")

csrc/attn/video_sparse_attn/MANIFEST.in

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

csrc/attn/video_sparse_attn/README.md

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

csrc/attn/video_sparse_attn/vsa/block_sparse_wrapper.py

@@ -1,185 +0,0 @@
-import torch
-try:
-    from vsa_cuda import block_sparse_fwd, block_sparse_bwd
-except ImportError:
-    block_sparse_fwd = None
-    block_sparse_bwd = None
-from vsa.block_sparse_attn_triton import triton_block_sparse_attn_forward, triton_block_sparse_attn_backward
-assert torch.__version__ >= "2.4.0", "VSA requires PyTorch 2.4.0 or higher"
-from vsa.index import map_to_index
-from typing import Tuple, Optional
-
-
-
-@torch.library.custom_op("vsa::block_sparse_attn_triton", mutates_args=(), device_types="cuda")
-def block_sparse_attn_triton(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    q = q.contiguous()
-    k = k.contiguous()
-    v = v.contiguous()
-    block_map = block_map.int()
-    q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
-    o, M = triton_block_sparse_attn_forward(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
-    return o, M
-
-
-@torch.library.register_fake("vsa::block_sparse_attn_triton")
-def _block_sparse_attn_triton_fake(
-    q: torch.Tensor,
-    k: torch.Tensor,
-    v: torch.Tensor,
-    block_map: torch.Tensor,
-    variable_block_sizes: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    q = q.contiguous()
-    k = k.contiguous()
-    v = v.contiguous()
-    o = torch.empty_like(q)
-    M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
-    return o, M
-
-
-@torch.library.custom_op("vsa::block_sparse_attn_backward_triton", mutates_args=(), device_types="cuda")
-def block_sparse_attn_backward_triton(
-    grad_output_padded: torch.Tensor, 
-    q_padded: torch.Tensor, 
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    o_padded: torch.Tensor, 
-    M: 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()
-    q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
-    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2))
-    dq, dk, dv = triton_block_sparse_attn_backward(grad_output_padded, q_padded, k_padded, v_padded, o_padded, M, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes)
-    return dq, dk, dv
-
-@torch.library.register_fake("vsa::block_sparse_attn_backward_triton")
-def _block_sparse_attn_backward_triton_fake(
-    grad_output_padded: torch.Tensor, 
-    q_padded: torch.Tensor, 
-    k_padded: torch.Tensor, 
-    v_padded: torch.Tensor, 
-    o_padded: torch.Tensor, 
-    M: 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()
-    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_triton(ctx, grad_output1, grad_output2):
-    q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes = ctx.saved_tensors
-    dq, dk, dv = block_sparse_attn_backward_triton(grad_output1, q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes)
-    return dq, dk, dv, None, None
-
-
-def setup_context_triton(ctx, inputs, output):
-    q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs
-    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)
-
-
-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)

csrc/attn/vsa/__init__.py

@@ -1,13 +1,12 @@
 import torch
 from typing import Tuple
 block_sparse_attn=None
-import torch
-major, minor = torch.cuda.get_device_capability(0)
-if major == 9 and minor == 0:# check if H100
+
+try:
     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
-else:
+except ImportError:
     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

@@ -568,6 +568,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
         __syncthreads(); // wait for sd_smem shared memory write
         warpgroup::mm_AtB(qg_reg, ds_smem_t[0], k_smem[0]); //delat dQ = dSK
         warpgroup::mma_commit_group();
+        tma::store_async_wait();
         warpgroup::mma_async_wait();
         // store qg to shared memory
         warpgroup::store(qg_smem, qg_reg);
@@ -577,7 +578,6 @@ 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();
         }
     }
 
@@ -624,6 +624,7 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals<D> g) {
         __syncthreads(); // wait for sd_smem shared memory write
         warpgroup::mm_AtB(qg_reg, ds_smem_t[0], k_smem[0]); //delat dQ = dSK
         warpgroup::mma_commit_group();
+        tma::store_async_wait();
         warpgroup::mma_async_wait();
         // store qg to shared memory
         warpgroup::store(qg_smem, qg_reg);
@@ -633,14 +634,13 @@ 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 +1174,4 @@ block_sparse_attention_backward(torch::Tensor q,
 
     return {qg, kg, vg};
     //cudadevicesynchronize();
-}
+}

csrc/attn/vsa/block_sparse_wrapper.py

@@ -0,0 +1,182 @@
+import torch
+try:
+    from vsa_cuda import block_sparse_fwd, block_sparse_bwd
+except ImportError:
+    block_sparse_fwd = None
+    block_sparse_bwd = None
+from vsa.block_sparse_attn_triton import triton_block_sparse_attn_forward, triton_block_sparse_attn_backward
+assert torch.__version__ >= "2.4.0", "VSA requires PyTorch 2.4.0 or higher"
+from vsa.index import map_to_index
+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,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    block_map = block_map.int()
+    q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
+    o, M = triton_block_sparse_attn_forward(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, variable_block_sizes)
+    return o, M
+
+
+@torch.library.register_fake("vsa::block_sparse_attn_triton")
+def _block_sparse_attn_triton_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    block_map: torch.Tensor,
+    variable_block_sizes: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    o = torch.empty_like(q)
+    M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32)
+    return o, M
+
+
+@torch.library.custom_op("vsa::block_sparse_attn_backward_triton", mutates_args=(), device_types="cuda")
+def block_sparse_attn_backward_triton(
+    grad_output_padded: torch.Tensor, 
+    q_padded: torch.Tensor, 
+    k_padded: torch.Tensor, 
+    v_padded: torch.Tensor, 
+    o_padded: torch.Tensor, 
+    M: 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()
+    q2k_block_sparse_index, q2k_block_sparse_num = map_to_index(block_map)
+    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(block_map.transpose(-1, -2))
+    dq, dk, dv = triton_block_sparse_attn_backward(grad_output_padded, q_padded, k_padded, v_padded, o_padded, M, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, variable_block_sizes)
+    return dq, dk, dv
+
+@torch.library.register_fake("vsa::block_sparse_attn_backward_triton")
+def _block_sparse_attn_backward_triton_fake(
+    grad_output_padded: torch.Tensor, 
+    q_padded: torch.Tensor, 
+    k_padded: torch.Tensor, 
+    v_padded: torch.Tensor, 
+    o_padded: torch.Tensor, 
+    M: 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()
+    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_triton(ctx, grad_output1, grad_output2):
+    q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes = ctx.saved_tensors
+    dq, dk, dv = block_sparse_attn_backward_triton(grad_output1, q_padded, k_padded, v_padded, o_padded, M, block_map, variable_block_sizes)
+    return dq, dk, dv, None, None
+
+
+def setup_context_triton(ctx, inputs, output):
+    q_padded, k_padded, v_padded, block_map, variable_block_sizes = inputs
+    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)

docker/Dockerfile.python3.10

@@ -1,9 +1,7 @@
-FROM nvidia/cuda:12.8.0-devel-ubuntu22.04
+FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
-SHELL ["/bin/bash", "-c"]
-
 WORKDIR /FastVideo
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -11,25 +9,17 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     git \
     ca-certificates \
     openssh-server \
-    zsh \
-    vim \
-    curl \
-    gcc-11 \
-    g++-11 \
-    clang-11 \
     && rm -rf /var/lib/apt/lists/*
 
-# Set up C++20 compilers for ThunderKittens
-RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+    rm Miniconda3-latest-Linux-x86_64.sh
 
-# Set CUDA environment variables
-ENV CUDA_HOME=/usr/local/cuda-12.8
-ENV PATH=${CUDA_HOME}/bin:${PATH}
-ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+ENV PATH=/opt/conda/bin:$PATH
 
-# Install uv and source its environment
-RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
-    echo 'source $HOME/.local/bin/env' >> /root/.bashrc
+RUN conda create --name fastvideo-dev python=3.10.0 -y
+
+SHELL ["/bin/bash", "-c"]
 
 # Copy just the pyproject.toml first to leverage Docker cache
 COPY pyproject.toml ./
@@ -37,36 +27,22 @@ COPY pyproject.toml ./
 # Create a dummy README to satisfy the installation
 RUN echo "# Placeholder" > README.md
 
-# Create and activate virtual environment with specific Python version and seed
-RUN source $HOME/.local/bin/env && \
-    uv venv --python 3.10 --seed /opt/venv && \
-    source /opt/venv/bin/activate && \
-    uv pip install --no-cache-dir --upgrade pip && \
-    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation
+RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
+    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
+    conda clean -afy
 
 COPY . .
 
-# Install dependencies using uv and set up shell configuration
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    uv pip install --no-cache-dir -e .[dev] && \
-    git config --unset-all http.https://github.com/.extraheader || true && \
-    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
-    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
 
-# Install STA (Sliding Tile Attention)
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn/sliding_tile_attn && \
-    git submodule update --init --recursive && \
-    python setup.py install
+# Remove authentication headers
+RUN git config --unset-all http.https://github.com/.extraheader || true
 
-# Install VSA 
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
-    git submodule update --init --recursive && \
-    python setup.py install
+# Set up automatic conda environment activation for all shells
+RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
+    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
+    # Ensure .bashrc is sourced for SSH login shells
+    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-EXPOSE 22
+EXPOSE 22

docker/Dockerfile.python3.11

@@ -1,9 +1,7 @@
-FROM nvidia/cuda:12.8.0-devel-ubuntu22.04
+FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
-SHELL ["/bin/bash", "-c"]
-
 WORKDIR /FastVideo
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -11,25 +9,17 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     git \
     ca-certificates \
     openssh-server \
-    zsh \
-    vim \
-    curl \
-    gcc-11 \
-    g++-11 \
-    clang-11 \
     && rm -rf /var/lib/apt/lists/*
 
-# Set up C++20 compilers for ThunderKittens
-RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+    rm Miniconda3-latest-Linux-x86_64.sh
 
-# Set CUDA environment variables
-ENV CUDA_HOME=/usr/local/cuda-12.8
-ENV PATH=${CUDA_HOME}/bin:${PATH}
-ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
+ENV PATH=/opt/conda/bin:$PATH
 
-# Install uv and source its environment
-RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
-    echo 'source $HOME/.local/bin/env' >> /root/.bashrc
+RUN conda create --name fastvideo-dev python=3.11.11 -y
+
+SHELL ["/bin/bash", "-c"]
 
 # Copy just the pyproject.toml first to leverage Docker cache
 COPY pyproject.toml ./
@@ -37,36 +27,22 @@ COPY pyproject.toml ./
 # Create a dummy README to satisfy the installation
 RUN echo "# Placeholder" > README.md
 
-# Create and activate virtual environment with specific Python version and seed
-RUN source $HOME/.local/bin/env && \
-    uv venv --python 3.11 --seed /opt/venv && \
-    source /opt/venv/bin/activate && \
-    uv pip install --no-cache-dir --upgrade pip && \
-    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation
+RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
+    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
+    conda clean -afy
 
 COPY . .
 
-# Install dependencies using uv and set up shell configuration
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    uv pip install --no-cache-dir -e .[dev] && \
-    git config --unset-all http.https://github.com/.extraheader || true && \
-    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
-    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
 
-# Install STA (Sliding Tile Attention)
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn/sliding_tile_attn && \
-    git submodule update --init --recursive && \
-    python setup.py install
+# Remove authentication headers
+RUN git config --unset-all http.https://github.com/.extraheader || true
 
-# Install VSA 
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
-    git submodule update --init --recursive && \
-    python setup.py install
+# Set up automatic conda environment activation for all shells
+RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
+    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
+    # Ensure .bashrc is sourced for SSH login shells
+    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-EXPOSE 22
+EXPOSE 22

docker/Dockerfile.python3.12

@@ -43,7 +43,7 @@ RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     uv pip install --no-cache-dir --upgrade pip && \
     uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation
+    uv pip install --no-cache-dir flash-attn==2.8.0.post2 --no-build-isolation
 
 COPY . .
 
@@ -58,15 +58,15 @@ RUN source $HOME/.local/bin/env && \
 # Install STA (Sliding Tile Attention)
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn/sliding_tile_attn && \
+    cd csrc/attn && \
     git submodule update --init --recursive && \
-    python setup.py install
+    python setup_sta.py install
 
 # Install VSA 
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
+    cd csrc/attn && \
     git submodule update --init --recursive && \
-    python setup.py install
+    python setup_vsa.py install
 
-EXPOSE 22
+EXPOSE 22

docker/Dockerfile.python3.12.cuda12.9.1

@@ -1,72 +0,0 @@
-FROM nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04
-
-ENV DEBIAN_FRONTEND=noninteractive
-
-SHELL ["/bin/bash", "-c"]
-
-WORKDIR /FastVideo
-
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    wget \
-    git \
-    ca-certificates \
-    openssh-server \
-    zsh \
-    vim \
-    curl \
-    gcc-11 \
-    g++-11 \
-    clang-11 \
-    && rm -rf /var/lib/apt/lists/*
-
-# Set up C++20 compilers for ThunderKittens
-RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-# Set CUDA environment variables
-ENV CUDA_HOME=/usr/local/cuda-12.9
-ENV PATH=${CUDA_HOME}/bin:${PATH}
-ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-
-# Install uv and source its environment
-RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
-    echo 'source $HOME/.local/bin/env' >> /root/.bashrc
-
-# Copy just the pyproject.toml first to leverage Docker cache
-COPY pyproject.toml ./
-
-# Create a dummy README to satisfy the installation
-RUN echo "# Placeholder" > README.md
-
-# Create and activate virtual environment with specific Python version and seed
-RUN source $HOME/.local/bin/env && \
-    uv venv --python 3.12 --seed /opt/venv && \
-    source /opt/venv/bin/activate && \
-    uv pip install --no-cache-dir --upgrade pip && \
-    uv pip install --no-cache-dir .[dev] && \
-    uv pip install --no-cache-dir flash-attn==2.8.3 --no-build-isolation
-
-COPY . .
-
-# Install dependencies using uv and set up shell configuration
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    uv pip install --no-cache-dir -e .[dev] && \
-    git config --unset-all http.https://github.com/.extraheader || true && \
-    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
-    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
-
-# Install STA (Sliding Tile Attention)
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn/sliding_tile_attn && \
-    git submodule update --init --recursive && \
-    python setup.py install
-
-# Install VSA 
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn/video_sparse_attn && \
-    git submodule update --init --recursive && \
-    python setup.py install
-
-EXPOSE 22

docs/source/conf.py

@@ -96,7 +96,8 @@ copybutton_prompt_is_regexp = True
 #
 html_title = project
 html_theme = 'sphinx_book_theme'
-html_logo = '../../assets/logos/icon_simple.svg'
+html_logo = '../../assets/logo.jpg'
+#html_favicon = 'assets/logos/vllm-logo-only-light.ico'
 html_theme_options = {
     'path_to_docs': 'docs/source',
     'repository_url': 'https://github.com/hao-ai-lab/FastVideo/',

docs/source/contributing/developer_env/docker.md

@@ -3,12 +3,12 @@
 
 If you prefer a containerized development environment or want to avoid managing dependencies manually, you can use our prebuilt Docker image:
 
-**Images:** [`ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest`](https://ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev)
+**Image:** [`ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest`](https://ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev)
 
 ## Starting the container
 
 ```bash
-docker run --gpus all -it ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest
+docker run --gpus all -it ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
 ```
 
 This will:

docs/source/contributing/developer_env/runpod.md

@@ -6,7 +6,7 @@ You can easily use the FastVideo Docker image as a custom container on [RunPod](
 
 ## Creating a new pod
 
-Choose a GPU that supports CUDA 12.8
+Choose a GPU that supports CUDA 12.4
 
 Pick 1 or 2 L40S GPU(s)
 

docs/source/contributing/overview.md

@@ -22,20 +22,10 @@ source ~/.bashrc
 Create and activate a Conda environment for FastVideo:
 
 ```
-conda create -n fastvideo python=3.12 -y
+conda create -n fastvideo python=3.10 -y
 conda activate fastvideo
 ```
 
-Install `uv` (optional, but recommended):
-
-From instructions on [uv](https://astral.sh/uv/):
-
-```
-curl -LsSf https://astral.sh/uv/install.sh | sh
-# or 
-wget -qO- https://astral.sh/uv/install.sh | sh
-```
-
 Clone the FastVideo repository and go to the FastVideo directory:
 
 ```
@@ -46,10 +36,10 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 Now you can install FastVideo and setup git hooks for running linting. By using `pre-commit`, the linters will run and have to pass before you'll be able to make a commit.
 
 ```bash
-uv pip install -e .[dev]
+pip install -e .[dev]
 
 # Can also install flash-attn (optional)
-uv pip install flash-attn --no-build-isolation 
+pip install flash-attn==2.7.4.post1 --no-build-isolation 
 
 # Linting, formatting and static type checking
 pre-commit install --hook-type pre-commit --hook-type commit-msg
@@ -60,14 +50,3 @@ pre-commit run --all-files
 # Unit tests
 pytest tests/
 ```
-
-If you are on a Hopper GPU, you should also install [FA3](https://github.com/Dao-AILab/flash-attention) for much better performance:
-
-```
-git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention/hopper
-
-# make sure you have ninja installed
-uv pip install ninja
-
-python setup.py install
-```

docs/source/distillation/data_preprocess.md

@@ -1,43 +0,0 @@
-(v0-data-preprocess)=
-
-# 🧱 Data Preprocess for Distillation
-
-For distillation, we use the same data preprocessing pipeline as training. Please refer to the [Training Data Preprocess](../training/data_preprocess.md) for general preprocessing steps.
-
-## Distillation-Specific Datasets
-
-### FastVideo 480P Synthetic Wan Dataset
-
-For Wan2.1 T2V distillation, we use the **FastVideo 480P Synthetic Wan dataset** ([FastVideo/Wan-Syn_77x448x832_600k](https://huggingface.co/datasets/FastVideo/Wan-Syn_77x448x832_600k)) which contains 600k synthetic latents.
-
-```bash
-# Download the preprocessed dataset
-python scripts/huggingface/download_hf.py \
-    --repo_id "FastVideo/Wan-Syn_77x448x832_600k" \
-    --local_dir "FastVideo/Wan-Syn_77x448x832_600k" \
-    --repo_type "dataset"
-```
-
-### Crush Smol Dataset
-
-For Wan2.2 TI2V distillation, we use the crush_smol dataset which includes both raw videos and preprocessed latents.
-
-```bash
-# Download dataset
-python scripts/huggingface/download_hf.py \
-    --repo_id=FastVideo/mini_i2v_dataset \
-    --local_dir=data/mini_i2v_dataset \
-    --repo_type=dataset
-```
-
-## Preprocessing for Distillation
-
-The preprocessing steps are identical to training. Run the appropriate preprocessing script based on your model:
-
-```bash
-# For Wan2.1 T2V
-bash scripts/preprocess/v1_preprocess_wan_data_t2v
-
-# For Wan2.2 TI2V  
-bash examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/preprocess_wan_data_ti2v_5b.sh
-```

docs/source/distillation/dmd.md

@@ -1,87 +0,0 @@
-# 🎯 Distillation
-
-We introduce a new finetuning strategy - **Sparse-distill**, which jointly integrates **[DMD](https://arxiv.org/abs/2405.14867)** and **[VSA](https://arxiv.org/abs/2505.13389)** in a single training process. This approach combines the benefits of both distillation to shorten diffusion steps and sparse attention to reduce attention computations, enabling much faster video generation.
-
-## 📊 Model Overview
-
-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 **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.
-
-### Download Dataset
-
-```bash
-# Download the preprocessed dataset
-python scripts/huggingface/download_hf.py \
-    --repo_id "FastVideo/Wan-Syn_77x448x832_600k" \
-    --local_dir "FastVideo/Wan-Syn_77x448x832_600k" \
-    --repo_type "dataset"
-```
-
-## 🚀 Training Scripts
-
-### Wan2.1 1.3B Model Sparse-Distill
-
-For the 1.3B model, we use **4 nodes with 32 H200 GPUs** (8 GPUs per node):
-
-```bash
-# Multi-node training (8 nodes, 64 GPUs total)
-sbatch examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm
-```
-
-**Key Configuration:**
-- Global batch size: 64
-- Gradient accumulation steps: 2
-- Learning rate: 1e-5
-- VSA attention sparsity: 0.8
-- Training steps: 4000 (~12 hours)
-
-### Wan2.1 14B Model Sparse-Distill
-
-For the 14B model, we use **8 nodes with 64 H200 GPUs** (8 GPUs per node):
-
-```bash
-# Multi-node training (8 nodes, 64 GPUs total)
-sbatch examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm
-```
-
-**Key Configuration:**
-- Global batch size: 64
-- Sequence parallel size: 4
-- Gradient accumulation steps: 4
-- Learning rate: 1e-5
-- VSA attention sparsity: 0.9
-- Training steps: 3000 (~52 hours)
-- HSDP shard dim: 8
-
-### Wan2.2 5B Model Sparse-Distill
-
-For the 5B model, we use **8 nodes with 64 H200 GPUs** (8 GPUs per node):
-
-```bash
-# 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/generate_examples.py

@@ -287,35 +287,11 @@ def create_nested_structures(
         relative_path = example.path.relative_to(category_dir)
         path_parts = relative_path.parts
 
-        if example.category == "training":
-            # For training examples like finetune/wan_i2v_14b_480p/crush_smol
-            if len(path_parts) >= 3:
-                method = path_parts[0]  # e.g., "finetune"
-                model = path_parts[1]  # e.g., "wan_i2v_14b_480p"
-                dataset = path_parts[2]  # e.g., "crush_smol"
-
-                # Initialize nested structure
-                if example.category not in nested_structures:
-                    nested_structures[example.category] = {}
-                if method not in nested_structures[example.category]:
-                    nested_structures[example.category][method] = {}
-                if model not in nested_structures[example.category][method]:
-                    nested_structures[example.category][method][model] = {}
-
-                # Store the nested structure
-                nested_structures[
-                    example.category][method][model][dataset] = NestedStructure(
-                        category=example.category,
-                        method=method,
-                        model=model,
-                        dataset=dataset,
-                        example=example)
-
-        elif example.category == "distillation" and len(path_parts) >= 2:
-            # For distillation examples like Wan2.1-T2V/Wan-Syn-Data-480P
-            model = path_parts[0]  # e.g., "Wan2.1-T2V"
-            dataset = path_parts[1]  # e.g., "Wan-Syn-Data-480P"
-            method = "DMD"  # Default method for distillation
+        # For nested examples like finetune/wan_i2v_14b_480p/crush_smol
+        if len(path_parts) >= 3:
+            method = path_parts[0]  # e.g., "finetune"
+            model = path_parts[1]  # e.g., "wan_i2v_14b_480p"
+            dataset = path_parts[2]  # e.g., "crush_smol"
 
             # Initialize nested structure
             if example.category not in nested_structures:

docs/source/getting_started/installation/gpu.md

@@ -6,7 +6,7 @@ Instructions to install FastVideo for NVIDIA CUDA GPUs.
 
 - **OS: Linux or Windows WSL**
 - **Python: 3.10-3.12**
-- **CUDA 12.8**
+- **CUDA 12.4**
 - **At least 1 NVIDIA GPU**
 
 ## Set up using Python
@@ -38,7 +38,6 @@ conda activate fastvideo
 
 :::{tip}
 We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
-Note that you can also use `uv` to install FastVideo in a Conda environment.
 :::
 
 Or you can create a new Python environment using [uv](https://docs.astral.sh/uv/), a very fast Python environment manager. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following command:
@@ -61,7 +60,7 @@ uv pip install fastvideo
 Also optionally install flash-attn:
 
 ```bash
-pip install flash-attn --no-build-isolation
+pip install flash-attn==2.7.4.post1 --no-build-isolation
 ```
 
 ### Installation from Source
@@ -88,7 +87,7 @@ uv pip install -e .
 #### Flash Attention
 
 ```bash
-pip install flash-attn --no-build-isolation
+pip install flash-attn==2.7.4.post1 --no-build-isolation
 ```
 
 ## Set up using Docker
@@ -103,7 +102,7 @@ If you're planning to contribute to FastVideo please see the following page:
 ## Hardware Requirements
 
 ### For Basic Inference
-- NVIDIA GPU with CUDA 12.8 support
+- NVIDIA GPU with CUDA 12.4 support
 
 ### For Lora Finetuning
 - 40GB GPU memory each for 2 GPUs with lora

docs/source/getting_started/installation/mps.md

@@ -39,7 +39,6 @@ conda activate fastvideo
 
 :::{tip}
 We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
-Note that you can also use `uv` to install FastVideo in a Conda environment.
 :::
 
 Or you can create a new Python environment using [uv](https://docs.astral.sh/uv/), a very fast Python environment manager. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following command:

docs/source/index.md

@@ -1,6 +1,6 @@
 # Welcome to FastVideo
 
-:::{figure} ../../assets/logos/logo.svg
+:::{figure} ../../assets/logo.jpg
 :align: center
 :alt: FastVideo
 :class: no-scaled-link
@@ -9,7 +9,7 @@
 
 :::{raw} html
 <p style="text-align:center">
-<strong>FastVideo is a unified inference and post-training framework for accelerated video generation.
+<strong>FastVideo is a unified framework for accelerated video generation.
 </strong>
 </p>
 
@@ -21,10 +21,11 @@
 </p>
 :::
 
-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.
+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.
 
 <div style="text-align: center;">
-  <img src=_static/images/fastwan.png width="100%"/>
+  <img src=_static/images/perf.png width="100%"/>
 </div>
 
 ## Key Features
@@ -34,11 +35,16 @@ FastVideo has the following features:
   - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
   - [TeaCache](https://arxiv.org/pdf/2411.19108)
   - [Sage Attention](https://arxiv.org/abs/2410.02367)
-- 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.
+- 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.
 
 ## Documentation
 
@@ -72,16 +78,18 @@ inference/add_pipeline
 
 training/examples/examples_training_index
 training/data_preprocess
+training/distillation
 <!-- training/finetune -->
 :::
 
-:::{toctree}
+<!-- :::{toctree}
 :caption: Distillation
 :maxdepth: 1
 
 distillation/examples/examples_distillation_index
 distillation/data_preprocess
-distillation/dmd
+distillation/dmd -->
+<!-- training/finetune -->
 :::
 
 % What is STA Kernel?
@@ -94,15 +102,6 @@ sliding_tile_attention/installation
 sliding_tile_attention/demo
 :::
 
-% What is VSA Kernel?
-
-:::{toctree}
-:caption: Video Sparse Attention
-:maxdepth: 1
-
-video_sparse_attention/installation
-:::
-
 :::{toctree}
 :caption: Design
 :maxdepth: 1

docs/source/inference/inference_quick_start.md

@@ -5,7 +5,7 @@ This page contains step-by-step instructions to get you quickly started with vid
 ## Requirements
 - **OS**: Linux (Tested on Ubuntu 22.04+)
 - **Python**: 3.10-3.12
-- **CUDA**: 12.8
+- **CUDA**: 12.4
 - **GPU**: At least one NVIDIA GPU
 
 ## Installation

docs/source/inference/optimizations.md

@@ -19,7 +19,6 @@ This page describes the various options for speeding up generation times in Fast
 - Torch SDPA: `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`
 - Flash Attention 2 and 3: `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN`
 - Sliding Tile Attention: `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
-- Video Sparse Attention: `FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN`
 - Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
 
 ### Configuring Backends
@@ -75,17 +74,6 @@ pip install st_attn==0.0.4
 
 Please see [this page](#sta-installation) for more installation instructions.
 
-(optimizations-vsa)=
-### Video Sparse Attention
-**`VIDEO_SPARSE_ATTN`**
-
-```bash
-git submodule update --init --recursive
-python setup_vsa.py install
-```
-
-Please see [this page](#vsa-installation) for more installation instructions.
-
 (optimizations-sage)=
 ### Sage Attention
 **`SAGE_ATTN`**

docs/source/inference/support_matrix.md

@@ -6,7 +6,6 @@ The symbols used have the following meanings:
 
 - ✅ = Full compatibility
 - ❌ = No compatibility
-- ⭕ = Does not apply to this model
 
 ## Models x Optimization
 The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods and FastVideo will use the optimal default parameters when initializing and generating videos.
@@ -38,94 +37,51 @@ The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods a
   * TeaCache
   * Sliding Tile Attn
   * Sage Attn
-  * Video Sparse Attention (VSA)
-- * FastWan2.1 T2V 1.3B
-  * `FastVideo/FastWan2.1-T2V-1.3B-Diffusers`
-  * 480P
-  * ⭕
-  * ⭕
-  * ⭕
-  * ✅
-- * FastWan2.2 TI2V 5B Full Attn*
-  * `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers`
-  * 720P
-  * ⭕
-  * ⭕
-  * ⭕
-  * ✅
-- * Wan2.2 TI2V 5B 
-  * `Wan-AI/Wan2.2-TI2V-5B-Diffusers`
-  * 720P
-  * ⭕
-  * ⭕
-  * ✅
-  * ⭕
-- * Wan2.2 T2V A14B
-  * `Wan-AI/Wan2.2-T2V-A14B-Diffusers`
-  * 480P<br>720P
-  * ❌
-  * ❌
-  * ✅
-  * ⭕
-- * Wan2.2 I2V A14B
-  * `Wan-AI/Wan2.2-I2V-A14B-Diffusers`
-  * 480P<br>720P
-  * ❌
-  * ❌
-  * ✅
-  * ⭕
 - * HunyuanVideo
   * `hunyuanvideo-community/HunyuanVideo`
   * 720px1280p<br>544px960p
   * ❌
   * ✅
   * ✅
-  * ⭕
 - * FastHunyuan
   * `FastVideo/FastHunyuan-diffusers`
   * 720px1280p<br>544px960p
   * ❌
   * ✅
   * ✅
-  * ⭕
-- * Wan2.1 T2V 1.3B
+- * Wan T2V 1.3B
   * `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`
   * 480P
   * ✅
   * ✅*
   * ✅
-  * ⭕
-- * Wan2.1 T2V 14B
+- * Wan T2V 14B
   * `Wan-AI/Wan2.1-T2V-14B-Diffusers`
   * 480P, 720P
   * ✅
   * ✅*
   * ✅
-  * ⭕
-- * Wan2.1 I2V 480P
+- * Wan I2V 480P
   * `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers`
   * 480P
   * ✅
   * ✅*
   * ✅
-  * ⭕
-- * Wan2.1 I2V 720P
+- * Wan I2V 720P
   * `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers`
   * 720P
   * ✅
+  * ✅*
   * ✅
-  * ✅
-  * ⭕
 - * StepVideo T2V
   * `FastVideo/stepvideo-t2v-diffusers`
   * 768px768px204f<br>544px992px204f<br>544px992px136f
   * ❌
   * ❌
   * ✅
-  * ⭕
 :::
 
-**Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
+**Note**: there are some known quality issues with Wan2.1 + Sliding Tile Attn. We are working on fixing this issue.
 
 ## Special requirements
 

docs/source/sliding_tile_attention/installation.md

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

@@ -0,0 +1,25 @@
+(v0-distill)=
+# 🎯 Distill
+Our distillation recipe is based on [Phased Consistency Model](https://github.com/G-U-N/Phased-Consistency-Model). We did not find significant improvement using multi-phase distillation, so we keep the one phase setup similar to the original latent consistency model's recipe.
+We use the [MixKit](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/all_mixkit) dataset for distillation. To avoid running the text encoder and VAE during training, we prprocess all data to generate text embeddings and VAE latents.
+Preprocessing instructions can be found [data_preprocess.md](#v0-data-preprocess). For convenience, we also provide preprocessed data that can be downloaded directly using the following command:
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/HD-Mixkit-Finetune-Hunyuan --local_dir=data/HD-Mixkit-Finetune-Hunyuan --repo_type=dataset
+```
+
+Next, download the original model weights with:
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model # original hunyuan
+python scripts/huggingface/download_hf.py --repo_id=genmo/mochi-1-preview --local_dir=data/mochi --repo_type=model # original mochi
+```
+
+To launch the distillation process, use the following commands:
+
+```
+bash scripts/distill/distill_hunyuan.sh # for hunyuan
+bash scripts/distill/distill_mochi.sh # for mochi
+```
+
+We also provide an optional script for distillation with adversarial loss, located at `fastvideo/distill_adv.py`. Although we tried adversarial loss, we did not observe significant improvements.

docs/source/training/finetune.md

@@ -7,7 +7,7 @@ Ensure your data is prepared and preprocessed in the format specified in [data_p
 python scripts/huggingface/download_hf.py --repo_id=FastVideo/Mochi-Black-Myth --local_dir=data/Mochi-Black-Myth --repo_type=dataset
 ```
 
-Download the original model weights as specified in the [Distillation Section](../distillation/dmd.md):
+Download the original model weights as specified in [Distill Section](#v0-distill):
 
 Then you can run the finetune with:
 

docs/source/video_sparse_attention/installation.md

@@ -1,66 +0,0 @@
-(vsa-installation)=
-
-# 🔧 Installation
-You can install the Video Sparse Attention package using
-
-```bash
-pip install vsa
-```
-
-# Building from Source
-We support H100 (via ThunderKittens) and any other GPU (via Triton) for VSA.
-
-First, install C++20 for ThunderKittens (if using H100):
-
-```bash
-sudo apt update
-sudo apt install gcc-11 g++-11
-
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-sudo apt update
-sudo apt install clang-11
-```
-
-Set up CUDA environment (if using CUDA 12.8):
-
-```bash
-export CUDA_HOME=/usr/local/cuda-12.8
-export PATH=${CUDA_HOME}/bin:${PATH} 
-export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-```
-
-Install VSA:
-
-```bash
-cd csrc/attn/video_sparse_attn/
-git submodule update --init --recursive
-python setup.py install
-```
-
-# 🧪 Test
-
-```bash
-python csrc/attn/tests/test_vsa.py
-```
-
-# 📋 Usage
-
-```python
-from vsa import video_sparse_attn
-
-# q, k, v: [batch_size, num_heads, seq_len, head_dim]
-# variable_block_sizes: [num_blocks] - number of valid tokens in each block
-# topk: int - number of top-k blocks to attend to
-# block_size: int or tuple of 3 ints - size of each block (default: 64 tokens)
-# compress_attn_weight: optional weight for compressed attention branch
-
-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/datasets/crush_smol/validation.json

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

examples/datasets/crush_smol/validation_single.json

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

examples/datasets/wan2.1_syn_data/download_dataset_720p.sh

@@ -1,4 +0,0 @@
-#!/bin/bash
-
-# 720P dataset
-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x768x1280_250k" --local_dir "FastVideo/Wan-Syn_77x768x1280_250k" --repo_type "dataset"

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

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

examples/distill/Wan2.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 "6.0" # not used for dmd inference
+  --validation_guidance_scale "1.0" # not used for dmd inference
 )
 
 # Optimizer arguments
@@ -102,6 +102,7 @@ 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 "6.0" # not used for dmd inference
+  --validation_guidance_scale "1.0" # not used for dmd inference
 )
 
 # Optimizer arguments
@@ -102,6 +102,7 @@ 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 "6.0" # not used for dmd inference
+  --validation_guidance_scale "1.0" # not used for dmd inference
 )
 
 # Optimizer arguments
@@ -102,6 +102,7 @@ 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

@@ -1,11 +0,0 @@
-# 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
-pip install vsa
-```
-
-### 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

@@ -1,144 +0,0 @@
-#!/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

@@ -1,145 +0,0 @@
-#!/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/crush_smol/README.md

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

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 "6.0" # not used for dmd inference
+  --validation_guidance_scale "1.0" # not used for dmd inference
 )
 
 # Optimizer arguments
@@ -77,6 +77,7 @@ 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

@@ -1,6 +1,6 @@
 from fastvideo import VideoGenerator
 
-# from fastvideo.configs.sample import SamplingParam
+from fastvideo.configs.sample import SamplingParam
 
 OUTPUT_PATH = "video_samples"
 def main():
@@ -8,29 +8,35 @@ def main():
     # model.
     # If a local path is provided, FastVideo will make a best effort
     # attempt to identify the optimal arguments.
+    model_name = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
     generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        model_name,
         # FastVideo will automatically handle distributed setup
         num_gpus=1,
         use_fsdp_inference=True,
-        dit_cpu_offload=False,
+        dit_cpu_offload=True,
         vae_cpu_offload=False,
         text_encoder_cpu_offload=True,
-        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        # 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"
+    sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.image_path = "test.jpg"
+    # sampling_param.num_inference_steps = 0
     # 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."
-    )
-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True)
+    i2v_prompt = "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+    i2v_prompt = "A little girl is packing a suitcase and the contents starts flying out of the suitcase everywhere."
+    prompt = i2v_prompt
+    # prompt = (
+    #     "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
+    #     "wide with interest. The playful yet serene atmosphere is complemented by soft "
+    #     "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
+    # )
+    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
     # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
+    return
 
     # Generate another video with a different prompt, without reloading the
     # model!

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=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        pin_cpu_memory=False,
         dit_cpu_offload=False,
         vae_cpu_offload=False,
         VSA_sparsity=0.8,
@@ -31,6 +31,7 @@ def main():
     prompt = (
         "A neon-lit alley in futuristic Tokyo during a heavy rainstorm at night. The puddles reflect glowing signs in kanji, advertising ramen, karaoke, and VR arcades. A woman in a translucent raincoat walks briskly with an LED umbrella. Steam rises from a street food cart, and a cat darts across the screen. Raindrops are visible on the camera lens, creating a cinematic bokeh effect."
     )
+    prompt = "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."
     start_time = time.perf_counter()
     video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
     end_time = time.perf_counter()
@@ -45,7 +46,7 @@ def main():
         "embodying the raw energy of the wild. Low angle, steady tracking shot, "
         "cinematic.")
     start_time = time.perf_counter()
-    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=False)
     end_time = time.perf_counter()
     gen_time2 = end_time - start_time
 

examples/inference/basic/basic_self_forcing_causal.py

@@ -1,31 +0,0 @@
-import os
-import time
-from fastvideo import VideoGenerator, SamplingParam
-
-OUTPUT_PATH = "video_samples_causal"
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    model_name = "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
-    generator = VideoGenerator.from_pretrained(
-        model_name,
-        # FastVideo will automatically handle distributed setup
-        num_gpus=1,
-        use_fsdp_inference=True,
-        text_encoder_cpu_offload=False,
-        dit_cpu_offload=False,
-    )
-
-    sampling_param = SamplingParam.from_pretrained(model_name)
-
-    prompt = (
-        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
-        "wide with interest. The playful yet serene atmosphere is complemented by soft "
-        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
-    )
-    video = generator.generate_video(prompt, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
-
-if __name__ == "__main__":
-    main()

examples/inference/basic/basic_wan2_2.py

@@ -1,48 +0,0 @@
-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

@@ -1,30 +0,0 @@
-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/basic/basic_wan2_2_ti2v.py

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

examples/inference/basic_i2v.py

@@ -0,0 +1,38 @@
+from fastvideo import VideoGenerator
+
+from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples"
+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.1-I2V-14B-480P-Diffusers",
+        # FastVideo will automatically handle distributed setup
+        num_gpus=2,
+        use_fsdp_inference=True,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        image_encoder_cpu_offload=False,
+    )
+
+    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-I2V-14B-480P-Diffusers")
+    sampling_param.num_frames = 61
+    sampling_param.num_inference_steps = 40
+    sampling_param.guidance_scale = 5.0
+    sampling_param.height = 448
+    sampling_param.width = 832
+    sampling_param.seed = 1024
+    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 = (
+        "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
+    )
+    video = generator.generate_video(prompt, sampling_param=sampling_param, output_path=OUTPUT_PATH, save_video=True)
+
+
+if __name__ == "__main__":
+    main()

examples/inference/gradio/README.md

@@ -0,0 +1,59 @@
+# 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/README_I2V.md

@@ -0,0 +1,182 @@
+# FastVideo Dual Model Setup (T2V + I2V)
+
+This document describes the dual model functionality that has been added to the FastVideo Gradio app, supporting both Text-to-Video (T2V) and Image-to-Video (I2V) generation modes with specialized models.
+
+## Overview
+
+The Gradio app now supports both Text-to-Video (T2V) and Image-to-Video (I2V) generation modes using specialized models:
+- **T2V Model**: `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` for text-to-video generation
+- **I2V Model**: `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` for image-to-video generation
+
+Users can switch between modes using the tabbed interface and upload images for I2V generation.
+
+## Model Configuration
+
+### T2V Model
+- **Model**: `FastVideo/FastWan2.1-T2V-1.3B-Diffusers`
+- **Purpose**: Text-to-video generation
+- **Default Parameters**: Optimized for text prompts
+
+### I2V Model  
+- **Model**: `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers`
+- **Purpose**: Image-to-video generation
+- **Default Parameters**: Optimized for image animation
+
+## Changes Made
+
+### Backend Changes (`ray_serve_backend.py`)
+
+1. **Added dual model support**:
+   - Separate model paths for T2V and I2V
+   - Automatic model selection based on request type
+   - Independent model initialization
+
+2. **Updated `VideoGenerationRequest`**:
+   - Added `model_type` field ("t2v" or "i2v")
+   - Added `image_path` field for I2V input
+
+3. **Enhanced `FastVideoAPI` class**:
+   - Dual model initialization (`t2v_generator` and `i2v_generator`)
+   - Separate default parameters for each model
+   - Automatic model selection in `generate_video` method
+
+### Frontend Changes (`gradio_frontend.py`)
+
+1. **Tabbed interface**:
+   - "Text-to-Video" tab for T2V generation
+   - "Image-to-Video" tab for I2V generation
+
+2. **Automatic model selection**:
+   - T2V tab uses T2V model automatically
+   - I2V tab uses I2V model automatically
+   - Model type sent in API requests
+
+3. **Separate event handlers**:
+   - `handle_t2v_generation` for text-to-video
+   - `handle_i2v_generation` for image-to-video
+
+## Usage
+
+### Starting the Application
+
+1. **Using the combined startup script (recommended)**:
+   ```bash
+   python start_ray_serve_app.py
+   ```
+
+2. **Manual startup**:
+   ```bash
+   # Start backend
+   python ray_serve_backend.py \
+     --t2v_model_path "FastVideo/FastWan2.1-T2V-1.3B-Diffusers" \
+     --i2v_model_path "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+   
+   # Start frontend
+   python gradio_frontend.py --backend_url "http://localhost:8000"
+   ```
+
+### Using T2V Mode
+
+1. Navigate to the "Text-to-Video" tab
+2. Enter a text prompt describing the video you want to generate
+3. Adjust advanced parameters if needed
+4. Click "Run" to generate the video
+
+### Using I2V Mode
+
+1. Navigate to the "Image-to-Video" tab
+2. Upload an image using the image upload component
+3. Enter a prompt describing how the image should animate
+4. Adjust advanced parameters if needed
+5. Click "Run" to generate the video
+
+### Example Prompts
+
+**T2V Examples**:
+- "A hand enters the frame, pulling a sheet of plastic wrap over three balls of dough placed on a wooden surface."
+- "A vintage train snakes through the mountains, its plume of white steam rising dramatically against the jagged peaks."
+
+**I2V Examples**:
+- "The image comes to life with subtle movement, the scene gently animating while maintaining the original composition and mood."
+- "The static image transforms into a dynamic scene with natural motion, preserving the original lighting and atmosphere."
+
+## Testing
+
+A comprehensive test script is provided to verify both T2V and I2V functionality:
+
+```bash
+python test_i2v.py
+```
+
+This script:
+- Tests backend health
+- Tests T2V functionality with text prompts
+- Tests I2V functionality with image uploads
+- Verifies response formats for both modes
+- Cleans up test files
+
+## Technical Details
+
+### Backend API Changes
+
+The `/generate_video` endpoint now accepts:
+
+```json
+{
+  "prompt": "Animation description",
+  "model_type": "t2v",  // or "i2v"
+  "image_path": "/path/to/input/image.png",  // for I2V
+  // ... other parameters
+}
+```
+
+### Model Selection Logic
+
+- **T2V Mode**: Uses `FastVideo/FastWan2.1-T2V-1.3B-Diffusers`
+- **I2V Mode**: Uses `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers`
+- **Automatic Selection**: Based on presence of `image_path` and `model_type`
+
+### Memory Management
+
+- Both models are loaded independently
+- Automatic cleanup prevents memory leaks
+- Temporary files are cleaned up after processing
+
+## Configuration Options
+
+### Command Line Arguments
+
+**Backend**:
+- `--t2v_model_path`: Path to T2V model
+- `--i2v_model_path`: Path to I2V model
+- `--output_path`: Output directory
+- `--host`, `--port`: Server configuration
+
+**Frontend**:
+- `--backend_url`: Backend API URL
+- `--t2v_model_path`, `--i2v_model_path`: Model paths (for reference)
+- `--host`, `--port`: Server configuration
+
+## Troubleshooting
+
+1. **Model Loading Issues**: Ensure both models are accessible
+2. **Memory Issues**: The backend includes automatic cleanup
+3. **Image Upload Failures**: Check image format and size
+4. **Generation Failures**: Check backend logs for detailed errors
+
+## Performance Considerations
+
+- **Model Loading**: Both models are loaded at startup
+- **Memory Usage**: Higher memory requirements due to dual models
+- **Generation Time**: I2V may take longer due to larger model size
+- **GPU Requirements**: Ensure sufficient VRAM for both models
+
+## Future Enhancements
+
+Potential improvements:
+- Model switching without restart
+- Batch processing for both modes
+- Advanced image preprocessing
+- Progress indicators
+- Result caching and history
+- Model-specific parameter optimization 

examples/inference/gradio/README_ray_serve.md

@@ -0,0 +1,222 @@
+# FastVideo with Ray Serve Backend and Gradio Frontend
+
+This setup provides a scalable web application for FastVideo inference using Ray Serve as the backend and Gradio as the frontend.
+
+## Architecture
+
+- **Backend**: Ray Serve handles video generation requests with GPU acceleration
+- **Frontend**: Gradio provides a user-friendly web interface
+- **Communication**: HTTP REST API between frontend and backend
+
+## Features
+
+- ✅ Scalable backend with Ray Serve
+- ✅ GPU-accelerated video generation
+- ✅ User-friendly Gradio interface
+- ✅ Health monitoring and error handling
+- ✅ All original functionality preserved
+- ✅ Easy deployment and management
+
+## Installation
+
+1. Install the additional dependencies:
+
+```bash
+pip install -r requirements_ray_serve.txt
+```
+
+2. Ensure you have the FastVideo model available (the default is `FastVideo/FastHunyuan-diffusers`)
+
+## Usage
+
+### Option 1: Start Both Services Together (Recommended)
+
+Use the startup script to launch both backend and frontend:
+
+```bash
+python start_ray_serve_app.py
+```
+
+This will:
+- Start the Ray Serve backend on port 8000
+- Start the Gradio frontend on port 7860
+- Monitor both services and provide unified logging
+- Handle graceful shutdown with Ctrl+C
+
+### Option 2: Start Services Separately
+
+#### Start Backend Only
+
+```bash
+python ray_serve_backend.py --model_path FastVideo/FastHunyuan-diffusers --output_path outputs
+```
+
+#### Start Frontend Only
+
+```bash
+python gradio_frontend.py --backend_url http://localhost:8000 --model_path FastVideo/FastHunyuan-diffusers
+```
+
+## Configuration
+
+### Command Line Arguments
+
+#### Startup Script (`start_ray_serve_app.py`)
+
+- `--model_path`: Path to the FastVideo model (default: `FastVideo/FastHunyuan-diffusers`)
+- `--output_path`: Directory to save generated videos (default: `outputs`)
+- `--backend_host`: Backend host to bind to (default: `0.0.0.0`)
+- `--backend_port`: Backend port (default: `8000`)
+- `--frontend_host`: Frontend host to bind to (default: `0.0.0.0`)
+- `--frontend_port`: Frontend port (default: `7860`)
+- `--skip_backend_check`: Skip backend health check
+
+#### Backend (`ray_serve_backend.py`)
+
+- `--model_path`: Path to the FastVideo model
+- `--output_path`: Directory to save generated videos
+- `--host`: Host to bind to
+- `--port`: Port to bind to
+
+#### Frontend (`gradio_frontend.py`)
+
+- `--backend_url`: URL of the Ray Serve backend
+- `--model_path`: Path to the model (for default parameters)
+- `--host`: Host to bind to
+- `--port`: Port to bind to
+
+### Environment Variables
+
+You can also set these environment variables:
+
+- `FASTVIDEO_MODEL_PATH`: Path to the FastVideo model
+- `FASTVIDEO_OUTPUT_PATH`: Directory to save generated videos
+- `RAY_SERVE_HOST`: Backend host
+- `RAY_SERVE_PORT`: Backend port
+- `GRADIO_HOST`: Frontend host
+- `GRADIO_PORT`: Frontend port
+
+## API Endpoints
+
+### Backend API (Ray Serve)
+
+- `GET /health`: Health check endpoint
+- `POST /generate_video`: Video generation endpoint
+
+#### Video Generation Request
+
+```json
+{
+  "prompt": "A beautiful sunset over the ocean",
+  "negative_prompt": "blurry, low quality",
+  "use_negative_prompt": true,
+  "seed": 42,
+  "guidance_scale": 7.5,
+  "num_frames": 21,
+  "height": 512,
+  "width": 512,
+  "num_inference_steps": 20,
+  "randomize_seed": false
+}
+```
+
+#### Video Generation Response
+
+```json
+{
+  "output_path": "/path/to/generated/video.mp4",
+  "seed": 42,
+  "success": true,
+  "error_message": null
+}
+```
+
+## Deployment
+
+### Local Development
+
+1. Start the application:
+   ```bash
+   python start_ray_serve_app.py
+   ```
+
+2. Access the frontend at: `http://localhost:7860`
+3. Access the backend API at: `http://localhost:8000`
+
+### Production Deployment
+
+For production deployment, consider:
+
+1. **Load Balancing**: Use a reverse proxy (nginx, traefik) in front of the services
+2. **Monitoring**: Add monitoring and logging (Prometheus, Grafana)
+3. **Scaling**: Configure Ray Serve for horizontal scaling
+4. **Security**: Add authentication and rate limiting
+5. **Storage**: Use shared storage for video outputs
+
+### Docker Deployment
+
+Create a Dockerfile for containerized deployment:
+
+```dockerfile
+FROM python:3.9-slim
+
+WORKDIR /app
+
+# Install dependencies
+COPY requirements_ray_serve.txt .
+RUN pip install -r requirements_ray_serve.txt
+
+# Copy application files
+COPY . .
+
+# Expose ports
+EXPOSE 8000 7860
+
+# Start the application
+CMD ["python", "start_ray_serve_app.py"]
+```
+
+## Troubleshooting
+
+### Common Issues
+
+1. **Backend not starting**: Check GPU availability and Ray installation
+2. **Frontend can't connect**: Verify backend URL and network connectivity
+3. **Video generation fails**: Check model path and GPU memory
+4. **Port conflicts**: Change ports using command line arguments
+
+### Logs
+
+- Backend logs are prefixed with `[BACKEND]`
+- Frontend logs are prefixed with `[FRONTEND]`
+- Use `--skip_backend_check` if you need to debug startup issues
+
+### Performance Tuning
+
+- Adjust `num_replicas` in the Ray Serve deployment for scaling
+- Configure `max_concurrent_queries` based on GPU memory
+- Use multiple GPUs by modifying `ray_actor_options`
+
+## Migration from Original Gradio Demo
+
+The new setup maintains full compatibility with the original functionality:
+
+1. All parameters and options are preserved
+2. The same example prompts are included
+3. The UI layout and behavior are identical
+4. Video generation quality is the same
+
+The main differences are:
+- Backend processing is now handled by Ray Serve
+- Better error handling and status monitoring
+- Scalable architecture for production use
+- Separation of concerns between frontend and backend
+
+## Contributing
+
+To extend this setup:
+
+1. Add new endpoints to `ray_serve_backend.py`
+2. Update the frontend in `gradio_frontend.py`
+3. Modify the startup script if needed
+4. Update this README with new features 

examples/inference/gradio/backend_test.py

@@ -0,0 +1,501 @@
+import argparse
+import os
+import time
+import json
+import statistics
+import asyncio
+import aiohttp
+from copy import deepcopy
+from typing import List, Dict, Any
+import threading
+
+import torch
+
+# All the prompts for stress testing
+STRESS_TEST_PROMPTS = [
+    "A person reading a book with words that float off the pages and form pictures.",
+    "A person diving into a pool of liquid crystal, creating ripples of light.",
+    "A handheld shot chasing after a group of friends laughing and playing on the beach at sunset.",
+    "A mysterious ancient temple hidden in the jungle.",
+    "A high-speed train navigating a steep descent.",
+    "a toy robot wearing blue jeans and a white t shirt taking a pleasant stroll in Antarctica during a winter storm",
+    "A cheetah accelerating to full speed while chasing its prey.",
+    "A serene orchard is in full bloom, with trees heavy with blossoms and bees buzzing around, darting from flower to flower in a display of natural harmony.",
+    "A little child let out a big yawn",
+    "Subtle reflections of a woman on the window of a train moving at hyper-speed in a Japanese city.",
+    "A truck left along the edge of a cliff, revealing the stunning coastal landscape below with waves crashing against the rocks.",
+    "A red bird transforms into a flag",
+    "A zoom-out from a single leaf on a tree to reveal the entire forest, showcasing the vastness and diversity of the woodland.",
+    "A slow-motion video of a liquid droplet bouncing on a water-repellent surface.",
+    "Static camera shot. A dinasour running near some lions and chasing them away.",
+    "an adorable kangaroo wearing purple overalls and cowboy boots taking a pleasant stroll in Mumbai India during a beautiful sunset",
+    "A zoom-in on an artist's brush touching the canvas, highlighting the texture of the paint and the strokes being made.",
+    "an old man wearing blue jeans and a white t shirt taking a pleasant stroll in Mumbai India during a colorful festival",
+    "A woman is ascending to the sky from the ground",
+    "View out a window of a giant strange creature walking in rundown city at night, one single street lamp dimly lighting the area.",
+    "An arc shot around a lone tree in a vast, foggy field at dawn, revealing the changing light and shadows.",
+    "A person sculpting a statue out of a waterfall, the water solidifying under their touch.",
+    "The person's forehead creased with concentration as she worked on a challenging puzzle.",
+    "The person's cheeks flushed with pleasure as she savored a delicious meal.",
+    "Hand-drawn simple line art, a young kid looking up into space with a wondrous expression on his face.",
+    "A crab made of different jewlery is walking on the beach. As it walks, it drops different jewelry pieces like diamonds, pearls, etc",
+    "Gold coins are falling out when elevator door opens",
+    "the scene transitions from huge waves into a snowy mountain at sunset",
+    "a giant cathedral is completely filled with cats. there are cats everywhere you look. a man enters the cathedral and bows before the giant cat king sitting on a throne.",
+    "A mother dog gently picks up a piece of meat and carefully places it in her puppy's bowl, her eyes filled with warmth and care as she watches her little one eat.",
+    "A soap bubble floating in the air, displaying iridescent colors that shift and change as it moves through different angles of light.",
+    "A truck left alongside a train moving through the countryside, matching its speed and revealing the changing landscape.",
+    "An astronaut walking between stone buildings.",
+    "A close-up shot of the person's face reveals his fear and desperation as he navigates the ship through the storm.",
+    "A frozen lake slowly cracking and thawing as spring arrives, with sheets of ice breaking apart and drifting across the surface.",
+    "A FPV shot zooming through a tunnel into a vibrant underwater space.",
+    "a toy robot wearing blue jeans and a white t shirt taking a pleasant stroll in Mumbai India during a colorful festival",
+    "A person sips on a smoothie, the cool and fruity flavors refreshing her mouth.",
+    "In a vibrant theater, a magician in dazzling attire stands center stage, pulling a comically oversized rubber chicken from an ornate, old-fashioned box. His costume shimmers under the stage lights, adding to the spectacle. The crowd erupts in laughter and applause, their faces filled with joy and amazement. The magician's expression hints at mischievous delight as he holds up the rubber chicken, his performance bringing cheer to the audience.",
+    "A hamster running on a spinning wheel.",
+    "A quaint village nestled in a valley is surrounded by blooming cherry blossoms, with petals drifting through the air as villagers go about their daily activities, adding life to the scene.",
+    "In a tranquil forest clearing, a sparkling waterfall cascades down into a clear pool, surrounded by lush greenery and flowers, with occasional birds fluttering by.",
+    "A woman beamed with pride as she watched her child perform on stage.",
+    "an adorable kangaroo wearing blue jeans and a white t shirt taking a pleasant stroll in Mumbai India during a winter storm",
+    "A man is eating salad",
+    "An Asian girl wearing a bright yellow T-shirt and white pants is Hip-Hop dancing",
+    "nighttime footage of a hermit crab using an incandescent lightbulb as its shell",
+    "a toy robot wearing a green dress and a sun hat taking a pleasant stroll in Antarctica during a beautiful sunset",
+    "A goat operating a food truck, serving gourmet grilled cheese sandwiches to a line of animals.",
+    "Macro shot. Man in an antique scuba helmet with dark glass walking out of a flower",
+    "A bustling train station in the heart of a vibrant city.",
+    "Light filtering through a canopy of autumn leaves, casting warm, dappled patterns of yellow, orange, and red onto the ground.",
+    "Chimneys in the setting sun",
+    "A longboarder accelerating downhill, carving through turns.",
+    "A couple runs through a sudden downpour, laughing and splashing in puddles as they try to find shelter.",
+    "A glass of iced coffee condensing water on the outside, with droplets forming and sliding down the glass in slow motion.",
+    "macro shot of a leaf showing tiny trains moving through its veins",
+    "A corgi wearing sunglasses walks on the beach of a tropical island",
+    "Borneo wildlife on the Kinabatangan River",
+    "A beautiful silhouette animation shows a wolf howling at the moon, feeling lonely, until it finds its pack.",
+    "an adorable kangaroo wearing blue jeans and a white t shirt taking a pleasant stroll in Johannesburg South Africa during a colorful festival",
+    "A green monster made of plants walks through an airport.",
+    "A close up view of a glass sphere that has a zen garden within it. There is a small dwarf in the sphere who is raking the zen garden and creating patterns in the sand.",
+    "A person on a scooter colliding with a park bench, the scooter tipping over.",
+    "A tilt-up from a city street, ascending to show the skyline with its mix of modern and historic architecture.",
+    "A chef tossing a pancake into the air and catching it.",
+    "A woman whispering a secret into a friend's ear.",
+    "A vulture circling high in the sky.",
+    "A medieval castle overlooking a bustling renaissance fair.",
+    "a toy robot wearing purple overalls and cowboy boots taking a pleasant stroll in Mumbai India during a beautiful sunset",
+    "A man standing in front of a burning building giving the 'thumbs up' sign.",
+    "The person's cheeks flushed with embarrassment as he told a funny story.",
+    "Llamas and Emus are playing chess",
+    "A woman sipping a steaming cup of tea.",
+    "A tree root bursting through the seat of an ancient, weathered bench, intertwining with the wood.",
+    "Smoke rises from the chimney of a cozy log cabin nestled in the woods, with soft light glowing from the windows, suggesting a warm and inviting atmosphere.",
+    "A close-up of sparkling water being poured into a glass, capturing the detailed flow and bubbles.",
+    "a woman wearing blue jeans and a white t shirt taking a pleasant stroll in Antarctica during a beautiful sunset",
+    "The Glenfinnan Viaduct is a historic railway bridge in Scotland, UK, that crosses over the west highland line between the towns of Mallaig and Fort William. It is a stunning sight as a steam train leaves the bridge, traveling over the arch-covered viaduct. The landscape is dotted with lush greenery and rocky mountains, creating a picturesque backdrop for the train journey. The sky is blue and the sun is shining, making for a beautiful day to explore this majestic spot.",
+    "A piece of elastic fabric being pulled and stretched, then returning to its original size when the tension is released.",
+    "a woman wearing a green dress and a sun hat taking a pleasant stroll in Antarctica during a beautiful sunset",
+    "A video of a water jet cutting through metal, showing the powerful and precise movement of water.",
+    "Car mirrors and sunsets",
+    "Giant Pandas are eating hot noodles in a Chinese restaurant",
+    "A rally car taking a fast turn on a track",
+    "a toy robot wearing purple overalls and cowboy boots taking a pleasant stroll in Mumbai India during a colorful festival",
+    "A crystal-clear icicle slowly dripping as it melts in the warmth of the midday sun, each drop sparkling as it falls.",
+    "A tilt-down from a chandelier in a grand hall, revealing the ornate decor and people mingling below.",
+    "A man is playing the drums under the water",
+    "A person playing an electric guitar made of lightning, with thunderous sound waves.",
+    "A person floating in a bubble, drifting over a bustling cityscape.",
+    "A tilt-down from a starry night sky, revealing a quiet forest clearing bathed in moonlight.",
+    "A pan right through a dense jungle, moving past lush vegetation and exotic wildlife.",
+    "Close-up of a man eating an apple.",
+    "A low-angle shot of a dancer leaping gracefully into the air, making their movement appear even more dynamic and powerful.",
+    "A woman is search her bag trying to find something.",
+    "A bulldozer clears debris from a demolished building, making way for new construction.",
+    "A man sighed in relief as the doctor delivered the good news.",
+    "A tsunami coming through an alley in Bulgaria, dynamic movement.",
+    "Blooming Flowers",
+    "A push-in through a dense crowd at a festival, moving towards a performer on stage who is captivating the audience.",
+    "A truck right through a tranquil garden, moving past blooming flowers, trees, and a small fountain.",
+    "The person's eyes sparkled with excitement as he greeted a friend.",
+    "A person playing chess with a robot on a floating platform above the ocean.",
+    "A gentle breeze rustles the leaves as someone walks down a serene forest path, sunlight filtering through the trees and shifting patterns on the ground as branches sway.",
+    "A rollercoaster ride from a city to a desert and then to an ice world",
+    "A pan left across an ancient library, moving from shelf to shelf, showcasing rows of leather-bound books.",
+    "A mother otter floating on her back in a river, cradling her pup on her stomach to keep it safe and warm in the gentle current.",
+    "an adorable kangaroo wearing purple overalls and cowboy boots taking a pleasant stroll in Johannesburg South Africa during a colorful festival",
+    "a woman wearing a green dress and a sun hat taking a pleasant stroll in Mumbai India during a colorful festival",
+    "A delicate layer of morning frost melting off a flower petal, the tiny droplets glistening like diamonds in the light.",
+    "A panda is cooking for her child, her child is next to her.",
+    "Macro shot of a man wearing an antique diving helmet with dark glass and a jetpack walking on the veins of a leaf. Realistic style",
+    "an old man wearing purple overalls and cowboy boots taking a pleasant stroll in Johannesburg South Africa during a beautiful sunset",
+    "A girl is unfolding a birthday gift.",
+    "A pencil drawing an architectural plan.",
+    "A handheld camera following a dog running through a park, bouncing and tilting as it captures the dog's joyful exploration.",
+    "A pan left across a serene beach at sunrise, moving from the darkened shore to the brightening horizon.",
+    "A group of people are clapping to celebrate",
+    "Vendors set up stalls at a bustling farmer's market, displaying fresh fruits and vegetables, while people stroll through, selecting produce and enjoying the lively atmosphere.",
+    "A police helicopter hovers above a high-speed chase, guiding officers on the ground to apprehend a suspect.",
+    "A paper origami dragon riding a boat in waves. Realistic style.",
+    "A close-up of a droplet of dew forming on a leaf, capturing the detailed surface tension.",
+    "a toy robot wearing blue jeans and a white t shirt taking a pleasant stroll in Mumbai India during a beautiful sunset",
+    "A dry rainbow rose is coming back to life.",
+    "A glass falling off a table and shattering on the floor.",
+    "A marathon runner crossing the finish line after a grueling race.",
+    "A zoom-in on a drop of morning dew on a leaf, showing the reflection of the surrounding world within it.",
+    "A child blowing on hot cocoa to cool it down.",
+    "A squad of futsal players showcasing their skills on an indoor court.",
+    "A princess is brushing her long golden hair in the garden.",
+    "A close-up of a pair of eyes, revealing the subtle emotions and reflections within them.",
+    "A tracking shot of a group of cyclists racing through a forest trail, with trees and foliage rushing by.",
+    "A woman yawning widely at the end of a long day.",
+    "an old man wearing a green dress and a sun hat taking a pleasant stroll in Johannesburg South Africa during a colorful festival",
+    "Hidden within a garden, an ancient fountain trickles with water, surrounded by vibrant flowers and lush greenery that seem to whisper secrets of the past.",
+    "A Chinese man sits at a table and eats noodles with chopsticks",
+    "A pink pig running fast toward the camera in an alley in Tokyo.",
+    "Strange creatures move through a mysterious, foggy marsh, their silhouettes barely visible through the dense mist as they navigate the eerie, otherworldly landscape.",
+    "Tour of an art gallery with many beautiful works of art in different styles.",
+    "FPV flying through a colorful coral lined streets of an underwater suburban neighborhood.",
+    "Aerial view of Santorini during the blue hour, showcasing the stunning architecture of white Cycladic buildings with blue domes. The caldera views are breathtaking, and the lighting creates a beautiful, serene atmosphere.",
+    "Camera zoom out. A couple walking along the beach as the sun sets over the ocean.",
+    "an extreme close up shot of a woman's eye, with her iris appearing as earth",
+    "a woman wearing purple overalls and cowboy boots taking a pleasant stroll in Mumbai India during a colorful festival",
+    "an old man wearing a green dress and a sun hat taking a pleasant stroll in Mumbai India during a winter storm",
+    "an adorable kangaroo wearing blue jeans and a white t shirt taking a pleasant stroll in Antarctica during a winter storm",
+    "A martial artist breaking a board with a powerful punch.",
+    "People gather on a peaceful beach at sunset, a bonfire crackling as they sit around, enjoying the warmth and the sight of the sun dipping below the horizon.",
+    "A close-up of a waterfall, showing the detailed movement of water as it crashes down.",
+    "A child is blowing bubbles",
+    "a woman wearing a green dress and a sun hat taking a pleasant stroll in Johannesburg South Africa during a winter storm",
+    "A wide-angle perspective of a serene lake surrounded by mountains, reflecting the sky and creating a sense of infinite space.",
+    "The person's eyebrows arched in skepticism as she listened to a dubious claim.",
+    "an old man wearing blue jeans and a white t shirt taking a pleasant stroll in Mumbai India during a beautiful sunset",
+    "a woman wearing a green dress and a sun hat taking a pleasant stroll in Johannesburg South Africa during a colorful festival",
+    "a woman wearing purple overalls and cowboy boots taking a pleasant stroll in Antarctica during a colorful festival",
+    "a toy robot wearing blue jeans and a white t shirt taking a pleasant stroll in Antarctica during a colorful festival",
+    "A chef flips a pancake and puts cream on it.",
+    "An astronaut runs on the surface of the moon, the low angle shot shows the vast background of the moon, the movement is smooth and appears lightweight",
+    "A man's face lit up with happiness as he received a heartfelt compliment.",
+    "A futuristic spaceport hums with activity as ships of various shapes and sizes take off and land on multiple platforms, their engines glowing with vibrant colors.",
+    "A person knitting a scarf using beams of light instead of yarn.",
+    "A pedestal up from the edge of a canyon, gradually revealing the expansive landscape and river below.",
+    "a woman wearing purple overalls and cowboy boots taking a pleasant stroll in Johannesburg South Africa during a colorful festival",
+    "an old man wearing blue jeans and a white t shirt taking a pleasant stroll in Johannesburg South Africa during a colorful festival",
+    "A person walking up a staircase made of clouds leading to a floating castle.",
+    "Monks meditate in a serene mountaintop temple, sitting in quiet reflection as the wind gently moves through the surrounding trees, creating a sense of peace and tranquility.",
+    "An aerial shot of a bustling city intersection at rush hour, capturing the organized chaos of cars and pedestrians.",
+    "A pair of hands skillfully knitting a colorful scarf, the yarn winding through their fingers with each stitch.",
+    "Close-up, a Chinese child is eating dumplings",
+    "A kite losing wind and falling to the ground.",
+    "Bioluminescent waves gently wash ashore on a deserted beach, illuminating the sand with each cresting wave as a figure walks along the water's edge, leaving glowing footprints.",
+    "A red panda taking a bite of a pizza",
+    "A close-up shot of a young woman driving a car, looking thoughtful, blurred green forest visible through the rainy car window.",
+    "A high-speed video of a splash created by a stone thrown into a pond.",
+    "A metal rod being bent slightly by a force and then springing back to its original straight shape when the force is removed.",
+    "A hedgehog in a knight's armor, riding a toy horse into a medieval castle.",
+    "A bird made of fresh oranges rushes out of the orange",
+    "A low altitude first person perspective camera tracking shot of a soccer player's feet dribbling the ball on the groud in a soccer field, Sports Videography, Motion Tracking camera shot",
+    "A tranquil island retreat features swaying palm trees and hammocks strung between them, inviting guests to relax and enjoy the serene beauty of the surroundings.",
+    "a spooky haunted mansion, with friendly jack o lanterns and ghost characters welcoming trick or treaters to the entrance, tilt shift photography",
+    "A coconut tree made of dollar bills at sunset, with bills falling off like leaves.",
+    "A motocross bike accelerating out of a tight turn on a dirt track.",
+    "A tranquil Zen garden with a gently flowing stream and koi fish.",
+    "A green monster made of leaves walks through the airport, carrying a suitcase.",
+    "A time-lapse of a frost-covered leaf gradually thawing in the morning sunlight, with tiny water droplets forming and trickling down.",
+    "A woman practicing her archery skills at a range.",
+    "A slow-motion video of ink being injected into a tank of water, creating intricate and beautiful patterns.",
+    "a woman wearing blue jeans and a white t shirt taking a pleasant stroll in Johannesburg South Africa during a winter storm",
+    "The person's forehead creased with worry as he listened to bad news.",
+    "An arc shot around a grand piano being played in an empty concert hall, the motion revealing the intricate details of the instrument.",
+    "A person conducting a symphony of animals in a forest clearing.",
+    "A truck right alongside a flowing river, capturing the movement of the water and the surrounding forest.",
+    "A rocket blasting off from the launch pad, accelerating rapidly into the sky.",
+    "Workers move through a picturesque vineyard during the harvest season, carefully picking grapes and placing them into baskets as the sun bathes the vines in a warm glow.",
+    "A person is eating an ice cream.",
+    "An over-the-shoulder perspective of a chef meticulously plating a dish in a bustling kitchen.",
+    "A man looked away in shame when confronted with his wrongdoing.",
+    "A person is savoring a slice of pizza at a pizzeria."
+]
+
+class BackendStressTest:
+    def __init__(self, output_path: str, 
+                 server_url: str = "http://localhost:8000", max_concurrent: int = 50):
+        self.output_path = output_path
+        self.server_url = server_url
+        self.max_concurrent = max_concurrent
+        
+        # Results storage
+        self.results = []
+        self.lock = threading.Lock()
+        
+    async def check_health(self) -> bool:
+        """Check if the Ray Serve backend is healthy"""
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.get(f"{self.server_url}/health", timeout=aiohttp.ClientTimeout(total=5)) as response:
+                    return response.status == 200
+        except Exception:
+            return False
+        
+    def build_request_params(self, prompt: str, **kwargs) -> Dict[str, Any]:
+        """Build request parameters for Ray Serve backend"""
+        # Default parameters matching the Ray Serve backend
+        default_params = {
+            'prompt': prompt,
+            'negative_prompt': None,
+            'use_negative_prompt': False,
+            'seed': 42,
+            'guidance_scale': 7.5,
+            'num_frames': 21,
+            'height': 448,
+            'width': 832,
+            'num_inference_steps': 20,
+            'randomize_seed': True,
+            'return_frames': False  # Don't return frames for stress testing to reduce overhead
+        }
+        
+        # Override with any provided kwargs
+        for key, value in kwargs.items():
+            if key in default_params:
+                default_params[key] = value
+        
+        # Randomize seed if requested
+        if default_params.get('randomize_seed', True):
+            default_params['seed'] = torch.randint(0, 1000000, (1,)).item()
+            
+        # Handle negative prompt
+        if not default_params.get('use_negative_prompt', False):
+            default_params['negative_prompt'] = None
+        
+        # NEW: Remove keys with None values to avoid sending nulls that may break validation
+        clean_params = {k: v for k, v in default_params.items() if v is not None}
+        return clean_params
+    
+    async def test_single_request(self, session: aiohttp.ClientSession, prompt: str, request_id: int) -> Dict[str, Any]:
+        """Test a single request and measure latency"""
+        start_time = time.time()
+        
+        try:
+            # Build request parameters
+            request_params = self.build_request_params(prompt)
+            
+            # Make request to Ray Serve backend
+            async with session.post(
+                f"{self.server_url}/generate_video",
+                json=request_params,
+                timeout=aiohttp.ClientTimeout(total=900)  # 15 minute timeout for video generation
+            ) as response:
+                
+                end_time = time.time()
+                latency = end_time - start_time
+                
+                if response.status == 200:
+                    response_data = await response.json()
+                    if response_data.get('success', False):
+                        result = {
+                            'request_id': request_id,
+                            'prompt': prompt,
+                            'latency': latency,
+                            'status': 'success',
+                            'response_time': latency,  # Use our own timing
+                            'timestamp': start_time,
+                            'output_path': response_data.get('output_path', ''),
+                            'used_seed': response_data.get('seed', request_params['seed'])
+                        }
+                    else:
+                        result = {
+                            'request_id': request_id,
+                            'prompt': prompt,
+                            'latency': latency,
+                            'status': 'error',
+                            'error': response_data.get('error_message', 'Unknown backend error'),
+                            'timestamp': start_time
+                        }
+                else:
+                    response_text = await response.text()
+                    result = {
+                        'request_id': request_id,
+                        'prompt': prompt,
+                        'latency': latency,
+                        'status': 'error',
+                        'error': f"HTTP {response.status}: {response_text}",
+                        'timestamp': start_time
+                    }
+                    
+        except Exception as e:
+            end_time = time.time()
+            latency = end_time - start_time
+            result = {
+                'request_id': request_id,
+                'prompt': prompt,
+                'latency': latency,
+                'status': 'error',
+                'error': str(e),
+                'timestamp': start_time
+            }
+        
+        # Thread-safe result storage
+        with self.lock:
+            self.results.append(result)
+            
+        return result
+    
+    async def run_stress_test(self, num_iterations: int = 1, concurrent_requests: int = None):
+        """Run the stress test with multiple iterations and concurrent requests"""
+        if concurrent_requests is None:
+            concurrent_requests = self.max_concurrent
+            
+        # Check backend health before starting
+        print(f"Testing Ray Serve backend at {self.server_url}...")
+        if not await self.check_health():
+            print(f"❌ Backend is not healthy at {self.server_url}")
+            print("Make sure the Ray Serve backend is running with:")
+            print("python ray_serve_backend.py")
+            return
+        print("✅ Backend is healthy and ready for stress testing")
+        
+        print(f"\nStarting stress test with {len(STRESS_TEST_PROMPTS)} prompts")
+        print(f"Running {num_iterations} iteration(s) with {concurrent_requests} concurrent requests")
+        print(f"Total requests: {len(STRESS_TEST_PROMPTS) * num_iterations}")
+        print(f"Backend URL: {self.server_url}")
+        print("-" * 80)
+        
+        all_prompts = STRESS_TEST_PROMPTS * num_iterations
+        request_id = 0
+        
+        # Create semaphore to limit concurrent requests
+        semaphore = asyncio.Semaphore(concurrent_requests)
+        
+        async def limited_request(session: aiohttp.ClientSession, prompt: str, req_id: int):
+            async with semaphore:
+                return await self.test_single_request(session, prompt, req_id)
+        
+        # Run concurrent requests using asyncio
+        async with aiohttp.ClientSession() as session:
+            # Create all tasks
+            tasks = [
+                limited_request(session, prompt, request_id + i)
+                for i, prompt in enumerate(all_prompts)
+            ]
+            
+            # Process completed requests as they finish
+            completed = 0
+            for coro in asyncio.as_completed(tasks):
+                try:
+                    result = await coro
+                    completed += 1
+                    prompt = result['prompt']
+                    status_icon = "✅" if result['status'] == 'success' else "❌"
+                    output_info = f" -> {result.get('output_path', 'N/A')}" if result['status'] == 'success' else ""
+                    print(f"{status_icon} [{completed}/{len(all_prompts)}] {result['latency']:.2f}s - {prompt[:50]}...{output_info}")
+                except Exception as e:
+                    completed += 1
+                    print(f"❌ [{completed}/{len(all_prompts)}] Exception: {e}")
+        
+        self.analyze_results()
+    
+    def analyze_results(self):
+        """Analyze and print test results"""
+        print("\n" + "=" * 80)
+        print("STRESS TEST RESULTS")
+        print("=" * 80)
+        
+        successful_requests = [r for r in self.results if r['status'] == 'success']
+        failed_requests = [r for r in self.results if r['status'] == 'error']
+        
+        print(f"Total Requests: {len(self.results)}")
+        print(f"Successful: {len(successful_requests)}")
+        print(f"Failed: {len(failed_requests)}")
+        print(f"Success Rate: {len(successful_requests)/len(self.results)*100:.1f}%")
+        
+        if successful_requests:
+            latencies = [r['latency'] for r in successful_requests]
+            print(f"\nLatency Statistics (seconds):")
+            print(f"  Min: {min(latencies):.2f}")
+            print(f"  Max: {max(latencies):.2f}")
+            print(f"  Mean: {statistics.mean(latencies):.2f}")
+            print(f"  Median: {statistics.median(latencies):.2f}")
+            print(f"  Std Dev: {statistics.stdev(latencies):.2f}")
+            
+            # Percentiles
+            sorted_latencies = sorted(latencies)
+            p50 = sorted_latencies[int(len(sorted_latencies) * 0.5)]
+            p90 = sorted_latencies[int(len(sorted_latencies) * 0.9)]
+            p95 = sorted_latencies[int(len(sorted_latencies) * 0.95)]
+            p99 = sorted_latencies[int(len(sorted_latencies) * 0.99)]
+            
+            print(f"  P50: {p50:.2f}")
+            print(f"  P90: {p90:.2f}")
+            print(f"  P95: {p95:.2f}")
+            print(f"  P99: {p99:.2f}")
+        
+        if failed_requests:
+            print(f"\nFailed Requests ({len(failed_requests)}):")
+            for req in failed_requests[:5]:  # Show first 5 failures
+                print(f"  - {req['error']}")
+            if len(failed_requests) > 5:
+                print(f"  ... and {len(failed_requests) - 5} more")
+        
+        # Save detailed results
+        results_file = os.path.join(self.output_path, "stress_test_results.json")
+        os.makedirs(self.output_path, exist_ok=True)
+        
+        with open(results_file, 'w') as f:
+            json.dump({
+                'summary': {
+                    'total_requests': len(self.results),
+                    'successful_requests': len(successful_requests),
+                    'failed_requests': len(failed_requests),
+                    'success_rate': len(successful_requests)/len(self.results)*100 if self.results else 0
+                },
+                'latency_stats': {
+                    'min': min(latencies) if successful_requests else 0,
+                    'max': max(latencies) if successful_requests else 0,
+                    'mean': statistics.mean(latencies) if successful_requests else 0,
+                    'median': statistics.median(latencies) if successful_requests else 0,
+                    'std_dev': statistics.stdev(latencies) if len(successful_requests) > 1 else 0
+                },
+                'detailed_results': self.results
+            }, f, indent=2)
+        
+        print(f"\nDetailed results saved to: {results_file}")
+
+async def main():
+    parser = argparse.ArgumentParser(description="FastVideo Ray Serve Backend Stress Test")
+    parser.add_argument("--output_path",
+                        type=str,
+                        default="outputs",
+                        help="Path to save test results")
+    parser.add_argument("--server_url",
+                        type=str,
+                        default="http://localhost:8000",
+                        help="Ray Serve backend URL")
+    parser.add_argument("--max_concurrent",
+                        type=int,
+                        default=50,
+                        help="Maximum concurrent requests")
+    parser.add_argument("--iterations",
+                        type=int,
+                        default=1,
+                        help="Number of iterations through all prompts")
+    parser.add_argument("--concurrent_requests",
+                        type=int,
+                        default=None,
+                        help="Number of concurrent requests (overrides max_concurrent)")
+    
+    args = parser.parse_args()
+    
+    # Create stress test instance
+    stress_test = BackendStressTest(
+        output_path=args.output_path,
+        server_url=args.server_url,
+        max_concurrent=args.max_concurrent
+    )
+    
+    # Run the stress test
+    await stress_test.run_stress_test(
+        num_iterations=args.iterations,
+        concurrent_requests=args.concurrent_requests
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/inference/gradio/gradio_demo.py

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

@@ -0,0 +1,182 @@
+import argparse
+import os
+import subprocess
+import sys
+import time
+import threading
+import signal
+import requests
+from pathlib import Path
+
+# Add the project root to the Python path
+project_root = Path(__file__).parent.parent.parent.parent
+sys.path.insert(0, str(project_root))
+
+
+def check_frontend_health(frontend_url: str, max_retries: int = 30) -> bool:
+    """Check if the frontend is healthy"""
+    for i in range(max_retries):
+        try:
+            response = requests.get(frontend_url, timeout=5)
+            if response.status_code == 200:
+                print(f"✅ Frontend is healthy at {frontend_url}")
+                return True
+        except requests.exceptions.RequestException:
+            pass
+        
+        if i < max_retries - 1:
+            print(f"⏳ Waiting for frontend to start... ({i+1}/{max_retries})")
+            time.sleep(2)
+    
+    print(f"❌ Frontend failed to start within {max_retries * 2} seconds")
+    return False
+
+
+def start_frontend_instance(args, instance_id: int, backend_url: str):
+    """Start a single frontend instance"""
+    frontend_script = Path(__file__).parent / "gradio_frontend.py"
+    frontend_port = args.frontend_base_port + instance_id
+    
+    cmd = [
+        sys.executable, str(frontend_script),
+        "--backend_url", backend_url,
+        "--t2v_model_path", args.t2v_model_path,
+        "--i2v_model_path", args.i2v_model_path,
+        "--host", args.frontend_host,
+        "--port", str(frontend_port)
+    ]
+    
+    print(f"🎨 Starting Frontend {instance_id + 1} on port {frontend_port}...")
+    print(f"Command: {' '.join(cmd)}")
+    
+    # Start the frontend process
+    frontend_process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        universal_newlines=True,
+        bufsize=1
+    )
+    
+    # Monitor frontend output
+    def monitor_frontend():
+        for line in frontend_process.stdout:
+            print(f"[FRONTEND-{instance_id + 1}] {line.rstrip()}")
+    
+    monitor_thread = threading.Thread(target=monitor_frontend, daemon=True)
+    monitor_thread.start()
+    
+    return frontend_process, frontend_port
+
+
+def main():
+    parser = argparse.ArgumentParser(description="FastVideo Multi-Frontend Launcher")
+    
+    # Model and output settings
+    parser.add_argument("--t2v_model_path",
+                        type=str,
+                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
+                        help="Path to the T2V model")
+    parser.add_argument("--i2v_model_path",
+                        type=str,
+                        default="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
+                        help="Path to the I2V model")
+    
+    # Frontend settings
+    parser.add_argument("--frontend_host",
+                        type=str,
+                        default="0.0.0.0",
+                        help="Frontend host to bind to")
+    parser.add_argument("--frontend_base_port",
+                        type=int,
+                        default=7860,
+                        help="Base port for frontend instances")
+    parser.add_argument("--num_frontends",
+                        type=int,
+                        default=2,
+                        help="Number of frontend instances to start")
+    
+    # Backend settings
+    parser.add_argument("--backend_url",
+                        type=str,
+                        default="http://localhost:8000",
+                        help="Backend URL for frontends to connect to")
+    
+    # Other settings
+    parser.add_argument("--skip_health_check",
+                        action="store_true",
+                        help="Skip frontend health check")
+    
+    args = parser.parse_args()
+    
+    print("🎬 FastVideo Multi-Frontend Launcher")
+    print("=" * 50)
+    print(f"T2V Model: {args.t2v_model_path}")
+    print(f"I2V Model: {args.i2v_model_path}")
+    print(f"Backend URL: {args.backend_url}")
+    print(f"Number of Frontends: {args.num_frontends}")
+    print(f"Frontend Base Port: {args.frontend_base_port}")
+    print("=" * 50)
+    
+    # Start multiple frontend instances
+    frontend_processes = []
+    frontend_urls = []
+    
+    for i in range(args.num_frontends):
+        process, port = start_frontend_instance(args, i, args.backend_url)
+        frontend_processes.append(process)
+        frontend_urls.append(f"http://{args.frontend_host}:{port}")
+    
+    # Wait for frontends to be ready
+    if not args.skip_health_check:
+        print("\n⏳ Waiting for frontends to start...")
+        for i, url in enumerate(frontend_urls):
+            if not check_frontend_health(url):
+                print(f"❌ Frontend {i + 1} failed to start. Terminating...")
+                for process in frontend_processes:
+                    process.terminate()
+                sys.exit(1)
+    
+    print("\n🎉 All frontend instances are starting up!")
+    for i, url in enumerate(frontend_urls):
+        print(f"📺 Frontend {i + 1}: {url}")
+    print("\nPress Ctrl+C to stop all frontend instances...")
+    
+    # Signal handler for graceful shutdown
+    def signal_handler(signum, frame):
+        print("\n🛑 Shutting down frontend instances...")
+        for process in frontend_processes:
+            process.terminate()
+        
+        # Wait for processes to terminate
+        try:
+            for process in frontend_processes:
+                process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            print("⚠️  Force killing processes...")
+            for process in frontend_processes:
+                process.kill()
+        
+        print("✅ Frontend instances stopped")
+        sys.exit(0)
+    
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
+    
+    # Monitor processes
+    try:
+        while True:
+            # Check if processes are still running
+            for i, process in enumerate(frontend_processes):
+                if process.poll() is not None:
+                    print(f"❌ Frontend {i + 1} process died unexpectedly")
+                    break
+            
+            time.sleep(1)
+    
+    except KeyboardInterrupt:
+        signal_handler(signal.SIGINT, None)
+
+
+if __name__ == "__main__":
+    main()

examples/inference/gradio/nginx.conf

@@ -0,0 +1,181 @@
+# Nginx configuration for FastVideo load balancing
+# This configuration implements the architecture:
+# ngrok -> nginx reverse proxy -> frontend1/frontend2 -> backend1×8/backend2×8
+
+events {
+    worker_connections 1024;
+}
+
+http {
+    # Basic settings
+    sendfile on;
+    tcp_nopush on;
+    tcp_nodelay on;
+    keepalive_timeout 65;
+    types_hash_max_size 2048;
+    client_max_body_size 100M;  # Allow large video uploads
+
+    # Logging
+    access_log /mnt/fast-disks/nfs/hao_lab/FastVideo/outputs/nginx_access.log;
+    error_log  /mnt/fast-disks/nfs/hao_lab/FastVideo/outputs/nginx_error.log;
+
+    # Gzip compression
+    gzip on;
+    gzip_vary on;
+    gzip_min_length 1024;
+    gzip_proxied any;
+    gzip_comp_level 6;
+    gzip_types
+        text/plain
+        text/css
+        text/xml
+        text/javascript
+        application/json
+        application/javascript
+        application/xml+rss
+        application/atom+xml
+        image/svg+xml;
+
+    # Upstream for frontend load balancing
+    upstream frontend_servers {
+        # Round-robin load balancing between frontends
+        server 127.0.0.1:7860 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:7861 weight=1 max_fails=3 fail_timeout=30s;
+    upstream frontend_servers {
+        # Round-robin load balancing between frontends
+        server 127.0.0.1:7860 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:7861 weight=1 max_fails=3 fail_timeout=30s;
+        
+        # Health check
+        keepalive 32;
+    }
+
+    # Upstream for backend1 load balancing
+    upstream backend1_servers {
+        server 127.0.0.1:8000 weight=1 max_fails=3 fail_timeout=30s; (8 replicas)
+    upstream backend1_servers {
+        # Round-robin load balancing for backend1 replicas
+        server 127.0.0.1:8000 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8001 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8002 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8003 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8004 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8005 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8006 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8007 weight=1 max_fails=3 fail_timeout=30s;
+        
+        keepalive 32;
+    }
+
+    # Upstream for backend2 load balancing
+    upstream backend2_servers {
+        server 127.0.0.1:8000 weight=1 max_fails=3 fail_timeout=30s; (8 replicas)
+    upstream backend2_servers {
+        # Round-robin load balancing for backend2 replicas
+        server 127.0.0.1:8010 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8011 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8012 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8013 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8014 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8015 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8016 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8017 weight=1 max_fails=3 fail_timeout=30s;
+        
+        keepalive 32;
+    }
+
+    # Main server block
+    server {
+        listen 80;
+        server_name localhost;
+
+        # Security headers
+        add_header X-Frame-Options "SAMEORIGIN" always;
+        add_header X-Content-Type-Options "nosniff" always;
+        add_header X-XSS-Protection "1; mode=block" always;
+        add_header Referrer-Policy "no-referrer-when-downgrade" always;
+
+        # Frontend routes (Gradio interfaces)
+        location / {
+            proxy_pass http://frontend_servers;
+            proxy_set_header Host $host;
+            proxy_set_header X-Real-IP $remote_addr;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header X-Forwarded-Proto $scheme;
+            
+            # WebSocket support for Gradio
+            proxy_http_version 1.1;
+            proxy_set_header Upgrade $http_upgrade;
+            proxy_set_header Connection "upgrade";
+            
+            # Timeouts
+            proxy_connect_timeout 60s;
+            proxy_send_timeout 60s;
+            proxy_read_timeout 60s;
+            
+            # Buffer settings
+            proxy_buffering on;
+            proxy_buffer_size 128k;
+            proxy_buffers 4 256k;
+            proxy_busy_buffers_size 256k;
+        }
+
+        # Backend API routes for frontend1
+        location /api/frontend1/ {
+            # Strip the /api/frontend1/ prefix
+            rewrite ^/api/frontend1/(.*) /$1 break;
+            proxy_pass http://backend1_servers;
+            proxy_set_header Host $host;
+            proxy_set_header X-Real-IP $remote_addr;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header X-Forwarded-Proto $scheme;
+            
+            # Timeouts for video generation
+            proxy_connect_timeout 300s;
+            proxy_send_timeout 300s;
+            proxy_read_timeout 300s;
+            
+            # Buffer settings for large responses
+            proxy_buffering on;
+            proxy_buffer_size 128k;
+            proxy_buffers 4 256k;
+            proxy_busy_buffers_size 256k;
+        }
+
+        # Backend API routes for frontend2
+        location /api/frontend2/ {
+            # Strip the /api/frontend2/ prefix
+            rewrite ^/api/frontend2/(.*) /$1 break;
+            proxy_pass http://backend2_servers;
+            proxy_set_header Host $host;
+            proxy_set_header X-Real-IP $remote_addr;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header X-Forwarded-Proto $scheme;
+            
+            # Timeouts for video generation
+            proxy_connect_timeout 300s;
+            proxy_send_timeout 300s;
+            proxy_read_timeout 300s;
+            
+            # Buffer settings for large responses
+            proxy_buffering on;
+            proxy_buffer_size 128k;
+            proxy_buffers 4 256k;
+            proxy_busy_buffers_size 256k;
+        }
+
+        # Health check endpoint
+        location /health {
+            access_log off;
+            return 200 "healthy\n";
+            add_header Content-Type text/plain;
+        }
+
+        # Static files (if needed)
+        location /static/ {
+            alias /var/www/static/;
+            expires 1y;
+            add_header Cache-Control "public, immutable";
+        }
+    }
+}

examples/inference/gradio/nginx.conf.backup

@@ -0,0 +1,173 @@
+# Nginx configuration for FastVideo load balancing
+# This configuration implements the architecture:
+# ngrok -> nginx reverse proxy -> frontend1/frontend2 -> backend1×8/backend2×8
+
+events {
+    worker_connections 1024;
+}
+
+http {
+    # Basic settings
+    sendfile on;
+    tcp_nopush on;
+    tcp_nodelay on;
+    keepalive_timeout 65;
+    types_hash_max_size 2048;
+    client_max_body_size 100M;  # Allow large video uploads
+
+    # Logging
+    access_log /var/log/nginx/access.log;
+    error_log /var/log/nginx/error.log;
+
+    # Gzip compression
+    gzip on;
+    gzip_vary on;
+    gzip_min_length 1024;
+    gzip_proxied any;
+    gzip_comp_level 6;
+    gzip_types
+        text/plain
+        text/css
+        text/xml
+        text/javascript
+        application/json
+        application/javascript
+        application/xml+rss
+        application/atom+xml
+        image/svg+xml;
+
+    # Upstream for frontend load balancing
+    upstream frontend_servers {
+        # Round-robin load balancing between frontends
+        server 127.0.0.1:7860 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:7861 weight=1 max_fails=3 fail_timeout=30s;
+        
+        # Health check
+        keepalive 32;
+    }
+
+    # Upstream for backend1 load balancing (8 replicas)
+    upstream backend1_servers {
+        # Round-robin load balancing for backend1 replicas
+        server 127.0.0.1:8000 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8001 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8002 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8003 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8004 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8005 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8006 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8007 weight=1 max_fails=3 fail_timeout=30s;
+        
+        keepalive 32;
+    }
+
+    # Upstream for backend2 load balancing (8 replicas)
+    upstream backend2_servers {
+        # Round-robin load balancing for backend2 replicas
+        server 127.0.0.1:8010 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8011 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8012 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8013 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8014 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8015 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8016 weight=1 max_fails=3 fail_timeout=30s;
+        server 127.0.0.1:8017 weight=1 max_fails=3 fail_timeout=30s;
+        
+        keepalive 32;
+    }
+
+    # Main server block
+    server {
+        listen 80;
+        server_name localhost;
+
+        # Security headers
+        add_header X-Frame-Options "SAMEORIGIN" always;
+        add_header X-Content-Type-Options "nosniff" always;
+        add_header X-XSS-Protection "1; mode=block" always;
+        add_header Referrer-Policy "no-referrer-when-downgrade" always;
+
+        # Frontend routes (Gradio interfaces)
+        location / {
+            proxy_pass http://frontend_servers;
+            proxy_set_header Host $host;
+            proxy_set_header X-Real-IP $remote_addr;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header X-Forwarded-Proto $scheme;
+            
+            # WebSocket support for Gradio
+            proxy_http_version 1.1;
+            proxy_set_header Upgrade $http_upgrade;
+            proxy_set_header Connection "upgrade";
+            
+            # Timeouts
+            proxy_connect_timeout 60s;
+            proxy_send_timeout 60s;
+            proxy_read_timeout 60s;
+            
+            # Buffer settings
+            proxy_buffering on;
+            proxy_buffer_size 128k;
+            proxy_buffers 4 256k;
+            proxy_busy_buffers_size 256k;
+        }
+
+        # Backend API routes for frontend1
+        location /api/frontend1/ {
+            # Strip the /api/frontend1/ prefix
+            rewrite ^/api/frontend1/(.*) /$1 break;
+            proxy_pass http://backend1_servers;
+            proxy_set_header Host $host;
+            proxy_set_header X-Real-IP $remote_addr;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header X-Forwarded-Proto $scheme;
+            
+            # Timeouts for video generation
+            proxy_connect_timeout 300s;
+            proxy_send_timeout 300s;
+            proxy_read_timeout 300s;
+            
+            # Buffer settings for large responses
+            proxy_buffering on;
+            proxy_buffer_size 128k;
+            proxy_buffers 4 256k;
+            proxy_busy_buffers_size 256k;
+        }
+
+        # Backend API routes for frontend2
+        location /api/frontend2/ {
+            # Strip the /api/frontend2/ prefix
+            rewrite ^/api/frontend2/(.*) /$1 break;
+            proxy_pass http://backend2_servers;
+            proxy_set_header Host $host;
+            proxy_set_header X-Real-IP $remote_addr;
+            proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+            proxy_set_header X-Forwarded-Proto $scheme;
+            
+            # Timeouts for video generation
+            proxy_connect_timeout 300s;
+            proxy_send_timeout 300s;
+            proxy_read_timeout 300s;
+            
+            # Buffer settings for large responses
+            proxy_buffering on;
+            proxy_buffer_size 128k;
+            proxy_buffers 4 256k;
+            proxy_busy_buffers_size 256k;
+        }
+
+        # Health check endpoint
+        location /health {
+            access_log off;
+            return 200 "healthy\n";
+            add_header Content-Type text/plain;
+        }
+
+        # Static files (if needed)
+        location /static/ {
+            alias /var/www/static/;
+            expires 1y;
+            add_header Cache-Control "public, immutable";
+        }
+    }
+}

examples/inference/gradio/ray_serve_backend.py

@@ -5,45 +5,16 @@ 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 fastapi import FastAPI, Request
 from pydantic import BaseModel
+from PIL import Image
 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):
@@ -55,297 +26,284 @@ class VideoGenerationRequest(BaseModel):
     num_frames: int = 21
     height: int = 448
     width: int = 832
+    num_inference_steps: int = 20
     randomize_seed: bool = False
-    return_frames: bool = False
-    model_path: Optional[str] = None
+    return_frames: bool = False  # Whether to return base64 encoded frames
+    image_path: Optional[str] = None  # Path to input image for I2V
+    model_type: str = "t2v"  # "t2v" or "i2v" to specify which model to use
 
 
 class VideoGenerationResponse(BaseModel):
-    video_data: Optional[str] = None
+    output_path: str
     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
+    frames: Optional[List[str]] = None  # Base64 encoded frames
 
 
-def encode_video_to_base64(frames: List[np.ndarray], fps: int = DEFAULT_FPS) -> str:
+def encode_frames_to_base64(frames: List[np.ndarray]) -> List[str]:
+    """Convert numpy frames (0-255) to base64-encoded PNG images"""
     if not frames:
-        return ""
+        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
+    encoded_frames = []
     
-    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 = []
+    for i, frame in enumerate(frames):
+        try:
+            # Ensure frame is numpy array
+            if not isinstance(frame, np.ndarray):
+                print(f"Warning: Frame {i} is not a numpy array, skipping")
+                continue
+                
+            # Ensure frame is uint8
+            if frame.dtype != np.uint8:
+                # Clip values to 0-255 range and convert to uint8
+                frame = np.clip(frame, 0, 255).astype(np.uint8)
+            
+            # Convert numpy array to PIL Image
+            if len(frame.shape) == 3 and frame.shape[2] == 3:
+                # RGB image
+                pil_image = Image.fromarray(frame, mode='RGB')
+            elif len(frame.shape) == 3 and frame.shape[2] == 4:
+                # RGBA image
+                pil_image = Image.fromarray(frame, mode='RGBA')
+            elif len(frame.shape) == 2:
+                # Grayscale image
+                pil_image = Image.fromarray(frame, mode='L')
+            else:
+                print(f"Warning: Frame {i} has unsupported shape {frame.shape}, skipping")
+                continue
+            
+            # Save to bytes buffer as PNG
+            buffer = io.BytesIO()
+            pil_image.save(buffer, format='PNG')
+            buffer.seek(0)
+            
+            # Encode to base64
+            img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+            encoded_frames.append(f"data:image/png;base64,{img_base64}")
+            
+        except Exception as e:
+            print(f"Warning: Failed to encode frame {i}: {e}")
+            continue
     
-    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")
+    return encoded_frames
 
 
+# Create FastAPI app with rate limiting
 app = FastAPI()
+
+# Initialize rate limiter
 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.deployment(
+    num_replicas=8,
+    # ray_actor_options={"num_cpus": 10, "num_gpus": 1, "runtime_env": {"conda": "fv", "working_dir": "/mnt/fast-disks/nfs/hao_lab/FastVideo"}},
+    ray_actor_options={"num_cpus": 10, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
+)
 @serve.ingress(app)
 class FastVideoAPI:
-
-    def __init__(self, t2v_deployments: Dict[str, DeploymentHandle]):
-        self.t2v_deployments = t2v_deployments
+    def __init__(self, t2v_model_path: str, i2v_model_path: str, output_path: str):
+        self.t2v_model_path = t2v_model_path
+        self.i2v_model_path = i2v_model_path
+        self.output_path = output_path
         
-        # 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)
+        # Initialize the video generators
+        self.t2v_generator = None # Initialize to None
+        self.i2v_generator = None # Initialize to None
+        self.t2v_default_params = None # Initialize to None
+        self.i2v_default_params = None # Initialize to None
         
-        if response and hasattr(response, 'generation_time') and response.generation_time:
-            self.video_generation_time.labels(model_type=model_name).observe(response.generation_time)
+        # Ensure output directory exists
+        os.makedirs(output_path, exist_ok=True)
+        time.sleep(10)
+        self._initialize_models() # Ensure models are initialized
+    
+    def _initialize_models(self):
+        # Set VSA environment variable
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+        # os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
+        
+        # Import only when needed - use direct imports to avoid module-level execution
+        from fastvideo.entrypoints.video_generator import VideoGenerator
+        from fastvideo.configs.sample.base import SamplingParam
+        
+        # Initialize T2V model
+        if self.t2v_generator is None:
+            print(f"Initializing T2V model: {self.t2v_model_path}")
+            self.t2v_generator = VideoGenerator.from_pretrained(
+                model_path=self.t2v_model_path, 
+                num_gpus=1,
+                use_fsdp_inference=True,
+                # Adjust these offload parameters if you have < 32GB of VRAM
+                text_encoder_cpu_offload=False,
+                dit_cpu_offload=False,
+                vae_cpu_offload=False,
+                VSA_sparsity=0.8,
+                # master_port=port,
+            )
+            self.t2v_default_params = SamplingParam.from_pretrained(self.t2v_model_path)
+            print("✅ T2V model initialized successfully")
+        
+        # Initialize I2V model
+        # if self.i2v_generator is None:
+        if False:
+            print(f"Initializing I2V model: {self.i2v_model_path}")
+            self.i2v_generator = VideoGenerator.from_pretrained(
+                model_path=self.i2v_model_path, 
+                num_gpus=1,
+                use_fsdp_inference=True,
+                # Adjust these offload parameters if you have < 32GB of VRAM
+                text_encoder_cpu_offload=False,
+                dit_cpu_offload=False,
+                vae_cpu_offload=False,
+                VSA_sparsity=0.8,
+                # master_port=port,
+            )
+            self.i2v_default_params = SamplingParam.from_pretrained(self.i2v_model_path)
+            print("✅ I2V model initialized successfully")
     
     @app.post("/generate_video", response_model=VideoGenerationResponse)
-    @limiter.limit("10/minute")
+    @limiter.limit("50/minute")  # Allow 2 requests per minute per IP
     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")
+            # Select the appropriate model and parameters based on model_type
+            if video_request.model_type.lower() == "i2v":
+                generator = self.i2v_generator
+                params = deepcopy(self.i2v_default_params)
+                print(f"Using I2V model for generation")
+            else:
+                generator = self.t2v_generator
+                params = deepcopy(self.t2v_default_params)
+                print(f"Using T2V model for generation")
             
-            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
+            # Update parameters with request values
+            params.prompt = video_request.prompt
+            # Only override negative prompt if user explicitly opts in
+            if video_request.use_negative_prompt:
+                params.negative_prompt = video_request.negative_prompt
 
-        except Exception as e:
-            self._record_metrics(model_name, "error", time.time() - start_time)
+            params.seed = video_request.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.num_inference_steps = video_request.num_inference_steps
             
+            # Handle seed randomization
+            if video_request.randomize_seed:
+                params.seed = torch.randint(0, 1000000, (1,)).item()
+            
+            # Ensure negative_prompt is a non-None string; FastVideo validation disallows None
+            if params.negative_prompt is None:
+                params.negative_prompt = ""  # empty string satisfies validator
+            
+            # Set up output path and video saving
+            params.save_video = True
+            params.output_path = self.output_path
+            # params.return_frames = False  # avoid keeping frames in memory
+            
+            # Create a clean filename from the prompt
+            safe_prompt = video_request.prompt[:100].replace(' ', '_').replace('/', '_').replace('\\', '_')
+
+            # Store desired video name inside the SamplingParam to avoid unknown kwarg errors
+            setattr(params, "output_video_name", safe_prompt)
+            
+            # Handle image_path for I2V
+            if video_request.image_path:
+                params.image_path = video_request.image_path
+            
+            # Generate the video with proper output path and filename
+            result = generator.generate_video(
+                prompt=video_request.prompt,
+                sampling_param=params,
+                save_video=True,  # Match the params.save_video setting
+            )
+            
+            # The actual output path where the video was saved
+            output_path = os.path.join(self.output_path, f"{safe_prompt}.mp4")
+            
+            # Verify the file exists
+            if not os.path.exists(output_path):
+                raise FileNotFoundError(f"Video was not saved to expected location: {output_path}")
+            
+            frames = result.get("frames", [])
+            
+            # Encode frames to base64 for web transmission only if requested
+            encoded_frames = None
+            if video_request.return_frames and frames:
+                try:
+                    encoded_frames = encode_frames_to_base64(frames)
+                except Exception as e:
+                    print(f"Warning: Failed to encode frames: {e}")
+                    encoded_frames = None
+            
+            response = VideoGenerationResponse(
+                output_path=output_path,
+                frames=encoded_frames,
+                seed=params.seed,
+                success=True
+            )
+            
+            # Memory cleanup to avoid OOM in repeated generations
+            import gc
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+            return response
+        except Exception as e:
             return VideoGenerationResponse(
-                video_data=None,
+                output_path="",
                 seed=video_request.seed,
                 success=False,
-                error_message=str(e),
-                generation_time=0,
-                inference_time=0,
-                encoding_time=0,
-                total_time=0,
+                error_message=str(e)
             )
     
     @app.get("/health")
-    @limiter.limit("10/minute")
-    async def health_check(self, request: Request) -> Dict[str, str]:
+    @limiter.limit("10/minute")  # Allow 10 health checks per minute per IP
+    async def health_check(self, request: Request):
         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,
+    t2v_model_path: str = "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+    i2v_model_path: str = "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
     output_path: str = "outputs",
     host: str = "0.0.0.0",
-    port: int = 8000,
-) -> None:
+    port: int = 8000
+):
+    """Start the Ray Serve backend"""
+    # Initialize Ray
     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")
-
+    
+    # Deploy the API
+    api = FastVideoAPI.bind(t2v_model_path, i2v_model_path, output_path)
+    serve.run(api, route_prefix="/", name="fast_video")  # detach
+    
     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"T2V Model: {t2v_model_path}")
+    print(f"I2V Model: {i2v_model_path}")
     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",
+    parser.add_argument("--t2v_model_path",
                         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",
+                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
+                        help="Path to the T2V model")
+    parser.add_argument("--i2v_model_path",
                         type=str,
-                        default="4,4",
-                        help="Comma separated list of number of replicas for the T2V model(s)")
+                        default="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+                        help="Path to the I2V model")
     parser.add_argument("--output_path",
                         type=str,
                         default="outputs",
@@ -360,20 +318,20 @@ if __name__ == "__main__":
                         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,
+        t2v_model_path=args.t2v_model_path,
+        i2v_model_path=args.i2v_model_path,
         output_path=args.output_path,
         host=args.host,
         port=args.port,
     )
 
-    setup_signal_handlers()
+    # ---- keep the process alive ---------------------------------
+    import signal, sys, time
+    signal.signal(signal.SIGINT,  lambda *_: sys.exit(0))   # Ctrl-C
+    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))   # docker stop etc.
+
     print("✅ FastVideo backend is running. Press Ctrl-C to stop.")
     while True:
         time.sleep(3600) 

examples/inference/gradio/ray_serve_backend_scalable.py

@@ -0,0 +1,334 @@
+import time
+import os
+import torch
+import base64
+import io
+from copy import deepcopy
+from typing import Dict, Any, Optional, List
+
+import ray
+from ray import serve
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
+from PIL import Image
+import numpy as np
+from slowapi import Limiter, _rate_limit_exceeded_handler
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+
+
+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
+    num_inference_steps: int = 20
+    randomize_seed: bool = False
+    return_frames: bool = False  # Whether to return base64 encoded frames
+    image_path: Optional[str] = None  # Path to input image for I2V
+    model_type: str = "t2v"  # "t2v" or "i2v" to specify which model to use
+
+
+class VideoGenerationResponse(BaseModel):
+    output_path: str
+    seed: int
+    success: bool
+    error_message: Optional[str] = None
+    frames: Optional[List[str]] = None  # Base64 encoded frames
+
+
+def encode_frames_to_base64(frames: List[np.ndarray]) -> List[str]:
+    """Convert numpy frames (0-255) to base64-encoded PNG images"""
+    if not frames:
+        return []
+    
+    encoded_frames = []
+    
+    for i, frame in enumerate(frames):
+        try:
+            # Ensure frame is numpy array
+            if not isinstance(frame, np.ndarray):
+                print(f"Warning: Frame {i} is not a numpy array, skipping")
+                continue
+                
+            # Ensure frame is uint8
+            if frame.dtype != np.uint8:
+                # Clip values to 0-255 range and convert to uint8
+                frame = np.clip(frame, 0, 255).astype(np.uint8)
+            
+            # Convert numpy array to PIL Image
+            if len(frame.shape) == 3 and frame.shape[2] == 3:
+                # RGB image
+                pil_image = Image.fromarray(frame, mode='RGB')
+            elif len(frame.shape) == 3 and frame.shape[2] == 4:
+                # RGBA image
+                pil_image = Image.fromarray(frame, mode='RGBA')
+            elif len(frame.shape) == 2:
+                # Grayscale image
+                pil_image = Image.fromarray(frame, mode='L')
+            else:
+                print(f"Warning: Frame {i} has unsupported shape {frame.shape}, skipping")
+                continue
+            
+            # Save to bytes buffer as PNG
+            buffer = io.BytesIO()
+            pil_image.save(buffer, format='PNG')
+            buffer.seek(0)
+            
+            # Encode to base64
+            img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
+            encoded_frames.append(f"data:image/png;base64,{img_base64}")
+            
+        except Exception as e:
+            print(f"Warning: Failed to encode frame {i}: {e}")
+            continue
+    
+    return encoded_frames
+
+
+# Create FastAPI app with rate limiting
+app = FastAPI()
+
+# Initialize rate limiter
+limiter = Limiter(key_func=get_remote_address)
+app.state.limiter = limiter
+app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+
+
+@serve.deployment(
+    num_replicas=3,  # Set to 3 for cluster 0
+    ray_actor_options={
+        "num_cpus": 10, 
+        "num_gpus": 1, 
+        "runtime_env": {"conda": "fv"},
+    },
+)
+@serve.ingress(app)
+class FastVideoMultiGPUAPI:
+    def __init__(self, t2v_model_path: str, i2v_model_path: str, output_path: str, gpu_id: int = 0):
+        self.t2v_model_path = t2v_model_path
+        self.i2v_model_path = i2v_model_path
+        self.output_path = output_path
+        self.gpu_id = gpu_id
+        
+        # Initialize the video generators
+        self.t2v_generator = None
+        self.i2v_generator = None
+        self.t2v_default_params = None
+        self.i2v_default_params = None
+        
+        # Ensure output directory exists
+        os.makedirs(output_path, exist_ok=True)
+        time.sleep(10)
+        self._initialize_models()
+    
+    def _initialize_models(self):
+        # Set VSA environment variable
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "VIDEO_SPARSE_ATTN"
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
+        
+        # Import only when needed
+        from fastvideo.entrypoints.video_generator import VideoGenerator
+        from fastvideo.configs.sample.base import SamplingParam
+        
+        # Initialize T2V model
+        if False:  # Disabled for now
+            print(f"Initializing T2V model on GPU {self.gpu_id}: {self.t2v_model_path}")
+            self.t2v_generator = VideoGenerator.from_pretrained(
+                model_path=self.t2v_model_path, 
+                num_gpus=1,
+                use_fsdp_inference=True,
+                text_encoder_cpu_offload=False,
+                dit_cpu_offload=False,
+                vae_cpu_offload=False,
+                VSA_sparsity=0.8,
+            )
+            self.t2v_default_params = SamplingParam.from_pretrained(self.t2v_model_path)
+            print(f"✅ T2V model initialized successfully on GPU {self.gpu_id}")
+        
+        # Initialize I2V model
+        if self.i2v_generator is None:
+            print(f"Initializing I2V model on GPU {self.gpu_id}: {self.i2v_model_path}")
+            self.i2v_generator = VideoGenerator.from_pretrained(
+                model_path=self.i2v_model_path, 
+                num_gpus=1,
+                use_fsdp_inference=True,
+                text_encoder_cpu_offload=False,
+                dit_cpu_offload=False,
+                vae_cpu_offload=False,
+                VSA_sparsity=0.8,
+            )
+            self.i2v_default_params = SamplingParam.from_pretrained(self.i2v_model_path)
+            print(f"✅ I2V model initialized successfully on GPU {self.gpu_id}")
+    
+    @app.post("/generate_video", response_model=VideoGenerationResponse)
+    @limiter.limit("2/minute")  # Allow 2 requests per minute per IP
+    async def generate_video(self, request: Request, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
+        try:
+            # Select the appropriate model and parameters based on model_type
+            if video_request.model_type.lower() == "i2v":
+                generator = self.i2v_generator
+                params = deepcopy(self.i2v_default_params)
+                print(f"Using I2V model for generation on GPU {self.gpu_id}")
+            else:
+                generator = self.t2v_generator
+                params = deepcopy(self.t2v_default_params)
+                print(f"Using T2V model for generation on GPU {self.gpu_id}")
+            
+            # Update parameters with request values
+            params.prompt = video_request.prompt
+            
+            # Handle seed randomization
+            if video_request.randomize_seed:
+                params.seed = torch.randint(0, 1000000, (1,)).item()
+            
+            # Ensure negative_prompt is a non-None string
+            if params.negative_prompt is None:
+                params.negative_prompt = ""
+            
+            # Set up output path and video saving
+            params.save_video = True
+            params.output_path = self.output_path
+            
+            # Create a clean filename from the prompt
+            safe_prompt = video_request.prompt[:100].replace(' ', '_').replace('/', '_').replace('\\', '_')
+            setattr(params, "output_video_name", safe_prompt)
+            
+            # Handle image_path for I2V
+            if video_request.image_path:
+                params.image_path = video_request.image_path
+            
+            # Generate the video
+            result = generator.generate_video(
+                prompt=video_request.prompt,
+                sampling_param=params,
+                save_video=True,
+            )
+            
+            frames = result.get("frames", [])
+            
+            # Encode frames to base64 for web transmission only if requested
+            encoded_frames = None
+            if video_request.return_frames and frames:
+                try:
+                    encoded_frames = encode_frames_to_base64(frames)
+                except Exception as e:
+                    print(f"Warning: Failed to encode frames: {e}")
+                    encoded_frames = None
+            
+            response = VideoGenerationResponse(
+                output_path="",
+                frames=encoded_frames,
+                seed=params.seed,
+                success=True
+            )
+            
+            # Memory cleanup to avoid OOM in repeated generations
+            import gc
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+            return response
+        except Exception as e:
+            return VideoGenerationResponse(
+                output_path="",
+                seed=video_request.seed,
+                success=False,
+                error_message=str(e)
+            )
+    
+    @app.get("/health")
+    @limiter.limit("10/minute")  # Allow 10 health checks per minute per IP
+    async def health_check(self, request: Request):
+        return {"status": "healthy", "gpu_id": self.gpu_id}
+
+
+def start_ray_serve_multi_gpu(
+    t2v_model_path: str = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
+    i2v_model_path: str = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
+    output_path: str = "outputs",
+    host: str = "0.0.0.0",
+    port: int = 8000,
+    num_gpus: int = 8,
+    cluster_id: int = 0
+):
+    """Start the Ray Serve backend with multiple GPU replicas"""
+    # Initialize Ray
+    if not ray.is_initialized():
+        ray.init()
+    
+    # Use unique application name based on cluster_id
+    app_name = f"fast_video_cluster_{cluster_id}"
+    
+    # Deploy the API
+    api = FastVideoMultiGPUAPI.bind(t2v_model_path, i2v_model_path, output_path)
+    serve.run(api, route_prefix=f"/cluster_{cluster_id}", name=app_name)
+    
+    print(f"Ray Serve multi-GPU backend started at http://{host}:{port}")
+    print(f"T2V Model: {t2v_model_path}")
+    print(f"I2V Model: {i2v_model_path}")
+    print(f"Number of GPU replicas: {num_gpus}")
+    print(f"Cluster ID: {cluster_id}")
+    print(f"Application name: {app_name}")
+    print(f"Health check: http://{host}:{port}/cluster_{cluster_id}/health")
+    print(f"Video generation endpoint: http://{host}:{port}/cluster_{cluster_id}/generate_video")
+
+
+if __name__ == "__main__":
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="FastVideo Ray Serve Multi-GPU Backend")
+    parser.add_argument("--t2v_model_path",
+                        type=str,
+                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
+                        help="Path to the T2V model")
+    parser.add_argument("--i2v_model_path",
+                        type=str,
+                        default="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
+                        help="Path to the I2V model")
+    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")
+    parser.add_argument("--num_gpus",
+                        type=int,
+                        default=8,
+                        help="Number of GPU replicas")
+    parser.add_argument("--cluster_id",
+                        type=int,
+                        default=0,
+                        help="Cluster ID for unique naming")
+    
+    args = parser.parse_args()
+    
+    start_ray_serve_multi_gpu(
+        t2v_model_path=args.t2v_model_path,
+        i2v_model_path=args.i2v_model_path,
+        output_path=args.output_path,
+        host=args.host,
+        port=args.port,
+        num_gpus=args.num_gpus,
+        cluster_id=args.cluster_id,
+    )
+
+    # Keep the process alive
+    import signal, sys, time
+    signal.signal(signal.SIGINT,  lambda *_: sys.exit(0))
+    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
+
+    print("✅ FastVideo multi-GPU backend is running. Press Ctrl-C to stop.")
+    while True:
+        time.sleep(3600)

examples/inference/gradio/start.sh

@@ -1,3 +0,0 @@
-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

@@ -12,245 +12,220 @@ import threading
 import signal
 import requests
 from pathlib import Path
-from typing import Dict, Any, Optional
+
+# Add the project root to the Python path
+project_root = Path(__file__).parent.parent.parent.parent
+sys.path.insert(0, str(project_root))
 
 
-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 check_backend_health(backend_url: str, max_retries: int = 100) -> bool:
+    """Check if the backend is healthy"""
+    health_url = f"{backend_url}/health"
     
-    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
-        
+    for i in range(max_retries):
         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:
+            response = requests.get(health_url, timeout=5)
+            if response.status_code == 200:
+                print(f"✅ Backend is healthy at {backend_url}")
+                return True
+        except requests.exceptions.RequestException:
             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)
+        if i < max_retries - 1:
+            print(f"⏳ Waiting for backend to start... ({i+1}/{max_retries})")
+            time.sleep(2)
     
-    signal.signal(signal.SIGINT, signal_handler)
-    signal.signal(signal.SIGTERM, signal_handler)
+    print(f"❌ Backend failed to start within {max_retries * 2} seconds")
+    return False
 
 
-def print_startup_info(args: argparse.Namespace) -> None:
+def start_backend(args):
+    """Start the Ray Serve backend"""
+    backend_script = Path(__file__).parent / "ray_serve_backend.py"
+    
+    cmd = [
+        sys.executable, str(backend_script),
+        "--t2v_model_path", args.t2v_model_path,
+        "--i2v_model_path", args.i2v_model_path,
+        "--output_path", args.output_path,
+        "--host", args.backend_host,
+        "--port", str(args.backend_port)
+    ]
+    
+    print(f"🚀 Starting Ray Serve backend...")
+    print(f"Command: {' '.join(cmd)}")
+    
+    # Start the backend process
+    backend_process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        universal_newlines=True,
+        bufsize=1
+    )
+    
+    # Monitor backend output
+    def monitor_backend():
+        for line in backend_process.stdout:
+            print(f"[BACKEND] {line.rstrip()}")
+    
+    monitor_thread = threading.Thread(target=monitor_backend, daemon=True)
+    monitor_thread.start()
+    
+    return backend_process
+
+
+def start_frontend(args):
+    """Start the Gradio frontend"""
+    frontend_script = Path(__file__).parent / "gradio_frontend.py"
+    backend_url = f"http://{args.backend_host}:{args.backend_port}"
+    
+    cmd = [
+        sys.executable, str(frontend_script),
+        "--backend_url", backend_url,
+        "--t2v_model_path", args.t2v_model_path,
+        "--i2v_model_path", args.i2v_model_path,
+        "--host", args.frontend_host,
+        "--port", str(args.frontend_port)
+    ]
+    
+    print(f"🎨 Starting Gradio frontend...")
+    print(f"Command: {' '.join(cmd)}")
+    
+    # Start the frontend process
+    frontend_process = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        universal_newlines=True,
+        bufsize=1
+    )
+    
+    # Monitor frontend output
+    def monitor_frontend():
+        for line in frontend_process.stdout:
+            print(f"[FRONTEND] {line.rstrip()}")
+    
+    monitor_thread = threading.Thread(target=monitor_frontend, daemon=True)
+    monitor_thread.start()
+    
+    return frontend_process
+
+
+def main():
+    parser = argparse.ArgumentParser(description="FastVideo Ray Serve App")
+    
+    # Model and output settings
+    parser.add_argument("--t2v_model_path",
+                        type=str,
+                        default="Wan-AI/Wan2.2-TI2V-5BDiffusers",
+                        help="Path to the T2V model")
+    parser.add_argument("--i2v_model_path",
+                        type=str,
+                        default="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+                        help="Path to the I2V model")
+    parser.add_argument("--output_path",
+                        type=str,
+                        default="outputs",
+                        help="Path to save generated videos")
+    
+    # Backend settings
+    parser.add_argument("--backend_host",
+                        type=str,
+                        default="0.0.0.0",
+                        help="Backend host to bind to")
+    parser.add_argument("--backend_port",
+                        type=int,
+                        default=8000,
+                        help="Backend port to bind to")
+    
+    # Frontend settings
+    parser.add_argument("--frontend_host",
+                        type=str,
+                        default="0.0.0.0",
+                        help="Frontend host to bind to")
+    parser.add_argument("--frontend_port",
+                        type=int,
+                        default=7860,
+                        help="Frontend port to bind to")
+    
+    # Other settings
+    parser.add_argument("--skip_backend_check",
+                        action="store_true",
+                        help="Skip backend health check")
+    
+    args = parser.parse_args()
+    
+    # Ensure output directory exists
+    os.makedirs(args.output_path, exist_ok=True)
+    
     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"T2V Model: {args.t2v_model_path}")
+    print(f"I2V Model: {args.i2v_model_path}")
     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")
+    # Start backend
+    backend_process = start_backend(args)
     
-    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")
+    # Wait for backend to be ready
+    backend_url = f"http://{args.backend_host}:{args.backend_port}"
     
-    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")
+    if not args.skip_backend_check:
+        if not check_backend_health(backend_url):
+            print("❌ Backend failed to start. Terminating...")
+            backend_process.terminate()
+            sys.exit(1)
     
-    parser.add_argument("--skip_backend_check", 
-                        action="store_true",
-                        help="Skip backend health check")
+    # Start frontend
+    frontend_process = start_frontend(args)
     
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_arguments()
+    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: {backend_url}")
+    print("\nPress Ctrl+C to stop both services...")
+    # return
     
-    os.makedirs(args.output_path, exist_ok=True)
-    print_startup_info(args)
+    # Signal handler for graceful shutdown
+    def signal_handler(signum, frame):
+        print("\n🛑 Shutting down services...")
+        frontend_process.terminate()
+        backend_process.terminate()
+        
+        # Wait for processes to terminate
+        try:
+            frontend_process.wait(timeout=5)
+            backend_process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            print("⚠️  Force killing processes...")
+            frontend_process.kill()
+            backend_process.kill()
+        
+        print("✅ Services stopped")
+        sys.exit(0)
     
-    service_manager = ServiceManager(args)
-    setup_signal_handlers(service_manager)
+    signal.signal(signal.SIGINT, signal_handler)
+    signal.signal(signal.SIGTERM, signal_handler)
     
+    # Monitor processes
     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()
+        while True:
+            # Check if processes are still running
+            if frontend_process.poll() is not None:
+                print("❌ Frontend process died unexpectedly")
+                break
+            
+            if backend_process.poll() is not None:
+                print("❌ Backend process died unexpectedly")
+                break
+            
+            time.sleep(1)
     
-    except Exception as e:
-        print(f"❌ Unexpected error: {e}")
-        service_manager.shutdown_services()
-        sys.exit(1)
+    except KeyboardInterrupt:
+        signal_handler(signal.SIGINT, None)
 
 
 if __name__ == "__main__":