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

.github/ISSUE_TEMPLATE/1-bug-report.yml

@@ -8,7 +8,7 @@ body:
   attributes:
     label: Environment
     description: |
-      Please share your environment with us. You can run the command **python fastvideo/utils/collect_env.py** and copy-paste its output below.
+      Please share your environment with us. You can run the command **python fastvideo/utils/env_utils.py** and copy-paste its output below.
     placeholder: FastVideo version, platform, python version, cuda version...
   validations:
     required: true

.github/workflows/build-image-template.yml

@@ -1,106 +0,0 @@
-name: Build Image Template
-
-on:
-  workflow_call:
-    inputs:
-      python_version:
-        required: true
-        type: string
-      dockerfile_path:
-        required: true
-        type: string
-      tag_suffix:
-        required: true
-        type: string
-
-jobs:
-  build-and-push:
-    runs-on: ubuntu-latest
-    permissions:
-      contents: read
-      packages: write
-    
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-      
-      - name: Free up disk space
-        run: |
-          # Display initial space
-          echo "Initial disk space:"
-          df -h
-          
-          # Remove large directories directly
-          sudo rm -rf /usr/share/dotnet
-          sudo rm -rf /usr/local/lib/android
-          sudo rm -rf /opt/ghc
-          sudo rm -rf /usr/local/share/boost
-          sudo rm -rf /usr/share/swift
-          sudo rm -rf /usr/local/lib/node_modules
-          sudo rm -rf /usr/local/share/powershell
-          sudo rm -rf /usr/share/rust
-          sudo rm -rf /usr/local/.ghcup
-          
-          # Remove cached files
-          sudo rm -rf /var/lib/apt/lists/*
-          sudo rm -rf /var/cache/apt/archives/*
-          
-          # Clean Docker
-          docker system prune -af --volumes
-          
-          # Display available space after cleanup
-          echo "Disk space after cleanup:"
-          df -h
-      
-      - name: Set up Docker Buildx
-        uses: docker/setup-buildx-action@v3
-      
-      - name: Login to GitHub Container Registry
-        uses: docker/login-action@v3
-        with:
-          registry: ghcr.io
-          username: ${{ github.repository_owner }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-      
-      - name: Prepare tags
-        id: prepare-tags
-        run: |
-          SHORT_SHA=$(echo ${{ github.sha }} | cut -c1-7)
-          
-          TAGS="type=raw,value=${{ inputs.tag_suffix }}-latest"
-          TAGS="${TAGS}\ntype=raw,value=${{ inputs.tag_suffix }}-sha-${SHORT_SHA}"
-          
-          # Set Python 3.10 as the default image
-          if [[ "${{ inputs.python_version }}" == "3.10" ]]; then
-            TAGS="${TAGS}\ntype=raw,value=latest"
-          fi
-          
-          {
-            echo "tags<<EOF"
-            echo -e "$TAGS"
-            echo "EOF" 
-          } >> $GITHUB_OUTPUT
-      
-      - name: Extract metadata for Docker
-        id: meta
-        uses: docker/metadata-action@v5
-        with:
-          images: ghcr.io/${{ github.repository }}/fastvideo-dev
-          tags: ${{ steps.prepare-tags.outputs.tags }}
-      
-      - name: Build and push Docker image
-        id: build-push
-        uses: docker/build-push-action@v6
-        with:
-          context: .
-          file: ${{ inputs.dockerfile_path }}
-          push: true
-          tags: ${{ steps.meta.outputs.tags }}
-          labels: ${{ steps.meta.outputs.labels }}
-          cache-from: type=gha
-          cache-to: type=gha,mode=max
-      
-      - name: Success message
-        run: |
-          echo "✅ Python ${{ inputs.python_version }} image successfully built and pushed to ghcr.io/${{ github.repository }}/fastvideo-dev:${{ inputs.tag_suffix }}-latest"
-          echo "To run tests with this image, manually trigger the 'Run Tests' workflow." 

.github/workflows/build-image.yml

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

.github/workflows/pr-test.yml

@@ -83,81 +83,193 @@ 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')
-    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 }}
+    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
 
   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')
-    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 }}
+    runs-on: ubuntu-latest
+    environment: runpod-runners
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Set up SSH key
+        run: |
+          mkdir -p ~/.ssh
+          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
+          chmod 600 ~/.ssh/id_rsa
+          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
+
+      - name: Install dependencies
+        run: pip install requests
+
+      - name: Run tests on RunPod
+        env:
+          JOB_ID: "vae-test"
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: 30
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "NVIDIA A40"
+          --gpu-count 1
+          --volume-size 100
+          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: "vae-test"
+        run: python .github/scripts/runpod_cleanup.py
 
   transformer-test:
     needs: change-filter
     if: >-
       (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.transformer-test == 'true') || 
       (github.event_name == 'workflow_dispatch' && github.event.inputs.run_transformer_test == 'true')
-    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 }}
+    runs-on: ubuntu-latest
+    environment: runpod-runners
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Set up SSH key
+        run: |
+          mkdir -p ~/.ssh
+          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
+          chmod 600 ~/.ssh/id_rsa
+          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
+
+      - name: Install dependencies
+        run: pip install requests
+
+      - name: Run tests on RunPod
+        env:
+          JOB_ID: "transformer-test"
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: 30
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "NVIDIA L40S"
+          --gpu-count 1
+          --volume-size 100
+          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: "transformer-test"
+        run: python .github/scripts/runpod_cleanup.py
       
   ssim-test:
     needs: change-filter
     if: >-
       (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
       (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
-    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 }}
+    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
 
   runpod-cleanup:
     needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
@@ -177,7 +289,7 @@ jobs:
 
       - name: Cleanup all RunPod instances
         env:
-          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12"]'
+          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test"]' # JSON array of job IDs
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           GITHUB_RUN_ID: ${{ github.run_id }}
         run: python .github/scripts/runpod_cleanup.py

.github/workflows/runpod-test.yml

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

.pre-commit-config.yaml

@@ -19,11 +19,9 @@ exclude: |
     fastvideo/sample/.*|
     fastvideo/train\.py|
     fastvideo/utils/.*|
-    examples/.*|
+    fastvideo/v1/examples/.*|
     .github/workflows/fastvideo-publish.yml|
-    .github/workflows/sta-publish.yml|
-    .github/workflows/build-image-template.yml|
-    docs/source/inference/support_matrix.md
+    .github/workflows/sta-publish.yml
   )
 repos:
 - repo: https://github.com/google/yapf

Dockerfile

@@ -29,7 +29,7 @@ RUN echo "# Placeholder" > README.md
 
 RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
     conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
-    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.0.post2 --no-build-isolation && \
     conda clean -afy
 
 COPY . .

README.md

@@ -2,109 +2,62 @@
 <img src=assets/logo.jpg width="30%"/>
 </div>
 
-**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.
+FastVideo is a lightweight framework for accelerating large video diffusion models.
 
 <p align="center">
-    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-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://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>
 
-<div align="center">
-<img src=assets/perf.png width="90%"/>
-</div>
+https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1
 
-## Key Features
+FastVideo currently offers: (with more to come)
 
-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.
+- [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.
 - 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)
 
-## Distillation and Finetuning
+### Inference
+- [Quick Start](https://hao-ai-lab.github.io/FastVideo/inference/examples/basic.html)
+- V1 Inference API Guide (Coming soon!)
+
+### 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 -->
-  - [x] Add StepVideo to V1
+  - [ ] Add StepVideo to V1
 - Optimization features
-  - [x] Teacache in V1
-  - [x] SageAttention in V1
+  - [ ] Teacache in V1
+  - [ ] SageAttention in V1
 - Code updates
-  - [x] V1 Configuration API
+  - [ ] V1 Configuration API
   - [ ] Support Training in V1
   <!-- - [ ] fp8 support -->
   <!-- - [ ] faster load model and save model support -->

csrc/sliding_tile_attention/test/test_sta.py

@@ -2,6 +2,7 @@ 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)
@@ -22,7 +23,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.linalg.norm(tensor, dim=-1, keepdim=True)
+    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
     scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
 
     return scaled_tensor.contiguous()

docker/Dockerfile.python3.11

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

docker/Dockerfile.python3.12

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

docs/requirements-docs.txt

@@ -1,15 +1,25 @@
-sphinx==7.4.7
-sphinx-argparse==0.5.2
-sphinx-autodoc2==0.5.0
-sphinx-book-theme==1.1.4
+sphinx==6.2.1
+sphinx-argparse==0.4.0
+sphinx-book-theme==1.0.1
 sphinx-copybutton==0.5.2
 sphinx-design==0.6.1
 sphinx-togglebutton==0.3.2
 myst-parser==3.0.1
 msgspec
-commonmark # Required by sphinx-argparse when using :markdownhelp:
+cloudpickle
 
 # packages to install to build the documentation
 cachetools
+pydantic >= 2.8
 -f https://download.pytorch.org/whl/cpu
-torch
+torch
+py-cpuinfo
+transformers
+mistral_common >= 1.5.4
+aiohttp
+starlette
+openai # Required by docs/source/serving/openai_compatible_server.md's vllm.entrypoints.openai.cli_args
+fastapi # Required by docs/source/serving/openai_compatible_server.md's vllm.entrypoints.openai.cli_args
+partial-json-parser # Required by docs/source/serving/openai_compatible_server.md's vllm.entrypoints.openai.cli_args
+requests
+zmq

docs/source/api/summary.md

@@ -1,19 +0,0 @@
-# Summary
-
-## Video Generator
-
-```{autodoc2-summary}
-    fastvideo.VideoGenerator
-```
-
-## Initialization Configuration
-
-```{autodoc2-summary}
-    fastvideo.v1.configs.pipelines.PipelineConfig
-```
-
-## Sampling Configuration
-
-```{autodoc2-summary}
-    fastvideo.v1.configs.sample.SamplingParam
-```

docs/source/autodoc2_docstring_parser.py

@@ -1,22 +0,0 @@
-# type: ignore
-# SPDX-License-Identifier: Apache-2.0
-from docutils import nodes
-from myst_parser.parsers.sphinx_ import MystParser
-from sphinx.ext.napoleon import docstring
-
-
-class NapoleonParser(MystParser):
-
-    def parse(self, input_string: str, document: nodes.document) -> None:
-        # Get the Sphinx configuration
-        config = document.settings.env.config
-
-        parsed_content = str(
-            docstring.GoogleDocstring(
-                str(docstring.NumpyDocstring(input_string, config)),
-                config,
-            ))
-        return super().parse(parsed_content, document)
-
-
-Parser = NapoleonParser

docs/source/conf.py

@@ -13,18 +13,17 @@
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 
 import datetime
+import inspect
 import logging
 import os
-import re
 import sys
-from pathlib import Path
+from typing import Optional
 
 import requests
+from sphinx.ext import autodoc
 
 logger = logging.getLogger(__name__)
-REPO_ROOT = Path(__file__).resolve().parent.parent.parent
-print(os.path.abspath(REPO_ROOT))
-sys.path.append(os.path.abspath(REPO_ROOT))
+sys.path.append(os.path.abspath("../.."))
 
 # -- Project information -----------------------------------------------------
 
@@ -42,7 +41,8 @@ extensions = [
     "sphinx.ext.linkcode",
     "sphinx.ext.intersphinx",
     "sphinx_copybutton",
-    "autodoc2",
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
     "myst_parser",
     "sphinxarg.ext",
     "sphinx_design",
@@ -50,31 +50,6 @@ extensions = [
 ]
 myst_enable_extensions = [
     "colon_fence",
-    "fieldlist",
-]
-autodoc2_packages = [
-    {
-        "path": "../../fastvideo",
-        "exclude_dirs": ["__pycache__", "third_party"],
-    },
-]
-autodoc2_output_dir = "api"
-autodoc2_render_plugin = "myst"
-autodoc2_hidden_objects = ["dunder", "private", "inherited"]
-autodoc2_docstring_parser_regexes = [
-    (".*", "docs.source.autodoc2_docstring_parser"),
-]
-autodoc2_sort_names = True
-autodoc2_index_template = None
-autodoc2_skip_module_regexes = [
-    "fastvideo.dataset",
-    "fastvideo.distill",
-    "fastvideo.data_preprocess",
-    "fastvideo.models",
-    "fastvideo.sample",
-    "fastvideo.utils",
-    "fastvideo.distill_adv",
-    "fastvideo.train",
 ]
 
 # Add any paths that contain templates here, relative to this directory.
@@ -103,11 +78,6 @@ html_theme_options = {
     'repository_url': 'https://github.com/hao-ai-lab/FastVideo/',
     'use_repository_button': True,
     'use_edit_page_button': True,
-    # Prevents the full API being added to the left sidebar of every page.
-    # Reduces build time by 2.5x and reduces build size from ~225MB to ~95MB.
-    'collapse_navbar': True,
-    # Makes API visible in the right sidebar on API reference pages.
-    'show_toc_level': 3,
 }
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
@@ -167,7 +137,8 @@ _cached_base: str = ""
 _cached_branch: str = ""
 
 
-def get_repo_base_and_branch(pr_number: str) -> tuple[str | None, str | None]:
+def get_repo_base_and_branch(
+        pr_number: str) -> tuple[Optional[str], Optional[str]]:
     global _cached_base, _cached_branch
     if _cached_base and _cached_branch:
         return _cached_base, _cached_branch
@@ -189,38 +160,38 @@ def linkcode_resolve(domain, info):
         return None
     if not info['module']:
         return None
+    module = info['module']
 
-    # Get path from module name
-    file = Path(f"{info['module'].replace('.', '/')}.py")
-    path = REPO_ROOT / file
-    if not path.exists():
-        path = REPO_ROOT / file.with_suffix("") / "__init__.py"
-    if not path.exists():
-        return None
+    # try to determine the correct file and line number to link to
+    obj = sys.modules[module]
 
-    # Get the line number of the object
-    with open(path) as f:
-        lines = f.readlines()
-    name = info['fullname'].split(".")[-1]
-    pattern = fr"^( {{4}})*((def|class) )?{name}\b.*"
-    for lineno, line in enumerate(lines, 1):
-        if not line or line.startswith("#"):
-            continue
-        if re.match(pattern, line):
-            break
+    # get as specific as we can
+    lineno: int = 0
+    filename: str = ""
+    try:
+        for part in info['fullname'].split('.'):
+            obj = getattr(obj, part)
 
-    # If the line number is not found, return None
-    if lineno == len(lines):
-        return None
+            if not (inspect.isclass(obj) or inspect.isfunction(obj)
+                    or inspect.ismethod(obj)):
+                obj = obj.__class__  # type: ignore[assignment]
 
-    # If the line number is found, create the URL
-    filename = path.relative_to(REPO_ROOT)
-    if "checkouts" in path.parts:
+            lineno = inspect.getsourcelines(obj)[1]
+            filename = (inspect.getsourcefile(obj)
+                        or f"{filename}.py").split("FastVideo/", 1)[1]
+    except Exception:
+        # For some things, like a class member, won't work, so
+        # we'll use the line number of the parent (the class)
+        pass
+
+    if filename.startswith("checkouts/"):
         # a PR build on readthedocs
-        pr_number = REPO_ROOT.name
+        pr_number = filename.split("/")[1]
+        filename = filename.split("/", 2)[2]
         base, branch = get_repo_base_and_branch(pr_number)
         if base and branch:
             return f"https://github.com/{base}/blob/{branch}/{filename}#L{lineno}"
+
     # Otherwise, link to the source file on the main branch
     return f"https://github.com/hao-ai-lab/FastVideo/blob/main/{filename}#L{lineno}"
 
@@ -232,8 +203,6 @@ autodoc_mock_imports = [
     "cpuinfo",
     "cv2",
     "torch",
-    "huggingface_hub",
-    "torchvision",
     "transformers",
     "psutil",
     "prometheus_client",
@@ -262,6 +231,18 @@ for mock_target in autodoc_mock_imports:
             "been loaded into sys.modules when the sphinx build starts.",
             mock_target)
 
+
+class MockedClassDocumenter(autodoc.ClassDocumenter):
+    """Remove note about base class when a class is derived from object."""
+
+    def add_line(self, line: str, source: str, *lineno: int) -> None:
+        if line == "   Bases: :py:class:`object`":
+            return
+        super().add_line(line, source, *lineno)
+
+
+autodoc.ClassDocumenter = MockedClassDocumenter
+
 intersphinx_mapping = {
     "python": ("https://docs.python.org/3", None),
     "typing_extensions":
@@ -273,4 +254,7 @@ intersphinx_mapping = {
     "psutil": ("https://psutil.readthedocs.io/en/stable", None),
 }
 
+autodoc_preserve_defaults = True
+autodoc_warningiserror = True
+
 navigation_with_keys = False

docs/source/contributing/add_pipeline.md

@@ -1,6 +1,6 @@
 (add-pipeline)=
 
-# 🏗️ Adding a New Pipeline
+# 🏗️ Adding a New Diffusion Pipeline
 
 This guide explains how to implement a custom diffusion pipeline in FastVideo, leveraging the framework's modular architecture for high-performance video generation.
 

docs/source/contributing/developer_env/docker.md

@@ -1,4 +1,3 @@
-(docker)=
 # 🐳 Using the FastVideo Docker Image
 
 If you prefer a containerized development environment or want to avoid managing dependencies manually, you can use our prebuilt Docker image:

docs/source/contributing/overview.md

@@ -39,7 +39,7 @@ Now you can install FastVideo and setup git hooks for running linting. By using
 pip install -e .[dev]
 
 # Can also install flash-attn (optional)
-pip install flash-attn==2.7.4.post1 --no-build-isolation 
+pip install flash-attn==2.7.0.post2 --no-build-isolation 
 
 # Linting, formatting and static type checking
 pre-commit install --hook-type pre-commit --hook-type commit-msg

docs/source/generate_examples.py

@@ -5,10 +5,11 @@ 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 = '../../../..'
-EXAMPLE_DIR = ROOT_DIR / "examples"
+EXAMPLE_DIR = ROOT_DIR / "fastvideo/v1/examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples"
 
 
@@ -88,7 +89,7 @@ class Example:
         generate() -> str: Generates the documentation content.
     """ # noqa: E501
     path: Path
-    category: str | None = None
+    category: Optional[str] = None
     main_file: Path = field(init=False)
     other_files: list[Path] = field(init=False)
     title: str = field(init=False)

docs/source/getting_started/installation.md

@@ -4,19 +4,32 @@
 
 FastVideo currently only supports Linux and NVIDIA CUDA GPUs.
 
+FastVideo has been tested on the following GPUs, but it should work on any GPUs that supports CUDA 12.4+, please create an issue if you discover any issues:
+- RTX 4090
+- A40
+- L40S
+- A100
+- H100
+
 ## Requirements
 
-- **OS: Linux**
-- **Python: 3.10-3.12**
-- **CUDA 12.4**
-- **At least 1 NVIDIA GPU**
+- OS: Linux
+- Python: 3.10-3.12
+- CUDA 12.4+
 
-## Set up using Python
-### Create a new Python environment
+## Installation Options
 
-#### 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)
+### Option 1: Quick Install
+
+```bash
+pip install fastvideo
+```
+
+### Option 2: Installation from Source
+
+We recommend using a Python environment such as Conda.
+
+#### 1. [Optional] Install Miniconda (if not already installed)
 
 ```bash
 wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
@@ -24,77 +37,38 @@ bash Miniconda3-latest-Linux-x86_64.sh
 source ~/.bashrc
 ```
 
-##### 2. Create and activate a Conda environment for FastVideo
+#### 2. [Optional] Create and activate a Conda environment for FastVideo
 
 ```bash
-# (Recommended) Create a new conda environment.
-conda create -n fastvideo python=3.12 -y
+conda create -n fastvideo python=3.10 -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:
-
-```bash
-pip install flash-attn==2.7.4.post1 --no-build-isolation
-```
-
-### Installation from Source
-
-#### 1. Clone the FastVideo repository
+#### 3. Clone the FastVideo repository
 
 ```bash
 git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
 ```
 
-#### 2. Install FastVideo
+#### 4. Install FastVideo
 
 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
+pip install flash-attn==2.7.0.post2 --no-build-isolation
 ```
 
-## Set up using Docker
-We also have prebuilt docker images with FastVideo dependencies pre-installed:
-[Docker Images](#docker)
+### Sliding Tile Attention (STA) (Requires CUDA 12.4+ and H100)
+
+To try Sliding Tile Attention (optional), please follow the instructions in [csrc/sliding_tile_attention/README.md](#sta-installation) to install STA.
 
 ## Development Environment Setup
 
@@ -104,7 +78,8 @@ If you're planning to contribute to FastVideo please see the following page:
 ## Hardware Requirements
 
 ### For Basic Inference
-- NVIDIA GPU with CUDA 12.4 support
+- NVIDIA GPU with CUDA support
+- Minimum 20GB VRAM for quantized models (e.g., single RTX 4090)
 
 ### For Lora Finetuning
 - 40GB GPU memory each for 2 GPUs with lora

docs/source/getting_started/v1_api.md

@@ -1,83 +0,0 @@
-# V1 API
-
-FastVideo's V1 API provides a streamlined interface for video generation tasks with powerful customization options. This page documents the primary components of the API.
-
-## Video Generator
-
-This class will be the primary Python API for generating videos and images.
-
-```{autodoc2-summary}
-    fastvideo.VideoGenerator
-```
-
-`````{py:class} VideoGenerator(fastvideo_args: fastvideo.v1.fastvideo_args.FastVideoArgs, executor_class: type[fastvideo.v1.worker.executor.Executor], log_stats: bool)
-:canonical: fastvideo.v1.entrypoints.video_generator.VideoGenerator
-
-```{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
-```

docs/source/index.md

@@ -9,7 +9,7 @@
 
 :::{raw} html
 <p style="text-align:center">
-<strong>FastVideo is a unified framework for accelerated video generation.
+<strong>FastVideo is a lightweight framework for accelerating large video diffusion models.
 </strong>
 </p>
 
@@ -21,31 +21,27 @@
 </p>
 :::
 
-It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
-With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
+FastVideo is a lightweight framework for accelerating large video diffusion models developed by the [Hao AI Lab](https://hao-ai-lab.github.io/).
 
 <div style="text-align: center;">
-  <img src=_static/images/perf.png width="100%"/>
+  <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>
 </div>
 
-## Key Features
+FastVideo currently offers: (with more to come)
 
-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.
+- [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.
 - 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?
@@ -55,20 +51,14 @@ FastVideo has the following features:
 :maxdepth: 1
 
 getting_started/installation
-<!-- getting_started/v1_api -->
+<!-- getting_started/examples/examples_index -->
 :::
 
 :::{toctree}
 :caption: Inference
 :maxdepth: 1
 
-inference/inference_quick_start
-inference/configuration
-inference/optimizations
-inference/support_matrix
 inference/examples/examples_inference_index
-inference/cli
-inference/add_pipeline
 inference/v0_inference
 :::
 
@@ -103,14 +93,7 @@ design/overview
 
 contributing/overview
 contributing/developer_env/index
-:::
-
-:::{toctree}
-:caption: API Reference
-:maxdepth: 2
-
-<!-- api/summary -->
-api/fastvideo/fastvideo
+contributing/add_pipeline
 :::
 
 ## Indices and tables

docs/source/inference/cli.md

@@ -1,151 +0,0 @@
-# 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.

docs/source/inference/configuration.md

@@ -1,77 +0,0 @@
-(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)

docs/source/inference/fasthunyuan.md

@@ -0,0 +1,33 @@
+(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).

docs/source/inference/fastmochi.md

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

docs/source/inference/hunyuanvideo.md

@@ -0,0 +1,18 @@
+(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.

docs/source/inference/inference_quick_start.md

@@ -1,124 +0,0 @@
-# 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).

docs/source/inference/optimizations.md

@@ -1,148 +0,0 @@
-(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`.

docs/source/inference/stepvideo.md

@@ -0,0 +1,16 @@
+(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
+```

docs/source/inference/support_matrix.md

@@ -1,92 +0,0 @@
-(support-matrix)=
-# Compatibility Matrix
-The table below shows every supported model and optimizations supported for them.
-
-The symbols used have the following meanings:
-
-- ✅ = Full compatibility
-- ❌ = No compatibility
-
-## 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.

docs/source/inference/wanvideo.md

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

examples/inference/basic/README.md

@@ -1,41 +1,3 @@
-# Basic Video Generation Tutorial
-The `VideoGenerator` class provides the primary Python interface for doing offline video generation, which is interacting with a diffusion pipeline without using a separate inference api server.
+# Basic
 
-## Requirements
-- At least a single NVIDIA GPU with CUDA 12.4.
-- Python 3.10-3.12
-
-## Installation
-If you have not installed FastVideo, please following these [instructions](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) first.
-
-## Usage
-The first script in this example shows the most basic usage of FastVideo. If you are new to Python and FastVideo, you should start here.
-
-```bash
-# if you have not cloned the directory:
-git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
-
-python examples/inference/basic/basic.py
-```
-
-## Basic Walkthrough
-
-All you need to generate videos using multi-gpus from state-of-the-art diffusion pipelines is the following few lines!
-
-```python
-from fastvideo import VideoGenerator
-
-def main():
-    generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        num_gpus=1,
-    )
-
-    prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
-             "wide with interest. The playful yet serene atmosphere is complemented by soft "
-             "natural light filtering through the petals. Mid-shot, warm and cheerful tones.")
-    video = generator.generate_video(prompt)
-
-if __name__ == "__main__":
-    main()
-```
+The class provides the main python interface for using FastVideo's inference pipeline.

examples/inference/basic/basic.py

@@ -1,41 +1 @@
-from fastvideo import VideoGenerator
-
-# from fastvideo.v1.configs.sample import SamplingParam
-
-
-def main():
-    # FastVideo will automatically use the optimal default arguments for the
-    # model.
-    # If a local path is provided, FastVideo will make a best effort
-    # attempt to identify the optimal arguments.
-    generator = VideoGenerator.from_pretrained(
-        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        # if num_gpus > 1, FastVideo will automatically handle distributed setup
-        num_gpus=1,
-    )
-
-    # sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
-    # sampling_param.num_frames = 45
-    # sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
-    # Generate videos with the same simple API, regardless of GPU count
-    prompt = (
-        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
-        "wide with interest. The playful yet serene atmosphere is complemented by soft "
-        "natural light filtering through the petals. Mid-shot, warm and cheerful tones."
-    )
-    video = generator.generate_video(prompt)
-    # video = generator.generate_video(prompt, sampling_param=sampling_param, output_path="wan_t2v_videos/")
-
-    # Generate another video with a different prompt, without reloading the
-    # model!
-    prompt2 = (
-        "A majestic lion strides across the golden savanna, its powerful frame "
-        "glistening under the warm afternoon sun. The tall grass ripples gently in "
-        "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
-        "embodying the raw energy of the wild. Low angle, steady tracking shot, "
-        "cinematic.")
-    video2 = generator.generate_video(prompt2)
-
-
-if __name__ == "__main__":
-    main()
+print('Hello, world!')

examples/inference/basic/default_args.py

@@ -1,62 +0,0 @@
-from fastvideo import VideoGenerator
-
-
-def main():
-
-    # This is the config class for the model initialization
-    config = PipelineConfig.from_pretrained("FastVideo/FastHunyuan-Diffusers")
-    # can be used to dump the config to a yaml file
-    config.dump_to_yaml("config.yaml")
-    print(config)
-    # {
-    #    'vae_config': {
-    #        'scale_factor': 8,
-    #        'sp': True,
-    #        'tiling': True,
-    #        'precision': 'fp16'
-    #    },
-    #    'text_encoder_config': {
-    #        'precision': 'fp16'
-    #    },
-    #    'dit_config': {
-    #        'precision': 'fp16'
-    #    },
-    #    'inference_args': {
-    #        'guidance_scale': 7.5,
-    #        'num_inference_steps': 5,
-    #        'seed': 1024,
-    #        'guidance_rescale': 0.0,
-    #        'flow_shift': 17,
-    #        'num_inference_steps': 5,
-    #    }
-    # }
-
-    config.vae_config.scale_factor = 16
-
-    # FastVideo will automatically used the optimal default arguments for the model
-    # If a local path is provided, FastVideo will make a best effort attempt to
-    # identify the optimal arguments.
-    generator = VideoGenerator.from_pretrained(
-        "FastVideo/FastHunyuan-Diffusers",
-        num_gpus=4,
-        config=config,
-        # or
-        config_path="config.yaml",
-    )
-
-    sampling_param = SamplingParam.from_pretrained(
-        "FastVideo/FastHunyuan-Diffusers")
-    sampling_param.num_inference_steps = 5
-
-    # Generate videos with the same simple API, regardless of GPU count
-    prompt = "A beautiful woman in a red dress walking down a street"
-    video = generator.generate_video(prompt,
-                                     sampling_param=sampling_param,
-                                     num_inference_steps=6)
-
-    video2 = generator.generate_video(prompt2)
-    prompt2 = "A beautiful woman in a blue dress walking down a street"
-
-
-if __name__ == "__main__":
-    main()

examples/inference/gradio/gradio_demo.py

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

examples/inference/optimizations/README.md

@@ -1,9 +0,0 @@
-# Optimization Examples
-
-```bash
-python examples/inference/optimizations/attention_example.py
-```
-
-```bash
-python examples/inference/optimizations/teacache_example.py
-```

examples/inference/optimizations/attention_example.py

@@ -1,33 +0,0 @@
-import os
-import time
-
-from fastvideo import VideoGenerator
-
-def main():
-    # set the attention backend 
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN"
-
-    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.perf_counter() - start_time
-    print(f"Model loading time: {load_time:.2f} seconds")
-
-    gen_start_time = time.perf_counter()
-
-    gen.generate_video(
-        prompt=
-        "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
-        seed=1024,
-        output_path="example_outputs/")
-    
-    generation_time = time.perf_counter() - gen_start_time
-    print(f"Video generation time: {generation_time:.2f} seconds")
-
-    total_time = time.perf_counter() - start_time
-    print(f"Total execution time: {total_time:.2f} seconds")
-
-if __name__ == "__main__":
-    main()

examples/inference/optimizations/teacache_example.py

@@ -1,44 +0,0 @@
-import time
-
-from fastvideo import VideoGenerator, SamplingParam
-
-
-def main():
-    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.perf_counter() - start_time
-    print(f"Model loading time: {load_time:.2f} seconds")
-
-    gen_start_time = time.perf_counter()
-
-    params = SamplingParam.from_pretrained(
-        model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-    )
-    # this controls the threshold for the tea cache
-    params.teacache_params.teacache_thresh = 0.08
-    gen.generate_video(
-        prompt=
-        "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
-        sampling_param=params,
-        height=480,
-        width=832,
-        num_frames=61,  # 85 ,77 
-        num_inference_steps=50,
-        enable_teacache=True,
-        seed=1024,
-        output_path="example_outputs/")
-    
-    generation_time = time.perf_counter() - gen_start_time
-    print(f"Video generation time: {generation_time:.2f} seconds")
-
-    total_time = time.perf_counter() - start_time
-    print(f"Total execution time: {total_time:.2f} seconds")
-
-
-if __name__ == "__main__":
-    main()

fastvideo/__init__.py

@@ -1,6 +1,5 @@
 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", "__version__"]
+__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam"]

fastvideo/distill.py

@@ -452,7 +452,7 @@ def main(args):
         return phase
 
     for step in range(init_steps + 1, args.max_train_steps + 1):
-        start_time = time.perf_counter()
+        start_time = time.time()
         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.perf_counter() - start_time
+        step_time = time.time() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 

fastvideo/distill_adv.py

@@ -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.perf_counter()
+        start_time = time.time()
         (
             generator_loss,
             generator_grad_norm,
@@ -547,7 +547,7 @@ def main(args):
             args.discriminator_head_stride,
         )
 
-        step_time = time.perf_counter() - start_time
+        step_time = time.time() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 

fastvideo/models/hunyuan/idle_config.py

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

fastvideo/models/hunyuan/inference.py

@@ -17,7 +17,7 @@ from fastvideo.models.hunyuan.vae import load_vae
 from fastvideo.utils.parallel_states import nccl_info
 
 
-class Inference:
+class Inference(object):
 
     def __init__(
         self,
@@ -452,7 +452,7 @@ class HunyuanVideoSampler(Inference):
         # ========================================================================
         # Pipeline inference
         # ========================================================================
-        start_time = time.perf_counter()
+        start_time = time.time()
         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.perf_counter() - start_time
+        gen_time = time.time() - start_time
         logger.info(f"Success, time: {gen_time}")
 
         return out_dict

fastvideo/models/hunyuan/prompt_rewrite.py

@@ -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 Master mode, but got {}".format(mode))
+        raise Exception("Only supports Normal and Normal", mode)
     return prompt
 
 

fastvideo/models/stepvideo/text_encoder/stepllm.py

@@ -267,25 +267,25 @@ class Step1Model(PreTrainedModel):
 class STEP1TextEncoder(torch.nn.Module):
 
     def __init__(self, model_dir, max_length=320):
-        super()
+        super(STEP1TextEncoder, self).__init__()
         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
-        if type(prompts) is str:
-            prompts = [prompts]
+        with torch.no_grad(), torch.cuda.amp.autocast(dtype=torch.bfloat16):
+            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

fastvideo/train.py

@@ -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.perf_counter()
+        start_time = time.time()
         loss, grad_norm = train_one_step(
             transformer,
             args.model_type,
@@ -383,7 +383,7 @@ def main(args):
             args.mode_scale,
         )
 
-        step_time = time.perf_counter() - start_time
+        step_time = time.time() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 

fastvideo/utils/collect_env.py

@@ -1,776 +0,0 @@
-# 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()

fastvideo/utils/env_utils.py

@@ -0,0 +1,38 @@
+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")

fastvideo/v1/attention/backends/abstract.py

@@ -3,7 +3,8 @@
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, fields
-from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
+from typing import (TYPE_CHECKING, Any, Dict, Generic, Optional, Protocol, Set,
+                    Type, TypeVar)
 
 if TYPE_CHECKING:
     from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -26,12 +27,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
@@ -45,7 +46,7 @@ class AttentionBackend(ABC):
 
     @staticmethod
     @abstractmethod
-    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
+    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
         raise NotImplementedError
 
 
@@ -55,8 +56,23 @@ class AttentionMetadata:
     # Current step of diffusion process
     current_timestep: int
 
+    # @property
+    # @abstractmethod
+    # def inference_metadata(self) -> Optional["AttentionMetadata"]:
+    #     """Return the attention metadata that's required to run prefill
+    #     attention."""
+    #     pass
+
+    # @property
+    # @abstractmethod
+    # def training_metadata(self) -> Optional["AttentionMetadata"]:
+    #     """Return the attention metadata that's required to run decode
+    #     attention."""
+    #     pass
+
     def asdict_zerocopy(self,
-                        skip_fields: set[str] | None = None) -> dict[str, Any]:
+                        skip_fields: Optional[Set[str]] = None
+                        ) -> Dict[str, Any]:
         """Similar to dataclasses.asdict, but avoids deepcopying."""
         if skip_fields is None:
             skip_fields = set()
@@ -70,6 +86,55 @@ class AttentionMetadata:
 
 T = TypeVar("T", bound=AttentionMetadata)
 
+# class AttentionState(ABC, Generic[T]):
+#     """Holds attention backend-specific objects reused during the
+#     lifetime of the model runner."""
+
+#     @abstractmethod
+#     def __init__(self, runner: "ModelRunnerBase"):
+#         ...
+
+#     @abstractmethod
+#     @contextmanager
+#     def graph_capture(self, max_batch_size: int):
+#         """Context manager used when capturing CUDA graphs."""
+#         yield
+
+#     @abstractmethod
+#     def graph_clone(self, batch_size: int) -> "AttentionState[T]":
+#         """Clone attention state to save in CUDA graph metadata."""
+#         ...
+
+#     @abstractmethod
+#     def graph_capture_get_metadata_for_batch(
+#             self,
+#             batch_size: int,
+#             is_encoder_decoder_model: bool = False) -> T:
+#         """Get attention metadata for CUDA graph capture of batch_size."""
+#         ...
+
+#     @abstractmethod
+#     def get_graph_input_buffers(
+#             self,
+#             attn_metadata: T,
+#             is_encoder_decoder_model: bool = False) -> Dict[str, Any]:
+#         """Get attention-specific input buffers for CUDA graph capture."""
+#         ...
+
+#     @abstractmethod
+#     def prepare_graph_input_buffers(
+#             self,
+#             input_buffers: Dict[str, Any],
+#             attn_metadata: T,
+#             is_encoder_decoder_model: bool = False) -> None:
+#         """In-place modify input buffers dict for CUDA graph replay."""
+#         ...
+
+#     @abstractmethod
+#     def begin_forward(self, model_input: "ModelRunnerInputBase") -> None:
+#         """Prepare state for forward pass."""
+#         ...
+
 
 class AttentionMetadataBuilder(ABC, Generic[T]):
     """Abstract class for attention metadata builders."""
@@ -122,7 +187,7 @@ class AttentionImpl(ABC, Generic[T]):
         head_size: int,
         softmax_scale: float,
         causal: bool = False,
-        num_kv_heads: int | None = None,
+        num_kv_heads: Optional[int] = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:

fastvideo/v1/attention/backends/flash_attn.py

@@ -1,5 +1,7 @@
 # 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
 
@@ -26,7 +28,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
@@ -34,15 +36,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
 
 
@@ -54,7 +56,7 @@ class FlashAttentionImpl(AttentionImpl):
         head_size: int,
         causal: bool,
         softmax_scale: float,
-        num_kv_heads: int | None = None,
+        num_kv_heads: Optional[int] = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:

fastvideo/v1/attention/backends/sage_attn.py

@@ -1,3 +1,5 @@
+from typing import List, Optional, Type
+
 import torch
 from sageattention import sageattn
 
@@ -15,7 +17,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
@@ -23,7 +25,7 @@ class SageAttentionBackend(AttentionBackend):
         return "SAGE_ATTN"
 
     @staticmethod
-    def get_impl_cls() -> type["SageAttentionImpl"]:
+    def get_impl_cls() -> Type["SageAttentionImpl"]:
         return SageAttentionImpl
 
     # @staticmethod
@@ -39,7 +41,7 @@ class SageAttentionImpl(AttentionImpl):
         head_size: int,
         causal: bool,
         softmax_scale: float,
-        num_kv_heads: int | None = None,
+        num_kv_heads: Optional[int] = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:

fastvideo/v1/attention/backends/sdpa.py

@@ -1,3 +1,5 @@
+from typing import List, Optional, Type
+
 import torch
 
 from fastvideo.v1.attention.backends.abstract import (
@@ -14,7 +16,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
@@ -22,7 +24,7 @@ class SDPABackend(AttentionBackend):
         return "SDPA"
 
     @staticmethod
-    def get_impl_cls() -> type["SDPAImpl"]:
+    def get_impl_cls() -> Type["SDPAImpl"]:
         return SDPAImpl
 
     # @staticmethod
@@ -38,7 +40,7 @@ class SDPAImpl(AttentionImpl):
         head_size: int,
         causal: bool,
         softmax_scale: float,
-        num_kv_heads: int | None = None,
+        num_kv_heads: Optional[int] = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:

fastvideo/v1/attention/backends/sliding_tile_attn.py

@@ -1,5 +1,6 @@
 import json
 from dataclasses import dataclass
+from typing import List, Optional, Type
 
 import torch
 from einops import rearrange
@@ -19,7 +20,7 @@ logger = init_logger(__name__)
 
 
 # TODO(will-refactor): move this to a utils file
-def dict_to_3d_list(mask_strategy) -> list[list[list[torch.Tensor | None]]]:
+def dict_to_3d_list(mask_strategy) -> List[List[List[Optional[torch.Tensor]]]]:
     indices = [tuple(map(int, key.split('_'))) for key in mask_strategy]
 
     max_timesteps_idx = max(
@@ -57,7 +58,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]
 
@@ -66,15 +67,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
 
 
@@ -109,7 +110,7 @@ class SlidingTileAttentionImpl(AttentionImpl):
         head_size: int,
         causal: bool,
         softmax_scale: float,
-        num_kv_heads: int | None = None,
+        num_kv_heads: Optional[int] = None,
         prefix: str = "",
         **extra_impl_args,
     ) -> None:

fastvideo/v1/attention/layer.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Optional, Tuple
+
 import torch
 import torch.nn as nn
 
@@ -20,11 +22,11 @@ class DistributedAttention(nn.Module):
     def __init__(self,
                  num_heads: int,
                  head_size: int,
-                 num_kv_heads: int | None = None,
-                 softmax_scale: float | None = None,
+                 num_kv_heads: Optional[int] = None,
+                 softmax_scale: Optional[float] = None,
                  causal: bool = False,
-                 supported_attention_backends: tuple[_Backend, ...]
-                 | None = None,
+                 supported_attention_backends: Optional[Tuple[_Backend,
+                                                              ...]] = None,
                  prefix: str = "",
                  **extra_impl_args) -> None:
         super().__init__()
@@ -60,10 +62,10 @@ class DistributedAttention(nn.Module):
         q: torch.Tensor,
         k: torch.Tensor,
         v: 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]:
+        replicated_q: Optional[torch.Tensor] = None,
+        replicated_k: Optional[torch.Tensor] = None,
+        replicated_v: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         """Forward pass for distributed attention.
         
         Args:
@@ -82,6 +84,9 @@ 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()
@@ -139,11 +144,11 @@ class LocalAttention(nn.Module):
     def __init__(self,
                  num_heads: int,
                  head_size: int,
-                 num_kv_heads: int | None = None,
-                 softmax_scale: float | None = None,
+                 num_kv_heads: Optional[int] = None,
+                 softmax_scale: Optional[float] = None,
                  causal: bool = False,
-                 supported_attention_backends: tuple[_Backend, ...]
-                 | None = None,
+                 supported_attention_backends: Optional[Tuple[_Backend,
+                                                              ...]] = None,
                  **extra_impl_args) -> None:
         super().__init__()
         if softmax_scale is None:

fastvideo/v1/attention/selector.py

@@ -2,10 +2,9 @@
 # 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 cast
+from typing import Generator, Optional, Tuple, Type, cast
 
 import torch
 
@@ -18,7 +17,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) -> _Backend | None:
+def backend_name_to_enum(backend_name: str) -> Optional[_Backend]:
     """
     Convert a string backend name to a _Backend enum value.
 
@@ -32,7 +31,7 @@ def backend_name_to_enum(backend_name: str) -> _Backend | None:
           None
 
 
-def get_env_variable_attn_backend() -> _Backend | None:
+def get_env_variable_attn_backend() -> Optional[_Backend]:
     '''
     Get the backend override specified by the FastVideo attention
     backend environment variable, if one is specified.
@@ -54,10 +53,10 @@ def get_env_variable_attn_backend() -> _Backend | None:
 #
 # THIS SELECTION TAKES PRECEDENCE OVER THE
 # FASTVIDEO ATTENTION BACKEND ENVIRONMENT VARIABLE
-forced_attn_backend: _Backend | None = None
+forced_attn_backend: Optional[_Backend] = None
 
 
-def global_force_attn_backend(attn_backend: _Backend | None) -> None:
+def global_force_attn_backend(attn_backend: Optional[_Backend]) -> None:
     '''
     Force all attention operations to use a specified backend.
 
@@ -72,7 +71,7 @@ def global_force_attn_backend(attn_backend: _Backend | None) -> None:
     forced_attn_backend = attn_backend
 
 
-def get_global_forced_attn_backend() -> _Backend | None:
+def get_global_forced_attn_backend() -> Optional[_Backend]:
     '''
     Get the currently-forced choice of attention backend,
     or None if auto-selection is currently enabled.
@@ -83,8 +82,8 @@ def get_global_forced_attn_backend() -> _Backend | None:
 def get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
-    supported_attention_backends: tuple[_Backend, ...] | None = None,
-) -> type[AttentionBackend]:
+    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
+) -> Type[AttentionBackend]:
     return _cached_get_attn_backend(head_size, dtype,
                                     supported_attention_backends)
 
@@ -93,8 +92,8 @@ def get_attn_backend(
 def _cached_get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
-    supported_attention_backends: tuple[_Backend, ...] | None = None,
-) -> type[AttentionBackend]:
+    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
+) -> Type[AttentionBackend]:
     # Check whether a particular choice of backend was
     # previously forced.
     #
@@ -103,13 +102,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: _Backend | None = (
+    backend_by_global_setting: Optional[_Backend] = (
         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: str | None = envs.FASTVIDEO_ATTENTION_BACKEND
+        backend_by_env_var: Optional[str] = envs.FASTVIDEO_ATTENTION_BACKEND
         if backend_by_env_var is not None:
             selected_backend = backend_name_to_enum(backend_by_env_var)
 
@@ -121,7 +120,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

fastvideo/v1/configs/fasthunyuan_t2v.json

@@ -1,48 +0,0 @@
-{
-  "embedded_cfg_scale": 6,
-  "flow_shift": 17,
-  "use_cpu_offload": false,
-  "disable_autocast": false,
-  "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,
-    "tile_sample_min_num_frames": 16,
-    "tile_sample_stride_height": 192,
-    "tile_sample_stride_width": 192,
-    "tile_sample_stride_num_frames": 12,
-    "blend_num_frames": 4,
-    "use_tiling": true,
-    "use_temporal_tiling": true,
-    "use_parallel_tiling": true
-  },
-  "dit_config": {
-    "prefix": "Hunyuan",
-    "quant_config": null
-  },
-  "text_encoder_precisions": [
-    "fp16",
-    "fp16"
-  ],
-  "text_encoder_configs": [
-    {
-      "prefix": "llama",
-      "quant_config": null,
-      "lora_config": null
-    },
-    {
-      "prefix": "clip",
-      "quant_config": null,
-      "lora_config": null,
-      "num_hidden_layers_override": null,
-      "require_post_norm": null
-    }
-  ],
-  "mask_strategy_file_path": null,
-  "enable_torch_compile": false
-}

fastvideo/v1/configs/models/base.py

@@ -1,5 +1,5 @@
-from dataclasses import dataclass, field, fields
-from typing import Any
+from dataclasses import dataclass, fields
+from typing import Any, Dict
 
 from fastvideo.v1.logger import init_logger
 
@@ -18,7 +18,7 @@ class ArchConfig:
 class ModelConfig:
     # Every model config parameter can be categorized into either ArchConfig or everything else
     # Diffuser/Transformer parameters
-    arch_config: ArchConfig = field(default_factory=ArchConfig)
+    arch_config: ArchConfig = ArchConfig()
 
     # FastVideo-specific parameters here
     # i.e. STA, quantization, teacache
@@ -30,18 +30,8 @@ class ModelConfig:
         raise AttributeError(
             f"'{type(self).__name__}' object has no attribute '{name}'")
 
-    def __getstate__(self):
-        # Return a dictionary of attributes to pickle
-        # Convert to dict and exclude any problematic attributes
-        state = self.__dict__.copy()
-        return state
-
-    def __setstate__(self, state):
-        # Restore instance attributes from the unpickled state
-        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 +45,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)}

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

@@ -1,5 +1,4 @@
 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", "StepVideoConfig"]
+__all__ = ["HunyuanVideoConfig", "WanVideoConfig"]

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

@@ -1,17 +1,16 @@
 from dataclasses import dataclass, field
-from typing import Any
+from typing import Optional, Tuple
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.layers.quantization import QuantizationConfig
+from fastvideo.v1.configs.quantization import QuantizationConfig
 from fastvideo.v1.platforms import _Backend
 
 
 @dataclass
 class DiTArchConfig(ArchConfig):
     _fsdp_shard_conditions: list = field(default_factory=list)
-    _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,
@@ -21,36 +20,11 @@ class DiTArchConfig(ArchConfig):
     num_attention_heads: int = 0
     num_channels_latents: int = 0
 
-    def __post_init__(self) -> None:
-        if not self._compile_conditions:
-            self._compile_conditions = self._fsdp_shard_conditions.copy()
-
 
 @dataclass
 class DiTConfig(ModelConfig):
-    arch_config: DiTArchConfig = field(default_factory=DiTArchConfig)
+    arch_config: DiTArchConfig = DiTArchConfig()
 
     # FastVideoDiT-specific parameters
     prefix: str = ""
-    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
+    quant_config: Optional[QuantizationConfig] = None

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

@@ -1,4 +1,5 @@
 from dataclasses import dataclass, field
+from typing import Optional, Tuple
 
 import torch
 
@@ -17,19 +18,12 @@ def is_refiner_block(n: str, m) -> bool:
     return "refiner" in n and str.isdigit(n.split(".")[-1])
 
 
-def is_txt_in(n: str, m) -> bool:
-    return n.split(".")[-1] == "txt_in"
-
-
 @dataclass
 class HunyuanVideoArchConfig(DiTArchConfig):
     _fsdp_shard_conditions: list = field(
         default_factory=lambda:
         [is_double_block, is_single_block, is_refiner_block])
 
-    _compile_conditions: list = field(
-        default_factory=lambda: [is_double_block, is_single_block, is_txt_in])
-
     _param_names_mapping: dict = field(
         default_factory=lambda: {
             # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
@@ -155,22 +149,21 @@ 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: torch.dtype | None = None
+    dtype: Optional[torch.dtype] = None
     text_embed_dim: int = 4096
     pooled_projection_dim: int = 768
     rope_theta: int = 256
     qk_norm: str = "rms_norm"
 
     def __post_init__(self):
-        super().__post_init__()
         self.hidden_size: int = self.attention_head_dim * self.num_attention_heads
         self.num_channels_latents: int = self.in_channels
 
 
 @dataclass
 class HunyuanVideoConfig(DiTConfig):
-    arch_config: DiTArchConfig = field(default_factory=HunyuanVideoArchConfig)
+    arch_config: DiTArchConfig = HunyuanVideoArchConfig()
 
     prefix: str = "Hunyuan"

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

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

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

@@ -1,4 +1,5 @@
 from dataclasses import dataclass, field
+from typing import Optional, Tuple
 
 from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
 
@@ -51,7 +52,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
@@ -64,12 +65,11 @@ class WanVideoArchConfig(DiTArchConfig):
     cross_attn_norm: bool = True
     qk_norm: str = "rms_norm_across_heads"
     eps: float = 1e-6
-    image_dim: int | None = None
-    added_kv_proj_dim: int | None = None
+    image_dim: Optional[int] = None
+    added_kv_proj_dim: Optional[int] = None
     rope_max_seq_len: int = 1024
 
     def __post_init__(self):
-        super().__post_init__()
         self.out_channels = self.out_channels or self.in_channels
         self.hidden_size = self.num_attention_heads * self.attention_head_dim
         self.num_channels_latents = self.in_channels if self.added_kv_proj_dim is None else self.out_channels
@@ -77,6 +77,6 @@ class WanVideoArchConfig(DiTArchConfig):
 
 @dataclass
 class WanVideoConfig(DiTConfig):
-    arch_config: DiTArchConfig = field(default_factory=WanVideoArchConfig)
+    arch_config: DiTArchConfig = WanVideoArchConfig()
 
     prefix: str = "Wan"

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

@@ -1,5 +1,4 @@
-from fastvideo.v1.configs.models.encoders.base import (BaseEncoderOutput,
-                                                       EncoderConfig,
+from fastvideo.v1.configs.models.encoders.base import (EncoderConfig,
                                                        ImageEncoderConfig,
                                                        TextEncoderConfig)
 from fastvideo.v1.configs.models.encoders.clip import (CLIPTextConfig,
@@ -9,6 +8,5 @@ from fastvideo.v1.configs.models.encoders.t5 import T5Config
 
 __all__ = [
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig",
-    "BaseEncoderOutput", "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig",
-    "T5Config"
+    "CLIPTextConfig", "CLIPVisionConfig", "LlamaConfig", "T5Config"
 ]

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

@@ -1,17 +1,15 @@
 from dataclasses import dataclass, field
-from typing import Any
-
-import torch
+from typing import Any, List, Optional, Tuple
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.layers.quantization import QuantizationConfig
+from fastvideo.v1.configs.quantization import QuantizationConfig
 from fastvideo.v1.platforms import _Backend
 
 
 @dataclass
 class EncoderArchConfig(ArchConfig):
-    architectures: list[str] = field(default_factory=lambda: [])
-    _supported_attention_backends: tuple[_Backend, ...] = (_Backend.FLASH_ATTN,
+    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,44 +30,26 @@ class TextEncoderArchConfig(EncoderArchConfig):
     scalable_attention: bool = True
     tie_word_embeddings: bool = False
 
-    tokenizer_kwargs: dict[str, Any] = field(default_factory=dict)
-
-    def __post_init__(self) -> None:
-        self.tokenizer_kwargs = {
-            "truncation": True,
-            "max_length": self.text_len,
-            "return_tensors": "pt",
-        }
-
 
 @dataclass
 class ImageEncoderArchConfig(EncoderArchConfig):
     pass
 
 
-@dataclass
-class BaseEncoderOutput:
-    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
 class EncoderConfig(ModelConfig):
-    arch_config: ArchConfig = field(default_factory=EncoderArchConfig)
+    arch_config: ArchConfig = EncoderArchConfig()
 
     prefix: str = ""
-    quant_config: QuantizationConfig | None = None
-    lora_config: Any | None = None
+    quant_config: Optional[QuantizationConfig] = None
+    lora_config: Optional[Any] = None
 
 
 @dataclass
 class TextEncoderConfig(EncoderConfig):
-    arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)
+    arch_config: ArchConfig = TextEncoderArchConfig()
 
 
 @dataclass
 class ImageEncoderConfig(EncoderConfig):
-    arch_config: ArchConfig = field(default_factory=ImageEncoderArchConfig)
+    arch_config: ArchConfig = ImageEncoderArchConfig()

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

@@ -1,4 +1,5 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
+from typing import Optional
 
 from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                        ImageEncoderConfig,
@@ -47,19 +48,17 @@ class CLIPVisionArchConfig(ImageEncoderArchConfig):
 
 @dataclass
 class CLIPTextConfig(TextEncoderConfig):
-    arch_config: TextEncoderArchConfig = field(
-        default_factory=CLIPTextArchConfig)
+    arch_config: TextEncoderArchConfig = CLIPTextArchConfig()
 
-    num_hidden_layers_override: int | None = None
-    require_post_norm: bool | None = None
+    num_hidden_layers_override: Optional[int] = None
+    require_post_norm: Optional[bool] = None
     prefix: str = "clip"
 
 
 @dataclass
 class CLIPVisionConfig(ImageEncoderConfig):
-    arch_config: ImageEncoderArchConfig = field(
-        default_factory=CLIPVisionArchConfig)
+    arch_config: ImageEncoderArchConfig = CLIPVisionArchConfig()
 
-    num_hidden_layers_override: int | None = None
-    require_post_norm: bool | None = None
+    num_hidden_layers_override: Optional[int] = None
+    require_post_norm: Optional[bool] = None
     prefix: str = "clip"

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

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

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

@@ -1,4 +1,5 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
+from typing import Optional
 
 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                        TextEncoderConfig)
@@ -11,7 +12,7 @@ class T5ArchConfig(TextEncoderArchConfig):
     d_kv: int = 64
     d_ff: int = 2048
     num_layers: int = 6
-    num_decoder_layers: int | None = None
+    num_decoder_layers: Optional[int] = None
     num_heads: int = 8
     relative_attention_num_buckets: int = 32
     relative_attention_max_distance: int = 128
@@ -30,25 +31,15 @@ class T5ArchConfig(TextEncoderArchConfig):
 
     # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py
     def __post_init__(self):
-        super().__post_init__()
         act_info = self.feed_forward_proj.split("-")
         self.dense_act_fn: str = act_info[-1]
         self.is_gated_act: bool = act_info[0] == "gated"
         if self.feed_forward_proj == "gated-gelu":
             self.dense_act_fn = "gelu_new"
 
-        self.tokenizer_kwargs = {
-            "padding": "max_length",
-            "truncation": True,
-            "max_length": self.text_len,
-            "add_special_tokens": True,
-            "return_attention_mask": True,
-            "return_tensors": "pt",
-        }
-
 
 @dataclass
 class T5Config(TextEncoderConfig):
-    arch_config: TextEncoderArchConfig = field(default_factory=T5ArchConfig)
+    arch_config: TextEncoderArchConfig = T5ArchConfig()
 
     prefix: str = "t5"

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

@@ -1,9 +1,7 @@
 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",
 ]

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

@@ -1,15 +1,14 @@
-from dataclasses import dataclass, field
-from typing import Any
+from dataclasses import dataclass
+from typing import Union
 
 import torch
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.utils import StoreBoolean
 
 
 @dataclass
 class VAEArchConfig(ArchConfig):
-    scaling_factor: float | torch.Tensor = 0
+    scaling_factor: Union[float, torch.tensor] = 0
 
     temporal_compression_ratio: int = 4
     spatial_compression_ratio: int = 8
@@ -17,7 +16,7 @@ class VAEArchConfig(ArchConfig):
 
 @dataclass
 class VAEConfig(ModelConfig):
-    arch_config: VAEArchConfig = field(default_factory=VAEArchConfig)
+    arch_config: VAEArchConfig = VAEArchConfig()
 
     # FastVideoVAE-specific parameters
     load_encoder: bool = True
@@ -34,97 +33,6 @@ 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

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

@@ -1,4 +1,5 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
+from typing import Tuple
 
 from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
 
@@ -8,19 +9,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
@@ -36,4 +37,4 @@ class HunyuanVAEArchConfig(VAEArchConfig):
 
 @dataclass
 class HunyuanVAEConfig(VAEConfig):
-    arch_config: VAEArchConfig = field(default_factory=HunyuanVAEArchConfig)
+    arch_config: VAEArchConfig = HunyuanVAEArchConfig()

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

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

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

@@ -1,4 +1,5 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
+from typing import Tuple
 
 import torch
 
@@ -9,12 +10,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,
@@ -32,7 +33,7 @@ class WanVAEArchConfig(VAEArchConfig):
         0.2503,
         -0.2921,
     )
-    latents_std: tuple[float, ...] = (
+    latents_std: Tuple[float, ...] = (
         2.8184,
         1.4541,
         2.3275,
@@ -54,15 +55,15 @@ 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)
 
 
 @dataclass
 class WanVAEConfig(VAEConfig):
-    arch_config: VAEArchConfig = field(default_factory=WanVAEArchConfig)
+    arch_config: VAEArchConfig = WanVAEArchConfig()
     use_feature_cache: bool = True
 
     use_tiling: bool = False

fastvideo/v1/configs/pipelines/__init__.py

@@ -4,15 +4,11 @@ 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,
-                                                WanT2V720PConfig)
+                                                WanT2V480PConfig)
 
 __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "PipelineConfig",
     "SlidingTileAttnConfig", "WanT2V480PConfig", "WanI2V480PConfig",
-    "WanT2V720PConfig", "WanI2V720PConfig", "StepVideoT2VConfig",
     "get_pipeline_config_cls_for_name"
 ]

fastvideo/v1/configs/pipelines/base.py

@@ -1,33 +1,21 @@
 import json
-from collections.abc import Callable
-from dataclasses import asdict, dataclass, field, fields
-from typing import Any, cast
-
-import torch
+from dataclasses import asdict, dataclass, fields
+from typing import Any, Dict, Optional
 
 from fastvideo.v1.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
                                          VAEConfig)
-from fastvideo.v1.configs.models.encoders import BaseEncoderOutput
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import shallow_asdict
 
 logger = init_logger(__name__)
 
 
-def preprocess_text(prompt: str) -> str:
-    return prompt
-
-
-def postprocess_text(output: BaseEncoderOutput) -> torch.Tensor:
-    raise NotImplementedError
-
-
 @dataclass
 class PipelineConfig:
     """Base configuration for all pipeline architectures."""
     # Video generation parameters
     embedded_cfg_scale: float = 6.0
-    flow_shift: float | None = None
+    flow_shift: Optional[float] = None
     use_cpu_offload: bool = False
     disable_autocast: bool = False
 
@@ -38,26 +26,17 @@ class PipelineConfig:
     vae_precision: str = "fp16"
     vae_tiling: bool = True
     vae_sp: bool = True
-    vae_config: VAEConfig = field(default_factory=VAEConfig)
+    vae_config: VAEConfig = VAEConfig()
 
     # DiT configuration
-    dit_config: DiTConfig = field(default_factory=DiTConfig)
+    dit_config: DiTConfig = DiTConfig()
 
     # Text encoder configuration
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("fp16", ))
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(
-        default_factory=lambda: (EncoderConfig(), ))
-    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
-        default_factory=lambda: (preprocess_text, ))
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda:
-                                               (postprocess_text, ))
+    text_encoder_precision: str = "fp16"
+    text_encoder_config: EncoderConfig = EncoderConfig()
 
     # STA (Spatial-Temporal Attention) parameters
-    mask_strategy_file_path: str | None = None
-
-    # Compilation
+    mask_strategy_file_path: Optional[str] = None
     enable_torch_compile: bool = False
 
     @classmethod
@@ -73,32 +52,16 @@ class PipelineConfig:
                 model_path)
             pipeline_config = cls()
 
-        return cast(PipelineConfig, pipeline_config)
+        return pipeline_config
 
     def dump_to_json(self, file_path: str):
         output_dict = shallow_asdict(self)
-        del_keys = []
         for key, value in output_dict.items():
             if isinstance(value, ModelConfig):
                 model_dict = asdict(value)
                 # Model Arch Config should be hidden away from the users
                 model_dict.pop("arch_config")
                 output_dict[key] = model_dict
-            elif isinstance(value, tuple) and all(
-                    isinstance(v, ModelConfig) for v in value):
-                model_dicts = []
-                for v in value:
-                    model_dict = asdict(v)
-                    # Model Arch Config should be hidden away from the users
-                    model_dict.pop("arch_config")
-                    model_dicts.append(model_dict)
-                output_dict[key] = model_dicts
-            elif isinstance(value, tuple) and all(callable(f) for f in value):
-                # Skip dumping functions
-                del_keys.append(key)
-
-        for key in del_keys:
-            output_dict.pop(key, None)
 
         with open(file_path, "w") as f:
             json.dump(output_dict, f, indent=2)
@@ -108,7 +71,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
@@ -119,15 +82,6 @@ class PipelineConfig:
                 # If it's a nested ModelConfig, update it recursively
                 if isinstance(current_value, ModelConfig):
                     current_value.update_model_config(new_value)
-                elif isinstance(current_value, tuple) and all(
-                        isinstance(v, ModelConfig) for v in current_value):
-                    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,
-                                                            strict=False):
-                        target_config.update_model_config(source_config)
                 else:
                     setattr(self, key, new_value)
 

fastvideo/v1/configs/pipelines/hunyuan.py

@@ -1,60 +1,11 @@
-from collections.abc import Callable
-from dataclasses import dataclass, field
-from typing import TypedDict
-
-import torch
+from dataclasses import dataclass
 
 from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.v1.configs.models.dits import HunyuanVideoConfig
-from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
-                                                  CLIPTextConfig, LlamaConfig)
+from fastvideo.v1.configs.models.encoders import CLIPTextConfig, LlamaConfig
 from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig
 from fastvideo.v1.configs.pipelines.base import PipelineConfig
 
-PROMPT_TEMPLATE_ENCODE_VIDEO = (
-    "<|start_header_id|>system<|end_header_id|>\n\nDescribe the video by detailing the following aspects: "
-    "1. The main content and theme of the video."
-    "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects."
-    "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects."
-    "4. background environment, light, style and atmosphere."
-    "5. camera angles, movements, and transitions used in the video:<|eot_id|>"
-    "<|start_header_id|>user<|end_header_id|>\n\n{}<|eot_id|>")
-
-
-class PromptTemplate(TypedDict):
-    template: str
-    crop_start: int
-
-
-prompt_template_video: PromptTemplate = {
-    "template": PROMPT_TEMPLATE_ENCODE_VIDEO,
-    "crop_start": 95,
-}
-
-
-def llama_preprocess_text(prompt: str) -> str:
-    return prompt_template_video["template"].format(prompt)
-
-
-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 +
-                                                      1)]
-    crop_start = prompt_template_video.get("crop_start", -1)
-    last_hidden_state = last_hidden_state[:, crop_start:]
-    return last_hidden_state
-
-
-def clip_preprocess_text(prompt: str) -> str:
-    return prompt
-
-
-def clip_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
-    pooler_output: torch.Tensor = outputs.pooler_output
-    return pooler_output
-
 
 @dataclass
 class HunyuanConfig(PipelineConfig):
@@ -62,31 +13,25 @@ class HunyuanConfig(PipelineConfig):
 
     # HunyuanConfig-specific parameters with defaults
     # DiT
-    dit_config: DiTConfig = field(default_factory=HunyuanVideoConfig)
+    dit_config: DiTConfig = HunyuanVideoConfig()
     # VAE
-    vae_config: VAEConfig = field(default_factory=HunyuanVAEConfig)
+    vae_config: VAEConfig = HunyuanVAEConfig()
     # Denoising stage
     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(
-        default_factory=lambda: (LlamaConfig(), CLIPTextConfig()))
-    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],
-        ...] = field(default_factory=lambda:
-                     (llama_postprocess_text, clip_postprocess_text))
+    text_encoder_config: EncoderConfig = LlamaConfig()
 
     # Precision for each component
     precision: str = "bf16"
     vae_precision: str = "fp16"
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("fp16", "fp16"))
+    text_encoder_precision: str = "fp16"
+
+    # HunyuanConfig-specific added parameters
+    # Secondary text encoder
+    text_encoder_config_2: EncoderConfig = CLIPTextConfig()
+    text_encoder_precision_2: str = "fp16"
 
     def __post_init__(self):
         self.vae_config.load_encoder = False

fastvideo/v1/configs/pipelines/registry.py

@@ -1,16 +1,13 @@
 """Registry for pipeline weight-specific configurations."""
 
 import os
-from collections.abc import Callable
+from typing import Callable, Dict, Optional, Type
 
 from fastvideo.v1.configs.pipelines.base import PipelineConfig
 from fastvideo.v1.configs.pipelines.hunyuan import (FastHunyuanConfig,
                                                     HunyuanConfig)
-from fastvideo.v1.configs.pipelines.stepvideo import StepVideoT2VConfig
 from fastvideo.v1.configs.pipelines.wan import (WanI2V480PConfig,
-                                                WanI2V720PConfig,
-                                                WanT2V480PConfig,
-                                                WanT2V720PConfig)
+                                                WanT2V480PConfig)
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import (maybe_download_model_index,
                                 verify_model_config_and_directory)
@@ -18,40 +15,35 @@ 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,
+    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PConfig
     # Add other specific weight variants
 }
 
 # For determining pipeline type from model ID
-PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = {
+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) -> type[PipelineConfig] | None:
+        pipeline_name_or_path: str) -> Optional[type[PipelineConfig]]:
     """Get the appropriate config class for specific pretrained weights."""
 
     if os.path.exists(pipeline_name_or_path):

fastvideo/v1/configs/pipelines/stepvideo.py

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

fastvideo/v1/configs/pipelines/wan.py

@@ -1,41 +1,21 @@
-from collections.abc import Callable
-from dataclasses import dataclass, field
-
-import torch
+from dataclasses import dataclass
 
 from fastvideo.v1.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.v1.configs.models.dits import WanVideoConfig
-from fastvideo.v1.configs.models.encoders import (BaseEncoderOutput,
-                                                  CLIPVisionConfig, T5Config)
+from fastvideo.v1.configs.models.encoders import CLIPVisionConfig, T5Config
 from fastvideo.v1.configs.models.vaes import WanVAEConfig
 from fastvideo.v1.configs.pipelines.base import PipelineConfig
 
 
-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, 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
-    ],
-                                                     dim=0)
-    return prompt_embeds_tensor
-
-
 @dataclass
 class WanT2V480PConfig(PipelineConfig):
     """Base configuration for Wan T2V 1.3B pipeline architecture."""
 
     # WanConfig-specific parameters with defaults
     # DiT
-    dit_config: DiTConfig = field(default_factory=WanVideoConfig)
+    dit_config: DiTConfig = WanVideoConfig()
     # VAE
-    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    vae_config: VAEConfig = WanVAEConfig()
     vae_tiling: bool = False
     vae_sp: bool = False
 
@@ -46,17 +26,12 @@ class WanT2V480PConfig(PipelineConfig):
     flow_shift: int = 3
 
     # Text encoding stage
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(
-        default_factory=lambda: (T5Config(), ))
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
-                                  ...] = field(default_factory=lambda:
-                                               (t5_postprocess_text, ))
+    text_encoder_config: EncoderConfig = T5Config()
 
     # Precision for each component
     precision: str = "bf16"
     vae_precision: str = "fp16"
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("fp32", ))
+    text_encoder_precision: str = "fp32"
 
     # WanConfig-specific added parameters
 
@@ -65,16 +40,6 @@ class WanT2V480PConfig(PipelineConfig):
         self.vae_config.load_decoder = True
 
 
-@dataclass
-class WanT2V720PConfig(WanT2V480PConfig):
-    """Base configuration for Wan T2V 14B 720P pipeline architecture."""
-
-    # WanConfig-specific parameters with defaults
-
-    # Denoising stage
-    flow_shift: int = 5
-
-
 @dataclass
 class WanI2V480PConfig(WanT2V480PConfig):
     """Base configuration for Wan I2V 14B 480P pipeline architecture."""
@@ -82,20 +47,9 @@ class WanI2V480PConfig(WanT2V480PConfig):
     # WanConfig-specific parameters with defaults
 
     # Precision for each component
-    image_encoder_config: EncoderConfig = field(
-        default_factory=CLIPVisionConfig)
+    image_encoder_config: EncoderConfig = CLIPVisionConfig()
     image_encoder_precision: str = "fp32"
 
     def __post_init__(self):
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
-
-
-@dataclass
-class WanI2V720PConfig(WanI2V480PConfig):
-    """Base configuration for Wan I2V 14B 720P pipeline architecture."""
-
-    # WanConfig-specific parameters with defaults
-
-    # Denoising stage
-    flow_shift: int = 5

fastvideo/v1/configs/quantization/__init__.py

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

fastvideo/v1/configs/quantization/base.py

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

fastvideo/v1/configs/sample/base.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Dict, List, Optional, Union
 
 from fastvideo.v1.logger import init_logger
 
@@ -8,19 +8,16 @@ logger = init_logger(__name__)
 
 @dataclass
 class SamplingParam:
-    """
-    Sampling parameters for video generation.
-    """
     # All fields below are copied from ForwardBatch
     data_type: str = "video"
 
     # Image inputs
-    image_path: str | None = None
+    image_path: Optional[str] = None
 
     # Text inputs
-    prompt: str | list[str] | None = None
-    negative_prompt: str | None = None
-    prompt_path: str | None = None
+    prompt: Optional[Union[str, List[str]]] = None
+    negative_prompt: Optional[str] = None
+    prompt_path: Optional[str] = None
     output_path: str = "outputs/"
 
     # Batch info
@@ -29,7 +26,6 @@ class SamplingParam:
 
     # Original dimensions (before VAE scaling)
     num_frames: int = 125
-    num_frames_round_down: bool = False  # Whether to round down num_frames if it's not divisible by num_gpus
     height: int = 720
     width: int = 1280
     fps: int = 24
@@ -39,9 +35,6 @@ class SamplingParam:
     guidance_scale: float = 1.0
     guidance_rescale: float = 0.0
 
-    # TeaCache parameters
-    enable_teacache: bool = False
-
     # Misc
     save_video: bool = True
     return_frames: bool = False
@@ -53,7 +46,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)
@@ -77,115 +70,3 @@ class SamplingParam:
             sampling_param = cls()
 
         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:
-    cache_type: str = "none"

fastvideo/v1/configs/sample/hunyuan.py

@@ -1,7 +1,6 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
 from fastvideo.v1.configs.sample.base import SamplingParam
-from fastvideo.v1.configs.sample.teacache import TeaCacheParams
 
 
 @dataclass
@@ -15,14 +14,6 @@ class HunyuanSamplingParam(SamplingParam):
 
     guidance_scale: float = 1.0
 
-    teacache_params: TeaCacheParams = field(
-        default_factory=lambda: TeaCacheParams(
-            teacache_thresh=0.15,
-            coefficients=[
-                7.33226126e+02, -4.01131952e+02, 6.75869174e+01,
-                -3.14987800e+00, 9.61237896e-02
-            ]))
-
 
 @dataclass
 class FastHunyuanSamplingParam(HunyuanSamplingParam):

fastvideo/v1/configs/sample/registry.py

@@ -1,53 +1,45 @@
 import os
-from collections.abc import Callable
-from typing import Any
+from typing import Any, Callable, Dict, Optional
 
 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,
-                                             WanT2V_14B_SamplingParam)
+from fastvideo.v1.configs.sample.wan import (WanI2V480PSamplingParam,
+                                             WanT2V480PSamplingParam)
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import (maybe_download_model_index,
                                 verify_model_config_and_directory)
 
 logger = init_logger(__name__)
 # Registry maps specific model weights to their config classes
-SAMPLING_PARAM_REGISTRY: dict[str, Any] = {
+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,
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PSamplingParam,
+    "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V480PSamplingParam
     # Add other specific weight variants
 }
 
 # For determining pipeline type from model ID
-SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = {
+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
+    WanT2V480PSamplingParam,  # Base Wan config as fallback for any Wan variant
+    "wanimagetovideo": WanI2V480PSamplingParam,
     # Other fallbacks by architecture
 }
 
 
-def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
+def get_sampling_param_cls_for_name(
+        pipeline_name_or_path: str) -> Optional[Any]:
     """Get the appropriate sampling param for specific pretrained weights."""
 
     if os.path.exists(pipeline_name_or_path):

fastvideo/v1/configs/sample/stepvideo.py

@@ -1,19 +0,0 @@
-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 = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。"

fastvideo/v1/configs/sample/teacache.py

@@ -1,40 +0,0 @@
-from dataclasses import dataclass, field
-
-from fastvideo.v1.configs.sample.base import CacheParams
-
-
-@dataclass
-class TeaCacheParams(CacheParams):
-    cache_type: str = "teacache"
-    teacache_thresh: float = 0.0
-    coefficients: list[float] = field(default_factory=list)
-
-
-@dataclass
-class WanTeaCacheParams(CacheParams):
-    # Unfortunately, TeaCache is very different for Wan than other models
-    cache_type: str = "teacache"
-    teacache_thresh: float = 0.0
-    use_ret_steps: bool = True
-    ret_steps_coeffs: list[float] = field(default_factory=list)
-    non_ret_steps_coeffs: list[float] = field(default_factory=list)
-
-    @property
-    def coefficients(self) -> list[float]:
-        if self.use_ret_steps:
-            return self.ret_steps_coeffs
-        else:
-            return self.non_ret_steps_coeffs
-
-    @property
-    def ret_steps(self) -> int:
-        if self.use_ret_steps:
-            return 5 * 2
-        else:
-            return 1 * 2
-
-    def get_cutoff_steps(self, num_inference_steps: int) -> int:
-        if self.use_ret_steps:
-            return num_inference_steps * 2
-        else:
-            return num_inference_steps * 2 - 2

fastvideo/v1/configs/sample/wan.py

@@ -1,11 +1,10 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
 from fastvideo.v1.configs.sample.base import SamplingParam
-from fastvideo.v1.configs.sample.teacache import WanTeaCacheParams
 
 
 @dataclass
-class WanT2V_1_3B_SamplingParam(SamplingParam):
+class WanT2V480PSamplingParam(SamplingParam):
     # Video parameters
     height: int = 480
     width: int = 832
@@ -17,79 +16,9 @@ class WanT2V_1_3B_SamplingParam(SamplingParam):
     negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
     num_inference_steps: int = 50
 
-    teacache_params: WanTeaCacheParams = field(
-        default_factory=lambda: WanTeaCacheParams(
-            teacache_thresh=0.08,
-            ret_steps_coeffs=[
-                -5.21862437e+04, 9.23041404e+03, -5.28275948e+02,
-                1.36987616e+01, -4.99875664e-02
-            ],
-            non_ret_steps_coeffs=[
-                2.39676752e+03, -1.31110545e+03, 2.01331979e+02,
-                -8.29855975e+00, 1.37887774e-01
-            ]))
-
 
 @dataclass
-class WanT2V_14B_SamplingParam(SamplingParam):
-    # Video parameters
-    height: int = 720
-    width: int = 1280
-    num_frames: int = 81
-    fps: int = 16
-
-    # Denoising stage
-    guidance_scale: float = 5.0
-    negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
-    num_inference_steps: int = 50
-
-    teacache_params: WanTeaCacheParams = field(
-        default_factory=lambda: WanTeaCacheParams(
-            teacache_thresh=0.20,
-            use_ret_steps=False,
-            ret_steps_coeffs=[
-                -3.03318725e+05, 4.90537029e+04, -2.65530556e+03,
-                5.87365115e+01, -3.15583525e-01
-            ],
-            non_ret_steps_coeffs=[
-                -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429,
-                -13.02252404
-            ]))
-
-
-@dataclass
-class WanI2V_14B_480P_SamplingParam(WanT2V_1_3B_SamplingParam):
+class WanI2V480PSamplingParam(WanT2V480PSamplingParam):
     # Denoising stage
     guidance_scale: float = 5.0
     num_inference_steps: int = 40
-
-    teacache_params: WanTeaCacheParams = field(
-        default_factory=lambda: WanTeaCacheParams(
-            teacache_thresh=0.26,
-            ret_steps_coeffs=[
-                -3.03318725e+05, 4.90537029e+04, -2.65530556e+03,
-                5.87365115e+01, -3.15583525e-01
-            ],
-            non_ret_steps_coeffs=[
-                -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429,
-                -13.02252404
-            ]))
-
-
-@dataclass
-class WanI2V_14B_720P_SamplingParam(WanT2V_14B_SamplingParam):
-    # Denoising stage
-    guidance_scale: float = 5.0
-    num_inference_steps: int = 40
-
-    teacache_params: WanTeaCacheParams = field(
-        default_factory=lambda: WanTeaCacheParams(
-            teacache_thresh=0.3,
-            ret_steps_coeffs=[
-                -3.03318725e+05, 4.90537029e+04, -2.65530556e+03,
-                5.87365115e+01, -3.15583525e-01
-            ],
-            non_ret_steps_coeffs=[
-                -5784.54975374, 5449.50911966, -1811.16591783, 256.27178429,
-                -13.02252404
-            ]))

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

@@ -26,16 +26,12 @@
     "prefix": "Wan",
     "quant_config": null
   },
-  "text_encoder_precisions": [
-    "fp32"
-  ],
-  "text_encoder_configs": [
-    {
-      "prefix": "t5",
-      "quant_config": null,
-      "lora_config": null
-    }
-  ],
+  "text_encoder_precision": "fp32",
+  "text_encoder_config": {
+    "prefix": "t5",
+    "quant_config": null,
+    "lora_config": null
+  },
   "mask_strategy_file_path": null,
   "enable_torch_compile": false
 }

fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json

@@ -26,16 +26,12 @@
     "prefix": "Wan",
     "quant_config": null
   },
-  "text_encoder_precisions": [
-    "fp32"
-  ],
-  "text_encoder_configs": [
-    {
-      "prefix": "t5",
-      "quant_config": null,
-      "lora_config": null
-    }
-  ],
+  "text_encoder_precision": "fp32",
+  "text_encoder_config": {
+    "prefix": "t5",
+    "quant_config": null,
+    "lora_config": null
+  },
   "mask_strategy_file_path": null,
   "enable_torch_compile": false,
   "image_encoder_config": {

fastvideo/v1/default_configs/v1_inference_hunyuan_config.yaml

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

fastvideo/v1/distributed/device_communicators/base_device_communicator.py

@@ -1,6 +1,8 @@
 # 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
@@ -16,8 +18,8 @@ class DeviceCommunicatorBase:
 
     def __init__(self,
                  cpu_group: ProcessGroup,
-                 device: torch.device | None = None,
-                 device_group: ProcessGroup | None = None,
+                 device: Optional[torch.device] = None,
+                 device_group: Optional[ProcessGroup] = None,
                  unique_name: str = ""):
         self.device = device or torch.device("cpu")
         self.cpu_group = cpu_group
@@ -64,7 +66,7 @@ class DeviceCommunicatorBase:
     def gather(self,
                input_: torch.Tensor,
                dst: int = 0,
-               dim: int = -1) -> torch.Tensor | None:
+               dim: int = -1) -> Optional[torch.Tensor]:
         """
         NOTE: We assume that the input tensor is on the same device across
         all the ranks.
@@ -168,7 +170,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: int | None = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: Optional[int] = 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:
@@ -178,7 +180,7 @@ class DeviceCommunicatorBase:
     def recv(self,
              size: torch.Size,
              dtype: torch.dtype,
-             src: int | None = None) -> torch.Tensor:
+             src: Optional[int] = None) -> torch.Tensor:
         """Receives a tensor from the source rank."""
         """NOTE: `src` is the local rank of the source rank."""
         if src is None:

fastvideo/v1/distributed/device_communicators/cuda_communicator.py

@@ -1,6 +1,8 @@
 # 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
 
@@ -12,15 +14,15 @@ class CudaCommunicator(DeviceCommunicatorBase):
 
     def __init__(self,
                  cpu_group: ProcessGroup,
-                 device: torch.device | None = None,
-                 device_group: ProcessGroup | None = None,
+                 device: Optional[torch.device] = None,
+                 device_group: Optional[ProcessGroup] = 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: PyNcclCommunicator | None = None
+        self.pynccl_comm: Optional[PyNcclCommunicator] = None
         if self.world_size > 1:
             self.pynccl_comm = PyNcclCommunicator(
                 group=self.cpu_group,
@@ -40,7 +42,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
             torch.distributed.all_reduce(out, group=self.device_group)
         return out
 
-    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: Optional[int] = 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:
@@ -55,7 +57,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
     def recv(self,
              size: torch.Size,
              dtype: torch.dtype,
-             src: int | None = None) -> torch.Tensor:
+             src: Optional[int] = None) -> torch.Tensor:
         """Receives a tensor from the source rank."""
         """NOTE: `src` is the local rank of the source rank."""
         if src is None:

fastvideo/v1/distributed/device_communicators/pynccl.py

@@ -1,6 +1,8 @@
 # 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
@@ -20,9 +22,9 @@ class PyNcclCommunicator:
 
     def __init__(
         self,
-        group: ProcessGroup | StatelessProcessGroup,
-        device: int | str | torch.device,
-        library_path: str | None = None,
+        group: Union[ProcessGroup, StatelessProcessGroup],
+        device: Union[int, str, torch.device],
+        library_path: Optional[str] = None,
     ):
         """
         Args:

fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py

@@ -27,7 +27,7 @@
 import ctypes
 import platform
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Dict, List, Optional
 
 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: str | None = None):
+    def __init__(self, so_file: Optional[str] = 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

fastvideo/v1/distributed/parallel_state.py

@@ -27,11 +27,10 @@ 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, Optional
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from unittest.mock import patch
 
 import torch
@@ -58,15 +57,15 @@ TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
 
 
 def _split_tensor_dict(
-    tensor_dict: dict[str, torch.Tensor | Any]
-) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:
+    tensor_dict: Dict[str, Union[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,
@@ -82,7 +81,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:
@@ -98,7 +97,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:
@@ -129,7 +128,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:
@@ -144,16 +143,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: Any | None  # shared memory broadcaster
+    mq_broadcaster: Optional[Any]  # shared memory broadcaster
 
     def __init__(
         self,
-        group_ranks: list[list[int]],
+        group_ranks: List[List[int]],
         local_rank: int,
-        torch_distributed_backend: str | Backend,
+        torch_distributed_backend: Union[str, Backend],
         use_device_communicator: bool,
         use_message_queue_broadcaster: bool = False,
-        group_name: str | None = None,
+        group_name: Optional[str] = None,
     ):
         group_name = group_name or "anonymous"
         self.unique_name = _get_unique_name(group_name)
@@ -244,8 +243,8 @@ class GroupCoordinator:
         return self.ranks[(rank_in_group - 1) % world_size]
 
     @contextmanager
-    def graph_capture(self,
-                      graph_capture_context: GraphCaptureContext | None = None):
+    def graph_capture(
+            self, graph_capture_context: Optional[GraphCaptureContext] = None):
         if graph_capture_context is None:
             stream = torch.cuda.Stream()
             graph_capture_context = GraphCaptureContext(stream)
@@ -302,7 +301,7 @@ class GroupCoordinator:
     def gather(self,
                input_: torch.Tensor,
                dst: int = 0,
-               dim: int = -1) -> torch.Tensor | None:
+               dim: int = -1) -> Optional[torch.Tensor]:
         """
         NOTE: We assume that the input tensor is on the same device across
         all the ranks.
@@ -338,7 +337,7 @@ class GroupCoordinator:
                                     group=self.device_group)
         return input_
 
-    def broadcast_object(self, obj: Any | None = None, src: int = 0):
+    def broadcast_object(self, obj: Optional[Any] = None, src: int = 0):
         """Broadcast the input object.
         NOTE: `src` is the local rank of the source rank.
         """
@@ -363,9 +362,9 @@ class GroupCoordinator:
             return recv[0]
 
     def broadcast_object_list(self,
-                              obj_list: list[Any],
+                              obj_list: List[Any],
                               src: int = 0,
-                              group: ProcessGroup | None = None):
+                              group: Optional[ProcessGroup] = None):
         """Broadcast the input object list.
         NOTE: `src` is the local rank of the source rank.
         """
@@ -445,11 +444,11 @@ class GroupCoordinator:
 
     def broadcast_tensor_dict(
         self,
-        tensor_dict: dict[str, torch.Tensor | Any] | None = None,
+        tensor_dict: Optional[Dict[str, Union[torch.Tensor, Any]]] = None,
         src: int = 0,
-        group: ProcessGroup | None = None,
-        metadata_group: ProcessGroup | None = None
-    ) -> dict[str, torch.Tensor | Any] | None:
+        group: Optional[ProcessGroup] = None,
+        metadata_group: Optional[ProcessGroup] = None
+    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
         """Broadcast the input tensor dictionary.
         NOTE: `src` is the local rank of the source rank.
         """
@@ -463,7 +462,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)}")
@@ -530,10 +529,10 @@ class GroupCoordinator:
 
     def send_tensor_dict(
         self,
-        tensor_dict: dict[str, torch.Tensor | Any],
-        dst: int | None = None,
+        tensor_dict: Dict[str, Union[torch.Tensor, Any]],
+        dst: Optional[int] = None,
         all_gather_group: Optional["GroupCoordinator"] = None,
-    ) -> dict[str, torch.Tensor | Any] | None:
+    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
         """Send the input tensor dictionary.
         NOTE: `dst` is the local rank of the source rank.
         """
@@ -553,7 +552,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)}"
@@ -584,9 +583,9 @@ class GroupCoordinator:
 
     def recv_tensor_dict(
         self,
-        src: int | None = None,
+        src: Optional[int] = None,
         all_gather_group: Optional["GroupCoordinator"] = None,
-    ) -> dict[str, torch.Tensor | Any] | None:
+    ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
         """Recv the input tensor dictionary.
         NOTE: `src` is the local rank of the source rank.
         """
@@ -607,7 +606,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,
@@ -657,7 +656,7 @@ class GroupCoordinator:
         """
         torch.distributed.barrier(group=self.cpu_group)
 
-    def send(self, tensor: torch.Tensor, dst: int | None = None) -> None:
+    def send(self, tensor: torch.Tensor, dst: Optional[int] = 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)
@@ -665,7 +664,7 @@ class GroupCoordinator:
     def recv(self,
              size: torch.Size,
              dtype: torch.dtype,
-             src: int | None = None) -> torch.Tensor:
+             src: Optional[int] = 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)
@@ -683,7 +682,7 @@ class GroupCoordinator:
             self.mq_broadcaster = None
 
 
-_WORLD: GroupCoordinator | None = None
+_WORLD: Optional[GroupCoordinator] = None
 
 
 def get_world_group() -> GroupCoordinator:
@@ -691,7 +690,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],
@@ -703,11 +702,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: str | None = None,
+    group_name: Optional[str] = None,
 ) -> GroupCoordinator:
 
     return GroupCoordinator(
@@ -720,7 +719,7 @@ def init_model_parallel_group(
     )
 
 
-_TP: GroupCoordinator | None = None
+_TP: Optional[GroupCoordinator] = None
 
 
 def get_tp_group() -> GroupCoordinator:
@@ -779,7 +778,7 @@ def init_distributed_environment(
             "world group already initialized with a different world size")
 
 
-_SP: GroupCoordinator | None = None
+_SP: Optional[GroupCoordinator] = None
 
 
 def get_sp_group() -> GroupCoordinator:
@@ -790,7 +789,7 @@ def get_sp_group() -> GroupCoordinator:
 def initialize_model_parallel(
     tensor_model_parallel_size: int = 1,
     sequence_model_parallel_size: int = 1,
-    backend: str | None = None,
+    backend: Optional[str] = None,
 ) -> None:
     """
     Initialize model parallel groups.
@@ -859,7 +858,7 @@ def get_sequence_model_parallel_rank() -> int:
 def ensure_model_parallel_initialized(
     tensor_model_parallel_size: int,
     sequence_model_parallel_size: int,
-    backend: str | None = None,
+    backend: Optional[str] = None,
 ) -> None:
     """Helper to initialize model parallel groups if they are not initialized,
     or ensure tensor-parallel, sequence-parallel sizes 
@@ -970,8 +969,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: ProcessGroup | StatelessProcessGroup,
-                        source_rank: int = 0) -> list[bool]:
+def in_the_same_node_as(pg: Union[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
@@ -1057,7 +1056,7 @@ def in_the_same_node_as(pg: ProcessGroup | StatelessProcessGroup,
 
 def initialize_tensor_parallel_group(
         tensor_model_parallel_size: int = 1,
-        backend: str | None = None,
+        backend: Optional[str] = None,
         group_name_suffix: str = "") -> GroupCoordinator:
     """Initialize a tensor parallel group for a specific model.
     
@@ -1121,7 +1120,7 @@ def initialize_tensor_parallel_group(
 
 def initialize_sequence_parallel_group(
         sequence_model_parallel_size: int = 1,
-        backend: str | None = None,
+        backend: Optional[str] = None,
         group_name_suffix: str = "") -> GroupCoordinator:
     """Initialize a sequence parallel group for a specific model.
     

fastvideo/v1/distributed/utils.py

@@ -9,8 +9,7 @@ import dataclasses
 import pickle
 import time
 from collections import deque
-from collections.abc import Sequence
-from typing import Any
+from typing import Any, Deque, Dict, Optional, Sequence, Tuple
 
 import torch
 from torch.distributed import TCPStore
@@ -73,15 +72,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
@@ -95,14 +94,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.perf_counter()))
+        self.entries.append((key, time.time()))
 
     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.perf_counter() - timestamp > self.data_expiration_seconds:
+            if time.time() - timestamp > self.data_expiration_seconds:
                 self.store.delete_key(key)
                 self.entries.popleft()
             else:
@@ -115,7 +114,7 @@ class StatelessProcessGroup:
         self.recv_src_counter[src] += 1
         return obj
 
-    def broadcast_obj(self, obj: Any | None, src: int) -> Any:
+    def broadcast_obj(self, obj: Optional[Any], 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.
@@ -126,7 +125,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.perf_counter()))
+            self.entries.append((key, time.time()))
             return obj
         else:
             key = (f"broadcast_from/{src}/"

fastvideo/v1/entrypoints/cli/generate.py

@@ -2,20 +2,13 @@
 # adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py
 
 import argparse
-import dataclasses
-import os
-from typing import Any, cast
+from typing import List, cast
 
-from fastvideo import PipelineConfig, VideoGenerator
-from fastvideo.v1.configs.sample.base import SamplingParam
+from fastvideo.v1.entrypoints.cli import utils
 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"""
@@ -23,80 +16,45 @@ 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', 'dispatch_function']
+        excluded_args = [
+            'subparser', 'config', 'num_gpus', 'master_port',
+            'dispatch_function'
+        ]
 
-        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 = {
+        # Create a filtered dictionary of arguments
+        filtered_args = {
             k: v
-            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
+            for k, v in vars(args).items()
+            if k not in excluded_args and v is not None
         }
 
-        pipeline_config = PipelineConfig.from_pretrained(
-            merged_args['model_path'])
+        main_args = []
 
-        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)
+        for key, value in filtered_args.items():
+            # Convert underscores to dashes in argument names
+            arg_name = f"--{key.replace('_', '-')}"
 
-        model_path = init_args.pop('model_path')
-        prompt = generation_args.pop('prompt')
+            # 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))
 
-        generator = VideoGenerator.from_pretrained(
-            model_path=model_path, **init_args, pipeline_config=pipeline_config)
-
-        generator.generate_video(prompt=prompt, **generation_args)
+        utils.launch_distributed(args.num_gpus,
+                                 main_args,
+                                 master_port=args.master_port)
 
     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.config and not os.path.exists(args.config):
-            raise ValueError(f"Config file not found: {args.config}")
+        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")
 
     def subparser_init(
             self,
@@ -105,7 +63,7 @@ class GenerateSubcommand(CLISubcommand):
             "generate",
             help="Run inference on a model",
             usage=
-            "fastvideo generate (--model-path MODEL_PATH_OR_ID --prompt PROMPT) | --config CONFIG_FILE [OPTIONS]"
+            "fastvideo generate --model-path MODEL_PATH_OR_ID --prompt PROMPT [OPTIONS]"
         )
 
         generate_parser.add_argument(
@@ -113,53 +71,17 @@ class GenerateSubcommand(CLISubcommand):
             type=str,
             default='',
             required=False,
-            help=
-            "Read CLI options from a config JSON or YAML file. If provided, --model-path and --prompt are optional."
-        )
+            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")
 
         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)

fastvideo/v1/entrypoints/cli/main.py

@@ -1,12 +1,14 @@
 # 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())