[V1] Model config (#358) (#368)

Co-authored-by: Will Lin <wlsaidhi@gmail.com>

.github/scripts/runpod_api.py

@@ -28,7 +28,7 @@ def parse_arguments():
     parser.add_argument(
         '--image',
         type=str,
-        default='runpod/pytorch:2.4.0-py3.11-cuda12.4.1-devel-ubuntu22.04',
+        required=True,
         help='Docker image to use')
     return parser.parse_args()
 
@@ -46,6 +46,16 @@ HEADERS = {
 
 def create_pod():
     """Create a RunPod instance"""
+    # Ensure image name is lowercase (Docker requirement)
+    image_name = args.image.lower()
+    print(f"Using specified image: {image_name}")
+
+    docker_start_cmd = [
+        "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"
+    ]
+    
     print(f"Creating RunPod instance with GPU: {args.gpu_type}...")
     payload = {
         "name": f"fastvideo-{JOB_ID}-{RUN_ID}",
@@ -53,8 +63,9 @@ def create_pod():
         "volumeInGb": args.volume_size,
         "gpuTypeIds": [args.gpu_type],
         "gpuCount": args.gpu_count,
-        "imageName": args.image,
-        "allowedCudaVersions": ["12.4"]
+        "imageName": image_name,
+        "allowedCudaVersions": ["12.4"],
+        "dockerStartCmd": docker_start_cmd
     }
 
     response = requests.post(PODS_API, headers=HEADERS, json=payload)
@@ -91,7 +102,7 @@ def wait_for_pod(pod_id):
             "Timed out waiting for RunPod to reach RUNNING state")
 
     # Wait for ports to be assigned
-    max_attempts = 6
+    max_attempts = 50
     attempts = 0
     while attempts < max_attempts:
         response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS)
@@ -108,7 +119,7 @@ def wait_for_pod(pod_id):
         print(
             f"Waiting for SSH port and public IP to be available... (attempt {attempts+1}/{max_attempts})"
         )
-        time.sleep(10)
+        time.sleep(20)
         attempts += 1
 
     if attempts >= max_attempts:
@@ -145,16 +156,15 @@ def execute_command(pod_id):
     ]
     subprocess.run(scp_command, check=True)
 
+    # For custom image, we can use the pre-configured environment
     setup_steps = [
-        "cd /workspace",
-        "wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh",
-        "bash Miniconda3-latest-Linux-x86_64.sh -b -p $HOME/miniconda3",
-        "source $HOME/miniconda3/bin/activate",
-        "conda create --name venv python=3.10.0 -y", "conda activate venv",
-        "mkdir -p /workspace/repo",
         "tar -xzf /tmp/repo.tar.gz --no-same-owner -C /workspace/",
-        f"cd /workspace/{repo_name}", args.test_command
+        f"cd /workspace/{repo_name}",
+        "source /opt/conda/etc/profile.d/conda.sh",
+        "conda activate fastvideo-dev",
+        args.test_command
     ]
+    
     remote_command = " && ".join(setup_steps)
 
     ssh_command = [

.github/workflows/build-image.yml

@@ -0,0 +1,78 @@
+name: Build and Push Docker Image
+
+on:
+  workflow_dispatch:  # Only manual triggers
+
+jobs:
+  build-and-push:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      
+      - name: Free up disk space
+        run: |
+          # Display initial space
+          echo "Initial disk space:"
+          df -h
+          
+          # Remove large directories directly
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          sudo rm -rf /usr/share/swift
+          sudo rm -rf /usr/local/lib/node_modules
+          sudo rm -rf /usr/local/share/powershell
+          sudo rm -rf /usr/share/rust
+          sudo rm -rf /usr/local/.ghcup
+          
+          # Remove cached files
+          sudo rm -rf /var/lib/apt/lists/*
+          sudo rm -rf /var/cache/apt/archives/*
+          
+          # Clean Docker
+          docker system prune -af --volumes
+          
+          # Display available space after cleanup
+          echo "Disk space after cleanup:"
+          df -h
+      
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+      
+      - name: Extract metadata for Docker
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: ghcr.io/${{ github.repository }}/fastvideo-dev
+          tags: |
+            type=raw,value=latest
+            type=sha,format=short
+      
+      - name: Build and push Docker image
+        id: build-push
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          push: true
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+      
+      - name: Success message
+        run: |
+          echo "✅ Image successfully built and pushed to ghcr.io/${{ github.repository }}/fastvideo-dev:latest"
+          echo "To run tests with this image, manually trigger the 'Run Tests' workflow." 

.github/workflows/docs.yml

@@ -8,12 +8,14 @@ on:
       - main
     paths:
       - "docs/**/*.md"
+      - "fastvideo/v1/examples/**/*.py"
   pull_request:
     branches: 
       - main
     types: [opened, ready_for_review, synchronize, reopened]
     paths:
       - "docs/**/*.md"
+      - "fastvideo/v1/examples/**/*.py"
 
   # Allows you to run this workflow manually from the Actions tab
   workflow_dispatch:
@@ -36,9 +38,13 @@ defaults:
     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

.github/workflows/fastvideo-publish.yml

@@ -6,6 +6,7 @@ on:
       - main
     paths:
       - 'pyproject.toml'  # Trigger when pyproject.toml changes
+  workflow_dispatch:
 
 jobs:
   check-version-change:
@@ -41,7 +42,7 @@ jobs:
 
   build-publish-main:
     needs: check-version-change
-    if: needs.check-version-change.outputs.version-changed == 'true'
+    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
     runs-on: ubuntu-latest
     permissions:
       id-token: write  # Needed for OIDC Trusted Publishing

.github/workflows/pr-test.yml

@@ -14,17 +14,35 @@ on:
       - ".github/workflows/pr-test.yml"
   workflow_dispatch:
     inputs:
+      custom_image:
+        description: "Custom image from this repository (default: fastvideo-dev:latest)"
+        required: false
+        default: "fastvideo-dev:latest"
+        type: string
       run_encoder_test:
         description: "Run encoder-test"
         required: false
         default: false
         type: boolean
+      run_vae_test:
+        description: "Run vae-test"
+        required: false
+        default: false
+        type: boolean
+      run_transformer_test:
+        description: "Run transformer-test"
+        required: false
+        default: false
+        type: boolean
       run_ssim_test:
         description: "Run ssim-test"
         required: false
         default: false
         type: boolean
 
+env:
+  PYTHONUNBUFFERED: "1"
+
 concurrency:
   group: pr-test-${{ github.ref }}
   cancel-in-progress: true
@@ -39,6 +57,8 @@ jobs:
     if: ${{ github.event.pull_request.draft == false || github.event_name == 'workflow_dispatch' }}
     outputs:
       encoder-test: ${{ steps.filter.outputs.encoder-test }}
+      vae-test: ${{ steps.filter.outputs.vae-test }}
+      transformer-test: ${{ steps.filter.outputs.transformer-test }}
     steps:
       - uses: actions/checkout@v4
       - uses: dorny/paths-filter@v3
@@ -49,6 +69,14 @@ jobs:
               - 'fastvideo/v1/models/encoders/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/encoders/**'
+            vae-test:
+              - 'fastvideo/v1/models/vaes/**'
+              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/tests/vaes/**'
+            transformer-test:
+              - 'fastvideo/v1/models/dits/**'
+              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/tests/transformers/**'
 
   encoder-test:
     needs: change-filter
@@ -87,9 +115,8 @@ jobs:
           --gpu-type "NVIDIA A40"
           --gpu-count 1
           --volume-size 100
-          --test-command "pip install -e .[test] &&
-          pip install flash-attn==2.7.0.post2 --no-build-isolation &&
-          pytest ./fastvideo/v1/tests/encoders -s"
+          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
 
       - name: Terminate RunPod Instances
         if: ${{ always() }}
@@ -99,6 +126,102 @@ jobs:
           JOB_ID: "encoder-test"
         run: python .github/scripts/runpod_cleanup.py
 
+  vae-test:
+    needs: change-filter
+    if: >-
+      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.vae-test == 'true') || 
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_vae_test == 'true')
+    runs-on: ubuntu-latest
+    environment: runpod-runners
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Set up SSH key
+        run: |
+          mkdir -p ~/.ssh
+          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
+          chmod 600 ~/.ssh/id_rsa
+          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
+
+      - name: Install dependencies
+        run: pip install requests
+
+      - name: Run tests on RunPod
+        env:
+          JOB_ID: "vae-test"
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: 30
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "NVIDIA A40"
+          --gpu-count 1
+          --volume-size 100
+          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: "vae-test"
+        run: python .github/scripts/runpod_cleanup.py
+
+  transformer-test:
+    needs: change-filter
+    if: >-
+      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.transformer-test == 'true') || 
+      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_transformer_test == 'true')
+    runs-on: ubuntu-latest
+    environment: runpod-runners
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Set up SSH key
+        run: |
+          mkdir -p ~/.ssh
+          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
+          chmod 600 ~/.ssh/id_rsa
+          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
+
+      - name: Install dependencies
+        run: pip install requests
+
+      - name: Run tests on RunPod
+        env:
+          JOB_ID: "transformer-test"
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: 30
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "NVIDIA L40S"
+          --gpu-count 1
+          --volume-size 100
+          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: "transformer-test"
+        run: python .github/scripts/runpod_cleanup.py
+      
   ssim-test:
     needs: change-filter
     if: >-
@@ -130,16 +253,15 @@ jobs:
           JOB_ID: "ssim-test"
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           GITHUB_RUN_ID: ${{ github.run_id }}
-        timeout-minutes: 30
+        timeout-minutes: 45
         run: >-
           python .github/scripts/runpod_api.py 
           --gpu-type "NVIDIA A40"
           --gpu-count 2
-          --disk-size 100
-          --volume-size 100
-          --test-command "pip install -e .[test] &&
-          pip install flash-attn==2.7.0.post2 --no-build-isolation &&
-          pytest ./fastvideo/v1/tests/ssim -vs"
+          --disk-size 200
+          --volume-size 200
+          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
 
       - name: Terminate RunPod Instances
         if: ${{ always() }}
@@ -150,7 +272,7 @@ jobs:
         run: python .github/scripts/runpod_cleanup.py
 
   runpod-cleanup:
-    needs: [encoder-test, ssim-test] # Add other jobs to this list as you create them
+    needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
     if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }}
     runs-on: ubuntu-latest
     steps:
@@ -167,7 +289,7 @@ jobs:
 
       - name: Cleanup all RunPod instances
         env:
-          JOB_IDS: '["encoder-test", "ssim-test"]' # JSON array of job IDs
+          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test"]' # JSON array of job IDs
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           GITHUB_RUN_ID: ${{ github.run_id }}
         run: python .github/scripts/runpod_cleanup.py

.github/workflows/sta-publish.yml

@@ -6,6 +6,7 @@ on:
       - main
     paths:
       - "csrc/sliding_tile_attention/setup.py"
+  workflow_dispatch:
 
 jobs:
   check-version-change:
@@ -43,7 +44,7 @@ jobs:
   build_wheels:
     name: Build Wheel
     needs: check-version-change
-    if: ${{ needs.check-version-change.outputs.version-changed == 'true' }}
+    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
     runs-on: ${{ matrix.os }}
 
     strategy:
@@ -57,6 +58,29 @@ jobs:
           cuda-version: ['12.4.1', '12.5.1', '12.6.3']
 
     steps:
+      - name: Free up disk space
+        run: |
+          echo "Initial disk space:"
+          df -h
+          
+          # Remove large directories
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /usr/local/lib/android
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          sudo rm -rf /usr/share/swift
+          sudo rm -rf /usr/local/lib/node_modules
+          sudo rm -rf /usr/local/share/powershell
+          sudo rm -rf /usr/share/rust
+          sudo rm -rf /usr/local/.ghcup
+          
+          # Remove cached files
+          sudo rm -rf /var/lib/apt/lists/*
+          sudo rm -rf /var/cache/apt/archives/*
+          
+          echo "Disk space after cleanup:"
+          df -h
+
       - name: Checkout
         uses: actions/checkout@v4
 
@@ -145,7 +169,7 @@ jobs:
   publish_package:
     name: Publish package
     needs: [build_wheels, check-version-change]
-    if: ${{ needs.check-version-change.outputs.version-changed == 'true' }}
+    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
     runs-on: ubuntu-22.04
     permissions:
       id-token: write  # Needed for OIDC Trusted Publishing

.gitignore

@@ -3,14 +3,11 @@ __pycache__
 *.pth
 UCF-101/
 results/
-build/
-fastvideo.egg-info/
 wandb/
 *.ipynb
 *.jpg
 *.safetensors
 *.mp4
-!fastvideo/v1/tests/ssim/reference_videos/**/*.mp4
 *.png
 *.gif
 *.pth
@@ -26,11 +23,40 @@ outputs_video
 sbatch.sh
 *.out
 env
-dist/
 *.o
 **/build/
-**.egg-info
 **.pyc
-**.egg
 **.txt
-**.json
+**.json
+
+# Distribution / packaging
+build/
+dist/
+*.egg-info/
+*.egg
+eggs/
+.eggs/
+
+# Sphinx documentation
+docs/_build/
+docs/source/getting_started/examples/
+docs/source/inference/examples/
+
+# VSCode
+.vscode/
+
+# DS Store
+.DS_Store
+
+# vim swap files
+*.swo
+*.swp
+
+# Python pickle files
+*.pkl
+
+# Reference videos
+!fastvideo/v1/tests/ssim/reference_videos/**/*.mp4
+
+# Static images
+!docs/source/_static/images/**/*.png

.pre-commit-config.yaml

@@ -19,6 +19,7 @@ exclude: |
     fastvideo/sample/.*|
     fastvideo/train\.py|
     fastvideo/utils/.*|
+    fastvideo/v1/examples/.*|
     .github/workflows/fastvideo-publish.yml|
     .github/workflows/sta-publish.yml
   )
@@ -41,7 +42,7 @@ repos:
     additional_dependencies: ['tomli']
     args: ['--toml', 'pyproject.toml']
 - repo: https://github.com/PyCQA/isort
-  rev: 0a0b7a830386ba6a31c2ec8316849ae4d1b8240d # 6.0.0
+  rev: 6.0.1
   hooks:
   - id: isort
 - repo: https://github.com/jackdewinter/pymarkdown
@@ -66,7 +67,7 @@ repos:
     entry: bash
     args:
       - -c
-      - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
+      - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
     language: system
     always_run: true
     pass_filenames: false

Dockerfile

@@ -0,0 +1,48 @@
+FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+    rm Miniconda3-latest-Linux-x86_64.sh
+
+ENV PATH=/opt/conda/bin:$PATH
+
+RUN conda create --name fastvideo-dev python=3.10.0 -y
+
+SHELL ["/bin/bash", "-c"]
+
+# Copy just the pyproject.toml first to leverage Docker cache
+COPY pyproject.toml ./
+
+# Create a dummy README to satisfy the installation
+RUN echo "# Placeholder" > README.md
+
+RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
+    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.0.post2 --no-build-isolation && \
+    conda clean -afy
+
+COPY . .
+
+RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
+
+# Remove authentication headers
+RUN git config --unset-all http.https://github.com/.extraheader || true
+
+# Set up automatic conda environment activation for all shells
+RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
+    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
+    # Ensure .bashrc is sourced for SSH login shells
+    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+
+EXPOSE 22

README.md

@@ -5,14 +5,15 @@
 FastVideo is a lightweight framework for accelerating large video diffusion models.
 
 <p align="center">
-    🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank">FastHunyuan</a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank">FastMochi</a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> Slack </a>
+    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
 </p>
 
 https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1
 
 FastVideo currently offers: (with more to come)
 
-- [NEW!] [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
+- [NEW!] V1 inference API available. Full announcement coming soon!
+- [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
 - FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
 - First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
 - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
@@ -26,211 +27,55 @@ Dev in progress and highly experimental.
 - ```2025/02/18```: Release the inference code and kernel for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
 - ```2025/01/13```: Support Lora finetuning for HunyuanVideo.
 - ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment.
-- ```2024/12/17```: `FastVideo` v1.0 is released.
+- ```2024/12/17```: `FastVideo` v0.0.1 is released.
 
-## 🔧 Installation from source
-The code is tested on Python 3.10.0, CUDA 12.4 and H100.
+## Getting Started
 
-```
-# Clone FastVideo
-git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
+- [Install FastVideo](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
+- [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)
 
-# Install FastVideo
-pip install -e .
+### Inference
+- [Quick Start](https://hao-ai-lab.github.io/FastVideo/inference/examples/basic.html)
+- V1 Inference API Guide (Coming soon!)
 
-# Install Flash Attention (optional)
-pip install flash-attn==2.7.0.post2
-```
+### Distillation and Finetuning
+- [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/training/distillation.html)
+- [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetuning.html)
 
-To try Sliding Tile Attention (optional), please follow the instruction in [csrc/sliding_tile_attention/README.md](csrc/sliding_tile_attention/README.md) to install STA.
-
-## 🚀 Inference
-### Inference StepVideo with Sliding Tile Attention
-First, download the model:
-
-```
-python scripts/huggingface/download_hf.py --repo_id=stepfun-ai/stepvideo-t2v --local_dir=data/stepvideo-t2v --repo_type=model
-```
-
-Use the following scripts to run inference for StepVideo. When using STA for inference, the generated videos will have dimensions of 204×768×768 (currently, this is the only supported shape).
-
-```bash
-sh scripts/inference/inference_stepvideo_STA.sh # Inference stepvideo with STA
-sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo
-```
-
-### Inference HunyuanVideo with Sliding Tile Attention
-First, download the model:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model
-```
-
-We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
-
-```bash
-sh scripts/inference/inference_hunyuan_STA.sh
-```
-
-### Video Demos using STA + Teacache
-Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100.
-
-### Inference FastHunyuan on single RTX4090
-We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model
-# CLI inference
-bash scripts/inference/inference_hunyuan_hf_quantization.sh
-```
-
-For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU):
-
-| Configuration                  | Memory to Init Transformer | Peak Memory After Init Pipeline (Denoise) | Diffusion Time | End-to-End Time |
-|--------------------------------|----------------------------|--------------------------------------------|----------------|-----------------|
-| BF16 + Pipeline CPU Offload    | 23.883G                   | 33.744G                                    | 81s            | 121.5s          |
-| INT8 + Pipeline CPU Offload    | 13.911G                   | 27.979G                                    | 88s            | 116.7s          |
-| NF4 + Pipeline CPU Offload     | 9.453G                    | 19.26G                                     | 78s            | 114.5s          |
-
-For improved quality in generated videos, we recommend using a GPU with 80GB of memory to run the BF16 model with the original Hunyuan pipeline. To execute the inference, use the following section:
-
-### FastHunyuan
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan --local_dir=data/FastHunyuan --repo_type=model
-# CLI inference
-bash scripts/inference/inference_hunyuan.sh
-```
-
-You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo).
-
-### FastMochi
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffusers --local_dir=data/FastMochi-diffusers --repo_type=model
-# CLI inference
-bash scripts/inference/inference_mochi_sp.sh
-```
-
-## 🎯 Distill
-Our distillation recipe is based on [Phased Consistency Model](https://github.com/G-U-N/Phased-Consistency-Model). We did not find significant improvement using multi-phase distillation, so we keep the one phase setup similar to the original latent consistency model's recipe.
-We use the [MixKit](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/all_mixkit) dataset for distillation. To avoid running the text encoder and VAE during training, we preprocess all data to generate text embeddings and VAE latents.
-Preprocessing instructions can be found [data_preprocess.md](docs/data_preprocess.md). For convenience, we also provide preprocessed data that can be downloaded directly using the following command:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/HD-Mixkit-Finetune-Hunyuan --local_dir=data/HD-Mixkit-Finetune-Hunyuan --repo_type=dataset
-```
-
-Next, download the original model weights with:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model # original hunyuan
-python scripts/huggingface/download_hf.py --repo_id=genmo/mochi-1-preview --local_dir=data/mochi --repo_type=model # original mochi
-```
-
-To launch the distillation process, use the following commands:
-
-```
-bash scripts/distill/distill_hunyuan.sh # for hunyuan
-bash scripts/distill/distill_mochi.sh # for mochi
-```
-
-We also provide an optional script for distillation with adversarial loss, located at `fastvideo/distill_adv.py`. Although we tried adversarial loss, we did not observe significant improvements.
-## Finetune
-### ⚡ Full Finetune
-Ensure your data is prepared and preprocessed in the format specified in [data_preprocess.md](docs/data_preprocess.md). For convenience, we also provide a mochi preprocessed Black Myth Wukong data that can be downloaded directly:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/Mochi-Black-Myth --local_dir=data/Mochi-Black-Myth --repo_type=dataset
-```
-
-Download the original model weights as specified in [Distill Section](#-distill):
-
-Then you can run the finetune with:
-
-```
-bash scripts/finetune/finetune_mochi.sh # for mochi
-```
-
-**Note that for finetuning, we did not tune the hyperparameters in the provided script.**
-### ⚡ Lora Finetune
-
-Hunyuan supports Lora fine-tuning of videos up to 720p. Demos and prompts of Black-Myth-Wukong can be found in [here](https://huggingface.co/FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight). You can download the Lora weight through:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight --local_dir=data/Hunyuan-Black-Myth-Wukong-lora-weight --repo_type=model
-```
-
-#### Minimum Hardware Requirement
-- 40 GB GPU memory each for 2 GPUs with lora.
-- 30 GB GPU memory each for 2 GPUs with CPU offload and lora.
-
-Currently, both Mochi and Hunyuan models support Lora finetuning through diffusers. To generate personalized videos from your own dataset, you'll need to follow three main steps: dataset preparation, finetuning, and inference.
-
-#### Dataset Preparation
-We provide scripts to better help you get started to train on your own characters!
-You can run this to organize your dataset to get the videos2caption.json before preprocess. Specify your video folder and corresponding caption folder (caption files should be .txt files and have the same name with its video):
-
-```
-python scripts/dataset_preparation/prepare_json_file.py --video_dir data/input_videos/ --prompt_dir data/captions/ --output_path data/output_folder/videos2caption.json --verbose
-```
-
-Also, we provide script to resize your videos:
-
-```
-python scripts/data_preprocess/resize_videos.py
-```
-
-#### Finetuning
-After basic dataset preparation and preprocess, you can start to finetune your model using Lora:
-
-```
-bash scripts/finetune/finetune_hunyuan_hf_lora.sh
-```
-
-#### Inference
-For inference with Lora checkpoint, you can run the following scripts with additional parameter `--lora_checkpoint_dir`:
-
-```
-bash scripts/inference/inference_hunyuan_hf.sh
-```
-
-**We also provide scripts for Mochi in the same directory.**
-
-#### Finetune with Both Image and Video
-Our codebase support finetuning with both image and video.
-
-```bash
-bash scripts/finetune/finetune_hunyuan.sh
-bash scripts/finetune/finetune_mochi_lora_mix.sh
-```
-
-For Image-Video Mixture Fine-tuning, make sure to enable the `--group_frame` option in your script.
+### Deprecated APIs
+- [V0 Inference (Deprecated)](https://hao-ai-lab.github.io/FastVideo/inference/v0_inference.html)
 
 ## 📑 Development Plan
 
-- More distillation methods
-  - [ ] Add Distribution Matching Distillation
+<!-- - More distillation methods -->
+  <!-- - [ ] Add Distribution Matching Distillation -->
 - More models support
-  - [ ] Add CogvideoX model
-- Code update
-  - [ ] fp8 support
-  - [ ] faster load model and save model support
+  <!-- - [ ] Add CogvideoX model -->
+  - [ ] Add StepVideo to V1
+- Optimization features
+  - [ ] Teacache in V1
+  - [ ] SageAttention in V1
+- Code updates
+  - [ ] V1 Configuration API
+  - [ ] Support Training in V1
+  <!-- - [ ] fp8 support -->
+  <!-- - [ ] faster load model and save model support -->
 
 ## 🤝 Contributing
 
-We welcome all contributions. Please run `bash format.sh --all` before submitting a pull request.
-
-## 🔧 Testing
-Run `pytest` to verify the data preprocessing, checkpoint saving, and sequence parallel pipelines. We recommend adding corresponding test cases in the `test` folder to support your contribution.
+We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/developer_guide/overview.html)
 
 ## Acknowledgement
-We learned and reused code from the following projects: [PCM](https://github.com/G-U-N/Phased-Consistency-Model), [diffusers](https://github.com/huggingface/diffusers), [OpenSoraPlan](https://github.com/PKU-YuanGroup/Open-Sora-Plan), and [xDiT](https://github.com/xdit-project/xDiT).
+We learned and reused code from the following projects:
+- [PCM](https://github.com/G-U-N/Phased-Consistency-Model)
+- [diffusers](https://github.com/huggingface/diffusers)
+- [OpenSoraPlan](https://github.com/PKU-YuanGroup/Open-Sora-Plan)
+- [xDiT](https://github.com/xdit-project/xDiT)
+- [vLLM](https://github.com/vllm-project/vllm)
+- [SGLang](https://github.com/sgl-project/sglang)
 
-We thank MBZUAI and Anyscale for their support throughout this project.
+We thank MBZUAI and [Anyscale](https://www.anyscale.com/) for their support throughout this project.
 
 ## Citation
 If you use FastVideo for your research, please cite our paper:

csrc/sliding_tile_attention/setup.py

@@ -9,7 +9,7 @@ target = target.lower()
 
 # Package metadata
 PACKAGE_NAME = "st_attn"
-VERSION = "0.0.2"
+VERSION = "0.0.4"
 AUTHOR = "Hao AI Lab"
 DESCRIPTION = "Sliding Tile Atteniton Kernel Used in FastVideo"
 URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/sliding_tile_attention"

csrc/sliding_tile_attention/st_attn.cpp

@@ -10,7 +10,7 @@
 
 #ifdef TK_COMPILE_ATTN
 extern torch::Tensor sta_forward(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_w_size, int kernel_h_size, int text_length, bool process_text, bool has_text
+    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_w_size, int kernel_h_size, int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag
 ); 
 #endif
 

csrc/sliding_tile_attention/st_attn/__init__.py

@@ -4,8 +4,13 @@ import torch
 from st_attn_cuda import sta_fwd
 
 
-def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True):
+def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True, img_latent_shape='30*48*80'):
     seq_length = q_all.shape[2]
+    img_latent_shape_mapping = {
+        '30x48x80':1,
+        '36x48x48':2,
+        '18x48x80':3, 
+    }
     if has_text:
         assert q_all.shape[
             2] >= 115200, "STA currently only supports video with latent size (30, 48, 80), which is 117 frames x 768 x 1280 pixels"
@@ -17,8 +22,14 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te
             k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2)
             v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2)
     else:
-        assert q_all.shape[2] == 82944
+        if img_latent_shape == '36x48x48': # Stepvideo 204x768x68
+            assert q_all.shape[2] == 82944
+        elif img_latent_shape == '18x48x80': # Wan 69x768x1280
+            assert q_all.shape[2] == 69120
+        else:
+            raise ValueError(f"Unsupported {img_latent_shape}, current shape is {q_all.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan")
 
+    kernel_aspect_ratio_flag = img_latent_shape_mapping[img_latent_shape]
     hidden_states = torch.empty_like(q_all)
     # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
     for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size):
@@ -29,7 +40,7 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te
                                                                                       head_index:head_index + 1],
                                               hidden_states[batch:batch + 1, head_index:head_index + 1])
 
-            _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text)
+            _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text, kernel_aspect_ratio_flag)
     if has_text:
-        _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True)
+        _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True, kernel_aspect_ratio_flag)
     return hidden_states[:, :, :seq_length]

