fix tensor type hint

update min python to 3.10
[Misc] Small fixes to Torch code (#395 )
2025-05-23 14:42:23 -07:00 · 2025-05-23 14:42:23 -07:00 · 2025-05-23 14:40:24 -07:00 · 2025-05-23 14:19:45 -07:00 · 2025-05-23 14:19:18 -07:00 · 2025-05-23 14:17:33 -07:00
155 changed files with 6225 additions and 1796 deletions
@@ -8,7 +8,7 @@ body:
  attributes:
    label: Environment
    description: |
-      Please share your environment with us. You can run the command **python fastvideo/utils/env_utils.py** and copy-paste its output below.
+      Please share your environment with us. You can run the command **python fastvideo/utils/collect_env.py** and copy-paste its output below.
    placeholder: FastVideo version, platform, python version, cuda version...
  validations:
    required: true
@@ -83,193 +83,81 @@ jobs:
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.encoder-test == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_encoder_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: "encoder-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/encoders -s"
-
-      - name: Terminate RunPod Instances
-        if: ${{ always() }}
-        env:
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          JOB_ID: "encoder-test"
-        run: python .github/scripts/runpod_cleanup.py
+    uses: ./.github/workflows/runpod-test.yml
+    with:
+      job_id: "encoder-test"
+      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/encoders -s"
+      timeout_minutes: 30
+    secrets:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

  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
+    uses: ./.github/workflows/runpod-test.yml
+    with:
+      job_id: "vae-test"
+      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"
+      timeout_minutes: 30
+    secrets:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

  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
+    uses: ./.github/workflows/runpod-test.yml
+    with:
+      job_id: "transformer-test"
+      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"
+      timeout_minutes: 30
+    secrets:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
      
  ssim-test:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_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: "ssim-test"
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-        timeout-minutes: 45
-        run: >-
-          python .github/scripts/runpod_api.py 
-          --gpu-type "NVIDIA A40"
-          --gpu-count 2
-          --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() }}
-        env:
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          JOB_ID: "ssim-test"
-        run: python .github/scripts/runpod_cleanup.py
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [
+          # {version: "3.10", tag: "latest"},
+          # {version: "3.11", tag: "py3.11-latest"},
+          {version: "3.12", tag: "py3.12-latest"}
+        ]
+    uses: ./.github/workflows/runpod-test.yml
+    with:
+      job_id: "ssim-test-py${{ matrix.python-version.version }}"
+      gpu_type: "NVIDIA A40"
+      gpu_count: 2
+      volume_size: 200
+      disk_size: 200
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
+      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
+      timeout_minutes: 60
+    secrets:
+      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}

  runpod-cleanup:
    needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
@@ -289,7 +177,7 @@ jobs:

      - name: Cleanup all RunPod instances
        env:
-          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test"]' # JSON array of job IDs
+          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12"]'
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
        run: python .github/scripts/runpod_cleanup.py
@@ -0,0 +1,91 @@
+name: RunPod Test
+
+on:
+  workflow_call:
+    inputs:
+      job_id:
+        required: true
+        type: string
+        description: "Unique identifier for this test job"
+      gpu_type:
+        required: true
+        type: string
+        description: "GPU type to use (e.g. NVIDIA A40, NVIDIA L40S)"
+      gpu_count:
+        required: true
+        type: number
+        description: "Number of GPUs to use"
+      volume_size:
+        required: false
+        type: number
+        default: 20
+        description: "Volume size in GB"
+      disk_size:
+        required: false
+        type: number
+        default: 20
+        description: "Disk size in GB"
+      image:
+        required: true
+        type: string
+        description: "Docker image to use"
+      test_command:
+        required: true
+        type: string
+        description: "Command to run tests"
+      timeout_minutes:
+        required: false
+        type: number
+        default: 30
+        description: "Timeout in minutes"
+    secrets:
+      RUNPOD_API_KEY:
+        required: true
+      RUNPOD_PRIVATE_KEY:
+        required: true
+
+jobs:
+  run-test:
+    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: ${{ inputs.job_id }}
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: ${{ inputs.timeout_minutes }}
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "${{ inputs.gpu_type }}"
+          --gpu-count ${{ inputs.gpu_count }}
+          --volume-size ${{ inputs.volume_size }}
+          --disk-size ${{ inputs.disk_size }}
+          --image "${{ inputs.image }}"
+          --test-command "${{ inputs.test_command }}"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: ${{ inputs.job_id }}
+        run: python .github/scripts/runpod_cleanup.py 
@@ -22,7 +22,8 @@ exclude: |
    examples/.*|
    .github/workflows/fastvideo-publish.yml|
    .github/workflows/sta-publish.yml|
-    .github/workflows/build-image-template.yml
+    .github/workflows/build-image-template.yml|
+    docs/source/inference/support_matrix.md
  )
 repos:
 - repo: https://github.com/google/yapf
@@ -2,62 +2,109 @@
 <img src=assets/logo.jpg width="30%"/>
 </div>

-FastVideo is a lightweight framework for accelerating large video diffusion models.
+**FastVideo is a unified framework for accelerated video generation.**
+
+It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
+With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.

 <p align="center">
-    | <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> |
+    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
 </p>

-https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1
+<div align="center">
+<img src=assets/perf.png width="90%"/>
+</div>

-FastVideo currently offers: (with more to come)
+## Key Features

