fixi2v demo bug

add full i2v demo
add demo prompts
2025-11-19 20:53:43 +00:00 · 2025-11-19 15:40:04 +00:00 · 2025-11-19 05:49:01 +00:00 · 2025-11-19 05:42:37 +00:00 · 2025-11-18 09:40:05 +00:00 · 2025-11-16 23:10:47 +00:00
262 changed files with 14695 additions and 2293 deletions
@@ -104,6 +104,18 @@ steps:
                  - TEST_TYPE=distillation_dmd
                agents:
                  queue: "default"
+            - path:
+                - "fastvideo/training/*self_forcing_distillation_pipeline.py"
+                - "fastvideo/tests/training/self-forcing/**"
+                - "pyproject.toml"
+                - "docker/Dockerfile.python3.12"
+              config:
+                command: "timeout 15m .buildkite/scripts/pr_test.sh"
+                label: "Self-Forcing Tests"
+                env:
+                  - TEST_TYPE=self_forcing
+                agents:
+                  queue: "default"
            - path:
                - "fastvideo/**"
                - "pyproject.toml"
@@ -31,9 +31,9 @@ log "Setting up Modal authentication from Buildkite secrets..."
 MODAL_TOKEN_ID=$(buildkite-agent secret get modal_token_id)
 MODAL_TOKEN_SECRET=$(buildkite-agent secret get modal_token_secret)

+# Retrieve other secrets
 WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
-
-WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
+HF_API_KEY=$(buildkite-agent secret get hf_api_key)

 if [ -n "$MODAL_TOKEN_ID" ] && [ -n "$MODAL_TOKEN_SECRET" ]; then
    log "Retrieved Modal credentials from Buildkite secrets"
@@ -63,15 +63,15 @@ MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUI
 case "$TEST_TYPE" in
    "encoder")
        log "Running encoder tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
        ;;
    "vae")
        log "Running VAE tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
        ;;
    "transformer")
        log "Running transformer tests..."
-        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
+        MODAL_COMMAND="$MODAL_ENV HF_API_KEY=$HF_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
        ;;
    "ssim")
        log "Running SSIM tests..."
@@ -110,6 +110,10 @@ case "$TEST_TYPE" in
        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_distill_dmd_tests"
        ;;
        # run_inference_tests_vmoba
+    "self_forcing")
+        log "Running self-forcing tests..."
+        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_self_forcing_tests"
+        ;;
    "inference_vmoba")
        log "Running V-MoBA inference tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_vmoba"
@@ -1,82 +1,55 @@
-# Sample workflow for building and deploying a Hugo site to GitHub Pages
-name: Deploy FastVideo Docs to Pages
+name: Deploy Documentation

 on:
-  # Runs on pushes targeting the default branch
  push:
-    branches:
-      - main
-    paths:
-      - "docs/**/*.md"
-      - "fastvideo/examples/**/*.py"
+    branches: [ main ]
  pull_request:
-    branches: 
-      - main
-    types: [opened, ready_for_review, synchronize, reopened]
-    paths:
-      - "docs/**/*.md"
-      - "fastvideo/examples/**/*.py"
+    branches: [ main ]

-  # Allows you to run this workflow manually from the Actions tab
-  workflow_dispatch:
-
-# Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages
 permissions:
  contents: read
  pages: write
  id-token: write

-# Allow only one concurrent deployment, skipping runs queued between the run in-progress and latest queued.
-# However, do NOT cancel in-progress runs as we want to allow these production deployments to complete.
 concurrency:
  group: "pages"
  cancel-in-progress: false

-# Default to bash
-defaults:
-  run:
-    shell: bash
-
 jobs:
-  pre-commit:
-    uses: ./.github/workflows/pre-commit.yml
-
-  # Build job
  build:
    runs-on: ubuntu-latest
-    needs: pre-commit
    steps:
      - name: Checkout
        uses: actions/checkout@v4
-      - name: Setup Pages
-        id: pages
-        uses: actions/configure-pages@v5
-      - name: Set up Python
+
+      - name: Setup Python
        uses: actions/setup-python@v5
        with:
-          python-version: "3.10"
+          python-version: '3.12'
+
      - name: Install dependencies
        run: |
-          cd docs
-          pip install -r requirements-docs.txt
-      - name: Build docs
-        run: |
-          cd docs
-          make clean
-          make html
+          python -m pip install --upgrade pip
+          pip install -r requirements-mkdocs.txt
+
+      - name: Setup Pages
+        uses: actions/configure-pages@v4
+
+      - name: Build documentation
+        run: mkdocs build
+
      - name: Upload artifact
        uses: actions/upload-pages-artifact@v3
        with:
-          path: ./docs/build/html
+          path: ./site

-  # Deployment job
  deploy:
    environment:
      name: github-pages
      url: ${{ steps.deployment.outputs.page_url }}
-    if: ${{ github.event_name == 'push' }}
    runs-on: ubuntu-latest
    needs: build
+    if: github.ref == 'refs/heads/main'
    steps:
      - name: Deploy to GitHub Pages
        id: deployment
@@ -352,7 +352,7 @@ jobs:
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/dataset/ -vs && pytest ./fastvideo/workflow/ -vs"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/dataset/ -vs && pytest ./fastvideo/workflow/ -vs && pytest ./fastvideo/entrypoints/ -vs"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -14,6 +14,8 @@ wandb/
 *.pt
 cache_dir/
 wandb/
+venv/
+.venv/
 runs/
 samples/
 *validation/
@@ -37,12 +39,13 @@ dist/
 eggs/
 .eggs/

-# Sphinx documentation
-docs/_build/
-docs/source/getting_started/examples/
-docs/source/inference/examples/
-docs/source/training/examples/
-docs/source/distillation/examples/
+# MkDocs documentation
+site/
+docs/getting_started/examples/
+docs/inference/examples/
+docs/training/examples/
+docs/distillation/examples/
+!requirements-mkdocs.txt

 # VSCode
 .vscode/
@@ -61,7 +64,7 @@ docs/source/distillation/examples/
 !fastvideo/tests/ssim/reference_videos/**/*.mp4

 # Static images
-!docs/source/_static/images/**/*.png
+!docs/assets/images/**/*.png
 !comfyui/assets/**/*.png
 !comfyui/assets/**/*.gif

@@ -10,11 +10,9 @@ exclude: |
    demo/.*|
    predict\.py|
    scripts/.*|
+    prompts/.*|
    fastvideo/data_preprocess/.*|
    fastvideo/dataset/.*|
-    fastvideo/distill/.*|
-    fastvideo/distill\.py|
-    fastvideo/distill_adv\.py|
    fastvideo/models/.*|
    fastvideo/sample/.*|
    fastvideo/train\.py|
@@ -44,10 +42,10 @@ repos:
  - id: codespell
    additional_dependencies: ['tomli']
    args: ['--toml', 'pyproject.toml']
- repo: https://github.com/PyCQA/isort
-  rev: 6.0.1
-  hooks:
-  - id: isort
+# - repo: https://github.com/PyCQA/isort
+#   rev: 6.0.1
+#   hooks:
+#   - id: isort
 - repo: https://github.com/jackdewinter/pymarkdown
  rev: v0.9.30
  hooks:
@@ -7,7 +7,7 @@
 FastVideo features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.

 <p align="center">
-    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/collections/FastVideo/fastwan-6886a305d9799c8cd1496408"  target="_blank"><b>FastWan</b></a>  | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-3csdw1isz-Euq8_Q8~baewG8hxjXs2gQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/S7HLCSTh" target="_blank"> <b> WeChat </b> </a> |
+    | 🕹️ <a href="https://fastwan.fastvideo.org/"<b>Online Demo</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/"><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-3f4lao1uq-u~Ipx6Lt4J27AlD2y~IdLQ" target="_blank"> <b>Slack</b> </a> |  🟣💬 <a href="https://ibb.co/XcY0Cpv" target="_blank"> <b> WeChat </b> </a> |
 </p>

 <div align="center">
@@ -49,10 +49,10 @@ conda activate fastvideo
 pip install fastvideo
 ```

-Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) for more detailed installation instructions.
+Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) 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/).
+For our sparse distillation techniques, please see our [distillation docs](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/) and check out our [blog](https://hao-ai-lab.github.io/blogs/fastvideo_post_training/).

 See below for recipes and datasets:

@@ -64,7 +64,7 @@ See below for recipes and datasets:

 ## 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. Make sure VSA kernels are [installed](https://hao-ai-lab.github.io/FastVideo/video_sparse_attention/installation/). Create a file called `example.py` with the following code:

 ```python
 import os
@@ -100,15 +100,15 @@ Run the script with:
 python example.py
 ```

-For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html).
+For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start/).

 ### Other docs:

- [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview.html)
- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
+- [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview/)
+- [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/)

 ## Distillation and Finetuning
- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd.html)
+- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/distillation/dmd/)
 <!-- - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html) -->

 ## 📑 Development Plan
@@ -127,7 +127,7 @@ See details in [development roadmap](https://github.com/hao-ai-lab/FastVideo/iss

 ## 🤝 Contributing

-We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview.html)
+We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview/)

 ## Acknowledgement
 We learned and reused code from the following projects:
@@ -155,8 +155,8 @@ If you find FastVideo useful, please considering citing our work:
 }

@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},
+  title={Vsa: Faster video diffusion with trainable sparse attention},
+  author={Zhang, Peiyuan and Chen, Yongqi and Huang, Haofeng and Lin, Will and Liu, Zhengzhong and Stoica, Ion and Xing, Eric and Zhang, Hao},
  journal={arXiv preprint arXiv:2505.13389},
  year={2025}
 }
@@ -0,0 +1,18 @@
+<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>
@@ -0,0 +1,6 @@
+<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>
@@ -72,6 +72,15 @@ 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

+
+## STA Configuration Logic
+Here is a diagram of how the window is configured and passed through the FastVideo pipeline:
+
+<div align="center">
+<img src="../../../docs/assets/images/STA_configuration.png" width="80%"/>
+</div>
+
+
 ## 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. 

@@ -247,11 +247,11 @@ def _attn_bwd_dq(dq, q, K, V,  #

    kv_blocks = tl.load(q2k_num  + meta_base)                 # int32
    kv_ptr    = q2k_index + meta_base * max_kv_blks           # ptr to list
+    block_size = tl.load(variable_block_sizes + q_blk)
    
    
    for blk_idx in range(kv_blocks*2):
        block_sparse_offset = (tl.load(kv_ptr + blk_idx//2).to(tl.int32)*2 + blk_idx%2) *step_n * stride_tok
-        block_size = tl.load(variable_block_sizes + blk_idx//2) - (blk_idx%2) * step_n
        kT = tl.load(kT_ptrs + block_sparse_offset)
        vT = tl.load(vT_ptrs + block_sparse_offset)
        qk = tl.dot(q, kT)
@@ -20,5 +20,7 @@ setup(
        "License :: OSI Approved :: Apache Software License",
    ],
    python_requires='>=3.12',
-    install_requires=[]
+    install_requires=[
+        "flash-attn >= 2.7.1",
+    ]
 )
@@ -6,8 +6,16 @@ import time
 import os
 import torch
 from typing import Tuple
-from flash_attn import flash_attn_varlen_func  # Use the new flash attention function
-from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
+try:
+    from flash_attn import flash_attn_varlen_func  # Use the new flash attention function
+    from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
+except ImportError:
+    def _unsupported(*args, **kwargs):
+        raise ImportError("flash-attn is not installed. Please install it, e.g., `pip install flash-attn`.")
+    _flash_attn_varlen_forward = _unsupported
+    _flash_attn_varlen_backward = _unsupported
+    flash_attn_varlen_func = _unsupported
+    
 from functools import lru_cache
 from einops import rearrange

@@ -1,26 +0,0 @@
-# Minimal makefile for Sphinx documentation
-#
-
-# You can set these variables from the command line, and also
-# from the environment for the first two.
-SPHINXOPTS    ?=
-SPHINXBUILD   ?= sphinx-build
-SOURCEDIR     = source
-BUILDDIR      = build
-
-# Put it first so that "make" without argument is like "make help".
-help:
-	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
-
-.PHONY: help Makefile
-
-# Catch-all target: route all unknown targets to Sphinx using the new
-# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
-%: Makefile
-	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
-
-clean:
-	@$(SPHINXBUILD) -M clean "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
-	rm -rf "$(SOURCEDIR)/getting_started/examples"
-	rm -rf "$(SOURCEDIR)/inference/examples"
-	rm -rf "$(SOURCEDIR)/training/examples"
@@ -1,20 +1,39 @@
-# FastVideo documents
+# FastVideo Documentation

-## Build the docs
+This directory contains the FastVideo documentation built with MkDocs.
+
+## Build the docs locally

 ```bash
-# Install dependencies.
-pip install -r requirements-docs.txt
+# Install dependencies
+pip install -r docs/requirements-mkdocs.txt

-# Build the docs.
-make clean
-make html
+# Serve docs with live reload (recommended for development)
+mkdocs serve
+
+# Or build static site
+mkdocs build
 ```

-## Open the docs with your browser
+## View the docs
+
+### Development server (with live reload)

 ```bash
-python -m http.server -d build/html/
+mkdocs serve
 ```

-Launch your browser and open localhost:8000.
+Then open your browser to: http://127.0.0.1:8000
+
+### Static build
+
+```bash
+mkdocs build
+python -m http.server -d site/
+```
+
+Then open your browser to: http://localhost:8000
+
+## Automatic Deployment
+
+Documentation is automatically built and deployed to GitHub Pages when changes are pushed to the `main` branch via the `.github/workflows/docs.yml` workflow.
@@ -0,0 +1,248 @@
+# FastVideo API Reference
+
+This page contains the complete API reference for the FastVideo library.
+
+## fastvideo
+
+### Modules
+
+| Name | Description |
+|------|-------------|
+| [attention](#fastvideoattention) | Attention mechanisms and backends for video generation |
+| [configs](#fastvideoconfigs) | Configuration classes for pipelines, models, and sampling |
+| [distributed](#fastvideodistributed) | Distributed execution and communication utilities |
+| [entrypoints](#fastvideoentrypoints) | Main API entry points for video generation |
+| [models](#fastvideomodels) | Model implementations (transformers, VAEs, schedulers) |
+| [pipelines](#fastvideopipelines) | Core pipeline classes for video diffusion |
+| [training](#fastvideotraining) | Training utilities and helpers |
+| [workflow](#fastvideoworkflow) | Workflow management and orchestration |
+| [dataset](#fastvideodataset) | Dataset handling and preprocessing |
+| [layers](#fastvideolayers) | Custom neural network layers |
+| [platforms](#fastvideoplatforms) | Platform-specific implementations |
+| [utils](#fastvideoutils) | Utility functions and helpers |
+| [worker](#fastvideoworker) | Execution workers for video generation |
+
+## fastvideo.attention
+
+::: fastvideo.attention
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.configs
+
+::: fastvideo.configs
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+### Submodules
+
+#### fastvideo.configs.pipelines
+
+::: fastvideo.configs.pipelines
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.configs.models
+
+::: fastvideo.configs.models
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.configs.sample
+
+::: fastvideo.configs.sample
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+## fastvideo.distributed
+
+::: fastvideo.distributed
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.entrypoints
+
+::: fastvideo.entrypoints
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.models
+
+::: fastvideo.models
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+### Submodules
+
+#### fastvideo.models.registry
+
+::: fastvideo.models.registry
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.models.loader
+
+::: fastvideo.models.loader
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+## fastvideo.pipelines
+
+::: fastvideo.pipelines
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+### Submodules
+
+#### fastvideo.pipelines.composed_pipeline_base
+
+::: fastvideo.pipelines.composed_pipeline_base
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.pipelines.lora_pipeline
+
+::: fastvideo.pipelines.lora_pipeline
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.pipelines.pipeline_batch_info
+
+::: fastvideo.pipelines.pipeline_batch_info
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.pipelines.pipeline_registry
+
+::: fastvideo.pipelines.pipeline_registry
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+#### fastvideo.pipelines.stages
+
+::: fastvideo.pipelines.stages
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 4
+
+## fastvideo.training
+
+::: fastvideo.training
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.workflow
+
+::: fastvideo.workflow
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.dataset
+
+::: fastvideo.dataset
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.layers
+
+::: fastvideo.layers
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.platforms
+
+::: fastvideo.platforms
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
+
+## fastvideo.utils
+
+::: fastvideo.utils
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      heading_level: 3
+
+## fastvideo.worker
+
+::: fastvideo.worker
+    options:
+      show_source: true
+      show_root_heading: true
+      show_root_toc_entry: true
+      show_submodules: true
+      heading_level: 3
@@ -0,0 +1,27 @@
+# API Summary
+
+This page provides a quick overview of the main FastVideo API components.
+
+## Video Generator
+
+::: fastvideo.VideoGenerator
+    options:
+      show_root_heading: false
+      show_source: false
+      heading_level: 3
+
+## Initialization Configuration
+
+::: fastvideo.PipelineConfig
+    options:
+      show_root_heading: false
+      show_source: false
+      heading_level: 3
+
+## Sampling Configuration
+
+::: fastvideo.SamplingParam
+    options:
+      show_root_heading: false
+      show_source: false
+      heading_level: 3
@@ -0,0 +1,41 @@
+.vertical-table-header th.head:not(.stub) {
+    writing-mode: sideways-lr;
+    white-space: nowrap;
+    max-width: 0;
+    p {
+       margin: 0;
+    }
+}
+
+/* Image sizing classes */
+.image-small {
+    max-width: 200px;
+    height: auto;
+}
+
+.image-medium {
+    max-width: 400px;
+    height: auto;
+}
+
+.image-large {
+    max-width: 600px;
+    height: auto;
+}
+
+.image-full {
+    max-width: 100%;
+    height: auto;
+}
+
+/* Responsive images */
+img {
+    max-width: 100%;
+    height: auto;
+}
+
+/* Center images */
+.image-center {
+    display: block;
+    margin: 0 auto;
+}
@@ -0,0 +1,6 @@
+<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>
@@ -0,0 +1,18 @@
+<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>
@@ -1,4 +1,4 @@
-(docker)=
+
 # 🐳 Using the FastVideo Docker Image

 If you prefer a containerized development environment or want to avoid managing dependencies manually, you can use our prebuilt Docker image:
@@ -3,11 +3,3 @@
 # 🧰 Developer Environment

 Accelerate your FastVideo development workflow by leveraging Docker images and cloud GPUs for efficient experimentation and reproducible environments.
-
-:::{toctree}
-:caption: Contents
-:maxdepth: 1
-
-docker
-runpod
-:::
@@ -1,4 +1,3 @@
-(runpod)=

 # 📦 Developing FastVideo on RunPod

@@ -10,7 +9,7 @@ Choose a GPU that supports CUDA 12.8

 Pick 1 or 2 L40S GPU(s)

-![RunPod CUDA selection](../../_static/images/runpod_cuda.png)
+![RunPod CUDA selection](../../assets/images/runpod_cuda.png)

 When creating your pod template, use this image:

@@ -24,11 +23,11 @@ Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.
 bash -c "apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-server -y;mkdir -p ~/.ssh;cd $_;chmod 700 ~/.ssh;echo \"$PUBLIC_KEY\" >> authorized_keys;chmod 700 authorized_keys;service ssh start;sleep infinity"
 ```

-![RunPod template configuration](../../_static/images/runpod_template.png)
+![RunPod template configuration](../../assets/images/runpod_template.png)

 After deploying, the pod will take a few minutes to pull the image and start the SSH service.

-![RunPod ssh](../../_static/images/runpod_ssh.png)
+![RunPod ssh](../../assets/images/runpod_ssh.png)

 ## Working with the pod

@@ -1,4 +1,3 @@
-(developer-overview)=

 # 🛠️ Contributing to FastVideo