csrc/sliding_tile_attention/st_attn/st_attn_h100.cu

@@ -359,7 +359,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
 #include <iostream>
 
 torch::Tensor 
-sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_h_size, int kernel_w_size, int text_length, bool process_text, bool has_text)
+sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_h_size, int kernel_w_size, int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag)
 {
     CHECK_INPUT(q);
     CHECK_INPUT(k);
@@ -446,7 +446,7 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
         auto threads  = NUM_WORKERS * kittens::WARP_THREADS;
         if (has_text) {
             // TORCH_CHECK(seq_len % (CONSUMER_WARPGROUPS*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 192");
-            dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4-2), qo_heads, batch);
+            dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4)-2, qo_heads, batch);
             dim3 grid_text(2, qo_heads, batch);
             if (!process_text) {
                 if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
@@ -558,123 +558,267 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
 
         } else {
             dim3 grid_image(seq_len/(CONSUMER_WARPGROUPS*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
+            if (kernel_aspect_ratio_flag == 2){
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 1, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            }  else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 6) {
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 1, 1, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false,1, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            } else if (kernel_t_size == 6 && kernel_h_size == 3 && kernel_w_size == 3) {
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 1, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            } else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 1, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                }else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 3){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            }else if (kernel_t_size ==3 && kernel_h_size == 6 && kernel_w_size == 3){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 1, 3, 1, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size ==6 && kernel_h_size == 3 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
 
-            } else if (kernel_t_size ==6 && kernel_h_size == 3 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 1, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                }  else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                }  else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else {
+                    // print error
+                    std::cout << "Invalid kernel size" << std::endl;
+                    //print kernel size
+                    std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+                }
+            }
+            else if (kernel_aspect_ratio_flag == 3) {
+                if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 3) {
 
-            } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 1){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 0, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            }  else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 0, 3, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            }  else if (kernel_t_size == 1 && kernel_h_size == 1 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 0, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 1){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 0, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 6 && kernel_h_size == 6 && kernel_w_size == 1){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 3, 0, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else if (kernel_t_size == 6 && kernel_h_size == 1 && kernel_w_size == 6){
-                cudaFuncSetAttribute(
-                    fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6>,
-                    cudaFuncAttributeMaxDynamicSharedMemorySize,
-                    mem_size
-                );
-                fwd_attend_ker<128, false, false, false, 3, 0, 3, 6, 6, 6><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
-            } else {
-                // print error
-                std::cout << "Invalid kernel size" << std::endl;
-                //print kernel size
-                std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 1, 1, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 5) {
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 2, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 1, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 2, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 0, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+
+                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 2, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 5 && kernel_w_size == 9){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 2, 4, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 6 && kernel_w_size == 3){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 3, 1, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 1){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 1, 0, 0, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 2, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 3, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 7){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 2, 3, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 5 && kernel_w_size == 9){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false, 0, 2, 4, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 1 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 0, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 0, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 3 && kernel_h_size == 3 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 1, 1, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,1, 1, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 3 && kernel_w_size == 10){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 1, 5, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,0, 1, 5, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else if (kernel_t_size == 1 && kernel_h_size == 6 && kernel_w_size == 5){
+                    cudaFuncSetAttribute(
+                        fwd_attend_ker<128, false, false, false, 0, 3, 2, 3, 6, 10>,
+                        cudaFuncAttributeMaxDynamicSharedMemorySize,
+                        mem_size
+                    );
+                    fwd_attend_ker<128, false, false, false,0, 3, 2, 3, 6, 10><<<grid_image, (32*NUM_WORKERS), mem_size, stream>>>(g);
+                } else {
+                    // print error
+                    std::cout << "Invalid kernel size" << std::endl;
+                    //print kernel size
+                    std::cout << "Kernel size: " << kernel_t_size << " " << kernel_h_size << " " << kernel_w_size << std::endl;
+                }
+            }
+
+            else {
+                std::cout << "Unsupported kernel_aspect_ratio_flag: " << kernel_aspect_ratio_flag << std::endl;
             }
 
         }

docs/Makefile

@@ -22,3 +22,4 @@ help:
 clean:
 	@$(SPHINXBUILD) -M clean "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 	rm -rf "$(SOURCEDIR)/getting_started/examples"
+	rm -rf "$(SOURCEDIR)/inference/examples"

docs/source/_templates/sections/header.html

@@ -34,6 +34,6 @@
   }
 </style>
 
-<div class="notification-bar">
+<!-- <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>
+</div> -->

docs/source/contributing/add_pipeline.md

@@ -0,0 +1,316 @@
+(add-pipeline)=
+
+# 🏗️ Adding a New Diffusion Pipeline
+
+This guide explains how to implement a custom diffusion pipeline in FastVideo, leveraging the framework's modular architecture for high-performance video generation.
+
+## Implementation Process Overview
+
+1. **Port Required Modules** - Identify and implement necessary model components
+2. **Create Directory Structure** - Set up pipeline files and folders
+3. **Implement Pipeline Class** - Build the pipeline using existing or custom stages
+4. **Register Your Pipeline** - Make it discoverable by the framework
+5. **Configure Your Pipeline** - (Coming soon)
+
+Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).
+
+## Step 1: Pipeline Modules
+
+### Identifying Required Modules
+
+FastVideo uses the Hugging Face Diffusers format for model organization:
+
+1. Examine the `model_index.json` in the HF model repository:
+
+```json
+{
+    "_class_name": "WanImageToVideoPipeline",
+    "_diffusers_version": "0.33.0.dev0",
+    "image_encoder": ["transformers", "CLIPVisionModelWithProjection"],
+    "image_processor": ["transformers", "CLIPImageProcessor"],
+    "scheduler": ["diffusers", "UniPCMultistepScheduler"],
+    "text_encoder": ["transformers", "UMT5EncoderModel"],
+    "tokenizer": ["transformers", "T5TokenizerFast"],
+    "transformer": ["diffusers", "WanTransformer3DModel"],
+    "vae": ["diffusers", "AutoencoderKLWan"]
+}
+```
+
+1. For each component:
+   - Note the originating library (`transformers` or `diffusers`)
+   - Identify the class name
+   - Check if it's already available in FastVideo
+
+2. Review config files in each component's directory for architecture details
+
+### Implementing Modules
+
+Place new modules in the appropriate directories:
+- Encoders: `fastvideo/v1/models/encoders/`
+- VAEs: `fastvideo/v1/models/vaes/`
+- Transformer models: `fastvideo/v1/models/dits/`
+- Schedulers: `fastvideo/v1/models/schedulers/`
+
+### Adapting Model Layers
+
+#### Layer Replacements
+Replace standard PyTorch layers with FastVideo optimized versions:
+- nn.LayerNorm → fastvideo.v1.layers.layernorm.RMSNorm
+- Embedding layers → fastvideo.v1.layers.vocab_parallel_embedding modules
+- Activation functions → versions from fastvideo.v1.layers.activation
+
+#### Distributed Linear Layers
+Use appropriate parallel layers for distribution:
+
+```python
+# Output dimension parallelism
+from fastvideo.v1.layers.linear import ColumnParallelLinear
+self.q_proj = ColumnParallelLinear(
+    input_size=hidden_size,
+    output_size=head_size * num_heads,
+    bias=bias,
+    gather_output=False
+)
+
+# Fused QKV projection
+from fastvideo.v1.layers.linear import QKVParallelLinear
+self.qkv_proj = QKVParallelLinear(
+    hidden_size=hidden_size,
+    head_size=attention_head_dim,
+    total_num_heads=num_attention_heads,
+    bias=True
+)
+
+# Input dimension parallelism
+from fastvideo.v1.layers.linear import RowParallelLinear
+self.out_proj = RowParallelLinear(
+    input_size=head_size * num_heads,
+    output_size=hidden_size,
+    bias=bias,
+    input_is_parallel=True
+)
+```
+
+### Attention Layers
+Replace standard attention with FastVideo's optimized attention:
+
+```python
+# Local attention patterns
+from fastvideo.v1.attention import LocalAttention
+from fastvideo.v1.attention.backends.abstract import _Backend
+self.attn = LocalAttention(
+    num_heads=num_heads,
+    head_size=head_dim,
+    dropout_rate=0.0,
+    softmax_scale=None,
+    causal=False,
+    supported_attention_backends=(_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+)
+
+# Distributed attention for long sequences
+from fastvideo.v1.attention import DistributedAttention
+self.attn = DistributedAttention(
+    num_heads=num_heads,
+    head_size=head_dim,
+    dropout_rate=0.0,
+    softmax_scale=None,
+    causal=False,
+    supported_attention_backends=(_Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+)
+```
+
+#### Define supported backend selection
+
+```python
+   _supported_attention_backends = (_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+```
+
+### Registering Models
+
+Register implemented modules in the model registry:
+
+```python
+# In fastvideo/v1/models/registry.py
+_TEXT_TO_VIDEO_DIT_MODELS = {
+    "YourTransformerModel": ("dits", "yourmodule", "YourTransformerClass"),
+}
+
+_VAE_MODELS = {
+    "YourVAEModel": ("vaes", "yourvae", "YourVAEClass"),
+}
+```
+
+## Step 2: Directory Structure
+
+Create a new directory for your pipeline:
+
+```
+fastvideo/v1/pipelines/
+├── your_pipeline/
+│   ├── __init__.py
+│   └── your_pipeline.py
+```
+
+## Step 3: Implement Pipeline Class
+
+Pipelines are composed of stages, each handling a specific part of the diffusion process:
+
+- **InputValidationStage**: Validates input parameters
+- **Text Encoding Stages**: Handle text encoding (CLIP/Llama/T5)
+- **CLIPImageEncodingStage**: Processes image inputs
+- **TimestepPreparationStage**: Prepares diffusion timesteps
+- **LatentPreparationStage**: Manages latent representations
+- **ConditioningStage**: Processes conditioning inputs
+- **DenoisingStage**: Performs denoising diffusion
+- **DecodingStage**: Converts latents to pixels
+
+### Creating Your Pipeline
+
+```python
+from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.v1.pipelines.stages import (
+    InputValidationStage, CLIPTextEncodingStage, TimestepPreparationStage,
+    LatentPreparationStage, DenoisingStage, DecodingStage
+)
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+import torch
+
+class MyCustomPipeline(ComposedPipelineBase):
+    """Custom diffusion pipeline implementation."""
+    
+    # Define required model components from model_index.json
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+    
+    @property
+    def required_config_modules(self) -> List[str]:
+        return self._required_config_modules
+        
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        """Initialize pipeline-specific components."""
+        pass
+        
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+        self.add_stage(
+            stage_name="input_validation_stage",
+            stage=InputValidationStage()
+        )
+        
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=CLIPTextEncodingStage(
+                text_encoder=self.get_module("text_encoder"),
+                tokenizer=self.get_module("tokenizer")
+            )
+        )
+        
+        self.add_stage(
+            stage_name="timestep_preparation_stage",
+            stage=TimestepPreparationStage(
+                scheduler=self.get_module("scheduler")
+            )
+        )
+        
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=LatentPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                vae=self.get_module("vae")
+            )
+        )
+        
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler")
+            )
+        )
+        
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=DecodingStage(
+                vae=self.get_module("vae")
+            )
+        )
+    
+# Register the pipeline class
+EntryClass = MyCustomPipeline
+```
+
+### Creating Custom Stages (Optional)
+
+If existing stages don't meet your needs, create custom ones:
+
+```python
+from fastvideo.v1.pipelines.stages.base import PipelineStage
+
+class MyCustomStage(PipelineStage):
+    """Custom processing stage for the pipeline."""
+    
+    def __init__(self, custom_module, other_param=None):
+        super().__init__()
+        self.custom_module = custom_module
+        self.other_param = other_param
+        
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        # Access input data
+        input_data = batch.some_attribute
+        
+        # Validate inputs
+        if input_data is None:
+            raise ValueError("Required input is missing")
+            
+        # Process with your module
+        result = self.custom_module(input_data)
+        
+        # Update batch with results
+        batch.some_output = result
+        
+        return batch
+```
+
+Add your custom stage to the pipeline:
+
+```python
+self.add_stage(
+    stage_name="my_custom_stage",
+    stage=MyCustomStage(
+        custom_module=self.get_module("custom_module"),
+        other_param="some_value"
+    )
+)
+```
+
+#### Stage Design Principles
+
+1. **Single Responsibility**: Focus on one specific task
+2. **Functional Pattern**: Receive and return a `ForwardBatch` object
+3. **Dependency Injection**: Pass dependencies through constructor
+4. **Input Validation**: Validate inputs for clear error messages
+
+## Step 4: Register Your Pipeline
+
+Define `EntryClass` at the end of your pipeline file:
+
+```python
+# Single pipeline class
+EntryClass = MyCustomPipeline
+
+# Or multiple pipeline classes
+EntryClass = [MyCustomPipeline, MyOtherPipeline]
+```
+
+The registry will automatically:
+1. Scan all packages under `fastvideo/v1/pipelines/`
+2. Look for `EntryClass` variables
+3. Register pipelines using their class names as identifiers
+
+## Best Practices
+
+- **Reuse Existing Components**: Leverage built-in stages and modules
+- **Follow Module Organization**: Place new modules in appropriate directories
+- **Match Model Patterns**: Follow existing code patterns and conventions

docs/source/contributing/developer_env/docker.md

@@ -0,0 +1,31 @@
+# 🐳 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:
+
+**Image:** [`ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest`](https://ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev)
+
+## Starting the container
+
+```bash
+docker run --gpus all -it ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
+```
+
+This will:
+
+- Start the container with GPU access  
+- Drop you into a shell with the `fastvideo-dev` Conda environment preconfigured
+
+## Using the container
+
+```bash
+# Conda environment should already be active
+# FastVideo package installed in editable mode
+
+# Pull the latest changes from remote
+cd /FastVideo
+git pull
+
+# Run linters and tests
+pre-commit run --all-files
+pytest tests/
+```

docs/source/contributing/developer_env/index.md

@@ -0,0 +1,13 @@
+(developer-env)
+
+# 🧰 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
+:::

docs/source/contributing/developer_env/runpod.md

@@ -0,0 +1,52 @@
+(runpod)=
+
+# 📦 Developing FastVideo on RunPod
+
+You can easily use the FastVideo Docker image as a custom container on [RunPod](https://www.runpod.io) for development or experimentation.
+
+## Creating a new pod
+
+Choose a GPU that supports CUDA 12.4
+
+![RunPod CUDA selection](../../_static/images/runpod_cuda.png)
+
+When creating your pod template, use this image:
+
+```
+ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
+```
+
+Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.io/pods/configuration/use-ssh)):
+
+```bash
+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)
+
+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)
+
+## Working with the pod
+
+After SSH'ing into your pod, you'll find the `fastvideo-dev` Conda environment already activated.
+
+To pull in the latest changes from the GitHub repo:
+
+```bash
+cd /FastVideo
+git pull
+```
+
+`If you have a persistent volume and want to keep your code changes, you can move /FastVideo to /workspace/FastVideo, or simply clone the repository there.`
+
+Run your development workflows as usual:
+
+```bash
+# Run linters
+pre-commit run --all-files
+
+# Run tests
+pytest tests/
+```

docs/source/contributing/overview.md

@@ -1,4 +1,6 @@
-# Contributing to FastVideo
+(developer-overview)=
+
+# 🛠️ Contributing to FastVideo
 
 Thank you for your interest in contributing to FastVideo. We want to make the process as smooth for you as possible and this is a guide to help get you started!
 
@@ -27,7 +29,7 @@ conda activate fastvideo
 Clone the FastVideo repository and go to the FastVideo directory:
 
 ```