- [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.
+FastVideo has the following features:
+- State-of-the-art performance optimizations for inference
+  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [TeaCache](https://arxiv.org/pdf/2411.19108)
+  - [Sage Attention](https://arxiv.org/abs/2410.02367)
+- Cutting edge models
+  - Wan2.1 T2V, I2V
+  - HunyuanVideo
+  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
+  - StepVideo T2V
+- Distillation support
+  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
 - Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
 - Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.

-Dev in progress and highly experimental.
-
-## Change Log
- ```2025/02/20```: FastVideo now supports STA on [StepVideo](https://github.com/stepfun-ai/Step-Video-T2V) with 3.4X speedup!
- ```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` v0.0.1 is released.
-
 ## Getting Started
+We recommend using an environment manager such as `Conda` to create a clean environment:
+
+```bash
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
+pip install fastvideo
+```
+
+Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) for more detailed installation instructions.
+
+## Inference
+### Generating Your First Video
+Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
+
+```python
+from fastvideo import VideoGenerator
+
+def main():
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
+
+    # Generate the video
+    video = generator.generate_video(
+        prompt,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Run the script with:
+
+```bash
+python example.py
+```
+
+For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html).
+
+### Other docs:

- [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)

-### Inference
- [Quick Start](https://hao-ai-lab.github.io/FastVideo/inference/examples/basic.html)
- V1 Inference API Guide (Coming soon!)
-
-### Distillation and Finetuning
+## 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)

-### 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 models support
  <!-- - [ ] Add CogvideoX model -->
-  - [ ] Add StepVideo to V1
+  - [x] Add StepVideo to V1
 - Optimization features
-  - [ ] Teacache in V1
-  - [ ] SageAttention in V1
+  - [x] Teacache in V1
+  - [x] SageAttention in V1
 - Code updates
-  - [ ] V1 Configuration API
+  - [x] V1 Configuration API
  - [ ] Support Training in V1
  <!-- - [ ] fp8 support -->
  <!-- - [ ] faster load model and save model support -->
@@ -2,7 +2,6 @@ import torch
 from flex_sta_ref import get_sliding_tile_attention_mask
 from st_attn import sliding_tile_attention
 from torch.nn.attention.flex_attention import flex_attention
-# from flash_attn_interface import flash_attn_func
 from tqdm import tqdm

 flex_attention = torch.compile(flex_attention, dynamic=False)
@@ -23,7 +22,7 @@ def h100_fwd_kernel_test(Q, K, V, kernel_size):
 def generate_tensor(shape, mean, std, dtype, device):
    tensor = torch.randn(shape, dtype=dtype, device=device)

-    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
+    magnitude = torch.linalg.norm(tensor, dim=-1, keepdim=True)
    scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude

    return scaled_tensor.contiguous()
@@ -18,7 +18,6 @@ import os
 import re
 import sys
 from pathlib import Path
-from typing import Optional

 import requests

@@ -168,8 +167,7 @@ _cached_base: str = ""
 _cached_branch: str = ""


-def get_repo_base_and_branch(
-        pr_number: str) -> tuple[Optional[str], Optional[str]]:
+def get_repo_base_and_branch(pr_number: str) -> tuple[str | None, str | None]:
    global _cached_base, _cached_branch
    if _cached_base and _cached_branch:
        return _cached_base, _cached_branch
@@ -5,7 +5,6 @@ import itertools
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Optional

 ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
 ROOT_DIR_RELATIVE = '../../../..'
@@ -89,7 +88,7 @@ class Example:
        generate() -> str: Generates the documentation content.
    """ # noqa: E501
    path: Path
-    category: Optional[str] = None
+    category: str | None = None
    main_file: Path = field(init=False)
    other_files: list[Path] = field(init=False)
    title: str = field(init=False)
@@ -11,11 +11,52 @@ FastVideo currently only supports Linux and NVIDIA CUDA GPUs.
 - **CUDA 12.4**
 - **At least 1 NVIDIA GPU**

-## Quick start
-We recommend using an environment manager such as [Conda](#miniconda-recommended)
+## Set up using Python
+### Create a new Python environment
+
+#### Conda
+You can create a new python environment using [Conda](https://docs.conda.io/projects/conda/en/stable/user-guide/getting-started.html)
+##### 1. 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. Create and activate a Conda environment for FastVideo
+
+```bash
+# (Recommended) Create a new conda environment.
+conda create -n fastvideo python=3.12 -y
+conda activate fastvideo
+```
+
+:::{note}
+[PyTorch has deprecated the conda release channel](https://github.com/pytorch/pytorch/issues/138506). If you use `conda`, please only use it to create Python environment rather than installing packages.
+:::
+
+#### uv
+
+:::{tip}
+We highly recommend using `uv` to install FastVideo. In our experience, `uv` speeds up installation by at least 3x.
+:::
+
+Or you can create a new Python environment using [uv](https://docs.astral.sh/uv/), a very fast Python environment manager. Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, you can create a new Python environment using the following command:
+
+```console
+# (Recommended) Create a new uv environment. Use `--seed` to install `pip` and `setuptools` in the environment.
+uv venv --python 3.12 --seed
+source .venv/bin/activate
+```
+
+### Installation

 ```bash
 pip install fastvideo
+
+# or if you are using uv
+uv pip install fastvideo
 ```

 Also optionally install flash-attn:
@@ -24,30 +65,9 @@ Also optionally install flash-attn:
 pip install flash-attn==2.7.4.post1 --no-build-isolation
 ```

-[Click here to install from source](#installation-from-source)
+### Installation from Source

-## Environment Setup
-
-We highly recommend using a `conda` or `virtualenv` environment:
-### Miniconda (Recommended)
-#### 1. 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. Create and activate a Conda environment for FastVideo
-
-```bash
-conda create -n fastvideo python=3.10 -y
-conda activate fastvideo
-```
-
-## Installation from Source
-
-### 1. Clone the FastVideo repository
+#### 1. Clone the FastVideo repository

 ```bash
 git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
@@ -59,23 +79,23 @@ Basic installation:

 ```bash
 pip install -e .
+
+# or if you are using uv
+uv pip install -e .
 ```

-## Optional Dependencies
+### Optional Dependencies

-### Flash Attention
+#### Flash Attention

 ```bash
 pip install flash-attn==2.7.4.post1 --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.
-
-## Docker Images
+## Set up using Docker
 We also have prebuilt docker images with FastVideo dependencies pre-installed:
 [Docker Images](#docker)
+
 ## Development Environment Setup

 If you're planning to contribute to FastVideo please see the following page:
@@ -84,7 +104,7 @@ If you're planning to contribute to FastVideo please see the following page:
 ## Hardware Requirements

 ### For Basic Inference
- NVIDIA GPU with CUDA 12.4support
+- NVIDIA GPU with CUDA 12.4 support

 ### For Lora Finetuning
 - 40GB GPU memory each for 2 GPUs with lora
@@ -1,143 +0,0 @@
-# Quick Start
-
-This page contains instructions and code to get you quickly started with video generation using FastVideo.
-
-## Requirements
- **OS: Linux**
- **Python: 3.10-3.12**
- **CUDA 12.4**
- **At least 1 NVIDIA GPU**
-
-## Installation
-
-We recommend using a environment manager such as `Conda`.
-
-```bash
-pip install fastvideo
-```
-
-Also see the [Installation Guide](installation.md).
-
-## Generating Your First Video
-Here's a minimal example to generate a video using the default settings. All of the following code snippets can be directly copied into a Python file and executed with
-
-```bash
-python example.py
-```
-
-```python
-from fastvideo import VideoGenerator
-
-def main():
-    # Create a video generator with a pre-trained model
-    generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        num_gpus=1,  # Adjust based on your hardware
-    )
-
-    # Define a prompt for your video
-    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
-
-    # Generate the video
-    video = generator.generate_video(prompt)
-if __name__ == '__main__':
-    main()
-```
-
-The generated video will be saved in the current directory under `outputs/` by default.
-
-## Customizing Generation
-
-You can customize various parameters when generating videos:
-
-```python
-from fastvideo import VideoGenerator, SamplingParam
-
-def main():
-    # Create the generator
-    generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        num_gpus=1,
-    )
-
-    # Create and customize sampling parameters
-    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
-    # How many frames to generate
-    sampling_param.num_frames = 45
-    # How many steps we denoise the video
-    sampling_param.num_inference_steps = 30
-    # How strongly does the video to conform to the prompt
-    sampling_param.guidance_scale = 7.5
-
-    # Optional: provide an initial image for image-to-video generation
-    sampling_param.image_path = "path/to/your/image.jpg"  # Optional
-
-    # Generate video with custom parameters
-    prompt = "A beautiful sunset over a calm ocean, with gentle waves."
-    video = generator.generate_video(
-        prompt, 
-        sampling_param=sampling_param, 
-        output_path="my_videos/", # controls where videos are saved
-        return_frames=True # also return frames from this call (defaults to False)
-    )
-
-    # `video` now contains frames
-if __name__ == '__main__':
-    main()
-```
-
-## Available Models
-
-FastVideo supports various models for text-to-video generation:
-
- `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` - A popular text-to-video model
- `FastVideo/FastHunyuan-Diffusers` - A high-performance model for video generation
-
-## Advanced Configuration
-
-You can use PipelineConfig for more advanced customization:
-
-```python
-from fastvideo import VideoGenerator, PipelineConfig
-
-# Load the default configuration for a model
-config = PipelineConfig.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
-
-# Modify configuration settings
-config.vae_config.scale_factor = 16
-config.vae_config.precision = "fp16"
-
-# Create generator with custom config
-generator = VideoGenerator.from_pretrained(
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-    num_gpus=4,
-    pipeline_config=config,
-)
-
-# Generate video
-prompt = "A futuristic cityscape with flying cars and neon signs."
-video = generator.generate_video(prompt)
-```
-
-## Multi-GPU Setup
-
-FastVideo automatically handles distributed setup when multiple GPUs are specified:
-
-```python
-generator = VideoGenerator.from_pretrained(
-    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-    num_gpus=4,  # Will use 4 GPUs in parallel
-)
-```
-
-## Hardware Requirements
-
- **Basic Inference**: Minimum 20GB VRAM for quantized models (e.g., single RTX 4090)
- **LoRA Finetuning**: 40GB GPU memory each for 2 GPUs with LoRA
- **Full Finetuning/Distillation**: Multiple high-memory GPUs recommended (e.g., H100)
-
-## Next Steps
-
- Explore the [API Reference](../api/index.md) for more details
- Learn about [Advanced Inference Options](../inference/overview_back.md)
- See [Examples](../examples/index.md) for more usage scenarios
@@ -10,9 +10,74 @@ This class will be the primary Python API for generating videos and images.
    fastvideo.VideoGenerator
 ```

-VideoGenerator.from_pretrained() should be the primary way of creating a new video generator.
+`````{py:class} VideoGenerator(fastvideo_args: fastvideo.v1.fastvideo_args.FastVideoArgs, executor_class: type[fastvideo.v1.worker.executor.Executor], log_stats: bool)
+:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator

-````{py:method} from_pretrained(model_path: str, device: typing.optional[str] = none, torch_dtype: typing.optional[torch.dtype] = none, pipeline_config: typing.optional[typing.union[str | fastvideo.v1.configs.pipelines.pipelineconfig]] = none, **kwargs) -> v1.entrypoints.video_generator.videogenerator
-:canonical: v1.entrypoints.video_generator.videogenerator.from_pretrained
+```{autodoc2-docstring} fastvideo.v1.entrypoints.video_generator.VideoGenerator
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+`VideoGenerator.from_pretrained()` should be the primary way of creating a new video generator.
+
+````{py:method} from_pretrained(model_path: str, device: typing.Optional[str] = None, torch_dtype: typing.Optional[torch.dtype] = None, pipeline_config: typing.Optional[typing.Union[str | fastvideo.v1.configs.pipelines.PipelineConfig]] = None, **kwargs) -> fastvideo.v1.entrypoints.video_generator.VideoGenerator
+:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator.from_pretrained
 :classmethod:
+
+```{autodoc2-docstring} fastvideo.v1.entrypoints.video_generator.VideoGenerator.from_pretrained
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+
+## Configuring FastVideo
+
+The follow two classes `PipelineConfig` and `SamplingParam` are used to configure initialization and sampling parameters, respectively.
+
+### PipelineConfig
+```{autodoc2-summary}
+    fastvideo.PipelineConfig
+```
+
+`````{py:class} PipelineConfig
+:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig
+
+```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+````{py:method} from_pretrained(model_path: str) -> fastvideo.v1.configs.pipelines.base.PipelineConfig
+:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig.from_pretrained
+:classmethod:
+
+```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig.from_pretrained
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+
+````{py:method} dump_to_json(file_path: str)
+:canonical: fastvideo.v1.configs.pipelines.base.PipelineConfig.dump_to_json
+
+```{autodoc2-docstring} fastvideo.v1.configs.pipelines.base.PipelineConfig.dump_to_json
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+
+### SamplingParam
+
+```{autodoc2-summary}
+    fastvideo.SamplingParam
+```
+
+`````{py:class} SamplingParam
+:canonical: fastvideo.v1.configs.sample.base.SamplingParam
+
+```{autodoc2-docstring} fastvideo.v1.configs.sample.base.SamplingParam
+:parser: docs.source.autodoc2_docstring_parser
+```
+
+````{py:method} from_pretrained(model_path: str) -> fastvideo.v1.configs.sample.base.SamplingParam
+:canonical: fastvideo.v1.configs.sample.base.SamplingParam.from_pretrained
+:classmethod:
+
+```{autodoc2-docstring} fastvideo.v1.configs.sample.base.SamplingParam.from_pretrained
+:parser: docs.source.autodoc2_docstring_parser
 ```
@@ -9,7 +9,7 @@

 :::{raw} html
 <p style="text-align:center">
-<strong>FastVideo is a lightweight framework for accelerating large video diffusion models.
+<strong>FastVideo is a unified framework for accelerated video generation.
 </strong>
 </p>

@@ -21,27 +21,31 @@
 </p>
 :::

-FastVideo is a lightweight framework for accelerating large video diffusion models developed by the [Hao AI Lab](https://hao-ai-lab.github.io/).
+It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
+With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.

 <div style="text-align: center;">
-  <video controls width="800">
-    <source src="https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1" type="video/mp4">
-    Your browser does not support the video tag.
-  </video>
+  <img src=_static/images/perf.png width="100%"/>
 </div>

-FastVideo currently offers: (with more to come)
+## Key Features

- [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.
+FastVideo has the following features:
+- State-of-the-art performance optimizations for inference
+  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [TeaCache](https://arxiv.org/pdf/2411.19108)
+  - [Sage Attention](https://arxiv.org/abs/2410.02367)
+- Cutting edge models
+  - Wan2.1 T2V, I2V
+  - HunyuanVideo
+  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
+  - StepVideo T2V
+- Distillation support
+  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
 - Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
 - Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.

-Dev in progress and highly experimental.
-
 ## Documentation

 % How to start using FastVideo?
@@ -51,17 +55,20 @@ Dev in progress and highly experimental.
 :maxdepth: 1

 getting_started/installation
-getting_started/quick_start
-getting_started/v1_api
+<!-- getting_started/v1_api -->
 :::

 :::{toctree}
 :caption: Inference
 :maxdepth: 1

-inference/examples/basic
-inference/examples/gradio
-inference/examples/optimizations
+inference/inference_quick_start
+inference/configuration
+inference/optimizations
+inference/support_matrix
+inference/examples/examples_inference_index
+inference/cli
+inference/add_pipeline
 inference/v0_inference
 :::

@@ -96,14 +103,13 @@ design/overview

 contributing/overview
 contributing/developer_env/index
-contributing/add_pipeline
 :::

 :::{toctree}
 :caption: API Reference
 :maxdepth: 2

-api/summary
+<!-- api/summary -->
 api/fastvideo/fastvideo
 :::

@@ -1,6 +1,6 @@
 (add-pipeline)=

-# 🏗️ Adding a New Diffusion Pipeline
+# 🏗️ Adding a New Pipeline

 This guide explains how to implement a custom diffusion pipeline in FastVideo, leveraging the framework's modular architecture for high-performance video generation.

@@ -0,0 +1,151 @@
+# FastVideo CLI Inference
+
+The FastVideo CLI provides a quick way to access the FastVideo inference pipeline for video generation. For more advanced usage,
+see the Python interface [here](https://hao-ai-lab.github.io/FastVideo/inference/examples/basic.html).
+
+## Basic Usage
+
+The basic command to generate a video is:
+
+```bash
+fastvideo generate --model-path {MODEL_PATH} --prompt {PROMPT}
+```
+
+### Required Parameters
+
+- `--model-path {MODEL_PATH}`: Path to the model or model ID
+- `--prompt {PROMPT}`: Text description for the video you want to generate
+
+## Common Arguments
+
+To see all the options, you can use the `--help` flag:
+
+```bash
+fastvideo generate --help
+```
+
+### Hardware Configuration
+
+- `--num-gpus {NUM_GPUS}`: Number of GPUs to use
+- `--tp-size {TP_SIZE}`: Tensor parallelism size (Typically should match the number of GPUs)
+- `--sp-size {SP_SIZE}`: Sequence parallelism size (Typically should match the number of GPUs)
+
+#### Video Configuration
+
+- `--height {HEIGHT}`: Height of the generated video
+- `--width {WIDTH}`: Width of the generated video
+- `--num-frames {NUM_FRAMES}`: Number of frames to generate
+- `--fps {FPS}`: Frames per second for the saved video
+
+#### Generation Parameters
+
+- `--num-inference-steps {STEPS}`: Number of denoising steps
+- `--negative-prompt {PROMPT}`: Negative prompt to guide generation away from certain concepts
+- `--seed {SEED}`: Random seed for reproducible generation
+
+#### Output Options
+
+- `--output-path {PATH}`: Directory to save the generated video
+- `--save-video`: Whether to save the video to disk
+- `--return-frames`: Whether to return the raw frames
+
+## Using Configuration Files
+
+Instead of specifying all parameters on the command line, you can use a configuration file:
+
+```bash
+fastvideo generate --config {CONFIG_FILE_PATH}
+```
+
+The config file should be in JSON or YAML format with the same parameter names as the CLI options. Command-line arguments will take precedence over settings in the configuration file, allowing you to override specific values while keeping the rest from the config file.
+
+Example configuration file (config.json):
+
+```json
+{
+    "model_path": "FastVideo/FastHunyuan-diffusers",
+    "prompt": "A beautiful woman in a red dress walking down a street",
+    "output_path": "outputs/",
+    "num_gpus": 2,
+    "sp_size": 2,
+    "tp_size": 2,
+    "num_frames": 45,
+    "height": 720,
+    "width": 1280,
+    "num_inference_steps": 6,
+    "seed": 1024,
+    "fps": 24,
+    "precision": "bf16",
+    "vae_precision": "fp16",
+    "vae_tiling": true,
+    "vae_sp": true,
+    "vae_config": {
+        "load_encoder": false,
+        "load_decoder": true,
+        "tile_sample_min_height": 256,
+        "tile_sample_min_width": 256
+    },
+    "text_encoder_precisions": [
+        "fp16",
+        "fp16"
+    ],
+    "mask_strategy_file_path": null,
+    "enable_torch_compile": false
+}
+```
+
+Or using YAML format (config.yaml):
+
+```yaml
+model_path: "FastVideo/FastHunyuan-diffusers"
+prompt: "A beautiful woman in a red dress walking down a street"
+output_path: "outputs/"
+num_gpus: 2
+sp_size: 2
+tp_size: 2
+num_frames: 45
+height: 720
+width: 1280
+num_inference_steps: 6
+seed: 1024
+fps: 24
+precision: "bf16"
+vae_precision: "fp16"
+vae_tiling: true
+vae_sp: true
+vae_config:
+  load_encoder: false
+  load_decoder: true
+  tile_sample_min_height: 256
+  tile_sample_min_width: 256
+text_encoder_precisions:
+  - "fp16"
+  - "fp16"
+mask_strategy_file_path: null
+enable_torch_compile: false
+```
+
+## Examples
+
+Generating a simple video:
+
+```bash
+fastvideo generate --model-path FastVideo/FastHunyuan-diffusers --prompt "A cat playing with a ball of yarn" --num-frames 45 --height 720 --width 1280 --num-inference-steps 6 --seed 1024 --output-path outputs/
+```
+
+Using a negative prompt to avoid certain elements:
+
+```bash
+fastvideo generate --model-path FastVideo/FastHunyuan-diffusers --prompt "A beautiful forest landscape" --negative-prompt "people, buildings, roads"
+```
+
+Combining command line arguments and a configuration file:
+
+```bash
+fastvideo generate --config config.json --prompt "A capybara lounging in a hammock"
+```
+
+## Troubleshooting
+
+- If you encounter CUDA out-of-memory errors, try reducing the video dimensions or number of frames, or the number of inference steps.
+- For reproducible results, set the same seed value between runs.
@@ -0,0 +1,77 @@
+(inference-configuration)=
+# Configuration
+
+## Multi-GPU Setup
+
+FastVideo automatically distributes the generation process when multiple GPUs are specified:
+
+```python
+# Will use 4 GPUs in parallel for faster generation
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=4,
+)
+```
+
+## Customizing Generation
+
+- `PipelineConfig`: Initialization time parameters
+- `SamplingParam`: Generation time parameters
+
+You can customize various parameters when generating videos using the `PipelineConfig` and `SamplingParam` class:
+
+```python
+from fastvideo import VideoGenerator, SamplingParam, PipelineConfig
+
+def main():
+    model_name = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    config = PipelineConfig.from_pretrained(model_name)
+    config.vae_precision = "fp16"
+    config.use_cpu_offload = True
+
+    # Create the generator
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        num_gpus=1,
+        pipeline_config=config
+    )
+
+    # Create and customize sampling parameters
+    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    
+    # How many frames to generate
+    sampling_param.num_frames = 45
+    
+    # Video resolution (width, height)
+    sampling_param.width = 1024
+    sampling_param.height = 576
+    
+    # How many steps we denoise the video (higher = better quality, slower generation)
+    sampling_param.num_inference_steps = 30
+    
+    # How strongly the video conforms to the prompt (higher = more faithful to prompt)
+    sampling_param.guidance_scale = 7.5
+    
+    # Random seed for reproducibility
+    sampling_param.seed = 42  # Optional, leave unset for random results
+
+    # Generate video with custom parameters
+    prompt = "A beautiful sunset over a calm ocean, with gentle waves."
+    video = generator.generate_video(
+        prompt, 
+        sampling_param=sampling_param, 
+        output_path="my_videos/",  # Controls where videos are saved
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        save_video=True
+    )
+
+    # If return_frames=True, video contains the generated frames as a NumPy array
+    print(f"Generated {len(video)} frames")
+
+if __name__ == '__main__':
+    main()
+```
+
+## Performance Optimization
+
+For configuring optimizations, please see our [optimizations guide](#inference-optimizations)
@@ -1,33 +0,0 @@
-(fasthunyuan)=
-
-# FastHunyuan
-## 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).
@@ -1,9 +0,0 @@
-(fastmochi)=
-
-# 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
@@ -1,18 +0,0 @@
-(hunyuanvideo)=
-
-# HunyuanVideo
-## 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.
@@ -0,0 +1,124 @@
+# Inference Quick Start
+
+This page contains step-by-step instructions to get you quickly started with video generation using FastVideo.
+
+## Requirements
+- **OS**: Linux (Tested on Ubuntu 22.04+)
+- **Python**: 3.10-3.12
+- **CUDA**: 12.4
+- **GPU**: At least one NVIDIA GPU
+
+## Installation
+
+We recommend using an environment manager such as `Conda` to create a clean environment:
+
+```bash
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
+pip install fastvideo
+```
+
+For advanced installation options, see the [Installation Guide](installation.md).
+
+## Generating Your First Video
+Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
+
+```python
+from fastvideo import VideoGenerator
+
+def main():
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
+
+    # Generate the video
+    video = generator.generate_video(
+        prompt,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Run the script with:
+
+```bash
+python example.py
+```
+
+The generated video will be saved in the current directory under `my_videos/`.
+
+## Available Models
+
+Please see the [support matrix](#support-matrix) for the list of supported models and their available optimizations.
+
+## Image-to-Video Generation
+
+You can generate a video starting from an initial image:
+
+```python
+from fastvideo import VideoGenerator, SamplingParam
+
+def main():
+    # Create the generator
+    model_name = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+    generator = VideoGenerator.from_pretrained(model_name, num_gpus=1)
+
+    # Set up parameters with an initial image
+    sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    sampling_param.num_frames = 107
+    sampling_param.image_strength = 0.8  # How much to preserve the original image (0-1)
+
+    # Generate video based on the image
+    prompt = "A photograph coming to life with gentle movement"
+    generator.generate_video(prompt, sampling_param=sampling_param,
+                             output_path="my_videos/",
+                             save_video=True)
+
+if __name__ == '__main__':
+    main()
+```
+
+## Troubleshooting
+
+Common issues and their solutions:
+
+### Out of Memory Errors
+If you encounter CUDA out of memory errors:
+- Reduce `num_frames` or video resolution
+- Enable memory optimization with `enable_model_cpu_offload`
+- Try a smaller model or use distilled versions
+- Use `num_gpus` > 1 if multiple GPUs are available
+
+### Slow Generation
+To speed up generation:
+- Reduce `num_inference_steps` (20-30 is usually sufficient)
+- Use half precision (`fp16`) for the VAE
+- Use multiple GPUs if available
+
+### Unexpected Results
+If the generated video doesn't match your prompt:
+- Try increasing `guidance_scale` (7.0-9.0 works well)
+- Make your prompt more detailed and specific
+- Experiment with different random seeds
+- Try a different model
+
+## Next Steps
+
+- Learn about [Advanced Inference Configurations](#inference-configuration)
+- Learn about using [Optimizations](#inference-optimizations)
+- See [Examples](../examples/examples_inference_index.md) for more usage scenarios
+- Join our [Community Discord](https://discord.gg/JA7cksDz86).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).
@@ -0,0 +1,148 @@
+(inference-optimizations)=
+# Optimizations
+
+This page describes the various options for speeding up generation times in FastVideo.
+
+## Table of Contents
+- Optimized Attention Backends
+  - [Flash Attention](#optimizations-flash)
+  - [Sliding Tile Attention](#optimizations-sta)
+  - [Sage Attention](#optimizations-sage)
+
+- Caching Techniques
+  - [TeaCache](#optimizations-teacache)
+
+(optimizations-backends)=
+## Attention Backends
+
+### Available Backends
+- Torch SDPA: `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`
+- Flash Attention 2 and 3: `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN`
+- Sliding Tile Attention: `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
+- Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
+
+### Configuring Backends
+
+There are two ways to configure the attention backend in FastVideo.
+
+#### 1. In Python
+In python, set the `FASTVIDEO_ATTENTION_BACKEND` environment variable before instantiating `VideoGenerator` like this:
+
+```python
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
+```
+
+#### 2. In CLI
+You can also set the environment variable on the command line:
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
+```
+
+(optimizations-flash)=
+### Flash Attention
+
+**`FLASH_ATTN`**
+
+We recommend always installing [Flash Attention 2](https://github.com/Dao-AILab/flash-attention):
+
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+And if using a Hopper+ GPU (ie H100), installing [Flash Attention 3](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#flashattention-3-beta-release) by compiling it from source (takes about 10 minutes for me):
+
+```bash
+git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention
+
+cd hopper
+pip install ninja 
+python setup.py install
+```
+
+:::{note}
+FastVideo will automatically detect and use `FA3` if it is installed when using `FLASH_ATTN` backend.
+:::
+
+(optimizations-sta)=
+### Sliding Tile Attention
+**`SLIDING_TILE_ATTN`**
+
+```bash
+pip install st_attn==0.0.4
+```
+
+Please see [this page](#sta-installation) for more installation instructions.
+
+(optimizations-sage)=
+### Sage Attention
+**`SAGE_ATTN`**
+
+To use [SageAttention](https://github.com/thu-ml/SageAttention) 2.1.1, please compile from source:
+
+```bash
+git clone https://github.com/thu-ml/SageAttention.git
+cd sageattention 
+python setup.py install  # or pip install -e .
+```
+
+(optimizations-teacache)=
+## Teacache
+TeaCache is an optimization technique supported in FastVideo that can significantly speed up video generation by skipping redundant calculations across diffusion steps. This guide explains how to enable and configure TeaCache for optimal performance in FastVideo.
+
+### What is TeaCache?
+
+See the official [TeaCache](https://github.com/ali-vilab/TeaCache) repo and their paper for more details.
+
+### How to Enable TeaCache
+
+Enabling TeaCache is straightforward - simply add the `enable_teacache=True` parameter to your `generate_video()` call:
+
+```python
+# ... previous code
+generator.generate_video(
+    prompt="Your prompt here",
+    sampling_param=params,
+    enable_teacache=True
+)
+# more code ...
+```
+
+### Complete Example
+
+At the bottom is a complete example of using TeaCache for faster video generation. You can run it using the following command:
+
+```bash
+python examples/inference/optimizations/teacache_example.py
+```
+
+### Advanced Configuration
+
+While TeaCache works well with default settings, you can fine-tune its behavior by adjusting the threshold value:
+
+1. Lower threshold values (e.g., 0.1) will result in more skipped calculations and faster generation with slightly more potential for quality degradation
+2. Higher threshold values (e.g., 0.15-0.23) will skip fewer calculations but maintain quality closer to the original
+
+Note that the optimal threshold depends on your specific model and content.
+
+## Benchmarking different optimizations
+
+To benchmark the performance improvement, try generating the same video with and without TeaCache enabled and compare the generation times:
+
+```python
+# Without TeaCache
+start_time = time.perf_counter()
+generator.generate_video(prompt="Your prompt", enable_teacache=False)
+standard_time = time.perf_counter() - start_time
+
+# With TeaCache
+start_time = time.perf_counter()
+generator.generate_video(prompt="Your prompt", enable_teacache=True)
+teacache_time = time.perf_counter() - start_time
+
+print(f"Standard generation: {standard_time:.2f} seconds")
+print(f"TeaCache generation: {teacache_time:.2f} seconds")
+print(f"Speedup: {standard_time/teacache_time:.2f}x")
+```
+
+Note: If you want to benchmark different attention backends, you'll need to reinstantiate `VideoGenerator`.
@@ -1,16 +0,0 @@
-(stepvideo)=
-
-# StepVideo
-## 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
-```
@@ -0,0 +1,92 @@
+(support-matrix)=
+# Compatibility Matrix
+The table below shows every supported model and optimizations supported for them.
+
+The symbols used have the following meanings:
+
+- ✅ = Full compatibility
+- ❌ = No compatibility
+
+## Models x Optimization
+The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods and FastVideo will use the optimal default parameters when initializing and generating videos.
+
+:::{raw} html
+<style>
+  /* Make smaller to try to improve readability  */
+  td {
+    font-size: 0.9rem;
+    text-align: center;
+  }
+
+  th {
+    text-align: center;
+    font-size: 0.9rem;
+  }
+</style>
+:::
+
+:::{list-table}
+:header-rows: 1
+:stub-columns: 3
+:widths: auto
+:class: vertical-table-header
+
+- * Model Name
+  * HuggingFace Model ID
+  * Resolutions
+  * TeaCache
+  * Sliding Tile Attn
+  * Sage Attn
+- * HunyuanVideo
+  * `hunyuanvideo-community/HunyuanVideo`
+  * 720px1280p<br>544px960p
+  * ❌
+  * ✅
+  * ✅
+- * FastHunyuan
+  * `FastVideo/FastHunyuan-diffusers`
+  * 720px1280p<br>544px960p
+  * ❌
+  * ✅
+  * ✅
+- * Wan T2V 1.3B
+  * `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`
+  * 480P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan T2V 14B
+  * `Wan-AI/Wan2.1-T2V-14B-Diffusers`
+  * 480P, 720P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan I2V 480P
+  * `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers`
+  * 480P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan I2V 720P
+  * `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers`
+  * 720P
+  * ✅
+  * ✅*
+  * ✅
+- * StepVideo T2V
+  * `FastVideo/stepvideo-t2v-diffusers`
+  * 768px768px204f<br>544px992px204f<br>544px992px136f
+  * ❌
+  * ❌
+  * ✅
+:::
+
+**Note**: there are some known quality issues with Wan2.1 + Sliding Tile Attn. We are working on fixing this issue.
+
+## Special requirements
+
+### StepVideo T2V
+- The self-attention in text-encoder (step_llm) only supports CUDA capabilities sm_80 sm_86 and sm_90
+
+### Sliding Tile Attention
+- Currently only Hopper GPUs (H100s) are supported.
@@ -1,44 +0,0 @@
-(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.
@@ -15,7 +15,7 @@ The first script in this example shows the most basic usage of FastVideo. If you
 # if you have not cloned the directory:
 git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo

-python fastvideo/v1/examples/inference/basic/basic.py
+python examples/inference/basic/basic.py
 ```

 ## Basic Walkthrough
@@ -13,26 +13,12 @@ The demo uses the FastVideo framework to generate videos based on text prompts.

 ---

-## Requirements
-
- Linux-based OS
- Python 3.10-3.12
- NVIDIA GPU
- Cuda 12.4
- FastVideo
-
-## Installation
-
-```bash
-pip install fastvideo
-```
-
 ## Usage

 Run the demo with:

 ```bash
-python fastvideo/v1/examples/inference/gradio/gradio_demo.py
+python examples/inference/gradio/gradio_demo.py
 ```

 This will start a web server at `http://0.0.0.0:7860` where you can access the interface.
@@ -1,104 +1,9 @@
-# Speeding Up Generation
+# Optimization Examples

-This page describes the various options for speeding up generation times.
-
-## Table of Contents
- Optimized Attention Backends
-    - [Flash Attention](#optimizations-flash)
-    - [Sliding Tile Attention](#optimizations-sta)
-    - [Sage Attention](#optimizations-sage)
-
- Caching Techniques
-    - [TeaCache](#optimizations-teacache)
-
-
-(optimizations-backends)=
-## Attention Backends
-
-`attention_example.py` shows how to set `FASTVIDEO_ATTENTION_BACKEND` env var to change attention backends. To run this example:
 ```bash
 python examples/inference/optimizations/attention_example.py
 ```

-In python, set the `FASTVIDEO_ATTENTION_BACKEND` before instantiating `VideoGenerator` like this:
-
-```python
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
-```
-
-You can also set the env var when running any of the other example like this:
-```bash
-FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
-```
-
-(optimizations-flash)=
-### Flash Attention
-
-(optimizations-sta)=
-### Sliding Tile Attention
-
-(optimizations-sage)=
-### Sage Attention
-
-
-(optimizations-teacache)=
-## Teacache 
-TeaCache is an optimization technique supported in FastVideo that can significantly speed up video generation by skipping redundant calculations across diffusion steps. This guide explains how to enable and configure TeaCache for optimal performance in FastVideo.
-
-### What is TeaCache?
-
-See the official [TeaCache](https://github.com/ali-vilab/TeaCache) repo and their paper for more details.
-
-
-### How to Enable TeaCache
-
-Enabling TeaCache is straightforward - simply add the `enable_teacache=True` parameter to your `generate_video()` call:
-
-```python
-# ... previous code
-generator.generate_video(
-    prompt="Your prompt here",
-    sampling_param=params,
-    enable_teacache=True
-)
-# more code ...
-```
-
-### Complete Example
-
-At the bottom is a complete example of using TeaCache for faster video generation. You can run it using the following command:
-
 ```bash
 python examples/inference/optimizations/teacache_example.py
 ```
-
-### Advanced Configuration
-
-While TeaCache works well with default settings, you can fine-tune its behavior by adjusting the threshold value:
-
-1. Lower threshold values (e.g., 0.1) will result in more skipped calculations and faster generation with slightly more potential for quality degradation
-2. Higher threshold values (e.g., 0.15-0.23) will skip fewer calculations but maintain quality closer to the original 
-
-Note that the optimal threshold depends on your specific model and content.
-
-## Benchmarking different optimizations
-
-To benchmark the performance improvement, try generating the same video with and without TeaCache enabled and compare the generation times:
-
-```python
-# Without TeaCache
-start_time = time.time()
-generator.generate_video(prompt="Your prompt", enable_teacache=False)
-standard_time = time.time() - start_time
-
-# With TeaCache
-start_time = time.time()
-generator.generate_video(prompt="Your prompt", enable_teacache=True)
-teacache_time = time.time() - start_time
-
-print(f"Standard generation: {standard_time:.2f} seconds")
-print(f"TeaCache generation: {teacache_time:.2f} seconds")
-print(f"Speedup: {standard_time/teacache_time:.2f}x")
-```
-
-Note: If you want to benchmark different attention backends, you'll need to reinstantiate `VideoGenerator`.
@@ -7,15 +7,15 @@ def main():
    # set the attention backend 
    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"

-    start_time = time.time()
+    start_time = time.perf_counter()
    gen = VideoGenerator.from_pretrained(
        model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
        num_gpus=1,
    )
-    load_time = time.time() - start_time
+    load_time = time.perf_counter() - start_time
    print(f"Model loading time: {load_time:.2f} seconds")

-    gen_start_time = time.time()
+    gen_start_time = time.perf_counter()

    gen.generate_video(
        prompt=
@@ -23,10 +23,10 @@ def main():
        seed=1024,
        output_path="example_outputs/")
    
-    generation_time = time.time() - gen_start_time
+    generation_time = time.perf_counter() - gen_start_time
    print(f"Video generation time: {generation_time:.2f} seconds")

-    total_time = time.time() - start_time
+    total_time = time.perf_counter() - start_time
    print(f"Total execution time: {total_time:.2f} seconds")

 if __name__ == "__main__":
@@ -4,17 +4,17 @@ from fastvideo import VideoGenerator, SamplingParam


 def main():
-    start_time = time.time()
+    start_time = time.perf_counter()

    gen = VideoGenerator.from_pretrained(
        model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
        num_gpus=1,
        use_cpu_offload=False,
    )
-    load_time = time.time() - start_time
+    load_time = time.perf_counter() - start_time
    print(f"Model loading time: {load_time:.2f} seconds")

-    gen_start_time = time.time()
+    gen_start_time = time.perf_counter()

    params = SamplingParam.from_pretrained(
        model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
@@ -33,10 +33,10 @@ def main():
        seed=1024,
        output_path="example_outputs/")
    
-    generation_time = time.time() - gen_start_time
+    generation_time = time.perf_counter() - gen_start_time
    print(f"Video generation time: {generation_time:.2f} seconds")

-    total_time = time.time() - start_time
+    total_time = time.perf_counter() - start_time
    print(f"Total execution time: {total_time:.2f} seconds")


@@ -1,5 +1,6 @@
 from fastvideo.v1.configs.pipelines import PipelineConfig
 from fastvideo.v1.configs.sample import SamplingParam
 from fastvideo.v1.entrypoints.video_generator import VideoGenerator
+from fastvideo.version import __version__

-__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam"]
+__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam", "__version__"]
@@ -452,7 +452,7 @@ def main(args):
        return phase

    for step in range(init_steps + 1, args.max_train_steps + 1):
-        start_time = time.time()
+        start_time = time.perf_counter()
        assert args.multi_phased_distill_schedule is not None
        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)

@@ -482,7 +482,7 @@ def main(args):
            args.hunyuan_teacher_disable_cfg,
        )

-        step_time = time.time() - start_time
+        step_time = time.perf_counter() - start_time
        step_times.append(step_time)
        avg_step_time = sum(step_times) / len(step_times)

@@ -517,7 +517,7 @@ def main(args):
    for step in range(init_steps + 1, args.max_train_steps + 1):
        assert args.multi_phased_distill_schedule is not None
        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
-        start_time = time.time()
+        start_time = time.perf_counter()
        (
            generator_loss,
            generator_grad_norm,
@@ -547,7 +547,7 @@ def main(args):
            args.discriminator_head_stride,
        )

-        step_time = time.time() - start_time
+        step_time = time.perf_counter() - start_time
        step_times.append(step_time)
        avg_step_time = sum(step_times) / len(step_times)

@@ -237,7 +237,7 @@ def add_inference_args(parser: argparse.ArgumentParser):
        type=str,
        default="540p",
        choices=["540p", "720p"],
-        help="Root path of all the models, including t2v models and extra models.",
+        help="The resolution of the model.",
    )
    group.add_argument(
        "--load-key",
@@ -361,7 +361,7 @@ def add_parallel_args(parser: argparse.ArgumentParser):
        "--ring-degree",
        type=int,
        default=1,
-        help="Ulysses degree.",
+        help="Ring degree.",
    )

    return parser
@@ -17,7 +17,7 @@ from fastvideo.models.hunyuan.vae import load_vae
 from fastvideo.utils.parallel_states import nccl_info


-class Inference(object):
+class Inference:

    def __init__(
        self,
@@ -452,7 +452,7 @@ class HunyuanVideoSampler(Inference):
        # ========================================================================
        # Pipeline inference
        # ========================================================================
-        start_time = time.time()
+        start_time = time.perf_counter()
        samples = self.pipeline(
            prompt=prompt,
            height=target_height,
@@ -476,7 +476,7 @@ class HunyuanVideoSampler(Inference):
        out_dict["samples"] = samples
        out_dict["prompts"] = prompt

-        gen_time = time.time() - start_time
+        gen_time = time.perf_counter() - start_time
        logger.info(f"Success, time: {gen_time}")

        return out_dict
@@ -41,7 +41,7 @@ def get_rewrite_prompt(ori_prompt, mode="Normal"):
    elif mode == "Master":
        prompt = master_mode_prompt.format(input=ori_prompt)
    else:
-        raise Exception("Only supports Normal and Normal", mode)
+        raise Exception("Only supports Normal and Master mode, but got {}".format(mode))
    return prompt


@@ -267,25 +267,25 @@ class Step1Model(PreTrainedModel):
 class STEP1TextEncoder(torch.nn.Module):

    def __init__(self, model_dir, max_length=320):
-        super(STEP1TextEncoder, self).__init__()
+        super()
        self.max_length = max_length
        self.text_tokenizer = Wrapped_StepChatTokenizer(os.path.join(model_dir, 'step1_chat_tokenizer.model'))
        text_encoder = Step1Model.from_pretrained(model_dir)
        self.text_encoder = text_encoder.eval().to(torch.bfloat16)

    @torch.no_grad
+    @torch.autocast(device_type='cuda', dtype=torch.bfloat16)
    def forward(self, prompts, with_mask=True, max_length=None):
        self.device = next(self.text_encoder.parameters()).device
-        with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
-            if type(prompts) is str:
-                prompts = [prompts]
+        if type(prompts) is str:
+            prompts = [prompts]

-            txt_tokens = self.text_tokenizer(prompts,
-                                             max_length=max_length or self.max_length,
-                                             padding="max_length",
-                                             truncation=True,
-                                             return_tensors="pt")
-            y = self.text_encoder(txt_tokens.input_ids.to(self.device),
+        txt_tokens = self.text_tokenizer(prompts,
+                                            max_length=max_length or self.max_length,
+                                            padding="max_length",
+                                            truncation=True,
+                                            return_tensors="pt")
+        y = self.text_encoder(txt_tokens.input_ids.to(self.device),
                                  attention_mask=txt_tokens.attention_mask.to(self.device) if with_mask else None)
-            y_mask = txt_tokens.attention_mask
+        y_mask = txt_tokens.attention_mask
        return y.transpose(0, 1), y_mask
@@ -364,7 +364,7 @@ def main(args):
    for i in range(init_steps):
        next(loader)
    for step in range(init_steps + 1, args.max_train_steps + 1):
-        start_time = time.time()
+        start_time = time.perf_counter()
        loss, grad_norm = train_one_step(
            transformer,
            args.model_type,
@@ -383,7 +383,7 @@ def main(args):
            args.mode_scale,
        )

-        step_time = time.time() - start_time
+        step_time = time.perf_counter() - start_time
        step_times.append(step_time)
        avg_step_time = sum(step_times) / len(step_times)

@@ -0,0 +1,776 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# type: ignore
+# ruff: noqa
+# code borrowed from https://github.com/pytorch/pytorch/blob/main/torch/utils/collect_env.py
+# and vllm: https://github.com/vllm-project/vllm/blob/main/vllm/collect_env.py
+
+import datetime
+import locale
+import os
+import re
+import subprocess
+import sys
+# Unlike the rest of the PyTorch this file must be python2 compliant.
+# This script outputs relevant system environment info
+# Run it with `python collect_env.py` or `python -m torch.utils.collect_env`
+from collections import namedtuple
+
+from fastvideo.v1.envs import environment_variables
+
+try:
+    import torch
+    TORCH_AVAILABLE = True
+except (ImportError, NameError, AttributeError, OSError):
+    TORCH_AVAILABLE = False
+
+# System Environment Information
+SystemEnv = namedtuple(
+    'SystemEnv',
+    [
+        'torch_version',
+        'is_debug_build',
+        'cuda_compiled_version',
+        'gcc_version',
+        'clang_version',
+        'cmake_version',
+        'os',
+        'libc_version',
+        'python_version',
+        'python_platform',
+        'is_cuda_available',
+        'cuda_runtime_version',
+        'cuda_module_loading',
+        'nvidia_driver_version',
+        'nvidia_gpu_models',
+        'cudnn_version',
+        'pip_version',  # 'pip' or 'pip3'
+        'pip_packages',
+        'conda_packages',
+        'hip_compiled_version',
+        'hip_runtime_version',
+        'miopen_runtime_version',
+        'caching_allocator_config',
+        'is_xnnpack_available',
+        'cpu_info',
+        'fastvideo_version',
+        'fastvideo_build_flags',
+        'gpu_topo',
+        'env_vars',
+    ])
+
+DEFAULT_CONDA_PATTERNS = {
+    "torch",
+    "numpy",
+    "cudatoolkit",
+    "soumith",
+    "mkl",
+    "magma",
+    "triton",
+    "optree",
+    "nccl",
+    "transformers",
+    "accelerate",
+    "peft",
+    "zmq",
+    "nvidia",
+    "pynvml",
+}
+
+DEFAULT_PIP_PATTERNS = {
+    "torch",
+    "numpy",
+    "mypy",
+    "flake8",
+    "triton",
+    "optree",
+    "onnx",
+    "nccl",
+    "transformers",
+    "accelerate",
+    "peft",
+    "zmq",
+    "nvidia",
+    "pynvml",
+}
+
+
+def run(command):
+    """Return (return-code, stdout, stderr)."""
+    shell = True if type(command) is str else False
+    p = subprocess.Popen(command,
+                         stdout=subprocess.PIPE,
+                         stderr=subprocess.PIPE,
+                         shell=shell)
+    raw_output, raw_err = p.communicate()
+    rc = p.returncode
+    if get_platform() == 'win32':
+        enc = 'oem'
+    else:
+        enc = locale.getpreferredencoding()
+    output = raw_output.decode(enc)
+    if command == 'nvidia-smi topo -m':
+        # don't remove the leading whitespace of `nvidia-smi topo -m`
+        #   because they are meaningful
+        output = output.rstrip()
+    else:
+        output = output.strip()
+    err = raw_err.decode(enc)
+    return rc, output, err.strip()
+
+
+def run_and_read_all(run_lambda, command):
+    """Run command using run_lambda; reads and returns entire output if rc is 0."""
+    rc, out, _ = run_lambda(command)
+    if rc != 0:
+        return None
+    return out
+
+
+def run_and_parse_first_match(run_lambda, command, regex):
+    """Run command using run_lambda, returns the first regex match if it exists."""
+    rc, out, _ = run_lambda(command)
+    if rc != 0:
+        return None
+    match = re.search(regex, out)
+    if match is None:
+        return None
+    return match.group(1)
+
+
+def run_and_return_first_line(run_lambda, command):
+    """Run command using run_lambda and returns first line if output is not empty."""
+    rc, out, _ = run_lambda(command)
+    if rc != 0:
+        return None
+    return out.split('\n')[0]
+
+
+def get_conda_packages(run_lambda, patterns=None):
+    if patterns is None:
+        patterns = DEFAULT_CONDA_PATTERNS
+    conda = os.environ.get('CONDA_EXE', 'conda')
+    out = run_and_read_all(run_lambda, "{} list".format(conda))
+    if out is None:
+        return out
+
+    return "\n".join(line for line in out.splitlines()
+                     if not line.startswith("#") and any(name in line
+                                                         for name in patterns))
+
+
+def get_gcc_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'gcc --version', r'gcc (.*)')
+
+
+def get_clang_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'clang --version',
+                                     r'clang version (.*)')
+
+
+def get_cmake_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'cmake --version',
+                                     r'cmake (.*)')
+
+
+def get_nvidia_driver_version(run_lambda):
+    if get_platform() == 'darwin':
+        cmd = 'kextstat | grep -i cuda'
+        return run_and_parse_first_match(run_lambda, cmd,
+                                         r'com[.]nvidia[.]CUDA [(](.*?)[)]')
+    smi = get_nvidia_smi()
+    return run_and_parse_first_match(run_lambda, smi, r'Driver Version: (.*?) ')
+
+
+def get_gpu_info(run_lambda):
+    if get_platform() == 'darwin' or (TORCH_AVAILABLE and hasattr(
+            torch.version, 'hip') and torch.version.hip is not None):
+        if TORCH_AVAILABLE and torch.cuda.is_available():
+            if torch.version.hip is not None:
+                prop = torch.cuda.get_device_properties(0)
+                if hasattr(prop, "gcnArchName"):
+                    gcnArch = " ({})".format(prop.gcnArchName)
+                else:
+                    gcnArch = "NoGCNArchNameOnOldPyTorch"
+            else:
+                gcnArch = ""
+            return torch.cuda.get_device_name(None) + gcnArch
+        return None
+    smi = get_nvidia_smi()
+    uuid_regex = re.compile(r' \(UUID: .+?\)')
+    rc, out, _ = run_lambda(smi + ' -L')
+    if rc != 0:
+        return None
+    # Anonymize GPUs by removing their UUID
+    return re.sub(uuid_regex, '', out)
+
+
+def get_running_cuda_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'nvcc --version',
+                                     r'release .+ V(.*)')
+
+
+def get_cudnn_version(run_lambda):
+    """Return a list of libcudnn.so; it's hard to tell which one is being used."""
+    if get_platform() == 'win32':
+        system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')
+        cuda_path = os.environ.get('CUDA_PATH', "%CUDA_PATH%")
+        where_cmd = os.path.join(system_root, 'System32', 'where')
+        cudnn_cmd = '{} /R "{}\\bin" cudnn*.dll'.format(where_cmd, cuda_path)
+    elif get_platform() == 'darwin':
+        # CUDA libraries and drivers can be found in /usr/local/cuda/. See
+        # https://docs.nvidia.com/cuda/cuda-installation-guide-mac-os-x/index.html#install
+        # https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#installmac
+        # Use CUDNN_LIBRARY when cudnn library is installed elsewhere.
+        cudnn_cmd = 'ls /usr/local/cuda/lib/libcudnn*'
+    else:
+        cudnn_cmd = 'ldconfig -p | grep libcudnn | rev | cut -d" " -f1 | rev'
+    rc, out, _ = run_lambda(cudnn_cmd)
+    # find will return 1 if there are permission errors or if not found
+    if len(out) == 0 or (rc != 1 and rc != 0):
+        l = os.environ.get('CUDNN_LIBRARY')
+        if l is not None and os.path.isfile(l):
+            return os.path.realpath(l)
+        return None
+    files_set = set()
+    for fn in out.split('\n'):
+        fn = os.path.realpath(fn)  # eliminate symbolic links
+        if os.path.isfile(fn):
+            files_set.add(fn)
+    if not files_set:
+        return None
+    # Alphabetize the result because the order is non-deterministic otherwise
+    files = sorted(files_set)
+    if len(files) == 1:
+        return files[0]
+    result = '\n'.join(files)
+    return 'Probably one of the following:\n{}'.format(result)
+
+
+def get_nvidia_smi():
+    # Note: nvidia-smi is currently available only on Windows and Linux
+    smi = 'nvidia-smi'
+    if get_platform() == 'win32':
+        system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')
+        program_files_root = os.environ.get('PROGRAMFILES', 'C:\\Program Files')
+        legacy_path = os.path.join(program_files_root, 'NVIDIA Corporation',
+                                   'NVSMI', smi)
+        new_path = os.path.join(system_root, 'System32', smi)
+        smis = [new_path, legacy_path]
+        for candidate_smi in smis:
+            if os.path.exists(candidate_smi):
+                smi = '"{}"'.format(candidate_smi)
+                break
+    return smi
+
+
+def get_fastvideo_version():
+    return ""
+    from fastvideo import __version__, __version_tuple__
+
+    if __version__ == "dev":
+        return "N/A (dev)"
+    version_str = __version_tuple__[-1]
+    if isinstance(version_str, str) and version_str.startswith('g'):
+        # it's a dev build
+        if '.' in version_str:
+            # it's a dev build containing local changes
+            git_sha = version_str.split('.')[0][1:]
+            date = version_str.split('.')[-1][1:]
+            return f"{__version__} (git sha: {git_sha}, date: {date})"
+        else:
+            # it's a dev build without local changes
+            git_sha = version_str[1:]  # type: ignore
+            return f"{__version__} (git sha: {git_sha})"
+    return __version__
+
+
+def summarize_fastvideo_build_flags():
+    # This could be a static method if the flags are constant, or dynamic if you need to check environment variables, etc.
+    return 'CUDA Archs: {}; ROCm: {}; Neuron: {}'.format(
+        os.environ.get('TORCH_CUDA_ARCH_LIST', 'Not Set'),
+        'Enabled' if os.environ.get('ROCM_HOME') else 'Disabled',
+        'Enabled' if os.environ.get('NEURON_CORES') else 'Disabled',
+    )
+
+
+def get_gpu_topo(run_lambda):
+    output = None
+
+    if get_platform() == 'linux':
+        output = run_and_read_all(run_lambda, 'nvidia-smi topo -m')
+        if output is None:
+            output = run_and_read_all(run_lambda, 'rocm-smi --showtopo')
+
+    return output
+
+
+# example outputs of CPU infos
+#  * linux
+#    Architecture:            x86_64
+#      CPU op-mode(s):        32-bit, 64-bit
+#      Address sizes:         46 bits physical, 48 bits virtual
+#      Byte Order:            Little Endian
+#    CPU(s):                  128
+#      On-line CPU(s) list:   0-127
+#    Vendor ID:               GenuineIntel
+#      Model name:            Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz
+#        CPU family:          6
+#        Model:               106
+#        Thread(s) per core:  2
+#        Core(s) per socket:  32
+#        Socket(s):           2
+#        Stepping:            6
+#        BogoMIPS:            5799.78
+#        Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr
+#                             sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl
+#                             xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16
+#                             pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand
+#                             hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced
+#                             fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap
+#                             avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1
+#                             xsaves wbnoinvd ida arat avx512vbmi pku ospke avx512_vbmi2 gfni vaes vpclmulqdq
+#                             avx512_vnni avx512_bitalg tme avx512_vpopcntdq rdpid md_clear flush_l1d arch_capabilities
+#    Virtualization features:
+#      Hypervisor vendor:     KVM
+#      Virtualization type:   full
+#    Caches (sum of all):
+#      L1d:                   3 MiB (64 instances)
+#      L1i:                   2 MiB (64 instances)
+#      L2:                    80 MiB (64 instances)
+#      L3:                    108 MiB (2 instances)
+#    NUMA:
+#      NUMA node(s):          2
+#      NUMA node0 CPU(s):     0-31,64-95
+#      NUMA node1 CPU(s):     32-63,96-127
+#    Vulnerabilities:
+#      Itlb multihit:         Not affected
+#      L1tf:                  Not affected
+#      Mds:                   Not affected
+#      Meltdown:              Not affected
+#      Mmio stale data:       Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
+#      Retbleed:              Not affected
+#      Spec store bypass:     Mitigation; Speculative Store Bypass disabled via prctl and seccomp
+#      Spectre v1:            Mitigation; usercopy/swapgs barriers and __user pointer sanitization
+#      Spectre v2:            Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
+#      Srbds:                 Not affected
+#      Tsx async abort:       Not affected
+#  * win32
+#    Architecture=9
+#    CurrentClockSpeed=2900
+#    DeviceID=CPU0
+#    Family=179
+#    L2CacheSize=40960
+#    L2CacheSpeed=
+#    Manufacturer=GenuineIntel
+#    MaxClockSpeed=2900
+#    Name=Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz
+#    ProcessorType=3
+#    Revision=27142
+#
+#    Architecture=9
+#    CurrentClockSpeed=2900
+#    DeviceID=CPU1
+#    Family=179
+#    L2CacheSize=40960
+#    L2CacheSpeed=
+#    Manufacturer=GenuineIntel
+#    MaxClockSpeed=2900
+#    Name=Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz
+#    ProcessorType=3
+#    Revision=27142
+
+
+def get_cpu_info(run_lambda):
+    rc, out, err = 0, '', ''
+    if get_platform() == 'linux':
+        rc, out, err = run_lambda('lscpu')
+    elif get_platform() == 'win32':
+        rc, out, err = run_lambda(
+            'wmic cpu get Name,Manufacturer,Family,Architecture,ProcessorType,DeviceID, \
+        CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision /VALUE'
+        )
+    elif get_platform() == 'darwin':
+        rc, out, err = run_lambda("sysctl -n machdep.cpu.brand_string")
+    cpu_info = 'None'
+    if rc == 0:
+        cpu_info = out
+    else:
+        cpu_info = err
+    return cpu_info
+
+
+def get_platform():
+    if sys.platform.startswith('linux'):
+        return 'linux'
+    elif sys.platform.startswith('win32'):
+        return 'win32'
+    elif sys.platform.startswith('cygwin'):
+        return 'cygwin'
+    elif sys.platform.startswith('darwin'):
+        return 'darwin'
+    else:
+        return sys.platform
+
+
+def get_mac_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'sw_vers -productVersion',
+                                     r'(.*)')
+
+
+def get_windows_version(run_lambda):
+    system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')
+    wmic_cmd = os.path.join(system_root, 'System32', 'Wbem', 'wmic')
+    findstr_cmd = os.path.join(system_root, 'System32', 'findstr')
+    return run_and_read_all(
+        run_lambda,
+        '{} os get Caption | {} /v Caption'.format(wmic_cmd, findstr_cmd))
+
+
+def get_lsb_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'lsb_release -a',
+                                     r'Description:\t(.*)')
+
+
+def check_release_file(run_lambda):
+    return run_and_parse_first_match(run_lambda, 'cat /etc/*-release',
+                                     r'PRETTY_NAME="(.*)"')
+
+
+def get_os(run_lambda):
+    from platform import machine
+    platform = get_platform()
+
+    if platform == 'win32' or platform == 'cygwin':
+        return get_windows_version(run_lambda)
+
+    if platform == 'darwin':
+        version = get_mac_version(run_lambda)
+        if version is None:
+            return None
+        return 'macOS {} ({})'.format(version, machine())
+
+    if platform == 'linux':
+        # Ubuntu/Debian based
+        desc = get_lsb_version(run_lambda)
+        if desc is not None:
+            return '{} ({})'.format(desc, machine())
+
+        # Try reading /etc/*-release
+        desc = check_release_file(run_lambda)
+        if desc is not None:
+            return '{} ({})'.format(desc, machine())
+
+        return '{} ({})'.format(platform, machine())
+
+    # Unknown platform
+    return platform
+
+
+def get_python_platform():
+    import platform
+    return platform.platform()
+
+
+def get_libc_version():
+    import platform
+    if get_platform() != 'linux':
+        return 'N/A'
+    return '-'.join(platform.libc_ver())
+
+
+def get_pip_packages(run_lambda, patterns=None):
+    """Return `pip list` output. Note: will also find conda-installed pytorch and numpy packages."""
+    if patterns is None:
+        patterns = DEFAULT_PIP_PATTERNS
+
+    def run_with_pip():
+        try:
+            import importlib.util
+            pip_spec = importlib.util.find_spec('pip')
+            pip_available = pip_spec is not None
+        except ImportError:
+            pip_available = False
+
+        if pip_available:
+            cmd = [sys.executable, '-mpip', 'list', '--format=freeze']
+        elif os.environ.get("UV") is not None:
+            print("uv is set")
+            cmd = ["uv", "pip", "list", "--format=freeze"]
+        else:
+            raise RuntimeError(
+                "Could not collect pip list output (pip or uv module not available)"
+            )
+
+        out = run_and_read_all(run_lambda, cmd)
+        return "\n".join(line for line in out.splitlines()
+                         if any(name in line for name in patterns))
+
+    pip_version = 'pip3' if sys.version[0] == '3' else 'pip'
+    out = run_with_pip()
+    return pip_version, out
+
+
+def get_cachingallocator_config():
+    ca_config = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', '')
+    return ca_config
+
+
+def get_cuda_module_loading_config():
+    if TORCH_AVAILABLE and torch.cuda.is_available():
+        torch.cuda.init()
+        config = os.environ.get('CUDA_MODULE_LOADING', '')
+        return config
+    else:
+        return "N/A"
+
+
+def is_xnnpack_available():
+    if TORCH_AVAILABLE:
+        import torch.backends.xnnpack
+        return str(torch.backends.xnnpack.enabled)  # type: ignore[attr-defined]
+    else:
+        return "N/A"
+
+
+def get_env_vars():
+    env_vars = ''
+    secret_terms = ('secret', 'token', 'api', 'access', 'password')
+    report_prefix = ("TORCH", "NCCL", "PYTORCH", "CUDA", "CUBLAS", "CUDNN",
+                     "OMP_", "MKL_", "NVIDIA")
+    for k, v in os.environ.items():
+        if any(term in k.lower() for term in secret_terms):
+            continue
+        if k in environment_variables:
+            env_vars = env_vars + "{}={}".format(k, v) + "\n"
+        if k.startswith(report_prefix):
+            env_vars = env_vars + "{}={}".format(k, v) + "\n"
+
+    return env_vars
+
+
+def get_env_info():
+    run_lambda = run
+    pip_version, pip_list_output = get_pip_packages(run_lambda)
+
+    if TORCH_AVAILABLE:
+        version_str = torch.__version__
+        debug_mode_str = str(torch.version.debug)
+        cuda_available_str = str(torch.cuda.is_available())
+        cuda_version_str = torch.version.cuda
+        if not hasattr(torch.version,
+                       'hip') or torch.version.hip is None:  # cuda version
+            hip_compiled_version = hip_runtime_version = miopen_runtime_version = 'N/A'
+        else:  # HIP version
+
+            def get_version_or_na(cfg, prefix):
+                _lst = [s.rsplit(None, 1)[-1] for s in cfg if prefix in s]
+                return _lst[0] if _lst else 'N/A'
+
+            cfg = torch._C._show_config().split('\n')
+            hip_runtime_version = get_version_or_na(cfg, 'HIP Runtime')
+            miopen_runtime_version = get_version_or_na(cfg, 'MIOpen')
+            cuda_version_str = 'N/A'
+            hip_compiled_version = torch.version.hip
+    else:
+        version_str = debug_mode_str = cuda_available_str = cuda_version_str = 'N/A'
+        hip_compiled_version = hip_runtime_version = miopen_runtime_version = 'N/A'
+
+    sys_version = sys.version.replace("\n", " ")
+
+    conda_packages = get_conda_packages(run_lambda)
+
+    fastvideo_version = get_fastvideo_version()
+    fastvideo_build_flags = summarize_fastvideo_build_flags()
+    gpu_topo = get_gpu_topo(run_lambda)
+
+    return SystemEnv(
+        torch_version=version_str,
+        is_debug_build=debug_mode_str,
+        python_version='{} ({}-bit runtime)'.format(
+            sys_version,
+            sys.maxsize.bit_length() + 1),
+        python_platform=get_python_platform(),
+        is_cuda_available=cuda_available_str,
+        cuda_compiled_version=cuda_version_str,
+        cuda_runtime_version=get_running_cuda_version(run_lambda),
+        cuda_module_loading=get_cuda_module_loading_config(),
+        nvidia_gpu_models=get_gpu_info(run_lambda),
+        nvidia_driver_version=get_nvidia_driver_version(run_lambda),
+        cudnn_version=get_cudnn_version(run_lambda),
+        hip_compiled_version=hip_compiled_version,
+        hip_runtime_version=hip_runtime_version,
+        miopen_runtime_version=miopen_runtime_version,
+        pip_version=pip_version,
+        pip_packages=pip_list_output,
+        conda_packages=conda_packages,
+        os=get_os(run_lambda),
+        libc_version=get_libc_version(),
+        gcc_version=get_gcc_version(run_lambda),
+        clang_version=get_clang_version(run_lambda),
+        cmake_version=get_cmake_version(run_lambda),
+        caching_allocator_config=get_cachingallocator_config(),
+        is_xnnpack_available=is_xnnpack_available(),
+        cpu_info=get_cpu_info(run_lambda),
+        fastvideo_version=fastvideo_version,
+        fastvideo_build_flags=fastvideo_build_flags,
+        gpu_topo=gpu_topo,
+        env_vars=get_env_vars(),
+    )
+
+
+env_info_fmt = """
+PyTorch version: {torch_version}
+Is debug build: {is_debug_build}
+CUDA used to build PyTorch: {cuda_compiled_version}
+ROCM used to build PyTorch: {hip_compiled_version}
+
+OS: {os}
+GCC version: {gcc_version}
+Clang version: {clang_version}
+CMake version: {cmake_version}
+Libc version: {libc_version}
+
+Python version: {python_version}
+Python platform: {python_platform}
+Is CUDA available: {is_cuda_available}
+CUDA runtime version: {cuda_runtime_version}
+CUDA_MODULE_LOADING set to: {cuda_module_loading}
+GPU models and configuration: {nvidia_gpu_models}
+Nvidia driver version: {nvidia_driver_version}
+cuDNN version: {cudnn_version}
+HIP runtime version: {hip_runtime_version}
+MIOpen runtime version: {miopen_runtime_version}
+Is XNNPACK available: {is_xnnpack_available}
+
+CPU:
+{cpu_info}
+
+Versions of relevant libraries:
+{pip_packages}
+{conda_packages}
+""".strip()
+
+# both the above code and the following code use `strip()` to
+# remove leading/trailing whitespaces, so we need to add a newline
+# in between to separate the two sections
+env_info_fmt += "\n"
+
+env_info_fmt += """
+FastVideo Version: {fastvideo_version}
+FastVideo Build Flags:
+{fastvideo_build_flags}
+GPU Topology:
+{gpu_topo}
+
+{env_vars}
+""".strip()
+
+
+def pretty_str(envinfo):
+
+    def replace_nones(dct, replacement='Could not collect'):
+        for key in dct.keys():
+            if dct[key] is not None:
+                continue
+            dct[key] = replacement
+        return dct
+
+    def replace_bools(dct, true='Yes', false='No'):
+        for key in dct.keys():
+            if dct[key] is True:
+                dct[key] = true
+            elif dct[key] is False:
+                dct[key] = false
+        return dct
+
+    def prepend(text, tag='[prepend]'):
+        lines = text.split('\n')
+        updated_lines = [tag + line for line in lines]
+        return '\n'.join(updated_lines)
+
+    def replace_if_empty(text, replacement='No relevant packages'):
+        if text is not None and len(text) == 0:
+            return replacement
+        return text
+
+    def maybe_start_on_next_line(string):
+        # If `string` is multiline, prepend a \n to it.
+        if string is not None and len(string.split('\n')) > 1:
+            return '\n{}\n'.format(string)
+        return string
+
+    mutable_dict = envinfo._asdict()
+
+    # If nvidia_gpu_models is multiline, start on the next line
+    mutable_dict['nvidia_gpu_models'] = \
+        maybe_start_on_next_line(envinfo.nvidia_gpu_models)
+
+    # If the machine doesn't have CUDA, report some fields as 'No CUDA'
+    dynamic_cuda_fields = [
+        'cuda_runtime_version',
+        'nvidia_gpu_models',
+        'nvidia_driver_version',
+    ]
+    all_cuda_fields = dynamic_cuda_fields + ['cudnn_version']
+    all_dynamic_cuda_fields_missing = all(mutable_dict[field] is None
+                                          for field in dynamic_cuda_fields)
+    if TORCH_AVAILABLE and not torch.cuda.is_available(
+    ) and all_dynamic_cuda_fields_missing:
+        for field in all_cuda_fields:
+            mutable_dict[field] = 'No CUDA'
+        if envinfo.cuda_compiled_version is None:
+            mutable_dict['cuda_compiled_version'] = 'None'
+
+    # Replace True with Yes, False with No
+    mutable_dict = replace_bools(mutable_dict)
+
+    # Replace all None objects with 'Could not collect'
+    mutable_dict = replace_nones(mutable_dict)
+
+    # If either of these are '', replace with 'No relevant packages'
+    mutable_dict['pip_packages'] = replace_if_empty(
+        mutable_dict['pip_packages'])
+    mutable_dict['conda_packages'] = replace_if_empty(
+        mutable_dict['conda_packages'])
+
+    # Tag conda and pip packages with a prefix
+    # If they were previously None, they'll show up as ie '[conda] Could not collect'
+    if mutable_dict['pip_packages']:
+        mutable_dict['pip_packages'] = prepend(
+            mutable_dict['pip_packages'], '[{}] '.format(envinfo.pip_version))
+    if mutable_dict['conda_packages']:
+        mutable_dict['conda_packages'] = prepend(mutable_dict['conda_packages'],
+                                                 '[conda] ')
+    mutable_dict['cpu_info'] = envinfo.cpu_info
+    return env_info_fmt.format(**mutable_dict)
+
+
+def get_pretty_env_info():
+    return pretty_str(get_env_info())
+
+
+def main():
+    print("Collecting environment information...")
+    output = get_pretty_env_info()
+    print(output)
+
+    if TORCH_AVAILABLE and hasattr(torch, 'utils') and hasattr(
+            torch.utils, '_crash_handler'):
+        minidump_dir = torch.utils._crash_handler.DEFAULT_MINIDUMP_DIR
+        if sys.platform == "linux" and os.path.exists(minidump_dir):
+            dumps = [
+                os.path.join(minidump_dir, dump)
+                for dump in os.listdir(minidump_dir)
+            ]
+            latest = max(dumps, key=os.path.getctime)
+            ctime = os.path.getctime(latest)
+            creation_time = datetime.datetime.fromtimestamp(ctime).strftime(
+                '%Y-%m-%d %H:%M:%S')
+            msg = "\n*** Detected a minidump at {} created on {}, ".format(latest, creation_time) + \
+                  "if this is related to your bug please include it when you file a report ***"
+            print(msg, file=sys.stderr)
+
+
+if __name__ == '__main__':
+    main()
@@ -1,38 +0,0 @@
-import platform
-
-import accelerate
-import peft
-import torch
-import transformers
-from transformers.utils import is_torch_cuda_available, is_torch_npu_available
-
-VERSION = "1.2.0"
-
-if __name__ == "__main__":
-    info = {
-        "FastVideo version": VERSION,
-        "Platform": platform.platform(),
-        "Python version": platform.python_version(),
-        "PyTorch version": torch.__version__,
-        "Transformers version": transformers.__version__,
-        "Accelerate version": accelerate.__version__,
-        "PEFT version": peft.__version__,
-    }
-
-    if is_torch_cuda_available():
-        info["PyTorch version"] += " (GPU)"
-        info["GPU type"] = torch.cuda.get_device_name()
-
-    if is_torch_npu_available():
-        info["PyTorch version"] += " (NPU)"
-        info["NPU type"] = torch.npu.get_device_name()
-        info["CANN version"] = torch.version.cann  # codespell:ignore
-
-    try:
-        import bitsandbytes
-
-        info["Bitsandbytes version"] = bitsandbytes.__version__
-    except Exception:
-        pass
-
-    print("\n" + "\n".join([f"- {key}: {value}" for key, value in info.items()]) + "\n")
@@ -3,8 +3,7 @@

 from abc import ABC, abstractmethod
 from dataclasses import dataclass, fields
-from typing import (TYPE_CHECKING, Any, Dict, Generic, Optional, Protocol, Set,
-                    Type, TypeVar)
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar

 if TYPE_CHECKING:
    from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -27,12 +26,12 @@ class AttentionBackend(ABC):

    @staticmethod
    @abstractmethod
-    def get_impl_cls() -> Type["AttentionImpl"]:
+    def get_impl_cls() -> type["AttentionImpl"]:
        raise NotImplementedError

    @staticmethod
    @abstractmethod
-    def get_metadata_cls() -> Type["AttentionMetadata"]:
+    def get_metadata_cls() -> type["AttentionMetadata"]:
        raise NotImplementedError

    # @staticmethod
@@ -46,7 +45,7 @@ class AttentionBackend(ABC):

    @staticmethod
    @abstractmethod
-    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
        raise NotImplementedError


@@ -57,8 +56,7 @@ class AttentionMetadata:
    current_timestep: int

    def asdict_zerocopy(self,
-                        skip_fields: Optional[Set[str]] = None
-                        ) -> Dict[str, Any]:
+                        skip_fields: set[str] | None = None) -> dict[str, Any]:
        """Similar to dataclasses.asdict, but avoids deepcopying."""
        if skip_fields is None:
            skip_fields = set()
@@ -124,7 +122,7 @@ class AttentionImpl(ABC, Generic[T]):
        head_size: int,
        softmax_scale: float,
        causal: bool = False,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0

-from typing import List, Optional, Type
-
 import torch
 from flash_attn import flash_attn_func as flash_attn_2_func

@@ -28,7 +26,7 @@ class FlashAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True

    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [32, 64, 96, 128, 160, 192, 224, 256]

    @staticmethod
@@ -36,15 +34,15 @@ class FlashAttentionBackend(AttentionBackend):
        return "FLASH_ATTN"

    @staticmethod
-    def get_impl_cls() -> Type["FlashAttentionImpl"]:
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
        return FlashAttentionImpl

    @staticmethod
-    def get_metadata_cls() -> Type["AttentionMetadata"]:
+    def get_metadata_cls() -> type["AttentionMetadata"]:
        raise NotImplementedError

    @staticmethod
-    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
        raise NotImplementedError


@@ -56,7 +54,7 @@ class FlashAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,5 +1,3 @@
-from typing import List, Optional, Type
-
 import torch
 from sageattention import sageattn

@@ -17,7 +15,7 @@ class SageAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True

    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [32, 64, 96, 128, 160, 192, 224, 256]

    @staticmethod
@@ -25,7 +23,7 @@ class SageAttentionBackend(AttentionBackend):
        return "SAGE_ATTN"

    @staticmethod
-    def get_impl_cls() -> Type["SageAttentionImpl"]:
+    def get_impl_cls() -> type["SageAttentionImpl"]:
        return SageAttentionImpl

    # @staticmethod
@@ -41,7 +39,7 @@ class SageAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,5 +1,3 @@
-from typing import List, Optional, Type
-
 import torch

 from fastvideo.v1.attention.backends.abstract import (
@@ -16,7 +14,7 @@ class SDPABackend(AttentionBackend):
    accept_output_buffer: bool = True

    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [32, 64, 96, 128, 160, 192, 224, 256]

    @staticmethod
@@ -24,7 +22,7 @@ class SDPABackend(AttentionBackend):
        return "SDPA"

    @staticmethod
-    def get_impl_cls() -> Type["SDPAImpl"]:
+    def get_impl_cls() -> type["SDPAImpl"]:
        return SDPAImpl

    # @staticmethod
@@ -40,7 +38,7 @@ class SDPAImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,6 +1,5 @@
 import json
 from dataclasses import dataclass
-from typing import List, Optional, Type

 import torch
 from einops import rearrange
@@ -20,7 +19,7 @@ logger = init_logger(__name__)


 # TODO(will-refactor): move this to a utils file
-def dict_to_3d_list(mask_strategy) -> List[List[List[Optional[torch.Tensor]]]]:
+def dict_to_3d_list(mask_strategy) -> list[list[list[torch.Tensor | None]]]:
    indices = [tuple(map(int, key.split('_'))) for key in mask_strategy]

    max_timesteps_idx = max(
@@ -58,7 +57,7 @@ class SlidingTileAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True

    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        # TODO(will-refactor): check this
        return [32, 64, 96, 128, 160, 192, 224, 256]

@@ -67,15 +66,15 @@ class SlidingTileAttentionBackend(AttentionBackend):
        return "SLIDING_TILE_ATTN"

    @staticmethod
-    def get_impl_cls() -> Type["SlidingTileAttentionImpl"]:
+    def get_impl_cls() -> type["SlidingTileAttentionImpl"]:
        return SlidingTileAttentionImpl

    @staticmethod
-    def get_metadata_cls() -> Type["SlidingTileAttentionMetadata"]:
+    def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]:
        return SlidingTileAttentionMetadata

    @staticmethod
-    def get_builder_cls() -> Type["SlidingTileAttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]:
        return SlidingTileAttentionMetadataBuilder


@@ -110,7 +109,7 @@ class SlidingTileAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0

-from typing import Optional, Tuple
-
 import torch
 import torch.nn as nn

@@ -22,11 +20,11 @@ class DistributedAttention(nn.Module):
    def __init__(self,
                 num_heads: int,
                 head_size: int,
-                 num_kv_heads: Optional[int] = None,
-                 softmax_scale: Optional[float] = None,
+                 num_kv_heads: int | None = None,
+                 softmax_scale: float | None = None,
                 causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[_Backend,
-                                                              ...]] = None,
+                 supported_attention_backends: tuple[_Backend, ...]
+                 | None = None,
                 prefix: str = "",
                 **extra_impl_args) -> None:
        super().__init__()
@@ -62,10 +60,10 @@ class DistributedAttention(nn.Module):
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
-        replicated_q: Optional[torch.Tensor] = None,
-        replicated_k: Optional[torch.Tensor] = None,
-        replicated_v: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        replicated_q: torch.Tensor | None = None,
+        replicated_k: torch.Tensor | None = None,
+        replicated_v: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Forward pass for distributed attention.
        
        Args:
@@ -84,9 +82,6 @@ class DistributedAttention(nn.Module):
        # Check input shapes
        assert q.dim() == 4 and k.dim() == 4 and v.dim(
        ) == 4, "Expected 4D tensors"
-        # assert bs = 1
-        assert q.shape[
-            0] == 1, "Batch size must be 1, and there should be no padding tokens"
        batch_size, seq_len, num_heads, head_dim = q.shape
        local_rank = get_sequence_model_parallel_rank()
        world_size = get_sequence_model_parallel_world_size()
@@ -144,11 +139,11 @@ class LocalAttention(nn.Module):
    def __init__(self,
                 num_heads: int,
                 head_size: int,
-                 num_kv_heads: Optional[int] = None,
-                 softmax_scale: Optional[float] = None,
+                 num_kv_heads: int | None = None,
+                 softmax_scale: float | None = None,
                 causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[_Backend,
-                                                              ...]] = None,
+                 supported_attention_backends: tuple[_Backend, ...]
+                 | None = None,
                 **extra_impl_args) -> None:
        super().__init__()
        if softmax_scale is None:
@@ -2,9 +2,10 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py

 import os
+from collections.abc import Generator
 from contextlib import contextmanager
 from functools import cache
-from typing import Generator, Optional, Tuple, Type, cast
+from typing import cast

 import torch

@@ -17,7 +18,7 @@ from fastvideo.v1.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
 logger = init_logger(__name__)


-def backend_name_to_enum(backend_name: str) -> Optional[_Backend]:
+def backend_name_to_enum(backend_name: str) -> _Backend | None:
    """
    Convert a string backend name to a _Backend enum value.

@@ -31,7 +32,7 @@ def backend_name_to_enum(backend_name: str) -> Optional[_Backend]:
          None


-def get_env_variable_attn_backend() -> Optional[_Backend]:
+def get_env_variable_attn_backend() -> _Backend | None:
    '''
    Get the backend override specified by the FastVideo attention
    backend environment variable, if one is specified.
@@ -53,10 +54,10 @@ def get_env_variable_attn_backend() -> Optional[_Backend]:
 #
 # THIS SELECTION TAKES PRECEDENCE OVER THE
 # FASTVIDEO ATTENTION BACKEND ENVIRONMENT VARIABLE
-forced_attn_backend: Optional[_Backend] = None
+forced_attn_backend: _Backend | None = None


-def global_force_attn_backend(attn_backend: Optional[_Backend]) -> None:
+def global_force_attn_backend(attn_backend: _Backend | None) -> None:
    '''
    Force all attention operations to use a specified backend.

@@ -71,7 +72,7 @@ def global_force_attn_backend(attn_backend: Optional[_Backend]) -> None:
    forced_attn_backend = attn_backend


-def get_global_forced_attn_backend() -> Optional[_Backend]:
+def get_global_forced_attn_backend() -> _Backend | None:
    '''
    Get the currently-forced choice of attention backend,
    or None if auto-selection is currently enabled.
@@ -82,8 +83,8 @@ def get_global_forced_attn_backend() -> Optional[_Backend]:
 def get_attn_backend(
    head_size: int,
    dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
-) -> Type[AttentionBackend]:
+    supported_attention_backends: tuple[_Backend, ...] | None = None,
+) -> type[AttentionBackend]:
    return _cached_get_attn_backend(head_size, dtype,
                                    supported_attention_backends)

@@ -92,8 +93,8 @@ def get_attn_backend(
 def _cached_get_attn_backend(
    head_size: int,
    dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
-) -> Type[AttentionBackend]:
+    supported_attention_backends: tuple[_Backend, ...] | None = None,
+) -> type[AttentionBackend]:
    # Check whether a particular choice of backend was
    # previously forced.
    #
@@ -102,13 +103,13 @@ def _cached_get_attn_backend(
    if not supported_attention_backends:
        raise ValueError("supported_attention_backends is empty")
    selected_backend = None
-    backend_by_global_setting: Optional[_Backend] = (
+    backend_by_global_setting: _Backend | None = (
        get_global_forced_attn_backend())
    if backend_by_global_setting is not None:
        selected_backend = backend_by_global_setting
    else:
        # Check the environment variable and override if specified
-        backend_by_env_var: Optional[str] = envs.FASTVIDEO_ATTENTION_BACKEND
+        backend_by_env_var: str | None = envs.FASTVIDEO_ATTENTION_BACKEND
        if backend_by_env_var is not None:
            selected_backend = backend_name_to_enum(backend_by_env_var)

@@ -120,7 +121,7 @@ def _cached_get_attn_backend(
    if not attention_cls:
        raise ValueError(
            f"Invalid attention backend for {current_platform.device_name}")
-    return cast(Type[AttentionBackend], resolve_obj_by_qualname(attention_cls))
+    return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls))


@contextmanager
@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field, fields
-from typing import Any, Dict
+from typing import Any

 from fastvideo.v1.logger import init_logger

@@ -41,7 +41,7 @@ class ModelConfig:
        self.__dict__.update(state)

    # This should be used only when loading from transformers/diffusers
-    def update_model_arch(self, source_model_dict: Dict[str, Any]) -> None:
+    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)}

@@ -55,7 +55,7 @@ class ModelConfig:
        if hasattr(arch_config, "__post_init__"):
            arch_config.__post_init__()

-    def update_model_config(self, source_model_dict: Dict[str, Any]) -> None:
+    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)}
@@ -1,4 +1,5 @@
 from fastvideo.v1.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
+from fastvideo.v1.configs.models.dits.stepvideo import StepVideoConfig
 from fastvideo.v1.configs.models.dits.wanvideo import WanVideoConfig

-__all__ = ["HunyuanVideoConfig", "WanVideoConfig"]
+__all__ = ["HunyuanVideoConfig", "WanVideoConfig", "StepVideoConfig"]
@@ -1,8 +1,8 @@
 from dataclasses import dataclass, field
-from typing import Optional, Tuple
+from typing import Any

 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.configs.quantization import QuantizationConfig
+from fastvideo.v1.layers.quantization import QuantizationConfig
 from fastvideo.v1.platforms import _Backend


@@ -11,7 +11,7 @@ class DiTArchConfig(ArchConfig):
    _fsdp_shard_conditions: list = field(default_factory=list)
    _compile_conditions: list = field(default_factory=list)
    _param_names_mapping: dict = field(default_factory=dict)
-    _supported_attention_backends: Tuple[_Backend,
+    _supported_attention_backends: tuple[_Backend,
                                         ...] = (_Backend.SLIDING_TILE_ATTN,
                                                 _Backend.SAGE_ATTN,
                                                 _Backend.FLASH_ATTN,
@@ -32,4 +32,25 @@ class DiTConfig(ModelConfig):

    # FastVideoDiT-specific parameters
    prefix: str = ""
-    quant_config: Optional[QuantizationConfig] = None
+    quant_config: QuantizationConfig | None = None
+
+    @staticmethod
+    def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any:
+        """Add CLI arguments for DiTConfig fields"""
+        parser.add_argument(
+            f"--{prefix}.prefix",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.prefix",
+            default=DiTConfig.prefix,
+            help="Prefix for the DiT model",
+        )
+
+        parser.add_argument(
+            f"--{prefix}.quant-config",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.quant_config",
+            default=None,
+            help="Quantization configuration for the DiT model",
+        )
+
+        return parser
@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from typing import Optional, Tuple

 import torch

@@ -156,9 +155,9 @@ class HunyuanVideoArchConfig(DiTArchConfig):
    num_layers: int = 20
    num_single_layers: int = 40
    num_refiner_layers: int = 2
-    rope_axes_dim: Tuple[int, int, int] = (16, 56, 56)
+    rope_axes_dim: tuple[int, int, int] = (16, 56, 56)
    guidance_embeds: bool = False
-    dtype: Optional[torch.dtype] = None
+    dtype: torch.dtype | None = None
    text_embed_dim: int = 4096
    pooled_projection_dim: int = 768
    rope_theta: int = 256
@@ -0,0 +1,64 @@
+from dataclasses import dataclass, field
+
+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 StepVideoArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    _param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # transformer block
+            r"^transformer_blocks\.(\d+)\.norm1\.(weight|bias)$":
+            r"transformer_blocks.\1.norm1.norm.\2",
+            r"^transformer_blocks\.(\d+)\.norm2\.(weight|bias)$":
+            r"transformer_blocks.\1.norm2.norm.\2",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.0\.proj\.weight$":
+            r"transformer_blocks.\1.ff.fc_in.weight",
+            r"^transformer_blocks\.(\d+)\.ff\.net\.2\.weight$":
+            r"transformer_blocks.\1.ff.fc_out.weight",
+
+            # adanorm block
+            r"^adaln_single\.emb\.timestep_embedder\.linear_1\.(weight|bias)$":
+            r"adaln_single.emb.mlp.fc_in.\1",
+            r"^adaln_single\.emb\.timestep_embedder\.linear_2\.(weight|bias)$":
+            r"adaln_single.emb.mlp.fc_out.\1",
+
+            # caption projection
+            r"^caption_projection\.linear_1\.(weight|bias)$":
+            r"caption_projection.fc_in.\1",
+            r"^caption_projection\.linear_2\.(weight|bias)$":
+            r"caption_projection.fc_out.\1",
+        })
+
+    num_attention_heads: int = 48
+    attention_head_dim: int = 128
+    in_channels: int = 64
+    out_channels: int | None = 64
+    num_layers: int = 48
+    dropout: float = 0.0
+    patch_size: int = 1
+    norm_type: str = "ada_norm_single"
+    norm_elementwise_affine: bool = False
+    norm_eps: float = 1e-6
+    caption_channels: int | list[int] | tuple[int, ...] | None = field(
+        default_factory=lambda: [6144, 1024])
+    attention_type: str | None = "torch"
+    use_additional_conditions: bool | None = False
+
+    def __post_init__(self):
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.out_channels = self.in_channels if self.out_channels is None else self.out_channels
+        self.num_channels_latents = self.out_channels
+
+
+@dataclass
+class StepVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=StepVideoArchConfig)
+
+    prefix: str = "StepVideo"
@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from typing import Optional, Tuple

 from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig

@@ -52,7 +51,7 @@ class WanVideoArchConfig(DiTArchConfig):
            r"blocks.\1.self_attn_residual_norm.norm.\2",
        })

-    patch_size: Tuple[int, int, int] = (1, 2, 2)
+    patch_size: tuple[int, int, int] = (1, 2, 2)
    text_len = 512
    num_attention_heads: int = 40
    attention_head_dim: int = 128
@@ -65,8 +64,8 @@ class WanVideoArchConfig(DiTArchConfig):
    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
+    image_dim: int | None = None
+    added_kv_proj_dim: int | None = None
    rope_max_seq_len: int = 1024

    def __post_init__(self):
@@ -1,17 +1,17 @@
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any

 import torch

 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.configs.quantization import QuantizationConfig
+from fastvideo.v1.layers.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,
+    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
@@ -32,7 +32,7 @@ class TextEncoderArchConfig(EncoderArchConfig):
    scalable_attention: bool = True
    tie_word_embeddings: bool = False

-    tokenizer_kwargs: Dict[str, Any] = field(default_factory=dict)
+    tokenizer_kwargs: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.tokenizer_kwargs = {
@@ -49,11 +49,11 @@ class ImageEncoderArchConfig(EncoderArchConfig):

@dataclass
 class BaseEncoderOutput:
-    last_hidden_state: Optional[torch.FloatTensor] = None
-    pooler_output: Optional[torch.FloatTensor] = None
-    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
-    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
-    attention_mask: Optional[torch.Tensor] = None
+    last_hidden_state: torch.FloatTensor | None = None
+    pooler_output: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    attention_mask: torch.Tensor | None = None


@dataclass
@@ -61,8 +61,8 @@ class EncoderConfig(ModelConfig):
    arch_config: ArchConfig = field(default_factory=EncoderArchConfig)

    prefix: str = ""
-    quant_config: Optional[QuantizationConfig] = None
-    lora_config: Optional[Any] = None
+    quant_config: QuantizationConfig | None = None
+    lora_config: Any | None = None


@dataclass
@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from typing import Optional

 from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                       ImageEncoderConfig,
@@ -51,8 +50,8 @@ class CLIPTextConfig(TextEncoderConfig):
    arch_config: TextEncoderArchConfig = field(
        default_factory=CLIPTextArchConfig)

-    num_hidden_layers_override: Optional[int] = None
-    require_post_norm: Optional[bool] = None
+    num_hidden_layers_override: int | None = None
+    require_post_norm: bool | None = None
    prefix: str = "clip"


@@ -61,6 +60,6 @@ class CLIPVisionConfig(ImageEncoderConfig):
    arch_config: ImageEncoderArchConfig = field(
        default_factory=CLIPVisionArchConfig)

-    num_hidden_layers_override: Optional[int] = None
-    require_post_norm: Optional[bool] = None
+    num_hidden_layers_override: int | None = None
+    require_post_norm: bool | None = None
    prefix: str = "clip"
@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from typing import Optional

 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                       TextEncoderConfig)
@@ -12,7 +11,7 @@ class LlamaArchConfig(TextEncoderArchConfig):
    intermediate_size: int = 11008
    num_hidden_layers: int = 32
    num_attention_heads: int = 32
-    num_key_value_heads: Optional[int] = None
+    num_key_value_heads: int | None = None
    hidden_act: str = "silu"
    max_position_embeddings: int = 2048
    initializer_range: float = 0.02
@@ -24,11 +23,11 @@ class LlamaArchConfig(TextEncoderArchConfig):
    pretraining_tp: int = 1
    tie_word_embeddings: bool = False
    rope_theta: float = 10000.0
-    rope_scaling: Optional[float] = None
+    rope_scaling: float | None = None
    attention_bias: bool = False
    attention_dropout: float = 0.0
    mlp_bias: bool = False
-    head_dim: Optional[int] = None
+    head_dim: int | None = None
    hidden_state_skip_layer: int = 2
    text_len: int = 256

@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from typing import Optional

 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                       TextEncoderConfig)
@@ -12,7 +11,7 @@ class T5ArchConfig(TextEncoderArchConfig):
    d_kv: int = 64
    d_ff: int = 2048
    num_layers: int = 6
-    num_decoder_layers: Optional[int] = None
+    num_decoder_layers: int | None = None
    num_heads: int = 8
    relative_attention_num_buckets: int = 32
    relative_attention_max_distance: int = 128
@@ -1,7 +1,9 @@
 from fastvideo.v1.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
+from fastvideo.v1.configs.models.vaes.stepvideovae import StepVideoVAEConfig
 from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig

 __all__ = [
    "HunyuanVAEConfig",
    "WanVAEConfig",
+    "StepVideoVAEConfig",
 ]
@@ -1,14 +1,15 @@
 from dataclasses import dataclass, field
-from typing import Union
+from typing import Any

 import torch

 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.v1.utils import StoreBoolean


@dataclass
 class VAEArchConfig(ArchConfig):
-    scaling_factor: Union[float, torch.tensor] = 0
+    scaling_factor: float | torch.Tensor = 0

    temporal_compression_ratio: int = 4
    spatial_compression_ratio: int = 8
@@ -33,6 +34,97 @@ class VAEConfig(ModelConfig):
    use_tiling: bool = True
    use_temporal_tiling: bool = True
    use_parallel_tiling: bool = True
+    use_temporal_scaling_frames: bool = True

    def __post_init__(self):
        self.blend_num_frames = self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
+
+    @staticmethod
+    def add_cli_args(parser: Any, prefix: str = "vae-config") -> Any:
+        """Add CLI arguments for VAEConfig fields"""
+        parser.add_argument(
+            f"--{prefix}.load-encoder",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.load_encoder",
+            default=VAEConfig.load_encoder,
+            help="Whether to load the VAE encoder",
+        )
+        parser.add_argument(
+            f"--{prefix}.load-decoder",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.load_decoder",
+            default=VAEConfig.load_decoder,
+            help="Whether to load the VAE decoder",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-min-height",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_min_height",
+            default=VAEConfig.tile_sample_min_height,
+            help="Minimum height for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-min-width",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_min_width",
+            default=VAEConfig.tile_sample_min_width,
+            help="Minimum width for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-min-num-frames",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_min_num_frames",
+            default=VAEConfig.tile_sample_min_num_frames,
+            help="Minimum number of frames for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-stride-height",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_stride_height",
+            default=VAEConfig.tile_sample_stride_height,
+            help="Stride height for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-stride-width",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_stride_width",
+            default=VAEConfig.tile_sample_stride_width,
+            help="Stride width for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.tile-sample-stride-num-frames",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.tile_sample_stride_num_frames",
+            default=VAEConfig.tile_sample_stride_num_frames,
+            help="Stride number of frames for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.blend-num-frames",
+            type=int,
+            dest=f"{prefix.replace('-', '_')}.blend_num_frames",
+            default=VAEConfig.blend_num_frames,
+            help="Number of frames to blend for VAE tile sampling",
+        )
+        parser.add_argument(
+            f"--{prefix}.use-tiling",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.use_tiling",
+            default=VAEConfig.use_tiling,
+            help="Whether to use tiling for VAE",
+        )
+        parser.add_argument(
+            f"--{prefix}.use-temporal-tiling",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.use_temporal_tiling",
+            default=VAEConfig.use_temporal_tiling,
+            help="Whether to use temporal tiling for VAE",
+        )
+        parser.add_argument(
+            f"--{prefix}.use-parallel-tiling",
+            action=StoreBoolean,
+            dest=f"{prefix.replace('-', '_')}.use_parallel_tiling",
+            default=VAEConfig.use_parallel_tiling,
+            help="Whether to use parallel tiling for VAE",
+        )
+
+        return parser
@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from typing import Tuple

 from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig

@@ -9,19 +8,19 @@ class HunyuanVAEArchConfig(VAEArchConfig):
    in_channels: int = 3
    out_channels: int = 3
    latent_channels: int = 16
-    down_block_types: Tuple[str, ...] = (
+    down_block_types: tuple[str, ...] = (
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
    )
-    up_block_types: Tuple[str, ...] = (
+    up_block_types: tuple[str, ...] = (
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
    )
-    block_out_channels: Tuple[int, ...] = (128, 256, 512, 512)
+    block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
    layers_per_block: int = 2
    act_fn: str = "silu"
    norm_num_groups: int = 32
@@ -0,0 +1,28 @@
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class StepVideoVAEArchConfig(VAEArchConfig):
+    in_channels: int = 3
+    out_channels: int = 3
+    z_channels: int = 64
+    num_res_blocks: int = 2
+    version: int = 2
+    frame_len: int = 17
+    world_size: int = 1
+
+    spatial_compression_ratio: int = 16
+    temporal_compression_ratio: int = 8
+
+    scaling_factor: float = 1.0
+
+
+@dataclass
+class StepVideoVAEConfig(VAEConfig):
+    arch_config: VAEArchConfig = field(default_factory=StepVideoVAEArchConfig)
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+    use_temporal_scaling_frames: bool = False
@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from typing import Tuple

 import torch

@@ -10,12 +9,12 @@ from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
 class WanVAEArchConfig(VAEArchConfig):
    base_dim: int = 96
    z_dim: int = 16
-    dim_mult: Tuple[int, ...] = (1, 2, 4, 4)
+    dim_mult: tuple[int, ...] = (1, 2, 4, 4)
    num_res_blocks: int = 2
-    attn_scales: Tuple[float, ...] = ()
-    temperal_downsample: Tuple[bool, ...] = (False, True, True)
+    attn_scales: tuple[float, ...] = ()
+    temperal_downsample: tuple[bool, ...] = (False, True, True)
    dropout: float = 0.0
-    latents_mean: Tuple[float, ...] = (
+    latents_mean: tuple[float, ...] = (
        -0.7571,
        -0.7089,
        -0.9113,
@@ -33,7 +32,7 @@ class WanVAEArchConfig(VAEArchConfig):
        0.2503,
        -0.2921,
    )
-    latents_std: Tuple[float, ...] = (
+    latents_std: tuple[float, ...] = (
        2.8184,
        1.4541,
        2.3275,
@@ -55,9 +54,9 @@ class WanVAEArchConfig(VAEArchConfig):
    spatial_compression_ratio = 8

    def __post_init__(self):
-        self.scaling_factor: torch.tensor = 1.0 / torch.tensor(
+        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(
+        self.shift_factor: torch.Tensor = torch.tensor(self.latents_mean).view(
            1, self.z_dim, 1, 1, 1)


@@ -4,6 +4,7 @@ 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.stepvideo import StepVideoT2VConfig
 from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
                                                WanI2V720PConfig,
                                                WanT2V480PConfig,
@@ -12,5 +13,6 @@ from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
 __all__ = [
    "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
    "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
-    "WanT2V720PConfig", "WanI2V720PConfig", "get_pipeline_config_cls_for_name"
+    "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig",
+    "get_pipeline_config_cls_for_name"
 ]
@@ -1,6 +1,7 @@
 import json
+from collections.abc import Callable
 from dataclasses import asdict, dataclass, field, fields
-from typing import Any, Callable, Dict, Optional, Tuple, cast
+from typing import Any, cast

 import torch

@@ -17,7 +18,7 @@ def preprocess_text(prompt: str) -> str:
    return prompt


-def postprocess_text(output: BaseEncoderOutput) -> torch.tensor:
+def postprocess_text(output: BaseEncoderOutput) -> torch.Tensor:
    raise NotImplementedError


@@ -26,7 +27,7 @@ class PipelineConfig:
    """Base configuration for all pipeline architectures."""
    # Video generation parameters
    embedded_cfg_scale: float = 6.0
-    flow_shift: Optional[float] = None
+    flow_shift: float | None = None
    use_cpu_offload: bool = False
    disable_autocast: bool = False

@@ -43,18 +44,18 @@ class PipelineConfig:
    dit_config: DiTConfig = field(default_factory=DiTConfig)

    # Text encoder configuration
-    text_encoder_precisions: Tuple[str, ...] = field(
+    text_encoder_precisions: tuple[str, ...] = field(
        default_factory=lambda: ("fp16", ))
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
        default_factory=lambda: (EncoderConfig(), ))
-    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
        default_factory=lambda: (preprocess_text, ))
-    postprocess_text_funcs: Tuple[Callable[[BaseEncoderOutput], torch.tensor],
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
                                  ...] = field(default_factory=lambda:
                                               (postprocess_text, ))

    # STA (Spatial-Temporal Attention) parameters
-    mask_strategy_file_path: Optional[str] = None
+    mask_strategy_file_path: str | None = None

    # Compilation
    enable_torch_compile: bool = False
@@ -107,7 +108,7 @@ class PipelineConfig:
            input_pipeline_dict = json.load(f)
        self.update_pipeline_config(input_pipeline_dict)

-    def update_pipeline_config(self, source_pipeline_dict: Dict[str,
+    def update_pipeline_config(self, source_pipeline_dict: dict[str,
                                                                Any]) -> None:
        for f in fields(self):
            key = f.name
@@ -123,8 +124,9 @@ class PipelineConfig:
                    assert len(current_value) == len(
                        new_value
                    ), "Users shouldn't delete or add text encoder config objects in your json"
-                    for target_config, source_config in zip(
-                            current_value, new_value):
+                    for target_config, source_config in zip(current_value,
+                                                            new_value,
+                                                            strict=False):
                        target_config.update_model_config(source_config)
                else:
                    setattr(self, key, new_value)
@@ -1,5 +1,6 @@
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Callable, Tuple, TypedDict
+from typing import TypedDict

 import torch

@@ -35,11 +36,11 @@ def llama_preprocess_text(prompt: str) -> str:
    return prompt_template_video["template"].format(prompt)


-def llama_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
+def llama_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
    hidden_state_skip_layer = 2
    assert outputs.hidden_states is not None
-    hidden_states: Tuple[torch.Tensor, ...] = outputs.hidden_states
-    last_hidden_state: torch.tensor = hidden_states[-(hidden_state_skip_layer +
+    hidden_states: tuple[torch.Tensor, ...] = outputs.hidden_states
+    last_hidden_state: torch.Tensor = hidden_states[-(hidden_state_skip_layer +
                                                      1)]
    crop_start = prompt_template_video.get("crop_start", -1)
    last_hidden_state = last_hidden_state[:, crop_start:]
@@ -50,8 +51,8 @@ def clip_preprocess_text(prompt: str) -> str:
    return prompt


-def clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
-    pooler_output: torch.tensor = outputs.pooler_output
+def clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    pooler_output: torch.Tensor = outputs.pooler_output
    return pooler_output


@@ -68,20 +69,23 @@ class HunyuanConfig(PipelineConfig):
    embedded_cfg_scale: int = 6
    flow_shift: int = 7

+    # Video parameters
+    use_cpu_offload: bool = True
+
    # Text encoding stage
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
        default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))
-    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
        default_factory=lambda: (llama_preprocess_text, clip_preprocess_text))
-    postprocess_text_funcs: Tuple[
-        Callable[[BaseEncoderOutput], torch.tensor],
+    postprocess_text_funcs: tuple[
+        Callable[[BaseEncoderOutput], torch.Tensor],
        ...] = field(default_factory=lambda:
                     (llama_postprocess_text, clip_postprocess_text))

    # Precision for each component
    precision: str = "bf16"
    vae_precision: str = "fp16"
-    text_encoder_precisions: Tuple[str, ...] = field(
+    text_encoder_precisions: tuple[str, ...] = field(
        default_factory=lambda: ("fp16", "fp16"))

    def __post_init__(self):
@@ -1,11 +1,12 @@
 """Registry for pipeline weight-specific configurations."""

 import os
-from typing import Callable, Dict, Optional, Type
+from collections.abc import Callable

 from fastvideo.v1.configs.pipelines.base import PipelineConfig
 from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
                                                    HunyuanConfig)
+from fastvideo.v1.configs.pipelines.stepvideo import StepVideoT2VConfig
 from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
                                                WanI2V720PConfig,
                                                WanT2V480PConfig,
@@ -17,37 +18,40 @@ from fastvideo.v1.utils import (maybe_download_model_index,
 logger = init_logger(__name__)

 # Registry maps specific model weights to their config classes
-WEIGHT_CONFIG_REGISTRY: Dict[str, Type[PipelineConfig]] = {
+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,
    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig,
+    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
    # Add other specific weight variants
 }

 # For determining pipeline type from model ID
-PIPELINE_DETECTOR: Dict[str, Callable[[str], bool]] = {
+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(),
+    "stepvideo": lambda id: "stepvideo" 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]] = {
+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,
+    "stepvideo": StepVideoT2VConfig
    # Other fallbacks by architecture
 }


 def get_pipeline_config_cls_for_name(
-        pipeline_name_or_path: str) -> Optional[type[PipelineConfig]]:
+        pipeline_name_or_path: str) -> type[PipelineConfig] | None:
    """Get the appropriate config class for specific pretrained weights."""

    if os.path.exists(pipeline_name_or_path):
@@ -0,0 +1,32 @@
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.models import DiTConfig, VAEConfig
+from fastvideo.v1.configs.models.dits import StepVideoConfig
+from fastvideo.v1.configs.models.vaes import StepVideoVAEConfig
+from fastvideo.v1.configs.pipelines.base import PipelineConfig
+
+
+@dataclass
+class StepVideoT2VConfig(PipelineConfig):
+    """Base configuration for StepVideo pipeline architecture."""
+
+    # WanConfig-specific parameters with defaults
+    # DiT
+    dit_config: DiTConfig = field(default_factory=StepVideoConfig)
+    # VAE
+    vae_config: VAEConfig = field(default_factory=StepVideoVAEConfig)
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # Video parameters
+    use_cpu_offload: bool = True
+
+    # Denoising stage
+    flow_shift: int = 13
+    timesteps_scale: bool = False
+    pos_magic: str = "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。"
+    neg_magic: str = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。"
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "bf16"
@@ -1,5 +1,5 @@
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Callable, Tuple

 import torch

@@ -11,13 +11,15 @@ from fastvideo.v1.configs.models.vaes import WanVAEConfig
 from fastvideo.v1.configs.pipelines.base import PipelineConfig


-def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.tensor:
-    mask: torch.tensor = outputs.attention_mask
-    hidden_state: torch.tensor = outputs.last_hidden_state
+def t5_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    mask: torch.Tensor = outputs.attention_mask
+    hidden_state: torch.Tensor = outputs.last_hidden_state
    seq_lens = mask.gt(0).sum(dim=1).long()
    assert torch.isnan(hidden_state).sum() == 0
-    prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens)]
-    prompt_embeds_tensor: torch.tensor = torch.stack([
+    prompt_embeds = [
+        u[:v] for u, v in zip(hidden_state, seq_lens, strict=False)
+    ]
+    prompt_embeds_tensor: torch.Tensor = torch.stack([
        torch.cat([u, u.new_zeros(512 - u.size(0), u.size(1))])
        for u in prompt_embeds
    ],
@@ -44,16 +46,16 @@ class WanT2V480PConfig(PipelineConfig):
    flow_shift: int = 3

    # Text encoding stage
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
        default_factory=lambda: (T5Config(), ))
-    postprocess_text_funcs: Tuple[Callable[[BaseEncoderOutput], torch.tensor],
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
                                  ...] = field(default_factory=lambda:
                                               (t5_postprocess_text, ))

    # Precision for each component
    precision: str = "bf16"
    vae_precision: str = "fp16"
-    text_encoder_precisions: Tuple[str, ...] = field(
+    text_encoder_precisions: tuple[str, ...] = field(
        default_factory=lambda: ("fp32", ))

    # WanConfig-specific added parameters
@@ -1,3 +0,0 @@
-from fastvideo.v1.configs.quantization.base import QuantizationConfig
-
-__all__ = ["QuantizationConfig"]
@@ -1,6 +0,0 @@
-from dataclasses import dataclass
-
-
-@dataclass
-class QuantizationConfig:
-    pass
@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+from typing import Any

 from fastvideo.v1.logger import init_logger

@@ -15,12 +15,12 @@ class SamplingParam:
    data_type: str = "video"

    # Image inputs
-    image_path: Optional[str] = None
+    image_path: str | None = None

    # Text inputs
-    prompt: Optional[Union[str, List[str]]] = None
-    negative_prompt: Optional[str] = None
-    prompt_path: Optional[str] = None
+    prompt: str | list[str] | None = None
+    negative_prompt: str | None = None
+    prompt_path: str | None = None
    output_path: str = "outputs/"

    # Batch info
@@ -53,7 +53,7 @@ class SamplingParam:
        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:
+    def update(self, source_dict: dict[str, Any]) -> None:
        for key, value in source_dict.items():
            if hasattr(self, key):
                setattr(self, key, value)
@@ -78,6 +78,113 @@ class SamplingParam:

        return sampling_param

+    @staticmethod
+    def add_cli_args(parser: Any) -> Any:
+        """Add CLI arguments for SamplingParam fields"""
+        parser.add_argument(
+            "--prompt",
+            type=str,
+            default=SamplingParam.prompt,
+            help="Text prompt for video generation",
+        )
+        parser.add_argument(
+            "--negative-prompt",
+            type=str,
+            default=SamplingParam.negative_prompt,
+            help="Negative text prompt for video generation",
+        )
+        parser.add_argument(
+            "--prompt-path",
+            type=str,
+            default=SamplingParam.prompt_path,
+            help="Path to a text file containing the prompt",
+        )
+        parser.add_argument(
+            "--output-path",
+            type=str,
+            default=SamplingParam.output_path,
+            help="Path to save the generated video",
+        )
+        parser.add_argument(
+            "--num-videos-per-prompt",
+            type=int,
+            default=SamplingParam.num_videos_per_prompt,
+            help="Number of videos to generate per prompt",
+        )
+        parser.add_argument(
+            "--seed",
+            type=int,
+            default=SamplingParam.seed,
+            help="Random seed for generation",
+        )
+        parser.add_argument(
+            "--num-frames",
+            type=int,
+            default=SamplingParam.num_frames,
+            help="Number of frames to generate",
+        )
+        parser.add_argument(
+            "--height",
+            type=int,
+            default=SamplingParam.height,
+            help="Height of generated video",
+        )
+        parser.add_argument(
+            "--width",
+            type=int,
+            default=SamplingParam.width,
+            help="Width of generated video",
+        )
+        parser.add_argument(
+            "--fps",
+            type=int,
+            default=SamplingParam.fps,
+            help="Frames per second for saved video",
+        )
+        parser.add_argument(
+            "--num-inference-steps",
+            type=int,
+            default=SamplingParam.num_inference_steps,
+            help="Number of denoising steps",
+        )
+        parser.add_argument(
+            "--guidance-scale",
+            type=float,
+            default=SamplingParam.guidance_scale,
+            help="Classifier-free guidance scale",
+        )
+        parser.add_argument(
+            "--guidance-rescale",
+            type=float,
+            default=SamplingParam.guidance_rescale,
+            help="Guidance rescale factor",
+        )
+        parser.add_argument(
+            "--save-video",
+            action="store_true",
+            default=SamplingParam.save_video,
+            help="Whether to save the video to disk",
+        )
+        parser.add_argument(
+            "--no-save-video",
+            action="store_false",
+            dest="save_video",
+            help="Don't save the video to disk",
+        )
+        parser.add_argument(
+            "--return-frames",
+            action="store_true",
+            default=SamplingParam.return_frames,
+            help="Whether to return the raw frames",
+        )
+        parser.add_argument(
+            "--image-path",
+            type=str,
+            default=SamplingParam.image_path,
+            help="Path to input image for image-to-video generation",
+        )
+        return parser
+

@dataclass
 class CacheParams:
@@ -1,8 +1,10 @@
 import os
-from typing import Any, Callable, Dict, Optional
+from collections.abc import Callable
+from typing import Any

 from fastvideo.v1.configs.sample.hunyuan import (FastHunyuanSamplingParam,
                                                 HunyuanSamplingParam)
+from fastvideo.v1.configs.sample.stepvideo import StepVideoT2VSamplingParam
 from fastvideo.v1.configs.sample.wan import (WanI2V_14B_480P_SamplingParam,
                                             WanI2V_14B_720P_SamplingParam,
                                             WanT2V_1_3B_SamplingParam,
@@ -13,37 +15,39 @@ from fastvideo.v1.utils import (maybe_download_model_index,

 logger = init_logger(__name__)
 # Registry maps specific model weights to their config classes
-SAMPLING_PARAM_REGISTRY: Dict[str, Any] = {
+SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
    "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam,
    "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam,
    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam,
    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam,
    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam,
    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam,
+    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam,
    # Add other specific weight variants
 }

 # For determining pipeline type from model ID
-SAMPLING_PARAM_DETECTOR: Dict[str, Callable[[str], bool]] = {
+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(),
+    "stepvideo": lambda id: "stepvideo" 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] = {
+SAMPLING_FALLBACK_PARAM: dict[str, Any] = {
    "hunyuan":
    HunyuanSamplingParam,  # Base Hunyuan config as fallback for any Hunyuan variant
    "wanpipeline":
    WanT2V_1_3B_SamplingParam,  # Base Wan config as fallback for any Wan variant
    "wanimagetovideo": WanI2V_14B_480P_SamplingParam,
+    "stepvideo": StepVideoT2VSamplingParam
    # Other fallbacks by architecture
 }


-def get_sampling_param_cls_for_name(
-        pipeline_name_or_path: str) -> Optional[Any]:
+def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
    """Get the appropriate sampling param for specific pretrained weights."""

    if os.path.exists(pipeline_name_or_path):
@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+from fastvideo.v1.configs.sample.base import SamplingParam
+
+
+@dataclass
+class StepVideoT2VSamplingParam(SamplingParam):
+    # Video parameters
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 81
+
+    # Denoising stage
+    guidance_scale: float = 9.0
+    num_inference_steps: int = 50
+
+    # neg magic and pos magic
+    # pos_magic: str = "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。"
+    # neg_magic: str = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。"
@@ -1,16 +0,0 @@
-num_gpus: 4
-model_path: FastVideo/FastHunyuan-diffusers
-master_port: 29503
-sp_size: 4 
-tp_size: 4 
-height: 720 
-width: 1280 
-num_frames: 125 
-num_inference_steps: 6 
-guidance_scale: 1 
-embedded_cfg_scale: 6 
-flow_shift: 17 
-prompt_path: ./assets/prompt.txt 
-seed: 1024 
-output_path: outputs_video/ 
-vae-sp: True
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py

-from typing import Optional
-
 import torch
 import torch.distributed as dist
 from torch.distributed import ProcessGroup
@@ -18,8 +16,8 @@ class DeviceCommunicatorBase:

    def __init__(self,
                 cpu_group: ProcessGroup,
-                 device: Optional[torch.device] = None,
-                 device_group: Optional[ProcessGroup] = None,
+                 device: torch.device | None = None,
+                 device_group: ProcessGroup | None = None,
                 unique_name: str = ""):
        self.device = device or torch.device("cpu")
        self.cpu_group = cpu_group
@@ -66,7 +64,7 @@ class DeviceCommunicatorBase:
    def gather(self,
               input_: torch.Tensor,
               dst: int = 0,
-               dim: int = -1) -> Optional[torch.Tensor]:
+               dim: int = -1) -> torch.Tensor | None:
        """
        NOTE: We assume that the input tensor is on the same device across
        all the ranks.
@@ -170,7 +168,7 @@ class DeviceCommunicatorBase:
            raise RuntimeError(
                "scatter_dim must be 1 or 2 and gather_dim must be 1 or 2")

-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
        """Sends a tensor to the destination rank in a non-blocking way"""
        """NOTE: `dst` is the local rank of the destination rank."""
        if dst is None:
@@ -180,7 +178,7 @@ class DeviceCommunicatorBase:
    def recv(self,
             size: torch.Size,
             dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
+             src: int | None = None) -> torch.Tensor:
        """Receives a tensor from the source rank."""
        """NOTE: `src` is the local rank of the source rank."""
        if src is None:
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py

-from typing import Optional
-
 import torch
 from torch.distributed import ProcessGroup

@@ -14,15 +12,15 @@ class CudaCommunicator(DeviceCommunicatorBase):

    def __init__(self,
                 cpu_group: ProcessGroup,
-                 device: Optional[torch.device] = None,
-                 device_group: Optional[ProcessGroup] = None,
+                 device: torch.device | None = None,
+                 device_group: ProcessGroup | None = None,
                 unique_name: str = ""):
        super().__init__(cpu_group, device, device_group, unique_name)

        from fastvideo.v1.distributed.device_communicators.pynccl import (
            PyNcclCommunicator)

-        self.pynccl_comm: Optional[PyNcclCommunicator] = None
+        self.pynccl_comm: PyNcclCommunicator | None = None
        if self.world_size > 1:
            self.pynccl_comm = PyNcclCommunicator(
                group=self.cpu_group,
@@ -42,7 +40,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
            torch.distributed.all_reduce(out, group=self.device_group)
        return out

-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
        """Sends a tensor to the destination rank in a non-blocking way"""
        """NOTE: `dst` is the local rank of the destination rank."""
        if dst is None:
@@ -57,7 +55,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
    def recv(self,
             size: torch.Size,
             dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
+             src: int | None = None) -> torch.Tensor:
        """Receives a tensor from the source rank."""
        """NOTE: `src` is the local rank of the source rank."""
        if src is None:
@@ -1,8 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py

-from typing import Optional, Union
-
 # ===================== import region =====================
 import torch
 import torch.distributed as dist
@@ -22,9 +20,9 @@ class PyNcclCommunicator:

    def __init__(
        self,
-        group: Union[ProcessGroup, StatelessProcessGroup],
-        device: Union[int, str, torch.device],
-        library_path: Optional[str] = None,
+        group: ProcessGroup | StatelessProcessGroup,
+        device: int | str | torch.device,
+        library_path: str | None = None,
    ):
        """
        Args:
@@ -27,7 +27,7 @@
 import ctypes
 import platform
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any

 import torch
 from torch.distributed import ReduceOp
@@ -124,7 +124,7 @@ class ncclRedOpTypeEnum:
 class Function:
    name: str
    restype: Any
-    argtypes: List[Any]
+    argtypes: list[Any]


 class NCCLLibrary:
@@ -212,13 +212,13 @@ class NCCLLibrary:

    # class attribute to store the mapping from the path to the library
    # to avoid loading the same library multiple times
-    path_to_library_cache: Dict[str, Any] = {}
+    path_to_library_cache: dict[str, Any] = {}

    # class attribute to store the mapping from library path
    #  to the corresponding dictionary
-    path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}
+    path_to_dict_mapping: dict[str, dict[str, Any]] = {}

-    def __init__(self, so_file: Optional[str] = None):
+    def __init__(self, so_file: str | None = None):

        so_file = so_file or find_nccl_library()

@@ -240,7 +240,7 @@ class NCCLLibrary:
            raise e

        if so_file not in NCCLLibrary.path_to_dict_mapping:
-            _funcs: Dict[str, Any] = {}
+            _funcs: dict[str, Any] = {}
            for func in NCCLLibrary.exported_functions:
                f = getattr(self.lib, func.name)
                f.restype = func.restype
@@ -27,10 +27,11 @@ import gc
 import pickle
 import weakref
 from collections import namedtuple
+from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import dataclass
 from multiprocessing import shared_memory
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Optional
 from unittest.mock import patch

 import torch
@@ -57,15 +58,15 @@ TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])


 def _split_tensor_dict(
-    tensor_dict: Dict[str, Union[torch.Tensor, Any]]
-) -> Tuple[List[Tuple[str, Any]], List[torch.Tensor]]:
+    tensor_dict: dict[str, torch.Tensor | Any]
+) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:
    """Split the tensor dictionary into two parts:
    1. A list of (key, value) pairs. If the value is a tensor, it is replaced
         by its metadata.
    2. A list of tensors.
    """
-    metadata_list: List[Tuple[str, Any]] = []
-    tensor_list: List[torch.Tensor] = []
+    metadata_list: list[tuple[str, Any]] = []
+    tensor_list: list[torch.Tensor] = []
    for key, value in tensor_dict.items():
        if isinstance(value, torch.Tensor):
            # Note: we cannot use `value.device` here,
@@ -81,7 +82,7 @@ def _split_tensor_dict(
    return metadata_list, tensor_list


-_group_name_counter: Dict[str, int] = {}
+_group_name_counter: dict[str, int] = {}


 def _get_unique_name(name: str) -> str:
@@ -97,7 +98,7 @@ def _get_unique_name(name: str) -> str:
    return newname


-_groups: Dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}
+_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}


 def _register_group(group: "GroupCoordinator") -> None:
@@ -128,7 +129,7 @@ class GroupCoordinator:

    # available attributes:
    rank: int  # global rank
-    ranks: List[int]  # global ranks in the group
+    ranks: list[int]  # global ranks in the group
    world_size: int  # size of the group
    # difference between `local_rank` and `rank_in_group`:
    # if we have a group of size 4 across two nodes:
@@ -143,16 +144,16 @@ class GroupCoordinator:
    device_group: ProcessGroup  # group for device communication
    use_device_communicator: bool  # whether to use device communicator
    device_communicator: DeviceCommunicatorBase  # device communicator
-    mq_broadcaster: Optional[Any]  # shared memory broadcaster
+    mq_broadcaster: Any | None  # shared memory broadcaster

    def __init__(
        self,
-        group_ranks: List[List[int]],
+        group_ranks: list[list[int]],
        local_rank: int,
-        torch_distributed_backend: Union[str, Backend],
+        torch_distributed_backend: str | Backend,
        use_device_communicator: bool,
        use_message_queue_broadcaster: bool = False,
-        group_name: Optional[str] = None,
+        group_name: str | None = None,
    ):
        group_name = group_name or "anonymous"
        self.unique_name = _get_unique_name(group_name)
@@ -243,8 +244,8 @@ class GroupCoordinator:
        return self.ranks[(rank_in_group - 1) % world_size]

    @contextmanager
-    def graph_capture(
-            self, graph_capture_context: Optional[GraphCaptureContext] = None):
+    def graph_capture(self,
+                      graph_capture_context: GraphCaptureContext | None = None):
        if graph_capture_context is None:
            stream = torch.cuda.Stream()
            graph_capture_context = GraphCaptureContext(stream)
@@ -301,7 +302,7 @@ class GroupCoordinator:
    def gather(self,
               input_: torch.Tensor,
               dst: int = 0,
-               dim: int = -1) -> Optional[torch.Tensor]:
+               dim: int = -1) -> torch.Tensor | None:
        """
        NOTE: We assume that the input tensor is on the same device across
        all the ranks.
@@ -337,7 +338,7 @@ class GroupCoordinator:
                                    group=self.device_group)
        return input_

-    def broadcast_object(self, obj: Optional[Any] = None, src: int = 0):
+    def broadcast_object(self, obj: Any | None = None, src: int = 0):
        """Broadcast the input object.
        NOTE: `src` is the local rank of the source rank.
        """
@@ -362,9 +363,9 @@ class GroupCoordinator:
            return recv[0]

    def broadcast_object_list(self,
-                              obj_list: List[Any],
+                              obj_list: list[Any],
                              src: int = 0,
-                              group: Optional[ProcessGroup] = None):
+                              group: ProcessGroup | None = None):
        """Broadcast the input object list.
        NOTE: `src` is the local rank of the source rank.
        """
@@ -444,11 +445,11 @@ class GroupCoordinator:

    def broadcast_tensor_dict(
        self,
-        tensor_dict: Optional[Dict[str, Union[torch.Tensor, Any]]] = None,
+        tensor_dict: dict[str, torch.Tensor | Any] | None = None,
        src: int = 0,
-        group: Optional[ProcessGroup] = None,
-        metadata_group: Optional[ProcessGroup] = None
-    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+        group: ProcessGroup | None = None,
+        metadata_group: ProcessGroup | None = None
+    ) -> dict[str, torch.Tensor | Any] | None:
        """Broadcast the input tensor dictionary.
        NOTE: `src` is the local rank of the source rank.
        """
@@ -462,7 +463,7 @@ class GroupCoordinator:

        rank_in_group = self.rank_in_group
        if rank_in_group == src:
-            metadata_list: List[Tuple[Any, Any]] = []
+            metadata_list: list[tuple[Any, Any]] = []
            assert isinstance(
                tensor_dict,
                dict), (f"Expecting a dictionary, got {type(tensor_dict)}")
@@ -529,10 +530,10 @@ class GroupCoordinator:

    def send_tensor_dict(
        self,
-        tensor_dict: Dict[str, Union[torch.Tensor, Any]],
-        dst: Optional[int] = None,
+        tensor_dict: dict[str, torch.Tensor | Any],
+        dst: int | None = None,
        all_gather_group: Optional["GroupCoordinator"] = None,
-    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+    ) -> dict[str, torch.Tensor | Any] | None:
        """Send the input tensor dictionary.
        NOTE: `dst` is the local rank of the source rank.
        """
@@ -552,7 +553,7 @@ class GroupCoordinator:
            dst = (self.rank_in_group + 1) % self.world_size
        assert dst < self.world_size, f"Invalid dst rank ({dst})"

-        metadata_list: List[Tuple[Any, Any]] = []
+        metadata_list: list[tuple[Any, Any]] = []
        assert isinstance(
            tensor_dict,
            dict), f"Expecting a dictionary, got {type(tensor_dict)}"
@@ -583,9 +584,9 @@ class GroupCoordinator:

    def recv_tensor_dict(
        self,
-        src: Optional[int] = None,
+        src: int | None = None,
        all_gather_group: Optional["GroupCoordinator"] = None,
-    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
+    ) -> dict[str, torch.Tensor | Any] | None:
        """Recv the input tensor dictionary.
        NOTE: `src` is the local rank of the source rank.
        """
@@ -606,7 +607,7 @@ class GroupCoordinator:
        assert src < self.world_size, f"Invalid src rank ({src})"

        recv_metadata_list = self.recv_object(src=src)
-        tensor_dict: Dict[str, Any] = {}
+        tensor_dict: dict[str, Any] = {}
        for key, value in recv_metadata_list:
            if isinstance(value, TensorMetadata):
                tensor = torch.empty(value.size,
@@ -656,7 +657,7 @@ class GroupCoordinator:
        """
        torch.distributed.barrier(group=self.cpu_group)

-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
        """Sends a tensor to the destination rank in a non-blocking way"""
        """NOTE: `dst` is the local rank of the destination rank."""
        self.device_communicator.send(tensor, dst)
@@ -664,7 +665,7 @@ class GroupCoordinator:
    def recv(self,
             size: torch.Size,
             dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
+             src: int | None = None) -> torch.Tensor:
        """Receives a tensor from the source rank."""
        """NOTE: `src` is the local rank of the source rank."""
        return self.device_communicator.recv(size, dtype, src)
@@ -682,7 +683,7 @@ class GroupCoordinator:
            self.mq_broadcaster = None


-_WORLD: Optional[GroupCoordinator] = None
+_WORLD: GroupCoordinator | None = None


 def get_world_group() -> GroupCoordinator:
@@ -690,7 +691,7 @@ def get_world_group() -> GroupCoordinator:
    return _WORLD


-def init_world_group(ranks: List[int], local_rank: int,
+def init_world_group(ranks: list[int], local_rank: int,
                     backend: str) -> GroupCoordinator:
    return GroupCoordinator(
        group_ranks=[ranks],
@@ -702,11 +703,11 @@ def init_world_group(ranks: List[int], local_rank: int,


 def init_model_parallel_group(
-    group_ranks: List[List[int]],
+    group_ranks: list[list[int]],
    local_rank: int,
    backend: str,
    use_message_queue_broadcaster: bool = False,
-    group_name: Optional[str] = None,
+    group_name: str | None = None,
 ) -> GroupCoordinator:

    return GroupCoordinator(
@@ -719,7 +720,7 @@ def init_model_parallel_group(
    )


-_TP: Optional[GroupCoordinator] = None
+_TP: GroupCoordinator | None = None


 def get_tp_group() -> GroupCoordinator:
@@ -778,7 +779,7 @@ def init_distributed_environment(
            "world group already initialized with a different world size")


-_SP: Optional[GroupCoordinator] = None
+_SP: GroupCoordinator | None = None


 def get_sp_group() -> GroupCoordinator:
@@ -789,7 +790,7 @@ def get_sp_group() -> GroupCoordinator:
 def initialize_model_parallel(
    tensor_model_parallel_size: int = 1,
    sequence_model_parallel_size: int = 1,
-    backend: Optional[str] = None,
+    backend: str | None = None,
 ) -> None:
    """
    Initialize model parallel groups.
@@ -858,7 +859,7 @@ def get_sequence_model_parallel_rank() -> int:
 def ensure_model_parallel_initialized(
    tensor_model_parallel_size: int,
    sequence_model_parallel_size: int,
-    backend: Optional[str] = None,
+    backend: str | None = None,
 ) -> None:
    """Helper to initialize model parallel groups if they are not initialized,
    or ensure tensor-parallel, sequence-parallel sizes 
@@ -969,8 +970,8 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
            "torch._C._host_emptyCache() only available in Pytorch >=2.5")


-def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],
-                        source_rank: int = 0) -> List[bool]:
+def in_the_same_node_as(pg: ProcessGroup | StatelessProcessGroup,
+                        source_rank: int = 0) -> list[bool]:
    """
    This is a collective operation that returns if each rank is in the same node
    as the source rank. It tests if processes are attached to the same
@@ -1056,7 +1057,7 @@ def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],

 def initialize_tensor_parallel_group(
        tensor_model_parallel_size: int = 1,
-        backend: Optional[str] = None,
+        backend: str | None = None,
        group_name_suffix: str = "") -> GroupCoordinator:
    """Initialize a tensor parallel group for a specific model.
    
@@ -1120,7 +1121,7 @@ def initialize_tensor_parallel_group(

 def initialize_sequence_parallel_group(
        sequence_model_parallel_size: int = 1,
-        backend: Optional[str] = None,
+        backend: str | None = None,
        group_name_suffix: str = "") -> GroupCoordinator:
    """Initialize a sequence parallel group for a specific model.
    
@@ -9,7 +9,8 @@ import dataclasses
 import pickle
 import time
 from collections import deque
-from typing import Any, Deque, Dict, Optional, Sequence, Tuple
+from collections.abc import Sequence
+from typing import Any

 import torch
 from torch.distributed import TCPStore
@@ -72,15 +73,15 @@ class StatelessProcessGroup:
    data_expiration_seconds: int = 3600  # 1 hour

    # dst rank -> counter
-    send_dst_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+    send_dst_counter: dict[int, int] = dataclasses.field(default_factory=dict)
    # src rank -> counter
-    recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
+    recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict)
    broadcast_send_counter: int = 0
-    broadcast_recv_src_counter: Dict[int, int] = dataclasses.field(
+    broadcast_recv_src_counter: dict[int, int] = dataclasses.field(
        default_factory=dict)

    # A deque to store the data entries, with key and timestamp.
-    entries: Deque[Tuple[str, float]] = dataclasses.field(default_factory=deque)
+    entries: deque[tuple[str, float]] = dataclasses.field(default_factory=deque)

    def __post_init__(self):
        assert self.rank < self.world_size
@@ -94,14 +95,14 @@ class StatelessProcessGroup:
        key = f"send_to/{dst}/{self.send_dst_counter[dst]}"
        self.store.set(key, pickle.dumps(obj))
        self.send_dst_counter[dst] += 1
-        self.entries.append((key, time.time()))
+        self.entries.append((key, time.perf_counter()))

    def expire_data(self) -> None:
        """Expire data that is older than `data_expiration_seconds` seconds."""
        while self.entries:
            # check the oldest entry
            key, timestamp = self.entries[0]
-            if time.time() - timestamp > self.data_expiration_seconds:
+            if time.perf_counter() - timestamp > self.data_expiration_seconds:
                self.store.delete_key(key)
                self.entries.popleft()
            else:
@@ -114,7 +115,7 @@ class StatelessProcessGroup:
        self.recv_src_counter[src] += 1
        return obj

-    def broadcast_obj(self, obj: Optional[Any], src: int) -> Any:
+    def broadcast_obj(self, obj: Any | None, src: int) -> Any:
        """Broadcast an object from a source rank to all other ranks.
        It does not clean up after all ranks have received the object.
        Use it for limited times, e.g., for initialization.
@@ -125,7 +126,7 @@ class StatelessProcessGroup:
                   f"{self.broadcast_send_counter}")
            self.store.set(key, pickle.dumps(obj))
            self.broadcast_send_counter += 1
-            self.entries.append((key, time.time()))
+            self.entries.append((key, time.perf_counter()))
            return obj
        else:
            key = (f"broadcast_from/{src}/"
@@ -2,13 +2,20 @@
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py

 import argparse
-from typing import List, cast
+import dataclasses
+import os
+from typing import Any, cast

-from fastvideo.v1.entrypoints.cli import utils
+from fastvideo import PipelineConfig, VideoGenerator
+from fastvideo.v1.configs.sample.base import SamplingParam
 from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
+from fastvideo.v1.entrypoints.cli.utils import RaiseNotImplementedAction
 from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import FlexibleArgumentParser

+logger = init_logger(__name__)
+

 class GenerateSubcommand(CLISubcommand):
    """The `generate` subcommand for the FastVideo CLI"""
@@ -16,45 +23,80 @@ class GenerateSubcommand(CLISubcommand):
    def __init__(self) -> None:
        self.name = "generate"
        super().__init__()
+        self.init_arg_names = self._get_init_arg_names()
+        self.generation_arg_names = self._get_generation_arg_names()
+
+    def _get_init_arg_names(self) -> list[str]:
+        """Get names of arguments for VideoGenerator initialization"""
+        return ["num_gpus", "tp_size", "sp_size", "model_path"]
+
+    def _get_generation_arg_names(self) -> list[str]:
+        """Get names of arguments for generate_video method"""
+        return [field.name for field in dataclasses.fields(SamplingParam)]

    def cmd(self, args: argparse.Namespace) -> None:
-        excluded_args = [
-            'subparser', 'config', 'num_gpus', 'master_port',
-            'dispatch_function'
-        ]
+        excluded_args = ['subparser', 'config', 'dispatch_function']

-        # Create a filtered dictionary of arguments
-        filtered_args = {
+        FastVideoArgs.from_cli_args(args)
+
+        provided_args = {}
+        for k, v in vars(args).items():
+            if (k not in excluded_args and v is not None
+                    and hasattr(args, '_provided') and k in args._provided):
+                provided_args[k] = v
+
+        if 'model_path' in vars(args) and args.model_path is not None:
+            provided_args['model_path'] = args.model_path
+
+        if 'prompt' in vars(args) and args.prompt is not None:
+            provided_args['prompt'] = args.prompt
+
+        merged_args = {**provided_args}
+
+        logger.info('CLI Args: %s', merged_args)
+
+        if 'model_path' not in merged_args or not merged_args['model_path']:
+            raise ValueError(
+                "model_path must be provided either in config file or via --model-path"
+            )
+
+        if 'prompt' not in merged_args or not merged_args['prompt']:
+            raise ValueError(
+                "prompt must be provided either in config file or via --prompt")
+
+        init_args = {
            k: v
-            for k, v in vars(args).items()
-            if k not in excluded_args and v is not None
+            for k, v in merged_args.items() if k in self.init_arg_names
+        }
+        generation_args = {
+            k: v
+            for k, v in merged_args.items() if k in self.generation_arg_names
        }

-        main_args = []
+        pipeline_config = PipelineConfig.from_pretrained(
+            merged_args['model_path'])

-        for key, value in filtered_args.items():
-            # Convert underscores to dashes in argument names
-            arg_name = f"--{key.replace('_', '-')}"
+        update_config_from_args(pipeline_config.dit_config, merged_args,
+                                "dit_config")
+        update_config_from_args(pipeline_config.vae_config, merged_args,
+                                "vae_config")
+        update_config_from_args(pipeline_config, merged_args)

-            # Handle boolean flags
-            if isinstance(value, bool):
-                if value:
-                    main_args.append(arg_name)
-            else:
-                main_args.append(arg_name)
-                main_args.append(str(value))
+        model_path = init_args.pop('model_path')
+        prompt = generation_args.pop('prompt')

-        utils.launch_distributed(args.num_gpus,
-                                 main_args,
-                                 master_port=args.master_port)
+        generator = VideoGenerator.from_pretrained(
+            model_path=model_path, **init_args, pipeline_config=pipeline_config)
+
+        generator.generate_video(prompt=prompt, **generation_args)

    def validate(self, args: argparse.Namespace) -> None:
+        """Validate the arguments for this command"""
        if args.num_gpus is not None and args.num_gpus <= 0:
            raise ValueError("Number of gpus must be positive")

-        if args.master_port is not None and (args.master_port < 1024
-                                             or args.master_port > 65535):
-            raise ValueError("Master port must be between 1024 and 65535")
+        if args.config and not os.path.exists(args.config):
+            raise ValueError(f"Config file not found: {args.config}")

    def subparser_init(
            self,
@@ -63,7 +105,7 @@ class GenerateSubcommand(CLISubcommand):
            "generate",
            help="Run inference on a model",
            usage=
-            "fastvideo generate --model-path MODEL_PATH_OR_ID --prompt PROMPT [OPTIONS]"
+            "fastvideo generate (--model-path MODEL_PATH_OR_ID --prompt PROMPT) | --config CONFIG_FILE [OPTIONS]"
        )

        generate_parser.add_argument(
@@ -71,17 +113,53 @@ class GenerateSubcommand(CLISubcommand):
            type=str,
            default='',
            required=False,
-            help="Read CLI options from a config YAML file.")
-
-        generate_parser.add_argument("--master-port",
-                                     type=int,
-                                     default=None,
-                                     help="Port for the master process")
+            help=
+            "Read CLI options from a config JSON or YAML file. If provided, --model-path and --prompt are optional."
+        )

        generate_parser = FastVideoArgs.add_cli_args(generate_parser)
+        generate_parser = SamplingParam.add_cli_args(generate_parser)
+
+        generate_parser.add_argument(
+            "--text-encoder-configs",
+            action=RaiseNotImplementedAction,
+            help=
+            "JSON array of text encoder configurations (NOT YET IMPLEMENTED)",
+        )

        return cast(FlexibleArgumentParser, generate_parser)


-def cmd_init() -> List[CLISubcommand]:
+def cmd_init() -> list[CLISubcommand]:
    return [GenerateSubcommand()]
+
+
+def update_config_from_args(config: Any,
+                            args_dict: dict[str, Any],
+                            prefix: str | None = None) -> None:
+    """
+    Update configuration object from arguments dictionary.
+    
+    Args:
+        config: The configuration object to update
+        args_dict: Dictionary containing arguments
+        prefix: Prefix for the configuration parameters in the args_dict.
+               If None, assumes direct attribute mapping without prefix.
+    """
+    # Handle top-level attributes (no prefix)
+    if prefix is None:
+        for key, value in args_dict.items():
+            if hasattr(config, key) and value is not None:
+                if key == "text_encoder_precisions" and isinstance(value, list):
+                    setattr(config, key, tuple(value))
+                else:
+                    setattr(config, key, value)
+        return
+
+    # Handle nested attributes with prefix
+    prefix_with_dot = f"{prefix}."
+    for key, value in args_dict.items():
+        if key.startswith(prefix_with_dot) and value is not None:
+            attr_name = key[len(prefix_with_dot):]
+            if hasattr(config, attr_name):
+                setattr(config, attr_name, value)
@@ -1,14 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/main.py

-from typing import List
-
 from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
 from fastvideo.v1.entrypoints.cli.generate import cmd_init as generate_cmd_init
 from fastvideo.v1.utils import FlexibleArgumentParser


-def cmd_init() -> List[CLISubcommand]:
+def cmd_init() -> list[CLISubcommand]:
    """Initialize all commands from separate modules"""
    commands = []
    commands.extend(generate_cmd_init())
@@ -1,18 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0

+import argparse
 import os
 import subprocess
 import sys
-from typing import List, Optional

 from fastvideo.v1.logger import init_logger

 logger = init_logger(__name__)


+class RaiseNotImplementedAction(argparse.Action):
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        raise NotImplementedError(
+            f"The {option_string} option is not yet implemented")
+
+
 def launch_distributed(num_gpus: int,
-                       args: List[str],
-                       master_port: Optional[int] = None) -> int:
+                       args: list[str],
+                       master_port: int | None = None) -> int:
    """
    Launch a distributed job with the given arguments
    
@@ -7,9 +7,10 @@ diffusion models.
 """

 import gc
+import math
 import os
 import time
-from typing import Any, Dict, List, Optional, Union
+from typing import Any

 import imageio
 import numpy as np
@@ -52,11 +53,9 @@ class VideoGenerator:
    @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,
+                        device: str | None = None,
+                        torch_dtype: torch.dtype | None = None,
+                        pipeline_config: str | PipelineConfig | None = None,
                        **kwargs) -> "VideoGenerator":
        """
        Create a video generator from a pretrained model.
@@ -72,7 +71,6 @@ class VideoGenerator:

        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):
@@ -128,9 +126,9 @@ class VideoGenerator:
    def generate_video(
        self,
        prompt: str,
-        sampling_param: Optional[SamplingParam] = None,
+        sampling_param: SamplingParam | None = None,
        **kwargs,
-    ) -> Union[Dict[str, Any], List[np.ndarray]]:
+    ) -> dict[str, Any] | list[np.ndarray]:
        """
        Generate a video based on the given prompt.
        
@@ -181,12 +179,46 @@ class VideoGenerator:
                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}"
-            )
+        temporal_scale_factor = fastvideo_args.vae_config.arch_config.temporal_compression_ratio
+        num_frames = sampling_param.num_frames
+        num_gpus = fastvideo_args.num_gpus
+        use_temporal_scaling_frames = fastvideo_args.vae_config.use_temporal_scaling_frames
+
+        # Adjust number of frames based on number of GPUs
+        if use_temporal_scaling_frames:
+            orig_latent_num_frames = (num_frames -
+                                      1) // temporal_scale_factor + 1
+        else:  # stepvideo only
+            orig_latent_num_frames = sampling_param.num_frames // 17 * 3
+
+        if orig_latent_num_frames % fastvideo_args.num_gpus != 0:
+            # Adjust latent frames to be divisible by number of GPUs
+            if sampling_param.num_frames_round_down:
+                # Ensure we have at least 1 batch per GPU
+                new_latent_num_frames = max(
+                    1, (orig_latent_num_frames // num_gpus)) * num_gpus
+            else:
+                new_latent_num_frames = math.ceil(
+                    orig_latent_num_frames / num_gpus) * num_gpus
+
+            if use_temporal_scaling_frames:
+                # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor
+                new_num_frames = (new_latent_num_frames -
+                                  1) * temporal_scale_factor + 1
+            else:  # stepvideo only
+                # Find the least common multiple of 3 and num_gpus
+                divisor = math.lcm(3, num_gpus)
+                # Round up to the nearest multiple of this LCM
+                new_latent_num_frames = (
+                    (new_latent_num_frames + divisor - 1) // divisor) * divisor
+                # Convert back to actual frames using the StepVideo formula
+                new_num_frames = new_latent_num_frames // 3 * 17
+
+            logger.info(
+                "Adjusting number of frames from %s to %s based on number of GPUs (%s)",
+                sampling_param.num_frames, new_num_frames,
+                fastvideo_args.num_gpus)
+            sampling_param.num_frames = new_num_frames

        # Calculate sizes
        target_height = align_to(sampling_param.height, 16)
@@ -225,11 +257,11 @@ class VideoGenerator:
        )

        # Run inference
-        start_time = time.time()
+        start_time = time.perf_counter()
        output_batch = self.executor.execute_forward(batch, fastvideo_args)
        samples = output_batch

-        gen_time = time.time() - start_time
+        gen_time = time.perf_counter() - start_time
        logger.info("Generated successfully in %.2f seconds", gen_time)

        # Process outputs
@@ -2,28 +2,29 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py

 import os
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any

 if TYPE_CHECKING:
    FASTVIDEO_RINGBUFFER_WARNING_INTERVAL: int = 60
-    FASTVIDEO_NCCL_SO_PATH: Optional[str] = None
-    LD_LIBRARY_PATH: Optional[str] = None
+    FASTVIDEO_NCCL_SO_PATH: str | None = None
+    LD_LIBRARY_PATH: str | None = None
    LOCAL_RANK: int = 0
-    CUDA_VISIBLE_DEVICES: Optional[str] = None
+    CUDA_VISIBLE_DEVICES: str | None = None
    FASTVIDEO_CACHE_ROOT: str = os.path.expanduser("~/.cache/fastvideo")
    FASTVIDEO_CONFIG_ROOT: str = os.path.expanduser("~/.config/fastvideo")
    FASTVIDEO_CONFIGURE_LOGGING: int = 1
    FASTVIDEO_LOGGING_LEVEL: str = "INFO"
    FASTVIDEO_LOGGING_PREFIX: str = ""
-    FASTVIDEO_LOGGING_CONFIG_PATH: Optional[str] = None
+    FASTVIDEO_LOGGING_CONFIG_PATH: str | None = None
    FASTVIDEO_TRACE_FUNCTION: int = 0
-    FASTVIDEO_ATTENTION_BACKEND: Optional[str] = None
-    FASTVIDEO_ATTENTION_CONFIG: Optional[str] = None
+    FASTVIDEO_ATTENTION_BACKEND: str | None = None
+    FASTVIDEO_ATTENTION_CONFIG: str | None = None
    FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "fork"
    FASTVIDEO_TARGET_DEVICE: str = "cuda"
-    MAX_JOBS: Optional[str] = None
-    NVCC_THREADS: Optional[str] = None
-    CMAKE_BUILD_TYPE: Optional[str] = None
+    MAX_JOBS: str | None = None
+    NVCC_THREADS: str | None = None
+    CMAKE_BUILD_TYPE: str | None = None
    VERBOSE: bool = False
    FASTVIDEO_SERVER_DEV_MODE: bool = False

@@ -42,7 +43,7 @@ def get_default_config_root() -> str:
    )


-def maybe_convert_int(value: Optional[str]) -> Optional[int]:
+def maybe_convert_int(value: str | None) -> int | None:
    if value is None:
        return None
    return int(value)
@@ -53,7 +54,7 @@ def maybe_convert_int(value: Optional[str]) -> Optional[int]:

 # begin-env-vars-definition

-environment_variables: Dict[str, Callable[[], Any]] = {
+environment_variables: dict[str, Callable[[], Any]] = {

    # ================== Installation Time Env Vars ==================

@@ -4,13 +4,14 @@

 import argparse
 import dataclasses
+from collections.abc import Callable
 from contextlib import contextmanager
 from dataclasses import field
-from typing import Any, Callable, List, Optional, Tuple
+from typing import Any

 from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import FlexibleArgumentParser
+from fastvideo.v1.utils import FlexibleArgumentParser, StoreBoolean

 logger = init_logger(__name__)

@@ -38,17 +39,17 @@ class FastVideoArgs:

    # HuggingFace specific parameters
    trust_remote_code: bool = False
-    revision: Optional[str] = None
+    revision: str | None = None

    # Parallelism
    num_gpus: int = 1
-    tp_size: Optional[int] = None
-    sp_size: Optional[int] = None
-    dist_timeout: Optional[int] = None  # timeout for torch.distributed
+    tp_size: int | None = None
+    sp_size: int | None = None
+    dist_timeout: int | None = None  # timeout for torch.distributed

    # Video generation parameters
    embedded_cfg_scale: float = 6.0
-    flow_shift: Optional[float] = None
+    flow_shift: float | None = None

    output_type: str = "pil"

@@ -68,27 +69,36 @@ class FastVideoArgs:
    image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)

    # Text encoder configuration
-    text_encoder_precisions: Tuple[str, ...] = field(
-        default_factory=lambda: ("fp16", ))
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+    DEFAULT_TEXT_ENCODER_PRECISIONS = (
+        "fp16",
+        "fp16",
+    )
+    text_encoder_precisions: tuple[str, ...] = field(
+        default_factory=lambda: FastVideoArgs.DEFAULT_TEXT_ENCODER_PRECISIONS)
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
        default_factory=lambda: (EncoderConfig(), ))
-    preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
        default_factory=lambda: (preprocess_text, ))
-    postprocess_text_funcs: Tuple[Callable[[Any], Any], ...] = field(
+    postprocess_text_funcs: tuple[Callable[[Any], Any], ...] = field(
        default_factory=lambda: (postprocess_text, ))

    # STA (Spatial-Temporal Attention) parameters
-    mask_strategy_file_path: Optional[str] = None
+    mask_strategy_file_path: str | None = None
    enable_torch_compile: bool = False

    use_cpu_offload: bool = False
    disable_autocast: bool = False

+    # StepVideo specific parameters
+    pos_magic: str | None = None
+    neg_magic: str | None = None
+    timesteps_scale: bool | None = None
+
    # Logging
    log_level: str = "info"

    # Inference parameters
-    device_str: Optional[str] = None
+    device_str: str | None = None
    device = None

    def __post_init__(self):
@@ -100,7 +110,6 @@ class FastVideoArgs:
        parser.add_argument(
            "--model-path",
            type=str,
-            required=True,
            help=
            "The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
        )
@@ -127,7 +136,7 @@ class FastVideoArgs:
        # HuggingFace specific parameters
        parser.add_argument(
            "--trust-remote-code",
-            action="store_true",
+            action=StoreBoolean,
            default=FastVideoArgs.trust_remote_code,
            help="Trust remote code when loading HuggingFace models",
        )
@@ -206,21 +215,21 @@ class FastVideoArgs:
        )
        parser.add_argument(
            "--vae-tiling",
-            action="store_true",
+            action=StoreBoolean,
            default=FastVideoArgs.vae_tiling,
            help="Enable VAE tiling",
        )
        parser.add_argument(
            "--vae-sp",
-            action="store_true",
+            action=StoreBoolean,
            help="Enable VAE spatial parallelism",
        )

        parser.add_argument(
-            "--text-encoder-precision",
+            "--text-encoder-precisions",
            nargs="+",
            type=str,
-            default=FastVideoArgs.text_encoder_precisions,
+            default=FastVideoArgs.DEFAULT_TEXT_ENCODER_PRECISIONS,
            choices=["fp32", "fp16", "bf16"],
            help="Precision for each text encoder",
        )
@@ -242,23 +251,42 @@ class FastVideoArgs:
        )
        parser.add_argument(
            "--enable-torch-compile",
-            action="store_true",
+            action=StoreBoolean,
            help=
            "Use torch.compile for speeding up STA inference without teacache",
        )

        parser.add_argument(
            "--use-cpu-offload",
-            action="store_true",
+            action=StoreBoolean,
            help="Use CPU offload for the model load",
        )
        parser.add_argument(
            "--disable-autocast",
-            action="store_true",
+            action=StoreBoolean,
            help=
            "Disable autocast for denoising loop and vae decoding in pipeline sampling",
        )

+        parser.add_argument(
+            "--pos_magic",
+            type=str,
+            default=FastVideoArgs.pos_magic,
+            help="Positive magic prompt for sampling",
+        )
+        parser.add_argument(
+            "--neg_magic",
+            type=str,
+            default=FastVideoArgs.neg_magic,
+            help="Negative magic prompt for sampling",
+        )
+        parser.add_argument(
+            "--timesteps_scale",
+            type=bool,
+            default=FastVideoArgs.timesteps_scale,
+            help="Bool for applying scheduler scale in set_timesteps",
+        )
+
        # Logging
        parser.add_argument(
            "--log-level",
@@ -267,6 +295,14 @@ class FastVideoArgs:
            help="The logging level of all loggers.",
        )

+        # Add VAE configuration arguments
+        from fastvideo.v1.configs.models.vaes.base import VAEConfig
+        VAEConfig.add_cli_args(parser)
+
+        # Add DiT configuration arguments
+        from fastvideo.v1.configs.models.dits.base import DiTConfig
+        DiTConfig.add_cli_args(parser)
+
        return parser

    @classmethod
@@ -341,7 +377,7 @@ class FastVideoArgs:
 _current_fastvideo_args = None


-def prepare_fastvideo_args(argv: List[str]) -> FastVideoArgs:
+def prepare_fastvideo_args(argv: list[str]) -> FastVideoArgs:
    """
    Prepare the inference arguments from the command line arguments.

@@ -5,7 +5,7 @@ import time
 from collections import defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING

 import torch

@@ -37,10 +37,10 @@ class ForwardContext:
    # attn_layers: Dict[str, Any]
    # TODO: extend to support per-layer dynamic forward context
    attn_metadata: "AttentionMetadata"  # set dynamically for each forward pass
-    forward_batch: Optional[ForwardBatch] = None
+    forward_batch: ForwardBatch | None = None


-_forward_context: Optional[ForwardContext] = None
+_forward_context: ForwardContext | None = None


 def get_forward_context() -> ForwardContext:
@@ -55,8 +55,8 @@ def get_forward_context() -> ForwardContext:
@contextmanager
 def set_forward_context(current_timestep,
                        attn_metadata,
-                        forward_batch: Optional[ForwardBatch] = None,
-                        fastvideo_args: Optional[FastVideoArgs] = None):
+                        forward_batch: ForwardBatch | None = None,
+                        fastvideo_args: FastVideoArgs | None = 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.
@@ -8,7 +8,7 @@ This module provides classes and functions for running inference with diffusion
 """

 import time
-from typing import Any, Dict
+from typing import Any

 import torch

@@ -83,7 +83,7 @@ class InferenceEngine:
        self,
        prompt: str,
        fastvideo_args: FastVideoArgs,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
        """
        Run inference with the pipeline.
        
@@ -96,7 +96,7 @@ class InferenceEngine:
        Returns:
            A dictionary containing the generated videos and metadata.
        """
-        out_dict: Dict[str, Any] = dict()
+        out_dict: dict[str, Any] = dict()

        num_videos_per_prompt = fastvideo_args.num_videos
        seed = fastvideo_args.seed
@@ -191,7 +191,7 @@ class InferenceEngine:
        # ========================================================================
        # Pipeline inference
        # ========================================================================
-        start_time = time.time()
+        start_time = time.perf_counter()
        samples = self.pipeline.forward(
            batch=batch,
            fastvideo_args=fastvideo_args,
@@ -201,7 +201,7 @@ class InferenceEngine:
        out_dict["samples"] = samples
        out_dict["prompts"] = prompt

-        gen_time = time.time() - start_time
+        gen_time = time.perf_counter() - start_time
        logger.info("Success, time: %s", gen_time)

        return out_dict
@@ -9,7 +9,6 @@ import torch.nn.functional as F

 # TODO (will): remove this dependency
 from fastvideo.v1.layers.custom_op import CustomOp
-from fastvideo.v1.platforms import current_platform


@CustomOp.register("silu_and_mul")
@@ -25,21 +24,12 @@ class SiluAndMul(CustomOp):

    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

    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
        """PyTorch-native implementation equivalent to forward()."""
        d = x.shape[-1] // 2
        return F.silu(x[..., :d]) * x[..., d:]

-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        output_shape = (x.shape[:-1] + (d, ))
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        self.op(out, x)
-        return out
-

@CustomOp.register("gelu_and_mul")
 class GeluAndMul(CustomOp):
@@ -57,24 +47,12 @@ class GeluAndMul(CustomOp):
        self.approximate = approximate
        if approximate not in ("none", "tanh"):
            raise ValueError(f"Unknown approximate mode: {approximate}")
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            if approximate == "none":
-                self.op = torch.ops._C.gelu_and_mul
-            elif approximate == "tanh":
-                self.op = torch.ops._C.gelu_tanh_and_mul

    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
        """PyTorch-native implementation equivalent to forward()."""
        d = x.shape[-1] // 2
        return F.gelu(x[..., :d], approximate=self.approximate) * x[..., d:]

-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        d = x.shape[-1] // 2
-        output_shape = (x.shape[:-1] + (d, ))
-        out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
-        self.op(out, x)
-        return out
-
    def extra_repr(self) -> str:
        return f'approximate={repr(self.approximate)}'

@@ -84,8 +62,6 @@ class NewGELU(CustomOp):

    def __init__(self):
        super().__init__()
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            self.op = torch.ops._C.gelu_new

    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
        """PyTorch-native implementation equivalent to forward()."""
@@ -93,32 +69,17 @@ class NewGELU(CustomOp):
        return 0.5 * x * (1.0 + torch.tanh(c *
                                           (x + 0.044715 * torch.pow(x, 3.0))))

-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        out = torch.empty_like(x)
-        self.op(out, x)
-        return out
-
-    def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
-        return self.op(x)
-

@CustomOp.register("quick_gelu")
 class QuickGELU(CustomOp):
    # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90
    def __init__(self):
        super().__init__()
-        if current_platform.is_cuda_alike() or current_platform.is_cpu():
-            self.op = torch.ops._C.gelu_quick

    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
        """PyTorch-native implementation equivalent to forward()."""
        return x * torch.sigmoid(1.702 * x)

-    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
-        out = torch.empty_like(x)
-        self.op(out, x)
-        return out
-

 _ACTIVATION_REGISTRY = {
    "gelu": nn.GELU,
@@ -1,7 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/custom_op.py

-from typing import Any, Callable, Dict, Type
+from collections.abc import Callable
+from typing import Any

 import torch.nn as nn

@@ -50,6 +51,10 @@ class CustomOp(nn.Module):
        return self.forward_native(*args, **kwargs)

    def dispatch_forward(self) -> Callable:
+        # FIXME(will): for now, we always use the native implementation, since
+        # forward_cuda is using vllm's custom ops and it doesn't support
+        # backwards. We should add our own custom ops that support backwards.
+        return self.forward_native
        # NOTE(woosuk): Here we assume that vLLM was built for only one
        # specific backend. Currently, we do not support dynamic dispatching.
        enabled = self.enabled()
@@ -77,7 +82,7 @@ class CustomOp(nn.Module):
    # Examples:
    # - MyOp.enabled()
    # - op_registry["my_op"].enabled()
-    op_registry: Dict[str, Type['CustomOp']] = {}
+    op_registry: dict[str, type['CustomOp']] = {}

    # Decorator to register custom ops.
    @classmethod
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/layernorm.py
 """Custom normalization layers."""
-from typing import Optional, Tuple, Union

 import torch
 import torch.nn as nn
@@ -22,7 +21,7 @@ class RMSNorm(CustomOp):
        hidden_size: int,
        eps: float = 1e-6,
        dtype: torch.dtype = torch.float32,
-        var_hidden_size: Optional[int] = None,
+        var_hidden_size: int | None = None,
        has_weight: bool = True,
    ) -> None:
        super().__init__()
@@ -40,8 +39,8 @@ class RMSNorm(CustomOp):
    def forward_native(
        self,
        x: torch.Tensor,
-        residual: Optional[torch.Tensor] = None,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        residual: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        """PyTorch-native implementation equivalent to forward()."""
        orig_dtype = x.dtype
        x = x.to(torch.float32)
@@ -75,33 +74,6 @@ class RMSNorm(CustomOp):
        else:
            return x, residual

-    def forward_cuda(
-        self,
-        x: torch.Tensor,
-        residual: Optional[torch.Tensor] = None,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        if self.variance_size_override is not None:
-            return self.forward_native(x, residual)
-
-        from vllm import _custom_ops as ops
-
-        if residual is not None:
-            ops.fused_add_rms_norm(
-                x,
-                residual,
-                self.weight.data,
-                self.variance_epsilon,
-            )
-            return x, residual
-        out = torch.empty_like(x)
-        ops.rms_norm(
-            out,
-            x,
-            self.weight.data,
-            self.variance_epsilon,
-        )
-        return out
-
    def extra_repr(self) -> str:
        s = f"hidden_size={self.weight.data.size(0)}"
        s += f", eps={self.variance_epsilon}"
@@ -157,7 +129,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module):

    def forward(self, residual: torch.Tensor, x: torch.Tensor,
                gate: torch.Tensor, shift: torch.Tensor,
-                scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+                scale: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Apply gated residual connection, followed by layernorm and 
        scale/shift in a single fused operation.
@@ -173,7 +145,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
        # Apply normalization
        normalized = self.norm(residual_output)
        # Apply scale and shift
-        modulated = normalized * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        modulated = normalized * (1.0 + scale) + shift
        return modulated, residual_output


@@ -209,4 +181,4 @@ class LayerNormScaleShift(nn.Module):
                scale: torch.Tensor) -> torch.Tensor:
        """Apply ln followed by scale and shift in a single fused operation."""
        normalized = self.norm(x)
-        return normalized * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        return normalized * (1.0 + scale) + shift
@@ -2,21 +2,18 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/linear.py

 from abc import abstractmethod
-from typing import Optional, Union

 import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter
-# TODO(will): remove this import by copying the definition from vLLM then
-# manually import each quantization method we want to use. Refer to SGLang
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig, QuantizeMethodBase)

 from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank,
                                      get_tensor_model_parallel_world_size,
                                      split_tensor_along_last_dim,
                                      tensor_model_parallel_all_gather,
                                      tensor_model_parallel_all_reduce)
+from fastvideo.v1.layers.quantization.base_config import (QuantizationConfig,
+                                                          QuantizeMethodBase)
 from fastvideo.v1.logger import init_logger
 # yapf: disable
 from fastvideo.v1.models.parameter import (BasevLLMParameter,
@@ -42,7 +39,7 @@ WEIGHT_LOADER_V2_SUPPORTED = [

 def adjust_scalar_to_fused_array(
        param: torch.Tensor, loaded_weight: torch.Tensor,
-        shard_id: Union[str, int]) -> tuple[torch.Tensor, torch.Tensor]:
+        shard_id: str | int) -> tuple[torch.Tensor, torch.Tensor]:
    """For fused modules (QKV and MLP) we have an array of length
    N that holds 1 scale for each "logical" matrix. So the param
    is an array of length N. The loaded_weight corresponds to 
@@ -93,7 +90,7 @@ class LinearMethodBase(QuantizeMethodBase):
    def apply(self,
              layer: torch.nn.Module,
              x: torch.Tensor,
-              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+              bias: torch.Tensor | None = None) -> torch.Tensor:
        """Apply the weights in layer to the input tensor.
        Expects create_weights to have been called before on the layer."""
        raise NotImplementedError
@@ -118,7 +115,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
    def apply(self,
              layer: torch.nn.Module,
              x: torch.Tensor,
-              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+              bias: torch.Tensor | None = None) -> torch.Tensor:

        return F.linear(x, layer.weight, bias)

@@ -140,8 +137,8 @@ class LinearBase(torch.nn.Module):
        input_size: int,
        output_size: int,
        skip_bias_add: bool = False,
-        params_dtype: Optional[torch.dtype] = None,
-        quant_config: Optional[QuantizationConfig] = None,
+        params_dtype: torch.dtype | None = None,
+        quant_config: QuantizationConfig | None = None,
        prefix: str = "",
    ):
        super().__init__()
@@ -154,14 +151,13 @@ class LinearBase(torch.nn.Module):
            params_dtype = torch.get_default_dtype()
        self.params_dtype = params_dtype
        if quant_config is None:
-            self.quant_method: Optional[
-                QuantizeMethodBase] = UnquantizedLinearMethod()
+            self.quant_method: QuantizeMethodBase | None = UnquantizedLinearMethod(
+            )
        else:
            self.quant_method = quant_config.get_quant_method(self,
                                                              prefix=prefix)

-    def forward(self,
-                x: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]:
        raise NotImplementedError


@@ -184,8 +180,8 @@ class ReplicatedLinear(LinearBase):
                 output_size: int,
                 bias: bool = True,
                 skip_bias_add: bool = False,
-                 params_dtype: Optional[torch.dtype] = None,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 params_dtype: torch.dtype | None = None,
+                 quant_config: QuantizationConfig | None = None,
                 prefix: str = ""):
        super().__init__(input_size,
                         output_size,
@@ -225,8 +221,7 @@ class ReplicatedLinear(LinearBase):
            f"to a parameter of size {param.size()}")
        param.data.copy_(loaded_weight)

-    def forward(self,
-                x: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]:
        bias = self.bias if not self.skip_bias_add else None
        assert self.quant_method is not None
        output = self.quant_method.apply(self, x, bias)
@@ -270,9 +265,9 @@ class ColumnParallelLinear(LinearBase):
                 bias: bool = True,
                 gather_output: bool = False,
                 skip_bias_add: bool = False,
-                 params_dtype: Optional[torch.dtype] = None,
-                 quant_config: Optional[QuantizationConfig] = None,
-                 output_sizes: Optional[list[int]] = None,
+                 params_dtype: torch.dtype | None = None,
+                 quant_config: QuantizationConfig | None = None,
+                 output_sizes: list[int] | None = None,
                 prefix: str = ""):
        # Divide the weight matrix along the last dimension.
        self.tp_size = get_tensor_model_parallel_world_size()
@@ -347,9 +342,8 @@ class ColumnParallelLinear(LinearBase):
            loaded_weight = loaded_weight.reshape(1)
        param.load_column_parallel_weight(loaded_weight=loaded_weight)

-    def forward(
-            self,
-            input_: torch.Tensor) -> tuple[torch.Tensor, Optional[Parameter]]:
+    def forward(self,
+                input_: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]:
        bias = self.bias if not self.skip_bias_add else None

        # Matrix multiply.
@@ -401,8 +395,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
                 bias: bool = True,
                 gather_output: bool = False,
                 skip_bias_add: bool = False,
-                 params_dtype: Optional[torch.dtype] = None,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 params_dtype: torch.dtype | None = None,
+                 quant_config: QuantizationConfig | None = None,
                 prefix: str = ""):
        self.output_sizes = output_sizes
        tp_size = get_tensor_model_parallel_world_size()
@@ -419,7 +413,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
    def weight_loader(self,
                      param: Parameter,
                      loaded_weight: torch.Tensor,
-                      loaded_shard_id: Optional[int] = None) -> None:
+                      loaded_shard_id: int | None = None) -> None:

        param_data = param.data
        output_dim = getattr(param, "output_dim", None)
@@ -512,10 +506,8 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
            # Special case for Quantization.
            # If quantized, we need to adjust the offset and size to account
            # for the packing.
-            if isinstance(
-                    param,
-                (PackedColumnParameter,
-                 PackedvLLMParameter)) and param.packed_dim == param.output_dim:
+            if isinstance(param, PackedColumnParameter | PackedvLLMParameter
+                          ) and param.packed_dim == param.output_dim:
                shard_size, shard_offset = \
                    param.adjust_shard_indexes_for_packing(
                    shard_size=shard_size, shard_offset=shard_offset)
@@ -527,7 +519,7 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
    def weight_loader_v2(self,
                         param: BasevLLMParameter,
                         loaded_weight: torch.Tensor,
-                         loaded_shard_id: Optional[int] = None) -> None:
+                         loaded_shard_id: int | None = None) -> None:
        if loaded_shard_id is None:
            if isinstance(param, PerTensorScaleParameter):
                param.load_merged_column_weight(loaded_weight=loaded_weight,
@@ -545,19 +537,21 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
        tp_size = get_tensor_model_parallel_world_size()

        if isinstance(param, BlockQuantScaleParameter):
-            from vllm.model_executor.layers.quantization.fp8 import (
-                Fp8LinearMethod, Fp8MoEMethod)
-            assert self.quant_method is not None
-            assert isinstance(self.quant_method,
-                              (Fp8LinearMethod, Fp8MoEMethod))
-            weight_block_size = self.quant_method.quant_config.weight_block_size
-            assert weight_block_size is not None
-            block_n, _ = weight_block_size[0], weight_block_size[1]
-            shard_offset = (
-                (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //
-                block_n) // tp_size
-            shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //
-                          block_n // tp_size)
+            raise NotImplementedError("FP8 is not implemented yet")
+            # FIXME(will): add fp8 support
+            # from vllm.model_executor.layers.quantization.fp8 import (
+            #     Fp8LinearMethod, Fp8MoEMethod)
+            # assert self.quant_method is not None
+            # assert isinstance(self.quant_method,
+            #                   (Fp8LinearMethod, Fp8MoEMethod))
+            # weight_block_size = self.quant_method.quant_config.weight_block_size
+            # assert weight_block_size is not None
+            # block_n, _ = weight_block_size[0], weight_block_size[1]
+            # shard_offset = (
+            #     (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //
+            #     block_n) // tp_size
+            # shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //
+            #               block_n // tp_size)
        else:
            shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size
            shard_size = self.output_sizes[loaded_shard_id] // tp_size
@@ -598,11 +592,11 @@ class QKVParallelLinear(ColumnParallelLinear):
                 hidden_size: int,
                 head_size: int,
                 total_num_heads: int,
-                 total_num_kv_heads: Optional[int] = None,
+                 total_num_kv_heads: int | None = None,
                 bias: bool = True,
                 skip_bias_add: bool = False,
-                 params_dtype: Optional[torch.dtype] = None,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 params_dtype: torch.dtype | None = None,
+                 quant_config: QuantizationConfig | None = None,
                 prefix: str = ""):
        self.hidden_size = hidden_size
        self.head_size = head_size
@@ -637,7 +631,7 @@ class QKVParallelLinear(ColumnParallelLinear):
                         quant_config=quant_config,
                         prefix=prefix)

-    def _get_shard_offset_mapping(self, loaded_shard_id: str) -> Optional[int]:
+    def _get_shard_offset_mapping(self, loaded_shard_id: str) -> int | None:
        shard_offset_mapping = {
            "q": 0,
            "k": self.num_heads * self.head_size,
@@ -646,7 +640,7 @@ class QKVParallelLinear(ColumnParallelLinear):
        }
        return shard_offset_mapping.get(loaded_shard_id)

-    def _get_shard_size_mapping(self, loaded_shard_id: str) -> Optional[int]:
+    def _get_shard_size_mapping(self, loaded_shard_id: str) -> int | None:
        shard_size_mapping = {
            "q": self.num_heads * self.head_size,
            "k": self.num_kv_heads * self.head_size,
@@ -679,10 +673,8 @@ class QKVParallelLinear(ColumnParallelLinear):
            # Special case for Quantization.
            # If quantized, we need to adjust the offset and size to account
            # for the packing.
-            if isinstance(
-                    param,
-                (PackedColumnParameter,
-                 PackedvLLMParameter)) and param.packed_dim == param.output_dim:
+            if isinstance(param, PackedColumnParameter | PackedvLLMParameter
+                          ) and param.packed_dim == param.output_dim:
                shard_size, shard_offset = \
                    param.adjust_shard_indexes_for_packing(
                    shard_size=shard_size, shard_offset=shard_offset)
@@ -694,7 +686,7 @@ class QKVParallelLinear(ColumnParallelLinear):
    def weight_loader_v2(self,
                         param: BasevLLMParameter,
                         loaded_weight: torch.Tensor,
-                         loaded_shard_id: Optional[str] = None):
+                         loaded_shard_id: str | None = None):
        if loaded_shard_id is None:  # special case for certain models
            if isinstance(param, PerTensorScaleParameter):
                param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=0)
@@ -720,7 +712,7 @@ class QKVParallelLinear(ColumnParallelLinear):
    def weight_loader(self,
                      param: Parameter,
                      loaded_weight: torch.Tensor,
-                      loaded_shard_id: Optional[str] = None):
+                      loaded_shard_id: str | None = None):

        param_data = param.data
        output_dim = getattr(param, "output_dim", None)
@@ -845,9 +837,9 @@ class RowParallelLinear(LinearBase):
                 bias: bool = True,
                 input_is_parallel: bool = True,
                 skip_bias_add: bool = False,
-                 params_dtype: Optional[torch.dtype] = None,
+                 params_dtype: torch.dtype | None = None,
                 reduce_results: bool = True,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 quant_config: QuantizationConfig | None = None,
                 prefix: str = ""):
        # Divide the weight matrix along the first dimension.
        self.tp_rank = get_tensor_model_parallel_rank()
@@ -921,7 +913,7 @@ class RowParallelLinear(LinearBase):

        param.load_row_parallel_weight(loaded_weight=loaded_weight)

-    def forward(self, input_) -> tuple[torch.Tensor, Optional[Parameter]]:
+    def forward(self, input_) -> tuple[torch.Tensor, Parameter | None]:
        if self.input_is_parallel:
            input_parallel = input_
        else:
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0

-from typing import Optional
-
 import torch
 import torch.nn as nn

@@ -18,10 +16,10 @@ class MLP(nn.Module):
        self,
        input_dim: int,
        mlp_hidden_dim: int,
-        output_dim: Optional[int] = None,
+        output_dim: int | None = None,
        bias: bool = True,
        act_type: str = "gelu_pytorch_tanh",
-        dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
        prefix: str = "",
    ):
        super().__init__()
@@ -0,0 +1,65 @@
+from typing import Literal, get_args
+
+from fastvideo.v1.layers.quantization.base_config import QuantizationConfig
+
+QuantizationMethods = Literal[None]
+
+QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
+
+# The customized quantization methods which will be added to this dict.
+_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {}
+
+
+def register_quantization_config(quantization: str):
+    """Register a customized vllm quantization config.
+
+    When a quantization method is not supported by vllm, you can register a customized
+    quantization config to support it.
+
+    Args:
+        quantization (str): The quantization method name.
+
+    Examples:
+        >>> from fastvideo.v1.layers.quantization import register_quantization_config
+        >>> from fastvideo.v1.layers.quantization import get_quantization_config
+        >>> from fastvideo.v1.layers.quantization.base_config import QuantizationConfig
+        >>>
+        >>> @register_quantization_config("my_quant")
+        ... class MyQuantConfig(QuantizationConfig):
+        ...     pass
+        >>>
+        >>> get_quantization_config("my_quant")
+        <class 'MyQuantConfig'>
+    """  # noqa: E501
+
+    def _wrapper(quant_config_cls):
+        if quantization in QUANTIZATION_METHODS:
+            raise ValueError(
+                f"The quantization method `{quantization}` is already exists.")
+        if not issubclass(quant_config_cls, QuantizationConfig):
+            raise ValueError("The quantization config must be a subclass of "
+                             "`QuantizationConfig`.")
+        _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls
+        QUANTIZATION_METHODS.append(quantization)
+        return quant_config_cls
+
+    return _wrapper
+
+
+def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
+    if quantization not in QUANTIZATION_METHODS:
+        raise ValueError(f"Invalid quantization method: {quantization}")
+
+    method_to_config: dict[str, type[QuantizationConfig]] = {}
+    # Update the `method_to_config` with customized quantization methods.
+    method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
+
+    return method_to_config[quantization]
+
+
+all = [
+    "QuantizationMethods",
+    "QuantizationConfig",
+    "get_quantization_config",
+    "QUANTIZATION_METHODS",
+]
@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/quantization/base_config.py
+
+import inspect
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from fastvideo.v1.layers.quantization import QuantizationMethods
+else:
+    QuantizationMethods = str
+
+
+class QuantizeMethodBase(ABC):
+    """Base class for different quantized methods."""
+
+    @abstractmethod
+    def create_weights(self, layer: torch.nn.Module, *weight_args,
+                       **extra_weight_attrs):
+        """Create weights for a layer.
+
+        The weights will be set as attributes of the layer."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor:
+        """Apply the weights in layer to the input tensor.
+
+        Expects create_weights to have been called before on the layer."""
+        raise NotImplementedError
+
+    # Not required functions
+    def embedding(self, layer: torch.nn.Module, *args,
+                  **kwargs) -> torch.Tensor:
+        """Gather embeddings in the layer based on indices in the input tensor.
+
+        Expects create_weights to have been called before on the layer."""
+        raise NotImplementedError
+
+    def process_weights_after_loading(self, layer: nn.Module) -> None:
+        """Process the weight after loading.
+
+        This can be used for example, to transpose weights for computation.
+        """
+        return
+
+
+def method_has_implemented_embedding(
+        method_class: type[QuantizeMethodBase]) -> bool:
+    """
+    Not all quant methods have embedding implemented, so we need to check that
+    it exists for our given method. We check this by making sure the function
+    has been changed from the base implementation.
+    """
+    base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding",
+                                            None)
+    class_embedding = inspect.getattr_static(method_class, "embedding", None)
+
+    return (class_embedding is not None
+            and class_embedding is not base_embedding)
+
+
+class QuantizationConfig(ABC):
+    """Base class for quantization configs."""
+
+    def __init__(self):
+        super().__init__()
+        # mapping is updated by models as they initialize
+        self.packed_modules_mapping: dict[str, list[str]] = dict()
+
+    @abstractmethod
+    def get_name(self) -> QuantizationMethods:
+        """Name of the quantization method."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_supported_act_dtypes(self) -> list[torch.dtype]:
+        """List of supported activation dtypes."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def get_min_capability(cls) -> int:
+        """Minimum GPU capability to support the quantization method.
+
+        E.g., 70 for Volta, 75 for Turing, 80 for Ampere.
+        This requirement is due to the custom CUDA kernels used by the
+        quantization method.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def get_config_filenames() -> list[str]:
+        """List of filenames to search for in the model directory."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig":
+        """Create a config class from the model's quantization config."""
+        raise NotImplementedError
+
+    @classmethod
+    def override_quantization_method(cls, hf_quant_cfg,
+                                     user_quant) -> QuantizationMethods | None:
+        """
+           Detects if this quantization method can support a given checkpoint
+           format by overriding the user specified quantization method -- 
+           this method should only be overwritten by subclasses in exceptional 
+           circumstances
+        """
+        return None
+
+    @staticmethod
+    def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any:
+        """Get a value from the model's quantization config."""
+        for key in keys:
+            if key in config:
+                return config[key]
+        raise ValueError(f"Cannot find any of {keys} in the model's "
+                         "quantization config.")
+
+    @staticmethod
+    def get_from_keys_or(config: dict[str, Any], keys: list[str],
+                         default: Any) -> Any:
+        """Get a optional value from the model's quantization config."""
+        try:
+            return QuantizationConfig.get_from_keys(config, keys)
+        except ValueError:
+            return default
+
+    @abstractmethod
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> QuantizeMethodBase | None:
+        """Get the quantize method to use for the quantized layer.
+        
+        Args:
+            layer: The layer for the quant method.
+            prefix: The full name of the layer in the state dict
+        Returns:
+            The quantize method. None if the given layer doesn't support quant
+            method.
+        """
+        raise NotImplementedError
+
+    def get_cache_scale(self, name: str) -> str | None:
+        return None
@@ -23,7 +23,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Rotary Positional Embeddings."""
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any

 import torch

@@ -84,7 +84,7 @@ class RotaryEmbedding(CustomOp):
        head_size: int,
        rotary_dim: int,
        max_position_embeddings: int,
-        base: Union[int, float],
+        base: int | float,
        is_neox_style: bool,
        dtype: torch.dtype,
    ) -> None:
@@ -101,7 +101,7 @@ class RotaryEmbedding(CustomOp):
        self.cos_sin_cache: torch.Tensor
        self.register_buffer("cos_sin_cache", cache, persistent=False)

-    def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
+    def _compute_inv_freq(self, base: int | float) -> torch.Tensor:
        """Compute the inverse frequency."""
        # NOTE(woosuk): To exactly match the HF implementation, we need to
        # use CPU to compute the cache and then move it to GPU. However, we
@@ -127,8 +127,8 @@ class RotaryEmbedding(CustomOp):
        positions: torch.Tensor,
        query: torch.Tensor,
        key: torch.Tensor,
-        offsets: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        offsets: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
        """A PyTorch-native implementation of forward()."""
        if offsets is not None:
            positions = positions + offsets
@@ -152,28 +152,6 @@ class RotaryEmbedding(CustomOp):
        key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
        return query, key

-    def forward_cuda(
-        self,
-        positions: torch.Tensor,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        offsets: Optional[torch.Tensor] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        from vllm import _custom_ops as ops
-
-        self.cos_sin_cache = self.cos_sin_cache.to(query.device,
-                                                   dtype=query.dtype)
-        # ops.rotary_embedding()/batched_rotary_embedding()
-        # are in-place operations that update the query and key tensors.
-        if offsets is not None:
-            ops.batched_rotary_embedding(positions, query, key, self.head_size,
-                                         self.cos_sin_cache, self.is_neox_style,
-                                         self.rotary_dim, offsets)
-        else:
-            ops.rotary_embedding(positions, query, key, self.head_size,
-                                 self.cos_sin_cache, self.is_neox_style)
-        return query, key
-
    def extra_repr(self) -> str:
        s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
        s += f", max_position_embeddings={self.max_position_embeddings}"
@@ -181,7 +159,7 @@ class RotaryEmbedding(CustomOp):
        return s


-def _to_tuple(x: Union[int, Tuple[int, ...]], dim: int = 2) -> Tuple[int, ...]:
+def _to_tuple(x: int | tuple[int, ...], dim: int = 2) -> tuple[int, ...]:
    if isinstance(x, int):
        return (x, ) * dim
    elif len(x) == dim:
@@ -190,8 +168,8 @@ def _to_tuple(x: Union[int, Tuple[int, ...]], dim: int = 2) -> Tuple[int, ...]:
        raise ValueError(f"Expected length {dim} or int, but got {x}")


-def get_meshgrid_nd(start: Union[int, Tuple[int, ...]],
-                    *args: Union[int, Tuple[int, ...]],
+def get_meshgrid_nd(start: int | tuple[int, ...],
+                    *args: int | tuple[int, ...],
                    dim: int = 2) -> torch.Tensor:
    """
    Get n-D meshgrid with start, stop and num.
@@ -239,12 +217,12 @@ def get_meshgrid_nd(start: Union[int, Tuple[int, ...]],

 def get_1d_rotary_pos_embed(
    dim: int,
-    pos: Union[torch.FloatTensor, int],
+    pos: torch.FloatTensor | int,
    theta: float = 10000.0,
    theta_rescale_factor: float = 1.0,
    interpolation_factor: float = 1.0,
    dtype: torch.dtype = torch.float32,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Precompute the frequency tensor for complex exponential (cis) with given dimensions.
    (Note: `cis` means `cos + i * sin`, where i is the imaginary unit.)
@@ -283,13 +261,13 @@ def get_nd_rotary_pos_embed(
    start,
    *args,
    theta=10000.0,
-    theta_rescale_factor: Union[float, List[float]] = 1.0,
-    interpolation_factor: Union[float, List[float]] = 1.0,
+    theta_rescale_factor: float | list[float] = 1.0,
+    interpolation_factor: float | list[float] = 1.0,
    shard_dim: int = 0,
    sp_rank: int = 0,
    sp_world_size: int = 1,
    dtype: torch.dtype = torch.float32,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
    """
    This is a n-d version of precompute_freqs_cis, which is a RoPE for tokens with n-d structure.
    Supports sequence parallelism by allowing sharding of a specific dimension.
@@ -346,7 +324,7 @@ def get_nd_rotary_pos_embed(
    else:
        grid = full_grid

-    if isinstance(theta_rescale_factor, (int, float)):
+    if isinstance(theta_rescale_factor, int | float):
        theta_rescale_factor = [theta_rescale_factor] * len(rope_dim_list)
    elif isinstance(theta_rescale_factor,
                    list) and len(theta_rescale_factor) == 1:
@@ -355,7 +333,7 @@ def get_nd_rotary_pos_embed(
        rope_dim_list
    ), "len(theta_rescale_factor) should equal to len(rope_dim_list)"

-    if isinstance(interpolation_factor, (int, float)):
+    if isinstance(interpolation_factor, int | float):
        interpolation_factor = [interpolation_factor] * len(rope_dim_list)
    elif isinstance(interpolation_factor,
                    list) and len(interpolation_factor) == 1:
@@ -392,7 +370,7 @@ def get_rotary_pos_embed(
    interpolation_factor=1.0,
    shard_dim: int = 0,
    dtype: torch.dtype = torch.float32,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Generate rotary positional embeddings for the given sizes.
    
@@ -439,17 +417,17 @@ def get_rotary_pos_embed(
    return freqs_cos, freqs_sin


-_ROPE_DICT: Dict[Tuple, RotaryEmbedding] = {}
+_ROPE_DICT: dict[tuple, RotaryEmbedding] = {}


 def get_rope(
    head_size: int,
    rotary_dim: int,
    max_position: int,
-    base: Union[int, float],
+    base: int | float,
    is_neox_style: bool = True,
-    rope_scaling: Optional[Dict[str, Any]] = None,
-    dtype: Optional[torch.dtype] = None,
+    rope_scaling: dict[str, Any] | None = None,
+    dtype: torch.dtype | None = None,
    partial_rotary_factor: float = 1.0,
 ) -> RotaryEmbedding:
    if dtype is None:
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/utils.py
 """Utility methods for model layers."""
-from typing import Tuple

 import torch

@@ -10,7 +9,7 @@ def get_token_bin_counts_and_mask(
    tokens: torch.Tensor,
    vocab_size: int,
    num_seqs: int,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
    # Compute the bin counts for the tokens.
    # vocab_size + 1 for padding.
    bin_counts = torch.zeros((num_seqs, vocab_size + 1),
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0

 import math
-from typing import Optional

 import torch
 import torch.nn as nn
@@ -36,7 +35,7 @@ class PatchEmbed(nn.Module):
                 prefix: str = ""):
        super().__init__()
        # Convert patch_size to 2-tuple
-        if isinstance(patch_size, (list, tuple)):
+        if isinstance(patch_size, list | tuple):
            if len(patch_size) == 1:
                patch_size = (patch_size[0], patch_size[0])
        else:
@@ -133,7 +132,7 @@ class ModulateProjection(nn.Module):
        hidden_size: int,
        factor: int = 2,
        act_layer: str = "silu",
-        dtype: Optional[torch.dtype] = None,
+        dtype: torch.dtype | None = None,
        prefix: str = "",
    ):
        super().__init__()
@@ -1,17 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0

+from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import List, Optional, Sequence, Tuple

 import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter, UninitializedParameter
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)

 from fastvideo.v1.distributed import (divide, get_tensor_model_parallel_rank,
                                      get_tensor_model_parallel_world_size,
                                      tensor_model_parallel_all_reduce)
+from fastvideo.v1.layers.quantization.base_config import (
+    QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
 from fastvideo.v1.models.parameter import BasevLLMParameter
 from fastvideo.v1.models.utils import set_weight_attrs
 from fastvideo.v1.platforms import current_platform
@@ -24,7 +24,7 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):

    def create_weights(self, layer: torch.nn.Module,
                       input_size_per_partition: int,
-                       output_partition_sizes: List[int], input_size: int,
+                       output_partition_sizes: list[int], input_size: int,
                       output_size: int, params_dtype: torch.dtype,
                       **extra_weight_attrs):
        """Create weights for embedding layer."""
@@ -39,7 +39,7 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
    def apply(self,
              layer: torch.nn.Module,
              x: torch.Tensor,
-              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+              bias: torch.Tensor | None = None) -> torch.Tensor:
        return F.linear(x, layer.weight, bias)

    def embedding(self, layer: torch.nn.Module,
@@ -139,7 +139,7 @@ def get_masked_input_and_mask(
        input_: torch.Tensor, org_vocab_start_index: int,
        org_vocab_end_index: int, num_org_vocab_padding: int,
        added_vocab_start_index: int,
-        added_vocab_end_index: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        added_vocab_end_index: int) -> tuple[torch.Tensor, torch.Tensor]:
    # torch.compile will fuse all of the pointwise ops below
    # into a single kernel, making it very fast
    org_vocab_mask = (input_ >= org_vocab_start_index) & (input_
@@ -197,10 +197,10 @@ class VocabParallelEmbedding(torch.nn.Module):
    def __init__(self,
                 num_embeddings: int,
                 embedding_dim: int,
-                 params_dtype: Optional[torch.dtype] = None,
-                 org_num_embeddings: Optional[int] = None,
+                 params_dtype: torch.dtype | None = None,
+                 org_num_embeddings: int | None = None,
                 padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
-                 quant_config: Optional[QuantizationConfig] = None,
+                 quant_config: QuantizationConfig | None = None,
                 prefix: str = ""):
        super().__init__()

@@ -296,7 +296,7 @@ class VocabParallelEmbedding(torch.nn.Module):
            org_vocab_start_index, org_vocab_end_index, added_vocab_start_index,
            added_vocab_end_index)

-    def get_sharded_to_full_mapping(self) -> Optional[List[int]]:
+    def get_sharded_to_full_mapping(self) -> list[int] | None:
        """Get a mapping that can be used to reindex the gathered
        logits for sampling.
        
@@ -310,9 +310,9 @@ class VocabParallelEmbedding(torch.nn.Module):
        if self.tp_size < 2:
            return None

-        base_embeddings: List[int] = []
-        added_embeddings: List[int] = []
-        padding: List[int] = []
+        base_embeddings: list[int] = []
+        added_embeddings: list[int] = []
+        padding: list[int] = []
        for tp_rank in range(self.tp_size):
            shard_indices = self._get_indices(self.num_embeddings_padded,
                                              self.org_vocab_size_padded,
@@ -11,7 +11,7 @@ from logging import Logger
 from logging.config import dictConfig
 from os import path
 from types import MethodType
-from typing import Any, Optional, cast
+from typing import Any, cast

 import fastvideo.v1.envs as envs

@@ -278,8 +278,7 @@ def _trace_calls(log_path, root_dir, frame, event, arg=None):
    return partial(_trace_calls, log_path, root_dir)


-def enable_trace_function_call(log_file_path: str,
-                               root_dir: Optional[str] = None):
+def enable_trace_function_call(log_file_path: str, root_dir: str | None = None):
    """
    Enable tracing of every function call in code under `root_dir`.
    This is useful for debugging hangs or crashes.
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
SolitaryThinker	d0c53871a4	fix tensor type hint	2025-05-23 14:42:23 -07:00
SolitaryThinker	0d5306f61f	update min python to 3.10	2025-05-23 14:42:23 -07:00
applesaucethebunandBrayden Zhong	12647457a7	[Misc] Small fixes to Torch code (#395 ) Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>	2025-05-23 14:40:24 -07:00
Kevin Lin	298f74f956	Set device for encode (#420 )	2025-05-23 14:19:45 -07:00
Wenxuan Tan	ee8babb298	Unify env report script in issue template (#423 )	2025-05-23 14:19:18 -07:00
Wenxuan Tan	60295cc03f	Use version.py (#424 )	2025-05-23 14:17:33 -07:00
William Lin	1572e13b6e	[Tests] don't run 3.10 and 3.11 for SSIM (#427 )	2025-05-23 12:59:50 -07:00
Wenxuan Tan	a157275b4c	Fix version number (#422 )	2025-05-22 12:37:34 -07:00
William Lin	c4dbe7dac3	[bug] fix bs > 1 (#418 )	2025-05-21 21:07:01 -07:00
Kevin Lin	d39591108e	Fulfill worker response on interrupt (#417 )	2025-05-21 20:59:11 -07:00
William Lin	ace6e971e5	[V1] Remove vLLM dependency (#413 )	2025-05-18 00:57:28 -07:00
William Lin	b4f6758253	[Teacache] allow None for forward_context batch when using teacache (#412 )	2025-05-17 18:43:14 -07:00
William Lin	535d29b392	[Docs] Fix image (#407 )	2025-05-12 14:30:05 -07:00
William Lin	b4255517e0	[Docs] Add CLI docs (#406 )	2025-05-12 14:16:48 -07:00
William Lin	6eeb60613f	Release 0.1.0 (#405 )	2025-05-12 11:54:23 -07:00
William Lin	53d2c7791f	[V1] Update where num_frame rounding is done (#403 )	2025-05-12 11:52:49 -07:00
William Lin	53cb693dca	[V1] Docs Update (#402 )	2025-05-12 11:52:09 -07:00
Kevin Lin	6f72d24876	[CLI] Default to pipeline config (#401 )	2025-05-12 00:22:53 -07:00
William Lin	d1459e9976	[V1] Update README (#400 )	2025-05-11 22:32:03 -07:00
William Lin	59ab481eb1	release 0.0.5 (#399 )	2025-05-11 16:23:11 -07:00
William Lin	0cf001986a	[Docs] More docs update (#394 )	2025-05-11 16:19:47 -07:00
applesaucethebunandBrayden Zhong	51956369a5	[Misc] Replace instances of `time.time()` with `time.perf_counter()` (#396 ) Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>	2025-05-11 15:56:39 -07:00
Kevin Lin	4b0970cbbf	[CI] Set volume_size required to false (#398 )	2025-05-11 15:55:34 -07:00
Kevin Lin	94bf47a572	[CI] Set default disk size (#397 )	2025-05-11 15:34:17 -07:00
River (Zihang He)andWill Lin	1a3ac9074b	Zihang stepvideo v1 (#389 ) Co-authored-by: Will Lin <wlsaidhi@gmail.com>	2025-05-11 14:41:42 -07:00
Kevin Lin	fb0581d5b0	[CLI] Update cli to support new api/model config (#384 )	2025-05-11 13:50:33 -07:00
Kevin Lin	6c74ab4132	[CI] Use python 3.10/3.11 for SSIM test (#392 )	2025-05-10 13:48:29 -07:00
William Lin	9a91021c56	[Docs] Add collect_env.py and various docs update (#393 )	2025-05-10 13:42:42 -07:00