@@ -0,0 +1,53 @@
+# Profiling FastVideo
+
+!!! warning
+    Profiling is only intended for FastVideo developers and maintainers to understand the proportion of time spent in different parts of the codebase. **FastVideo end-users should never turn on profiling** as it will significantly slow down the inference.
+
+## Profiling with PyTorch
+
+FastVideo exposes a process-wide torch profiler that you can enable via environment variables. Set `FASTVIDEO_TORCH_PROFILER_DIR` to an absolute directory path to start collecting traces, and specify the regions you want recorded with `FASTVIDEO_TORCH_PROFILE_REGIONS`:
+
+```bash
+FASTVIDEO_TORCH_PROFILER_DIR=/mnt/traces/fastvideo \
+FASTVIDEO_TORCH_PROFILE_REGIONS="profiler_region_model_loading,profiler_region_training_step"
+```
+
+All profiled regions must be registered in `fastvideo.profiler`; the current list includes:
+
+- `profiler_region_model_loading` — pipeline/module loading
+- `profiler_region_inference_pre_denoising`
+- `profiler_region_inference_denoising`
+- `profiler_region_inference_post_denoising`
+- `profiler_region_training_checkpoint_saving`
+- `profiler_region_training_dit`
+- `profiler_region_training_validation`
+- `profiler_region_training_epoch`
+- `profiler_region_training_step`
+- `profiler_region_training_backward`
+- `profiler_region_training_optimizer`
+- `profiler_region_distillation_teacher_forward`
+- `profiler_region_distillation_student_forward`
+- `profiler_region_distillation_loss`
+- `profiler_region_distillation_update`
+
+While profiling is enabled, FastVideo records additional annotations:
+
+- `fastvideo.region::<name>` spans are emitted when entering a region.
+- `fastvideo.profiler.enable_collection` / `fastvideo.profiler.disable_collection` events mark when torch profiler collection is toggled on or off.
+
+Only one profiler instance is created per process; subsequent pipelines reuse the same controller. If you set `FASTVIDEO_TORCH_PROFILE_REGIONS` incorrectly (e.g. misspelled name), FastVideo logs a warning and ignores that entry.
+
+Additional knobs:
+
+- `FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES`
+- `FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY`
+- `FASTVIDEO_TORCH_PROFILER_WITH_STACK`
+- `FASTVIDEO_TORCH_PROFILER_WITH_FLOPS`
+
+Traces can be visualized using <https://ui.perfetto.dev/>.
+
+### Best Practices
+
+- Keep the profiled step count small; traces can be large and slow down job shutdown while the profiler flushes data.
+- After profiling, clean up trace directories to avoid filling disks.
+- When adding new regions, register them in `fastvideo.profiler` and wrap the corresponding code block with `with self.profiler_controller.region("your_region"):` or the `@profile_region` decorator.
@@ -29,7 +29,6 @@ FastVideo separates model components from execution logic with these principles:
 - **Custom Attention Backends**: Components can support and use different Attention implementations
 - **Pipeline Abstraction**: Consistent interface across diffusion models

-(design-fastvideo-args)=
 ## FastVideoArgs

 The `FastVideoArgs` class in `fastvideo/fastvideo_args.py` serves as the central configuration system for FastVideo. It contains all parameters needed to control model loading, inference configuration, performance optimization settings, and more.
@@ -61,7 +60,6 @@ with set_current_fastvideo_args(fastvideo_args):
    result = generate_video()
 ```

-(design-pipeline-system)=
 ## Pipeline System

 ### `ComposedPipelineBase`
@@ -108,7 +106,8 @@ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> Forward
    return batch
 ```

-(design-forwardbatch)=
+![Pipeline execution and data flow](../assets/images/pipeline.png)
+
 ### ForwardBatch

 Defined in `fastvideo/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:
@@ -120,12 +119,10 @@ Defined in `fastvideo/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsul

 This structure facilitates clear state transitions between stages.

-(design-model-components)=
 ## Model Components

 The `fastvideo/models/` directory contains implementations of the core neural network models used in video diffusion:

-(design-transformer-models)=
 ### Transformer Models

 Transformer networks perform the actual denoising during diffusion:
@@ -152,7 +149,6 @@ def forward(
    return noise_pred  # Predicted noise residual
 ```

-(design-vae-variational-auto-encoder)=
 ### VAE (Variational Auto-Encoder)

 VAEs handle conversion between pixel space and latent space:
@@ -170,7 +166,6 @@ FastVideo's VAE implementations include:
 - Optional tiling for large frames
 - Distributed weight support

-(design-text-and-image-encoders)=
 ### Text and Image Encoders

 Encoders process conditioning inputs into embeddings:
@@ -188,7 +183,6 @@ FastVideo implements optimizations such as:
 - Caching for common prompts
 - Precision-tuned computation

-(design-schedulers)=
 ### Schedulers

 Schedulers manage the diffusion sampling process:
@@ -216,7 +210,10 @@ def step(
    return prev_sample
 ```

-(design-optimized-attention)=
+This diagram shows how models are discovered, validated, and loaded across entrypoints, executors, pipelines, and model loaders.
+
+![Model loading flow](../assets/images/load_models.png)
+
 ## Optimized Attention

 The `fastvideo/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:
@@ -240,17 +237,17 @@ self.attn = LocalAttention(
 # export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
 ```

+![Attention backend selector design](../assets/images/attention_backend.png)
+
 ### Attention Patterns
 Supports various patterns with memory optimization techniques:
 - **Cross/Self/Temporal/Global-Local Attention**
 - Chunking, progressive computation, optimized masking

-(design-distributed-processing)=
 ## Distributed Processing

 The `fastvideo/distributed/` directory contains implementations for distributed model execution:

-(design-tensor-parallelism)=
 ### Tensor Parallelism

 Tensor parallelism splits model weights across devices:
@@ -307,7 +304,6 @@ Efficient communication primitives minimize distributed overhead:
 - **Tensor-Parallel AllReduce**: Combines partial results
 - **Distributed Synchronization**: Coordinates execution

-(design-forwardcontext)=
 ## Forward Context Management

 ### ForwardContext
@@ -330,7 +326,6 @@ with set_forward_context(current_timestep, attn_metadata, fastvideo_args):
    output = model(inputs)
 ```

-(design-executor-and-worker-abstractions)=
 ## Executor and Worker System

 The `fastvideo/worker/` directory contains the distributed execution framework:
@@ -357,7 +352,6 @@ Each GPU worker:

 This design allows FastVideo to efficiently utilize multiple GPUs while providing a simple, unified interface for model execution.

-(design-platforms)=
 ## Platforms

 The `fastvideo/platforms/` directory provides hardware platform abstractions that enable FastVideo to run efficiently on different hardware configurations:
@@ -388,7 +382,6 @@ else:

 The platform system is designed to be extensible for future hardware targets.

-(design-logger)=
 ## Logger
 See [PR](https://github.com/hao-ai-lab/FastVideo/pull/356)

@@ -1,4 +1,3 @@
-(v0-data-preprocess)=

 # 🧱 Data Preprocess for Distillation

@@ -0,0 +1,12 @@
+# 💡 Examples
+
+A collection of examples demonstrating usage of FastVideo.
+
+All documented examples are autogenerated using [generate_examples.py](https://github.com/hao-ai-lab/FastVideo/blob/main/docs/generate_examples.py) from examples found in the [examples](https://github.com/hao-ai-lab/FastVideo/tree/main/examples) directory.
+
+## Examples
+
+- [Examples Distillation Index](distillation/examples/examples_distillation_index.md)
+- [Examples Training Index](training/examples/examples_training_index.md)
+- [Examples Inference Index](inference/examples/examples_inference_index.md)
+
@@ -6,10 +6,11 @@ import re
 from dataclasses import dataclass, field
 from pathlib import Path

-ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
-ROOT_DIR_RELATIVE = '../../../..'
+ROOT_DIR = Path(__file__).parent.parent.resolve()
+ROOT_DIR_RELATIVE = '../..'
 EXAMPLE_DIR = ROOT_DIR / "examples"
-EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples"
+EXAMPLE_DOC_DIR = ROOT_DIR / "docs/getting_started/examples"
+GITHUB_REPO = "hao-ai-lab/FastVideo"  # Update this to your repo


 def fix_case(text: str) -> str:
@@ -71,9 +72,16 @@ class Index:

    def generate(self) -> str:
        content = f"# {self.title}\n\n{self.description}\n\n"
-        content += ":::{toctree}\n"
-        content += f":caption: {self.caption}\n:maxdepth: {self.maxdepth}\n"
-        content += "\n".join(self.documents) + "\n:::\n"
+        if self.caption:
+            content += f"## {self.caption}\n\n"
+        # Generate a simple list of links for MkDocs
+        for doc in self.documents:
+            # Convert document path to proper link
+            doc_link = doc.replace("\\", "/")
+            # Get just the filename for the link text
+            doc_title = fix_case(Path(doc).stem.replace("_", " ").title())
+            content += f"- [{doc_title}]({doc_link}.md)\n"
+        content += "\n"
        return content


@@ -142,30 +150,66 @@ class Example:
        return fix_case(self.path.stem.replace("_", " ").title())

    def generate(self) -> str:
-        # Convert the path to a relative path from __file__
-        make_relative = lambda path: ROOT_DIR_RELATIVE / path.relative_to(
-            ROOT_DIR)
+        # Create GitHub link to source
+        github_path = str(self.path.relative_to(ROOT_DIR)).replace("\\", "/")
+        github_url = f"https://github.com/{GITHUB_REPO}/blob/main/{github_path}"
+        content = f"**Source:** [{github_path}]({github_url})\n\n"

-        content = f"Source <gh-file:{self.path.relative_to(ROOT_DIR)}>.\n\n"
-        include = "include" if self.main_file.suffix == ".md" else \
-            "literalinclude"
-        if include == "literalinclude":
+        # Add title for code files
+        if self.main_file.suffix != ".md":
            content += f"# {self.title}\n\n"
-        content += f":::{{{include}}} {make_relative(self.main_file)}\n"  # type: ignore[no-untyped-call]
-        if include == "literalinclude":
-            content += f":language: {self.main_file.suffix[1:]}\n"
-        content += ":::\n\n"
+
+        # Include main file content
+        if self.main_file.suffix == ".md":
+            # For markdown files, include the content directly
+            with open(self.main_file, encoding='utf-8') as f:
+                content += f.read() + "\n\n"
+        else:
+            # For code files, use code blocks
+            language = self.main_file.suffix[1:] if self.main_file.suffix else ""
+            with open(self.main_file, encoding='utf-8') as f:
+                file_content = f.read()
+            content += f"```{language}\n{file_content}\n```\n\n"

        if not self.other_files:
            return content

-        content += "## Example materials\n\n"
+        content += "## Additional Files\n\n"
+        # Define binary/non-text file extensions to skip
+        binary_extensions = {
+            '.mp4', '.avi', '.mov', '.mkv', '.gif', '.jpg', '.jpeg', '.png',
+            '.webp', '.bmp', '.pdf', '.zip', '.tar', '.gz', '.mp3', '.wav'
+        }
+
        for file in sorted(self.other_files):
-            include = "include" if file.suffix == ".md" else "literalinclude"
-            content += f":::{{admonition}} {file.relative_to(self.path)}\n"
-            content += ":class: dropdown\n\n"
-            content += f":::{{{include}}} {make_relative(file)}\n:::\n"  # type: ignore[no-untyped-call]
-            content += ":::\n\n"
+            # Skip binary files
+            if file.suffix.lower() in binary_extensions:
+                continue
+
+            file_rel_path = file.relative_to(self.path)
+            # Use collapsible admonition syntax for MkDocs
+            content += f"??? note \"{file_rel_path}\"\n\n"
+
+            try:
+                if file.suffix == ".md":
+                    # Include markdown content with indentation
+                    with open(file, encoding='utf-8') as f:
+                        for line in f:
+                            content += f"    {line}"
+                else:
+                    # Include code with proper formatting
+                    language = file.suffix[1:] if file.suffix else ""
+                    with open(file, encoding='utf-8') as f:
+                        file_content = f.read()
+                    # Indent the code block for the admonition
+                    content += f"    ```{language}\n"
+                    for line in file_content.split('\n'):
+                        content += f"    {line}\n"
+                    content += "    ```\n"
+                content += "\n"
+            except UnicodeDecodeError:
+                # Skip files that can't be decoded as UTF-8
+                continue

        return content

@@ -195,7 +239,7 @@ class NestedStructure:

 def create_category_indices() -> dict[str, Index]:
    """Create category indices with their respective configurations."""
-    main_index_dir = ROOT_DIR / "docs/source/examples"
+    main_index_dir = ROOT_DIR / "docs/examples"
    if not main_index_dir.exists():
        main_index_dir.mkdir(parents=True)

@@ -203,17 +247,16 @@ def create_category_indices() -> dict[str, Index]:
        "inference":
        Index(
            path=ROOT_DIR /
-            "docs/source/inference/examples/examples_inference_index.md",
+            "docs/inference/examples/examples_inference_index.md",
            title="🚀 Examples",
            description=
-            "Inference examples demonstrate how to use FastVideo inference. We recommend starting with <project:basic.md>.",
+            "Inference examples demonstrate how to use FastVideo inference. We recommend starting with [basic.md](basic.md).",
            caption="Examples",
            maxdepth=1,
        ),
        "training":
        Index(
-            path=ROOT_DIR /
-            "docs/source/training/examples/examples_training_index.md",
+            path=ROOT_DIR / "docs/training/examples/examples_training_index.md",
            title="🚀 Examples",
            description=
            "Training examples demonstrate how to use FastVideo training.",
@@ -223,7 +266,7 @@ def create_category_indices() -> dict[str, Index]:
        "distillation":
        Index(
            path=ROOT_DIR /
-            "docs/source/distillation/examples/examples_distillation_index.md",
+            "docs/distillation/examples/examples_distillation_index.md",
            title="🚀 Examples",
            description=
            "Distillation examples demonstrate how to use FastVideo distillation.",
@@ -246,9 +289,21 @@ def find_examples(category_indices: dict[str, Index],
    examples = []
    glob_patterns = ["*.py", "*.md", "*.sh"]

+    # Map category names to actual directory names
+    category_dir_mapping = {
+        "distillation": "distill",  # examples/distill/ -> distillation category
+    }
+
    # Find categorised examples
    for category in category_indices:
-        category_dir = EXAMPLE_DIR / category
+        # Use mapped directory name if available, otherwise use category name
+        dir_name = category_dir_mapping.get(category, category)
+        category_dir = EXAMPLE_DIR / dir_name
+
+        # Skip if directory doesn't exist
+        if not category_dir.exists():
+            continue
+
        globs = [category_dir.glob(pattern) for pattern in glob_patterns]
        for path in itertools.chain(*globs):
            examples.append(Example(path, category))
@@ -279,11 +334,18 @@ def create_nested_structures(
                                                dict[str,
                                                     NestedStructure]]]] = {}

+    # Map category names to actual directory names
+    category_dir_mapping = {
+        "distillation": "distill",
+    }
+
    for example in examples:
        if example.category not in ["training", "distillation"]:
            continue

-        category_dir = EXAMPLE_DIR / example.category
+        # Use mapped directory name if available
+        dir_name = category_dir_mapping.get(example.category, example.category)
+        category_dir = EXAMPLE_DIR / dir_name
        relative_path = example.path.relative_to(category_dir)
        path_parts = relative_path.parts

@@ -415,7 +477,7 @@ def generate_nested_examples(nested_structures: dict[str, dict[str, dict[
            category_index.documents.append(method)


-def generate_examples(generate_main_index=False):
+def generate_examples(generate_main_index: bool = False) -> None:
    """
    Generate example documentation.
    
@@ -429,12 +491,14 @@ def generate_examples(generate_main_index=False):
    # Create the main examples index only if requested
    examples_index = None
    if generate_main_index:
-        main_index_dir = ROOT_DIR / "docs/source/examples"
+        main_index_dir = ROOT_DIR / "docs/examples"
        examples_index = Index(
            path=main_index_dir / "examples_index.md",
            title="💡 Examples",
            description=
-            "A collection of examples demonstrating usage of FastVideo.\nAll documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.",
+            "A collection of examples demonstrating usage of FastVideo.\n\n"
+            f"All documented examples are autogenerated using [generate_examples.py](https://github.com/{GITHUB_REPO}/blob/main/docs/generate_examples.py) "
+            f"from examples found in the [examples](https://github.com/{GITHUB_REPO}/tree/main/examples) directory.",
            caption="Examples",
            maxdepth=2)

@@ -471,3 +535,19 @@ def generate_examples(generate_main_index=False):
    if generate_main_index and examples_index:
        with open(examples_index.path, "w+") as f:
            f.write(examples_index.generate())
+
+
+def on_pre_build_hook(config, **kwargs):
+    """
+    MkDocs hook to generate examples before building the documentation.
+    This function is called automatically by the mkdocs-simple-hooks plugin.
+    """
+    print("Generating example documentation...")
+    generate_examples(generate_main_index=True)
+    print("Example documentation generated successfully!")
+
+
+if __name__ == "__main__":
+    print("Generating example documentation...")
+    generate_examples(generate_main_index=True)
+    print("Example documentation generated successfully!")
@@ -0,0 +1,41 @@
+
+# 🔧 Installation
+
+FastVideo supports the following hardware platforms:
+
+- [NVIDIA CUDA](installation/gpu.md)
+- [Apple silicon](installation/mps.md)
+
+## Quick Installation
+
+### Using pip
+
+```bash
+pip install fastvideo
+```
+
+### Using conda
+
+```bash
+conda install -c conda-forge fastvideo
+```
+
+### From source
+
+```bash
+git clone https://github.com/hao-ai-lab/FastVideo.git
+cd FastVideo
+pip install -e .
+```
+
+## Hardware Requirements
+
+- **NVIDIA GPUs**: CUDA 11.8+ with compute capability 7.0+
+- **Apple Silicon**: macOS 12.0+ with M1/M2/M3 chips
+- **CPU**: x86_64 architecture (for CPU-only inference)
+
+## Next Steps
+
+- [Quick Start Guide](quick_start.md) - Get started with your first video generation
+- [Configuration](../inference/configuration.md) - Learn about configuration options
+- [Examples](../inference/examples/) - Explore example scripts and notebooks
@@ -30,17 +30,8 @@ conda create -n fastvideo python=3.12 -y
 conda activate fastvideo
 ```

-:::{note}
-[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
-:::
-
 #### uv

-:::{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:

 ```console
@@ -31,17 +31,8 @@ conda create -n fastvideo python=3.12.4 -y
 conda activate fastvideo
 ```

-:::{note}
-[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
-:::
-
 #### uv

-:::{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:

 ```console
@@ -0,0 +1,55 @@
+# 🚀 Quick Start
+
+Get up and running with FastVideo in minutes!
+
+## Installation
+
+First, install FastVideo:
+
+```bash
+pip install fastvideo
+```
+
+## Basic Usage
+
+### Text-to-Video Generation
+
+```python
+from fastvideo import FastVideoPipeline
+
+# Initialize the pipeline
+pipe = FastVideoPipeline.from_pretrained("wan2.1-t2v-1.3B")
+
+# Generate a video
+prompt = "A cat playing with a ball of yarn"
+video = pipe(prompt, num_frames=16, height=512, width=512)
+
+# Save the video
+video.save("output.mp4")
+```
+
+### Image-to-Video Generation
+
+```python
+from fastvideo import FastVideoPipeline
+from PIL import Image
+
+# Load an image
+image = Image.open("input.jpg")
+
+# Initialize the pipeline
+pipe = FastVideoPipeline.from_pretrained("wan2.1-i2v-14B-480p")
+
+# Generate a video from the image
+video = pipe(image, num_frames=16, height=480, width=480)
+
+# Save the video
+video.save("output.mp4")
+```
+
+## Next Steps
+
+- [Installation Guide](installation.md) - Detailed installation instructions
+- [Configuration](../inference/configuration.md) - Learn about configuration options
+- [Examples](../inference/examples/) - Explore more examples
+- [Optimizations](../inference/optimizations.md) - Performance optimization tips
@@ -0,0 +1,42 @@
+# V1 API
+
+FastVideo's V1 API provides a streamlined interface for video generation tasks with powerful customization options. This page documents the primary components of the API.
+
+## Video Generator
+
+This class will be the primary Python API for generating videos and images.
+
+::: fastvideo.entrypoints.video_generator.VideoGenerator
+    options:
+      show_root_heading: true
+      show_source: false
+      members:
+        - from_pretrained
+      heading_level: 3
+
+`VideoGenerator.from_pretrained()` should be the primary way of creating a new video generator.
+
+## Configuring FastVideo
+
+The following two classes `PipelineConfig` and `SamplingParam` are used to configure initialization and sampling parameters, respectively.
+
+### PipelineConfig
+
+::: fastvideo.configs.pipelines.base.PipelineConfig
+    options:
+      show_root_heading: true
+      show_source: false
+      members:
+        - from_pretrained
+        - dump_to_json
+      heading_level: 4
+
+### SamplingParam
+
+::: fastvideo.configs.sample.base.SamplingParam
+    options:
+      show_root_heading: true
+      show_source: false
+      members:
+        - from_pretrained
+      heading_level: 4
@@ -1,30 +1,24 @@
 # Welcome to FastVideo

-:::{figure} ../../assets/logos/logo.svg
-:align: center
-:alt: FastVideo
-:class: no-scaled-link
-:width: 60%
-:::
+<div style="text-align: center;">
+  <img src="assets/logos/logo.svg" alt="FastVideo" style="width: 60%;" />
+</div>

-:::{raw} html
-<p style="text-align:center">
-<strong>FastVideo is a unified inference and post-training framework for accelerated video generation.
-</strong>
-</p>
+<div style="text-align: center;">
+<strong>FastVideo is a unified inference and post-training framework for accelerated video generation.</strong>
+</div>

-<p style="text-align:center">
+<div style="text-align: center;">
 <script async defer src="https://buttons.github.io/buttons.js"></script>
 <a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/" data-show-count="true" data-size="large" aria-label="Star">Star</a>
 <a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/subscription" data-icon="octicon-eye" data-size="large" aria-label="Watch">Watch</a>
 <a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/fork" data-icon="octicon-repo-forked" data-size="large" aria-label="Fork">Fork</a>
-</p>
-:::
+</div>

 FastVideo is an inference and post-training framework for diffusion models. It features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.

 <div style="text-align: center;">
-  <img src=_static/images/fastwan.png width="100%"/>
+  <img src="assets/images/fastwan.png" style="width: 100%;"/>
 </div>

 ## Key Features
@@ -42,90 +36,16 @@ FastVideo has the following features:

 ## Documentation

-% How to start using FastVideo?
+Welcome to FastVideo! This documentation will help you get started with our unified inference and post-training framework for accelerated video generation.

-:::{toctree}
-:caption: Getting Started
-:maxdepth: 1
+Use the navigation menu on the left to explore different sections:

-getting_started/installation
-<!-- getting_started/v1_api -->
-:::
-
-:::{toctree}
-:caption: Inference
-:maxdepth: 1
-
-inference/inference_quick_start
-inference/examples/examples_inference_index
-inference/configuration
-inference/optimizations
-inference/comfyui
-inference/support_matrix
-inference/cli
-inference/add_pipeline
-:::
-
-:::{toctree}
-:caption: Training
-:maxdepth: 1
-
-training/examples/examples_training_index
-training/data_preprocess
-<!-- training/finetune -->
-:::
-
-:::{toctree}
-:caption: Distillation
-:maxdepth: 1
-
-distillation/examples/examples_distillation_index
-distillation/data_preprocess
-distillation/dmd
-:::
-
-% What is STA Kernel?
-
-:::{toctree}
-:caption: Sliding Tile Attention
-:maxdepth: 1
-
-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
-design/overview
-:::
-
-:::{toctree}
-:caption: Developer Guide
-:maxdepth: 2
-
-contributing/overview
-contributing/developer_env/index
-:::
-
-:::{toctree}
-:caption: API Reference
-:maxdepth: 2
-
-<!-- api/summary -->
-api/fastvideo/fastvideo
-:::
-
-## Indices and tables
-
- {ref}`genindex`
- {ref}`modindex`
+- **Getting Started**: Installation and quick start guides
+- **Inference**: Learn how to use FastVideo for video generation
+- **Training**: Data preprocessing and fine-tuning workflows  
+- **Distillation**: Post-training optimization techniques
+- **Sliding Tile Attention**: Advanced attention mechanisms
+- **Video Sparse Attention**: Efficient attention for video models
+- **Design**: Framework architecture and design principles
+- **Developer Guide**: Contributing and development setup
+- **API Reference**: Complete API documentation
@@ -1,4 +1,3 @@
-(add-pipeline)=

 # 🏗️ Adding a New Pipeline

@@ -1,4 +1,4 @@
-(inference-configuration)=
+
 # Configuration

 ## Multi-GPU Setup
@@ -1,32 +1,37 @@
-(inference-optimizations)=
+
 # Optimizations

 This page describes the various options for speeding up generation times in FastVideo.

 ## Table of Contents
+
 - Optimized Attention Backends
+
  - [Flash Attention](#optimizations-flash)
  - [Sliding Tile Attention](#optimizations-sta)
  - [Sage Attention](#optimizations-sage)
+  - [Sage Attention 3](#optimizations-sage3)

 - Caching Techniques
  - [TeaCache](#optimizations-teacache)

-(optimizations-backends)=
 ## Attention Backends

 ### Available Backends
+
 - 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`
+- Sage Attention 3: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN_THREE`

 ### Configuring Backends

 There are two ways to configure the attention backend in FastVideo.

 #### 1. In Python
+
 In python, set the `FASTVIDEO_ATTENTION_BACKEND` environment variable before instantiating `VideoGenerator` like this:

 ```python
@@ -34,13 +39,13 @@ os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
 ```

 #### 2. In CLI
+
 You can also set the environment variable on the command line:

 ```bash
 FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
 ```

-(optimizations-flash)=
 ### Flash Attention

 **`FLASH_ATTN`**
@@ -57,16 +62,12 @@ And if using a Hopper+ GPU (ie H100), installing [Flash Attention 3](https://git
 git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention

 cd hopper
-pip install ninja 
+pip install ninja
 python setup.py install
 ```

-:::{note}
-FastVideo will automatically detect and use `FA3` if it is installed when using `FLASH_ATTN` backend.
-:::
-
-(optimizations-sta)=
 ### Sliding Tile Attention
+
 **`SLIDING_TILE_ATTN`**

 ```bash
@@ -75,8 +76,8 @@ 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
@@ -86,20 +87,40 @@ python setup_vsa.py install

 Please see [this page](#vsa-installation) for more installation instructions.

-(optimizations-sage)=
 ### Sage Attention
+
 **`SAGE_ATTN`**

 To use [SageAttention](https://github.com/thu-ml/SageAttention) 2.1.1, please compile from source:

 ```bash
 git clone https://github.com/thu-ml/SageAttention.git
-cd sageattention 
+cd sageattention
 python setup.py install  # or pip install -e .
 ```

-(optimizations-teacache)=
+### Sage Attention 3
+
+**`SAGE_ATTN_THREE`**
+
+[SageAttention 3](https://huggingface.co/jt-zhang/SageAttention3) is an advanced attention mechanism that leverages FP4 quantization and Blackwell GPU Tensor Cores for significant performance improvements.
+
+#### Hardware Requirements
+
+- RTX5090
+
+#### Installation
+
+Note that Sage Attention 3 requires `python>=3.13`, `torch>=2.8.0`, `CUDA >=12.8`. If you are using `uv` and using `torch==2.8.0` make sure that `sentencepiece==0.2.1` in the pyproject.toml file.
+
+To use Sage Attention 3 in FastVideo, first get access to the SageAttention3 code, then move `sageattn/` and `setup.py` to the directory `fastvideo/attention/backends`, then install from using:
+
+```bash
+python setup.py install
+```
+
 ## Teacache
+
 TeaCache is an optimization technique supported in FastVideo that can significantly speed up video generation by skipping redundant calculations across diffusion steps. This guide explains how to enable and configure TeaCache for optimal performance in FastVideo.

 ### What is TeaCache?
@@ -0,0 +1,66 @@
+# Compatibility Matrix
+
+The table below shows every supported model and optimizations supported for them.
+
+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.
+
+<style>
+  /* Target tables in this section */
+  #models-x-optimization + p + table {
+    display: block;
+    overflow-x: auto;
+    width: 100%;
+    font-size: 0.85rem;
+  }
+  
+  #models-x-optimization + p + table td,
+  #models-x-optimization + p + table th {
+    text-align: center;
+    white-space: nowrap;
+    padding: 0.5em;
+  }
+  
+  /* First two columns can wrap */
+  #models-x-optimization + p + table td:nth-child(1),
+  #models-x-optimization + p + table td:nth-child(2) {
+    white-space: normal;
+    min-width: 120px;
+  }
+  
+  #models-x-optimization + p + table td:nth-child(2) code {
+    font-size: 0.75rem;
+  }
+</style>
+
+| Model Name | HuggingFace Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | 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-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ |
+| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480P, 720P | ✅ | ✅* | ✅ | ⭕ |
+| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480P | ✅ | ✅* | ✅ | ⭕ |
+| Wan2.1 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.
+
+## Special requirements
+
+### StepVideo T2V
+- The self-attention in text-encoder (step_llm) only supports CUDA capabilities sm_80 sm_86 and sm_90
+
+### Sliding Tile Attention
+- Currently only Hopper GPUs (H100s) are supported.
@@ -1,35 +0,0 @@
-@ECHO OFF
-
-pushd %~dp0
-
-REM Command file for Sphinx documentation
-
-if "%SPHINXBUILD%" == "" (
-	set SPHINXBUILD=sphinx-build
-)
-set SOURCEDIR=source
-set BUILDDIR=build
-
-%SPHINXBUILD% >NUL 2>NUL
-if errorlevel 9009 (
-	echo.
-	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
-	echo.installed, then set the SPHINXBUILD environment variable to point
-	echo.to the full path of the 'sphinx-build' executable. Alternatively you
-	echo.may add the Sphinx directory to PATH.
-	echo.
-	echo.If you don't have Sphinx installed, grab it from
-	echo.https://www.sphinx-doc.org/
-	exit /b 1
-)
-
-if "%1" == "" goto help
-
-%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
-goto end
-
-:help
-%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
-
-:end
-popd
@@ -1,15 +0,0 @@
-sphinx==7.4.7
-sphinx-argparse==0.5.2
-sphinx-autodoc2==0.5.0
-sphinx-book-theme==1.1.4
-sphinx-copybutton==0.5.2
-sphinx-design==0.6.1
-sphinx-togglebutton==0.3.2
-myst-parser==3.0.1
-msgspec
-commonmark # Required by sphinx-argparse when using :markdownhelp:
-
-# packages to install to build the documentation
-cachetools
-# -f https://download.pytorch.org/whl/cpu
-torch
@@ -1,4 +1,3 @@
-(sta-demo)=

 # 🔍 Demo
 This is is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
@@ -1,4 +1,3 @@
-(sta-installation)=

 # 🔧 Installation
 You can install the Sliding Tile Attention package using
@@ -1,8 +0,0 @@
-.vertical-table-header th.head:not(.stub) {
-    writing-mode: sideways-lr;
-    white-space: nowrap;
-    max-width: 0;
-    p {
-       margin: 0;
-    }
-}
@@ -1,39 +0,0 @@
-<style>
-  .notification-bar {
-    width: 100vw;
-    display: flex;
-    justify-content: center;
-    align-items: center;
-    font-size: 16px;
-    padding: 0 6px 0 6px;
-  }
-  .notification-bar p {
-    margin: 0;
-  }
-  .notification-bar a {
-    font-weight: bold;
-    text-decoration: none;
-  }
-
-  /* Light mode styles (default) */
-  .notification-bar {
-    background-color: #fff3cd;
-    color: #856404;
-  }
-  .notification-bar a {
-    color: #d97706;
-  }
-
-  /* Dark mode styles */
-  html[data-theme=dark] .notification-bar {
-    background-color: #333;
-    color: #ddd;
-  }
-  html[data-theme=dark] .notification-bar a {
-    color: #ffa500; /* Brighter color for visibility */
-  }
-</style>
-
-<!-- <div class="notification-bar">
-  <p>You are viewing the latest developer preview docs. <a href="https://docs.vllm.ai/en/stable/">Click here</a> to view docs for the latest stable release.</p>
-</div> -->
@@ -1,19 +0,0 @@
-# Summary
-
-## Video Generator
-
-```{autodoc2-summary}
-    fastvideo.VideoGenerator
-```
-
-## Initialization Configuration
-
-```{autodoc2-summary}
-    fastvideo.configs.pipelines.PipelineConfig
-```
-
-## Sampling Configuration
-
-```{autodoc2-summary}
-    fastvideo.configs.sample.SamplingParam
-```
@@ -1,22 +0,0 @@
-# type: ignore
-# SPDX-License-Identifier: Apache-2.0
-from docutils import nodes
-from myst_parser.parsers.sphinx_ import MystParser
-from sphinx.ext.napoleon import docstring
-
-
-class NapoleonParser(MystParser):
-
-    def parse(self, input_string: str, document: nodes.document) -> None:
-        # Get the Sphinx configuration
-        config = document.settings.env.config
-
-        parsed_content = str(
-            docstring.GoogleDocstring(
-                str(docstring.NumpyDocstring(input_string, config)),
-                config,
-            ))
-        return super().parse(parsed_content, document)
-
-
-Parser = NapoleonParser
@@ -1,275 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-# Configuration file for the Sphinx documentation builder.
-#
-# This file only contains a selection of the most common options. For a full
-# list see the documentation:
-# https://www.sphinx-doc.org/en/master/usage/configuration.html
-
-# -- Path setup --------------------------------------------------------------
-
-# If extensions (or modules to document with autodoc) are in another directory,
-# add these directories to sys.path here. If the directory is relative to the
-# documentation root, use os.path.abspath to make it absolute, like shown here.
-
-import datetime
-import logging
-import os
-import re
-import sys
-from pathlib import Path
-
-import requests
-
-logger = logging.getLogger(__name__)
-REPO_ROOT = Path(__file__).resolve().parent.parent.parent
-print(os.path.abspath(REPO_ROOT))
-sys.path.append(os.path.abspath(REPO_ROOT))
-
-# -- Project information -----------------------------------------------------
-
-project = 'FastVideo'
-copyright = f'{datetime.datetime.now().year}, FastVideo Team'
-author = 'the FastVideo Team'
-
-# -- General configuration ---------------------------------------------------
-
-# Add any Sphinx extension module names here, as strings. They can be
-# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
-# ones.
-extensions = [
-    "sphinx.ext.napoleon",
-    "sphinx.ext.linkcode",
-    "sphinx.ext.intersphinx",
-    "sphinx_copybutton",
-    "autodoc2",
-    "myst_parser",
-    "sphinxarg.ext",
-    "sphinx_design",
-    "sphinx_togglebutton",
-]
-myst_enable_extensions = [
-    "colon_fence",
-    "fieldlist",
-]
-autodoc2_packages = [
-    {
-        "path": "../../fastvideo",
-        "exclude_dirs": ["__pycache__", "third_party"],
-    },
-]
-autodoc2_output_dir = "api"
-autodoc2_render_plugin = "myst"
-autodoc2_hidden_objects = ["dunder", "private", "inherited"]
-autodoc2_docstring_parser_regexes = [
-    (".*", "docs.source.autodoc2_docstring_parser"),
-]
-autodoc2_sort_names = True
-autodoc2_index_template = None
-autodoc2_skip_module_regexes = [
-    "fastvideo.dataset",
-    "fastvideo.distill",
-    "fastvideo.data_preprocess",
-    "fastvideo.models",
-    "fastvideo.sample",
-    "fastvideo.utils",
-    "fastvideo.distill_adv",
-    "fastvideo.train",
-]
-
-# Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
-
-# List of patterns, relative to source directory, that match files and
-# directories to ignore when looking for source files.
-# This pattern also affects html_static_path and html_extra_path.
-exclude_patterns: list[str] = ["**/*.template.md", "**/*.inc.md"]
-
-# Exclude the prompt "$" when copying code
-copybutton_prompt_text = r"\$ "
-copybutton_prompt_is_regexp = True
-
-# -- Options for HTML output -------------------------------------------------
-
-# The theme to use for HTML and HTML Help pages.  See the documentation for
-# a list of builtin themes.
-#
-html_title = project
-html_theme = 'sphinx_book_theme'
-html_logo = '../../assets/logos/icon_simple.svg'
-html_theme_options = {
-    'path_to_docs': 'docs/source',
-    'repository_url': 'https://github.com/hao-ai-lab/FastVideo/',
-    'use_repository_button': True,
-    'use_edit_page_button': True,
-    # Prevents the full API being added to the left sidebar of every page.
-    # Reduces build time by 2.5x and reduces build size from ~225MB to ~95MB.
-    'collapse_navbar': True,
-    # Makes API visible in the right sidebar on API reference pages.
-    'show_toc_level': 3,
-}
-# Add any paths that contain custom static files (such as style sheets) here,
-# relative to this directory. They are copied after the builtin static files,
-# so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ["_static"]
-html_js_files = ["custom.js"]
-html_css_files = ["custom.css"]
-
-myst_url_schemes = {
-    'http': None,
-    'https': None,
-    'mailto': None,
-    'ftp': None,
-    "gh-issue": {
-        "url":
-        "https://github.com/hao-ai-lab/FastVideo/issues/{{path}}#{{fragment}}",
-        "title": "Issue #{{path}}",
-        "classes": ["github"],
-    },
-    "gh-pr": {
-        "url":
-        "https://github.com/hao-ai-lab/FastVideo/pull/{{path}}#{{fragment}}",
-        "title": "Pull Request #{{path}}",
-        "classes": ["github"],
-    },
-    "gh-dir": {
-        "url": "https://github.com/hao-ai-lab/FastVideo/tree/main/{{path}}",
-        "title": "{{path}}",
-        "classes": ["github"],
-    },
-    "gh-file": {
-        "url": "https://github.com/hao-ai-lab/FastVideo/blob/main/{{path}}",
-        "title": "{{path}}",
-        "classes": ["github"],
-    },
-}
-
-# see https://docs.readthedocs.io/en/stable/reference/environment-variables.html # noqa
-READTHEDOCS_VERSION_TYPE = os.environ.get('READTHEDOCS_VERSION_TYPE')
-if READTHEDOCS_VERSION_TYPE == "tag":
-    # remove the warning banner if the version is a tagged release
-    header_file = os.path.join(os.path.dirname(__file__),
-                               "_templates/sections/header.html")
-    # The file might be removed already if the build is triggered multiple times
-    # (readthedocs build both HTML and PDF versions separately)
-    if os.path.exists(header_file):
-        os.remove(header_file)
-
-
-# Generate additional rst documentation here.
-def setup(app):
-    from docs.source.generate_examples import generate_examples
-    generate_examples()
-
-
-_cached_base: str = ""
-_cached_branch: str = ""
-
-
-def get_repo_base_and_branch(pr_number: str) -> tuple[str | None, str | None]:
-    global _cached_base, _cached_branch
-    if _cached_base and _cached_branch:
-        return _cached_base, _cached_branch
-
-    url = f"https://api.github.com/repos/hao-ai-lab/FastVideo/pulls/{pr_number}"
-    response = requests.get(url)
-    if response.status_code == 200:
-        data = response.json()
-        _cached_base = data['head']['repo']['full_name']
-        _cached_branch = data['head']['ref']
-        return _cached_base, _cached_branch
-    else:
-        logger.error("Failed to fetch PR details: %s", response)
-        return None, None
-
-
-def linkcode_resolve(domain, info):
-    if domain != 'py':
-        return None
-    if not info['module']:
-        return None
-
-    # Get path from module name
-    file = Path(f"{info['module'].replace('.', '/')}.py")
-    path = REPO_ROOT / file
-    if not path.exists():
-        path = REPO_ROOT / file.with_suffix("") / "__init__.py"
-    if not path.exists():
-        return None
-
-    # Get the line number of the object
-    with open(path) as f:
-        lines = f.readlines()
-    name = info['fullname'].split(".")[-1]
-    pattern = fr"^( {{4}})*((def|class) )?{name}\b.*"
-    for lineno, line in enumerate(lines, 1):
-        if not line or line.startswith("#"):
-            continue
-        if re.match(pattern, line):
-            break
-
-    # If the line number is not found, return None
-    if lineno == len(lines):
-        return None
-
-    # If the line number is found, create the URL
-    filename = path.relative_to(REPO_ROOT)
-    if "checkouts" in path.parts:
-        # a PR build on readthedocs
-        pr_number = REPO_ROOT.name
-        base, branch = get_repo_base_and_branch(pr_number)
-        if base and branch:
-            return f"https://github.com/{base}/blob/{branch}/{filename}#L{lineno}"
-    # Otherwise, link to the source file on the main branch
-    return f"https://github.com/hao-ai-lab/FastVideo/blob/main/{filename}#L{lineno}"
-
-
-# Mock out external dependencies here, otherwise the autodoc pages may be blank.
-autodoc_mock_imports = [
-    "blake3",
-    "compressed_tensors",
-    "cpuinfo",
-    "cv2",
-    "torch",
-    "huggingface_hub",
-    "torchvision",
-    "transformers",
-    "psutil",
-    "prometheus_client",
-    "sentencepiece",
-    "vllm._C",
-    "PIL",
-    "numpy",
-    'triton',
-    "tqdm",
-    "tensorizer",
-    "pynvml",
-    "outlines",
-    "xgrammar",
-    "librosa",
-    "soundfile",
-    "gguf",
-    "lark",
-    "decord",
-]
-
-for mock_target in autodoc_mock_imports:
-    if mock_target in sys.modules:
-        logger.info(
-            "Potentially problematic mock target (%s) found; "
-            "autodoc_mock_imports cannot mock modules that have already "
-            "been loaded into sys.modules when the sphinx build starts.",
-            mock_target)
-
-intersphinx_mapping = {
-    "python": ("https://docs.python.org/3", None),
-    "typing_extensions":
-    ("https://typing-extensions.readthedocs.io/en/latest", None),
-    "aiohttp": ("https://docs.aiohttp.org/en/stable", None),
-    "pillow": ("https://pillow.readthedocs.io/en/stable", None),
-    "numpy": ("https://numpy.org/doc/stable", None),
-    "torch": ("https://pytorch.org/docs/stable", None),
-    "psutil": ("https://psutil.readthedocs.io/en/stable", None),
-}
-
-navigation_with_keys = False
@@ -1,18 +0,0 @@
-(installation-index)=
-
-# 🔧 Installation
-
-FastVideo supports the following hardware platforms:
-
-:::{toctree}
-:maxdepth: 1
-:hidden:
-
-installation/gpu
-installation/mps
-:::
-
- <project:installation/gpu.md>
-  - NVIDIA CUDA
- <project:installation/mps.md>
-  - Apple silicon
@@ -1,83 +0,0 @@
-# V1 API
-
-FastVideo's V1 API provides a streamlined interface for video generation tasks with powerful customization options. This page documents the primary components of the API.
-
-## Video Generator
-
-This class will be the primary Python API for generating videos and images.
-
-```{autodoc2-summary}
-    fastvideo.VideoGenerator
-```
-
-`````{py:class} VideoGenerator(fastvideo_args: fastvideo.fastvideo_args.FastVideoArgs, executor_class: type[fastvideo.worker.executor.Executor], log_stats: bool)
-:canonical: fastvideo.entrypoints.video_generator.VideoGenerator
-
-```{autodoc2-docstring} fastvideo.entrypoints.video_generator.VideoGenerator
-:parser: docs.source.autodoc2_docstring_parser
-```
-
-`VideoGenerator.from_pretrained()` should be the primary way of creating a new video generator.
-
-````{py:method} from_pretrained(model_path: str, device: typing.Optional[str] = None, torch_dtype: typing.Optional[torch.dtype] = None, pipeline_config: typing.Optional[typing.Union[str | fastvideo.configs.pipelines.PipelineConfig]] = None, **kwargs) -> fastvideo.entrypoints.video_generator.VideoGenerator
-:canonical: fastvideo.entrypoints.video_generator.VideoGenerator.from_pretrained
-:classmethod:
-
-```{autodoc2-docstring} fastvideo.entrypoints.video_generator.VideoGenerator.from_pretrained
-:parser: docs.source.autodoc2_docstring_parser
-```
-
-
-## Configuring FastVideo
-
-The follow two classes `PipelineConfig` and `SamplingParam` are used to configure initialization and sampling parameters, respectively.
-
-### PipelineConfig
-```{autodoc2-summary}
-    fastvideo.PipelineConfig
-```
-
-`````{py:class} PipelineConfig
-:canonical: fastvideo.configs.pipelines.base.PipelineConfig
-
-```{autodoc2-docstring} fastvideo.configs.pipelines.base.PipelineConfig
-:parser: docs.source.autodoc2_docstring_parser
-```
-
-````{py:method} from_pretrained(model_path: str) -> fastvideo.configs.pipelines.base.PipelineConfig
-:canonical: fastvideo.configs.pipelines.base.PipelineConfig.from_pretrained
-:classmethod:
-
-```{autodoc2-docstring} fastvideo.configs.pipelines.base.PipelineConfig.from_pretrained
-:parser: docs.source.autodoc2_docstring_parser
-```
-
-
-````{py:method} dump_to_json(file_path: str)
-:canonical: fastvideo.configs.pipelines.base.PipelineConfig.dump_to_json
-
-```{autodoc2-docstring} fastvideo.configs.pipelines.base.PipelineConfig.dump_to_json
-:parser: docs.source.autodoc2_docstring_parser
-```
-
-
-### SamplingParam
-
-```{autodoc2-summary}
-    fastvideo.SamplingParam
-```
-
-`````{py:class} SamplingParam
-:canonical: fastvideo.configs.sample.base.SamplingParam
-
-```{autodoc2-docstring} fastvideo.configs.sample.base.SamplingParam
-:parser: docs.source.autodoc2_docstring_parser
-```
-
-````{py:method} from_pretrained(model_path: str) -> fastvideo.configs.sample.base.SamplingParam
-:canonical: fastvideo.configs.sample.base.SamplingParam.from_pretrained
-:classmethod:
-
-```{autodoc2-docstring} fastvideo.configs.sample.base.SamplingParam.from_pretrained
-:parser: docs.source.autodoc2_docstring_parser
-```
@@ -1,136 +0,0 @@
-(support-matrix)=
-# Compatibility Matrix
-The table below shows every supported model and optimizations supported for them.
-
-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.
-
-:::{raw} html
-<style>
-  /* Make smaller to try to improve readability  */
-  td {
-    font-size: 0.9rem;
-    text-align: center;
-  }
-
-  th {
-    text-align: center;
-    font-size: 0.9rem;
-  }
-</style>
-:::
-
-:::{list-table}
-:header-rows: 1
-:stub-columns: 3
-:widths: auto
-:class: vertical-table-header
-
- * Model Name
-  * HuggingFace Model ID
-  * Resolutions
-  * 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-AI/Wan2.1-T2V-1.3B-Diffusers`
-  * 480P
-  * ✅
-  * ✅*
-  * ✅
-  * ⭕
- * Wan2.1 T2V 14B
-  * `Wan-AI/Wan2.1-T2V-14B-Diffusers`
-  * 480P, 720P
-  * ✅
-  * ✅*
-  * ✅
-  * ⭕
- * Wan2.1 I2V 480P
-  * `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers`
-  * 480P
-  * ✅
-  * ✅*
-  * ✅
-  * ⭕
- * Wan2.1 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.
-
-## Special requirements
-
-### StepVideo T2V
- The self-attention in text-encoder (step_llm) only supports CUDA capabilities sm_80 sm_86 and sm_90
-
-### Sliding Tile Attention
- Currently only Hopper GPUs (H100s) are supported.
@@ -1,4 +1,3 @@
-(v0-data-preprocess)=

 # 🧱 Data Preprocess

@@ -1,4 +1,4 @@
-(v0-finetune)=
+
 # 🧠 Finetune
 ## ⚡ Full Finetune
 Ensure your data is prepared and preprocessed in the format specified in [data_preprocess.md](#v0-data-preprocess). For convenience, we also provide a mochi preprocessed Black Myth Wukong data that can be downloaded directly:
@@ -1,4 +1,3 @@
-(vsa-installation)=

 # 🔧 Installation
 You can install the Video Sparse Attention package using
@@ -0,0 +1,140 @@
+#!/bin/bash
+#SBATCH --job-name=t2v
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:1
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=dmd_t2v_output/t2v_%j.out
+#SBATCH --error=dmd_t2v_output/t2v_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+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=29503
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY=your_wandb_api_key
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+# Configs
+NUM_GPUS=1
+
+# Model paths for Self-Forcing DMD distillation:
+GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Critic model
+
+DATA_DIR=your_data_dir
+VALIDATION_DATASET_FILE=your_validation_data_dir
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+training_args=(
+  --tracker_project_name SFwan_t2v_distill_self_forcing_dmd  
+  --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 21
+  --num_height 480
+  --num_width 832
+  --enable_gradient_checkpointing_type "full"
+  --log_visualization
+  --simulate_generator_forward
+  --num_frames 81
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+)
+
+parallel_args=(
+  --num_gpus $NUM_GPUS # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim $NUM_GPUS
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH  # TODO: check if you can remove this in this script
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 1.0
+)
+
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 100
+  --init_weights_from_safetensors your_ode_init_weights_path
+)
+
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 3.0
+  --fake_score_learning_rate 8e-6
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+)
+
+torchrun \
+--nnodes 1 \
+--master_port $MASTER_PORT \
+--nproc_per_node $NUM_GPUS \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -0,0 +1,24 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_TYPE="wan"
+DATA_MERGE_PATH="data/crush-smol/merge.txt"
+OUTPUT_DIR="data/crush-smol_processed_t2v/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 8 \
+    --seed 42 \
+    --max_height 480 \
+    --max_width 832 \
+    --num_frames 81 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --train_fps 16 \
+    --samples_per_file 8 \
+    --flush_frequency 8 \
+    --video_length_tolerance_range 5 \
+    --preprocess_task "t2v" 
@@ -0,0 +1,157 @@
+#!/bin/bash
+#SBATCH --job-name=t2v
+#SBATCH --partition=main
+#SBATCH --nodes=4
+#SBATCH --ntasks=4
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=dmd_t2v_output/t2v_%j.out
+#SBATCH --error=dmd_t2v_output/t2v_%j.err
+#SBATCH --exclusive
+
+# Basic Info
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export NCCL_DEBUG_SUBSYS=INIT,NET
+# 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 TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f"
+# export WANDB_API_KEY='your_wandb_api_key_here'
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+
+# Configs
+NUM_GPUS=8
+
+# Model paths for Self-Forcing DMD distillation with Wan2.2:
+# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers"  # Updated to Wan2.2
+# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model
+# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model
+GENERATOR_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Updated to Wan2.2
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"  # Teacher model
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
+
+# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
+DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing"
+# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset
+# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/"
+# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
+VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json"
+
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+training_args=(
+  --tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd  # Updated for Wan2.2
+  --output_dir "/mnt/sharefs/users/hao.zhang/SFwan2.2_t2v_finetune"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --max_train_steps 4000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 448  # Updated to match Wan2.2 config
+  --num_width 832   # Updated to match Wan2.2 config
+  --enable_gradient_checkpointing_type "full"
+  --simulate_generator_forward
+  # --log_visualization
+  --num_frames 81
+  --num_frame_per_block 3  # Frame generation block size for self-forcing
+  --enable_gradient_masking
+  --gradient_mask_last_n_frames 21
+  # --init_weights_from_safetensors /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/high/3k/
+  # --init_weights_from_safetensors_2 /mnt/sharefs/users/hao.zhang/wl/models/sf_ode_init_wan22_checkpoints/low/3k/
+)
+
+parallel_args=(
+  --num_gpus 32 # 64
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1 # 64
+  --hsdp_shard_dim 32
+)
+
+model_args=(
+  --model_path $GENERATOR_MODEL_PATH
+  --pretrained_model_name_or_path $GENERATOR_MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
+)
+
+dataset_args=(
+  --data_path "$DATA_DIR"
+  --dataloader_num_workers 4
+)
+
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 20
+  --validation_sampling_steps "4"
+  --validation_guidance_scale "6.0" # not used for dmd inference
+)
+
+optimizer_args=(
+  --learning_rate 1e-5
+  --mixed_precision "bf16"
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
+  --weight_decay 0.01
+  --betas '0.0,0.999'
+  --max_grad_norm 1.0
+)
+
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.0
+  --dit_precision "fp32"
+  --flow_shift 5
+  --seed 1000
+  --use_ema True
+  --ema_decay 0.99
+  --ema_start_step 100
+)
+
+dmd_args=(
+  --dmd_denoising_steps '1000,750,500,250'
+  --min_timestep_ratio 0.02
+  --max_timestep_ratio 0.98
+  --dfake_gen_update_ratio 5
+  --real_score_guidance_scale 3.0
+  --fake_score_learning_rate 8e-6
+  --fake_score_betas '0.0,0.999'
+  --warp_denoising_step
+)
+
+self_forcing_args=(
+  --independent_first_frame False  # Whether to treat first frame independently
+  --same_step_across_blocks True  # Whether to use same denoising step across all blocks
+  --last_step_only False  # Whether to only use the last denoising step
+  --context_noise 0  # Amount of noise to add during context caching (0 = no noise)
+)
+
+srun torchrun \
+--nnodes $SLURM_JOB_NUM_NODES \
+--nproc_per_node $NUM_GPUS \
+--node_rank $SLURM_PROCID \
+--rdzv_backend=c10d \
+--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/wan_self_forcing_distillation_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}" \
+    "${dmd_args[@]}" \
+    "${self_forcing_args[@]}"
@@ -39,15 +39,18 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
+OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir"checkpoints/wan_t2v_finetune"
+  --output_dir $OUTPUT_DIR
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -72,6 +75,8 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -91,7 +96,7 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 2e-6
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
  --weight_only_checkpointing_steps 500
@@ -134,4 +139,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -39,15 +39,18 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
+OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir "checkpoints/wan_t2v_finetune"
+  --output_dir "$OUTPUT_DIR"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -72,6 +75,8 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -91,7 +96,7 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 2e-6
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
  --weight_only_checkpointing_steps 500
@@ -134,4 +139,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -39,15 +39,18 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DATASET_FILE=your_validation_dataset_file
+OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd
-  --output_dir "checkpoints/wan_t2v_finetune"
+  --output_dir "$OUTPUT_DIR"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -72,6 +75,8 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -91,7 +96,7 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 1e-5
+  --learning_rate 2e-6
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
  --weight_only_checkpointing_steps 500
@@ -133,4 +138,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -1,3 +1,3 @@
 #!/bin/bash

-python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "FastVideo/Wan-Syn_77x448x832_600k" --repo_type "dataset"
@@ -40,15 +40,18 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR=your_data_dir
 VALIDATION_DIR=your_validation_path  #(example:validation_64.json)
+OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name Wan_distillation
-  --output_dir "your_output_dir"
+  --output_dir "$OUTPUT_DIR"
  --max_train_steps 4000
  --train_batch_size 1
  --train_sp_batch_size 1
@@ -73,6 +76,8 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -92,11 +97,11 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-5
+  --learning_rate 4e-6
  --lr_scheduler "cosine_with_min_lr"
  --min_lr_ratio 0.5
  --lr_warmup_steps 100
-  --fake_score_learning_rate 1e-5
+  --fake_score_learning_rate 2e-6
  --fake_score_lr_scheduler "cosine_with_min_lr"
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
@@ -141,4 +146,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -40,6 +40,8 @@ echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+FAKE_SCORE_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
@@ -73,6 +75,8 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -92,11 +96,11 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate 2e-5
+  --learning_rate 4e-6
  --lr_scheduler "cosine_with_min_lr"
  --min_lr_ratio 0.5
  --lr_warmup_steps 100
-  --fake_score_learning_rate 1e-5
+  --fake_score_learning_rate 2e-6
  --fake_score_lr_scheduler "cosine_with_min_lr"
  --mixed_precision "bf16"
  --training_state_checkpointing_steps 500
@@ -142,4 +146,4 @@ srun torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -14,26 +14,29 @@ export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
 # Configs
 NUM_GPUS=1
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers" 
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
+OUTPUT_DIR="checkpoints/wan_t2v_finetune"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]

 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_distill_dmd_VSA
-  --output_dir="checkpoints/wan_t2v_finetune"
-  --max_train_steps=4000
-  --train_batch_size=1
+  --output_dir "$OUTPUT_DIR"
+  --max_train_steps 4000
+  --train_batch_size 1
  --train_sp_batch_size 1
-  --gradient_accumulation_steps=1
+  --gradient_accumulation_steps 1
  --num_latent_t 31
  --num_height 704
  --num_width 1280
  --num_frames 121
  --enable_gradient_checkpointing_type "full"
-  --training_state_checkpointing_steps=500
-  --weight_only_checkpointing_steps=500
+  --training_state_checkpointing_steps 500
+  --weight_only_checkpointing_steps 500
 )

 # Parallel arguments
@@ -49,6 +52,8 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -68,8 +73,8 @@ validation_args=(

 # Optimizer arguments
 optimizer_args=(
-  --learning_rate=1e-5
-  --mixed_precision="bf16"
+  --learning_rate 2e-6
+  --mixed_precision "bf16"
  --weight_decay 0.01
  --max_grad_norm 1.0
 )
@@ -107,4 +112,4 @@ torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -14,6 +14,8 @@ export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN
 # Configs
 NUM_GPUS=1
 MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
+FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json"
 # export CUDA_VISIBLE_DEVICES=4,5
@@ -51,6 +53,8 @@ parallel_args=(
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
+  --real_score_model_path $REAL_SCORE_MODEL_PATH
+  --fake_score_model_path $FAKE_SCORE_MODEL_PATH
 )

 # Dataset arguments
@@ -109,4 +113,4 @@ torchrun \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}" \
-    "${dmd_args[@]}"
+    "${dmd_args[@]}"
@@ -1,3 +1,3 @@
 #!/bin/bash

-python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
+python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -21,4 +21,4 @@ torchrun --nproc_per_node=$GPU_NUM \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
-    --preprocess_task "t2v" 
+    --preprocess_task "t2v" 
@@ -0,0 +1,44 @@
+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-T2V-1.3B-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=True,
+        pin_cpu_memory=True, # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
+        distributed_executor_backend="ray",
+        # image_encoder_cpu_offload=False,
+    )
+
+    # 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)
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+    prompt2 = (
+        "A majestic lion strides across the golden savanna, its powerful frame "
+        "glistening under the warm afternoon sun. The tall grass ripples gently in "
+        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
+        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
+        "cinematic.")
+    video2 = generator.generate_video(prompt2, output_path=OUTPUT_PATH, save_video=True)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,44 @@
+# NOTE: This is still a work in progress, and the checkpoints are not released yet.
+
+from fastvideo import VideoGenerator, SamplingParam
+import json
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_self_forcing_causal_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(
+        "FastVideo/SFWan2.2-I2V-A14B-Preview-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
+        dit_precision="fp32",
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125],
+        # 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("FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers")
+    sampling_param.num_frames = 73
+    sampling_param.width = 832
+    sampling_param.height = 480
+    sampling_param.seed = 1000
+
+    with open("prompts/mixkit_i2v.jsonl", "r") as f:
+        prompt_image_pairs = json.load(f)
+
+    for prompt_image_pair in prompt_image_pairs:
+        prompt = prompt_image_pair["prompt"]
+        image_path = prompt_image_pair["image_path"]
+        _ = generator.generate_video(prompt, image_path=image_path, output_path=OUTPUT_PATH, save_video=True, sampling_param=sampling_param)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,43 @@
+# NOTE: This is still a work in progress, and the checkpoints are not released yet.
+
+from fastvideo import VideoGenerator
+
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_self_forcing_causal_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(
+        "rand0nmr/SFWan2.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,
+        dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125],
+        # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer
+        pin_cpu_memory=True,
+        init_weights_from_safetensors="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_inference_transformer/",
+        init_weights_from_safetensors_2="/mnt/sharefs/users/hao.zhang/wei/SFwan2.2_distill_self_forcing_release_cfg2/checkpoint-246_weight_only/generator_2_inference_transformer/",
+        num_frame_per_block=7,
+        # 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, num_frames=81)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,36 @@
+from fastvideo import VideoGenerator
+
+# from fastvideo.configs.sample import SamplingParam
+
+OUTPUT_PATH = "video_samples_wan2_1_Fun"
+OUTPUT_NAME = "wan2.1_test"
+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(
+        "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers",
+        # "alibaba-pai/Wan2.2-Fun-A14B-Control",
+        # 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 = "一位年轻女性穿着一件粉色的连衣裙，裙子上有白色的装饰和粉色的纽扣。她的头发是紫色的，头上戴着一个红色的大蝴蝶结，显得非常可爱和精致。她还戴着一个红色的领结，整体造型充满了少女感和活力。她的表情温柔，双手轻轻交叉放在身前，姿态优雅。背景是简单的灰色，没有任何多余的装饰，使得人物更加突出。她的妆容清淡自然，突显了她的清新气质。整体画面给人一种甜美、梦幻的感觉，仿佛置身于童话世界中。"
+    negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+    # prompt                  = "A young woman with beautiful, clear eyes and blonde hair stands in the forest, wearing a white dress and a crown. Her expression is serene, reminiscent of a movie star, with fair and youthful skin. Her brown long hair flows in the wind. The video quality is very high, with a clear view. High quality, masterpiece, best quality, high resolution, ultra-fine, fantastical."
+    # negative_prompt         = "Twisted body, limb deformities, text captions, comic, static, ugly, error, messy code."
+    image_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/8.png"
+    control_video_path = "https://pai-aigc-photog.oss-cn-hangzhou.aliyuncs.com/wan_fun/asset_Wan2_2/v1.0/pose.mp4"
+
+    video = generator.generate_video(prompt, negative_prompt=negative_prompt, image_path=image_path, video_path=control_video_path, output_path=OUTPUT_PATH, output_video_name=OUTPUT_NAME, save_video=True)
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,56 @@
+# FastVideo Gradio Local 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/local/gradio_local_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, *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()`)
+
+## 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
+  - 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
+- **Negative Prompt**: Specify what you don't want to see in the video
+- **Seed**: Control randomness for reproducible results
@@ -1,87 +1,20 @@
 import argparse
 import os
-import requests
 import base64
 import time

 import gradio as gr
-
+from fastvideo.entrypoints.video_generator import VideoGenerator
 from fastvideo.configs.sample.base import SamplingParam
+from copy import deepcopy


 MODEL_PATH_MAPPING = {
    "FastWan2.1-T2V-1.3B": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
-    "FastWan2.2-TI2V-5B-FullAttn": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+    # "FastWan2.2-TI2V-5B-FullAttn": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
 }

-
-class RayServeClient:
-    def __init__(self, backend_url: str):
-        self.backend_url = backend_url
-        self.session = requests.Session()
-    
-    def check_health(self) -> bool:
-        try:
-            response = self.session.get(f"{self.backend_url}/health", timeout=5)
-            return response.status_code == 200
-        except requests.exceptions.RequestException:
-            return False
-    
-    def generate_video(self, request_data: dict) -> dict:
-        start_time = time.time()
-        
-        try:
-            headers = {"Content-Type": "application/json"}
-            
-            response = self.session.post(
-                f"{self.backend_url}/generate_video",
-                json=request_data,
-                headers=headers,
-                timeout=300
-            )
-            
-            round_trip_time = time.time() - start_time
-            
-            if response.status_code == 200:
-                result = response.json()
-                backend_total = result.get("total_time", 0)
-                network_time = round_trip_time - backend_total
-                result["network_time"] = network_time
-                return result
-            else:
-                return {"success": False, "error_message": f"HTTP {response.status_code}: {response.text}"}
-                
-        except requests.exceptions.RequestException as e:
-            return {"success": False, "error_message": f"Request failed: {str(e)}"}
-
-
-def save_video_from_base64(video_data: str, output_dir: str, prompt: str) -> str:
-    if not video_data:
-        return None
-    
-    try:
-        if video_data.startswith('data:video/'):
-            video_data = video_data.split(',')[1]
-        
-        video_bytes = base64.b64decode(video_data)
-        
-        safe_prompt = prompt[:50].replace(' ', '_').replace('/', '_').replace('\\', '_')
-        video_filename = f"{safe_prompt}.mp4"
-        video_path = os.path.join(output_dir, video_filename)
-        
-        os.makedirs(output_dir, exist_ok=True)
-        
-        with open(video_path, 'wb') as f:
-            f.write(video_bytes)
-        
-        return video_path
-        
-    except Exception as e:
-        print(f"Failed to save video: {e}")
-        return None
-
-
-def create_timing_display(inference_time, encoding_time, network_time, total_time, stage_execution_times, num_frames):
+def create_timing_display(inference_time, total_time, stage_execution_times, num_frames):
    dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A"
    
    timing_html = f"""
@@ -101,12 +34,12 @@ def create_timing_display(inference_time, encoding_time, network_time, total_tim
            <div class="timing-card">
                <div style="font-size: 20px;">🎬</div>
                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Video Encoding</div>
-                <div style="font-size: 16px; color: #dc2626;">{encoding_time:.2f}s</div>
+                <div style="font-size: 16px; color: #dc2626;">N/A</div>
            </div>
            <div class="timing-card">
                <div style="font-size: 20px;">🌐</div>
                <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Network Transfer</div>
-                <div style="font-size: 16px; color: #059669;">{network_time:.2f}s</div>
+                <div style="font-size: 16px; color: #059669;">N/A</div>
            </div>
            <div class="timing-card">
                <div style="font-size: 20px;">📊</div>
@@ -124,7 +57,12 @@ def create_timing_display(inference_time, encoding_time, network_time, total_tim
        </div>"""
    
    return timing_html + "</div>"
-
+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 load_example_prompts():
    def contains_chinese(text):
@@ -144,7 +82,7 @@ def load_example_prompts():
            print(f"Warning: Could not read {filepath}: {e}")
        return prompts, labels
    
-    examples, example_labels = load_from_file("prompts/prompts_final.txt")
+    examples, example_labels = load_from_file("examples/inference/gradio/local/prompts_final.txt")
    
    if not examples:
        examples = ["A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background."]
@@ -153,81 +91,71 @@ def load_example_prompts():
    return examples, example_labels


-def create_gradio_interface(backend_url: str, default_params: dict[str, SamplingParam]):
-    
-    client = RayServeClient(backend_url)
-    
+def create_gradio_interface(default_params: dict[str, SamplingParam], generators: dict[str, VideoGenerator]):
    def generate_video(
        prompt, negative_prompt, use_negative_prompt, seed, guidance_scale,
        num_frames, height, width, randomize_seed, model_selection, progress
    ):
-        if not client.check_health():
-            return None, f"Backend is not available. Please check if Ray Serve is running at {backend_url}", ""
-        
-        # Validate dimensions
-        max_pixels = 720 * 1280
-        if height * width > max_pixels:
-            return None, f"Video dimensions too large. Maximum: 720x1280 pixels", ""
-        
-        if progress:
-            progress(0.1, desc="Checking backend health...")
-        
-        request_data = {
-            "prompt": prompt,
-            "negative_prompt": negative_prompt,
-            "use_negative_prompt": use_negative_prompt,
-            "seed": seed,
-            "guidance_scale": guidance_scale,
-            "num_frames": num_frames,
-            "height": height,
-            "width": width,
-            "randomize_seed": randomize_seed,
-            "return_frames": False,
-            "image_path": None,
-            "model_path": MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers")
-        }
-        
-        if progress:
-            progress(0.4, desc="Generating video...")
-        
-        response = client.generate_video(request_data)
-        
-        if progress:
-            progress(0.8, desc="Processing response...")
-        
-        if response.get("success", False):
-            video_data = response.get("video_data", "")
-            used_seed = response.get("seed", seed)
-            inference_time = response.get("inference_time", 0.0)
-            encoding_time = response.get("encoding_time", 0.0)
-            total_time = response.get("total_time", 0.0)
-            network_time = response.get("network_time", 0.0)
-            stage_execution_times = response.get("stage_execution_times", [])
+        model_path = MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers")
+        setup_model_environment(model_path)
+        try:
+            if progress:
+                progress(0.1, desc="Loading model for local inference...")
            
-            timing_details = create_timing_display(
-                inference_time, encoding_time, network_time, total_time, 
-                stage_execution_times, num_frames
-            )
-            
-            if video_data:
-                if progress:
-                    progress(0.9, desc="Saving video...")
-                
-                video_path = save_video_from_base64(video_data, "outputs", prompt)
-                
-                if progress:
-                    progress(1.0, desc="Generation complete!")
-                
-                if video_path and os.path.exists(video_path):
-                    return video_path, used_seed, timing_details
-                else:
-                    return None, "Failed to save video", ""
+            generator = generators[model_path]
+            params = deepcopy(default_params[model_path])
+            total_start_time = time.time()
+            if progress:
+                progress(0.2, desc="Configuring parameters...")
+
+            params.prompt = prompt
+            params.seed = int(seed)
+            params.guidance_scale = guidance_scale
+            params.num_frames = int(num_frames)
+            params.height = int(height)
+            params.width = int(width)
+
+            if randomize_seed:
+                params.seed = torch.randint(0, 1000000, (1, )).item()
+
+            if use_negative_prompt and negative_prompt:
+                params.negative_prompt = negative_prompt
            else:
-                return None, "No video data received from backend", ""
-        else:
-            error_msg = response.get("error_message", "Unknown error occurred")
-            return None, f"Generation failed: {error_msg}", ""
-    
+                params.negative_prompt = default_params[model_path].negative_prompt
+
+            if progress:
+                progress(0.4, desc="Generating video locally...")
+
+            output_dir = "outputs/"
+            os.makedirs(output_dir, exist_ok=True)
+            start_time = time.time()
+            result = generator.generate_video(prompt=prompt, sampling_param=params, save_video=True, return_frames=False)
+            inference_time = time.time() - start_time
+            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 = []
+            total_time = time.time() - total_start_time
+            timing_details=create_timing_display(inference_time=inference_time, total_time=total_time, stage_execution_times=stage_execution_times, num_frames=params.num_frames)
+            safe_prompt = params.prompt[:100].replace(' ', '_').replace('/', '_').replace('\\', '_')
+            video_filename = f"{params.prompt[:100]}.mp4"
+            output_path = os.path.join(output_dir, video_filename)
+
+            if progress:
+                progress(1.0, desc="Generation complete!")
+
+            return output_path, params.seed, timing_details
+
+        except Exception as e:
+            print(f"An error occurred during local generation: {e}")
+            return None, f"Generation failed: {str(e)}", ""
+
    examples, example_labels = load_example_prompts()
    
    theme = gr.themes.Base().set(
@@ -261,7 +189,8 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling
    initial_values = get_default_values("FastWan2.1-T2V-1.3B")
    
    with gr.Blocks(title="FastWan", theme=theme) as demo:
-        gr.Image("assets/logos/logo.svg", show_label=False, container=False, height=80)
+        gr.Image("assets/full.svg", show_label=False, container=False, height=80)
+
        gr.HTML("""
        <div style="text-align: center; margin-bottom: 10px;">
            <p style="font-size: 18px;"> Make Video Generation Go Blurrrrrrr </p>
@@ -522,7 +451,7 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling
        
        gr.HTML("""
        <div style="text-align: center; margin-top: 10px; margin-bottom: 15px;">
-            <p style="font-size: 16px; margin: 0;">The compute for this demo is generously provided by <a href="https://www.gmicloud.ai/" target="_blank">GMI Cloud</a>. Note that this demo is meant to showcase FastWan's quality and that under a large number of requests, generation speed may be affected. We are also rate-limiting users to 3 requests per minute.</p>
+            <p style="font-size: 16px; margin: 0;">Note that this demo is meant to showcase FastWan's quality and that under a large number of requests, generation speed may be affected.</p>
        </div>
        """)
        
@@ -569,7 +498,6 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling
                prompt, negative_prompt, use_negative_prompt, seed, guidance_scale, 
                num_frames, height, width, randomize_seed, model_selection, progress
            )
-            
            if result_path and os.path.exists(result_path):
                return (
                    result_path, 
@@ -607,28 +535,25 @@ def create_gradio_interface(backend_url: str, default_params: dict[str, Sampling


 def main():
-    parser = argparse.ArgumentParser(description="FastVideo Gradio Frontend")
-    parser.add_argument("--backend_url", type=str, default="http://localhost:8000",
-                        help="URL of the Ray Serve backend")
+    parser = argparse.ArgumentParser(description="FastVideo Gradio Local Demo")
    parser.add_argument("--t2v_model_paths", type=str,
-                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.1-T2V-14B-Diffusers",
+                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
                        help="Comma separated list of paths to the T2V model(s)")
    parser.add_argument("--host", type=str, default="0.0.0.0",
                        help="Host to bind to")
    parser.add_argument("--port", type=int, default=7860,
                        help="Port to bind to")
-    
    args = parser.parse_args()
-    
+    generators = {}
    default_params = {}
    model_paths = args.t2v_model_paths.split(",")
    for model_path in model_paths:
+        print(f"Loading model: {model_path}")
+        setup_model_environment(model_path)
+        generators[model_path] = VideoGenerator.from_pretrained(model_path)
        default_params[model_path] = SamplingParam.from_pretrained(model_path)
-    
-    demo = create_gradio_interface(args.backend_url, default_params)
-    
+    demo = create_gradio_interface(default_params, generators)
    print(f"Starting Gradio frontend at http://{args.host}:{args.port}")
-    print(f"Backend URL: {args.backend_url}")
    print(f"T2V Models: {args.t2v_model_paths}")
    
    from fastapi import FastAPI, Request, HTTPException
@@ -637,10 +562,10 @@ def main():
    
    app = FastAPI()
    
-    @app.get("/logo.svg")
+    @app.get("/logo.png")
    def get_logo():
        return FileResponse(
-            "assets/logos/logo.svg", 
+            "assets/full.svg",
            media_type="image/svg+xml",
            headers={
                "Cache-Control": "public, max-age=3600",
@@ -650,7 +575,7 @@ def main():
    
    @app.get("/favicon.ico")
    def get_favicon():
-        favicon_path = "assets/logos/icon_simple.svg"
+        favicon_path = "assets/icon-simple.svg"
        
        if os.path.exists(favicon_path):
            return FileResponse(
@@ -683,7 +608,7 @@ def main():
            <meta property="og:url" content="{base_url}/">
            <meta property="og:title" content="FastWan">
            <meta property="og:description" content="Make video generation go blurrrrrrr">
-            <meta property="og:image" content="{base_url}/logo.svg">
+            <meta property="og:image" content="{base_url}/logo.png">
            <meta property="og:image:width" content="1200">
            <meta property="og:image:height" content="630">
            <meta property="og:site_name" content="FastWan">
@@ -692,7 +617,7 @@ def main():
            <meta property="twitter:url" content="{base_url}/">
            <meta property="twitter:title" content="FastWan">
            <meta property="twitter:description" content="Make video generation go blurrrrrrr">
-            <meta property="twitter:image" content="{base_url}/logo.svg">
+            <meta property="twitter:image" content="{base_url}/logo.png">
            <link rel="icon" type="image/png" sizes="32x32" href="/favicon.ico">
            <link rel="icon" type="image/png" sizes="16x16" href="/favicon.ico">
            <link rel="apple-touch-icon" href="/favicon.ico">
@@ -727,4 +652,5 @@ def main():


 if __name__ == "__main__":
+    
    main() 
@@ -0,0 +1,11 @@
+A dynamic shot of a sleek black motorcycle accelerating down an empty highway at sunset. The bike's engine roars as it gains speed, smoke trailing from the tires. The rider, wearing a black leather jacket and helmet, leans forward with determination, gripping the handlebars tightly. The camera follows the motorcycle from a distance, capturing the dust kicked up behind it, then zooms in to show the intense focus on the rider's face. The background showcases the endless road stretching into the horizon with vibrant orange and pink hues of the setting sun. Medium shot transitioning to close-up.
+A Jedi Master Yoda, recognizable by his green skin, large ears, and wise wrinkles, is performing on a small stage, strumming a guitar with great concentration. Yoda wears a casual robe and sits on a stool, his eyes closed as he plays, fully immersed in the music. The stage is dimly lit with spotlights highlighting Yoda, creating a mystical atmosphere. The background shows a live audience watching intently. Medium close-up shot focusing on Yoda's expressive face and hands moving gracefully over the guitar strings.
+A cute, fluffy panda bear is preparing a meal in a cozy, modern kitchen. The panda is standing at a wooden countertop, wearing a white chef’s hat and apron. It skillfully stirs a pot on the stove with one hand while holding a spatula in the other. The kitchen is well-lit, with appliances and cabinets in pastel colors, creating a warm and inviting atmosphere. The panda moves gracefully, with a focused and determined expression, as steam rises from the pot. Medium shot focusing on the panda’s actions at the stove.
+In a futuristic Tokyo rooftop during a heavy rainstorm, a robotic DJ stands behind a turntable, spinning vinyl records in a cyberpunk night setting. The robot has metallic, sleek body parts with glowing blue LED lights, and it moves gracefully with the beat. Raindrops create a shimmering effect as they hit the ground and the DJ. The surrounding environment features neon signs, towering skyscrapers, and a dark, misty atmosphere. The camera starts with a wide shot of the city skyline before zooming in on the DJ performing. Sci-fi, fantasy.
+A realistic animated scene featuring a polar bear playing a guitar. The polar bear is standing upright, wearing a cozy fur vest and fingerless gloves. It holds the guitar with both hands, strumming the strings with one hand while plucking them with the other, showcasing natural, fluid motions. The polar bear's expressive face shows concentration and joy as it plays. The background is a snowy Arctic landscape with icebergs and a clear blue sky. The scene captures the bear from a mid-shot angle, focusing on its interaction with the guitar.
+The scene opens to a breathtaking view of a tranquil ocean horizon at dusk, displaying a vibrant tapestry of oranges, pinks, and purples as the sun sets. In the foreground, tall, swaying palm trees frame the scene, their silhouettes stark against the colorful sky. The ocean itself shimmers with reflections of the sunset, creating a peaceful, almost ethereal atmosphere. A small boat can be seen in the distance, centered on the horizon, adding a sense of scale and solitude to the scene. The waves gently lap the shore, creating faint patterns on the sandy beach, which stretches across the foreground. Above, the sky is dotted with scattered clouds that catch the last light of the day, enhancing the drama and beauty of the scene. The overall mood is serene and contemplative, capturing a perfect moment of nature’s grandeur.
+A large, modern semi-truck accelerating down an empty highway, gaining speed with each second. The truck's powerful engine roars as it moves forward, smoke billowing from the tires. The camera starts from a wide shot, capturing the truck in the distance, then smoothly zooms in to follow the vehicle as it speeds up. The truck's headlights illuminate the road ahead, casting a bright glow. The truck driver can be seen through the windshield, focused and determined. The background shows the vast openness of the highway stretching into the horizon under a clear blue sky. Medium to close-up shots of the truck as it accelerates.
+Soft blue light pulses from the blade’s rune-etched hilt, illuminating nearby moss-covered roots and ferns. The surrounding trees are tall and gnarled, their branches curling like claws overhead. Fog swirls gently at ground level, parting slightly as a figure in a cloak approaches from the distance. Medium shot slowly zooming toward the sword, emphasizing its mystical aura.
+The video opens with a tranquil scene in the heart of a dense forest, emphasizing two large, textured tree trunks in the foreground framing the view. Sunlight filters through the canopy above, casting intricate patterns of light and shadow on the trees and the ground. Between the tree trunks, a clear view of a calm, muddy river unfolds, its surface shimmering under the gentle sunlight. The riverbank is decorated with a variety of small bushes and vibrant foliage, subtly transitioning into the deep greens of tall, leafy plants. In the background, the dense forest looms, filled with dark, towering trees, their branches intertwining to form an intricate canopy. The scene is bathed in the soft glow of the sun, creating a serene and picturesque setting. Occasional sunbeams pierce through the foliage, adding a magical aura to the landscape. The vibrant reds and oranges of the smaller plants add contrast, bringing warmth to the earthy tones of the scenery. Overall, this harmonious blend of natural elements creates a peaceful and idyllic forest setting.
+A lone figure stands on a large, moss-covered rock, surrounded by the soft rush of a nearby stream. The figure is wearing white sneakers and shorts, with a plaid shirt that hangs loosely in the breeze. The lighting creates dramatic shadows, enhancing the textures of the rock and the subtle movement of the water below. In the background, a waterfall cascades into the stream, completing this tranquil and serene nature scene.
+In an industrial setting, a person leans casually against a railing, exuding a sense of confidence and composure. They are wearing a striking outfit, consisting of a vibrant, patterned jacket over a simple white crop top, creating a bold contrast. The atmosphere is infused with warm, ambient lighting that casts soft shadows on the concrete walls and metallic surfaces. Intricate wiring and pipes form an intricate backdrop, enhancing the urban aesthetic. Their relaxed posture and direct, engaging gaze suggest a sense of ease in this industrial environment. This scene encapsulates a blend of modern fashion and gritty, urban architecture, creating a visually compelling narrative.
@@ -0,0 +1,774 @@
+import argparse
+import os
+import requests
+import base64
+import time
+import json
+from pathlib import Path
+import tempfile
+from io import BytesIO
+
+import gradio as gr
+from PIL import Image
+
+from fastvideo.configs.sample.base import SamplingParam
+
+
+MODEL_PATH_MAPPING = {
+    "FastWan2.1-T2V-1.3B": "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
+    "FastWan2.2-TI2V-5B-FullAttn": "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+    "CausalWan2.2-I2V-A14B-Preview": "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers",
+}
+
+
+class RayServeClient:
+    def __init__(self, backend_url: str):
+        self.backend_url = backend_url
+        self.session = requests.Session()
+    
+    def check_health(self) -> bool:
+        try:
+            response = self.session.get(f"{self.backend_url}/health", timeout=5)
+            return response.status_code == 200
+        except requests.exceptions.RequestException:
+            return False
+    
+    def generate_video(self, request_data: dict) -> dict:
+        start_time = time.time()
+        
+        try:
+            headers = {"Content-Type": "application/json"}
+            
+            response = self.session.post(
+                f"{self.backend_url}/generate_video",
+                json=request_data,
+                headers=headers,
+                timeout=900  # 15 minutes timeout for longer video generation
+            )
+            
+            round_trip_time = time.time() - start_time
+            
+            if response.status_code == 200:
+                result = response.json()
+                backend_total = result.get("total_time", 0)
+                network_time = round_trip_time - backend_total
+                result["network_time"] = network_time
+                return result
+            else:
+                return {"success": False, "error_message": f"HTTP {response.status_code}: {response.text}"}
+                
+        except requests.exceptions.RequestException as e:
+            return {"success": False, "error_message": f"Request failed: {str(e)}"}
+
+
+def save_video_from_base64(video_data: str, output_dir: str, prompt: str) -> str:
+    if not video_data:
+        return None
+    
+    try:
+        if video_data.startswith('data:video/'):
+            video_data = video_data.split(',')[1]
+        
+        video_bytes = base64.b64decode(video_data)
+        
+        safe_prompt = prompt[:50].replace(' ', '_').replace('/', '_').replace('\\', '_')
+        video_filename = f"{safe_prompt}.mp4"
+        video_path = os.path.join(output_dir, video_filename)
+        
+        os.makedirs(output_dir, exist_ok=True)
+        
+        with open(video_path, 'wb') as f:
+            f.write(video_bytes)
+        
+        return video_path
+        
+    except Exception as e:
+        print(f"Failed to save video: {e}")
+        return None
+
+
+def encode_image_to_base64(image_input) -> str:
+    """Encode an image file path or in-memory image to a base64 string."""
+    if image_input is None:
+        return None
+    
+    mime_types = {
+        '.jpg': 'image/jpeg',
+        '.jpeg': 'image/jpeg',
+        '.png': 'image/png',
+        '.gif': 'image/gif',
+        '.webp': 'image/webp',
+    }
+    
+    try:
+        if isinstance(image_input, str):
+            if not os.path.exists(image_input):
+                return None
+            
+            with open(image_input, 'rb') as f:
+                image_bytes = f.read()
+            
+            ext = os.path.splitext(image_input)[1].lower()
+            mime_type = mime_types.get(ext, 'image/jpeg')
+        elif isinstance(image_input, Image.Image):
+            buffer = BytesIO()
+            image_to_save = image_input.convert("RGB")
+            image_to_save.save(buffer, format="PNG")
+            image_bytes = buffer.getvalue()
+            mime_type = 'image/png'
+        else:
+            return None
+        
+        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
+        return f"data:{mime_type};base64,{image_base64}"
+        
+    except Exception as e:
+        print(f"Failed to encode image: {e}")
+        return None
+
+
+# def create_timing_display(inference_time, encoding_time, network_time, total_time, stage_execution_times, num_frames):
+#     dit_denoising_time = f"{stage_execution_times[5]:.2f}s" if len(stage_execution_times) > 5 else "N/A"
+#     
+#     timing_html = f"""
+#     <div style="margin: 10px 0;">
+#         <h3 style="text-align: center; margin-bottom: 10px;">⏱️ Timing Breakdown</h3>
+#         <div style="display: grid; grid-template-columns: repeat(5, 1fr); gap: 10px; margin-bottom: 10px;">
+#             <div class="timing-card timing-card-highlight">
+#                 <div style="font-size: 20px;">🚀</div>
+#                 <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">DiT Denoising</div>
+#                 <div style="font-size: 16px; color: #ffa200; font-weight: bold;">{dit_denoising_time}</div>
+#             </div>
+#             <div class="timing-card">
+#                 <div style="font-size: 20px;">🧠</div>
+#                 <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">E2E (w. vae/text encoder)</div>
+#                 <div style="font-size: 16px; color: #2563eb;">{inference_time:.2f}s</div>
+#             </div>
+#             <div class="timing-card">
+#                 <div style="font-size: 20px;">🎬</div>
+#                 <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Video Encoding</div>
+#                 <div style="font-size: 16px; color: #dc2626;">{encoding_time:.2f}s</div>
+#             </div>
+#             <div class="timing-card">
+#                 <div style="font-size: 20px;">🌐</div>
+#                 <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Network Transfer</div>
+#                 <div style="font-size: 16px; color: #059669;">{network_time:.2f}s</div>
+#             </div>
+#             <div class="timing-card">
+#                 <div style="font-size: 20px;">📊</div>
+#                 <div style="font-weight: bold; margin: 3px 0; font-size: 14px;">Total Processing</div>
+#                 <div style="font-size: 18px; color: #0277bd;">{total_time:.2f}s</div>
+#             </div>
+#         </div>"""
+#     
+#     if inference_time > 0:
+#         fps = num_frames / inference_time
+#         timing_html += f"""
+#         <div class="performance-card" style="margin-top: 15px;">
+#             <span style="font-weight: bold;">Generation Speed: </span>
+#             <span style="font-size: 18px; color: #6366f1; font-weight: bold;">{fps:.1f} frames/second</span>
+#         </div>"""
+#     
+#     return timing_html + "</div>"
+
+
+def load_example_prompts():
+    def contains_chinese(text):
+        return any('\u4e00' <= char <= '\u9fff' for char in text)
+    
+    def load_from_file(filepath):
+        prompts, labels = [], []
+        try:
+            with open(filepath, "r", encoding='utf-8') as f:
+                for line in f:
+                    line = line.strip()
+                    if line and not contains_chinese(line):
+                        label = line[:100] + "..." if len(line) > 100 else line
+                        labels.append(label)
+                        prompts.append(line)
+        except Exception as e:
+            print(f"Warning: Could not read {filepath}: {e}")
+        return prompts, labels
+    
+    # Load prompts from prompts.txt
+    examples, example_labels = load_from_file("examples/inference/gradio/serving/prompts.txt")
+    
+    if not examples:
+        examples = ["A crowded rooftop bar buzzes with energy, the city skyline twinkling like a field of stars in the background."]
+        example_labels = ["Crowded rooftop bar at night"]
+    
+    # Load image mappings from JSON file
+    prompt_to_image = {}
+    # Try to find the JSON file relative to project root
+    possible_json_paths = [
+        Path("prompts/mixkit_i2v.jsonl"),
+        Path(__file__).parent.parent.parent.parent / "prompts" / "mixkit_i2v.jsonl",
+    ]
+    json_path = None
+    for path in possible_json_paths:
+        if path.exists():
+            json_path = path
+            break
+    
+    if json_path and json_path.exists():
+        try:
+            with open(json_path, "r", encoding='utf-8') as f:
+                data = json.load(f)
+                # Get the project root directory (parent of prompts directory)
+                project_root = json_path.parent.parent
+                for item in data:
+                    prompt_text = item.get("prompt", "").strip()
+                    image_path = item.get("image_path", "")
+                    if prompt_text and image_path:
+                        # Resolve image path relative to project root
+                        full_image_path = project_root / image_path
+                        if full_image_path.exists():
+                            prompt_to_image[prompt_text] = str(full_image_path.absolute())
+        except Exception as e:
+            print(f"Warning: Could not load image mappings from {json_path}: {e}")
+    
+    # Create image paths list matching the prompts
+    example_images = []
+    for prompt in examples:
+        # Try exact match first
+        image_path = prompt_to_image.get(prompt)
+        if not image_path:
+            # Try fuzzy match (case-insensitive, whitespace normalized)
+            normalized_prompt = " ".join(prompt.split())
+            for json_prompt, img_path in prompt_to_image.items():
+                normalized_json = " ".join(json_prompt.split())
+                if normalized_prompt.lower() == normalized_json.lower():
+                    image_path = img_path
+                    break
+        example_images.append(image_path if image_path and os.path.exists(image_path) else None)
+    
+    return examples, example_labels, example_images
+
+
+def create_gradio_interface(backend_url: str, default_params: dict[str, SamplingParam]):
+    
+    client = RayServeClient(backend_url)
+    
+    def is_i2v_model(model_name: str) -> bool:
+        """Check if the model is an I2V model."""
+        return "I2V" in model_name
+    
+    def generate_video(
+        prompt, negative_prompt, use_negative_prompt, guidance_scale,
+        num_frames, height, width, model_selection, input_image, progress
+    ):
+        # Use default seed value (randomize_seed disabled)
+        seed = 1000
+        randomize_seed = False
+        if not client.check_health():
+            return None, f"Backend is not available. Please check if Ray Serve is running at {backend_url}", ""
+        
+        # Check if I2V model requires an image
+        if is_i2v_model(model_selection) and not input_image:
+            return None, "I2V models require an input image. Please upload an image.", ""
+        
+        # Validate dimensions
+        max_pixels = 720 * 1280
+        if height * width > max_pixels:
+            return None, f"Video dimensions too large. Maximum: 720x1280 pixels", ""
+        
+        if progress:
+            progress(0.1, desc="Checking backend health...")
+        
+        # Encode image if provided
+        image_data = None
+        if input_image:
+            if progress:
+                progress(0.2, desc="Encoding input image...")
+            image_data = encode_image_to_base64(input_image)
+            if not image_data:
+                return None, "Failed to encode input image", ""
+        
+        request_data = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "use_negative_prompt": use_negative_prompt,
+            "seed": seed,
+            "guidance_scale": guidance_scale,
+            "num_frames": num_frames,
+            "height": height,
+            "width": width,
+            "randomize_seed": randomize_seed,
+            "return_frames": False,
+            "image_data": image_data,
+            "model_path": MODEL_PATH_MAPPING.get(model_selection, "FastVideo/FastWan2.1-T2V-1.3B-Diffusers")
+        }
+        
+        if progress:
+            progress(0.4, desc="Generating video...")
+        
+        response = client.generate_video(request_data)
+        
+        if progress:
+            progress(0.8, desc="Processing response...")
+        
+        if response.get("success", False):
+            video_data = response.get("video_data", "")
+            used_seed = response.get("seed", seed)
+            # inference_time = response.get("inference_time", 0.0)
+            # encoding_time = response.get("encoding_time", 0.0)
+            # total_time = response.get("total_time", 0.0)
+            # network_time = response.get("network_time", 0.0)
+            # stage_execution_times = response.get("stage_execution_times", [])
+            
+            # timing_details = create_timing_display(
+            #     inference_time, encoding_time, network_time, total_time, 
+            #     stage_execution_times, num_frames
+            # )
+            
+            if video_data:
+                if progress:
+                    progress(0.9, desc="Saving video...")
+                
+                video_path = save_video_from_base64(video_data, "outputs", prompt)
+                
+                if progress:
+                    progress(1.0, desc="Generation complete!")
+                
+                if video_path and os.path.exists(video_path):
+                    return video_path, used_seed, ""
+                else:
+                    return None, "Failed to save video", ""
+            else:
+                return None, "No video data received from backend", ""
+        else:
+            error_msg = response.get("error_message", "Unknown error occurred")
+            return None, f"Generation failed: {error_msg}", ""
+    
+    examples, example_labels, example_images = load_example_prompts()
+    
+    theme = gr.themes.Base().set(
+        button_primary_background_fill="#2563eb",
+        button_primary_background_fill_hover="#1d4ed8",
+        button_primary_text_color="white",
+        slider_color="#2563eb",
+        checkbox_background_color_selected="#2563eb",
+    )
+    
+    def get_default_values(model_name):
+        # model_path = MODEL_PATH_MAPPING.get(model_name)
+        # if model_path and model_path in default_params:
+        #     params = default_params[model_path]
+        #     return {
+        #         'height': params.height,
+        #         'width': params.width,
+        #         'num_frames': params.num_frames,
+        #         'guidance_scale': params.guidance_scale,
+        #     }
+        
+        return {
+            'height': 480,
+            'width': 832,
+            'num_frames': 73,
+        }
+    
+    # Get available models based on what's loaded
+    available_models = []
+    for model_name, model_path in MODEL_PATH_MAPPING.items():
+        if model_path in default_params:
+            available_models.append(model_name)
+    
+    # Select first available model as default
+    default_model = available_models[0] if available_models else "FastWan2.1-T2V-1.3B"
+    initial_values = get_default_values(default_model)
+    initial_show_image = is_i2v_model(default_model)
+    
+    with gr.Blocks(title="CausalWan", theme=theme) as demo:
+        gr.Image("assets/logos/logo.svg", show_label=False, container=False, height=80)
+        gr.HTML("""
+        <div style="text-align: center; margin-bottom: 10px;">
+            <p style="font-size: 18px;"> Make Video Generation Go Blurrrrrrr </p>
+            <p style="font-size: 18px;"> <a href="https://github.com/hao-ai-lab/FastVideo/tree/main" target="_blank">Code</a> | <a href="https://hao-ai-lab.github.io/blogs/fastvideo_causalwan_preview/" target="_blank">Blog</a> | <a href="https://hao-ai-lab.github.io/FastVideo/" target="_blank">Docs</a>  </p>
+        </div>
+        """)
+        
+        with gr.Accordion("🎥 What Is FastVideo?", open=False):
+            gr.HTML("""
+            <div style="padding: 20px; line-height: 1.6;">
+                <p style="font-size: 16px; margin-bottom: 15px;">
+                    FastVideo is an inference and post-training framework for diffusion models. It features an end-to-end unified pipeline for accelerating diffusion models, starting from data preprocessing to model training, finetuning, distillation, and inference. FastVideo is designed to be modular and extensible, allowing users to easily add new optimizations and techniques. Whether it is training-free optimizations or post-training optimizations, FastVideo has you covered.
+                </p>
+            </div>
+            """)
+        
+        with gr.Row():
+            model_selection = gr.Dropdown(
+                choices=available_models,
+                value=default_model,
+                label="Select Model",
+                interactive=True
+            )
+
+        with gr.Row():
+            example_dropdown = gr.Dropdown(
+                choices=example_labels,
+                label="Example Prompts",
+                value=None,
+                interactive=True,
+                allow_custom_value=False
+            )
+        
+        with gr.Row():
+            with gr.Column(scale=6):
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=3,
+                    placeholder="Describe your scene...",
+                    container=False,
+                    lines=3,
+                    autofocus=True,
+                )
+            with gr.Column(scale=1, min_width=120, elem_classes="center-button"):
+                run_button = gr.Button("Run", variant="primary", size="lg")
+        
+        with gr.Row():
+            with gr.Column():
+                error_output = gr.Text(label="Error", visible=False)
+                # timing_display = gr.Markdown(label="Timing Breakdown", visible=False)
+
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=1):
+                with gr.Tabs():
+                    with gr.Tab("Input Image", visible=initial_show_image) as image_tab:
+                        gr.Markdown("**Please make sure you upload a 480x832 image**")
+                        input_image = gr.Image(
+                            label="",
+                            type="pil",
+                            height=400,
+                        )
+                    
+                    with gr.Tab("Advanced Options"):
+                        with gr.Group():
+                            with gr.Row():
+                                height = gr.Number(
+                                    label="Height",
+                                    value=initial_values['height'],
+                                    interactive=False,
+                                    container=True
+                                )
+                                width = gr.Number(
+                                    label="Width",
+                                    value=initial_values['width'],
+                                    interactive=False,
+                                    container=True
+                                )
+                            
+                            with gr.Row():
+                                num_frames = gr.Number(
+                                    label="Number of Frames",
+                                    value=initial_values['num_frames'],
+                                    interactive=False,
+                                    container=True
+                                )
+                                guidance_scale = gr.Number(
+                                    label="Guidance Scale",
+                                    value=1.0,
+                                    interactive=False,
+                                    container=True
+                                )
+                            
+                            with gr.Row():
+                                use_negative_prompt = gr.Checkbox(
+                                    label="Use negative prompt", value=False)
+                                negative_prompt = gr.Text(
+                                    label="Negative prompt",
+                                    max_lines=3,
+                                    lines=3,
+                                    placeholder="Enter a negative prompt",
+                                    visible=False,
+                                )
+
+                            # randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+                            seed_output = gr.Number(label="Used Seed", value=1000)
+        
+            with gr.Column(scale=1):
+                result = gr.Video(
+                    label="Generated Video", 
+                    show_label=True,
+                    height=500,
+                    container=True,
+                    autoplay=True,
+                )
+        
+        gr.HTML("""
+        <style>
+        .center-button {
+            display: flex !important;
+            justify-content: center !important;
+            height: 100% !important;
+            padding-top: 1.4em !important;
+        }
+        
+        .gradio-container {
+            max-width: 1400px !important;
+            margin: 0 auto !important;
+        }
+        
+        .gr-number input[readonly] {
+            background-color: var(--background-fill-secondary) !important;
+            border: 1px solid var(--border-color-primary) !important;
+            color: var(--body-text-color-subdued) !important;
+            cursor: default !important;
+            text-align: center !important;
+            font-weight: 500 !important;
+        }
+        </style>
+        """)
+        
+        def on_example_select(example_label):
+            if example_label and example_label in example_labels:
+                index = example_labels.index(example_label)
+                selected_prompt = examples[index]
+                selected_image_path = example_images[index] if index < len(example_images) else None
+                
+                if selected_image_path and os.path.exists(selected_image_path):
+                    try:
+                        with Image.open(selected_image_path) as img:
+                            selected_image = img.convert("RGB")
+                    except Exception:
+                        selected_image = None
+                else:
+                    selected_image = None
+                
+                return selected_prompt, selected_image
+            return "", None
+        
+        example_dropdown.change(
+            fn=on_example_select,
+            inputs=example_dropdown,
+            outputs=[prompt, input_image],
+        )
+        
+        gr.HTML("""
+        <div style="text-align: center; margin-top: 10px; margin-bottom: 15px;">
+            <p style="font-size: 16px; margin: 0;">The compute for this demo is generously provided by <a href="https://www.gmicloud.ai/" target="_blank">GMI Cloud</a>.  Note that this demo is meant as a preview of our distilled I2V model. Outside of few-step distillation, we have not yet fully optimized it for speed. Stay tuned for updates!</p>
+        </div>
+        """)
+        
+        use_negative_prompt.change(
+            fn=lambda x: gr.update(visible=x),
+            inputs=use_negative_prompt,
+            outputs=negative_prompt,
+        )
+        
+        def on_model_selection_change(selected_model):
+            if not selected_model:
+                selected_model = "FastWan2.1-T2V-1.3B"
+            
+            model_path = MODEL_PATH_MAPPING.get(selected_model)
+            show_image_input = is_i2v_model(selected_model)
+            
+            if model_path and model_path in default_params:
+                params = default_params[model_path]
+                return (
+                    gr.update(value=params.height),
+                    gr.update(value=params.width),
+                    gr.update(value=params.num_frames),
+                    gr.update(value=params.guidance_scale),
+                    gr.update(visible=show_image_input),
+                )
+            
+            return (
+                gr.update(value=448),
+                gr.update(value=832),
+                gr.update(value=20),
+                gr.update(value=3.0),
+                gr.update(visible=show_image_input),
+            )
+        
+        model_selection.change(
+            fn=on_model_selection_change,
+            inputs=model_selection,
+            outputs=[height, width, num_frames, guidance_scale, image_tab],
+        )
+        
+        def handle_generation(*args, progress=None, request: gr.Request = None):
+            model_selection, prompt, negative_prompt, use_negative_prompt, guidance_scale, num_frames, height, width, input_image = args
+            
+            result_path, seed_or_error, _ = generate_video(
+                prompt, negative_prompt, use_negative_prompt, guidance_scale, 
+                num_frames, height, width, model_selection, input_image, progress
+            )
+            
+            if result_path and os.path.exists(result_path):
+                return (
+                    result_path, 
+                    seed_or_error, 
+                    gr.update(visible=False),
+                )
+            else:
+                return (
+                    None, 
+                    seed_or_error, 
+                    gr.update(visible=True, value=seed_or_error),
+                )
+        
+        run_button.click(
+            fn=handle_generation,
+            inputs=[
+                model_selection,
+                prompt,
+                negative_prompt,
+                use_negative_prompt,
+                guidance_scale,
+                num_frames,
+                height,
+                width,
+                # randomize_seed,
+                input_image,
+            ],
+            outputs=[result, seed_output, error_output],  # timing_display removed
+            concurrency_limit=20,
+        )
+    
+    return demo
+
+
+def main():
+    parser = argparse.ArgumentParser(description="FastVideo Gradio Frontend")
+    parser.add_argument("--backend_url", type=str, default="http://localhost:8000",
+                        help="URL of the Ray Serve backend")
+    parser.add_argument("--t2v_model_paths", type=str,
+                        default="",
+                        help="Comma separated list of paths to the T2V model(s)")
+    parser.add_argument("--i2v_model_paths", type=str,
+                        default="FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers",
+                        help="Comma separated list of paths to the I2V model(s)")
+    parser.add_argument("--host", type=str, default="0.0.0.0",
+                        help="Host to bind to")
+    parser.add_argument("--port", type=int, default=7860,
+                        help="Port to bind to")
+    
+    args = parser.parse_args()
+    
+    default_params = {}
+    
+    # Load T2V model params
+    t2v_paths = [p.strip() for p in args.t2v_model_paths.split(",") if p.strip()]
+    for model_path in t2v_paths:
+        default_params[model_path] = SamplingParam.from_pretrained(model_path)
+    
+    # Load I2V model params
+    i2v_paths = [p.strip() for p in args.i2v_model_paths.split(",") if p.strip()]
+    for model_path in i2v_paths:
+        default_params[model_path] = SamplingParam.from_pretrained(model_path)
+    
+    demo = create_gradio_interface(args.backend_url, default_params)
+    
+    print(f"Starting Gradio frontend at http://{args.host}:{args.port}")
+    print(f"Backend URL: {args.backend_url}")
+    print(f"T2V Models: {args.t2v_model_paths}")
+    if args.i2v_model_paths:
+        print(f"I2V Models: {args.i2v_model_paths}")
+    
+    from fastapi import FastAPI, Request, HTTPException
+    from fastapi.responses import HTMLResponse, FileResponse
+    import uvicorn
+    
+    app = FastAPI()
+    
+    @app.get("/logo.svg")
+    def get_logo():
+        return FileResponse(
+            "assets/logos/logo.svg", 
+            media_type="image/svg+xml",
+            headers={
+                "Cache-Control": "public, max-age=3600",
+                "Access-Control-Allow-Origin": "*"
+            }
+        )
+    
+    @app.get("/favicon.ico")
+    def get_favicon():
+        favicon_path = "assets/logos/icon_simple.svg"
+        
+        if os.path.exists(favicon_path):
+            return FileResponse(
+                favicon_path, 
+                media_type="image/svg+xml",
+                headers={
+                    "Cache-Control": "public, max-age=3600",
+                    "Access-Control-Allow-Origin": "*"
+                }
+            )
+        else:
+            raise HTTPException(status_code=404, detail="Favicon not found")
+    
+    @app.get("/", response_class=HTMLResponse)
+    def index(request: Request):
+        base_url = str(request.base_url).rstrip('/')
+        return f"""
+        <!DOCTYPE html>
+        <html lang="en">
+        <head>
+            <meta charset="UTF-8" />
+            <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+            
+            <title>CausalWan</title>
+            <meta name="title" content="CausalWan">
+            <meta name="description" content="Make video generation go blurrrrrrr">
+            <meta name="keywords" content="FastVideo, video generation, AI, machine learning, CausalWan">
+            
+            <meta property="og:type" content="website">
+            <meta property="og:url" content="{base_url}/">
+            <meta property="og:title" content="CausalWan">
+            <meta property="og:description" content="Make video generation go blurrrrrrr">
+            <meta property="og:image" content="{base_url}/logo.svg">
+            <meta property="og:image:width" content="1200">
+            <meta property="og:image:height" content="630">
+            <meta property="og:site_name" content="CausalWan">
+            
+            <meta property="twitter:card" content="summary_large_image">
+            <meta property="twitter:url" content="{base_url}/">
+            <meta property="twitter:title" content="CausalWan">
+            <meta property="twitter:description" content="Make video generation go blurrrrrrr">
+            <meta property="twitter:image" content="{base_url}/logo.svg">
+            <link rel="icon" type="image/png" sizes="32x32" href="/favicon.ico">
+            <link rel="icon" type="image/png" sizes="16x16" href="/favicon.ico">
+            <link rel="apple-touch-icon" href="/favicon.ico">
+            <style>
+                body, html {{
+                    margin: 0;
+                    padding: 0;
+                    height: 100%;
+                    overflow: hidden;
+                }}
+                iframe {{
+                    width: 100%;
+                    height: 100vh;
+                    border: none;
+                }}
+            </style>
+        </head>
+        <body>
+            <iframe src="/gradio" width="100%" height="100%" style="border: none;"></iframe>
+        </body>
+        </html>
+        """
+    
+    app = gr.mount_gradio_app(
+        app, 
+        demo, 
+        path="/gradio",
+        root_path="/gradio",
+        allowed_paths=[
+            os.path.abspath("outputs"), 
+            os.path.abspath("fastvideo-logos"),
+            os.path.abspath("prompts"),
+            os.path.abspath("images"),
+            os.path.abspath(tempfile.gettempdir()),
+            os.path.abspath(os.path.join(tempfile.gettempdir(), "gradio")),
+        ]
+    )
+    
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main() 
@@ -0,0 +1,15 @@
+Some friends dancing and having fun together in circles, at a party surrounded by colored lights at a party, in a fancy old place, in a view from below them.
+A man wearing grey shorts jumps rope in a gym, weights and gym equipment in the background.
+Flying over a peninsula covered in bushy trees, while discovering the sea around it, painted a beautiful turquoise blue, on a sunny day.
+Skillful cyclist doing a wheelie on a bike while riding through a forest, on a dirt road, surrounded by many trees, in the morning.
+In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.
+A silver SUV drives along a winding, snow-covered mountain road, with dense pine trees blanketed in snow lining both sides. The scene is serene, with the vehicle moving smoothly, possibly on a winter journey or vacation. As the SUV disappears around the bend, another, darker SUV follows, creating a sense of motion and perspective on the snow-dusted asphalt. The towering, snow-laden rock formation to the right contrasts with the dark green of the pines, highlighting the peacefulness of the wintry landscape.
+A man and a woman playing in a field with grass, during a bright afternoon, while cars pass by in the distance.
+A saxophonist wearing a blazer dances while playing a song in a park.
+Romantic couple embracing and looking at each other in the middle of a forest, during a break on a road trip through nature.
+Young woman cleaning her house decorated with plants and decorations, while dancing happily to music in her headphones.
+Man dressed in 80's style dances very happily in his kitchen while listening to music on his radio and drinking wine.
+Pair of jazz musicians performing a song with their saxophone and trombone on an abandoned train.
+A young woman with short hair wearing pink sunglasses chews gum and makes a bubble gum with the city in the background.
+Natural aerial landscape with a relief covered with abundant trees and vegetation and a thick layer of mist.
+Loving couple sitting on a log on the shore of a lake outside, sharing an affectionate hug.
@@ -26,6 +26,7 @@ SEED_RANGE_MAX = 1_000_000
 SUPPORTED_MODELS = [
    "FastVideo/FastWan2.1-T2V-1.3B-Diffusers", 
    "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+    "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers",
 ]

 MODEL_CONFIGS = {
@@ -42,6 +43,13 @@ MODEL_CONFIGS = {
        "dit_cpu_offload": True,
        "vae_cpu_offload": False,
        "VSA_sparsity": 0.9,
+    },
+    "I2V-A14B": {
+        "num_cpus": 15,
+        "text_encoder_cpu_offload": True,
+        "dit_cpu_offload": True,
+        "vae_cpu_offload": False,
+        "VSA_sparsity": 0.0,
    }
 }

@@ -58,6 +66,7 @@ class VideoGenerationRequest(BaseModel):
    randomize_seed: bool = False
    return_frames: bool = False
    model_path: Optional[str] = None
+    image_data: Optional[str] = None  # Base64 encoded image for I2V


 class VideoGenerationResponse(BaseModel):
@@ -91,11 +100,38 @@ def encode_video_to_base64(frames: List[np.ndarray], fps: int = DEFAULT_FPS) ->
        return ""


+def save_image_from_base64(image_data: str, output_dir: str) -> Optional[str]:
+    """Save base64 image data to a temporary file and return the path."""
+    if not image_data:
+        return None
+    
+    try:
+        # Remove data URL prefix if present
+        if image_data.startswith('data:image/'):
+            image_data = image_data.split(',')[1]
+        
+        image_bytes = base64.b64decode(image_data)
+        
+        # Save to temporary file
+        os.makedirs(output_dir, exist_ok=True)
+        temp_image_path = os.path.join(output_dir, f"temp_input_{int(time.time() * 1000)}.png")
+        
+        with open(temp_image_path, 'wb') as f:
+            f.write(image_bytes)
+        
+        return temp_image_path
+        
+    except Exception as e:
+        print(f"Warning: Failed to save image: {e}")
+        return None
+
+
 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"
+    # 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_ATTENTION_BACKEND"] = "FLASH_ATTN"
    os.environ["FASTVIDEO_STAGE_LOGGING"] = "1"


@@ -157,22 +193,41 @@ class BaseModelDeployment:
            num_gpus=1,
            use_fsdp_inference=True,
            text_encoder_cpu_offload=config["text_encoder_cpu_offload"],
+            dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125], # TODO: hardocde for I2V
+            dit_precision="fp32",  # TODO: hardocde for I2V
            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)
+        self.default_params.seed = 1000
+        self.default_params.num_frames = 73
+        self.default_params.width = 832
+        self.default_params.height = 480

    def generate_video(self, video_request: VideoGenerationRequest) -> VideoGenerationResponse:
        total_start_time = time.time()
        
        params = prepare_sampling_params(video_request, self.default_params)

+        # Save image if provided (for I2V)
+        image_path = None
+        if video_request.image_data:
+            image_path = save_image_from_base64(video_request.image_data, self.output_path)
+            if image_path is None:
+                return VideoGenerationResponse(
+                    video_data=None,
+                    seed=params.seed,
+                    success=False,
+                    error_message="Failed to save input image",
+                )
+
        inference_start_time = time.time()
        result = self.generator.generate_video(
            prompt=video_request.prompt,
            sampling_param=params,
+            image_path=image_path,
            save_video=False,
            return_frames=False,
        )
@@ -185,6 +240,13 @@ class BaseModelDeployment:
        encoding_time = time.time() - encoding_start_time
        
        total_time = time.time() - total_start_time
+        
+        # Clean up temporary image file
+        if image_path and os.path.exists(image_path):
+            try:
+                os.remove(image_path)
+            except Exception as e:
+                print(f"Warning: Failed to remove temporary image file {image_path}: {e}")

        return VideoGenerationResponse(
            video_data=video_data,
@@ -200,7 +262,7 @@ class BaseModelDeployment:


@serve.deployment(
-    ray_actor_options={"num_cpus": 2, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
+    ray_actor_options={"num_cpus": 2, "num_gpus": 1, "runtime_env": {"conda": "demo-fv"}},
 )
 class T2VModelDeployment(BaseModelDeployment):
    def __init__(self, t2v_model_path: str, output_path: str = "outputs"):
@@ -210,7 +272,7 @@ class T2VModelDeployment(BaseModelDeployment):


@serve.deployment(
-    ray_actor_options={"num_cpus": 16, "num_gpus": 1, "runtime_env": {"conda": "fv"}},
+    ray_actor_options={"num_cpus": 16, "num_gpus": 1, "runtime_env": {"conda": "demo-fv"}},
 )
 class T2V14BModelDeployment(BaseModelDeployment):
    def __init__(self, t2v_14b_model_path: str, output_path: str = "outputs"):
@@ -221,18 +283,32 @@ class T2V14BModelDeployment(BaseModelDeployment):
        print("✅ T2V 14B model initialized successfully")


+@serve.deployment(
+    ray_actor_options={"num_cpus": 15, "num_gpus": 1, "runtime_env": {"conda": "demo-fv"}},
+)
+class I2VModelDeployment(BaseModelDeployment):
+    def __init__(self, i2v_model_path: str, output_path: str = "outputs"):
+        super().__init__(i2v_model_path, output_path)
+        # Override environment for I2V model
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
+        self._initialize_generator(MODEL_CONFIGS["I2V-A14B"])
+        print("✅ I2V model initialized successfully")
+
+
 app = FastAPI()
 limiter = Limiter(key_func=get_remote_address)
 app.state.limiter = limiter
 app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)


-@serve.deployment(num_replicas=50, ray_actor_options={"num_cpus": 2})
+@serve.deployment(num_replicas=1, ray_actor_options={"num_cpus": 1})
@serve.ingress(app)
 class FastVideoAPI:

-    def __init__(self, t2v_deployments: Dict[str, DeploymentHandle]):
+    def __init__(self, t2v_deployments: Dict[str, DeploymentHandle], i2v_deployments: Dict[str, DeploymentHandle] = None):
        self.t2v_deployments = t2v_deployments
+        self.i2v_deployments = i2v_deployments or {}
+        self.all_deployments = {**self.t2v_deployments, **self.i2v_deployments}
        
        # Initialize Prometheus metrics
        self.request_count = Counter('fastvideo_requests_total', 'Total FastVideo requests', ['model_type', 'status'])
@@ -257,10 +333,10 @@ class FastVideoAPI:
        model_name = self._get_model_name(video_request.model_path)
        
        try:
-            if video_request.model_path not in self.t2v_deployments:
+            if video_request.model_path not in self.all_deployments:
                raise ValueError(f"Model {video_request.model_path} not found")
            
-            response_ref = self.t2v_deployments[video_request.model_path].generate_video.remote(video_request)
+            response_ref = self.all_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)
@@ -291,18 +367,21 @@ class FastVideoAPI:


 def validate_configuration(model_paths: List[str], replicas: List[int]) -> None:
+    assert len(model_paths) > 0, "At least one model must be specified"
    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 model in SUPPORTED_MODELS, f"Model {model} not supported. Supported models: {SUPPORTED_MODELS}"
        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_paths: str = "",
+    t2v_model_replicas: str = "",
+    i2v_model_paths: str = "",
+    i2v_model_replicas: str = "",
    output_path: str = "outputs",
    host: str = "0.0.0.0",
    port: int = 8000,
@@ -310,21 +389,39 @@ def start_ray_serve(
    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)
+    # Parse T2V models
+    t2v_paths = [p.strip() for p in t2v_model_paths.split(",") if p.strip()]
+    t2v_reps = [int(r.strip()) for r in t2v_model_replicas.split(",") if r.strip()] if t2v_model_replicas else []
+    
+    # Parse I2V models
+    i2v_paths = [p.strip() for p in i2v_model_paths.split(",") if p.strip()]
+    i2v_reps = [int(r.strip()) for r in i2v_model_replicas.split(",") if r.strip()] if i2v_model_replicas else []
+    
+    # Validate configurations
+    all_paths = t2v_paths + i2v_paths
+    all_replicas = t2v_reps + i2v_reps
+    validate_configuration(all_paths, all_replicas)

+    # Create T2V deployments
    t2v_deps = {}
-    for model_path, replica_count in zip(model_paths, replicas):
+    for model_path, replica_count in zip(t2v_paths, t2v_reps):
        t2v_dep = T2VModelDeployment.options(num_replicas=replica_count).bind(model_path, output_path)
        t2v_deps[model_path] = t2v_dep

-    api = FastVideoAPI.bind(t2v_deps)
+    # Create I2V deployments
+    i2v_deps = {}
+    for model_path, replica_count in zip(i2v_paths, i2v_reps):
+        i2v_dep = I2VModelDeployment.options(num_replicas=replica_count).bind(model_path, output_path)
+        i2v_deps[model_path] = i2v_dep
+
+    api = FastVideoAPI.bind(t2v_deps, i2v_deps)
    serve.run(api, route_prefix="/", name="fast_video")

    print(f"Ray Serve backend started at http://{host}:{port}")
-    for model_path, replica_count in zip(model_paths, replicas):
+    for model_path, replica_count in zip(t2v_paths, t2v_reps):
        print(f"T2V Model: {model_path} | Replicas: {replica_count}")
+    for model_path, replica_count in zip(i2v_paths, i2v_reps):
+        print(f"I2V Model: {model_path} | Replicas: {replica_count}")
    print(f"Health check: http://{host}:{port}/health")
    print(f"Video generation endpoint: http://{host}:{port}/generate_video")

@@ -340,12 +437,20 @@ if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="FastVideo Ray Serve Backend")
    parser.add_argument("--t2v_model_paths",
                        type=str,
-                        default="FastVideo/FastWan2.1-T2V-1.3B-Diffusers,FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
+                        default="",
                        help="Comma separated list of paths to the T2V model(s)")
    parser.add_argument("--t2v_model_replicas",
                        type=str,
-                        default="4,4",
+                        default="",
                        help="Comma separated list of number of replicas for the T2V model(s)")
+    parser.add_argument("--i2v_model_paths",
+                        type=str,
+                        default="FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers",
+                        help="Comma separated list of paths to the I2V model(s)")
+    parser.add_argument("--i2v_model_replicas",
+                        type=str,
+                        default="1",
+                        help="Comma separated list of number of replicas for the I2V model(s)")
    parser.add_argument("--output_path",
                        type=str,
                        default="outputs",
@@ -361,13 +466,21 @@ if __name__ == "__main__":
    
    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)
+    # Parse and validate all models
+    t2v_paths = [p.strip() for p in args.t2v_model_paths.split(",") if p.strip()]
+    t2v_reps = [int(r.strip()) for r in args.t2v_model_replicas.split(",") if r.strip()] if args.t2v_model_replicas else []
+    i2v_paths = [p.strip() for p in args.i2v_model_paths.split(",") if p.strip()]
+    i2v_reps = [int(r.strip()) for r in args.i2v_model_replicas.split(",") if r.strip()] if args.i2v_model_replicas else []
+    
+    all_paths = t2v_paths + i2v_paths
+    all_replicas = t2v_reps + i2v_reps
+    validate_configuration(all_paths, all_replicas)
    
    start_ray_serve(
        t2v_model_paths=args.t2v_model_paths,
        t2v_model_replicas=args.t2v_model_replicas,
+        i2v_model_paths=args.i2v_model_paths,
+        i2v_model_replicas=args.i2v_model_replicas,
        output_path=args.output_path,
        host=args.host,
        port=args.port,
@@ -376,4 +489,4 @@ if __name__ == "__main__":
    setup_signal_handlers()
    print("✅ FastVideo backend is running. Press Ctrl-C to stop.")
    while True:
-        time.sleep(3600) 
+        time.sleep(3600) 
@@ -0,0 +1,5 @@
+python examples/inference/gradio/serving/start_ray_serve_app.py \
+--t2v_model_paths "" \
+--t2v_model_replicas "" \
+--i2v_model_paths "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers" \
+--i2v_model_replicas "1"
@@ -20,8 +20,10 @@ 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"
+DEFAULT_T2V_MODELS = ""
+DEFAULT_T2V_REPLICAS = ""
+DEFAULT_I2V_MODELS = "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers"
+DEFAULT_I2V_REPLICAS = "1"

 HEALTH_CHECK_TIMEOUT = 5
 HEALTH_CHECK_MAX_RETRIES = 100
@@ -100,6 +102,12 @@ class ServiceManager:
            "port": self.args.backend_port
        }
        
+        # Add I2V parameters if provided
+        if self.args.i2v_model_paths:
+            backend_args["i2v_model_paths"] = self.args.i2v_model_paths
+        if self.args.i2v_model_replicas:
+            backend_args["i2v_model_replicas"] = self.args.i2v_model_replicas
+        
        self.backend_process = self._start_service("ray_serve_backend.py", backend_args, "backend")
        return self.backend_process
    
@@ -111,6 +119,10 @@ class ServiceManager:
            "port": self.args.frontend_port
        }
        
+        # Add I2V parameters if provided
+        if self.args.i2v_model_paths:
+            frontend_args["i2v_model_paths"] = self.args.i2v_model_paths
+        
        self.frontend_process = self._start_service("gradio_frontend.py", frontend_args, "frontend")
        return self.frontend_process
    
@@ -173,6 +185,9 @@ def print_startup_info(args: argparse.Namespace) -> None:
    print("=" * 50)
    print(f"T2V Models: {args.t2v_model_paths}")
    print(f"T2V Model Replicas: {args.t2v_model_replicas}")
+    if args.i2v_model_paths:
+        print(f"I2V Models: {args.i2v_model_paths}")
+        print(f"I2V Model Replicas: {args.i2v_model_replicas}")
    print(f"Output: {args.output_path}")
    print(f"Backend: http://{args.backend_host}:{args.backend_port}")
    print(f"Frontend: http://{args.frontend_host}:{args.frontend_port}")
@@ -190,6 +205,14 @@ def parse_arguments() -> argparse.Namespace:
                        type=str, 
                        default=DEFAULT_T2V_REPLICAS,
                        help="Comma separated list of number of replicas for the T2V model(s)")
+    parser.add_argument("--i2v_model_paths", 
+                        type=str,
+                        default=DEFAULT_I2V_MODELS,
+                        help="Comma separated list of paths to the I2V model(s)")
+    parser.add_argument("--i2v_model_replicas", 
+                        type=str, 
+                        default=DEFAULT_I2V_REPLICAS,
+                        help="Comma separated list of number of replicas for the I2V model(s)")
    parser.add_argument("--output_path", 
                        type=str, 
                        default=DEFAULT_OUTPUT_PATH,
@@ -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"
@@ -0,0 +1,47 @@
+A watermelon wearing a helmet is crushed by a hydraulic press, causing it to flatten and burst open.
+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.
+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.
+A red toy car is being crushed by a large hydraulic press, which is flattening objects as if they were under a hydraulic press.
+A large, cylindrical object is seen pressing down on a small orange ball, causing it to flatten as if it were under a hydraulic press. The background features a green wall with yellow and red warning signs.
+The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is shown compressing a wooden object, which shatters into small pieces. The background features a green wall with a yellow sign displaying a lightning bolt.
+A large metal cylinder is seen descending, flattening objects as if they were under a hydraulic press. The cylinder compresses a stack of matches and boxes, causing them to crumble into small pieces. The scene is set against a green background with yellow and red signs.
+A large metal press is shown compressing a pile of colorful macarons, flattening them as if they were under a hydraulic press. The press moves down, crushing the macarons into a pile of crumbs and squishing the colorful filling out.
+The video shows a metal press flattening objects as if they were under a hydraulic press. The press is pressing down on a pile of colorful gummy candies, squishing them into a pile of squiggly shapes. The press is made of metal and has a large base, and the gummy candies are of various colors, including red, green, and orange. The background is a green wall, and the press is placed on a metal surface.
+A pile of colorful candies is being flattened by a hydraulic press, causing them to crumble into small pieces.
+The video shows a stack of colorful sponges being flattened as if they were under a hydraulic press. The sponges, which are pink, white, blue, and green, are compressed into a smaller size, demonstrating the press's power. The background features a green wall with a yellow and red sign, adding context to the setting.
+A bowling ball is placed on a metal platform, and a large metal cylinder descends from above, flattening the ball as if it were under a hydraulic press. The ball is crushed into a flat, round shape, leaving a pile of debris around it.
+A large metal cylinder with yellow and black stripes is seen pressing down on a pile of popcorn, flattening the objects as if they were under a hydraulic press.
+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.
+The video shows a close-up of a metal cylinder pressing down on a yellow object, which is being flattened as if it were under a hydraulic press. The cylinder is positioned above the object, and the force is causing the object to compress and spread out, creating a visible deformation. The background is blurred, focusing attention on the action of the cylinder and the object being flattened.
+A colorful puzzle ball is being crushed by a large metal cylinder, which flattens the objects as if they were under a hydraulic press.
+The video shows a hydraulic press flattening objects as if they were under a hydraulic press. The press is shown in action, compressing two colorful objects that resemble sandwiches. The press is yellow and black striped, and the objects being flattened are placed on a metal plate. The background is green, and the press is moving down, compressing the objects.
+The scene shows a metal press with a yellow and black striped pattern, holding a container filled with chocolate. A metal cylinder is descending, flattening the chocolate as if it were under a hydraulic press. The background is a green wall, and the press is mounted on a sturdy metal frame.
+The video shows a colorful sponge being flattened as if it were under a hydraulic press, with the sponge being compressed and eventually flattened into a thin layer.
+The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is pressing down on a stack of wooden blocks, causing them to crumble and break apart. The press is black and yellow striped, and the wooden blocks are small and rectangular. The background is green, and the press is sitting on a metal table.
+A pile of colorful candies is being flattened by a hydraulic press, causing them to crumble into small pieces.
+The video shows a stack of colorful sponges being flattened by a large, cylindrical object, which appears to be a hydraulic press. The sponges, which are pink, blue, white, and green, are compressed into a single layer, demonstrating the press's powerful force. The background features a green wall with a yellow and red sign, adding context to the industrial setting.
+A bowling ball is placed on a metal platform, and a large metal cylinder descends from above, flattening the ball as if it were under a hydraulic press. The ball is crushed into a flat, round shape, demonstrating the immense pressure applied by the cylinder.
+A large metal cylinder with yellow and black stripes is seen pressing down on a pile of popcorn, flattening the objects as if they were under a hydraulic press. The popcorn is crushed and scattered around the base of the cylinder, creating a satisfying visual effect.
+The video shows a hydraulic press in action, flattening objects as if they were under a hydraulic press. The press is composed of a large, cylindrical metal cylinder with yellow and black stripes, and a metal base. The objects being flattened are two cylindrical blocks of cotton candy, one pink and one blue. The press is positioned on a metal table, and the background features a green wall with a yellow and red sign.
+The video shows a large 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.
+The video shows a cylindrical object being pressed down onto a flat surface, causing the objects beneath it to be flattened as if they were under a hydraulic press. The objects being flattened appear to be yellow and are being crushed into a pile of debris. The background is a greenish-gray color, and the surface on which the objects are being flattened is metallic and shiny.
+A green and blue object with a spiky texture is being flattened by a large, cylindrical metal press, demonstrating its resilience and durability.
+The video shows a stack of caramelized sugar cubes being flattened as if they were under a hydraulic press, resulting in a messy pile of broken sugar on the table.
+A large metal cylinder is seen pressing down on a pile of colorful jelly beans, flattening them as if they were under a hydraulic press.
+The video shows a machine with a yellow and black striped cylinder pressing down on a stack of colorful sponges, flattening them as if they were under a hydraulic press. The machine is situated in a green-walled room with warning signs in the background.
+The video shows a machine with a yellow and black striped cylinder, which is pressing down on two colorful objects, flattening them as if they were under a hydraulic press. The machine appears to be in a workshop or industrial setting, with a green wall in the background. The objects being flattened are green and orange, and the machine is covered in dirt and grime, indicating it has been used frequently.
+The video shows a large, industrial press flattening objects as if they were under a hydraulic press. The press is shown in action, compressing a pile of pink objects into a pile of crumbs. The press is large and metallic, with a yellow and black striped pattern on its side. The background is a green wall with a yellow warning sign.
+The video shows a pink, sparkly ball being crushed by a large, rusty cylinder, which flattens the objects as if they were under a hydraulic press.
+A lime is being crushed by a hydraulic press, causing it to flatten and burst open, releasing its juice and segments.
+The video shows a machine with a yellow and black striped cylinder, which is flattening objects as if they were under a hydraulic press. The machine is pressing down on two colorful objects, causing them to compress and flatten. The background is a green wall, and the machine appears to be in a workshop or industrial setting.
+The video shows a large, yellow and black striped cylinder flattening objects as if they were under a hydraulic press. The objects being flattened are pink and are being crushed into small pieces. The background is a green wall with a yellow sign.
+The video shows a machine with a yellow and black striped cylinder pressing down on two colorful objects, which are flattened as if they were under a hydraulic press. The machine is positioned on a metal platform, and the background is a green wall.
+A green cube is being compressed by a hydraulic press, which flattens the object as if it were under a hydraulic press. The press is shown in action, with the cube being squeezed into a smaller shape.
+A pink, sparkly ball is being crushed by a large, rusty cylinder, which flattens the objects as if they were under a hydraulic press.
+A red cabbage is being crushed by a hydraulic press, which flattens the objects as if they were under a hydraulic press. The press is shown in action, compressing the cabbage into a smaller, more compact form.
+A lime is being crushed by a hydraulic press, causing it to flatten and burst open, releasing its juice and pulp.
+A large metal press is shown compressing a stack of burgers, causing them to be flattened and crushed into a pile of ground meat.
+A pizza is being crushed by a hydraulic press, causing the toppings to spread out and the crust to crumble.
+A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.
+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.
+A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.
@@ -0,0 +1,94 @@
+#!/bin/bash
+
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export TOKENIZERS_PARALLELISM=false
+
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
+VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json"
+NUM_GPUS=1
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir "wan_ode_init_crush_smol"
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "wan_ode_init_crush_smol"
+  --max_train_steps 6000
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 77
+  --warp_denoising_step
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim 1
+  --hsdp_shard_dim 1
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR" 
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 6e-6
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 1000
+  --training_state_checkpointing_steps 1000
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+torchrun \
+  --nnodes 1 \
+  --nproc_per_node $NUM_GPUS \
+    fastvideo/training/ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -0,0 +1,136 @@
+#!/bin/bash
+#SBATCH --job-name=2e6B8_16kFV_ode_vidprom
+#SBATCH --partition=main
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --gres=gpu:8
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=1440G
+#SBATCH --output=ode_vidprom16k/ode_vidprom8b16k_2e-6.out
+#SBATCH --error=ode_vidprom16k/ode_vidprom8b16k_2e-6.err
+#SBATCH --exclusive
+set -e -x
+
+# Environment Setup
+source ~/conda/miniconda/bin/activate
+conda activate your-conda-env
+
+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_API_KEY=your-wandb-api-key
+# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "NODE_RANK: $NODE_RANK"
+
+
+MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
+DATA_DIR="your-data-dir"
+VALIDATION_DATASET_FILE="examples/training/consistency_finetune/causal_ode_init/validation.json"
+OUTPUT_DIR="your-output-dir"
+INIT_WEIGHTS_FROM_SAFETENSORS="your-init-weights-from-safetensors" # bidirectional weights from Wan2.1-T2V-1.3B-Diffusers
+NUM_GPUS=8
+# export CUDA_VISIBLE_DEVICES=4,5
+# IP=[MASTER NODE IP]
+
+# Training arguments
+training_args=(
+  --tracker_project_name "wan_ode_init"
+  --output_dir $OUTPUT_DIR
+  --override_transformer_cls_name "CausalWanTransformer3DModel"
+  --wandb_run_name "vidprom_8b16k_ode_init_2e-6"
+  # --resume_from_checkpoint "ode_init_diffusers/"
+  --warp_denoising_step
+  --log_visualization
+  --max_train_steps 6001
+  --train_batch_size 1
+  --train_sp_batch_size 1
+  --gradient_accumulation_steps 1
+  --num_latent_t 21
+  --num_height 480
+  --num_width 832
+  --num_frames 81
+  --dmd_denoising_steps "1000,750,500,250"
+  --enable_gradient_checkpointing_type "full"
+)
+
+# Parallel arguments
+parallel_args=(
+  --num_gpus $NUM_GPUS
+  --sp_size 1
+  --tp_size 1
+  --hsdp_replicate_dim $NUM_GPUS
+  --hsdp_shard_dim 1
+)
+
+# Model arguments
+model_args=(
+  --model_path $MODEL_PATH
+  --pretrained_model_name_or_path $MODEL_PATH
+)
+
+# Dataset arguments
+dataset_args=(
+  --data_path "$DATA_DIR" 
+  --dataloader_num_workers 1
+)
+
+# Validation arguments
+validation_args=(
+  --log_validation
+  --validation_dataset_file "$VALIDATION_DATASET_FILE"
+  --validation_steps 50
+  --validation_sampling_steps "50"
+  --validation_guidance_scale "6.0"
+  --init_weights_from_safetensors $INIT_WEIGHTS_FROM_SAFETENSORS
+)
+
+# Optimizer arguments
+optimizer_args=(
+  --learning_rate 2e-6
+  --mixed_precision "bf16"
+  --weight_only_checkpointing_steps 500
+  --training_state_checkpointing_steps 500
+  --weight_decay 1e-4
+  --max_grad_norm 1.0
+)
+
+# Miscellaneous arguments
+miscellaneous_args=(
+  --inference_mode False
+  --checkpoints_total_limit 3
+  --training_cfg_rate 0.1
+  --multi_phased_distill_schedule "4000-1"
+  --not_apply_cfg_solver
+  --dit_precision "fp32"
+  --num_euler_timesteps 50
+  --ema_start_step 0
+  # --enable_gradient_checkpointing_type "full"
+)
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+srun torchrun \
+  --nnodes $SLURM_JOB_NUM_NODES \
+  --nproc_per_node $NUM_GPUS \
+  --node_rank $SLURM_PROCID \
+  --rdzv_backend=c10d \
+  --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
+    fastvideo/training/ode_causal_pipeline.py \
+    "${parallel_args[@]}" \
+    "${model_args[@]}" \
+    "${dataset_args[@]}" \
+    "${training_args[@]}" \
+    "${optimizer_args[@]}" \
+    "${validation_args[@]}" \
+    "${miscellaneous_args[@]}"
@@ -0,0 +1,25 @@
+#!/bin/bash
+
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_TYPE="wan"
+DATA_MERGE_PATH="$(dirname "$0")/crush_smol_prompts.txt"
+OUTPUT_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/pipelines/preprocess/v1_preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size 1 \
+    --seed 42 \
+    --max_height 480 \
+    --max_width 832 \
+    --num_frames 81 \
+    --flow_shift 5.0 \
+    --dataloader_num_workers 0 \
+    --output_dir=$OUTPUT_DIR \
+    --train_fps 16 \
+    --samples_per_file 8 \
+    --flush_frequency 8 \
+    --video_length_tolerance_range 5 \
+    --preprocess_task "ode_trajectory" 
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Y-aang	c1d7c51bc7	fixi2v demo bug	2025-11-19 20:53:43 +00:00
Y-aang	5de2d80ebd	add full i2v demo	2025-11-19 15:40:04 +00:00
RandNMR73	6cd227349c	add demo prompts	2025-11-19 05:49:01 +00:00
RandNMR73	7e15e5dab1	add demo images	2025-11-19 05:42:37 +00:00
RandNMR73	c9ee4f8bf8	add i2v demo	2025-11-18 09:40:05 +00:00
SolitaryThinker	f764b43aaa	nit	2025-11-16 23:10:47 +00:00
SolitaryThinker	592af8c954	fix lint	2025-11-16 23:09:04 +00:00
SolitaryThinker	4225a96a40	use new ckpt	2025-11-16 12:49:06 +00:00
JerryZhou54	733c14a2cb	Change to fp32 inference	2025-11-16 04:53:57 +00:00
JerryZhou54	81b302409e	Add i2v images	2025-11-16 04:51:14 +00:00
JerryZhou54	7fbd0cfec7	Fi lint	2025-11-15 23:50:04 +00:00
RandNMR73	85b9934079	Add inference for MoE SF	2025-11-15 23:44:21 +00:00
Y-aang	c30779184f	fix: incorrect dv in vsa Triton kernel causing test_vsa error (#879 )	2025-11-14 22:00:39 -08:00
William Lin	9d188c0b6c	[misc] update wechat and slack invite links (#875 )	2025-11-12 23:03:56 -08:00
Mihir Jagtap	9dd7c54221	[docs] Update Home Readme.md with fixed links (#873 )	2025-11-12 13:32:44 -08:00
William Lin	62b95d8287	[feat] prepare for wan2.2 SF (#861 )	2025-11-04 18:06:48 -08:00
Kaiqin Kong	fdf21702f5	[Docs] add diagrams to docs (#863 )	2025-11-04 16:29:07 -08:00
Ohm-Rishabh	2972fc9449	Improve FSDP loading with size-based filtering (#853 )	2025-11-04 15:31:07 -08:00
Mihir Jagtap	8f5712629f	[docs] port to mkdocs (#855 )	2025-11-04 14:31:56 -08:00
Kevin Lin	436c701b9f	[bugfix] Add Cosmos2 sampling params to registry (#862 )	2025-11-02 00:09:17 -07:00
Kevin Lin	543fea88e3	[Feature] Add Cosmos2 i2v pipeline (#837 )	2025-10-30 20:03:57 -07:00
Kaiqin Kong	bdec816b31	move STA_configuration.py to fastvideo/attention/backends (#856 )	2025-10-29 13:54:13 -07:00
William Lin	2cd2e57d2e	[ci] fix causal ssim test (#848 )	2025-10-26 19:33:07 -07:00
William Linandainsley	9370234294	[feat] Add gradio local inference demo (#847 ) Co-authored-by: ainsley <jzhang2765@wisc.edu>	2025-10-26 07:01:33 -07:00
Jinzhe Pan	50da62e722	[bugfix] always force spawn instead of fork (#852 )	2025-10-23 16:36:50 -07:00
William Lin	4f3e8751db	[bugfix] [misc] Use training_state_checkpointing_steps in scripts/ (#846 )	2025-10-19 20:20:53 -07:00
Jinzhe PanandXingyu Long	f4c58894d9	[Feat] add ray support (#838 ) Co-authored-by: Xingyu Long <xingyulong97@gmail.com>	2025-10-16 23:17:54 -07:00
Ohm-Rishabh	01c94ef385	[feat] unified trainer logging (#841 )	2025-10-16 23:16:16 -07:00
Zhang Peiyuan	2415226d25	Update WeChat Link	2025-10-13 21:02:46 -07:00
Jiali Chen	404314d00f	[Feature]Add video-to-video (V2V) pipeline (#829 )	2025-10-12 21:53:05 -07:00
zyang6andkiritorl	87489f0872	Add wan2.1 functionality support for Ascend NPU platform (#810 ) Co-authored-by: kiritorl <1021709528@qq.com>	2025-10-09 16:25:08 -07:00
Zhang Peiyuan	9ce7c8039e	Update Wechat link	2025-10-06 15:01:19 -07:00
William Lin	e1e25e95f9	[feature] Add torch profiler (#827 )	2025-10-06 07:59:46 -07:00
William Lin	490bde90e1	[bugfix] Allow overriding dit checkpoint for inference and Lower VSA LR in example scripts (#831 )	2025-10-05 01:44:49 -07:00
RandNMR73 SolitaryThinker gemini-code-assist[bot]	dc7596b973	[self-forcing][8/n] Self-Forcing For Wan2.2-A14B + torch.compile training and distillation support (#818 ) Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2025-10-02 15:01:45 -07:00
William Lin	335afa4457	[bugfix] Use training_state_checkpointing_steps instead of checkpointing_steps (#821 )	2025-09-28 15:22:43 -07:00
Yongqi Chen	3f77a6805a	[Feature]Update count trainable param for FSDP2 (#820 )	2025-09-28 15:22:04 -07:00
RandNMR73	13d0aae706	Add Sage Attention 3 Backend (#815 )	2025-09-24 15:11:38 -07:00
William Lin	404cbf4f3c	[self-forcing] [6/n] Add Ode Init training (#811 )	2025-09-22 17:58:19 -07:00
William Lin	958ffec844	[bugfix] Update learning rates for sparse distillation recipe (#812 )	2025-09-22 12:07:03 -07:00
Wei Zhou RandNMR73 SolitaryThinker	31f000d1cc	[self-forcing] [5/n] Add Self-Forcing distillation pipeline (#808 ) Co-authored-by: RandNMR73 <notomatthew31@gmail.com> Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>	2025-09-20 19:32:10 -07:00
Yongqi Chen	cd32b3e02f	Update example files and readme (#809 )	2025-09-20 18:15:59 -07:00
Zhang Peiyuan	bf27908095	Update WeChat Link	2025-09-20 14:16:20 -07:00
William Lin	c5f9ea53b2	[self-forcing] [4/n] Preprocessing for collecting ODE trajectory (#788 )	2025-09-15 17:54:42 -07:00
William Lin	d32a7184da	[bugfix] Wan2.2 Boundary ratio (#804 )	2025-09-15 11:17:35 -07:00
Wenxuan Tanandgemini-code-assist[bot]	2930abe456	[Bugfix] Fix VMoba requirements (#802 ) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>	2025-09-14 18:28:52 -07:00