-git clone https://github.com/vllm-project/vllm.git && cd vllm
+git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 
 ```
 

docs/source/design/overview.md

@@ -0,0 +1,410 @@
+# 🔍 FastVideo Overview
+
+This document outlines FastVideo's architecture for developers interested in framework internals or contributions. It serves as an onboarding guide for new contributors by providing an overview of the most important directories and files within the `fastvideo/v1/` codebase.
+
+## Table of Contents - V1 Directory Structure and Files
+
+- [`fastvideo/v1/pipelines/`](#design-pipeline-system) - Core diffusion pipeline components
+- [`fastvideo/v1/models/`](#design-model-components) - Model implementations
+  - [`dits/`](#design-transformer-models) - Transformer-based diffusion models
+  - [`vaes/`](#design-vae-variational-auto-encoder) - Variational autoencoders
+  - [`encoders/`](#design-text-and-image-encoders) - Text and image encoders
+  - [`schedulers/`](#design-schedulers) - Diffusion schedulers
+- [`fastvideo/v1/attention/`](#design-optimized-attention) - Optimized attention implementations
+- [`fastvideo/v1/distributed/`](#design-distributed-processing) - Distributed computing utilities
+- [`fastvideo/v1/layers/`](#design-tensor-parallelism) - Custom neural network layers
+- [`fastvideo/v1/platforms/`](#design-platforms) - Hardware platform abstractions
+- [`fastvideo/v1/worker/`](#design-executor-and-worker-abstractions) - Multi-GPU process management
+- [`fastvideo/v1/fastvideo_args.py`](#design-fastvideo-args) - Argument handling
+- [`fastvideo/v1/forward_context.py`](#design-forwardcontext) - Forward pass context management
+- `fastvideo/v1/utils.py` - Utility functions
+- [`fastvideo/v1/logger.py`](#design-logger) - Logging infrastructure
+
+## Core Architecture
+
+FastVideo separates model components from execution logic with these principles:
+- **Component Isolation**: Models (encoders, VAEs, transformers) are isolated from execution (pipelines, stages, distributed processing)
+- **Modular Design**: Components can be independently replaced
+- **Distributed Execution**: Supports various parallelism strategies (Tensor, Sequence)
+- **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/v1/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.
+
+Key features include:
+- **Command-line Interface**: Automatic conversion between CLI arguments and dataclass fields
+- **Configuration Groups**: Organized by functional areas (model loading, video params, optimization settings)
+- **Context Management**: Global access to current settings via `get_current_fastvideo_args()`
+- **Parameter Validation**: Ensures valid combinations of settings
+
+Common configuration areas:
+- **Model paths and loading options**: `model_path`, `trust_remote_code`, `revision`
+- **Distributed execution settings**: `num_gpus`, `tp_size`, `sp_size`
+- **Video generation parameters**: `height`, `width`, `num_frames`, `num_inference_steps`
+- **Precision settings**: Control computation precision for different components
+
+Example usage:
+
+```python
+# Load arguments from command line
+fastvideo_args = prepare_fastvideo_args(sys.argv[1:])
+
+# Access parameters
+model = load_model(fastvideo_args.model_path)
+
+# Set as global context
+with set_current_fastvideo_args(fastvideo_args):
+    # Code that requires access to these arguments
+    result = generate_video()
+```
+
+(design-pipeline-system)=
+## Pipeline System
+
+### `ComposedPipelineBase`
+
+This foundational class provides:
+
+- **Model Loading**: Automatically loads components from HuggingFace-Diffusers-compatible model directories
+- **Stage Management**: Creates and orchestrates processing stages
+- **Data Flow Coordination**: Ensures proper state flow between stages
+
+```python
+class MyCustomPipeline(ComposedPipelineBase):
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae", "transformer", "scheduler"
+    ]
+    
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+        # Pipeline-specific initialization
+        pass
+        
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        self.add_stage("input_validation_stage", InputValidationStage())
+        self.add_stage("text_encoding_stage", CLIPTextEncodingStage(
+            text_encoder=self.get_module("text_encoder"),
+            tokenizer=self.get_module("tokenizer")
+        ))
+        # Additional stages...
+```
+
+### Pipeline Stages
+Each stage handles a specific diffusion process component:
+- **Input Validation**: Parameter verification
+- **Text Encoding**: CLIP, LLaMA, or T5-based encoding
+- **Image Encoding**: Image input processing
+- **Timestep & Latent Preparation**: Setup for diffusion
+- **Denoising**: Core diffusion loop
+- **Decoding**: Latent-to-pixel conversion
+
+Each stage implements a standard interface:
+
+```python
+def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+    # Process batch and update state
+    return batch
+```
+
+(design-forwardbatch)=
+### ForwardBatch
+
+Defined in `fastvideo/v1/pipelines/pipeline_batch_info.py`, `ForwardBatch` encapsulates the data payload passed between pipeline stages. It typically holds:
+
+- **Input Data**: Prompts, images, generation parameters
+- **Intermediate State**: Embeddings, latents, timesteps, accumulated during stage execution
+- **Output Storage**: Generated results and metadata
+- **Configuration**: Sampling parameters, precision settings
+
+This structure facilitates clear state transitions between stages.
+
+(design-model-components)=
+## Model Components
+
+The `fastvideo/v1/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:
+
+- **Location**: `fastvideo/v1/models/dits/`
+- **Examples**:
+  - `WanTransformer3DModel`
+  - `HunyuanVideoTransformer3DModel`
+
+Features include:
+- Text/image conditioning
+- Standardized interface for model-specific optimizations
+
+```python
+def forward(
+    self, 
+    latents,                    # [B, T, C, H, W]
+    encoder_hidden_states,      # Text embeddings
+    timestep,                   # Current diffusion timestep
+    encoder_hidden_states_image=None,  # Optional image embeddings
+    **kwargs
+):
+    # Perform denoising computation
+    return noise_pred  # Predicted noise residual
+```
+
+(design-vae-variational-auto-encoder)=
+### VAE (Variational Auto-Encoder)
+
+VAEs handle conversion between pixel space and latent space:
+
+- **Location**: `fastvideo/v1/models/vaes/`
+- **Examples**:
+  - `AutoencoderKLWan`
+  - `AutoencoderKLHunyuanVideo`
+
+These models compress image/video data to a more efficient latent representation (typically 4x-8x smaller in each dimension).
+
+FastVideo's VAE implementations include:
+- Efficient video batch processing
+- Memory optimization
+- Optional tiling for large frames
+- Distributed weight support
+
+(design-text-and-image-encoders)=
+### Text and Image Encoders
+
+Encoders process conditioning inputs into embeddings:
+
+- **Location**: `fastvideo/v1/models/encoders/`
+- **Text Encoders**:
+  - `CLIPTextModel`
+  - `LlamaModel`
+  - `UMT5EncoderModel`
+- **Image Encoders**:
+  - `CLIPVisionModel`
+
+FastVideo implements optimizations such as:
+- Vocab parallelism for distributed processing
+- Caching for common prompts
+- Precision-tuned computation
+
+(design-schedulers)=
+### Schedulers
+
+Schedulers manage the diffusion sampling process:
+
+- **Location**: `fastvideo/v1/models/schedulers/`
+- **Examples**:
+  - `UniPCMultistepScheduler`
+  - `FlowMatchEulerDiscreteScheduler`
+
+These components control:
+- Diffusion timestep sequences
+- Noise prediction to latent update conversions
+- Quality/speed trade-offs
+
+```python
+def step(
+    self, 
+    model_output: torch.Tensor,
+    timestep: torch.LongTensor,
+    sample: torch.Tensor,
+    **kwargs
+) -> torch.Tensor:
+    # Process model output and update latents
+    # Return updated latents
+    return prev_sample
+```
+
+(design-optimized-attention)=
+## Optimized Attention
+
+The `fastvideo/v1/attention/` directory contains optimized attention implementations crucial for efficient video diffusion:
+
+### Attention Backends
+Multiple implementations with automatic selection:
+- **FLASH_ATTN**: Optimized for supporting hardware
+- **TORCH_SDPA**: Built-in PyTorch scaled dot-product attention
+- **SLIDING_TILE_ATTN**: For very long sequences
+
+```python
+# Configure available attention backends for this layer
+self.attn = LocalAttention(
+    num_heads=num_heads,
+    head_size=head_dim,
+    causal=False,
+    supported_attention_backends=(_Backend.FLASH_ATTN, _Backend.TORCH_SDPA)
+)
+
+# Override via environment variable
+# export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
+```
+
+### 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/v1/distributed/` directory contains implementations for distributed model execution:
+
+(design-tensor-parallelism)=
+### Tensor Parallelism
+
+Tensor parallelism splits model weights across devices:
+
+- **Implementation**: Through `RowParallelLinear` and `ColumnParallelLinear` layers
+- **Use cases**: Will be used by encoder models as their sequence lengths are shorter and enables efficient sharding.
+
+```python
+# Tensor-parallel layers in a transformer block
+from fastvideo.v1.layers.linear import ColumnParallelLinear, RowParallelLinear
+
+# Split along output dimension
+self.qkv_proj = ColumnParallelLinear(
+    input_size=hidden_size,
+    output_size=3 * hidden_size,
+    bias=True,
+    gather_output=False
+)
+
+# Split along input dimension
+self.out_proj = RowParallelLinear(
+    input_size=hidden_size,
+    output_size=hidden_size,
+    bias=True,
+    input_is_parallel=True
+)
+```
+
+### Sequence Parallelism
+
+Sequence parallelism splits sequences across devices:
+
+- **Implementation**: Through `DistributedAttention` and sequence splitting
+- **Use cases**: Long video sequences or high-resolution processing. Used by DiT models.
+
+```python
+# Distributed attention for long sequences
+from fastvideo.v1.attention import DistributedAttention
+
+self.attn = DistributedAttention(
+    num_heads=num_heads,
+    head_size=head_dim,
+    causal=False,
+    supported_attention_backends=(_Backend.SLIDING_TILE_ATTN, _Backend.FLASH_ATTN)
+)
+```
+
+### Communication Primitives
+Efficient distributed operations via AllGather, AllReduce, and synchronization mechanisms.
+
+Efficient communication primitives minimize distributed overhead:
+
+- **Sequence-Parallel AllGather**: Collects sequence chunks
+- **Tensor-Parallel AllReduce**: Combines partial results
+- **Distributed Synchronization**: Coordinates execution
+
+(design-forwardcontext)=
+## Forward Context Management
+
+### ForwardContext
+
+Defined in `fastvideo/v1/forward_context.py`, `ForwardContext` manages execution-specific state *within* a forward pass, particularly for low-level optimizations. It is accessed via `get_forward_context()`.
+
+- **Attention Metadata**: Configuration for optimized attention kernels (`attn_metadata`)
+- **Profiling Data**: Potential hooks for performance metrics collection
+
+This context-based approach enables:
+- Dynamic optimization based on execution state (e.g., attention backend selection)
+- Step-specific customizations within model components
+
+Usage example:
+
+```python
+with set_forward_context(current_timestep, attn_metadata, fastvideo_args):
+    # During this forward pass, components can access context
+    # through get_forward_context()
+    output = model(inputs)
+```
+
+(design-executor-and-worker-abstractions)=
+## Executor and Worker System
+
+The `fastvideo/v1/worker/` directory contains the distributed execution framework:
+
+### Executor Abstraction
+
+FastVideo implements a flexible execution model for distributed processing:
+
+- **Executor Base Class**: An abstract base class defining the interface for all executors
+- **MultiProcExecutor**: Primary implementation that spawns and manages worker processes
+- **GPU Workers**: Handle actual model execution on individual GPUs
+
+The MultiProcExecutor implementation:
+1. Spawns worker processes for each GPU
+2. Establishes communication channels via pipes
+3. Coordinates distributed operations across workers
+4. Handles graceful startup and shutdown of the process group
+
+Each GPU worker:
+1. Initializes the distributed environment
+2. Builds the pipeline for the specified model
+3. Executes requested operations on its assigned GPU
+4. Manages local resources and communicates results back to the executor
+
+This design allows FastVideo to efficiently utilize multiple GPUs while providing a simple, unified interface for model execution.
+
+(design-platforms)=
+## Platforms
+
+The `fastvideo/v1/platforms/` directory provides hardware platform abstractions that enable FastVideo to run efficiently on different hardware configurations:
+
+### Platform Abstraction
+
+FastVideo's platform abstraction layer enables:
+- **Hardware Detection**: Automatic detection of available hardware
+- **Backend Selection**: Appropriate selection of compute kernels
+- **Memory Management**: Efficient utilization of hardware-specific memory features
+
+The primary components include:
+- **Platform Interface**: Defines the common API for all platform implementations
+- **CUDA Platform**: Optimized implementation for NVIDIA GPUs
+- **Backend Enum**: Used throughout the codebase for feature selection
+
+Usage example:
+
+```python
+from fastvideo.v1.platforms import current_platform, _Backend
+
+# Check hardware capabilities
+if current_platform.supports_backend(_Backend.FLASH_ATTN):
+    # Use FlashAttention implementation
+else:
+    # Fall back to standard implementation
+```
+
+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)
+
+*TODO*: (help wanted) Add an environment variable that disables process-aware logging.
+
+## Contributing to FastVideo
+
+If you're a new contributor, here are some common areas to explore:
+
+1. **Adding a new model**: Implement new model types in the appropriate subdirectory of `fastvideo/v1/models/`
+2. **Optimizing performance**: Look at attention implementations or memory management
+3. **Adding a new pipeline**: Create a new pipeline subclass in `fastvideo/v1/pipelines/`
+4. **Hardware support**: Extend the `platforms` module for new hardware targets
+
+When adding code, follow these practices:
+- Use type hints for better code readability
+- Add appropriate docstrings
+- Maintain the separation between model components and execution logic
+- Follow existing patterns for distributed processing

docs/source/generate_examples.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/docs/source/generate_examples.py
 
 import itertools
 import re
@@ -8,7 +9,7 @@ from typing import Optional
 
 ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
 ROOT_DIR_RELATIVE = '../../../..'
-EXAMPLE_DIR = ROOT_DIR / "examples"
+EXAMPLE_DIR = ROOT_DIR / "fastvideo/v1/examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples"
 
 
@@ -161,52 +162,54 @@ class Example:
         return content
 
 
-def generate_examples():
-    # Create the EXAMPLE_DOC_DIR if it doesn't exist
-    if not EXAMPLE_DOC_DIR.exists():
-        EXAMPLE_DOC_DIR.mkdir(parents=True)
+def generate_examples(generate_main_index=False):
+    """
+    Generate example documentation.
+    
+    Args:
+        generate_main_index (bool): Whether to generate the main examples index.
+            If False, only category-specific indices will be generated.
+    """
+    # Create empty indices with dynamic paths
+    main_index_dir = ROOT_DIR / "docs/source/examples"
+    if not main_index_dir.exists():
+        main_index_dir.mkdir(parents=True)
 
-    # Create empty indices
-    examples_index = Index(
-        path=EXAMPLE_DOC_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>.",  # noqa: E501
-        caption="Examples",
-        maxdepth=2)
-    # Category indices stored in reverse order because they are inserted into
-    # examples_index.documents at index 0 in order
+    # Create the main examples index only if requested
+    examples_index = None
+    if generate_main_index:
+        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>.",  # noqa: E501
+            caption="Examples",
+            maxdepth=2)
+
+    # Category indices with dynamic paths based on category names
     category_indices = {
-        "other":
+        "inference":
         Index(
-            path=EXAMPLE_DOC_DIR / "examples_other_index.md",
-            title="Other",
+            path=ROOT_DIR /
+            "docs/source/inference/examples/examples_inference_index.md",
+            title="🚀 Examples",
             description=
-            "Other examples that don't strongly fit into the online or offline serving categories.",  # noqa: E501
-            caption="Examples",
-        ),
-        "online_serving":
-        Index(
-            path=EXAMPLE_DOC_DIR / "examples_online_serving_index.md",
-            title="Online Serving",
-            description=
-            "Online serving examples demonstrate how to use FastVideo in an online setting, where the model is queried for predictions in real-time.",  # noqa: E501
-            caption="Examples",
-        ),
-        "offline_inference":
-        Index(
-            path=EXAMPLE_DOC_DIR / "examples_offline_inference_index.md",
-            title="Offline Inference",
-            description=
-            "Offline inference examples demonstrate how to use FastVideo in an offline setting, where the model is queried for predictions in batches. We recommend starting with <project:basic.md>.",  # noqa: E501
+            "Inference examples demonstrate how to use FastVideo in an offline setting, where the model is queried for predictions in batches. We recommend starting with <project:basic.md>.",  # noqa: E501
             caption="Examples",
         ),
     }
 
+    # Ensure all category doc directories exist
+    for category, index in category_indices.items():
+        category_dir = index.path.parent
+        if not category_dir.exists():
+            category_dir.mkdir(parents=True)
+
     examples = []
     glob_patterns = ["*.py", "*.md", "*.sh"]
     # Find categorised examples
     for category in category_indices:
+        print(category)
         category_dir = EXAMPLE_DIR / category
         globs = [category_dir.glob(pattern) for pattern in glob_patterns]
         for path in itertools.chain(*globs):
@@ -214,33 +217,58 @@ def generate_examples():
         # Find examples in subdirectories
         for path in category_dir.glob("*/*.md"):
             examples.append(Example(path.parent, category))
-    # Find uncategorised examples
-    globs = [EXAMPLE_DIR.glob(pattern) for pattern in glob_patterns]
-    for path in itertools.chain(*globs):
-        examples.append(Example(path))
-    # Find examples in subdirectories
-    for path in EXAMPLE_DIR.glob("*/*.md"):
-        # Skip categorised examples
-        if path.parent.name in category_indices:
-            continue
-        examples.append(Example(path.parent))
 
-    # Generate the example documentation
+    # Find uncategorised examples only if we're generating a main index
+    if generate_main_index:
+        globs = [EXAMPLE_DIR.glob(pattern) for pattern in glob_patterns]
+        for path in itertools.chain(*globs):
+            examples.append(Example(path))
+        # Find examples in subdirectories
+        for path in EXAMPLE_DIR.glob("*/*.md"):
+            # Skip categorised examples
+            if path.parent.name in category_indices:
+                continue
+            examples.append(Example(path.parent))
+
+    # Create document directories for each category based on category name and generate files
     for example in sorted(examples, key=lambda e: e.path.stem):
-        doc_path = EXAMPLE_DOC_DIR / f"{example.path.stem}.md"
+        print(example)
+
+        # Determine which index to use for this example
+        if example.category is not None and example.category in category_indices:
+            index = category_indices[example.category]
+        elif generate_main_index:
+            assert examples_index is not None
+            index = examples_index  # Default to main index if available
+        else:
+            # Skip examples without a category if no main index
+            print(f"Skipping {example.path} (no category and no main index)")
+            continue
+
+        # Place generated example markdown in the same directory as its index
+        doc_path = index.path.parent / f"{example.path.stem}.md"
         with open(doc_path, "w+") as f:
             f.write(example.generate())
-        # Add the example to the appropriate index
-        assert example.category is not None
-        index = category_indices.get(example.category, examples_index)
+        # Add the example to the index
         index.documents.append(example.path.stem)
 
-    # Generate the index files
+    # Generate the index files for categories
     for category_index in category_indices.values():
         if category_index.documents:
-            examples_index.documents.insert(0, category_index.path.name)
+            # Add to main index if it exists
+            if generate_main_index:
+                rel_path = category_index.path.relative_to(
+                    main_index_dir.parent)
+                assert examples_index is not None
+                examples_index.documents.insert(
+                    0,
+                    str(rel_path).replace(".md", ""))
+
+            # Write the category index file
             with open(category_index.path, "w+") as f:
                 f.write(category_index.generate())
 
-    with open(examples_index.path, "w+") as f:
-        f.write(examples_index.generate())
+    # Write the main index file if requested
+    if generate_main_index and examples_index:
+        with open(examples_index.path, "w+") as f:
+            f.write(examples_index.generate())

docs/source/getting_started/examples/examples_index.md

@@ -1,10 +0,0 @@
-# Examples
-
-A collection of examples demonstrating usage of FastVideo.
-All documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.
-
-:::{toctree}
-:caption: Examples
-:maxdepth: 2
-
-:::

docs/source/getting_started/installation.md

@@ -1,10 +1,95 @@
 (fastvideo-installation)=
 
 # 🔧 Installation
-The code is tested on Python 3.10.0, CUDA 12.4 and H100.
 
-```
-./env_setup.sh fastvideo
+FastVideo currently only supports Linux and NVIDIA CUDA GPUs.
+
+FastVideo has been tested on the following GPUs, but it should work on any GPUs that supports CUDA 12.4+, please create an issue if you discover any issues:
+- RTX 4090
+- A40
+- L40S
+- A100
+- H100
+
+## Requirements
+
+- OS: Linux
+- Python: 3.10-3.12
+- CUDA 12.4+
+
+## Installation Options
+
+### Option 1: Quick Install
+
+```bash
+pip install fastvideo
 ```
 
-To try Sliding Tile Attention (optional), please follow the instruction in [here](#sta-installation) to install STA.
+### Option 2: Installation from Source
+
+We recommend using a Python environment such as Conda.
+
+#### 1. [Optional] Install Miniconda (if not already installed)
+
+```bash
+wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
+bash Miniconda3-latest-Linux-x86_64.sh
+source ~/.bashrc
+```
+
+#### 2. [Optional] Create and activate a Conda environment for FastVideo
+
+```bash
+conda create -n fastvideo python=3.10 -y
+conda activate fastvideo
+```
+
+#### 3. Clone the FastVideo repository
+
+```bash
+git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
+```
+
+#### 4. Install FastVideo
+
+Basic installation:
+
+```bash
+pip install -e .
+```
+
+## Optional Dependencies
+
+### Flash Attention
+
+```bash
+pip install flash-attn==2.7.0.post2 --no-build-isolation
+```
+
+### Sliding Tile Attention (STA) (Requires CUDA 12.4+ and H100)
+
+To try Sliding Tile Attention (optional), please follow the instructions in [csrc/sliding_tile_attention/README.md](#sta-installation) to install STA.
+
+## Development Environment Setup
+
+If you're planning to contribute to FastVideo please see the following page:
+[Contributor Guide](#developer-overview)
+
+## Hardware Requirements
+
+### For Basic Inference
+- NVIDIA GPU with CUDA support
+- Minimum 20GB VRAM for quantized models (e.g., single RTX 4090)
+
+### For Lora Finetuning
+- 40GB GPU memory each for 2 GPUs with lora
+- 30GB GPU memory each for 2 GPUs with CPU offload and lora
+
+### For Full Finetuning/Distillation
+- Multiple high-memory GPUs recommended (e.g., H100)
+
+## Troubleshooting
+
+If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).
+
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg) for additional support.

docs/source/index.md

@@ -32,7 +32,8 @@ FastVideo is a lightweight framework for accelerating large video diffusion mode
 
 FastVideo currently offers: (with more to come)
 
-- [NEW!] [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
+- [NEW!] V1 inference API available. Full announcement coming soon!
+- [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
 - FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
 - First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
 - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
@@ -43,14 +44,31 @@ Dev in progress and highly experimental.
 
 ## Documentation
 
-% How to start using vLLM?
+% How to start using FastVideo?
 
 :::{toctree}
 :caption: Getting Started
 :maxdepth: 1
 
 getting_started/installation
-getting_started/examples/examples_index
+<!-- getting_started/examples/examples_index -->
+:::
+
+:::{toctree}
+:caption: Inference
+:maxdepth: 1
+
+inference/examples/examples_inference_index
+inference/v0_inference
+:::
+
+:::{toctree}
+:caption: Training
+:maxdepth: 1
+
+training/data_preprocess
+training/distillation
+training/finetune
 :::
 
 % What is STA Kernel?
@@ -60,26 +78,22 @@ getting_started/examples/examples_index
 :maxdepth: 1
 
 sliding_tile_attention/installation
-sliding_tile_attention/usage
-sliding_tile_attention/test
 sliding_tile_attention/demo
 :::
 
 :::{toctree}
-:caption: Inference
+:caption: Design
 :maxdepth: 1
-
-inference/stepvideo
-inference/hunyuanvideo
-inference/fasthunyuan
-inference/fastmochi
+design/overview
 :::
 
 :::{toctree}
 :caption: Developer Guide
-:maxdepth: 1
+:maxdepth: 2
 
-developer_guide/overview
+contributing/overview
+contributing/developer_env/index
+contributing/add_pipeline
 :::
 
 ## Indices and tables

docs/source/inference/v0_inference.md

@@ -0,0 +1,74 @@
+(v0-inference)=
+
+# [Deprecated] V0 Inference
+The following commands and APIs are deprecated but still supported until V1's API can completely replace all the features in this page.
+
+## Inference StepVideo with Sliding Tile Attention
+First, download the model:
+
+```
+python scripts/huggingface/download_hf.py --repo_id=stepfun-ai/stepvideo-t2v --local_dir=data/stepvideo-t2v --repo_type=model
+```
+
+Use the following scripts to run inference for StepVideo. When using STA for inference, the generated videos will have dimensions of 204×768×768 (currently, this is the only supported shape).
+
+```bash
+sh scripts/inference/inference_stepvideo_STA.sh # Inference stepvideo with STA
+sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo
+```
+
+## Inference HunyuanVideo with Sliding Tile Attention
+First, download the model:
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model
+```
+
+We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
+
+```bash
+sh scripts/inference/inference_hunyuan_STA.sh
+```
+
+## Video Demos using STA + Teacache
+Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100.
+
+## Inference FastHunyuan on single RTX4090
+We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
+
+```bash
+# Download the model weight
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model
+# CLI inference
+bash scripts/inference/inference_hunyuan_hf_quantization.sh
+```
+
+For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU):
+
+| Configuration                  | Memory to Init Transformer | Peak Memory After Init Pipeline (Denoise) | Diffusion Time | End-to-End Time |
+|--------------------------------|----------------------------|--------------------------------------------|----------------|-----------------|
+| BF16 + Pipeline CPU Offload    | 23.883G                   | 33.744G                                    | 81s            | 121.5s          |
+| INT8 + Pipeline CPU Offload    | 13.911G                   | 27.979G                                    | 88s            | 116.7s          |
+| NF4 + Pipeline CPU Offload     | 9.453G                    | 19.26G                                     | 78s            | 114.5s          |
+
+For improved quality in generated videos, we recommend using a GPU with 80GB of memory to run the BF16 model with the original Hunyuan pipeline. To execute the inference, use the following section:
+
+## FastHunyuan
+
+```bash
+# Download the model weight
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan --local_dir=data/FastHunyuan --repo_type=model
+# CLI inference
+bash scripts/inference/inference_hunyuan.sh
+```
+
+You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo).
+
+## FastMochi
+
+```bash
+# Download the model weight
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffusers --local_dir=data/FastMochi-diffusers --repo_type=model
+# CLI inference
+bash scripts/inference/inference_mochi_sp.sh
+```

docs/source/inference/wanvideo.md

@@ -0,0 +1,44 @@
+(wanvideo)=
+
+# WanVideo
+## Inference T2V with WanVideo
+First, download the model:
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-T2V-1.3B-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model 
+```
+
+or
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-T2V-14B-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model 
+```
+
+Then run the inference using:
+
+```bash
+sh scripts/inference/v1_inference_wan.sh
+```
+
+Remember to set `MODEL_BASE` and `num_gpus` accordingly.
+
+## Inference I2V with WanVideo
+First, download the model:
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-I2V-14B-480P-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model
+```
+
+or
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=Wan-AI/Wan2.1-I2V-14B-720P-Diffusers --local_dir=YOUR_LOCAL_DIR --repo_type=model
+```
+
+Then run the inference using:
+
+```bash
+sh scripts/inference/v1_inference_wan_i2v.sh
+```
+
+Remember to set `MODEL_BASE` and `num_gpus` accordingly.

docs/source/sliding_tile_attention/demo.md

@@ -1,6 +1,6 @@
 (sta-demo)=
 
-# Demo
+# 🔍 Demo
 There is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
 
 <div style="text-align: center;">

docs/source/sliding_tile_attention/installation.md

@@ -1,10 +1,18 @@
 (sta-installation)=
 
-# Installation
+# 🔧 Installation
+You can install the Sliding Tile Attention package using
+
+```
+pip install st_attn==0.0.4
+``` 
+
+# Building from Source
 We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
 First, install C++20 for ThunderKittens:
 
 ```bash
+cd csrc/sliding_tile_attention/
 sudo apt update
 sudo apt install gcc-11 g++-11
 
@@ -23,3 +31,25 @@ export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 git submodule update --init --recursive
 python setup.py install
 ```
+
+# 🧪 Test
+
+```bash
+python test/test_sta.py
+```
+
+# 📋 Usage
+
+```python
+from st_attn import sliding_tile_attention
+# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
+# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256
+# a tile is a cube of size (6, 8, 8)
+# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video.
+# text_length: int ranging from 0 to 256
+# If your attention contains text token (Hunyuan)
+out = sliding_tile_attention(q, k, v, window_size, text_length)
+# If your attention does not contain text token (StepVideo)
+out = sliding_tile_attention(q, k, v, window_size, 0, False)
+
+```

docs/source/sliding_tile_attention/test.md

@@ -1,7 +0,0 @@
-(sta-test)=
-
-# Test
-
-```bash
-python test/test_sta.py
-```

docs/source/sliding_tile_attention/usage.md

@@ -1,17 +0,0 @@
-(sta-usage)=
-
-# Usage
-
-```python
-from st_attn import sliding_tile_attention
-# assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 
-# q, k, v: [batch_size, num_heads, seq_length, head_dim], seq_length = T*H*W + 256
-# a tile is a cube of size (6, 8, 8)
-# window_size in tiles: [(window_t, window_h, window_w), (..)...]. For example, window size (3, 3, 3) means a query can attend to (3x6, 3x8, 3x8) = (18, 24, 24) tokens out of the total 30x48x80 video.
-# text_length: int ranging from 0 to 256
-# If your attention contains text token (Hunyuan)
-out = sliding_tile_attention(q, k, v, window_size, text_length)
-# If your attention does not contain text token (StepVideo)
-out = sliding_tile_attention(q, k, v, window_size, 0, False)
-
-```

docs/source/training/data_preprocess.md

@@ -1,5 +1,6 @@
+(v0-data-preprocess)=
 
-## 🧱 Data Preprocess
+# 🧱 Data Preprocess
 
 To save GPU memory, we precompute text embeddings and VAE latents to eliminate the need to load the text encoder and VAE during training.
 
@@ -18,10 +19,11 @@ bash scripts/preprocess/preprocess_hunyuan_data.sh # for hunyuan
 
 The preprocessed dataset will be stored in `Image-Vid-Finetune-Mochi` or `Image-Vid-Finetune-HunYuan` correspondingly.
 
-### Process your own dataset
+## Process your own dataset
 
 If you wish to create your own dataset for finetuning or distillation, please structure you video dataset in the following format:
 
+```
 path_to_dataset_folder/
 ├── media/
 │   ├── 0.jpg
@@ -29,6 +31,7 @@ path_to_dataset_folder/
 │   ├── 2.jpg
 ├── video2caption.json
 └── merge.txt
+```
 
 Format the JSON file as a list, where each item represents a media source:
 

docs/source/training/distillation.md

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

docs/source/training/finetune.md

@@ -0,0 +1,71 @@
+(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:
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/Mochi-Black-Myth --local_dir=data/Mochi-Black-Myth --repo_type=dataset
+```
+
+Download the original model weights as specified in [Distill Section](#v0-distill):
+
+Then you can run the finetune with:
+
+```
+bash scripts/finetune/finetune_mochi.sh # for mochi
+```
+
+**Note that for finetuning, we did not tune the hyperparameters in the provided script.**
+## ⚡ Lora Finetune
+
+Hunyuan supports Lora fine-tuning of videos up to 720p. Demos and prompts of Black-Myth-Wukong can be found in [here](https://huggingface.co/FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight). You can download the Lora weight through:
+
+```bash
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight --local_dir=data/Hunyuan-Black-Myth-Wukong-lora-weight --repo_type=model
+```
+
+### Minimum Hardware Requirement
+- 40 GB GPU memory each for 2 GPUs with lora.
+- 30 GB GPU memory each for 2 GPUs with CPU offload and lora.
+
+Currently, both Mochi and Hunyuan models support Lora finetuning through diffusers. To generate personalized videos from your own dataset, you'll need to follow three main steps: dataset preparation, finetuning, and inference.
+
+### Dataset Preparation
+We provide scripts to better help you get started to train on your own characters!
+You can run this to organize your dataset to get the videos2caption.json before preprocess. Specify your video folder and corresponding caption folder (caption files should be .txt files and have the same name with its video):
+
+```
+python scripts/dataset_preparation/prepare_json_file.py --video_dir data/input_videos/ --prompt_dir data/captions/ --output_path data/output_folder/videos2caption.json --verbose
+```
+
+Also, we provide script to resize your videos:
+
+```
+python scripts/data_preprocess/resize_videos.py
+```
+
+### Finetuning
+After basic dataset preparation and preprocess, you can start to finetune your model using Lora:
+
+```
+bash scripts/finetune/finetune_hunyuan_hf_lora.sh
+```
+
+### Inference
+For inference with Lora checkpoint, you can run the following scripts with additional parameter `--lora_checkpoint_dir`:
+
+```
+bash scripts/inference/inference_hunyuan_hf.sh
+```
+
+**We also provide scripts for Mochi in the same directory.**
+
+### Finetune with Both Image and Video
+Our codebase support finetuning with both image and video.
+
+```bash
+bash scripts/finetune/finetune_hunyuan.sh
+bash scripts/finetune/finetune_mochi_lora_mix.sh
+```
+
+For Image-Video Mixture Fine-tuning, make sure to enable the `--group_frame` option in your script.

examples/inference/basic/README.md

@@ -0,0 +1,3 @@
+# Basic
+
+The class provides the main python interface for using FastVideo's inference pipeline.

examples/inference/basic/basic.py

@@ -0,0 +1 @@
+print('Hello, world!')

fastvideo/__init__.py

@@ -0,0 +1,5 @@
+from fastvideo.v1.configs.pipelines import PipelineConfig
+from fastvideo.v1.configs.sample import SamplingParam
+from fastvideo.v1.entrypoints.video_generator import VideoGenerator
+
+__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam"]

fastvideo/v1/attention/backends/abstract.py

@@ -7,7 +7,7 @@ from typing import (TYPE_CHECKING, Any, Dict, Generic, Optional, Protocol, Set,
                     Type, TypeVar)
 
 if TYPE_CHECKING:
-    from fastvideo.v1.inference_args import InferenceArgs
+    from fastvideo.v1.fastvideo_args import FastVideoArgs
     from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 
 import torch
@@ -154,7 +154,7 @@ class AttentionMetadataBuilder(ABC, Generic[T]):
         self,
         current_timestep: int,
         forward_batch: "ForwardBatch",
-        inference_args: "InferenceArgs",
+        fastvideo_args: "FastVideoArgs",
     ) -> T:
         """Build attention metadata with on-device tensors."""
         raise NotImplementedError
@@ -186,9 +186,10 @@ class AttentionImpl(ABC, Generic[T]):
         num_heads: int,
         head_size: int,
         softmax_scale: float,
-        dropout_rate: float = 0.0,
         causal: bool = False,
         num_kv_heads: Optional[int] = None,
+        prefix: str = "",
+        **extra_impl_args,
     ) -> None:
         raise NotImplementedError
 

fastvideo/v1/attention/backends/flash_attn.py

@@ -3,7 +3,16 @@
 from typing import List, Optional, Type
 
 import torch
-from flash_attn import flash_attn_func
+from flash_attn import flash_attn_func as flash_attn_2_func
+
+try:
+    from flash_attn_interface import flash_attn_func as flash_attn_3_func
+
+    # flash_attn 3 has slightly different API: it returns lse by default
+    flash_attn_func = lambda q, k, v, softmax_scale, causal: flash_attn_3_func(
+        q, k, v, softmax_scale, causal)[0]
+except ImportError:
+    flash_attn_func = flash_attn_2_func
 
 from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
                                                       AttentionImpl,
@@ -45,12 +54,12 @@ class FlashAttentionImpl(AttentionImpl):
         self,
         num_heads: int,
         head_size: int,
-        dropout_rate: float,
         causal: bool,
         softmax_scale: float,
         num_kv_heads: Optional[int] = None,
+        prefix: str = "",
+        **extra_impl_args,
     ) -> None:
-        self.dropout_rate = dropout_rate
         self.causal = causal
         self.softmax_scale = softmax_scale
 
@@ -61,10 +70,10 @@ class FlashAttentionImpl(AttentionImpl):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata,
     ):
-        output = flash_attn_func(query,
-                                 key,
-                                 value,
-                                 dropout_p=self.dropout_rate,
-                                 softmax_scale=self.softmax_scale,
-                                 causal=self.causal)
+        output = flash_attn_func(
+            query,  # type: ignore[no-untyped-call]
+            key,
+            value,
+            softmax_scale=self.softmax_scale,
+            causal=self.causal)
         return output

fastvideo/v1/attention/backends/sage_attn.py

@@ -0,0 +1,66 @@
+from typing import List, Optional, Type
+
+import torch
+from sageattention import sageattn
+
+from fastvideo.v1.attention.backends.abstract import (
+    AttentionBackend)  # FlashAttentionMetadata,
+from fastvideo.v1.attention.backends.abstract import (AttentionImpl,
+                                                      AttentionMetadata)
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class SageAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @staticmethod
+    def get_supported_head_sizes() -> List[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @staticmethod
+    def get_name() -> str:
+        return "SAGE_ATTN"
+
+    @staticmethod
+    def get_impl_cls() -> Type["SageAttentionImpl"]:
+        return SageAttentionImpl
+
+    # @staticmethod
+    # def get_metadata_cls() -> Type["AttentionMetadata"]:
+    #     return FlashAttentionMetadata
+
+
+class SageAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        causal: bool,
+        softmax_scale: float,
+        num_kv_heads: Optional[int] = None,
+        prefix: str = "",
+        **extra_impl_args,
+    ) -> None:
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> torch.Tensor:
+        output = sageattn(
+            query,
+            key,
+            value,
+            # since input is (batch_size, seq_len, head_num, head_dim)
+            tensor_layout="NHD",
+            is_causal=self.causal)
+        return output

fastvideo/v1/attention/backends/sdpa.py

@@ -38,14 +38,15 @@ class SDPAImpl(AttentionImpl):
         self,
         num_heads: int,
         head_size: int,
-        dropout_rate: float,
         causal: bool,
         softmax_scale: float,
         num_kv_heads: Optional[int] = None,
+        prefix: str = "",
+        **extra_impl_args,
     ) -> None:
-        self.dropout_rate = dropout_rate
         self.causal = causal
         self.softmax_scale = softmax_scale
+        self.dropout = extra_impl_args.get("dropout_p", 0.0)
 
     def forward(
         self,
@@ -60,7 +61,7 @@ class SDPAImpl(AttentionImpl):
         value = value.transpose(1, 2)
         attn_kwargs = {
             "attn_mask": None,
-            "dropout_p": self.dropout_rate,
+            "dropout_p": self.dropout,
             "is_causal": self.causal,
             "scale": self.softmax_scale
         }

fastvideo/v1/attention/backends/sliding_tile_attn.py

@@ -12,7 +12,7 @@ from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
                                                       AttentionMetadata,
                                                       AttentionMetadataBuilder)
 from fastvideo.v1.distributed import get_sp_group
-from fastvideo.v1.inference_args import InferenceArgs
+from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 
@@ -20,20 +20,39 @@ logger = init_logger(__name__)
 
 
 # TODO(will-refactor): move this to a utils file
-def dict_to_3d_list(mask_strategy,
-                    t_max=50,
-                    l_max=60,
-                    h_max=24) -> List[List[List[Optional[torch.Tensor]]]]:
-    result = [[[None for _ in range(h_max)] for _ in range(l_max)]
-              for _ in range(t_max)]
-    if mask_strategy is None:
-        return result
+def dict_to_3d_list(mask_strategy) -> List[List[List[Optional[torch.Tensor]]]]:
+    indices = [tuple(map(int, key.split('_'))) for key in mask_strategy]
+
+    max_timesteps_idx = max(
+        timesteps_idx for timesteps_idx, layer_idx, head_idx in indices) + 1
+    max_layer_idx = max(layer_idx
+                        for timesteps_idx, layer_idx, head_idx in indices) + 1
+    max_head_idx = max(head_idx
+                       for timesteps_idx, layer_idx, head_idx in indices) + 1
+
+    result = [[[None for _ in range(max_head_idx)]
+               for _ in range(max_layer_idx)] for _ in range(max_timesteps_idx)]
+
     for key, value in mask_strategy.items():
-        t, layer, h = map(int, key.split('_'))
-        result[t][layer][h] = value
+        timesteps_idx, layer_idx, head_idx = map(int, key.split('_'))
+        result[timesteps_idx][layer_idx][head_idx] = value
+
     return result
 
 
+class RangeDict(dict):
+
+    def __getitem__(self, item):
+        for key in self.keys():
+            if isinstance(key, tuple):
+                low, high = key
+                if low <= item <= high:
+                    return super().__getitem__(key)
+            elif key == item:
+                return super().__getitem__(key)
+        raise KeyError(f"seq_len {item} not supported for STA")
+
+
 class SlidingTileAttentionBackend(AttentionBackend):
 
     accept_output_buffer: bool = True
@@ -62,7 +81,7 @@ class SlidingTileAttentionBackend(AttentionBackend):
 
 @dataclass
 class SlidingTileAttentionMetadata(AttentionMetadata):
-    text_length: int
+    current_timestep: int
 
 
 class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder):
@@ -77,13 +96,10 @@ class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder):
         self,
         current_timestep: int,
         forward_batch: ForwardBatch,
-        inference_args: InferenceArgs,
+        fastvideo_args: FastVideoArgs,
     ) -> SlidingTileAttentionMetadata:
 
-        return SlidingTileAttentionMetadata(
-            current_timestep=current_timestep,
-            text_length=forward_batch.attention_mask.sum(),
-        )
+        return SlidingTileAttentionMetadata(current_timestep=current_timestep, )
 
 
 class SlidingTileAttentionImpl(AttentionImpl):
@@ -92,10 +108,11 @@ class SlidingTileAttentionImpl(AttentionImpl):
         self,
         num_heads: int,
         head_size: int,
-        dropout_rate: float,
         causal: bool,
         softmax_scale: float,
         num_kv_heads: Optional[int] = None,
+        prefix: str = "",
+        **extra_impl_args,
     ) -> None:
         # TODO(will-refactor): for now this is the mask strategy, but maybe we should
         # have a more general config for STA?
@@ -105,52 +122,73 @@ class SlidingTileAttentionImpl(AttentionImpl):
 
         with open(config_file) as f:
             mask_strategy = json.load(f)
-
         mask_strategy = dict_to_3d_list(mask_strategy)
 
+        self.prefix = prefix
         self.mask_strategy = mask_strategy
         sp_group = get_sp_group()
         self.sp_size = sp_group.world_size
+        # STA config
+        self.STA_base_tile_size = [6, 8, 8]
+        self.img_latent_shape_mapping = RangeDict({
+            (115200, 115456): '30x48x80',
+            82944: '36x48x48',
+            69120: '18x48x80',
+        })
+        self.full_window_mapping = {
+            '30x48x80': [5, 6, 10],
+            '36x48x48': [6, 6, 6],
+            '18x48x80': [3, 6, 10]
+        }
 
     def tile(self, x: torch.Tensor) -> torch.Tensor:
         x = rearrange(x,
                       "b (sp t h w) head d -> b (t sp h w) head d",
                       sp=self.sp_size,
-                      t=30 // self.sp_size,
-                      h=48,
-                      w=80)
+                      t=self.img_latent_shape_int[0] // self.sp_size,
+                      h=self.img_latent_shape_int[1],
+                      w=self.img_latent_shape_int[2])
         return rearrange(
             x,
             "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
-            n_t=5,
-            n_h=6,
-            n_w=10,
-            ts_t=6,
-            ts_h=8,
-            ts_w=8)
+            n_t=self.full_window_size[0],
+            n_h=self.full_window_size[1],
+            n_w=self.full_window_size[2],
+            ts_t=self.STA_base_tile_size[0],
+            ts_h=self.STA_base_tile_size[1],
+            ts_w=self.STA_base_tile_size[2])
 
     def untile(self, x: torch.Tensor) -> torch.Tensor:
         x = rearrange(
             x,
             "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
-            n_t=5,
-            n_h=6,
-            n_w=10,
-            ts_t=6,
-            ts_h=8,
-            ts_w=8)
+            n_t=self.full_window_size[0],
+            n_h=self.full_window_size[1],
+            n_w=self.full_window_size[2],
+            ts_t=self.STA_base_tile_size[0],
+            ts_h=self.STA_base_tile_size[1],
+            ts_w=self.STA_base_tile_size[2])
         return rearrange(x,
                          "b (t sp h w) head d -> b (sp t h w) head d",
                          sp=self.sp_size,
-                         t=30 // self.sp_size,
-                         h=48,
-                         w=80)
+                         t=self.img_latent_shape_int[0] // self.sp_size,
+                         h=self.img_latent_shape_int[1],
+                         w=self.img_latent_shape_int[2])
 
     def preprocess_qkv(
         self,
         qkv: torch.Tensor,
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
+        img_sequence_length = qkv.shape[1]
+        self.img_latent_shape_str = self.img_latent_shape_mapping[
+            img_sequence_length]
+        self.full_window_size = self.full_window_mapping[
+            self.img_latent_shape_str]
+        self.img_latent_shape_int = list(
+            map(int, self.img_latent_shape_str.split('x')))
+        self.img_seq_length = self.img_latent_shape_int[
+            0] * self.img_latent_shape_int[1] * self.img_latent_shape_int[2]
         return self.tile(qkv)
 
     def postprocess_output(
@@ -172,24 +210,32 @@ class SlidingTileAttentionImpl(AttentionImpl):
         assert self.mask_strategy[
             0] is not None, "mask_strategy[0] cannot be None for SlidingTileAttention"
 
-        text_length = attn_metadata.text_length
+        timestep = attn_metadata.current_timestep
+        # pattern:'.double_blocks.0.attn.impl' or '.single_blocks.0.attn.impl'
+        layer_idx = int(self.prefix.split('.')[-3])
 
-        query = q.transpose(1, 2)
-        key = k.transpose(1, 2)
-        value = v.transpose(1, 2)
+        # TODO: remove hardcode
+
+        text_length = q.shape[1] - self.img_seq_length
+        has_text = text_length > 0
+
+        query = q.transpose(1, 2).contiguous()
+        key = k.transpose(1, 2).contiguous()
+        value = v.transpose(1, 2).contiguous()
 
         head_num = query.size(1)
         sp_group = get_sp_group()
         current_rank = sp_group.rank_in_group
         start_head = current_rank * head_num
         windows = [
-            self.mask_strategy[head_idx + start_head]
+            self.mask_strategy[timestep][layer_idx][head_idx + start_head]
             for head_idx in range(head_num)
         ]
-
-        hidden_states = sliding_tile_attention(query, key, value, windows,
-                                               text_length).transpose(1, 2)
-
-        hidden_states = hidden_states.transpose(1, 2)
+        # if has_text is False:
+        #     from IPython import embed
+        #     embed()
+        hidden_states = sliding_tile_attention(
+            query, key, value, windows, text_length, has_text,
+            self.img_latent_shape_str).transpose(1, 2)
 
         return hidden_states

fastvideo/v1/attention/layer.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -12,6 +12,7 @@ from fastvideo.v1.distributed.communication_op import (
 from fastvideo.v1.distributed.parallel_state import (
     get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size)
 from fastvideo.v1.forward_context import ForwardContext, get_forward_context
+from fastvideo.v1.platforms import _Backend
 
 
 class DistributedAttention(nn.Module):
@@ -22,13 +23,13 @@ class DistributedAttention(nn.Module):
                  num_heads: int,
                  head_size: int,
                  num_kv_heads: Optional[int] = None,
-                 dropout_rate: float = 0.0,
                  softmax_scale: Optional[float] = None,
                  causal: bool = False,
+                 supported_attention_backends: Optional[Tuple[_Backend,
+                                                              ...]] = None,
+                 prefix: str = "",
                  **extra_impl_args) -> None:
         super().__init__()
-        # self.dropout_rate = dropout_rate
-        # self.causal = causal
         if softmax_scale is None:
             self.softmax_scale = head_size**-0.5
         else:
@@ -38,14 +39,17 @@ class DistributedAttention(nn.Module):
             num_kv_heads = num_heads
 
         dtype = torch.get_default_dtype()
-        attn_backend = get_attn_backend(head_size, dtype, distributed=True)
+        attn_backend = get_attn_backend(
+            head_size,
+            dtype,
+            supported_attention_backends=supported_attention_backends)
         impl_cls = attn_backend.get_impl_cls()
         self.impl = impl_cls(num_heads=num_heads,
                              head_size=head_size,
-                             dropout_rate=dropout_rate,
                              causal=causal,
                              softmax_scale=self.softmax_scale,
                              num_kv_heads=num_kv_heads,
+                             prefix=f"{prefix}.impl",
                              **extra_impl_args)
         self.num_heads = num_heads
         self.head_size = head_size
@@ -97,7 +101,6 @@ class DistributedAttention(nn.Module):
         qkv = sequence_model_parallel_all_to_all_4D(qkv,
                                                     scatter_dim=2,
                                                     gather_dim=1)
-
         # Apply backend-specific preprocess_qkv
         qkv = self.impl.preprocess_qkv(qkv, ctx_attn_metadata)
 
@@ -124,8 +127,7 @@ class DistributedAttention(nn.Module):
             output = output[:, :seq_len * world_size]
             # TODO: make this asynchronous
             replicated_output = sequence_model_parallel_all_gather(
-                replicated_output, dim=2)
-
+                replicated_output.contiguous(), dim=2)
         # Apply backend-specific postprocess_output
         output = self.impl.postprocess_output(output, ctx_attn_metadata)
 
@@ -143,13 +145,12 @@ class LocalAttention(nn.Module):
                  num_heads: int,
                  head_size: int,
                  num_kv_heads: Optional[int] = None,
-                 dropout_rate: float = 0.0,
                  softmax_scale: Optional[float] = None,
                  causal: bool = False,
+                 supported_attention_backends: Optional[Tuple[_Backend,
+                                                              ...]] = None,
                  **extra_impl_args) -> None:
         super().__init__()
-        # self.dropout_rate = dropout_rate
-        # self.causal = causal
         if softmax_scale is None:
             self.softmax_scale = head_size**-0.5
         else:
@@ -158,11 +159,13 @@ class LocalAttention(nn.Module):
             num_kv_heads = num_heads
 
         dtype = torch.get_default_dtype()
-        attn_backend = get_attn_backend(head_size, dtype, distributed=False)
+        attn_backend = get_attn_backend(
+            head_size,
+            dtype,
+            supported_attention_backends=supported_attention_backends)
         impl_cls = attn_backend.get_impl_cls()
         self.impl = impl_cls(num_heads=num_heads,
                              head_size=head_size,
-                             dropout_rate=dropout_rate,
                              softmax_scale=self.softmax_scale,
                              num_kv_heads=num_kv_heads,
                              causal=causal,

fastvideo/v1/attention/selector.py

@@ -4,7 +4,7 @@
 import os
 from contextlib import contextmanager
 from functools import cache
-from typing import Generator, Optional, Type, cast
+from typing import Generator, Optional, Tuple, Type, cast
 
 import torch
 
@@ -82,29 +82,25 @@ def get_global_forced_attn_backend() -> Optional[_Backend]:
 def get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
-    distributed: bool,
+    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
 ) -> Type[AttentionBackend]:
-    """Selects which attention backend to use and lazily imports it."""
-    # Accessing envs.* behind an @lru_cache decorator can cause the wrong
-    # value to be returned from the cache if the value changes between calls.
-    return _cached_get_attn_backend(
-        head_size=head_size,
-        dtype=dtype,
-        distributed=distributed,
-    )
+    return _cached_get_attn_backend(head_size, dtype,
+                                    supported_attention_backends)
 
 
 @cache
 def _cached_get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
-    distributed: bool,
+    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
 ) -> Type[AttentionBackend]:
     # Check whether a particular choice of backend was
     # previously forced.
     #
     # THIS SELECTION OVERRIDES THE FASTVIDEO_ATTENTION_BACKEND
     # ENVIRONMENT VARIABLE.
+    if not supported_attention_backends:
+        raise ValueError("supported_attention_backends is empty")
     selected_backend = None
     backend_by_global_setting: Optional[_Backend] = (
         get_global_forced_attn_backend())
@@ -117,8 +113,10 @@ def _cached_get_attn_backend(
             selected_backend = backend_name_to_enum(backend_by_env_var)
 
     # get device-specific attn_backend
+    if selected_backend not in supported_attention_backends:
+        selected_backend = None
     attention_cls = current_platform.get_attn_backend_cls(
-        selected_backend, head_size, dtype, distributed)
+        selected_backend, head_size, dtype)
     if not attention_cls:
         raise ValueError(
             f"Invalid attention backend for {current_platform.device_name}")

fastvideo/v1/configs/models/__init__.py

@@ -0,0 +1,6 @@
+from fastvideo.v1.configs.models.base import ModelConfig
+from fastvideo.v1.configs.models.dits.base import DiTConfig
+from fastvideo.v1.configs.models.encoders.base import EncoderConfig
+from fastvideo.v1.configs.models.vaes.base import VAEConfig
+
+__all__ = ["ModelConfig", "VAEConfig", "DiTConfig", "EncoderConfig"]

fastvideo/v1/configs/models/base.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass, fields
+from typing import Any, Dict
+
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+# 1. ArchConfig contains all fields from diffuser's/transformer's config.json (i.e. all fields related to the architecture of the model)
+# 2. ArchConfig should be inherited & overridden by each model arch_config
+# 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
+@dataclass
+class ArchConfig:
+    pass
+
+
+@dataclass
+class ModelConfig:
+    # Every model config parameter can be categorized into either ArchConfig or everything else
+    # Diffuser/Transformer parameters
+    arch_config: ArchConfig = ArchConfig()
+
+    # FastVideo-specific parameters here
+    # i.e. STA, quantization, teacache
+
+    def __getattr__(self, name):
+        # Only called if 'name' is not found in ModelConfig directly
+        if hasattr(self.arch_config, name):
+            return getattr(self.arch_config, name)
+        raise AttributeError(
+            f"'{type(self).__name__}' object has no attribute '{name}'")
+
+    # This should be used only when loading from transformers/diffusers
+    def update_model_arch(self, source_model_dict: Dict[str, Any]) -> None:
+        arch_config = self.arch_config
+        valid_fields = {f.name for f in fields(arch_config)}
+
+        for key, value in source_model_dict.items():
+            if key in valid_fields:
+                setattr(arch_config, key, value)
+            else:
+                raise AttributeError(
+                    f"{type(arch_config).__name__} has no field '{key}'")
+
+        if hasattr(arch_config, "__post_init__"):
+            arch_config.__post_init__()
+
+    def update_model_config(self, source_model_dict: Dict[str, Any]) -> None:
+        assert "arch_config" not in source_model_dict, "Source model config shouldn't contain arch_config."
+
+        valid_fields = {f.name for f in fields(self)}
+
+        for key, value in source_model_dict.items():
+            if key in valid_fields:
+                setattr(self, key, value)
+            else:
+                logger.warning("%s does not contain field '%s'!",
+                               type(self).__name__, key)
+                raise AttributeError(f"Invalid field: {key}")
+
+        if hasattr(self, "__post_init__"):
+            self.__post_init__()

fastvideo/v1/configs/models/dits/__init__.py

@@ -0,0 +1,4 @@
+from fastvideo.v1.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
+from fastvideo.v1.configs.models.dits.wanvideo import WanVideoConfig
+
+__all__ = ["HunyuanVideoConfig", "WanVideoConfig"]

fastvideo/v1/configs/models/dits/base.py

@@ -0,0 +1,30 @@
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.v1.configs.quantization import QuantizationConfig
+from fastvideo.v1.platforms import _Backend
+
+
+@dataclass
+class DiTArchConfig(ArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=list)
+    _param_names_mapping: dict = field(default_factory=dict)
+    _supported_attention_backends: Tuple[_Backend,
+                                         ...] = (_Backend.SLIDING_TILE_ATTN,
+                                                 _Backend.SAGE_ATTN,
+                                                 _Backend.FLASH_ATTN,
+                                                 _Backend.TORCH_SDPA)
+
+    hidden_size: int = 0
+    num_attention_heads: int = 0
+    num_channels_latents: int = 0
+
+
+@dataclass
+class DiTConfig(ModelConfig):
+    arch_config: DiTArchConfig = DiTArchConfig()
+
+    # FastVideoDiT-specific parameters
+    prefix: str = ""
+    quant_config: Optional[QuantizationConfig] = None

fastvideo/v1/configs/models/dits/hunyuanvideo.py

@@ -0,0 +1,169 @@
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+import torch
+
+from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_double_block(n: str, m) -> bool:
+    return "double" in n and str.isdigit(n.split(".")[-1])
+
+
+def is_single_block(n: str, m) -> bool:
+    return "single" in n and str.isdigit(n.split(".")[-1])
+
+
+def is_refiner_block(n: str, m) -> bool:
+    return "refiner" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class HunyuanVideoArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda:
+        [is_double_block, is_single_block, is_refiner_block])
+
+    _param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
+            r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
+            r"txt_in.t_embedder.mlp.fc_in.\1",
+            r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_2\.(.*)$":
+            r"txt_in.t_embedder.mlp.fc_out.\1",
+            r"^context_embedder\.proj_in\.(.*)$":
+            r"txt_in.input_embedder.\1",
+            r"^context_embedder\.time_text_embed\.text_embedder\.linear_1\.(.*)$":
+            r"txt_in.c_embedder.fc_in.\1",
+            r"^context_embedder\.time_text_embed\.text_embedder\.linear_2\.(.*)$":
+            r"txt_in.c_embedder.fc_out.\1",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm1\.(.*)$":
+            r"txt_in.refiner_blocks.\1.norm1.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm2\.(.*)$":
+            r"txt_in.refiner_blocks.\1.norm2.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_q\.(.*)$":
+            (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 0, 3),
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_k\.(.*)$":
+            (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 1, 3),
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_v\.(.*)$":
+            (r"txt_in.refiner_blocks.\1.self_attn_qkv.\2", 2, 3),
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
+            r"txt_in.refiner_blocks.\1.self_attn_proj.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
+            r"txt_in.refiner_blocks.\1.mlp.fc_in.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
+            r"txt_in.refiner_blocks.\1.mlp.fc_out.\2",
+            r"^context_embedder\.token_refiner\.refiner_blocks\.(\d+)\.norm_out\.linear\.(.*)$":
+            r"txt_in.refiner_blocks.\1.adaLN_modulation.linear.\2",
+
+            # 3. x_embedder mapping:
+            r"^x_embedder\.proj\.(.*)$":
+            r"img_in.proj.\1",
+
+            # 4. Top-level time_text_embed mappings:
+            r"^time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
+            r"time_in.mlp.fc_in.\1",
+            r"^time_text_embed\.timestep_embedder\.linear_2\.(.*)$":
+            r"time_in.mlp.fc_out.\1",
+            r"^time_text_embed\.guidance_embedder\.linear_1\.(.*)$":
+            r"guidance_in.mlp.fc_in.\1",
+            r"^time_text_embed\.guidance_embedder\.linear_2\.(.*)$":
+            r"guidance_in.mlp.fc_out.\1",
+            r"^time_text_embed\.text_embedder\.linear_1\.(.*)$":
+            r"vector_in.fc_in.\1",
+            r"^time_text_embed\.text_embedder\.linear_2\.(.*)$":
+            r"vector_in.fc_out.\1",
+
+            # 5. transformer_blocks mapping:
+            r"^transformer_blocks\.(\d+)\.norm1\.linear\.(.*)$":
+            r"double_blocks.\1.img_mod.linear.\2",
+            r"^transformer_blocks\.(\d+)\.norm1_context\.linear\.(.*)$":
+            r"double_blocks.\1.txt_mod.linear.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
+            r"double_blocks.\1.img_attn_q_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
+            r"double_blocks.\1.img_attn_k_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
+            (r"double_blocks.\1.img_attn_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 0, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 1, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$":
+            (r"double_blocks.\1.txt_attn_qkv.\2", 2, 3),
+            r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$":
+            r"double_blocks.\1.img_attn_proj.\2",
+            # Corrected: merge attn.to_add_out into the main projection.
+            r"^transformer_blocks\.(\d+)\.attn\.to_add_out\.(.*)$":
+            r"double_blocks.\1.txt_attn_proj.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_added_q\.(.*)$":
+            r"double_blocks.\1.txt_attn_q_norm.\2",
+            r"^transformer_blocks\.(\d+)\.attn\.norm_added_k\.(.*)$":
+            r"double_blocks.\1.txt_attn_k_norm.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.img_mlp.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.img_mlp.fc_out.\2",
+            r"^transformer_blocks\.(\d+)\.ff_context\.net\.0(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.txt_mlp.fc_in.\2",
+            r"^transformer_blocks\.(\d+)\.ff_context\.net\.2(?:\.proj)?\.(.*)$":
+            r"double_blocks.\1.txt_mlp.fc_out.\2",
+
+            # 6. single_transformer_blocks mapping:
+            r"^single_transformer_blocks\.(\d+)\.attn\.norm_q\.(.*)$":
+            r"single_blocks.\1.q_norm.\2",
+            r"^single_transformer_blocks\.(\d+)\.attn\.norm_k\.(.*)$":
+            r"single_blocks.\1.k_norm.\2",
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$":
+            (r"single_blocks.\1.linear1.\2", 0, 4),
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$":
+            (r"single_blocks.\1.linear1.\2", 1, 4),
+            r"^single_transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$":
+            (r"single_blocks.\1.linear1.\2", 2, 4),
+            r"^single_transformer_blocks\.(\d+)\.proj_mlp\.(.*)$":
+            (r"single_blocks.\1.linear1.\2", 3, 4),
+            # Corrected: map proj_out to modulation.linear rather than a separate proj_out branch.
+            r"^single_transformer_blocks\.(\d+)\.proj_out\.(.*)$":
+            r"single_blocks.\1.linear2.\2",
+            r"^single_transformer_blocks\.(\d+)\.norm\.linear\.(.*)$":
+            r"single_blocks.\1.modulation.linear.\2",
+
+            # 7. Final layers mapping:
+            r"^norm_out\.linear\.(.*)$":
+            r"final_layer.adaLN_modulation.linear.\1",
+            r"^proj_out\.(.*)$":
+            r"final_layer.linear.\1",
+        })
+
+    patch_size: int = 2
+    patch_size_t: int = 1
+    in_channels: int = 16
+    out_channels: int = 16
+    num_attention_heads: int = 24
+    attention_head_dim: int = 128
+    mlp_ratio: float = 4.0
+    num_layers: int = 20
+    num_single_layers: int = 40
+    num_refiner_layers: int = 2
+    rope_axes_dim: Tuple[int, int, int] = (16, 56, 56)
+    guidance_embeds: bool = False
+    dtype: Optional[torch.dtype] = None
+    text_embed_dim: int = 4096
+    pooled_projection_dim: int = 768
+    rope_theta: int = 256
+    qk_norm: str = "rms_norm"
+
+    def __post_init__(self):
+        self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
+        self.num_channels_latents: int = self.in_channels
+
+
+@dataclass
+class HunyuanVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = HunyuanVideoArchConfig()
+
+    prefix: str = "Hunyuan"

fastvideo/v1/configs/models/dits/wanvideo.py

@@ -0,0 +1,82 @@
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class WanVideoArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    _param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^patch_embedding\.(.*)$":
+            r"patch_embedding.proj.\1",
+            r"^condition_embedder\.text_embedder\.linear_1\.(.*)$":
+            r"condition_embedder.text_embedder.fc_in.\1",
+            r"^condition_embedder\.text_embedder\.linear_2\.(.*)$":
+            r"condition_embedder.text_embedder.fc_out.\1",
+            r"^condition_embedder\.time_embedder\.linear_1\.(.*)$":
+            r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^condition_embedder\.time_embedder\.linear_2\.(.*)$":
+            r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^condition_embedder\.time_proj\.(.*)$":
+            r"condition_embedder.time_modulation.linear.\1",
+            r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$":
+            r"condition_embedder.image_embedder.ff.fc_in.\1",
+            r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$":
+            r"condition_embedder.image_embedder.ff.fc_out.\1",
+            r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$":
+            r"blocks.\1.to_q.\2",
+            r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$":
+            r"blocks.\1.to_k.\2",
+            r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$":
+            r"blocks.\1.to_v.\2",
+            r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$":
+            r"blocks.\1.to_out.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$":
+            r"blocks.\1.norm_q.\2",
+            r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$":
+            r"blocks.\1.norm_k.\2",
+            r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$":
+            r"blocks.\1.attn2.to_out.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$":
+            r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
+            r"blocks.\1.ffn.fc_out.\2",
+            r"blocks\.(\d+)\.norm2\.(.*)$":
+            r"blocks.\1.self_attn_residual_norm.norm.\2",
+        })
+
+    patch_size: Tuple[int, int, int] = (1, 2, 2)
+    text_len = 512
+    num_attention_heads: int = 40
+    attention_head_dim: int = 128
+    in_channels: int = 16
+    out_channels: int = 16
+    text_dim: int = 4096
+    freq_dim: int = 256
+    ffn_dim: int = 13824
+    num_layers: int = 40
+    cross_attn_norm: bool = True
+    qk_norm: str = "rms_norm_across_heads"
+    eps: float = 1e-6
+    image_dim: Optional[int] = None
+    added_kv_proj_dim: Optional[int] = None
+    rope_max_seq_len: int = 1024
+
+    def __post_init__(self):
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.in_channels if self.added_kv_proj_dim is None else self.out_channels
+
+
+@dataclass
+class WanVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = WanVideoArchConfig()
+
+    prefix: str = "Wan"

fastvideo/v1/configs/models/encoders/__init__.py

@@ -0,0 +1,12 @@
+from fastvideo.v1.configs.models.encoders.base import (EncoderConfig,
+                                                       ImageEncoderConfig,
+                                                       TextEncoderConfig)
+from fastvideo.v1.configs.models.encoders.clip import (CLIPTextConfig,
+                                                       CLIPVisionConfig)
+from fastvideo.v1.configs.models.encoders.llama import LlamaConfig
+from fastvideo.v1.configs.models.encoders.t5 import T5Config
+
+__all__ = [
+    "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
+    "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig", "T5Config"
+]

fastvideo/v1/configs/models/encoders/base.py

@@ -0,0 +1,55 @@
+from dataclasses import dataclass, field
+from typing import Any, List, Optional, Tuple
+
+from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.v1.configs.quantization import QuantizationConfig
+from fastvideo.v1.platforms import _Backend
+
+
+@dataclass
+class EncoderArchConfig(ArchConfig):
+    architectures: List[str] = field(default_factory=lambda: [])
+    _supported_attention_backends: Tuple[_Backend, ...] = (_Backend.FLASH_ATTN,
+                                                           _Backend.TORCH_SDPA)
+    output_hidden_states: bool = False
+    use_return_dict: bool = True
+
+
+@dataclass
+class TextEncoderArchConfig(EncoderArchConfig):
+    vocab_size: int = 0
+    hidden_size: int = 0
+    num_hidden_layers: int = 0
+    num_attention_heads: int = 0
+    pad_token_id: int = 0
+    eos_token_id: int = 0
+    text_len: int = 0
+    hidden_state_skip_layer: int = 0
+    decoder_start_token_id: int = 0
+    output_past: bool = True
+    scalable_attention: bool = True
+    tie_word_embeddings: bool = False
+
+
+@dataclass
+class ImageEncoderArchConfig(EncoderArchConfig):
+    pass
+
+
+@dataclass
+class EncoderConfig(ModelConfig):
+    arch_config: ArchConfig = EncoderArchConfig()
+
+    prefix: str = ""
+    quant_config: Optional[QuantizationConfig] = None
+    lora_config: Optional[Any] = None
+
+
+@dataclass
+class TextEncoderConfig(EncoderConfig):
+    arch_config: ArchConfig = TextEncoderArchConfig()
+
+
+@dataclass
+class ImageEncoderConfig(EncoderConfig):
+    arch_config: ArchConfig = ImageEncoderArchConfig()

fastvideo/v1/configs/models/encoders/clip.py

@@ -0,0 +1,64 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
+                                                       ImageEncoderConfig,
+                                                       TextEncoderArchConfig,
+                                                       TextEncoderConfig)
+
+
+@dataclass
+class CLIPTextArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 49408
+    hidden_size: int = 512
+    intermediate_size: int = 2048
+    projection_dim: int = 512
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 8
+    max_position_embeddings: int = 77
+    hidden_act: str = "quick_gelu"
+    layer_norm_eps: float = 1e-5
+    dropout: float = 0.0
+    attention_dropout: float = 0.0
+    initializer_range: float = 0.02
+    initializer_factor: float = 1.0
+    pad_token_id: int = 1
+    bos_token_id: int = 49406
+    eos_token_id: int = 49407
+    text_len: int = 77
+
+
+@dataclass
+class CLIPVisionArchConfig(ImageEncoderArchConfig):
+    hidden_size: int = 768
+    intermediate_size: int = 3072
+    projection_dim: int = 512
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 12
+    num_channels: int = 3
+    image_size: int = 224
+    patch_size: int = 32
+    hidden_act: str = "quick_gelu"
+    layer_norm_eps: float = 1e-5
+    dropout: float = 0.0
+    attention_dropout: float = 0.0
+    initializer_range: float = 0.02
+    initializer_factor: float = 1.0
+
+
+@dataclass
+class CLIPTextConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = CLIPTextArchConfig()
+
+    num_hidden_layers_override: Optional[int] = None
+    require_post_norm: Optional[bool] = None
+    prefix: str = "clip"
+
+
+@dataclass
+class CLIPVisionConfig(ImageEncoderConfig):
+    arch_config: ImageEncoderArchConfig = CLIPVisionArchConfig()
+
+    num_hidden_layers_override: Optional[int] = None
+    require_post_norm: Optional[bool] = None
+    prefix: str = "clip"

fastvideo/v1/configs/models/encoders/llama.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
+                                                       TextEncoderConfig)
+
+
+@dataclass
+class LlamaArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 32000
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: Optional[int] = None
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 2048
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = 0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    pretraining_tp: int = 1
+    tie_word_embeddings: bool = False
+    rope_theta: float = 10000.0
+    rope_scaling: Optional[float] = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    mlp_bias: bool = False
+    head_dim: Optional[int] = None
+    hidden_state_skip_layer: int = 2
+    text_len: int = 256
+
+
+@dataclass
+class LlamaConfig(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = LlamaArchConfig()
+
+    prefix: str = "llama"

fastvideo/v1/configs/models/encoders/t5.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
+                                                       TextEncoderConfig)
+
+
+@dataclass
+class T5ArchConfig(TextEncoderArchConfig):
+    vocab_size: int = 32128
+    d_model: int = 512
+    d_kv: int = 64
+    d_ff: int = 2048
+    num_layers: int = 6
+    num_decoder_layers: Optional[int] = None
+    num_heads: int = 8
+    relative_attention_num_buckets: int = 32
+    relative_attention_max_distance: int = 128
+    dropout_rate: float = 0.1
+    layer_norm_epsilon: float = 1e-6
+    initializer_factor: float = 1.0
+    feed_forward_proj: str = "relu"
+    dense_act_fn: str = ""
+    is_gated_act: bool = False
+    is_encoder_decoder: bool = True
+    use_cache: bool = True
+    pad_token_id: int = 0
+    eos_token_id: int = 1
+    classifier_dropout: float = 0.0
+    text_len: int = 512
+
+    # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py
+    def __post_init__(self):
+        act_info = self.feed_forward_proj.split("-")
+        self.dense_act_fn: str = act_info[-1]
+        self.is_gated_act: bool = act_info[0] == "gated"
+        if self.feed_forward_proj == "gated-gelu":
+            self.dense_act_fn = "gelu_new"
+
+
+@dataclass
+class T5Config(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = T5ArchConfig()
+
+    prefix: str = "t5"

fastvideo/v1/configs/models/vaes/__init__.py

@@ -0,0 +1,7 @@
+from fastvideo.v1.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
+from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig
+
+__all__ = [
+    "HunyuanVAEConfig",
+    "WanVAEConfig",
+]

fastvideo/v1/configs/models/vaes/base.py

@@ -0,0 +1,38 @@
+from dataclasses import dataclass
+from typing import Union
+
+import torch
+
+from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
+
+
+@dataclass
+class VAEArchConfig(ArchConfig):
+    scaling_factor: Union[float, torch.tensor] = 0
+
+    temporal_compression_ratio: int = 4
+    spatial_compression_ratio: int = 8
+
+
+@dataclass
+class VAEConfig(ModelConfig):
+    arch_config: VAEArchConfig = VAEArchConfig()
+
+    # FastVideoVAE-specific parameters
+    load_encoder: bool = True
+    load_decoder: bool = True
+
+    tile_sample_min_height: int = 256
+    tile_sample_min_width: int = 256
+    tile_sample_min_num_frames: int = 16
+    tile_sample_stride_height: int = 192
+    tile_sample_stride_width: int = 192
+    tile_sample_stride_num_frames: int = 12
+    blend_num_frames: int = 0
+
+    use_tiling: bool = True
+    use_temporal_tiling: bool = True
+    use_parallel_tiling: bool = True
+
+    def __post_init__(self):
+        self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames

fastvideo/v1/configs/models/vaes/hunyuanvae.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from typing import Tuple
+
+from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class HunyuanVAEArchConfig(VAEArchConfig):
+    in_channels: int = 3
+    out_channels: int = 3
+    latent_channels: int = 16
+    down_block_types: Tuple[str, ...] = (
+        "HunyuanVideoDownBlock3D",
+        "HunyuanVideoDownBlock3D",
+        "HunyuanVideoDownBlock3D",
+        "HunyuanVideoDownBlock3D",
+    )
+    up_block_types: Tuple[str, ...] = (
+        "HunyuanVideoUpBlock3D",
+        "HunyuanVideoUpBlock3D",
+        "HunyuanVideoUpBlock3D",
+        "HunyuanVideoUpBlock3D",
+    )
+    block_out_channels: Tuple[int, ...] = (128, 256, 512, 512)
+    layers_per_block: int = 2
+    act_fn: str = "silu"
+    norm_num_groups: int = 32
+    scaling_factor: float = 0.476986
+    spatial_compression_ratio: int = 8
+    temporal_compression_ratio: int = 4
+    mid_block_add_attention: bool = True
+
+    def __post_init__(self):
+        self.spatial_compression_ratio: int = 2**(len(self.block_out_channels) -
+                                                  1)
+
+
+@dataclass
+class HunyuanVAEConfig(VAEConfig):
+    arch_config: VAEArchConfig = HunyuanVAEArchConfig()

fastvideo/v1/configs/models/vaes/wanvae.py

@@ -0,0 +1,75 @@
+from dataclasses import dataclass
+from typing import Tuple
+
+import torch
+
+from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class WanVAEArchConfig(VAEArchConfig):
+    base_dim: int = 96
+    z_dim: int = 16
+    dim_mult: Tuple[int, ...] = (1, 2, 4, 4)
+    num_res_blocks: int = 2
+    attn_scales: Tuple[float, ...] = ()
+    temperal_downsample: Tuple[bool, ...] = (False, True, True)
+    dropout: float = 0.0
+    latents_mean: Tuple[float, ...] = (
+        -0.7571,
+        -0.7089,
+        -0.9113,
+        0.1075,
+        -0.1745,
+        0.9653,
+        -0.1517,
+        1.5508,
+        0.4134,
+        -0.0715,
+        0.5517,
+        -0.3632,
+        -0.1922,
+        -0.9497,
+        0.2503,
+        -0.2921,
+    )
+    latents_std: Tuple[float, ...] = (
+        2.8184,
+        1.4541,
+        2.3275,
+        2.6558,
+        1.2196,
+        1.7708,
+        2.6052,
+        2.0743,
+        3.2687,
+        2.1526,
+        2.8652,
+        1.5579,
+        1.6382,
+        1.1253,
+        2.8251,
+        1.9160,
+    )
+    temporal_compression_ratio = 4
+    spatial_compression_ratio = 8
+
+    def __post_init__(self):
+        self.scaling_factor: torch.tensor = 1.0 / torch.tensor(
+            self.latents_std).view(1, self.z_dim, 1, 1, 1)
+        self.shift_factor: torch.tensor = torch.tensor(self.latents_mean).view(
+            1, self.z_dim, 1, 1, 1)
+
+
+@dataclass
+class WanVAEConfig(VAEConfig):
+    arch_config: VAEArchConfig = WanVAEArchConfig()
+    use_feature_cache: bool = True
+
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+
+    def __post_init__(self):
+        self.blend_num_frames = (self.tile_sample_min_num_frames -
+                                 self.tile_sample_stride_num_frames) * 2

fastvideo/v1/configs/pipelines/__init__.py

@@ -0,0 +1,14 @@
+from fastvideo.v1.configs.pipelines.base import (PipelineConfig,
+                                                 SlidingTileAttnConfig)
+from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
+                                                    HunyuanConfig)
+from fastvideo.v1.configs.pipelines.registry import (
+    get_pipeline_config_cls_for_name)
+from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
+                                                WanT2V480PConfig)
+
+__all__ = [
+    "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
+    "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
+    "get_pipeline_config_cls_for_name"
+]

fastvideo/v1/configs/pipelines/base.py

@@ -0,0 +1,107 @@
+import json
+from dataclasses import asdict, dataclass, fields
+from typing import Any, Dict, Optional
+
+from fastvideo.v1.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
+                                         VAEConfig)
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.utils import shallow_asdict
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class PipelineConfig:
+    """Base configuration for all pipeline architectures."""
+    # Video generation parameters
+    embedded_cfg_scale: float = 6.0
+    flow_shift: Optional[float] = None
+    use_cpu_offload: bool = False
+    disable_autocast: bool = False
+
+    # Model configuration
+    precision: str = "bf16"
+
+    # VAE configuration
+    vae_precision: str = "fp16"
+    vae_tiling: bool = True
+    vae_sp: bool = True
+    vae_config: VAEConfig = VAEConfig()
+
+    # DiT configuration
+    dit_config: DiTConfig = DiTConfig()
+
+    # Text encoder configuration
+    text_encoder_precision: str = "fp16"
+    text_encoder_config: EncoderConfig = EncoderConfig()
+
+    # STA (Spatial-Temporal Attention) parameters
+    mask_strategy_file_path: Optional[str] = None
+    enable_torch_compile: bool = False
+
+    @classmethod
+    def from_pretrained(cls, model_path: str) -> "PipelineConfig":
+        from fastvideo.v1.configs.pipelines.registry import (
+            get_pipeline_config_cls_for_name)
+        pipeline_config_cls = get_pipeline_config_cls_for_name(model_path)
+        if pipeline_config_cls is not None:
+            pipeline_config = pipeline_config_cls()
+        else:
+            logger.warning(
+                "Couldn't find an optimal sampling param for %s. Using the default sampling param.",
+                model_path)
+            pipeline_config = cls()
+
+        return pipeline_config
+
+    def dump_to_json(self, file_path: str):
+        output_dict = shallow_asdict(self)
+        for key, value in output_dict.items():
+            if isinstance(value, ModelConfig):
+                model_dict = asdict(value)
+                # Model Arch Config should be hidden away from the users
+                model_dict.pop("arch_config")
+                output_dict[key] = model_dict
+
+        with open(file_path, "w") as f:
+            json.dump(output_dict, f, indent=2)
+
+    def load_from_json(self, file_path: str):
+        with open(file_path) as f:
+            input_pipeline_dict = json.load(f)
+        self.update_pipeline_config(input_pipeline_dict)
+
+    def update_pipeline_config(self, source_pipeline_dict: Dict[str,
+                                                                Any]) -> None:
+        for f in fields(self):
+            key = f.name
+            if key in source_pipeline_dict:
+                current_value = getattr(self, key)
+                new_value = source_pipeline_dict[key]
+
+                # If it's a nested ModelConfig, update it recursively
+                if isinstance(current_value, ModelConfig):
+                    current_value.update_model_config(new_value)
+                else:
+                    setattr(self, key, new_value)
+
+        if hasattr(self, "__post_init__"):
+            self.__post_init__()
+
+
+@dataclass
+class SlidingTileAttnConfig(PipelineConfig):
+    """Configuration for sliding tile attention."""
+
+    # Override any BaseConfig defaults as needed
+    # Add sliding tile specific parameters
+    window_size: int = 16
+    stride: int = 8
+
+    # You can provide custom defaults for inherited fields
+    height: int = 576
+    width: int = 1024
+
+    # Additional configuration specific to sliding tile attention
+    pad_to_square: bool = False
+    use_overlap_optimization: bool = True

fastvideo/v1/configs/pipelines/hunyuan.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass
+
+from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.v1.configs.models.dits import HunyuanVideoConfig
+from fastvideo.v1.configs.models.encoders import CLIPTextConfig, LlamaConfig
+from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig
+from fastvideo.v1.configs.pipelines.base import PipelineConfig
+
+
+@dataclass
+class HunyuanConfig(PipelineConfig):
+    """Base configuration for HunYuan pipeline architecture."""
+
+    # HunyuanConfig-specific parameters with defaults
+    # DiT
+    dit_config: DiTConfig = HunyuanVideoConfig()
+    # VAE
+    vae_config: VAEConfig = HunyuanVAEConfig()
+    # Denoising stage
+    embedded_cfg_scale: int = 6
+    flow_shift: int = 7
+
+    # Text encoding stage
+    text_encoder_config: EncoderConfig = LlamaConfig()
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "fp16"
+    text_encoder_precision: str = "fp16"
+
+    # HunyuanConfig-specific added parameters
+    # Secondary text encoder
+    text_encoder_config_2: EncoderConfig = CLIPTextConfig()
+    text_encoder_precision_2: str = "fp16"
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True
+
+
+@dataclass
+class FastHunyuanConfig(HunyuanConfig):
+    """Configuration specifically optimized for FastHunyuan weights."""
+
+    # Override HunyuanConfig defaults
+    flow_shift: int = 17
+
+    # No need to re-specify guidance_scale or embedded_cfg_scale as they
+    # already have the desired values from HunyuanConfig

fastvideo/v1/configs/pipelines/registry.py

@@ -0,0 +1,78 @@
+"""Registry for pipeline weight-specific configurations."""
+
+import os
+from typing import Callable, Dict, Optional, Type
+
+from fastvideo.v1.configs.pipelines.base import PipelineConfig
+from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
+                                                    HunyuanConfig)
+from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
+                                                WanT2V480PConfig)
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.utils import (maybe_download_model_index,
+                                verify_model_config_and_directory)
+
+logger = init_logger(__name__)
+
+# Registry maps specific model weights to their config classes
+WEIGHT_CONFIG_REGISTRY: Dict[str, Type[PipelineConfig]] = {
+    "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
+    "hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
+    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig
+    # Add other specific weight variants
+}
+
+# For determining pipeline type from model ID
+PIPELINE_DETECTOR: Dict[str, Callable[[str], bool]] = {
+    "hunyuan": lambda id: "hunyuan" in id.lower(),
+    "wanpipeline": lambda id: "wanpipeline" in id.lower(),
+    "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
+    # Add other pipeline architecture detectors
+}
+
+# Fallback configs when exact match isn't found but architecture is detected
+PIPELINE_FALLBACK_CONFIG: Dict[str, Type[PipelineConfig]] = {
+    "hunyuan":
+    HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
+    "wanpipeline":
+    WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
+    "wanimagetovideo": WanI2V480PConfig,
+    # Other fallbacks by architecture
+}
+
+
+def get_pipeline_config_cls_for_name(
+        pipeline_name_or_path: str) -> Optional[type[PipelineConfig]]:
+    """Get the appropriate config class for specific pretrained weights."""
+
+    if os.path.exists(pipeline_name_or_path):
+        config = verify_model_config_and_directory(pipeline_name_or_path)
+        logger.warning(
+            "FastVideo may not correctly identify the optimal config for this model, as the local directory may have been renamed."
+        )
+    else:
+        config = maybe_download_model_index(pipeline_name_or_path)
+
+    pipeline_name = config["_class_name"]
+
+    # First try exact match for specific weights
+    if pipeline_name_or_path in WEIGHT_CONFIG_REGISTRY:
+        return WEIGHT_CONFIG_REGISTRY[pipeline_name_or_path]
+
+    # Try partial matches (for local paths that might include the weight ID)
+    for registered_id, config_class in WEIGHT_CONFIG_REGISTRY.items():
+        if registered_id in pipeline_name_or_path:
+            return config_class
+
+    # If no match, try to use the fallback config
+    fallback_config = None
+    # Try to determine pipeline architecture for fallback
+    for pipeline_type, detector in PIPELINE_DETECTOR.items():
+        if detector(pipeline_name.lower()):
+            fallback_config = PIPELINE_FALLBACK_CONFIG.get(pipeline_type)
+            break
+
+    logger.warning("No match found for pipeline %s, using fallback config %s.",
+                   pipeline_name_or_path, fallback_config)
+    return fallback_config

fastvideo/v1/configs/pipelines/wan.py

@@ -0,0 +1,55 @@
+from dataclasses import dataclass
+
+from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.v1.configs.models.dits import WanVideoConfig
+from fastvideo.v1.configs.models.encoders import CLIPVisionConfig, T5Config
+from fastvideo.v1.configs.models.vaes import WanVAEConfig
+from fastvideo.v1.configs.pipelines.base import PipelineConfig
+
+
+@dataclass
+class WanT2V480PConfig(PipelineConfig):
+    """Base configuration for Wan T2V 1.3B pipeline architecture."""
+
+    # WanConfig-specific parameters with defaults
+    # DiT
+    dit_config: DiTConfig = WanVideoConfig()
+    # VAE
+    vae_config: VAEConfig = WanVAEConfig()
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # Video parameters
+    use_cpu_offload: bool = True
+
+    # Denoising stage
+    flow_shift: int = 3
+
+    # Text encoding stage
+    text_encoder_config: EncoderConfig = T5Config()
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "fp16"
+    text_encoder_precision: str = "fp32"
+
+    # WanConfig-specific added parameters
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True
+
+
+@dataclass
+class WanI2V480PConfig(WanT2V480PConfig):
+    """Base configuration for Wan I2V 14B 480P pipeline architecture."""
+
+    # WanConfig-specific parameters with defaults
+
+    # Precision for each component
+    image_encoder_config: EncoderConfig = CLIPVisionConfig()
+    image_encoder_precision: str = "fp32"
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True

fastvideo/v1/configs/quantization/__init__.py

@@ -0,0 +1,3 @@
+from fastvideo.v1.configs.quantization.base import QuantizationConfig
+
+__all__ = ["QuantizationConfig"]

fastvideo/v1/configs/quantization/base.py

@@ -0,0 +1,6 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class QuantizationConfig:
+    pass

fastvideo/v1/configs/sample/__init__.py

@@ -0,0 +1,3 @@
+from fastvideo.v1.configs.sample.base import SamplingParam
+
+__all__ = ["SamplingParam"]

fastvideo/v1/configs/sample/base.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class SamplingParam:
+    # All fields below are copied from ForwardBatch
+    data_type: str = "video"
+
+    # Image inputs
+    image_path: Optional[str] = None
+
+    # Text inputs
+    prompt: Optional[Union[str, List[str]]] = None
+    negative_prompt: Optional[str] = None
+    prompt_path: Optional[str] = None
+    output_path: str = "outputs/"
+
+    # Batch info
+    num_videos_per_prompt: int = 1
+    seed: int = 1024
+
+    # Original dimensions (before VAE scaling)
+    num_frames: int = 125
+    height: int = 720
+    width: int = 1280
+    fps: int = 24
+
+    # Denoising parameters
+    num_inference_steps: int = 50
+    guidance_scale: float = 1.0
+    guidance_rescale: float = 0.0
+
+    # Misc
+    save_video: bool = True
+    return_frames: bool = False
+
+    def __post_init__(self) -> None:
+        self.data_type = "video" if self.num_frames > 1 else "image"
+
+    def check_sampling_param(self):
+        if self.prompt_path and not self.prompt_path.endswith(".txt"):
+            raise ValueError("prompt_path must be a txt file")
+
+    def update(self, source_dict: Dict[str, Any]) -> None:
+        for key, value in source_dict.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+            else:
+                logger.exception("%s has no attribute %s",
+                                 type(self).__name__, key)
+
+        self.__post_init__()
+
+    @classmethod
+    def from_pretrained(cls, model_path: str) -> "SamplingParam":
+        from fastvideo.v1.configs.sample.registry import (
+            get_sampling_param_cls_for_name)
+        sampling_cls = get_sampling_param_cls_for_name(model_path)
+        if sampling_cls is not None:
+            sampling_param: SamplingParam = sampling_cls()
+        else:
+            logger.warning(
+                "Couldn't find an optimal sampling param for %s. Using the default sampling param.",
+                model_path)
+            sampling_param = cls()
+
+        return sampling_param

fastvideo/v1/configs/sample/hunyuan.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+
+from fastvideo.v1.configs.sample.base import SamplingParam
+
+
+@dataclass
+class HunyuanSamplingParam(SamplingParam):
+    num_inference_steps: int = 50
+
+    num_frames: int = 125
+    height: int = 720
+    width: int = 1280
+    fps: int = 24
+
+    guidance_scale: float = 1.0
+
+
+@dataclass
+class FastHunyuanSamplingParam(HunyuanSamplingParam):
+    num_inference_steps: int = 6

fastvideo/v1/configs/sample/registry.py

@@ -0,0 +1,75 @@
+import os
+from typing import Any, Callable, Dict, Optional
+
+from fastvideo.v1.configs.sample.hunyuan import (FastHunyuanSamplingParam,
+                                                 HunyuanSamplingParam)
+from fastvideo.v1.configs.sample.wan import (WanI2V480PSamplingParam,
+                                             WanT2V480PSamplingParam)
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.utils import (maybe_download_model_index,
+                                verify_model_config_and_directory)
+
+logger = init_logger(__name__)
+# Registry maps specific model weights to their config classes
+SAMPLING_PARAM_REGISTRY: Dict[str, Any] = {
+    "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
+    "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PSamplingParam,
+    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PSamplingParam
+    # Add other specific weight variants
+}
+
+# For determining pipeline type from model ID
+SAMPLING_PARAM_DETECTOR: Dict[str, Callable[[str], bool]] = {
+    "hunyuan": lambda id: "hunyuan" in id.lower(),
+    "wanpipeline": lambda id: "wanpipeline" in id.lower(),
+    "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(),
+    # Add other pipeline architecture detectors
+}
+
+# Fallback configs when exact match isn't found but architecture is detected
+SAMPLING_FALLBACK_PARAM: Dict[str, Any] = {
+    "hunyuan":
+    HunyuanSamplingParam,  # Base Hunyuan config as fallback for any Hunyuan variant
+    "wanpipeline":
+    WanT2V480PSamplingParam,  # Base Wan config as fallback for any Wan variant
+    "wanimagetovideo": WanI2V480PSamplingParam,
+    # Other fallbacks by architecture
+}
+
+
+def get_sampling_param_cls_for_name(
+        pipeline_name_or_path: str) -> Optional[Any]:
+    """Get the appropriate sampling param for specific pretrained weights."""
+
+    if os.path.exists(pipeline_name_or_path):
+        config = verify_model_config_and_directory(pipeline_name_or_path)
+        logger.warning(
+            "FastVideo may not correctly identify the optimal sampling param for this model, as the local directory may have been renamed."
+        )
+    else:
+        config = maybe_download_model_index(pipeline_name_or_path)
+
+    pipeline_name = config["_class_name"]
+
+    # First try exact match for specific weights
+    if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY:
+        return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path]
+
+    # Try partial matches (for local paths that might include the weight ID)
+    for registered_id, config_class in SAMPLING_PARAM_REGISTRY.items():
+        if registered_id in pipeline_name_or_path:
+            return config_class
+
+    # If no match, try to use the fallback config
+    fallback_config = None
+    # Try to determine pipeline architecture for fallback
+    for pipeline_type, detector in SAMPLING_PARAM_DETECTOR.items():
+        if detector(pipeline_name.lower()):
+            fallback_config = SAMPLING_FALLBACK_PARAM.get(pipeline_type)
+            break
+
+    logger.warning(
+        "No match found for pipeline %s, using fallback sampling param %s.",
+        pipeline_name_or_path, fallback_config)
+    return fallback_config

fastvideo/v1/configs/sample/wan.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+
+from fastvideo.v1.configs.sample.base import SamplingParam
+
+
+@dataclass
+class WanT2V480PSamplingParam(SamplingParam):
+    # Video parameters
+    height: int = 480
+    width: int = 832
+    num_frames: int = 81
+    fps: int = 16
+
+    # Denoising stage
+    guidance_scale: float = 3.0
+    negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+    num_inference_steps: int = 50
+
+
+@dataclass
+class WanI2V480PSamplingParam(WanT2V480PSamplingParam):
+    # Denoising stage
+    guidance_scale: float = 5.0
+    num_inference_steps: int = 40

fastvideo/v1/configs/wan_1.3B_t2v_pipeline.json

@@ -0,0 +1,37 @@
+{
+  "embedded_cfg_scale": 6.0,
+  "flow_shift": 3,
+  "use_cpu_offload": true,
+  "disable_autocast": false,
+  "precision": "bf16",
+  "vae_precision": "fp16",
+  "vae_tiling": false,
+  "vae_sp": false,
+  "vae_config": {
+    "load_encoder": false,
+    "load_decoder": true,
+    "tile_sample_min_height": 256,
+    "tile_sample_min_width": 256,
+    "tile_sample_min_num_frames": 16,
+    "tile_sample_stride_height": 192,
+    "tile_sample_stride_width": 192,
+    "tile_sample_stride_num_frames": 12,
+    "blend_num_frames": 8,
+    "use_tiling": false,
+    "use_temporal_tiling": false,
+    "use_parallel_tiling": false,
+    "use_feature_cache": true
+  },
+  "dit_config": {
+    "prefix": "Wan",
+    "quant_config": null
+  },
+  "text_encoder_precision": "fp32",
+  "text_encoder_config": {
+    "prefix": "t5",
+    "quant_config": null,
+    "lora_config": null
+  },
+  "mask_strategy_file_path": null,
+  "enable_torch_compile": false
+}

fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json

@@ -0,0 +1,45 @@
+{
+  "embedded_cfg_scale": 6.0,
+  "flow_shift": 3,
+  "use_cpu_offload": true,
+  "disable_autocast": false,
+  "precision": "bf16",
+  "vae_precision": "fp16",
+  "vae_tiling": false,
+  "vae_sp": false,
+  "vae_config": {
+    "load_encoder": true,
+    "load_decoder": true,
+    "tile_sample_min_height": 256,
+    "tile_sample_min_width": 256,
+    "tile_sample_min_num_frames": 16,
+    "tile_sample_stride_height": 192,
+    "tile_sample_stride_width": 192,
+    "tile_sample_stride_num_frames": 12,
+    "blend_num_frames": 8,
+    "use_tiling": false,
+    "use_temporal_tiling": false,
+    "use_parallel_tiling": false,
+    "use_feature_cache": true
+  },
+  "dit_config": {
+    "prefix": "Wan",
+    "quant_config": null
+  },
+  "text_encoder_precision": "fp32",
+  "text_encoder_config": {
+    "prefix": "t5",
+    "quant_config": null,
+    "lora_config": null
+  },
+  "mask_strategy_file_path": null,
+  "enable_torch_compile": false,
+  "image_encoder_config": {
+    "prefix": "clip",
+    "quant_config": null,
+    "lora_config": null,
+    "num_hidden_layers_override": null,
+    "require_post_norm": null
+  },
+  "image_encoder_precision": "fp32"
+}

fastvideo/v1/distributed/__init__.py

@@ -1,5 +1,20 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from fastvideo.v1.distributed.communication_op import *
-from fastvideo.v1.distributed.parallel_state import *
+from fastvideo.v1.distributed.parallel_state import (
+    cleanup_dist_env_and_memory, get_sequence_model_parallel_rank,
+    get_sequence_model_parallel_world_size, get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size, get_world_group,
+    init_distributed_environment, initialize_model_parallel)
 from fastvideo.v1.distributed.utils import *
+
+__all__ = [
+    "init_distributed_environment",
+    "initialize_model_parallel",
+    "get_sequence_model_parallel_rank",
+    "get_sequence_model_parallel_world_size",
+    "get_tensor_model_parallel_rank",
+    "get_tensor_model_parallel_world_size",
+    "cleanup_dist_env_and_memory",
+    "get_world_group",
+]

fastvideo/v1/distributed/device_communicators/base_device_communicator.py

@@ -190,5 +190,5 @@ class DeviceCommunicatorBase:
         torch.distributed.recv(tensor, self.ranks[src], self.device_group)
         return tensor
 
-    def destroy(self):
+    def destroy(self) -> None:
         pass

fastvideo/v1/distributed/parallel_state.py

@@ -845,12 +845,12 @@ def initialize_model_parallel(
                                     group_name="sp")
 
 
-def get_sequence_model_parallel_world_size():
+def get_sequence_model_parallel_world_size() -> int:
     """Return world size for the sequence model parallel group."""
     return get_sp_group().world_size
 
 
-def get_sequence_model_parallel_rank():
+def get_sequence_model_parallel_rank() -> int:
     """Return my rank for the sequence model parallel group."""
     return get_sp_group().rank_in_group
 

fastvideo/v1/entrypoints/cli/cli_types.py

@@ -9,8 +9,7 @@ from fastvideo.v1.utils import FlexibleArgumentParser
 class CLISubcommand:
     """Base class for CLI subcommands"""
 
-    def __init__(self):
-        self.name = ""
+    name: str
 
     def cmd(self, args: argparse.Namespace) -> None:
         """Execute the command with the given arguments"""

fastvideo/v1/entrypoints/cli/generate.py

@@ -2,11 +2,11 @@
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py
 
 import argparse
-from typing import List
+from typing import List, cast
 
 from fastvideo.v1.entrypoints.cli import utils
 from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
-from fastvideo.v1.inference_args import InferenceArgs
+from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.utils import FlexibleArgumentParser
 
 
@@ -73,18 +73,14 @@ class GenerateSubcommand(CLISubcommand):
             required=False,
             help="Read CLI options from a config YAML file.")
 
-        generate_parser.add_argument("--num-gpus",
-                                     type=int,
-                                     default=1,
-                                     help="Number of GPUs to use")
         generate_parser.add_argument("--master-port",
                                      type=int,
                                      default=None,
                                      help="Port for the master process")
 
-        generate_parser = InferenceArgs.add_cli_args(generate_parser)
+        generate_parser = FastVideoArgs.add_cli_args(generate_parser)
 
-        return generate_parser
+        return cast(FlexibleArgumentParser, generate_parser)
 
 
 def cmd_init() -> List[CLISubcommand]:

fastvideo/v1/entrypoints/cli/utils.py

@@ -3,13 +3,16 @@
 import os
 import subprocess
 import sys
+from typing import List, Optional
 
 from fastvideo.v1.logger import init_logger
 
 logger = init_logger(__name__)
 
 
-def launch_distributed(num_gpus=None, args=None, master_port=None):
+def launch_distributed(num_gpus: int,
+                       args: List[str],
+                       master_port: Optional[int] = None) -> int:
     """
     Launch a distributed job with the given arguments
     

fastvideo/v1/entrypoints/video_generator.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+VideoGenerator module for FastVideo.
+
+This module provides a consolidated interface for generating videos using
+diffusion models.
+"""
+
+import os
+import time
+from dataclasses import asdict
+from typing import Any, Dict, List, Optional, Union
+
+import imageio
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+
+from fastvideo.v1.configs.pipelines import (PipelineConfig,
+                                            get_pipeline_config_cls_for_name)
+from fastvideo.v1.configs.sample import SamplingParam
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines import ForwardBatch
+from fastvideo.v1.utils import align_to, shallow_asdict
+from fastvideo.v1.worker.executor import Executor
+
+logger = init_logger(__name__)
+
+
+class VideoGenerator:
+    """
+    A unified class for generating videos using diffusion models.
+    
+    This class provides a simple interface for video generation with rich
+    customization options, similar to popular frameworks like HF Diffusers.
+    """
+
+    def __init__(self, fastvideo_args: FastVideoArgs,
+                 executor_class: type[Executor], log_stats: bool):
+        """
+        Initialize the video generator.
+        
+        Args:
+            pipeline: The pipeline to use for inference
+            fastvideo_args: The inference arguments
+        """
+        self.fastvideo_args = fastvideo_args
+        self.executor = executor_class(fastvideo_args)
+
+    @classmethod
+    def from_pretrained(cls,
+                        model_path: str,
+                        device: Optional[str] = None,
+                        torch_dtype: Optional[torch.dtype] = None,
+                        pipeline_config: Optional[
+                            Union[str
+                                  | PipelineConfig]] = None,
+                        **kwargs) -> "VideoGenerator":
+        """
+        Create a video generator from a pretrained model.
+        
+        Args:
+            model_path: Path or identifier for the pretrained model
+            device: Device to load the model on (e.g., "cuda", "cuda:0", "cpu")
+            torch_dtype: Data type for model weights (e.g., torch.float16)
+            **kwargs: Additional arguments to customize model loading
+                
+        Returns:
+            The created video generator
+
+        Priority level: Default pipeline config < User's pipeline config < User's kwargs
+        """
+
+        config = None
+        # 1. If users provide a pipeline config, it will override the default pipeline config
+        if isinstance(pipeline_config, PipelineConfig):
+            config = pipeline_config
+        else:
+            config_cls = get_pipeline_config_cls_for_name(model_path)
+            if config_cls is not None:
+                config = config_cls()
+                if isinstance(pipeline_config, str):
+                    config.load_from_json(pipeline_config)
+
+        # 2. If users also provide some kwargs, it will override the pipeline config.
+        # The user kwargs shouldn't contain model config parameters!
+        if config is None:
+            logger.warning("No config found for model %s, using default config",
+                           model_path)
+            config_args = kwargs
+        else:
+            config_args = shallow_asdict(config)
+            config_args.update(kwargs)
+
+        fastvideo_args = FastVideoArgs(
+            model_path=model_path,
+            device_str=device or "cuda" if torch.cuda.is_available() else "cpu",
+            **config_args)
+        fastvideo_args.check_fastvideo_args()
+
+        return cls.from_fastvideo_args(fastvideo_args)
+
+    @classmethod
+    def from_fastvideo_args(cls,
+                            fastvideo_args: FastVideoArgs) -> "VideoGenerator":
+        """
+        Create a video generator with the specified arguments.
+        
+        Args:
+            fastvideo_args: The inference arguments
+                
+        Returns:
+            The created video generator
+        """
+        # Initialize distributed environment if needed
+        # initialize_distributed_and_parallelism(fastvideo_args)
+
+        executor_class = Executor.get_class(fastvideo_args)
+
+        return cls(
+            fastvideo_args=fastvideo_args,
+            executor_class=executor_class,
+            log_stats=False,  # TODO: implement
+        )
+
+    def generate_video(
+        self,
+        prompt: str,
+        sampling_param: Optional[SamplingParam] = None,
+        **kwargs,
+    ) -> Union[Dict[str, Any], List[np.ndarray]]:
+        """
+        Generate a video based on the given prompt.
+        
+        Args:
+            prompt: The prompt to use for generation
+            negative_prompt: The negative prompt to use (overrides the one in fastvideo_args)
+            output_path: Path to save the video (overrides the one in fastvideo_args)
+            save_video: Whether to save the video to disk
+            return_frames: Whether to return the raw frames
+            num_inference_steps: Number of denoising steps (overrides fastvideo_args)
+            guidance_scale: Classifier-free guidance scale (overrides fastvideo_args)
+            num_frames: Number of frames to generate (overrides fastvideo_args)
+            height: Height of generated video (overrides fastvideo_args)
+            width: Width of generated video (overrides fastvideo_args)
+            fps: Frames per second for saved video (overrides fastvideo_args)
+            seed: Random seed for generation (overrides fastvideo_args)
+            callback: Callback function called after each step
+            callback_steps: Number of steps between each callback
+            
+        Returns:
+            Either the output dictionary or the list of frames depending on return_frames
+        """
+        # Create a copy of inference args to avoid modifying the original
+        fastvideo_args = self.fastvideo_args
+
+        # Validate inputs
+        if not isinstance(prompt, str):
+            raise TypeError(
+                f"`prompt` must be a string, but got {type(prompt)}")
+        prompt = prompt.strip()
+
+        if sampling_param is None:
+            sampling_param = SamplingParam.from_pretrained(
+                fastvideo_args.model_path)
+        kwargs["prompt"] = prompt
+        sampling_param.update(kwargs)
+
+        # Process negative prompt
+        if sampling_param.negative_prompt is not None:
+            sampling_param.negative_prompt = sampling_param.negative_prompt.strip(
+            )
+
+        # Validate dimensions
+        if (sampling_param.height <= 0 or sampling_param.width <= 0
+                or sampling_param.num_frames <= 0):
+            raise ValueError(
+                f"Height, width, and num_frames must be positive integers, got "
+                f"height={sampling_param.height}, width={sampling_param.width}, "
+                f"num_frames={sampling_param.num_frames}")
+
+        if (
+                sampling_param.num_frames - 1
+        ) % fastvideo_args.vae_config.arch_config.temporal_compression_ratio != 0:
+            raise ValueError(
+                f"num_frames-1 must be a multiple of {fastvideo_args.vae_config.arch_config.temporal_compression_ratio}, got {sampling_param.num_frames}"
+            )
+
+        # Calculate sizes
+        target_height = align_to(sampling_param.height, 16)
+        target_width = align_to(sampling_param.width, 16)
+
+        # Calculate latent sizes
+        latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
+                        sampling_param.height // 8, sampling_param.width // 8]
+        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
+
+        # Log parameters
+        debug_str = f"""
+                      height: {target_height}
+                       width: {target_width}
+                video_length: {sampling_param.num_frames}
+                      prompt: {prompt}
+                  neg_prompt: {sampling_param.negative_prompt}
+                        seed: {sampling_param.seed}
+                 infer_steps: {sampling_param.num_inference_steps}
+       num_videos_per_prompt: {sampling_param.num_videos_per_prompt}
+              guidance_scale: {sampling_param.guidance_scale}
+                    n_tokens: {n_tokens}
+                  flow_shift: {fastvideo_args.flow_shift}
+     embedded_guidance_scale: {fastvideo_args.embedded_cfg_scale}"""
+        logger.info(debug_str)
+
+        # Prepare batch
+        batch = ForwardBatch(
+            **asdict(sampling_param),
+            eta=0.0,
+            n_tokens=n_tokens,
+            extra={},
+        )
+
+        # Run inference
+        start_time = time.time()
+        output_batch = self.executor.execute_forward(batch, fastvideo_args)
+        samples = output_batch
+
+        gen_time = time.time() - start_time
+        logger.info("Generated successfully in %.2f seconds", gen_time)
+
+        # Process outputs
+        videos = rearrange(samples, "b c t h w -> t b c h w")
+        frames = []
+        for x in videos:
+            x = torchvision.utils.make_grid(x, nrow=6)
+            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+            frames.append((x * 255).numpy().astype(np.uint8))
+
+        # Save video if requested
+        if batch.save_video:
+            save_path = batch.output_path
+            if save_path:
+                os.makedirs(os.path.dirname(save_path), exist_ok=True)
+                video_path = os.path.join(save_path, f"{prompt[:100]}.mp4")
+                imageio.mimsave(video_path, frames, fps=batch.fps)
+                logger.info("Saved video to %s", video_path)
+            else:
+                logger.warning("No output path provided, video not saved")
+
+        if batch.return_frames:
+            return frames
+        else:
+            return {
+                "samples": samples,
+                "prompts": prompt,
+                "size": (target_height, target_width, batch.num_frames),
+                "generation_time": gen_time
+            }

fastvideo/v1/envs.py

@@ -8,8 +8,6 @@ if TYPE_CHECKING:
     FASTVIDEO_RINGBUFFER_WARNING_INTERVAL: int = 60
     FASTVIDEO_NCCL_SO_PATH: Optional[str] = None
     LD_LIBRARY_PATH: Optional[str] = None
-    FASTVIDEO_USE_TRITON_FLASH_ATTN: bool = False
-    FASTVIDEO_FLASH_ATTN_VERSION: Optional[int] = None
     LOCAL_RANK: int = 0
     CUDA_VISIBLE_DEVICES: Optional[str] = None
     FASTVIDEO_CACHE_ROOT: str = os.path.expanduser("~/.cache/fastvideo")
@@ -127,18 +125,6 @@ environment_variables: Dict[str, Callable[[], Any]] = {
     "LD_LIBRARY_PATH":
     lambda: os.environ.get("LD_LIBRARY_PATH", None),
 
-    # flag to control if fastvideo should use triton flash attention
-    "FASTVIDEO_USE_TRITON_FLASH_ATTN":
-    lambda:
-    (os.environ.get("FASTVIDEO_USE_TRITON_FLASH_ATTN", "True").lower() in
-     ("true", "1")),
-
-    # Force fastvideo to use a specific flash-attention version (2 or 3), only valid
-    # when using the flash-attention backend.
-    "FASTVIDEO_FLASH_ATTN_VERSION":
-    lambda: maybe_convert_int(
-        os.environ.get("FASTVIDEO_FLASH_ATTN_VERSION", None)),
-
     # Internal flag to enable Dynamo fullgraph capture
     "FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE":
     lambda: bool(
@@ -184,7 +170,8 @@ environment_variables: Dict[str, Callable[[], Any]] = {
     # Available options:
     # - "TORCH_SDPA": use torch.nn.MultiheadAttention
     # - "FLASH_ATTN": use FlashAttention
-    # - "STA" : use sliding tile attention
+    # - "SLIDING_TILE_ATTN" : use Sliding Tile Attention
+    # - "SAGE_ATTN": use Sage Attention
     "FASTVIDEO_ATTENTION_BACKEND":
     lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),
 

fastvideo/v1/examples/inference/basic/README.md

@@ -0,0 +1,28 @@
+# Basic Video Generation Tutorial
+The `VideoGenerator` class provides the primary Python interface for doing offline video generation, which is interacting with a diffusion pipeline without using a separate inference api server.
+
+
+## Usage
+The first script in this example shows the most basic usage of FastVideo. If you are new to Python and FastVideo, you should start here.
+
+```bash
+python fastvideo/v1/examples/inference/basic/basic.py
+```
+
+## Basic Walkthrough
+
+All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines! 
+
+```python
+from fastvideo import VideoGenerator
+
+generator = VideoGenerator.from_pretrained(
+    "FastVideo/FastHunyuan-Diffusers",
+    num_gpus=2,
+)
+
+prompt = "A beautiful woman in a red dress walking down a street"
+video = generator.generate_video(prompt)
+```
+
+More to come! These examples and APIs are still under construction!

fastvideo/v1/examples/inference/basic/basic.py

@@ -0,0 +1,30 @@
+from fastvideo import VideoGenerator
+# from fastvideo.v1.configs.sample import SamplingParam
+
+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/FastHunyuan-Diffusers",
+        # if num_gpus > 1, FastVideo will automatically handle distributed setup
+        num_gpus=4,
+    )
+
+    # sampling_param = SamplingParam.from_pretrained("/workspace/data/Wan-AI/Wan2.1-I2V-14B-480P-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 beautiful woman in a red dress walking down a street"
+    video = generator.generate_video(prompt)
+    # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
+
+    # Generate another video with a different prompt, without reloading the
+    # model!
+    prompt2 = "A beautiful woman in a blue dress walking down a street"
+    video2 = generator.generate_video(prompt2)
+
+
+if __name__ == "__main__":
+    main()

fastvideo/v1/examples/inference/gradio/README.md

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

fastvideo/v1/examples/inference/gradio/gradio_demo.py

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

fastvideo/v1/fastvideo_args.py

@@ -4,68 +4,69 @@
 
 import argparse
 import dataclasses
+from contextlib import contextmanager
 from typing import List, Optional
 
+from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
+from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import FlexibleArgumentParser
 
+logger = init_logger(__name__)
+
 
 @dataclasses.dataclass
-class InferenceArgs:
+class FastVideoArgs:
     # Model and path configuration
     model_path: str
 
+    # Distributed executor backend
+    distributed_executor_backend: str = "mp"
+
+    inference_mode: bool = True  # if False == training mode
+
     # HuggingFace specific parameters
     trust_remote_code: bool = False
     revision: Optional[str] = None
 
     # Parallelism
-    tp_size: int = 1
-    sp_size: int = 1
+    num_gpus: int = 1
+    tp_size: Optional[int] = None
+    sp_size: Optional[int] = None
     dist_timeout: Optional[int] = None  # timeout for torch.distributed
 
     # Video generation parameters
-    height: int = 720
-    width: int = 1280
-    num_frames: int = 117
-    num_inference_steps: int = 50
-    guidance_scale: float = 1.0
-    guidance_rescale: float = 0.0
     embedded_cfg_scale: float = 6.0
-    flow_shift: int = 7
+    flow_shift: Optional[float] = None
 
     output_type: str = "pil"
 
-    # Model configuration
+    # DiT configuration
+    dit_config: DiTConfig = DiTConfig()
     precision: str = "bf16"
 
     # VAE configuration
     vae_precision: str = "fp16"
-    vae_tiling: bool = True
-    vae_sp: bool = False
+    vae_tiling: bool = True  # Might change in between forward passes
+    vae_sp: bool = False  # Might change in between forward passes
+    # vae_scale_factor: Optional[int] = None # Deprecated
+    vae_config: VAEConfig = VAEConfig()
+
+    # Image encoder configuration
+    image_encoder_precision: str = "fp32"
+    image_encoder_config: EncoderConfig = EncoderConfig()
 
     # Text encoder configuration
     text_encoder_precision: str = "fp16"
-    text_len: int = 256
-    hidden_state_skip_layer: int = 2
+    text_encoder_config: EncoderConfig = EncoderConfig()
 
     # Secondary text encoder
+    text_encoder_config_2: EncoderConfig = EncoderConfig()
     text_encoder_precision_2: str = "fp16"
-    text_len_2: int = 77
-
-    # Flow Matching parameters
-    flow_solver: str = "euler"
-    denoise_type: str = "flow"
 
     # STA (Spatial-Temporal Attention) parameters
     mask_strategy_file_path: Optional[str] = None
     enable_torch_compile: bool = False
 
-    # Scheduler options
-    scheduler_type: str = "euler"
-
-    neg_prompt: Optional[str] = None
-    num_videos: int = 1
-    fps: int = 24
     use_cpu_offload: bool = False
     disable_autocast: bool = False
 
@@ -73,10 +74,6 @@ class InferenceArgs:
     log_level: str = "info"
 
     # Inference parameters
-    prompt: Optional[str] = None
-    prompt_path: Optional[str] = None
-    output_path: str = "outputs/"
-    seed: int = 1024
     device_str: Optional[str] = None
     device = None
 
@@ -104,97 +101,75 @@ class InferenceArgs:
             help="Directory containing StepVideo model",
         )
 
+        # distributed_executor_backend
+        parser.add_argument(
+            "--distributed-executor-backend",
+            type=str,
+            choices=["mp"],
+            default=FastVideoArgs.distributed_executor_backend,
+            help="The distributed executor backend to use",
+        )
+
         # HuggingFace specific parameters
         parser.add_argument(
             "--trust-remote-code",
             action="store_true",
-            default=InferenceArgs.trust_remote_code,
+            default=FastVideoArgs.trust_remote_code,
             help="Trust remote code when loading HuggingFace models",
         )
         parser.add_argument(
             "--revision",
             type=str,
-            default=InferenceArgs.revision,
+            default=FastVideoArgs.revision,
             help=
             "The specific model version to use (can be a branch name, tag name, or commit id)",
         )
 
         # Parallelism
+        parser.add_argument(
+            "--num-gpus",
+            type=int,
+            default=FastVideoArgs.num_gpus,
+            help="The number of GPUs to use.",
+        )
         parser.add_argument(
             "--tensor-parallel-size",
             "--tp-size",
             type=int,
-            default=InferenceArgs.tp_size,
+            default=FastVideoArgs.tp_size,
             help="The tensor parallelism size.",
         )
         parser.add_argument(
             "--sequence-parallel-size",
             "--sp-size",
             type=int,
-            default=InferenceArgs.sp_size,
+            default=FastVideoArgs.sp_size,
             help="The sequence parallelism size.",
         )
         parser.add_argument(
             "--dist-timeout",
             type=int,
-            default=InferenceArgs.dist_timeout,
+            default=FastVideoArgs.dist_timeout,
             help="Set timeout for torch.distributed initialization.",
         )
 
-        # Video generation parameters
-        parser.add_argument(
-            "--height",
-            type=int,
-            default=InferenceArgs.height,
-            help="Height of generated video",
-        )
-        parser.add_argument(
-            "--width",
-            type=int,
-            default=InferenceArgs.width,
-            help="Width of generated video",
-        )
-        parser.add_argument(
-            "--num-frames",
-            type=int,
-            default=InferenceArgs.num_frames,
-            help="Number of frames to generate",
-        )
-        parser.add_argument(
-            "--num-inference-steps",
-            type=int,
-            default=InferenceArgs.num_inference_steps,
-            help="Number of inference steps",
-        )
-        parser.add_argument(
-            "--guidance-scale",
-            type=float,
-            default=InferenceArgs.guidance_scale,
-            help="Guidance scale for classifier-free guidance",
-        )
-        parser.add_argument(
-            "--guidance-rescale",
-            type=float,
-            default=InferenceArgs.guidance_rescale,
-            help="Guidance rescale for classifier-free guidance",
-        )
         parser.add_argument(
             "--embedded-cfg-scale",
             type=float,
-            default=InferenceArgs.embedded_cfg_scale,
+            default=FastVideoArgs.embedded_cfg_scale,
             help="Embedded CFG scale",
         )
         parser.add_argument(
             "--flow-shift",
             "--shift",
-            type=int,
-            default=InferenceArgs.flow_shift,
+            type=float,
+            default=FastVideoArgs.flow_shift,
             help="Flow shift parameter",
         )
         parser.add_argument(
             "--output-type",
             type=str,
-            default=InferenceArgs.output_type,
+            default=FastVideoArgs.output_type,
             choices=["pil"],
             help="Output type for the generated video",
         )
@@ -202,7 +177,7 @@ class InferenceArgs:
         parser.add_argument(
             "--precision",
             type=str,
-            default=InferenceArgs.precision,
+            default=FastVideoArgs.precision,
             choices=["fp32", "fp16", "bf16"],
             help="Precision for the model",
         )
@@ -211,14 +186,14 @@ class InferenceArgs:
         parser.add_argument(
             "--vae-precision",
             type=str,
-            default=InferenceArgs.vae_precision,
+            default=FastVideoArgs.vae_precision,
             choices=["fp32", "fp16", "bf16"],
             help="Precision for VAE",
         )
         parser.add_argument(
             "--vae-tiling",
             action="store_true",
-            default=InferenceArgs.vae_tiling,
+            default=FastVideoArgs.vae_tiling,
             help="Enable VAE tiling",
         )
         parser.add_argument(
@@ -230,45 +205,29 @@ class InferenceArgs:
         parser.add_argument(
             "--text-encoder-precision",
             type=str,
-            default=InferenceArgs.text_encoder_precision,
+            default=FastVideoArgs.text_encoder_precision,
             choices=["fp32", "fp16", "bf16"],
             help="Precision for text encoder",
         )
+
+        # Image encoder config
         parser.add_argument(
-            "--text-len",
-            type=int,
-            default=InferenceArgs.text_len,
-            help="Maximum text length",
+            "--image-encoder-precision",
+            type=str,
+            default=FastVideoArgs.image_encoder_precision,
+            choices=["fp32", "fp16", "bf16"],
+            help="Precision for image encoder",
         )
+
         # Secondary text encoder
 
         parser.add_argument(
             "--text-encoder-precision-2",
             type=str,
-            default=InferenceArgs.text_encoder_precision_2,
+            default=FastVideoArgs.text_encoder_precision_2,
             choices=["fp32", "fp16", "bf16"],
             help="Precision for secondary text encoder",
         )
-        parser.add_argument(
-            "--text-len-2",
-            type=int,
-            default=InferenceArgs.text_len_2,
-            help="Maximum secondary text length",
-        )
-
-        # Flow Matching parameters
-        parser.add_argument(
-            "--flow-solver",
-            type=str,
-            default=InferenceArgs.flow_solver,
-            help="Solver for flow matching",
-        )
-        parser.add_argument(
-            "--denoise-type",
-            type=str,
-            default=InferenceArgs.denoise_type,
-            help="Denoise type for noised inputs",
-        )
 
         # STA (Spatial-Temporal Attention) parameters
         parser.add_argument(
@@ -283,33 +242,6 @@ class InferenceArgs:
             "Use torch.compile for speeding up STA inference without teacache",
         )
 
-        # Scheduler options
-        parser.add_argument(
-            "--scheduler-type",
-            type=str,
-            default=InferenceArgs.scheduler_type,
-            help="Type of scheduler to use",
-        )
-
-        # HunYuan specific parameters
-        parser.add_argument(
-            "--neg-prompt",
-            type=str,
-            default=InferenceArgs.neg_prompt,
-            help="Negative prompt for sampling",
-        )
-        parser.add_argument(
-            "--num-videos",
-            type=int,
-            default=InferenceArgs.num_videos,
-            help="Number of videos to generate per prompt",
-        )
-        parser.add_argument(
-            "--fps",
-            type=int,
-            default=InferenceArgs.fps,
-            help="Frames per second for output video",
-        )
         parser.add_argument(
             "--use-cpu-offload",
             action="store_true",
@@ -326,40 +258,14 @@ class InferenceArgs:
         parser.add_argument(
             "--log-level",
             type=str,
-            default=InferenceArgs.log_level,
+            default=FastVideoArgs.log_level,
             help="The logging level of all loggers.",
         )
 
-        # Inference parameters
-        prompt_group = parser.add_mutually_exclusive_group(required=True)
-        prompt_group.add_argument(
-            "--prompt",
-            type=str,
-            help="Text prompt for video generation",
-        )
-        prompt_group.add_argument(
-            "--prompt-path",
-            type=str,
-            help="Path to a text file containing the prompt",
-        )
-
-        parser.add_argument(
-            "--output-path",
-            type=str,
-            default=InferenceArgs.output_path,
-            help="Directory to save generated videos",
-        )
-        parser.add_argument(
-            "--seed",
-            type=int,
-            default=InferenceArgs.seed,
-            help="Random seed for reproducibility",
-        )
-
         return parser
 
     @classmethod
-    def from_cli_args(cls, args: argparse.Namespace) -> "InferenceArgs":
+    def from_cli_args(cls, args: argparse.Namespace) -> "FastVideoArgs":
         args.tp_size = args.tensor_parallel_size
         args.sp_size = args.sequence_parallel_size
         args.flow_shift = getattr(args, "shift", args.flow_shift)
@@ -384,22 +290,32 @@ class InferenceArgs:
 
         return cls(**kwargs)
 
-    def check_inference_args(self) -> None:
+    def check_fastvideo_args(self) -> None:
         """Validate inference arguments for consistency"""
+        if self.tp_size is None:
+            self.tp_size = self.num_gpus
+        if self.sp_size is None:
+            self.sp_size = self.num_gpus
+
+        if self.num_gpus < max(self.tp_size, self.sp_size):
+            self.num_gpus = max(self.tp_size, self.sp_size)
+
+        if self.tp_size != self.sp_size:
+            raise ValueError(
+                f"tp_size ({self.tp_size}) must be equal to sp_size ({self.sp_size})"
+            )
 
         # Validate VAE spatial parallelism with VAE tiling
         if self.vae_sp and not self.vae_tiling:
             raise ValueError(
                 "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True."
             )
-        if self.prompt_path and not self.prompt_path.endswith(".txt"):
-            raise ValueError("prompt_path must be a text file")
 
 
-_inference_args = None
+_current_fastvideo_args = None
 
 
-def prepare_inference_args(argv: List[str]) -> InferenceArgs:
+def prepare_fastvideo_args(argv: List[str]) -> FastVideoArgs:
     """
     Prepare the inference arguments from the command line arguments.
 
@@ -411,26 +327,38 @@ def prepare_inference_args(argv: List[str]) -> InferenceArgs:
         The inference arguments.
     """
     parser = FlexibleArgumentParser()
-    InferenceArgs.add_cli_args(parser)
+    FastVideoArgs.add_cli_args(parser)
     raw_args = parser.parse_args(argv)
-    inference_args = InferenceArgs.from_cli_args(raw_args)
-    inference_args.check_inference_args()
-    global _inference_args
-    _inference_args = inference_args
-    return inference_args
+    fastvideo_args = FastVideoArgs.from_cli_args(raw_args)
+    fastvideo_args.check_fastvideo_args()
+    global _current_fastvideo_args
+    _current_fastvideo_args = fastvideo_args
+    return fastvideo_args
 
 
-def get_inference_args() -> InferenceArgs:
-    global _inference_args
-    if _inference_args is None:
-        raise ValueError("Inference arguments not set")
-    return _inference_args
+@contextmanager
+def set_current_fastvideo_args(fastvideo_args: FastVideoArgs):
+    """
+    Temporarily set the current fastvideo config.
+    Used during model initialization.
+    We save the current fastvideo config in a global variable,
+    so that all modules can access it, e.g. custom ops
+    can access the fastvideo config to determine how to dispatch.
+    """
+    global _current_fastvideo_args
+    old_fastvideo_args = _current_fastvideo_args
+    try:
+        _current_fastvideo_args = fastvideo_args
+        yield
+    finally:
+        _current_fastvideo_args = old_fastvideo_args
 
 
-class DeprecatedAction(argparse.Action):
-
-    def __init__(self, option_strings, dest, nargs=0, **kwargs):
-        super().__init__(option_strings, dest, nargs=nargs, **kwargs)
-
-    def __call__(self, parser, namespace, values, option_string=None):
-        raise ValueError(self.help)
+def get_current_fastvideo_args() -> FastVideoArgs:
+    if _current_fastvideo_args is None:
+        # in ci, usually when we test custom ops/modules directly,
+        # we don't set the fastvideo config. In that case, we set a default
+        # config.
+        # TODO(will): may need to handle this for CI.
+        raise ValueError("Current fastvideo args is not set.")
+    return _current_fastvideo_args

fastvideo/v1/forward_context.py

@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Optional
 
 import torch
 
-from fastvideo.v1.inference_args import InferenceArgs
+from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 
 if TYPE_CHECKING:
@@ -52,7 +52,7 @@ def get_forward_context() -> ForwardContext:
 @contextmanager
 def set_forward_context(current_timestep,
                         attn_metadata,
-                        inference_args: InferenceArgs = None):
+                        fastvideo_args: Optional[FastVideoArgs] = None):
     """A context manager that stores the current forward context,
     can be attention metadata, etc.
     Here we can inject common logic for every model forward pass.

fastvideo/v1/inference_engine.py

@@ -1,3 +1,5 @@
+# type: ignore
+
 # SPDX-License-Identifier: Apache-2.0
 """
 Inference module for diffusion models.
@@ -10,7 +12,7 @@ from typing import Any, Dict
 
 import torch
 
-from fastvideo.v1.inference_args import InferenceArgs
+from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch,
                                     build_pipeline)
@@ -28,29 +30,29 @@ class InferenceEngine:
     def __init__(
         self,
         pipeline: ComposedPipelineBase,
-        inference_args: InferenceArgs,
+        fastvideo_args: FastVideoArgs,
     ):
         """
         Initialize the inference engine.
         
         Args:
             pipeline: The pipeline to use for inference.
-            inference_args: The inference arguments.
+            fastvideo_args: The inference arguments.
             default_negative_prompt: The default negative prompt to use.
         """
         self.pipeline = pipeline
-        self.inference_args = inference_args
+        self.fastvideo_args = fastvideo_args
 
     @classmethod
     def create_engine(
         cls,
-        inference_args: InferenceArgs,
+        fastvideo_args: FastVideoArgs,
     ) -> "InferenceEngine":
         """
         Create an inference engine with the specified arguments.
         
         Args:
-            inference_args: The inference arguments.
+            fastvideo_args: The inference arguments.
             model_loader_cls: The model loader class to use. If None, it will be
                 determined from the model type.
             pipeline_type: The type of pipeline to create. If None, it will be
@@ -71,16 +73,16 @@ class InferenceEngine:
         # this way for training we can just do pipeline_cls.from_pretrained(
         # checkpoint_path) and have it handle everything.
         # TODO(Peiyuan): Then maybe we should only pass in model path and device, not the entire inference args?
-        pipeline = build_pipeline(inference_args)
+        pipeline = build_pipeline(fastvideo_args)
         logger.info("Pipeline Ready")
 
         # Create the inference engine
-        return cls(pipeline, inference_args)
+        return cls(pipeline, fastvideo_args)
 
     def run(
         self,
         prompt: str,
-        inference_args: InferenceArgs,
+        fastvideo_args: FastVideoArgs,
     ) -> Dict[str, Any]:
         """
         Run inference with the pipeline.
@@ -96,16 +98,17 @@ class InferenceEngine:
         """
         out_dict: Dict[str, Any] = dict()
 
-        num_videos_per_prompt = inference_args.num_videos
-        seed = inference_args.seed
-        height = inference_args.height
-        width = inference_args.width
-        video_length = inference_args.num_frames
-        negative_prompt = inference_args.neg_prompt
-        infer_steps = inference_args.num_inference_steps
-        guidance_scale = inference_args.guidance_scale
-        flow_shift = inference_args.flow_shift
-        embedded_guidance_scale = inference_args.embedded_cfg_scale
+        num_videos_per_prompt = fastvideo_args.num_videos
+        seed = fastvideo_args.seed
+        height = fastvideo_args.height
+        width = fastvideo_args.width
+        video_length = fastvideo_args.num_frames
+        negative_prompt = fastvideo_args.neg_prompt
+        infer_steps = fastvideo_args.num_inference_steps
+        guidance_scale = fastvideo_args.guidance_scale
+        flow_shift = fastvideo_args.flow_shift
+        embedded_guidance_scale = fastvideo_args.embedded_cfg_scale
+        image_path = fastvideo_args.image_path
 
         # ========================================================================
         # Arguments: target_width, target_height, target_video_length
@@ -160,20 +163,21 @@ class InferenceEngine:
         # return
         # sp_group = get_sp_group()
         # local_rank = sp_group.rank
-        device = torch.device(inference_args.device_str)
+        device = torch.device(fastvideo_args.device_str)
         batch = ForwardBatch(
+            image_path=image_path,
             prompt=prompt,
             negative_prompt=negative_prompt,
             num_videos_per_prompt=num_videos_per_prompt,
-            height=inference_args.height,
-            width=inference_args.width,
-            num_frames=inference_args.num_frames,
-            num_inference_steps=inference_args.num_inference_steps,
-            guidance_scale=inference_args.guidance_scale,
+            height=fastvideo_args.height,
+            width=fastvideo_args.width,
+            num_frames=fastvideo_args.num_frames,
+            num_inference_steps=fastvideo_args.num_inference_steps,
+            guidance_scale=fastvideo_args.guidance_scale,
             # generator=generator,
             eta=0.0,
             n_tokens=n_tokens,
-            data_type="video" if inference_args.num_frames > 1 else "image",
+            data_type="video" if fastvideo_args.num_frames > 1 else "image",
             device=device,
             extra={},  # Any additional parameters
         )
@@ -182,7 +186,7 @@ class InferenceEngine:
         print(batch)
         print('===============================================')
         print('===============================================')
-        print(inference_args)
+        print(fastvideo_args)
 
         # ========================================================================
         # Pipeline inference
@@ -190,7 +194,7 @@ class InferenceEngine:
         start_time = time.time()
         samples = self.pipeline.forward(
             batch=batch,
-            inference_args=inference_args,
+            fastvideo_args=fastvideo_args,
         ).output
         # TODO(will): fix and move to hunyuan stage
         # out_dict["seeds"] = batch.seeds

fastvideo/v1/layers/activation.py

@@ -23,7 +23,7 @@ class SiluAndMul(CustomOp):
         return: (num_tokens, d) or (batch_size, seq_len, d)
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         if current_platform.is_cuda_alike() or current_platform.is_cpu():
             self.op = torch.ops._C.silu_and_mul

fastvideo/v1/layers/layernorm.py

@@ -113,7 +113,7 @@ class ScaleResidual(nn.Module):
     Applies gated residual connection.
     """
 
-    def __init__(self):
+    def __init__(self, prefix: str = ""):
         super().__init__()
 
     def forward(self, residual: torch.Tensor, x: torch.Tensor,
@@ -139,6 +139,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
         eps: float = 1e-6,
         elementwise_affine: bool = False,
         dtype: torch.dtype = torch.float32,
+        prefix: str = "",
     ):
         super().__init__()
         if norm_type == "rms":
@@ -189,6 +190,7 @@ class LayerNormScaleShift(nn.Module):
         eps: float = 1e-6,
         elementwise_affine: bool = False,
         dtype: torch.dtype = torch.float32,
+        prefix: str = "",
     ):
         super().__init__()
         if norm_type == "rms":

fastvideo/v1/layers/linear.py

@@ -778,12 +778,13 @@ class QKVParallelLinear(ColumnParallelLinear):
             # no need to narrow
             is_sharded_weight = is_sharded_weight
 
+            shard_idx = 0
             param_data = param_data.narrow(output_dim, shard_offset, shard_size)
             if loaded_shard_id == "q":
-                shard_id = tp_rank
+                shard_idx = tp_rank
             else:
-                shard_id = tp_rank // self.num_kv_head_replicas
-            start_idx = shard_id * shard_size
+                shard_idx = tp_rank // self.num_kv_head_replicas
+            start_idx = shard_idx * shard_size
 
             if not is_sharded_weight:
                 loaded_weight = loaded_weight.narrow(output_dim, start_idx,

fastvideo/v1/layers/mlp.py

@@ -12,7 +12,6 @@ from fastvideo.v1.layers.linear import ReplicatedLinear
 class MLP(nn.Module):
     """
     MLP for DiT blocks, NO gated linear units
-    TODO: add Tensor Parallel
     """
 
     def __init__(
@@ -23,6 +22,7 @@ class MLP(nn.Module):
         bias: bool = True,
         act_type: str = "gelu_pytorch_tanh",
         dtype: Optional[torch.dtype] = None,
+        prefix: str = "",
     ):
         super().__init__()
         self.fc_in = ReplicatedLinear(

fastvideo/v1/layers/rotary_embedding.py

@@ -84,7 +84,7 @@ class RotaryEmbedding(CustomOp):
         head_size: int,
         rotary_dim: int,
         max_position_embeddings: int,
-        base: int,
+        base: Union[int, float],
         is_neox_style: bool,
         dtype: torch.dtype,
     ) -> None:
@@ -446,7 +446,7 @@ def get_rope(
     head_size: int,
     rotary_dim: int,
     max_position: int,
-    base: int,
+    base: Union[int, float],
     is_neox_style: bool = True,
     rope_scaling: Optional[Dict[str, Any]] = None,
     dtype: Optional[torch.dtype] = None,