mp training

Co-authored-by: Kevin Lin <42618777+kevin314@users.noreply.github.com>
Co-authored-by: William Lin <SolitaryThinker@users.noreply.github.com>
Co-authored-by: Will Lin <wlsaidhi@gmail.com>

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

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

.github/workflows/pr-test.yml

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

.github/workflows/runpod-test.yml

@@ -0,0 +1,91 @@
+name: RunPod Test
+
+on:
+  workflow_call:
+    inputs:
+      job_id:
+        required: true
+        type: string
+        description: "Unique identifier for this test job"
+      gpu_type:
+        required: true
+        type: string
+        description: "GPU type to use (e.g. NVIDIA A40, NVIDIA L40S)"
+      gpu_count:
+        required: true
+        type: number
+        description: "Number of GPUs to use"
+      volume_size:
+        required: false
+        type: number
+        default: 20
+        description: "Volume size in GB"
+      disk_size:
+        required: false
+        type: number
+        default: 20
+        description: "Disk size in GB"
+      image:
+        required: true
+        type: string
+        description: "Docker image to use"
+      test_command:
+        required: true
+        type: string
+        description: "Command to run tests"
+      timeout_minutes:
+        required: false
+        type: number
+        default: 30
+        description: "Timeout in minutes"
+    secrets:
+      RUNPOD_API_KEY:
+        required: true
+      RUNPOD_PRIVATE_KEY:
+        required: true
+
+jobs:
+  run-test:
+    runs-on: ubuntu-latest
+    environment: runpod-runners
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.10"
+
+      - name: Set up SSH key
+        run: |
+          mkdir -p ~/.ssh
+          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
+          chmod 600 ~/.ssh/id_rsa
+          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
+
+      - name: Install dependencies
+        run: pip install requests
+
+      - name: Run tests on RunPod
+        env:
+          JOB_ID: ${{ inputs.job_id }}
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+        timeout-minutes: ${{ inputs.timeout_minutes }}
+        run: >-
+          python .github/scripts/runpod_api.py 
+          --gpu-type "${{ inputs.gpu_type }}"
+          --gpu-count ${{ inputs.gpu_count }}
+          --volume-size ${{ inputs.volume_size }}
+          --disk-size ${{ inputs.disk_size }}
+          --image "${{ inputs.image }}"
+          --test-command "${{ inputs.test_command }}"
+
+      - name: Terminate RunPod Instances
+        if: ${{ always() }}
+        env:
+          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
+          GITHUB_RUN_ID: ${{ github.run_id }}
+          JOB_ID: ${{ inputs.job_id }}
+        run: python .github/scripts/runpod_cleanup.py 

.pre-commit-config.yaml

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

README.md

@@ -2,62 +2,109 @@
 <img src=assets/logo.jpg width="30%"/>
 </div>
 
-FastVideo is a lightweight framework for accelerating large video diffusion models.
+**FastVideo is a unified framework for accelerated video generation.**
+
+It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
+With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
 
 <p align="center">
-    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
+    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
 </p>
 
-https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1
+<div align="center">
+<img src=assets/perf.png width="90%"/>
+</div>
 
-FastVideo currently offers: (with more to come)
+## Key Features
 
-- [NEW!] V1 inference API available. Full announcement coming soon!
-- [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
-- First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
-- Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
+FastVideo has the following features:
+- State-of-the-art performance optimizations for inference
+  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [TeaCache](https://arxiv.org/pdf/2411.19108)
+  - [Sage Attention](https://arxiv.org/abs/2410.02367)
+- Cutting edge models
+  - Wan2.1 T2V, I2V
+  - HunyuanVideo
+  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
+  - StepVideo T2V
+- Distillation support
+  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
 - Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
 - Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
 
-Dev in progress and highly experimental.
-
-## Change Log
-- ```2025/02/20```: FastVideo now supports STA on [StepVideo](https://github.com/stepfun-ai/Step-Video-T2V) with 3.4X speedup!
-- ```2025/02/18```: Release the inference code and kernel for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- ```2025/01/13```: Support Lora finetuning for HunyuanVideo.
-- ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment.
-- ```2024/12/17```: `FastVideo` v0.0.1 is released.
-
 ## Getting Started
+We recommend using an environment manager such as `Conda` to create a clean environment:
+
+```bash
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
+pip install fastvideo
+```
+
+Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html) for more detailed installation instructions.
+
+## Inference
+### Generating Your First Video
+Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
+
+```python
+from fastvideo import VideoGenerator
+
+def main():
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
+
+    # Generate the video
+    video = generator.generate_video(
+        prompt,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Run the script with:
+
+```bash
+python example.py
+```
+
+For a more detailed guide, please see our [inference quick start](https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html).
+
+### Other docs:
 
-- [Install FastVideo](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
 - [Design Overview](https://hao-ai-lab.github.io/FastVideo/design/overview.html)
 - [Contribution Guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation.html)
 
-### Inference
-- [Quick Start](https://hao-ai-lab.github.io/FastVideo/inference/examples/basic.html)
-- V1 Inference API Guide (Coming soon!)
-
-### Distillation and Finetuning
+## Distillation and Finetuning
 - [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/training/distillation.html)
 - [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetuning.html)
 
-### Deprecated APIs
-- [V0 Inference (Deprecated)](https://hao-ai-lab.github.io/FastVideo/inference/v0_inference.html)
-
 ## 📑 Development Plan
 
 <!-- - More distillation methods -->
   <!-- - [ ] Add Distribution Matching Distillation -->
 - More models support
   <!-- - [ ] Add CogvideoX model -->
-  - [ ] Add StepVideo to V1
+  - [x] Add StepVideo to V1
 - Optimization features
-  - [ ] Teacache in V1
-  - [ ] SageAttention in V1
+  - [x] Teacache in V1
+  - [x] SageAttention in V1
 - Code updates
-  - [ ] V1 Configuration API
+  - [x] V1 Configuration API
   - [ ] Support Training in V1
   <!-- - [ ] fp8 support -->
   <!-- - [ ] faster load model and save model support -->

collect_env.py

@@ -0,0 +1,772 @@
+# 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",
+    "zmq",
+    "nvidia",
+    "pynvml",
+}
+
+DEFAULT_PIP_PATTERNS = {
+    "torch",
+    "numpy",
+    "mypy",
+    "flake8",
+    "triton",
+    "optree",
+    "onnx",
+    "nccl",
+    "transformers",
+    "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()

distill_v1.sh

@@ -0,0 +1,46 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+
+DATA_DIR=./data
+# IP=[MASTER NODE IP]
+
+# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
+    # --gradient_checkpointing\
+    # --pretrained_model_name_or_path hunyuanvideo-community/HunyuanVideo \
+    # --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+torchrun --nnodes 1 --nproc_per_node 4\
+    fastvideo/v1/pipelines/training_pipeline.py\
+    --inference_mode False\
+    --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
+    --cache_dir "/home/ray/.cache"\
+    --data_json_path "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/videos2caption.json"\
+    --validation_prompt_dir "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/validation"\
+    --train_batch_size=1\
+    --num_latent_t 1 \
+    --sp_size 4 \
+    --tp_size 4 \
+    --train_sp_batch_size 1\
+    --dataloader_num_workers 4\
+    --gradient_accumulation_steps=1\
+    --max_train_steps=320\
+    --learning_rate=1e-6\
+    --mixed_precision="bf16"\
+    --checkpointing_steps=64\
+    --validation_steps 64\
+    --validation_sampling_steps "2,4,8" \
+    --checkpoints_total_limit 3\
+    --allow_tf32\
+    --ema_start_step 0\
+    --cfg 0.0\
+    --log_validation\
+    --output_dir="$DATA_DIR/outputs/hy_phase1_shift17_bs_16_HD"\
+    --tracker_project_name Hunyuan_Distill \
+    --num_height 720 \
+    --num_width 1280 \
+    --num_frames  125 \
+    --shift 17 \
+    --validation_guidance_scale "1.0" \
+    --num_euler_timesteps 50 \
+    --multi_phased_distill_schedule "4000-1" \
+    --not_apply_cfg_solver \
+    --master_weight_type "bf16"

docker/Dockerfile.python3.10

@@ -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.0.post2 --no-build-isolation && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
     conda clean -afy
 
 COPY . .

docker/Dockerfile.python3.11

@@ -0,0 +1,48 @@
+FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+    rm Miniconda3-latest-Linux-x86_64.sh
+
+ENV PATH=/opt/conda/bin:$PATH
+
+RUN conda create --name fastvideo-dev python=3.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

@@ -0,0 +1,48 @@
+FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+WORKDIR /FastVideo
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    wget \
+    git \
+    ca-certificates \
+    openssh-server \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+    rm Miniconda3-latest-Linux-x86_64.sh
+
+ENV PATH=/opt/conda/bin:$PATH
+
+RUN conda create --name fastvideo-dev python=3.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,25 +1,15 @@
-sphinx==6.2.1
-sphinx-argparse==0.4.0
-sphinx-book-theme==1.0.1
+sphinx==7.4.7
+sphinx-argparse==0.5.2
+sphinx-autodoc2==0.5.0
+sphinx-book-theme==1.1.4
 sphinx-copybutton==0.5.2
 sphinx-design==0.6.1
 sphinx-togglebutton==0.3.2
 myst-parser==3.0.1
 msgspec
-cloudpickle
+commonmark # Required by sphinx-argparse when using :markdownhelp:
 
 # packages to install to build the documentation
 cachetools
-pydantic >= 2.8
 -f https://download.pytorch.org/whl/cpu
-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
+torch

docs/source/api/summary.md

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

@@ -0,0 +1,22 @@
+# 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,17 +13,19 @@
 # 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__)
-sys.path.append(os.path.abspath("../.."))
+REPO_ROOT = Path(__file__).resolve().parent.parent.parent
+print(os.path.abspath(REPO_ROOT))
+sys.path.append(os.path.abspath(REPO_ROOT))
 
 # -- Project information -----------------------------------------------------
 
@@ -41,8 +43,7 @@ extensions = [
     "sphinx.ext.linkcode",
     "sphinx.ext.intersphinx",
     "sphinx_copybutton",
-    "sphinx.ext.autodoc",
-    "sphinx.ext.autosummary",
+    "autodoc2",
     "myst_parser",
     "sphinxarg.ext",
     "sphinx_design",
@@ -50,6 +51,31 @@ 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.
@@ -78,6 +104,11 @@ 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,
@@ -160,38 +191,38 @@ def linkcode_resolve(domain, info):
         return None
     if not info['module']:
         return None
-    module = info['module']
 
-    # try to determine the correct file and line number to link to
-    obj = sys.modules[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
 
-    # get as specific as we can
-    lineno: int = 0
-    filename: str = ""
-    try:
-        for part in info['fullname'].split('.'):
-            obj = getattr(obj, part)
+    # Get the line number of the object
+    with open(path) as f:
+        lines = f.readlines()
+    name = info['fullname'].split(".")[-1]
+    pattern = fr"^( {{4}})*((def|class) )?{name}\b.*"
+    for lineno, line in enumerate(lines, 1):
+        if not line or line.startswith("#"):
+            continue
+        if re.match(pattern, line):
+            break
 
-            if not (inspect.isclass(obj) or inspect.isfunction(obj)
-                    or inspect.ismethod(obj)):
-                obj = obj.__class__  # type: ignore[assignment]
+    # If the line number is not found, return None
+    if lineno == len(lines):
+        return None
 
-            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/"):
+    # If the line number is found, create the URL
+    filename = path.relative_to(REPO_ROOT)
+    if "checkouts" in path.parts:
         # a PR build on readthedocs
-        pr_number = filename.split("/")[1]
-        filename = filename.split("/", 2)[2]
+        pr_number = REPO_ROOT.name
         base, branch = get_repo_base_and_branch(pr_number)
         if base and branch:
             return f"https://github.com/{base}/blob/{branch}/{filename}#L{lineno}"
-
     # Otherwise, link to the source file on the main branch
     return f"https://github.com/hao-ai-lab/FastVideo/blob/main/{filename}#L{lineno}"
 
@@ -203,6 +234,8 @@ autodoc_mock_imports = [
     "cpuinfo",
     "cv2",
     "torch",
+    "huggingface_hub",
+    "torchvision",
     "transformers",
     "psutil",
     "prometheus_client",
@@ -231,18 +264,6 @@ 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":
@@ -254,7 +275,4 @@ intersphinx_mapping = {
     "psutil": ("https://psutil.readthedocs.io/en/stable", None),
 }
 
-autodoc_preserve_defaults = True
-autodoc_warningiserror = True
-
 navigation_with_keys = False

docs/source/contributing/developer_env/docker.md

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

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.0.post2 --no-build-isolation 
+pip install flash-attn==2.7.4.post1 --no-build-isolation 
 
 # Linting, formatting and static type checking
 pre-commit install --hook-type pre-commit --hook-type commit-msg

docs/source/generate_examples.py

@@ -9,7 +9,7 @@ from typing import Optional
 
 ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
 ROOT_DIR_RELATIVE = '../../../..'
-EXAMPLE_DIR = ROOT_DIR / "fastvideo/v1/examples"
+EXAMPLE_DIR = ROOT_DIR / "examples"
 EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples"
 
 

docs/source/getting_started/installation.md

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

docs/source/getting_started/v1_api.md

@@ -0,0 +1,83 @@
+# 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 lightweight framework for accelerating large video diffusion models.
+<strong>FastVideo is a unified framework for accelerated video generation.
 </strong>
 </p>
 
@@ -21,27 +21,31 @@
 </p>
 :::
 
-FastVideo is a lightweight framework for accelerating large video diffusion models developed by the [Hao AI Lab](https://hao-ai-lab.github.io/).
+It features a clean, consistent API that works across popular video models, making it easier for developers to author new models and incorporate system- or kernel-level optimizations.
+With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
 
 <div style="text-align: center;">
-  <video controls width="800">
-    <source src="https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1" type="video/mp4">
-    Your browser does not support the video tag.
-  </video>
+  <img src=_static/images/perf.png width="100%"/>
 </div>
 
-FastVideo currently offers: (with more to come)
+## Key Features
 
-- [NEW!] V1 inference API available. Full announcement coming soon!
-- [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
-- First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
-- Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
+FastVideo has the following features:
+- State-of-the-art performance optimizations for inference
+  - [Sliding Tile Attention](https://arxiv.org/pdf/2502.04507)
+  - [TeaCache](https://arxiv.org/pdf/2411.19108)
+  - [Sage Attention](https://arxiv.org/abs/2410.02367)
+- Cutting edge models
+  - Wan2.1 T2V, I2V
+  - HunyuanVideo
+  - FastHunyuan: consistency distilled video diffusion models for 8x inference speedup.
+  - StepVideo T2V
+- Distillation support
+  - Recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
+  - Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
 - Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
 - Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
 
-Dev in progress and highly experimental.
-
 ## Documentation
 
 % How to start using FastVideo?
@@ -51,14 +55,20 @@ Dev in progress and highly experimental.
 :maxdepth: 1
 
 getting_started/installation
-<!-- getting_started/examples/examples_index -->
+<!-- getting_started/v1_api -->
 :::
 
 :::{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
 :::
 
@@ -93,7 +103,14 @@ design/overview
 
 contributing/overview
 contributing/developer_env/index
-contributing/add_pipeline
+:::
+
+:::{toctree}
+:caption: API Reference
+:maxdepth: 2
+
+<!-- api/summary -->
+api/fastvideo/fastvideo
 :::
 
 ## Indices and tables

docs/source/inference/add_pipeline.md

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

docs/source/inference/cli.md

@@ -0,0 +1,151 @@
+# FastVideo CLI Inference
+
+The FastVideo CLI provides a quick way to access the FastVideo inference pipeline for video generation. For more advanced usage,
+see the Python interface [here](https://hao-ai-lab.github.io/FastVideo/inference/examples/basic.html).
+
+## Basic Usage
+
+The basic command to generate a video is:
+
+```bash
+fastvideo generate --model-path {MODEL_PATH} --prompt {PROMPT}
+```
+
+### Required Parameters
+
+- `--model-path {MODEL_PATH}`: Path to the model or model ID
+- `--prompt {PROMPT}`: Text description for the video you want to generate
+
+## Common Arguments
+
+To see all the options, you can use the `--help` flag:
+
+```bash
+fastvideo generate --help
+```
+
+### Hardware Configuration
+
+- `--num-gpus {NUM_GPUS}`: Number of GPUs to use
+- `--tp-size {TP_SIZE}`: Tensor parallelism size (Typically should match the number of GPUs)
+- `--sp-size {SP_SIZE}`: Sequence parallelism size (Typically should match the number of GPUs)
+
+#### Video Configuration
+
+- `--height {HEIGHT}`: Height of the generated video
+- `--width {WIDTH}`: Width of the generated video
+- `--num-frames {NUM_FRAMES}`: Number of frames to generate
+- `--fps {FPS}`: Frames per second for the saved video
+
+#### Generation Parameters
+
+- `--num-inference-steps {STEPS}`: Number of denoising steps
+- `--negative-prompt {PROMPT}`: Negative prompt to guide generation away from certain concepts
+- `--seed {SEED}`: Random seed for reproducible generation
+
+#### Output Options
+
+- `--output-path {PATH}`: Directory to save the generated video
+- `--save-video`: Whether to save the video to disk
+- `--return-frames`: Whether to return the raw frames
+
+## Using Configuration Files
+
+Instead of specifying all parameters on the command line, you can use a configuration file:
+
+```bash
+fastvideo generate --config {CONFIG_FILE_PATH}
+```
+
+The config file should be in JSON or YAML format with the same parameter names as the CLI options. Command-line arguments will take precedence over settings in the configuration file, allowing you to override specific values while keeping the rest from the config file.
+
+Example configuration file (config.json):
+
+```json
+{
+    "model_path": "FastVideo/FastHunyuan-diffusers",
+    "prompt": "A beautiful woman in a red dress walking down a street",
+    "output_path": "outputs/",
+    "num_gpus": 2,
+    "sp_size": 2,
+    "tp_size": 2,
+    "num_frames": 45,
+    "height": 720,
+    "width": 1280,
+    "num_inference_steps": 6,
+    "seed": 1024,
+    "fps": 24,
+    "precision": "bf16",
+    "vae_precision": "fp16",
+    "vae_tiling": true,
+    "vae_sp": true,
+    "vae_config": {
+        "load_encoder": false,
+        "load_decoder": true,
+        "tile_sample_min_height": 256,
+        "tile_sample_min_width": 256
+    },
+    "text_encoder_precisions": [
+        "fp16",
+        "fp16"
+    ],
+    "mask_strategy_file_path": null,
+    "enable_torch_compile": false
+}
+```
+
+Or using YAML format (config.yaml):
+
+```yaml
+model_path: "FastVideo/FastHunyuan-diffusers"
+prompt: "A beautiful woman in a red dress walking down a street"
+output_path: "outputs/"
+num_gpus: 2
+sp_size: 2
+tp_size: 2
+num_frames: 45
+height: 720
+width: 1280
+num_inference_steps: 6
+seed: 1024
+fps: 24
+precision: "bf16"
+vae_precision: "fp16"
+vae_tiling: true
+vae_sp: true
+vae_config:
+  load_encoder: false
+  load_decoder: true
+  tile_sample_min_height: 256
+  tile_sample_min_width: 256
+text_encoder_precisions:
+  - "fp16"
+  - "fp16"
+mask_strategy_file_path: null
+enable_torch_compile: false
+```
+
+## Examples
+
+Generating a simple video:
+
+```bash
+fastvideo generate --model-path FastVideo/FastHunyuan-diffusers --prompt "A cat playing with a ball of yarn" --num-frames 45 --height 720 --width 1280 --num-inference-steps 6 --seed 1024 --output-path outputs/
+```
+
+Using a negative prompt to avoid certain elements:
+
+```bash
+fastvideo generate --model-path FastVideo/FastHunyuan-diffusers --prompt "A beautiful forest landscape" --negative-prompt "people, buildings, roads"
+```
+
+Combining command line arguments and a configuration file:
+
+```bash
+fastvideo generate --config config.json --prompt "A capybara lounging in a hammock"
+```
+
+## Troubleshooting
+
+- If you encounter CUDA out-of-memory errors, try reducing the video dimensions or number of frames, or the number of inference steps.
+- For reproducible results, set the same seed value between runs.

docs/source/inference/configuration.md

@@ -0,0 +1,77 @@
+(inference-configuration)=
+# Configuration
+
+## Multi-GPU Setup
+
+FastVideo automatically distributes the generation process when multiple GPUs are specified:
+
+```python
+# Will use 4 GPUs in parallel for faster generation
+generator = VideoGenerator.from_pretrained(
+    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    num_gpus=4,
+)
+```
+
+## Customizing Generation
+
+- `PipelineConfig`: Initialization time parameters
+- `SamplingParam`: Generation time parameters
+
+You can customize various parameters when generating videos using the `PipelineConfig` and `SamplingParam` class:
+
+```python
+from fastvideo import VideoGenerator, SamplingParam, PipelineConfig
+
+def main():
+    model_name = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    config = PipelineConfig.from_pretrained(model_name)
+    config.vae_precision = "fp16"
+    config.use_cpu_offload = True
+
+    # Create the generator
+    generator = VideoGenerator.from_pretrained(
+        model_name,
+        num_gpus=1,
+        pipeline_config=config
+    )
+
+    # Create and customize sampling parameters
+    sampling_param = SamplingParam.from_pretrained("Wan-AI/Wan2.1-T2V-1.3B-Diffusers")
+    
+    # How many frames to generate
+    sampling_param.num_frames = 45
+    
+    # Video resolution (width, height)
+    sampling_param.width = 1024
+    sampling_param.height = 576
+    
+    # How many steps we denoise the video (higher = better quality, slower generation)
+    sampling_param.num_inference_steps = 30
+    
+    # How strongly the video conforms to the prompt (higher = more faithful to prompt)
+    sampling_param.guidance_scale = 7.5
+    
+    # Random seed for reproducibility
+    sampling_param.seed = 42  # Optional, leave unset for random results
+
+    # Generate video with custom parameters
+    prompt = "A beautiful sunset over a calm ocean, with gentle waves."
+    video = generator.generate_video(
+        prompt, 
+        sampling_param=sampling_param, 
+        output_path="my_videos/",  # Controls where videos are saved
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        save_video=True
+    )
+
+    # If return_frames=True, video contains the generated frames as a NumPy array
+    print(f"Generated {len(video)} frames")
+
+if __name__ == '__main__':
+    main()
+```
+
+## Performance Optimization
+
+For configuring optimizations, please see our [optimizations guide](#inference-optimizations)

docs/source/inference/fasthunyuan.md

@@ -1,33 +0,0 @@
-(fasthunyuan)=
-
-# FastHunyuan
-## Inference FastHunyuan on single RTX4090
-We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model
-# CLI inference
-bash scripts/inference/inference_hunyuan_hf_quantization.sh
-```
-
-For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU):
-
-| Configuration                  | Memory to Init Transformer | Peak Memory After Init Pipeline (Denoise) | Diffusion Time | End-to-End Time |
-|--------------------------------|----------------------------|--------------------------------------------|----------------|-----------------|
-| BF16 + Pipeline CPU Offload    | 23.883G                   | 33.744G                                    | 81s            | 121.5s          |
-| INT8 + Pipeline CPU Offload    | 13.911G                   | 27.979G                                    | 88s            | 116.7s          |
-| NF4 + Pipeline CPU Offload     | 9.453G                    | 19.26G                                     | 78s            | 114.5s          |
-  
-For improved quality in generated videos, we recommend using a GPU with 80GB of memory to run the BF16 model with the original Hunyuan pipeline. To execute the inference, use the following section:
-
-## FastHunyuan
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan --local_dir=data/FastHunyuan --repo_type=model
-# CLI inference
-bash scripts/inference/inference_hunyuan.sh
-```
-
-You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo).

docs/source/inference/fastmochi.md

@@ -1,9 +0,0 @@
-(fastmochi)=
-
-# FastMochi
-
-```bash
-# Download the model weight
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffusers --local_dir=data/FastMochi-diffusers --repo_type=model
-# CLI inference
-bash scripts/inference/inference_mochi_sp.sh

docs/source/inference/hunyuanvideo.md

@@ -1,18 +0,0 @@
-(hunyuanvideo)=
-
-# HunyuanVideo
-## Inference HunyuanVideo with Sliding Tile Attention
-First, download the model:
-
-```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model 
-```
-
-We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
-
-```bash
-sh scripts/inference/inference_hunyuan_STA.sh
-```
-
-## Video Demos using STA + Teacache
-Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100.

docs/source/inference/inference_quick_start.md

@@ -0,0 +1,124 @@
+# Inference Quick Start
+
+This page contains step-by-step instructions to get you quickly started with video generation using FastVideo.
+
+## Requirements
+- **OS**: Linux (Tested on Ubuntu 22.04+)
+- **Python**: 3.10-3.12
+- **CUDA**: 12.4
+- **GPU**: At least one NVIDIA GPU
+
+## Installation
+
+We recommend using an environment manager such as `Conda` to create a clean environment:
+
+```bash
+# Create and activate a new conda environment
+conda create -n fastvideo python=3.12
+conda activate fastvideo
+
+# Install FastVideo
+pip install fastvideo
+```
+
+For advanced installation options, see the [Installation Guide](installation.md).
+
+## Generating Your First Video
+Here's a minimal example to generate a video using the default settings. Create a file called `example.py` with the following code:
+
+```python
+from fastvideo import VideoGenerator
+
+def main():
+    # Create a video generator with a pre-trained model
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        num_gpus=1,  # Adjust based on your hardware
+    )
+
+    # Define a prompt for your video
+    prompt = "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest."
+
+    # Generate the video
+    video = generator.generate_video(
+        prompt,
+        return_frames=True,  # Also return frames from this call (defaults to False)
+        output_path="my_videos/",  # Controls where videos are saved
+        save_video=True
+    )
+
+if __name__ == '__main__':
+    main()
+```
+
+Run the script with:
+
+```bash
+python example.py
+```
+
+The generated video will be saved in the current directory under `my_videos/`.
+
+## Available Models
+
+Please see the [support matrix](#support-matrix) for the list of supported models and their available optimizations.
+
+## Image-to-Video Generation
+
+You can generate a video starting from an initial image:
+
+```python
+from fastvideo import VideoGenerator, SamplingParam
+
+def main():
+    # Create the generator
+    model_name = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
+    generator = VideoGenerator.from_pretrained(model_name, num_gpus=1)
+
+    # Set up parameters with an initial image
+    sampling_param = SamplingParam.from_pretrained(model_name)
+    sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
+    sampling_param.num_frames = 107
+    sampling_param.image_strength = 0.8  # How much to preserve the original image (0-1)
+
+    # Generate video based on the image
+    prompt = "A photograph coming to life with gentle movement"
+    generator.generate_video(prompt, sampling_param=sampling_param,
+                             output_path="my_videos/",
+                             save_video=True)
+
+if __name__ == '__main__':
+    main()
+```
+
+## Troubleshooting
+
+Common issues and their solutions:
+
+### Out of Memory Errors
+If you encounter CUDA out of memory errors:
+- Reduce `num_frames` or video resolution
+- Enable memory optimization with `enable_model_cpu_offload`
+- Try a smaller model or use distilled versions
+- Use `num_gpus` > 1 if multiple GPUs are available
+
+### Slow Generation
+To speed up generation:
+- Reduce `num_inference_steps` (20-30 is usually sufficient)
+- Use half precision (`fp16`) for the VAE
+- Use multiple GPUs if available
+
+### Unexpected Results
+If the generated video doesn't match your prompt:
+- Try increasing `guidance_scale` (7.0-9.0 works well)
+- Make your prompt more detailed and specific
+- Experiment with different random seeds
+- Try a different model
+
+## Next Steps
+
+- Learn about [Advanced Inference Configurations](#inference-configuration)
+- Learn about using [Optimizations](#inference-optimizations)
+- See [Examples](../examples/examples_inference_index.md) for more usage scenarios
+- Join our [Community Discord](https://discord.gg/JA7cksDz86).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).

docs/source/inference/optimizations.md

@@ -0,0 +1,148 @@
+(inference-optimizations)=
+# Optimizations
+
+This page describes the various options for speeding up generation times in FastVideo.
+
+## Table of Contents
+- Optimized Attention Backends
+  - [Flash Attention](#optimizations-flash)
+  - [Sliding Tile Attention](#optimizations-sta)
+  - [Sage Attention](#optimizations-sage)
+
+- Caching Techniques
+  - [TeaCache](#optimizations-teacache)
+
+(optimizations-backends)=
+## Attention Backends
+
+### Available Backends
+- Torch SDPA: `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA`
+- Flash Attention 2 and 3: `FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN`
+- Sliding Tile Attention: `FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN`
+- Sage Attention: `FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN`
+
+### Configuring Backends
+
+There are two ways to configure the attention backend in FastVideo.
+
+#### 1. In Python
+In python, set the `FASTVIDEO_ATTENTION_BACKEND` environment variable before instantiating `VideoGenerator` like this:
+
+```python
+os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLIDING_TILE_ATTN"
+```
+
+#### 2. In CLI
+You can also set the environment variable on the command line:
+
+```bash
+FASTVIDEO_ATTENTION_BACKEND=SAGE_ATTN python example.py
+```
+
+(optimizations-flash)=
+### Flash Attention
+
+**`FLASH_ATTN`**
+
+We recommend always installing [Flash Attention 2](https://github.com/Dao-AILab/flash-attention):
+
+```bash
+pip install flash-attn==2.7.4.post1 --no-build-isolation
+```
+
+And if using a Hopper+ GPU (ie H100), installing [Flash Attention 3](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#flashattention-3-beta-release) by compiling it from source (takes about 10 minutes for me):
+
+```bash
+git clone https://github.com/Dao-AILab/flash-attention.git && cd flash-attention
+
+cd hopper
+pip install ninja 
+python setup.py install
+```
+
+:::{note}
+FastVideo will automatically detect and use `FA3` if it is installed when using `FLASH_ATTN` backend.
+:::
+
+(optimizations-sta)=
+### Sliding Tile Attention
+**`SLIDING_TILE_ATTN`**
+
+```bash
+pip install st_attn==0.0.4
+```
+
+Please see [this page](#sta-installation) for more installation instructions.
+
+(optimizations-sage)=
+### Sage Attention
+**`SAGE_ATTN`**
+
+To use [SageAttention](https://github.com/thu-ml/SageAttention) 2.1.1, please compile from source:
+
+```bash
+git clone https://github.com/thu-ml/SageAttention.git
+cd sageattention 
+python setup.py install  # or pip install -e .
+```
+
+(optimizations-teacache)=
+## Teacache
+TeaCache is an optimization technique supported in FastVideo that can significantly speed up video generation by skipping redundant calculations across diffusion steps. This guide explains how to enable and configure TeaCache for optimal performance in FastVideo.
+
+### What is TeaCache?
+
+See the official [TeaCache](https://github.com/ali-vilab/TeaCache) repo and their paper for more details.
+
+### How to Enable TeaCache
+
+Enabling TeaCache is straightforward - simply add the `enable_teacache=True` parameter to your `generate_video()` call:
+
+```python
+# ... previous code
+generator.generate_video(
+    prompt="Your prompt here",
+    sampling_param=params,
+    enable_teacache=True
+)
+# more code ...
+```
+
+### Complete Example
+
+At the bottom is a complete example of using TeaCache for faster video generation. You can run it using the following command:
+
+```bash
+python examples/inference/optimizations/teacache_example.py
+```
+
+### Advanced Configuration
+
+While TeaCache works well with default settings, you can fine-tune its behavior by adjusting the threshold value:
+
+1. Lower threshold values (e.g., 0.1) will result in more skipped calculations and faster generation with slightly more potential for quality degradation
+2. Higher threshold values (e.g., 0.15-0.23) will skip fewer calculations but maintain quality closer to the original
+
+Note that the optimal threshold depends on your specific model and content.
+
+## Benchmarking different optimizations
+
+To benchmark the performance improvement, try generating the same video with and without TeaCache enabled and compare the generation times:
+
+```python
+# Without TeaCache
+start_time = time.perf_counter()
+generator.generate_video(prompt="Your prompt", enable_teacache=False)
+standard_time = time.perf_counter() - start_time
+
+# With TeaCache
+start_time = time.perf_counter()
+generator.generate_video(prompt="Your prompt", enable_teacache=True)
+teacache_time = time.perf_counter() - start_time
+
+print(f"Standard generation: {standard_time:.2f} seconds")
+print(f"TeaCache generation: {teacache_time:.2f} seconds")
+print(f"Speedup: {standard_time/teacache_time:.2f}x")
+```
+
+Note: If you want to benchmark different attention backends, you'll need to reinstantiate `VideoGenerator`.

docs/source/inference/stepvideo.md

@@ -1,16 +0,0 @@
-(stepvideo)=
-
-# StepVideo
-## Inference StepVideo with Sliding Tile Attention
-First, download the model:
-
-```
-python scripts/huggingface/download_hf.py --repo_id=stepfun-ai/stepvideo-t2v --local_dir=data/stepvideo-t2v --repo_type=model 
-```
-
-Use the following scripts to run inference for StepVideo. When using STA for inference, the generated videos will have dimensions of 204×768×768 (currently, this is the only supported shape).
-
-```bash
-sh scripts/inference/inference_stepvideo_STA.sh # Inference stepvideo with STA
-sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo
-```

docs/source/inference/support_matrix.md

@@ -0,0 +1,92 @@
+(support-matrix)=
+# Compatibility Matrix
+The table below shows every supported model and optimizations supported for them.
+
+The symbols used have the following meanings:
+
+- ✅ = Full compatibility
+- ❌ = No compatibility
+
+## Models x Optimization
+The `HuggingFace Model ID` can be directly pass to `from_pretrained()` methods and FastVideo will use the optimal default parameters when initializing and generating videos.
+
+:::{raw} html
+<style>
+  /* Make smaller to try to improve readability  */
+  td {
+    font-size: 0.9rem;
+    text-align: center;
+  }
+
+  th {
+    text-align: center;
+    font-size: 0.9rem;
+  }
+</style>
+:::
+
+:::{list-table}
+:header-rows: 1
+:stub-columns: 3
+:widths: auto
+:class: vertical-table-header
+
+- * Model Name
+  * HuggingFace Model ID
+  * Resolutions
+  * TeaCache
+  * Sliding Tile Attn
+  * Sage Attn
+- * HunyuanVideo
+  * `hunyuanvideo-community/HunyuanVideo`
+  * 720px1280p<br>544px960p
+  * ❌
+  * ✅
+  * ✅
+- * FastHunyuan
+  * `FastVideo/FastHunyuan-diffusers`
+  * 720px1280p<br>544px960p
+  * ❌
+  * ✅
+  * ✅
+- * Wan T2V 1.3B
+  * `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`
+  * 480P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan T2V 14B
+  * `Wan-AI/Wan2.1-T2V-14B-Diffusers`
+  * 480P, 720P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan I2V 480P
+  * `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers`
+  * 480P
+  * ✅
+  * ✅*
+  * ✅
+- * Wan I2V 720P
+  * `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers`
+  * 720P
+  * ✅
+  * ✅*
+  * ✅
+- * StepVideo T2V
+  * `FastVideo/stepvideo-t2v-diffusers`
+  * 768px768px204f<br>544px992px204f<br>544px992px136f
+  * ❌
+  * ❌
+  * ✅
+:::
+
+**Note**: there are some known quality issues with Wan2.1 + Sliding Tile Attn. We are working on fixing this issue.
+
+## Special requirements
+
+### StepVideo T2V
+- The self-attention in text-encoder (step_llm) only supports CUDA capabilities sm_80 sm_86 and sm_90
+
+### Sliding Tile Attention
+- Currently only Hopper GPUs (H100s) are supported.

docs/source/inference/wanvideo.md

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

examples/inference/basic/README.md

@@ -1,3 +1,41 @@
-# Basic
+# 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.
 
-The class provides the main python interface for using FastVideo's inference pipeline.
+## 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()
+```

examples/inference/basic/basic.py

@@ -1 +1,41 @@
-print('Hello, world!')
+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()

examples/inference/basic/default_args.py

@@ -0,0 +1,62 @@
+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/README.md

@@ -1,4 +1,4 @@
-# FastVideo VideoGenerator Gradio Demo
+# FastVideo Gradio Demo
 
 This is a Gradio-based web interface for generating videos using the FastVideo framework. The demo allows users to create videos from text prompts with various customization options.
 
@@ -13,25 +13,12 @@ The demo uses the FastVideo framework to generate videos based on text prompts.
 
 ---
 
-## Requirements
-
-- Linux-based OS
-- Python 3.10
-- Cuda 12.4
-- FastVideo
-
-## Installation
-
-```bash
-pip install fastvideo
-```
-
 ## Usage
 
 Run the demo with:
 
 ```bash
-python fastvideo/v1/examples/inference/gradio/gradio_demo.py
+python examples/inference/gradio/gradio_demo.py
 ```
 
 This will start a web server at `http://0.0.0.0:7860` where you can access the interface.
@@ -40,15 +27,6 @@ This will start a web server at `http://0.0.0.0:7860` where you can access the i
 
 ## Model Initialization
 
-```python
-args = FastVideoArgs(model_path="FastVideo/FastHunyuan-Diffusers", num_gpus=2)
-
-generator = VideoGenerator.from_pretrained(
-    model_path=args.model_path,
-    num_gpus=args.num_gpus
-)
-```
-
 This demo initializes a `VideoGenerator` with the minimum required arguments for inference. Users can seamlessly adjust inference options between generations, including prompts, resolution, video length, or even the number of inference steps, *without ever needing to reload the model*.
 
 ## Video Generation
@@ -78,4 +56,4 @@ The interface is built with several components:
 - **Guidance Scale**: Control how closely the generation follows the prompt
 - **Inference Steps**: More steps can improve quality but take longer
 - **Negative Prompt**: Specify what you don't want to see in the video
-- **Seed**: Control randomness for reproducible results
+- **Seed**: Control randomness for reproducible results

examples/inference/gradio/gradio_demo.py

@@ -0,0 +1,169 @@
+import argparse
+import os
+from copy import deepcopy
+
+import gradio as gr
+import torch
+
+from fastvideo import VideoGenerator
+from fastvideo.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

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

examples/inference/optimizations/attention_example.py

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

@@ -0,0 +1,44 @@
+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,3 +1,5 @@
+from fastvideo.v1.configs.pipelines import PipelineConfig
+from fastvideo.v1.configs.sample import SamplingParam
 from fastvideo.v1.entrypoints.video_generator import VideoGenerator
 
-__all__ = ["VideoGenerator"]
+__all__ = ["VideoGenerator", "PipelineConfig", "SamplingParam"]

fastvideo/data_preprocess/preprocess.py

@@ -0,0 +1,122 @@
+import argparse
+import json
+import os
+
+import torch
+import torch.distributed as dist
+
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.utils import maybe_download_model, shallow_asdict
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.configs.models.vaes import WanVAEConfig
+from fastvideo import PipelineConfig
+from fastvideo.v1.pipelines.preprocess_pipeline import PreprocessPipeline
+
+logger = init_logger(__name__)
+
+BASE_MODEL_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+
+def main(args):
+    # Assume using torchrun
+    local_rank = int(os.getenv("RANK", 0))
+    rank = int(os.environ.get("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
+    torch.cuda.set_device(local_rank)
+    if not dist.is_initialized():
+        dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
+
+    pipeline_config = PipelineConfig.from_pretrained(MODEL_PATH)
+    kwargs = {
+        "use_cpu_offload": False,
+        "vae_precision": "fp32",
+        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
+    }
+    pipeline_config_args = shallow_asdict(pipeline_config)
+    pipeline_config_args.update(kwargs)
+    fastvideo_args = FastVideoArgs(model_path=MODEL_PATH,
+                                   num_gpus=world_size,
+                                   device_str="cuda",
+                                   **pipeline_config_args,
+                                   )
+    fastvideo_args.check_fastvideo_args()
+    fastvideo_args.device = torch.device(f"cuda:{local_rank}")
+
+    pipeline = PreprocessPipeline(MODEL_PATH, fastvideo_args)
+    pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # dataset & dataloader
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--model_type", type=str, default="mochi")
+    parser.add_argument("--data_merge_path", type=str, required=True)
+    parser.add_argument("--validation_prompt_txt", type=str)
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=1,
+        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--preprocess_video_batch_size",
+        type=int,
+        default=2,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--preprocess_text_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--samples_per_file",
+        type=int,
+        default=64
+    )
+    parser.add_argument(
+        "--flush_frequency",
+        type=int,
+        default=256,
+        help="how often to save to parquet files"
+    )
+    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
+    parser.add_argument("--max_height", type=int, default=480)
+    parser.add_argument("--max_width", type=int, default=848)
+    parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
+    parser.add_argument("--dataset", default="t2v")
+    parser.add_argument("--train_fps", type=int, default=30)
+    parser.add_argument("--use_image_num", type=int, default=0)
+    parser.add_argument("--text_max_length", type=int, default=256)
+    parser.add_argument("--speed_factor", type=float, default=1.0)
+    parser.add_argument("--drop_short_ratio", type=float, default=1.0)
+    # text encoder & vae & diffusion model
+    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+    parser.add_argument("--cfg", type=float, default=0.0)
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
+
+    args = parser.parse_args()
+    main(args)

fastvideo/data_preprocess/preprocess_text_embeddings_v1.py

@@ -0,0 +1,199 @@
+import argparse
+import json
+import os
+
+import torch
+import torch.distributed as dist
+from accelerate.logging import get_logger
+from diffusers.utils import export_to_video
+from diffusers.video_processor import VideoProcessor
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+
+# from fastvideo.utils.load import load_text_encoder, load_vae
+
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.loader.component_loader import VAELoader, TextEncoderLoader, TokenizerLoader
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.configs.models.vaes import WanVAEConfig
+from fastvideo.v1.configs.models.encoders.t5 import T5Config
+
+logger = get_logger(__name__)
+
+class T5dataset(Dataset):
+
+    def __init__(
+        self,
+        json_path,
+        vae_debug,
+    ):
+        self.json_path = json_path
+        self.vae_debug = vae_debug
+        with open(self.json_path, "r") as f:
+            train_dataset = json.load(f)
+            self.train_dataset = sorted(train_dataset, key=lambda x: x["latent_path"])
+
+    def __getitem__(self, idx):
+        caption = self.train_dataset[idx]["caption"]
+        filename = self.train_dataset[idx]["latent_path"].split(".")[0]
+        length = self.train_dataset[idx]["length"]
+        if self.vae_debug:
+            latents = torch.load(
+                os.path.join(args.output_dir, "latent", self.train_dataset[idx]["latent_path"]),
+                map_location="cpu",
+            )
+        else:
+            latents = []
+
+        return dict(caption=caption, latents=latents, filename=filename, length=length)
+
+    def __len__(self):
+        return len(self.train_dataset)
+
+
+def main(args):
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    rank = int(os.getenv("RANK", 0))
+    init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
+    print("world_size", world_size, "local rank", local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    torch.cuda.set_device(device)
+    world_group = get_world_group()
+
+    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # torch.cuda.set_device(local_rank)
+    # if not dist.is_initialized():
+    #     dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
+
+    videoprocessor = VideoProcessor(vae_scale_factor=8)
+    os.makedirs(args.output_dir, exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, "video"), exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, "prompt_embed"), exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, "prompt_attention_mask"), exist_ok=True)
+
+    vae_precision = "fp16" 
+    text_encoder_precision = "fp32"
+    fastvideo_args = FastVideoArgs(model_path=args.model_path,
+                         use_cpu_offload=False,
+                         vae_precision=vae_precision,
+                         text_encoder_precisions=(text_encoder_precision,))
+    fastvideo_args.device = device
+    fastvideo_args.device_str = f"cuda:{local_rank}"
+
+    # fastvideo_args.dit_config = HunyuanVideoConfig()
+    fastvideo_args.vae_config = WanVAEConfig()
+    fastvideo_args.text_encoder_configs = (T5Config(),)
+
+    # vae_loader = VAELoader()
+    # vae = vae_loader.load_vae()
+    text_encoder_loader = TextEncoderLoader()
+    tokenizer_loader = TokenizerLoader()
+
+    model_path = args.model_path
+    path = maybe_download_model(model_path)
+    # PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    ENCODER_PATH = os.path.join(path, "text_encoder")
+    TOKENIZER_PATH = os.path.join(path, "tokenizer")
+    print(ENCODER_PATH)
+    text_encoder = text_encoder_loader.load(ENCODER_PATH, "text_encoder", fastvideo_args)
+    tokenizer = tokenizer_loader.load(TOKENIZER_PATH, "tokenizer", fastvideo_args)
+
+    latents_json_path = os.path.join(args.output_dir, "videos2caption_temp.json")
+    train_dataset = T5dataset(latents_json_path, args.vae_debug)
+    # text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
+    # vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
+    # vae.enable_tiling()
+    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
+    train_dataloader = DataLoader(
+        train_dataset,
+        sampler=sampler,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+    )
+
+    json_data = []
+    for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
+        with torch.inference_mode():
+            # with torch.autocast("cuda", dtype=torch.float32):
+            print(data["caption"])
+            text_inputs = tokenizer(data["caption"], **fastvideo_args.text_encoder_configs[0].tokenizer_kwargs).to(
+                fastvideo_args.device)
+            input_ids = text_inputs["input_ids"]
+            attention_mask = text_inputs["attention_mask"]
+            outputs = text_encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+            from fastvideo.v1.configs.pipelines.wan import t5_postprocess_text
+            post_process_func = t5_postprocess_text
+            prompt_embeds = post_process_func(outputs)
+            prompt_attention_mask = attention_mask
+            if args.vae_debug:
+                latents = data["latents"]
+                video = vae.decode(latents.to(device), return_dict=False)[0]
+                video = videoprocessor.postprocess_video(video)
+            for idx, video_name in enumerate(data["filename"]):
+                prompt_embed_path = os.path.join(args.output_dir, "prompt_embed", video_name + ".pt")
+                video_path = os.path.join(args.output_dir, "video", video_name + ".mp4")
+                prompt_attention_mask_path = os.path.join(args.output_dir, "prompt_attention_mask",
+                                                            video_name + ".pt")
+                # save latent
+                torch.save(prompt_embeds[idx], prompt_embed_path)
+                torch.save(prompt_attention_mask[idx], prompt_attention_mask_path)
+                print(f"sample {video_name} saved")
+                if args.vae_debug:
+                    export_to_video(video[idx], video_path, fps=16)
+                item = {}
+                item["length"] = int(data["length"][idx])
+                item["latent_path"] = video_name + ".pt"
+                item["prompt_embed_path"] = video_name + ".pt"
+                item["prompt_attention_mask"] = video_name + ".pt"
+                item["caption"] = data["caption"][idx]
+                json_data.append(item)
+    dist.barrier()
+    local_data = json_data
+    gathered_data = [None] * world_size
+    dist.all_gather_object(gathered_data, local_data)
+    if local_rank == 0:
+        # os.remove(latents_json_path)
+        all_json_data = [item for sublist in gathered_data for item in sublist]
+        with open(os.path.join(args.output_dir, "videos2caption.json"), "w") as f:
+            json.dump(all_json_data, f, indent=4)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # dataset & dataloader
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    # parser.add_argument("--model_type", type=str, default="mochi")
+    # text encoder & vae & diffusion model
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=1,
+        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--train_batch_size",
+        type=int,
+        default=1,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument("--vae_debug", action="store_true")
+    args = parser.parse_args()
+    main(args)

fastvideo/data_preprocess/preprocess_vae_latents_v1.py

@@ -0,0 +1,151 @@
+import argparse
+import json
+import os
+
+import torch
+# import torch.distributed as dist
+# from accelerate.logging import get_logger
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from tqdm import tqdm
+
+from fastvideo.dataset import getdataset
+# from fastvideo.utils.load import load_vae
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.loader.component_loader import VAELoader
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.configs.models.vaes import WanVAEConfig
+
+logger = init_logger(__name__)
+
+model_path = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+path = maybe_download_model(model_path)
+# PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+VAE_PATH = os.path.join(path, "vae")
+print(VAE_PATH)
+
+
+def main(args):
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    rank = int(os.getenv("RANK", 0))
+    init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
+    print("world_size", world_size, "local rank", local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    torch.cuda.set_device(device)
+    world_group = get_world_group()
+
+    vae_precision = "fp16" 
+    fastvideo_args = FastVideoArgs(model_path=VAE_PATH,
+                         use_cpu_offload=False,
+                         vae_precision=vae_precision)
+    fastvideo_args.device = device
+    # fastvideo_args.dit_config = HunyuanVideoConfig()
+    fastvideo_args.vae_config = WanVAEConfig()
+
+    train_dataset = getdataset(args)
+    sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
+    train_dataloader = DataLoader(
+        train_dataset,
+        sampler=sampler,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+    )
+
+
+    # encoder_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # torch.cuda.set_device(local_rank)
+    # if not dist.is_initialized():
+    #     dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
+    vae_loader = VAELoader()
+    vae = vae_loader.load(VAE_PATH, "vae", fastvideo_args)
+    # vae, autocast_type, fps = load_vae(args.model_type, args.model_path)
+    # vae.enable_tiling()
+    os.makedirs(args.output_dir, exist_ok=True)
+    os.makedirs(os.path.join(args.output_dir, "latent"), exist_ok=True)
+
+    json_data = []
+    for _, data in tqdm(enumerate(train_dataloader), disable=local_rank != 0):
+        with torch.inference_mode():
+            with torch.autocast("cuda", dtype=torch.float16):
+                latents = vae.encode(data["pixel_values"].to(device)).sample()
+            for idx, video_path in enumerate(data["path"]):
+                video_name = os.path.basename(video_path).split(".")[0]
+                latent_path = os.path.join(args.output_dir, "latent", video_name + ".pt")
+                torch.save(latents[idx].to(torch.bfloat16), latent_path)
+                item = {}
+                item["length"] = latents[idx].shape[1]
+                item["latent_path"] = video_name + ".pt"
+                item["caption"] = data["text"][idx]
+                json_data.append(item)
+                print(f"{video_name} processed")
+    world_group.barrier()
+    local_data = json_data
+    gathered_data = [None] * world_size
+    for i in range(world_size):
+        if local_rank == i:
+            world_group.broadcast_object(local_data, src=i)
+        else:
+            gathered_data[i] = world_group.broadcast_object(None, src=i)
+    gathered_data[local_rank] = json_data
+    print(gathered_data)
+    if local_rank == 0:
+        all_json_data = [item for sublist in gathered_data for item in sublist]
+        with open(os.path.join(args.output_dir, "videos2caption_temp.json"), "w") as f:
+            json.dump(all_json_data, f, indent=4)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # dataset & dataloader
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    # parser.add_argument("--model_type", type=str, default="mochi")
+    parser.add_argument("--data_merge_path", type=str, required=True)
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=1,
+        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--train_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
+    parser.add_argument("--max_height", type=int, default=480)
+    parser.add_argument("--max_width", type=int, default=848)
+    parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
+    parser.add_argument("--dataset", default="t2v")
+    parser.add_argument("--train_fps", type=int, default=30)
+    parser.add_argument("--use_image_num", type=int, default=0)
+    parser.add_argument("--text_max_length", type=int, default=256)
+    parser.add_argument("--speed_factor", type=float, default=1.0)
+    parser.add_argument("--drop_short_ratio", type=float, default=1.0)
+    # text encoder & vae & diffusion model
+    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+    parser.add_argument("--cfg", type=float, default=0.0)
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
+
+    args = parser.parse_args()
+    main(args)

fastvideo/data_preprocess/preprocess_validation_text_embeddings_v1.py

@@ -0,0 +1,115 @@
+import argparse
+import os
+
+import torch
+# import torch.distributed as dist
+from accelerate.logging import get_logger
+
+# from fastvideo.utils.load import load_text_encoder
+
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.models.loader.component_loader import VAELoader, TextEncoderLoader, TokenizerLoader
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, get_world_group
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.configs.models.vaes import WanVAEConfig
+from fastvideo.v1.configs.models.encoders.t5 import T5Config
+
+logger = get_logger(__name__)
+
+
+def main(args):
+    local_rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    rank = int(os.getenv("RANK", 0))
+    init_distributed_environment(rank=rank, world_size=world_size, local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
+    print("world_size", world_size, "local rank", local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    torch.cuda.set_device(device)
+    world_group = get_world_group()
+
+    vae_precision = "fp16" 
+    text_encoder_precision = "fp32"
+    fastvideo_args = FastVideoArgs(model_path=args.model_path,
+                         use_cpu_offload=False,
+                         vae_precision=vae_precision,
+                         text_encoder_precisions=(text_encoder_precision,))
+    fastvideo_args.device = device
+    fastvideo_args.device_str = f"cuda:{local_rank}"
+
+    # fastvideo_args.dit_config = HunyuanVideoConfig()
+    fastvideo_args.vae_config = WanVAEConfig()
+    fastvideo_args.text_encoder_configs = (T5Config(),)
+
+    # vae_loader = VAELoader()
+    # vae = vae_loader.load_vae()
+    text_encoder_loader = TextEncoderLoader()
+    tokenizer_loader = TokenizerLoader()
+
+    model_path = args.model_path
+    path = maybe_download_model(model_path)
+    # PIPELINE_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    ENCODER_PATH = os.path.join(path, "text_encoder")
+    TOKENIZER_PATH = os.path.join(path, "tokenizer")
+    print(ENCODER_PATH)
+    text_encoder = text_encoder_loader.load(ENCODER_PATH, "text_encoder", fastvideo_args)
+    tokenizer = tokenizer_loader.load(TOKENIZER_PATH, "tokenizer", fastvideo_args)
+
+    # text_encoder = load_text_encoder(args.model_type, args.model_path, device=device)
+    # autocast_type = torch.float16 if args.model_type == "hunyuan" else torch.bfloat16
+    # output_dir/validation/prompt_attention_mask
+    # output_dir/validation/prompt_embed
+    os.makedirs(os.path.join(args.output_dir, "validation"), exist_ok=True)
+    os.makedirs(
+        os.path.join(args.output_dir, "validation", "prompt_attention_mask"),
+        exist_ok=True,
+    )
+    os.makedirs(os.path.join(args.output_dir, "validation", "prompt_embed"), exist_ok=True)
+
+    with open(args.validation_prompt_txt, "r", encoding="utf-8") as file:
+        lines = file.readlines()
+    prompts = [line.strip() for line in lines]
+    for prompt in prompts:
+        with torch.inference_mode():
+            # with torch.autocast("cuda", dtype=autocast_type):
+            text_inputs = tokenizer(prompt, **fastvideo_args.text_encoder_configs[0].tokenizer_kwargs).to(
+                fastvideo_args.device)
+            input_ids = text_inputs["input_ids"]
+            attention_mask = text_inputs["attention_mask"]
+            outputs = text_encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+            from fastvideo.v1.configs.pipelines.wan import t5_postprocess_text
+            post_process_func = t5_postprocess_text
+            prompt_embeds = post_process_func(outputs)
+            prompt_attention_mask = attention_mask
+
+            file_name = prompt.split(".")[0]
+            prompt_embed_path = os.path.join(args.output_dir, "validation", "prompt_embed", f"{file_name}.pt")
+            prompt_attention_mask_path = os.path.join(
+                args.output_dir,
+                "validation",
+                "prompt_attention_mask",
+                f"{file_name}.pt",
+            )
+            torch.save(prompt_embeds[0], prompt_embed_path)
+            torch.save(prompt_attention_mask[0], prompt_attention_mask_path)
+            print(f"sample {file_name} saved")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # dataset & dataloader
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--validation_prompt_txt", type=str)
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    args = parser.parse_args()
+    main(args)

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.time()
+        start_time = time.perf_counter()
         assert args.multi_phased_distill_schedule is not None
         num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
 
@@ -482,7 +482,7 @@ def main(args):
             args.hunyuan_teacher_disable_cfg,
         )
 
-        step_time = time.time() - start_time
+        step_time = time.perf_counter() - start_time
         step_times.append(step_time)
         avg_step_time = sum(step_times) / len(step_times)
 

fastvideo/distill/solver.py

@@ -7,7 +7,7 @@ from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.schedulers.scheduling_utils import SchedulerMixin
 from diffusers.utils import BaseOutput, logging
 
-from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
+# from fastvideo.models.mochi_hf.pipeline_mochi import linear_quadratic_schedule
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -38,6 +38,7 @@ class PCMFMScheduler(SchedulerMixin, ConfigMixin):
         linear_range=0.5,
     ):
         if linear_quadratic:
+            raise NotImplementedError("Linear quadratic schedule is not implemented")
             linear_steps = int(num_train_timesteps * linear_range)
             sigmas = linear_quadratic_schedule(num_train_timesteps, linear_quadratic_threshold, linear_steps)
             sigmas = torch.tensor(sigmas).to(dtype=torch.float32)

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

fastvideo/distill_wan.py

@@ -0,0 +1,870 @@
+# !/bin/python3
+# isort: skip_file
+import argparse
+import math
+import os
+import time
+from collections import deque
+
+import torch
+import torch.distributed as dist
+import wandb
+from accelerate.utils import set_seed
+from diffusers.optimization import get_scheduler
+from diffusers.utils import check_min_version
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import ShardingStrategy
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from tqdm.auto import tqdm
+
+from fastvideo.dataset.latent_datasets import (LatentDataset,
+                                               latent_collate_function)
+from fastvideo.distill.solver import EulerSolver, extract_into_tensor
+from fastvideo.utils.checkpoint import (save_checkpoint, save_lora_checkpoint)
+from fastvideo.utils.communications import (broadcast,
+                                            sp_parallel_dataloader_wrapper)
+from fastvideo.utils.dataset_utils import LengthGroupedSampler
+from fastvideo.utils.parallel_states import (destroy_sequence_parallel_group,
+                                             get_sequence_parallel_state)
+from fastvideo.utils.validation import log_validation
+
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.configs.models.dits import WanVideoConfig
+from fastvideo.v1.models.loader.component_loader import TransformerLoader, SchedulerLoader
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.31.0")
+
+logger = init_logger(__name__)
+
+BASE_MODEL_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+SCHEDULER_PATH = os.path.join(MODEL_PATH, "scheduler")
+
+
+def reshard_fsdp(model):
+    for m in FSDP.fsdp_modules(model):
+        if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD:
+            torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True)
+
+
+def get_norm(model_pred, norms, gradient_accumulation_steps):
+    fro_norm = (
+        torch.linalg.matrix_norm(model_pred, ord="fro") /  # codespell:ignore
+        gradient_accumulation_steps)
+    largest_singular_value = (torch.linalg.matrix_norm(model_pred, ord=2) /
+                              gradient_accumulation_steps)
+    absolute_mean = torch.mean(
+        torch.abs(model_pred)) / gradient_accumulation_steps
+    absolute_max = torch.max(
+        torch.abs(model_pred)) / gradient_accumulation_steps
+    dist.all_reduce(fro_norm, op=dist.ReduceOp.AVG)
+    dist.all_reduce(largest_singular_value, op=dist.ReduceOp.AVG)
+    dist.all_reduce(absolute_mean, op=dist.ReduceOp.AVG)
+    norms["fro"] += torch.mean(fro_norm).item()  # codespell:ignore
+    norms["largest singular value"] += torch.mean(largest_singular_value).item()
+    norms["absolute mean"] += absolute_mean.item()
+    norms["absolute max"] += absolute_max.item()
+
+
+def distill_one_step(
+    transformer,
+    model_type,
+    teacher_transformer,
+    ema_transformer,
+    optimizer,
+    lr_scheduler,
+    loader,
+    noise_scheduler,
+    solver,
+    noise_random_generator,
+    gradient_accumulation_steps,
+    sp_size,
+    max_grad_norm,
+    uncond_prompt_embed,
+    uncond_prompt_mask,
+    num_euler_timesteps,
+    multiphase,
+    not_apply_cfg_solver,
+    distill_cfg,
+    ema_decay,
+    pred_decay_weight,
+    pred_decay_type,
+    hunyuan_teacher_disable_cfg,
+):
+    total_loss = 0.0
+    optimizer.zero_grad()
+    model_pred_norm = {
+        "fro": 0.0,  # codespell:ignore
+        "largest singular value": 0.0,
+        "absolute mean": 0.0,
+        "absolute max": 0.0,
+    }
+    for _ in range(gradient_accumulation_steps):
+        (
+            latents,
+            encoder_hidden_states,
+            latents_attention_mask,
+            encoder_attention_mask,
+        ) = next(loader)
+        # model_input = normalize_dit_input(model_type, latents)
+        model_input = latents
+        noise = torch.randn_like(model_input)
+        bsz = model_input.shape[0]
+        index = torch.randint(0,
+                              num_euler_timesteps, (bsz, ),
+                              device=model_input.device).long()
+        if sp_size > 1:
+            broadcast(index)
+        # Add noise according to flow matching.
+        # sigmas = get_sigmas(start_timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
+        sigmas = extract_into_tensor(solver.sigmas, index, model_input.shape)
+        sigmas_prev = extract_into_tensor(solver.sigmas_prev, index,
+                                          model_input.shape)
+
+        timesteps = (sigmas *
+                     noise_scheduler.config.num_train_timesteps).view(-1)
+        # if squeeze to [], unsqueeze to [1]
+
+        timesteps_prev = (sigmas_prev *
+                          noise_scheduler.config.num_train_timesteps).view(-1)
+        noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
+        noisy_model_input = noisy_model_input.to(torch.bfloat16)
+
+        forward_batch = ForwardBatch(data_type="video", enable_teacache=False)
+        # Predict the noise residual
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            teacher_kwargs = {
+                "hidden_states": noisy_model_input,
+                "encoder_hidden_states": encoder_hidden_states,
+                "timestep": timesteps,
+                "encoder_attention_mask": encoder_attention_mask,  # B, L
+                "return_dict": False,
+            }
+            if hunyuan_teacher_disable_cfg:
+                teacher_kwargs["guidance"] = torch.tensor(
+                    [1000.0],
+                    device=noisy_model_input.device,
+                    dtype=torch.bfloat16)
+            with set_forward_context(current_timestep=0,
+                                     attn_metadata=None,
+                                     forward_batch=forward_batch):
+                with torch.autograd.graph.save_on_cpu(pin_memory=True):
+                    model_pred = transformer(**teacher_kwargs)
+
+        # if accelerator.is_main_process:
+        model_pred, end_index = solver.euler_style_multiphase_pred(
+            noisy_model_input, model_pred, index, multiphase)
+        with torch.no_grad():
+            w = distill_cfg
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                with set_forward_context(current_timestep=0,
+                                         attn_metadata=None,
+                                         forward_batch=forward_batch):
+                    cond_teacher_output = teacher_transformer(
+                        noisy_model_input,
+                        encoder_hidden_states,
+                        timesteps,
+                        encoder_attention_mask,  # B, L
+                        return_dict=False,
+                    ).float()
+            if not_apply_cfg_solver:
+                uncond_teacher_output = cond_teacher_output
+            else:
+                # Get teacher model prediction on noisy_latents and unconditional embedding
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    with set_forward_context(current_timestep=0,
+                                             attn_metadata=None,
+                                             forward_batch=forward_batch):
+                        uncond_teacher_output = teacher_transformer(
+                            noisy_model_input,
+                            uncond_prompt_embed.unsqueeze(0).expand(
+                                bsz, -1, -1),
+                            timesteps,
+                            uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
+                            return_dict=False,
+                        ).float()
+            teacher_output = uncond_teacher_output + w * (cond_teacher_output -
+                                                          uncond_teacher_output)
+            x_prev = solver.euler_step(noisy_model_input, teacher_output,
+                                       index).to(torch.bfloat16)
+
+        # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
+        with torch.no_grad():
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                if ema_transformer is not None:
+                    target_pred = ema_transformer(
+                        x_prev.float(),
+                        encoder_hidden_states,
+                        timesteps_prev,
+                        encoder_attention_mask,  # B, L
+                        return_dict=False,
+                    )[0]
+                else:
+                    with set_forward_context(current_timestep=0,
+                                             attn_metadata=None,
+                                             forward_batch=forward_batch):
+                        with torch.autograd.graph.save_on_cpu(pin_memory=True):
+                            target_pred = transformer(
+                                x_prev,
+                                encoder_hidden_states,
+                                timesteps_prev,
+                                encoder_attention_mask,  # B, L
+                                return_dict=False,
+                            )
+
+            target, end_index = solver.euler_style_multiphase_pred(
+                x_prev, target_pred, index, multiphase, True)
+
+        huber_c = 0.001
+        # loss = loss.mean()
+        loss = (torch.mean(
+            torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) -
+            huber_c) / gradient_accumulation_steps)
+        if pred_decay_weight > 0:
+            if pred_decay_type == "l1":
+                pred_decay_loss = (
+                    torch.mean(torch.sqrt(model_pred.float()**2)) *
+                    pred_decay_weight / gradient_accumulation_steps)
+                loss += pred_decay_loss
+            elif pred_decay_type == "l2":
+                # essnetially k2?
+                pred_decay_loss = (torch.mean(model_pred.float()**2) *
+                                   pred_decay_weight /
+                                   gradient_accumulation_steps)
+                loss += pred_decay_loss
+            else:
+                assert NotImplementedError("pred_decay_type is not implemented")
+
+        # calculate model_pred norm and mean
+        get_norm(model_pred.detach().float(), model_pred_norm,
+                 gradient_accumulation_steps)
+        loss.backward()
+
+        avg_loss = loss.detach().clone()
+        dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
+        total_loss += avg_loss.item()
+
+    # update ema
+    if ema_transformer is not None:
+        reshard_fsdp(ema_transformer)
+        for p_averaged, p_model in zip(ema_transformer.parameters(),
+                                       transformer.parameters()):
+            with torch.no_grad():
+                p_averaged.copy_(
+                    torch.lerp(p_averaged.detach(), p_model.detach(),
+                               1 - ema_decay))
+
+    # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
+    grad_norm = torch.nn.utils.clip_grad_norm_(transformer.parameters(),
+                                               max_norm=max_grad_norm)
+    optimizer.step()
+    lr_scheduler.step()
+
+    return total_loss, grad_norm.item(), model_pred_norm
+
+
+def main(args):
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    rank = int(os.environ.get("RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    torch.cuda.set_device(rank)
+    init_distributed_environment(world_size=world_size,
+                                 rank=rank,
+                                 local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=args.sp_size,
+                              sequence_model_parallel_size=args.sp_size)
+    fastvideo_args = FastVideoArgs(
+        model_path=MODEL_PATH,
+        num_gpus=world_size,
+        use_cpu_offload=False,
+        precision=args.master_weight_type,
+        dit_config=WanVideoConfig(),
+        device_str="cuda",
+    )
+    fastvideo_args.check_fastvideo_args()
+
+    device_str = f"cuda:{rank}"
+    device = torch.device(device_str)
+    fastvideo_args.device = device
+
+    # If passed along, set the training seed now. On GPU...
+    if args.seed is not None:
+        # TODO: t within the same seq parallel group should be the same. Noise should be different.
+        set_seed(args.seed + rank)
+    # We use different seeds for the noise generation in each process to ensure that the noise is different in a batch.
+    noise_random_generator = None
+
+    # Handle the repository creation
+    if rank <= 0 and args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    # For mixed precision training we cast all non-trainable weights to half-precision
+    # as these weights are only used for inference, keeping weights in full precision is not required.
+
+    # Create model:
+
+    logger.info("--> loading model from %s", TRANSFORMER_PATH)
+
+    fastvideo_args.device = device
+    transformer_loader = TransformerLoader()
+    transformer = transformer_loader.load(TRANSFORMER_PATH, "", fastvideo_args)
+    transformer = transformer.train()
+    transformer.requires_grad_(True)
+
+    teacher_loader = TransformerLoader()
+    teacher_transformer = teacher_loader.load(TRANSFORMER_PATH, "",
+                                              fastvideo_args)
+    if args.use_ema:
+        ema_transformer = teacher_loader.load(TRANSFORMER_PATH, "",
+                                              fastvideo_args)
+    else:
+        ema_transformer = None
+
+    logger.info(
+        "  Total training parameters = %s M",
+        sum(p.numel()
+            for p in transformer.parameters() if p.requires_grad) / 1e6)
+    logger.info("--> model loaded")
+
+    teacher_transformer.requires_grad_(False)
+    if args.use_ema:
+        ema_transformer.requires_grad_(False)
+
+    # scheduler
+    noise_scheduler_loader = SchedulerLoader()
+    noise_scheduler = noise_scheduler_loader.load(SCHEDULER_PATH, "",
+                                                  fastvideo_args)
+    solver = EulerSolver(
+        noise_scheduler.sigmas.numpy()[::-1],
+        noise_scheduler.config.num_train_timesteps,
+        euler_timesteps=args.num_euler_timesteps,
+    )
+    solver.to(device)
+    params_to_optimize = transformer.parameters()
+    params_to_optimize = list(
+        filter(lambda p: p.requires_grad, params_to_optimize))
+
+    optimizer = torch.optim.AdamW(
+        params_to_optimize,
+        lr=args.learning_rate,
+        betas=(0.9, 0.999),
+        weight_decay=args.weight_decay,
+        eps=1e-8,
+    )
+
+    init_steps = 0
+    logger.info("optimizer: %s", optimizer)
+
+    # todo add lr scheduler
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * world_size,
+        num_training_steps=args.max_train_steps * world_size,
+        num_cycles=args.lr_num_cycles,
+        power=args.lr_power,
+        last_epoch=init_steps - 1,
+    )
+
+    train_dataset = LatentDataset(args.data_json_path, args.num_latent_t,
+                                  args.cfg)
+    uncond_prompt_embed = train_dataset.uncond_prompt_embed
+    uncond_prompt_mask = train_dataset.uncond_prompt_mask
+    sampler = (LengthGroupedSampler(
+        args.train_batch_size,
+        rank=rank,
+        world_size=world_size,
+        lengths=train_dataset.lengths,
+        group_frame=args.group_frame,
+        group_resolution=args.group_resolution,
+    ) if (args.group_frame or args.group_resolution) else DistributedSampler(
+        train_dataset, rank=rank, num_replicas=world_size, shuffle=False))
+
+    train_dataloader = DataLoader(
+        train_dataset,
+        sampler=sampler,
+        collate_fn=latent_collate_function,
+        pin_memory=True,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+        drop_last=True,
+    )
+
+    num_update_steps_per_epoch = math.ceil(
+        len(train_dataloader) / args.gradient_accumulation_steps *
+        args.sp_size / args.train_sp_batch_size)
+    args.num_train_epochs = math.ceil(args.max_train_steps /
+                                      num_update_steps_per_epoch)
+
+    if rank <= 0:
+        project = args.tracker_project_name or "fastvideo"
+        wandb.init(project=project, config=args)
+
+    # Train!
+    total_batch_size = (world_size * args.gradient_accumulation_steps /
+                        args.sp_size * args.train_sp_batch_size)
+    logger.info("***** Running training *****")
+    logger.info("  Num examples = %s", len(train_dataset))
+    logger.info("  Dataloader size = %s", len(train_dataloader))
+    logger.info("  Num Epochs = %s", args.num_train_epochs)
+    logger.info("  Resume training from step %s", init_steps)
+    logger.info("  Instantaneous batch size per device = %s",
+                args.train_batch_size)
+    logger.info(
+        "  Total train batch size (w. data & sequence parallel, accumulation) = %s",
+        total_batch_size)
+    logger.info("  Gradient Accumulation steps = %s",
+                args.gradient_accumulation_steps)
+    logger.info("  Total optimization steps = %s", args.max_train_steps)
+    logger.info(
+        "  Total training parameters per FSDP shard = %s B",
+        sum(p.numel()
+            for p in transformer.parameters() if p.requires_grad) / 1e9)
+    # print dtype
+    logger.info("  Master weight dtype: %s",
+                transformer.parameters().__next__().dtype)
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        assert NotImplementedError(
+            "resume_from_checkpoint is not supported now.")
+        # TODO
+
+    progress_bar = tqdm(
+        range(0, args.max_train_steps),
+        initial=init_steps,
+        desc="Steps",
+        # Only show the progress bar once on each machine.
+        disable=local_rank > 0,
+    )
+
+    loader = sp_parallel_dataloader_wrapper(
+        train_dataloader,
+        device,
+        args.train_batch_size,
+        args.sp_size,
+        args.train_sp_batch_size,
+    )
+
+    step_times = deque(maxlen=100)
+
+    # todo future
+    for i in range(init_steps):
+        next(loader)
+
+    # log_validation(args, transformer, device,
+    #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
+    def get_num_phases(multi_phased_distill_schedule, step):
+        # step-phase,step-phase
+        multi_phases = multi_phased_distill_schedule.split(",")
+        phase = multi_phases[-1].split("-")[-1]
+        for step_phases in multi_phases:
+            phase_step, phase = step_phases.split("-")
+            if step <= int(phase_step):
+                return int(phase)
+        return phase
+
+    for step in range(init_steps + 1, args.max_train_steps + 1):
+        start_time = time.time()
+        assert args.multi_phased_distill_schedule is not None
+        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
+
+        loss, grad_norm, pred_norm = distill_one_step(
+            transformer,
+            args.model_type,
+            teacher_transformer,
+            ema_transformer,
+            optimizer,
+            lr_scheduler,
+            loader,
+            noise_scheduler,
+            solver,
+            noise_random_generator,
+            args.gradient_accumulation_steps,
+            args.sp_size,
+            args.max_grad_norm,
+            uncond_prompt_embed,
+            uncond_prompt_mask,
+            args.num_euler_timesteps,
+            num_phases,
+            args.not_apply_cfg_solver,
+            args.distill_cfg,
+            args.ema_decay,
+            args.pred_decay_weight,
+            args.pred_decay_type,
+            args.hunyuan_teacher_disable_cfg,
+        )
+
+        step_time = time.time() - start_time
+        step_times.append(step_time)
+        avg_step_time = sum(step_times) / len(step_times)
+
+        progress_bar.set_postfix({
+            "loss": f"{loss:.4f}",
+            "step_time": f"{step_time:.2f}s",
+            "grad_norm": grad_norm,
+            "phases": num_phases,
+        })
+        progress_bar.update(1)
+        if rank <= 0:
+            wandb.log(
+                {
+                    "train_loss":
+                    loss,
+                    "learning_rate":
+                    lr_scheduler.get_last_lr()[0],
+                    "step_time":
+                    step_time,
+                    "avg_step_time":
+                    avg_step_time,
+                    "grad_norm":
+                    grad_norm,
+                    "pred_fro_norm":
+                    pred_norm["fro"],  # codespell:ignore
+                    "pred_largest_singular_value":
+                    pred_norm["largest singular value"],
+                    "pred_absolute_mean":
+                    pred_norm["absolute mean"],
+                    "pred_absolute_max":
+                    pred_norm["absolute max"],
+                },
+                step=step,
+            )
+        if step % args.checkpointing_steps == 0:
+            if args.use_lora:
+                # Save LoRA weights
+                save_lora_checkpoint(transformer, optimizer, rank,
+                                     args.output_dir, step)
+            else:
+                # Your existing checkpoint saving code
+                if args.use_ema:
+                    save_checkpoint(ema_transformer, rank, args.output_dir,
+                                    step)
+                else:
+                    save_checkpoint(transformer, rank, args.output_dir, step)
+            dist.barrier()
+        if args.log_validation and step % args.validation_steps == 0:
+            log_validation(
+                args,
+                transformer,
+                device,
+                torch.bfloat16,
+                step,
+                scheduler_type=args.scheduler_type,
+                shift=args.shift,
+                num_euler_timesteps=args.num_euler_timesteps,
+                linear_quadratic_threshold=args.linear_quadratic_threshold,
+                linear_range=args.linear_range,
+                ema=False,
+            )
+            if args.use_ema:
+                log_validation(
+                    args,
+                    ema_transformer,
+                    device,
+                    torch.bfloat16,
+                    step,
+                    scheduler_type=args.scheduler_type,
+                    shift=args.shift,
+                    num_euler_timesteps=args.num_euler_timesteps,
+                    linear_quadratic_threshold=args.linear_quadratic_threshold,
+                    linear_range=args.linear_range,
+                    ema=True,
+                )
+
+    if args.use_lora:
+        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir,
+                             args.max_train_steps)
+    else:
+        save_checkpoint(transformer, rank, args.output_dir,
+                        args.max_train_steps)
+
+    if get_sequence_parallel_state():
+        destroy_sequence_parallel_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--model_type",
+                        type=str,
+                        default="mochi",
+                        help="The type of model to train.")
+
+    # dataset & dataloader
+    parser.add_argument("--data_json_path", type=str, required=True)
+    parser.add_argument("--num_height", type=int, default=480)
+    parser.add_argument("--num_width", type=int, default=848)
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=10,
+        help=
+        "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--train_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument("--num_latent_t",
+                        type=int,
+                        default=28,
+                        help="Number of latent timesteps.")
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
+
+    # text encoder & vae & diffusion model
+    parser.add_argument("--pretrained_model_name_or_path", type=str)
+    parser.add_argument("--dit_model_name_or_path", type=str)
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+
+    # diffusion setting
+    parser.add_argument("--ema_decay", type=float, default=0.95)
+    parser.add_argument("--ema_start_step", type=int, default=0)
+    parser.add_argument("--cfg", type=float, default=0.1)
+
+    # validation & logs
+    parser.add_argument("--validation_prompt_dir", type=str)
+    parser.add_argument("--validation_sampling_steps", type=str, default="64")
+    parser.add_argument("--validation_guidance_scale", type=str, default="4.5")
+
+    parser.add_argument("--validation_steps", type=float, default=64)
+    parser.add_argument("--log_validation", action="store_true")
+    parser.add_argument("--tracker_project_name", type=str, default=None)
+    parser.add_argument("--seed",
+                        type=int,
+                        default=None,
+                        help="A seed for reproducible training.")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help=
+        "The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=("Max number of checkpoints to store."),
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help=
+        ("Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+         " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+         " training using `--resume_from_checkpoint`."),
+    )
+    parser.add_argument("--shift", type=float, default=1.0)
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=
+        ("Whether training should be resumed from a previous checkpoint. Use a path saved by"
+         ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+         ),
+    )
+    parser.add_argument(
+        "--resume_from_lora_checkpoint",
+        type=str,
+        default=None,
+        help=
+        ("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
+         ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
+         ),
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=
+        ("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+         " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
+
+    # optimizer & scheduler & Training
+    parser.add_argument("--num_train_epochs", type=int, default=100)
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help=
+        "Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help=
+        "Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-4,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help=
+        "Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_warmup_steps",
+        type=int,
+        default=10,
+        help="Number of steps for the warmup in the lr scheduler.",
+    )
+    parser.add_argument("--max_grad_norm",
+                        default=1.0,
+                        type=float,
+                        help="Max gradient norm.")
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help=
+        "Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument("--selective_checkpointing", type=float, default=1.0)
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=
+        ("Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+         " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
+         ),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=
+        ("Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+         " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+         " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+         ),
+    )
+    parser.add_argument(
+        "--use_cpu_offload",
+        action="store_true",
+        help=
+        "Whether to use CPU offload for param & gradient & optimizer states.",
+    )
+
+    parser.add_argument("--sp_size",
+                        type=int,
+                        default=1,
+                        help="For sequence parallel")
+    parser.add_argument(
+        "--train_sp_batch_size",
+        type=int,
+        default=1,
+        help="Batch size for sequence parallel training",
+    )
+
+    parser.add_argument(
+        "--use_lora",
+        action="store_true",
+        default=False,
+        help="Whether to use LoRA for finetuning.",
+    )
+    parser.add_argument("--lora_alpha",
+                        type=int,
+                        default=256,
+                        help="Alpha parameter for LoRA.")
+    parser.add_argument("--lora_rank",
+                        type=int,
+                        default=128,
+                        help="LoRA rank parameter. ")
+    parser.add_argument("--fsdp_sharding_startegy", default="full")
+
+    # lr_scheduler
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=
+        ('The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+         ' "constant", "constant_with_warmup"]'),
+    )
+    parser.add_argument("--num_euler_timesteps", type=int, default=100)
+    parser.add_argument(
+        "--lr_num_cycles",
+        type=int,
+        default=1,
+        help="Number of cycles in the learning rate scheduler.",
+    )
+    parser.add_argument(
+        "--lr_power",
+        type=float,
+        default=1.0,
+        help="Power factor of the polynomial scheduler.",
+    )
+    parser.add_argument(
+        "--not_apply_cfg_solver",
+        action="store_true",
+        help="Whether to apply the cfg_solver.",
+    )
+    parser.add_argument("--distill_cfg",
+                        type=float,
+                        default=3.0,
+                        help="Distillation coefficient.")
+    # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
+    parser.add_argument("--scheduler_type",
+                        type=str,
+                        default="pcm",
+                        help="The scheduler type to use.")
+    parser.add_argument(
+        "--linear_quadratic_threshold",
+        type=float,
+        default=0.025,
+        help="Threshold for linear quadratic scheduler.",
+    )
+    parser.add_argument(
+        "--linear_range",
+        type=float,
+        default=0.5,
+        help="Range for linear quadratic scheduler.",
+    )
+    parser.add_argument("--weight_decay",
+                        type=float,
+                        default=0.001,
+                        help="Weight decay to apply.")
+    parser.add_argument("--use_ema",
+                        action="store_true",
+                        help="Whether to use EMA.")
+    parser.add_argument("--multi_phased_distill_schedule",
+                        type=str,
+                        default=None)
+    parser.add_argument("--pred_decay_weight", type=float, default=0.0)
+    parser.add_argument("--pred_decay_type", default="l1")
+    parser.add_argument("--hunyuan_teacher_disable_cfg", action="store_true")
+    parser.add_argument(
+        "--master_weight_type",
+        type=str,
+        default="fp32",
+        help="Weight type to use - fp32 or bf16.",
+    )
+    args = parser.parse_args()
+    main(args)

fastvideo/models/hunyuan/inference.py

@@ -452,7 +452,7 @@ class HunyuanVideoSampler(Inference):
         # ========================================================================
         # Pipeline inference
         # ========================================================================
-        start_time = time.time()
+        start_time = time.perf_counter()
         samples = self.pipeline(
             prompt=prompt,
             height=target_height,
@@ -476,7 +476,7 @@ class HunyuanVideoSampler(Inference):
         out_dict["samples"] = samples
         out_dict["prompts"] = prompt
 
-        gen_time = time.time() - start_time
+        gen_time = time.perf_counter() - start_time
         logger.info(f"Success, time: {gen_time}")
 
         return out_dict

fastvideo/models/mochi_hf/mochi_latents_utils.py

@@ -31,7 +31,7 @@ mochi_latents_std = torch.tensor([
 mochi_scaling_factor = 1.0
 
 
-def normalize_dit_input(model_type, latents):
+def normalize_dit_input(model_type, latents, args=None):
     if model_type == "mochi":
         latents_mean = mochi_latents_mean.to(latents.device, latents.dtype)
         latents_std = mochi_latents_std.to(latents.device, latents.dtype)
@@ -41,5 +41,16 @@ def normalize_dit_input(model_type, latents):
         return latents * 0.476986
     elif model_type == "hunyuan":
         return latents * 0.476986
+    elif model_type == "wan":
+        from fastvideo.v1.configs.models.vaes.wanvae import WanVAEConfig
+        vae_config = WanVAEConfig()
+        latents_mean = torch.tensor(vae_config.arch_config.latents_mean)
+        latents_std = 1.0 / torch.tensor(vae_config.arch_config.latents_std)
+
+
+        latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(device=latents.device)
+        latents_std = latents_std.view(1, -1, 1, 1, 1).to(device=latents.device)
+        latents = ((latents.float() - latents_mean) * latents_std).to(latents)
+        return latents
     else:
         raise NotImplementedError(f"model_type {model_type} not supported")

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

fastvideo/utils/checkpoint.py

@@ -11,6 +11,7 @@ from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_
 from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import StateDictType
+import dataclasses
 
 from fastvideo.utils.logging_ import main_print
 
@@ -44,13 +45,50 @@ def save_checkpoint_optimizer(model, optimizer, rank, output_dir, step, discrimi
         optimizer_path = os.path.join(save_dir, "optimizer.pt")
         torch.save(optim_state, optimizer_path)
     else:
-        weight_path = os.path.join(save_dir, "discriminator_pytorch_model.safetensors")
+        weight_path = os.path.join(save_dstate_dictir, "discriminator_pytorch_model.safetensors")
         save_file(cpu_state, weight_path)
         optimizer_path = os.path.join(save_dir, "discriminator_optimizer.pt")
         torch.save(optim_state, optimizer_path)
     main_print(f"--> checkpoint saved at step {step}")
 
 
+def save_checkpoint_v1(transformer, rank, output_dir, step):
+    # from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+# from torch.distributed.fsdp import StateDictType, FullStateDictConfig
+
+    # Configure FSDP to save full state dict
+    FSDP.set_state_dict_type(
+        transformer,
+        state_dict_type=StateDictType.FULL_STATE_DICT,
+        state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
+    )
+
+    # Now get the state dict
+    cpu_state = transformer.state_dict()
+
+    # Save it (only on rank 0 since we used rank0_only=True)
+    # if torch.distributed.get_rank() == 0:
+    #     torch.save(state_dict, "model_checkpoint.pt")
+    if rank <= 0:
+        save_dir = os.path.join(output_dir, f"checkpoint-{step}")
+        os.makedirs(save_dir, exist_ok=True)
+        # save using safetensors
+        # weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors")
+        weight_path = os.path.join(save_dir, "diffusion_pytorch_model.pt")
+        print(weight_path)
+        # save_file(cpu_state, weight_path)
+        torch.save(cpu_state, weight_path)
+        config_dict = transformer.hf_config
+        if "dtype" in config_dict:
+            del config_dict["dtype"]  # TODO
+        config_path = os.path.join(save_dir, "config.json")
+        # save dict as json
+        with open(config_path, "w") as f:
+            json.dump(config_dict, f, indent=4)
+    main_print(f"--> checkpoint saved at step {step}")
+
+
+
 def save_checkpoint(transformer, rank, output_dir, step):
     main_print(f"--> saving checkpoint at step {step}")
     with FSDP.state_dict_type(

fastvideo/v1/attention/backends/abstract.py

@@ -56,20 +56,6 @@ 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: Optional[Set[str]] = None
                         ) -> Dict[str, Any]:
@@ -86,55 +72,6 @@ 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."""

fastvideo/v1/attention/layer.py

@@ -84,9 +84,6 @@ class DistributedAttention(nn.Module):
         # Check input shapes
         assert q.dim() == 4 and k.dim() == 4 and v.dim(
         ) == 4, "Expected 4D tensors"
-        # assert bs = 1
-        assert q.shape[
-            0] == 1, "Batch size must be 1, and there should be no padding tokens"
         batch_size, seq_len, num_heads, head_dim = q.shape
         local_rank = get_sequence_model_parallel_rank()
         world_size = get_sequence_model_parallel_world_size()

fastvideo/v1/configs/__init__.py

@@ -1,9 +0,0 @@
-from fastvideo.v1.configs.base import BaseConfig, SlidingTileAttnConfig
-from fastvideo.v1.configs.hunyuan import FastHunyuanConfig, HunyuanConfig
-from fastvideo.v1.configs.registry import get_pipeline_config_cls_for_name
-from fastvideo.v1.configs.wan import WanI2V480PConfig, WanT2V480PConfig
-
-__all__ = [
-    "HunyuanConfig", "FastHunyuanConfig", "BaseConfig", "SlidingTileAttnConfig",
-    "WanT2V480PConfig", "WanI2V480PConfig", "get_pipeline_config_cls_for_name"
-]

fastvideo/v1/configs/base.py

@@ -1,67 +0,0 @@
-from dataclasses import dataclass
-from typing import Optional
-
-
-@dataclass
-class BaseConfig:
-    """Base configuration for all pipeline architectures."""
-
-    # Video parameters
-    height: int = 720
-    width: int = 1280
-    num_frames: int = 125
-    fps: int = 24
-
-    # Video generation parameters
-    num_inference_steps: int = 50
-    guidance_scale: float = 1.0
-    seed: int = 1024
-    guidance_rescale: float = 0.0
-    embedded_cfg_scale: float = 6.0
-    flow_shift: Optional[float] = None
-    use_cpu_offload: bool = False
-    disable_autocast: bool = False
-
-    # Model configuration
-    precision: str = "bf16"
-
-    # VAE configuration
-    vae_precision: str = "fp16"
-    vae_tiling: bool = True
-    vae_sp: bool = True
-    vae_scale_factor: Optional[int] = None
-
-    # DiT configuration
-    num_channels_latents: Optional[int] = None
-
-    # Image encoder configuration
-    image_encoder_precision: str = "fp32"
-
-    # Text encoder configuration
-    text_encoder_precision: str = "fp16"
-    text_len: int = -1
-    hidden_state_skip_layer: int = 0
-
-    # STA (Spatial-Temporal Attention) parameters
-    mask_strategy_file_path: Optional[str] = None
-    enable_torch_compile: bool = False
-
-    neg_prompt: Optional[str] = None
-
-
-@dataclass
-class SlidingTileAttnConfig(BaseConfig):
-    """Configuration for sliding tile attention."""
-
-    # Override any BaseConfig defaults as needed
-    # Add sliding tile specific parameters
-    window_size: int = 16
-    stride: int = 8
-
-    # You can provide custom defaults for inherited fields
-    height: int = 576
-    width: int = 1024
-
-    # Additional configuration specific to sliding tile attention
-    pad_to_square: bool = False
-    use_overlap_optimization: bool = True

fastvideo/v1/configs/fasthunyuan_t2v.json

@@ -0,0 +1,48 @@
+{
+  "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/hunyuan.py

@@ -1,40 +0,0 @@
-from dataclasses import dataclass
-
-from fastvideo.v1.configs.base import BaseConfig
-
-
-@dataclass
-class HunyuanConfig(BaseConfig):
-    """Base configuration for HunYuan pipeline architecture."""
-
-    # HunyuanConfig-specific parameters with defaults
-    # Denoising stage
-    embedded_cfg_scale: int = 6
-    flow_shift: int = 7
-    num_inference_steps: int = 50
-
-    # Text encoding stage
-    hidden_state_skip_layer: int = 2
-    text_len: int = 256
-
-    # Precision for each component
-    precision: str = "bf16"
-    vae_precision: str = "fp16"
-    text_encoder_precision: str = "fp16"
-
-    # HunyuanConfig-specific added parameters
-    # Secondary text encoder
-    text_encoder_precision_2: str = "fp16"
-    text_len_2: int = 77
-
-
-@dataclass
-class FastHunyuanConfig(HunyuanConfig):
-    """Configuration specifically optimized for FastHunyuan weights."""
-
-    # Override HunyuanConfig defaults
-    num_inference_steps: int = 6
-    flow_shift: int = 17
-
-    # No need to re-specify guidance_scale or embedded_cfg_scale as they
-    # already have the desired values from HunyuanConfig

fastvideo/v1/configs/models/__init__.py

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

fastvideo/v1/configs/models/base.py

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

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

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

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

@@ -0,0 +1,56 @@
+from dataclasses import dataclass, field
+from typing import Any, Optional, Tuple
+
+from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.v1.layers.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,
+                                         ...] = (_Backend.SLIDING_TILE_ATTN,
+                                                 _Backend.SAGE_ATTN,
+                                                 _Backend.FLASH_ATTN,
+                                                 _Backend.TORCH_SDPA)
+
+    hidden_size: int = 0
+    num_attention_heads: int = 0
+    num_channels_latents: int = 0
+
+    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)
+
+    # FastVideoDiT-specific parameters
+    prefix: str = ""
+    quant_config: Optional[QuantizationConfig] = 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

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

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

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

@@ -0,0 +1,65 @@
+from dataclasses import dataclass, field
+from typing import List, Optional, Tuple, Union
+
+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: Optional[int] = 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: Optional[Union[int, List[int], Tuple[int, ...]]] = field(
+        default_factory=lambda: [6144, 1024])
+    attention_type: Optional[str] = "torch"
+    use_additional_conditions: Optional[bool] = 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

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

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

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

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

@@ -0,0 +1,75 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.v1.layers.quantization import QuantizationConfig
+from fastvideo.v1.platforms import _Backend
+
+
+@dataclass
+class EncoderArchConfig(ArchConfig):
+    architectures: List[str] = field(default_factory=lambda: [])
+    _supported_attention_backends: Tuple[_Backend, ...] = (_Backend.FLASH_ATTN,
+                                                           _Backend.TORCH_SDPA)
+    output_hidden_states: bool = False
+    use_return_dict: bool = True
+
+
+@dataclass
+class TextEncoderArchConfig(EncoderArchConfig):
+    vocab_size: int = 0
+    hidden_size: int = 0
+    num_hidden_layers: int = 0
+    num_attention_heads: int = 0
+    pad_token_id: int = 0
+    eos_token_id: int = 0
+    text_len: int = 0
+    hidden_state_skip_layer: int = 0
+    decoder_start_token_id: int = 0
+    output_past: bool = True
+    scalable_attention: bool = True
+    tie_word_embeddings: bool = False
+
+    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: Optional[torch.FloatTensor] = None
+    pooler_output: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attention_mask: Optional[torch.Tensor] = None
+
+
+@dataclass
+class EncoderConfig(ModelConfig):
+    arch_config: ArchConfig = field(default_factory=EncoderArchConfig)
+
+    prefix: str = ""
+    quant_config: Optional[QuantizationConfig] = None
+    lora_config: Optional[Any] = None
+
+
+@dataclass
+class TextEncoderConfig(EncoderConfig):
+    arch_config: ArchConfig = field(default_factory=TextEncoderArchConfig)
+
+
+@dataclass
+class ImageEncoderConfig(EncoderConfig):
+    arch_config: ArchConfig = field(default_factory=ImageEncoderArchConfig)

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,130 @@
+from dataclasses import dataclass, field
+from typing import Any, Union
+
+import torch
+
+from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
+from fastvideo.v1.utils import StoreBoolean
+
+
+@dataclass
+class VAEArchConfig(ArchConfig):
+    scaling_factor: Union[float, torch.tensor] = 0
+
+    temporal_compression_ratio: int = 4
+    spatial_compression_ratio: int = 8
+
+
+@dataclass
+class VAEConfig(ModelConfig):
+    arch_config: VAEArchConfig = field(default_factory=VAEArchConfig)
+
+    # FastVideoVAE-specific parameters
+    load_encoder: bool = True
+    load_decoder: bool = True
+
+    tile_sample_min_height: int = 256
+    tile_sample_min_width: int = 256
+    tile_sample_min_num_frames: int = 16
+    tile_sample_stride_height: int = 192
+    tile_sample_stride_width: int = 192
+    tile_sample_stride_num_frames: int = 12
+    blend_num_frames: int = 0
+
+    use_tiling: bool = True
+    use_temporal_tiling: bool = True
+    use_parallel_tiling: bool = True
+    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

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

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

@@ -0,0 +1,28 @@
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class StepVideoVAEArchConfig(VAEArchConfig):
+    in_channels: int = 3
+    out_channels: int = 3
+    z_channels: int = 64
+    num_res_blocks: int = 2
+    version: int = 2
+    frame_len: int = 17
+    world_size: int = 1
+
+    spatial_compression_ratio: int = 16
+    temporal_compression_ratio: int = 8
+
+    scaling_factor: float = 1.0
+
+
+@dataclass
+class StepVideoVAEConfig(VAEConfig):
+    arch_config: VAEArchConfig = field(default_factory=StepVideoVAEArchConfig)
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False
+    use_temporal_scaling_frames: bool = False

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

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

fastvideo/v1/configs/pipelines/__init__.py

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

fastvideo/v1/configs/pipelines/base.py

@@ -0,0 +1,151 @@
+import json
+from dataclasses import asdict, dataclass, field, fields
+from typing import Any, Callable, Dict, Optional, Tuple, cast
+
+import torch
+
+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: Optional[float] = None
+    use_cpu_offload: bool = False
+    disable_autocast: bool = False
+
+    # Model configuration
+    precision: str = "bf16"
+
+    # VAE configuration
+    vae_precision: str = "fp16"
+    vae_tiling: bool = True
+    vae_sp: bool = True
+    vae_config: VAEConfig = field(default_factory=VAEConfig)
+
+    # DiT configuration
+    dit_config: DiTConfig = field(default_factory=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, ))
+
+    # STA (Spatial-Temporal Attention) parameters
+    mask_strategy_file_path: Optional[str] = None
+
+    # Compilation
+    enable_torch_compile: bool = False
+
+    @classmethod
+    def from_pretrained(cls, model_path: str) -> "PipelineConfig":
+        from fastvideo.v1.configs.pipelines.registry import (
+            get_pipeline_config_cls_for_name)
+        pipeline_config_cls = get_pipeline_config_cls_for_name(model_path)
+        if pipeline_config_cls is not None:
+            pipeline_config = pipeline_config_cls()
+        else:
+            logger.warning(
+                "Couldn't find an optimal sampling param for %s. Using the default sampling param.",
+                model_path)
+            pipeline_config = cls()
+
+        return cast(PipelineConfig, 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)
+
+    def load_from_json(self, file_path: str):
+        with open(file_path) as f:
+            input_pipeline_dict = json.load(f)
+        self.update_pipeline_config(input_pipeline_dict)
+
+    def update_pipeline_config(self, source_pipeline_dict: Dict[str,
+                                                                Any]) -> None:
+        for f in fields(self):
+            key = f.name
+            if key in source_pipeline_dict:
+                current_value = getattr(self, key)
+                new_value = source_pipeline_dict[key]
+
+                # If it's a nested ModelConfig, update it recursively
+                if isinstance(current_value, ModelConfig):
+                    current_value.update_model_config(new_value)
+                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):
+                        target_config.update_model_config(source_config)
+                else:
+                    setattr(self, key, new_value)
+
+        if hasattr(self, "__post_init__"):
+            self.__post_init__()
+
+
+@dataclass
+class SlidingTileAttnConfig(PipelineConfig):
+    """Configuration for sliding tile attention."""
+
+    # Override any BaseConfig defaults as needed
+    # Add sliding tile specific parameters
+    window_size: int = 16
+    stride: int = 8
+
+    # You can provide custom defaults for inherited fields
+    height: int = 576
+    width: int = 1024
+
+    # Additional configuration specific to sliding tile attention
+    pad_to_square: bool = False
+    use_overlap_optimization: bool = True

fastvideo/v1/configs/pipelines/hunyuan.py

@@ -0,0 +1,103 @@
+from dataclasses import dataclass, field
+from typing import Callable, Tuple, TypedDict
+
+import torch
+
+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.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):
+    """Base configuration for HunYuan pipeline architecture."""
+
+    # HunyuanConfig-specific parameters with defaults
+    # DiT
+    dit_config: DiTConfig = field(default_factory=HunyuanVideoConfig)
+    # VAE
+    vae_config: VAEConfig = field(default_factory=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))
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "fp16"
+    text_encoder_precisions: Tuple[str, ...] = field(
+        default_factory=lambda: ("fp16", "fp16"))
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = False
+        self.vae_config.load_decoder = True
+
+
+@dataclass
+class FastHunyuanConfig(HunyuanConfig):
+    """Configuration specifically optimized for FastHunyuan weights."""
+
+    # Override HunyuanConfig defaults
+    flow_shift: int = 17
+
+    # No need to re-specify guidance_scale or embedded_cfg_scale as they
+    # already have the desired values from HunyuanConfig

fastvideo/v1/configs/pipelines/registry.py

@@ -3,9 +3,14 @@
 import os
 from typing import Callable, Dict, Optional, Type
 
-from fastvideo.v1.configs.base import BaseConfig
-from fastvideo.v1.configs.hunyuan import FastHunyuanConfig, HunyuanConfig
-from fastvideo.v1.configs.wan import WanI2V480PConfig, WanT2V480PConfig
+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)
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.utils import (maybe_download_model_index,
                                 verify_model_config_and_directory)
@@ -13,11 +18,14 @@ 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[BaseConfig]] = {
-    "FastVideo/FastHunyuan-Diffusers": FastHunyuanConfig,
+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-480P-Diffusers": WanI2V480PConfig,
+    "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V720PConfig,
+    "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V720PConfig,
+    "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VConfig,
     # Add other specific weight variants
 }
 
@@ -26,22 +34,24 @@ 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[BaseConfig]] = {
+PIPELINE_FALLBACK_CONFIG: Dict[str, Type[PipelineConfig]] = {
     "hunyuan":
     HunyuanConfig,  # Base Hunyuan config as fallback for any Hunyuan variant
     "wanpipeline":
     WanT2V480PConfig,  # Base Wan config as fallback for any Wan variant
     "wanimagetovideo": WanI2V480PConfig,
+    "stepvideo": StepVideoT2VConfig
     # Other fallbacks by architecture
 }
 
 
 def get_pipeline_config_cls_for_name(
-        pipeline_name_or_path: str) -> Optional[type[BaseConfig]]:
+        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):
@@ -65,7 +75,6 @@ def get_pipeline_config_cls_for_name(
 
     # If no match, try to use the fallback config
     fallback_config = None
-    print(pipeline_name)
     # Try to determine pipeline architecture for fallback
     for pipeline_type, detector in PIPELINE_DETECTOR.items():
         if detector(pipeline_name.lower()):

fastvideo/v1/configs/pipelines/stepvideo.py

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

fastvideo/v1/configs/pipelines/wan.py

@@ -0,0 +1,99 @@
+from dataclasses import dataclass, field
+from typing import Callable, Tuple
+
+import torch
+
+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.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)]
+    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)
+    # VAE
+    vae_config: VAEConfig = field(default_factory=WanVAEConfig)
+    vae_tiling: bool = False
+    vae_sp: bool = False
+
+    # Video parameters
+    use_cpu_offload: bool = True
+
+    # Denoising stage
+    flow_shift: int = 3
+
+    # Text encoding stage
+    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (T5Config(), ))
+    postprocess_text_funcs: Tuple[Callable[[BaseEncoderOutput], torch.tensor],
+                                  ...] = field(default_factory=lambda:
+                                               (t5_postprocess_text, ))
+
+    # Precision for each component
+    precision: str = "bf16"
+    vae_precision: str = "fp16"
+    text_encoder_precisions: Tuple[str, ...] = field(
+        default_factory=lambda: ("fp32", ))
+
+    # WanConfig-specific added parameters
+
+    def __post_init__(self):
+        self.vae_config.load_encoder = False
+        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."""
+
+    # WanConfig-specific parameters with defaults
+
+    # Precision for each component
+    image_encoder_config: EncoderConfig = field(
+        default_factory=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/sample/__init__.py

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

fastvideo/v1/configs/sample/base.py

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

@@ -0,0 +1,29 @@
+from dataclasses import dataclass, field
+
+from fastvideo.v1.configs.sample.base import SamplingParam
+from fastvideo.v1.configs.sample.teacache import TeaCacheParams
+
+
+@dataclass
+class HunyuanSamplingParam(SamplingParam):
+    num_inference_steps: int = 50
+
+    num_frames: int = 125
+    height: int = 720
+    width: int = 1280
+    fps: int = 24
+
+    guidance_scale: float = 1.0
+
+    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):
+    num_inference_steps: int = 6

fastvideo/v1/configs/sample/registry.py

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

fastvideo/v1/configs/sample/stepvideo.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass
+
+from fastvideo.v1.configs.sample.base import SamplingParam
+
+
+@dataclass
+class StepVideoT2VSamplingParam(SamplingParam):
+    # Video parameters
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 81
+
+    # Denoising stage
+    guidance_scale: float = 9.0
+    num_inference_steps: int = 50
+
+    # neg magic and pos magic
+    # pos_magic: str = "超高清、HDR 视频、环境光、杜比全景声、画面稳定、流畅动作、逼真的细节、专业级构图、超现实主义、自然、生动、超细节、清晰。"
+    # neg_magic: str = "画面暗、低分辨率、不良手、文本、缺少手指、多余的手指、裁剪、低质量、颗粒状、签名、水印、用户名、模糊。"

fastvideo/v1/configs/sample/teacache.py

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

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

@@ -1,45 +0,0 @@
-from dataclasses import dataclass
-
-from fastvideo.v1.configs.base import BaseConfig
-
-
-@dataclass
-class WanT2V480PConfig(BaseConfig):
-    """Base configuration for Wan T2V 1.3B pipeline architecture."""
-
-    # WanConfig-specific parameters with defaults
-    # Video parameters
-    height: int = 480
-    width: int = 832
-    num_frames: int = 81
-    fps: int = 16
-    use_cpu_offload: bool = True
-
-    # Denoising stage
-    guidance_scale: float = 3.0
-    neg_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"
-    flow_shift: int = 3
-    num_inference_steps: int = 50
-
-    # Text encoding stage
-    text_len: int = 512
-
-    # Precision for each component
-    precision: str = "bf16"
-    vae_precision: str = "fp16"
-    text_encoder_precision: str = "fp32"
-
-    # WanConfig-specific added parameters
-
-
-@dataclass
-class WanI2V480PConfig(WanT2V480PConfig):
-    """Base configuration for Wan I2V 14B 480P pipeline architecture."""
-
-    # WanConfig-specific parameters with defaults
-    # Denoising stage
-    guidance_scale: float = 5.0
-    num_inference_steps: int = 40
-
-    # Precision for each component
-    image_encoder_precision: str = "fp32"

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

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

fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json

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

fastvideo/v1/dataset/__init__.py

@@ -2,11 +2,12 @@ from torchvision import transforms
 from torchvision.transforms import Lambda
 from transformers import AutoTokenizer
 
-from fastvideo.dataset.t2v_datasets import T2V_dataset
-from fastvideo.dataset.transform import CenterCropResizeVideo, Normalize255, TemporalRandomCrop
+from fastvideo.v1.dataset.t2v_datasets import T2V_dataset
+from fastvideo.v1.dataset.transform import (CenterCropResizeVideo, Normalize255,
+                                            TemporalRandomCrop)
 
 
-def getdataset(args):
+def getdataset(args, start_idx=0):
     temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
     norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
     resize_topcrop = [
@@ -25,15 +26,15 @@ def getdataset(args):
         norm_fun,
     ])
     # tokenizer = AutoTokenizer.from_pretrained("/storage/ongoing/new/Open-Sora-Plan/cache_dir/mt5-xxl", cache_dir=args.cache_dir)
-    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name, cache_dir=args.cache_dir)
+    tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_name,
+                                              cache_dir=args.cache_dir)
     if args.dataset == "t2v":
-        return T2V_dataset(
-            args,
-            transform=transform,
-            temporal_sample=temporal_sample,
-            tokenizer=tokenizer,
-            transform_topcrop=transform_topcrop,
-        )
+        return T2V_dataset(args,
+                           transform=transform,
+                           temporal_sample=temporal_sample,
+                           tokenizer=tokenizer,
+                           transform_topcrop=transform_topcrop,
+                           start_idx=start_idx)
 
     raise NotImplementedError(args.dataset)
 
@@ -44,7 +45,7 @@ if __name__ == "__main__":
     from accelerate import Accelerator
     from tqdm import tqdm
 
-    from fastvideo.dataset.t2v_datasets import dataset_prog
+    from fastvideo.v1.dataset.t2v_datasets import dataset_prog
 
     args = type(
         "args",
@@ -63,7 +64,8 @@ if __name__ == "__main__":
             "interpolation_scale_h": 1,
             "interpolation_scale_w": 1,
             "cache_dir": "../cache_dir",
-            "image_data": "/storage/ongoing/new/Open-Sora-Plan-bak/7.14bak/scripts/train_data/image_data.txt",
+            "image_data":
+            "/storage/ongoing/new/Open-Sora-Plan-bak/7.14bak/scripts/train_data/image_data.txt",
             "video_data": "1",
             "train_fps": 24,
             "drop_short_ratio": 1.0,
@@ -80,7 +82,10 @@ if __name__ == "__main__":
     zero = 0
     for idx in tqdm(range(num)):
         image_data = dataset_prog.img_cap_list[idx]
-        caps = [i["cap"] if isinstance(i["cap"], list) else [i["cap"]] for i in image_data]
+        caps = [
+            i["cap"] if isinstance(i["cap"], list) else [i["cap"]]
+            for i in image_data
+        ]
         try:
             caps = [[random.choice(i)] for i in caps]
         except Exception as e:

fastvideo/v1/dataset/dataloader/schema.py

@@ -0,0 +1,44 @@
+# schema.py
+"""
+Unified data schema and format for saving and loading image/video data after
+preprocessing.
+
+It uses apache arrow in-memory format that can be consumed by modern data
+frameworks that can handle parquet or lance file.
+"""
+
+import pyarrow as pa
+
+pyarrow_schema = pa.schema([
+    pa.field("id", pa.string()),
+    # --- Image/Video VAE latents ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("vae_latent_bytes", pa.binary()),
+    # e.g., [C, T, H, W] or [C, H, W]
+    pa.field("vae_latent_shape", pa.list_(pa.int64())),
+    # e.g., 'float32'
+    pa.field("vae_latent_dtype", pa.string()),
+    # --- Text encoder output tensor ---
+    # Tensors are stored as raw bytes with shape and dtype info for loading
+    pa.field("text_embedding_bytes", pa.binary()),
+    # e.g., [SeqLen, Dim]
+    pa.field("text_embedding_shape", pa.list_(pa.int64())),
+    # e.g., 'bfloat16' or 'float32'
+    pa.field("text_embedding_dtype", pa.string()),
+    pa.field("text_attention_mask_bytes", pa.binary()),
+    # e.g., [SeqLen]
+    pa.field("text_attention_mask_shape", pa.list_(pa.int64())),
+    # e.g., 'bool' or 'int8'
+    pa.field("text_attention_mask_dtype", pa.string()),
+    # --- Metadata ---
+    pa.field("file_name", pa.string()),
+    pa.field("caption", pa.string()),
+    pa.field("media_type", pa.string()),  # 'image' or 'video'
+    pa.field("width", pa.int64()),
+    pa.field("height", pa.int64()),
+    # -- Video-specific (can be null/default for images) ---
+    # Number of frames processed (e.g., 1 for image, N for video)
+    pa.field("num_frames", pa.int64()),
+    pa.field("duration_sec", pa.float64()),
+    pa.field("fps", pa.float64()),
+])

fastvideo/v1/dataset/latent_datasets.py

@@ -20,9 +20,11 @@ class LatentDataset(Dataset):
         self.datase_dir_path = os.path.dirname(json_path)
         self.video_dir = os.path.join(self.datase_dir_path, "video")
         self.latent_dir = os.path.join(self.datase_dir_path, "latent")
-        self.prompt_embed_dir = os.path.join(self.datase_dir_path, "prompt_embed")
-        self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path, "prompt_attention_mask")
-        with open(self.json_path, "r") as f:
+        self.prompt_embed_dir = os.path.join(self.datase_dir_path,
+                                             "prompt_embed")
+        self.prompt_attention_mask_dir = os.path.join(self.datase_dir_path,
+                                                      "prompt_attention_mask")
+        with open(self.json_path) as f:
             self.data_anno = json.load(f)
         # json.load(f) already keeps the order
         # self.data_anno = sorted(self.data_anno, key=lambda x: x['latent_path'])
@@ -31,12 +33,16 @@ class LatentDataset(Dataset):
         self.uncond_prompt_embed = torch.zeros(256, 4096).to(torch.float32)
         # 256 zeros
         self.uncond_prompt_mask = torch.zeros(256).bool()
-        self.lengths = [data_item["length"] if "length" in data_item else 1 for data_item in self.data_anno]
+        self.lengths = [
+            data_item["length"] if "length" in data_item else 1
+            for data_item in self.data_anno
+        ]
 
     def __getitem__(self, idx):
         latent_file = self.data_anno[idx]["latent_path"]
         prompt_embed_file = self.data_anno[idx]["prompt_embed_path"]
-        prompt_attention_mask_file = self.data_anno[idx]["prompt_attention_mask"]
+        prompt_attention_mask_file = self.data_anno[idx][
+            "prompt_attention_mask"]
         # load
         latent = torch.load(
             os.path.join(self.latent_dir, latent_file),
@@ -54,7 +60,8 @@ class LatentDataset(Dataset):
                 weights_only=True,
             )
             prompt_attention_mask = torch.load(
-                os.path.join(self.prompt_attention_mask_dir, prompt_attention_mask_file),
+                os.path.join(self.prompt_attention_mask_dir,
+                             prompt_attention_mask_file),
                 map_location="cpu",
                 weights_only=True,
             )
@@ -104,8 +111,12 @@ def latent_collate_function(batch):
 
 
 if __name__ == "__main__":
-    dataset = LatentDataset("data/Mochi-Synthetic-Data/merge.txt", num_latent_t=28)
-    dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=False, collate_fn=latent_collate_function)
+    dataset = LatentDataset("data/Mochi-Synthetic-Data/merge.txt",
+                            num_latent_t=28)
+    dataloader = torch.utils.data.DataLoader(dataset,
+                                             batch_size=2,
+                                             shuffle=False,
+                                             collate_fn=latent_collate_function)
     for latent, prompt_embed, latent_attn_mask, prompt_attention_mask in dataloader:
         print(
             latent.shape,

fastvideo/v1/dataset/parquet_datasets.py

@@ -0,0 +1,418 @@
+import argparse
+import os
+import random
+import time
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+import tqdm
+from torch import distributed as dist
+from torch.utils.data import IterableDataset, get_worker_info
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+# Path to your dataset
+dataset_path = "/mnt/sharefs/users/hao.zhang/Vchitect-2M/Vchitect-2M-laten-93x512x512/train/"
+
+
+class ParquetVideoTextDataset(IterableDataset):
+    """Efficient loader for video-text data from a directory of Parquet files."""
+
+    def __init__(self,
+                 path: str,
+                 batch_size: int = 1024,
+                 rank: int = 0,
+                 world_size: int = 1,
+                 cfg_rate: float = 0.0,
+                 num_latent_t: int = 2):
+        super().__init__()
+        self.path = str(path)
+        self.batch_size = batch_size
+        self.rank = rank
+        self.world_size = world_size
+        self.cfg_rate = cfg_rate
+        self.num_latent_t = num_latent_t
+
+        # Find all parquet files recursively
+        print(f"Scanning for parquet files in {self.path}")
+        self.parquet_files = []
+        for root, _, files in os.walk(self.path):
+            for file in files:
+                if file.endswith('.parquet'):
+                    self.parquet_files.append(os.path.join(root, file))
+        # Sort files for consistent ordering
+        self.parquet_files.sort()
+
+        # Distribute files among workers
+        # drop last unenven files
+        print(f"Total files: {len(self.parquet_files)}")
+        total_files = len(self.parquet_files)
+        base_count = total_files // world_size
+        extra_files = total_files % world_size
+
+        if rank < extra_files:
+            start_idx = rank * (base_count + 1)
+            end_idx = start_idx + base_count + 1
+        else:
+            start_idx = rank * base_count + extra_files
+            end_idx = start_idx + base_count
+
+        self.parquet_files = self.parquet_files[start_idx:end_idx]
+
+        print(f"Files assigned to rank {rank}: {len(self.parquet_files)}")
+        if len(self.parquet_files) > 0:
+            print(f"First file: {self.parquet_files[0]}")
+            print(f"Last file: {self.parquet_files[-1]}")
+
+        # Initialize current file index
+        self.current_file_idx = 0
+        self.current_reader = None
+        self.current_batches = None
+        self.total_samples = 0
+
+    def _open_next_file(self):
+        """Open the next parquet file for reading."""
+        num_workers = get_worker_info().num_workers
+        worker_id = get_worker_info().id
+        total_files = len(self.parquet_files)
+        base_count = total_files // num_workers
+        extra_files = total_files % num_workers
+
+        if worker_id < extra_files:
+            start_idx = worker_id * (base_count + 1)
+            end_idx = start_idx + base_count + 1
+        else:
+            start_idx = worker_id * base_count + extra_files
+            end_idx = start_idx + base_count
+
+        worker_parquet_files = self.parquet_files[start_idx:end_idx]
+        if self.current_file_idx >= len(worker_parquet_files):
+            print(
+                f"Rank {self.rank}, Worker {worker_id}: No more files to open (current_idx={self.current_file_idx}, total_files={len(worker_parquet_files)})"
+            )
+            return False
+
+        if self.current_reader is not None:
+            self.current_reader.close()
+
+        file_path = worker_parquet_files[self.current_file_idx]
+        print(
+            f"Rank {self.rank}, Worker {worker_id}: Opening file {self.current_file_idx + 1}/{len(worker_parquet_files)}: {file_path}"
+        )
+
+        try:
+            self.current_reader = pq.ParquetFile(file_path)
+            self.current_batches = self.current_reader.iter_batches(
+                batch_size=self.batch_size)
+            self.current_file_idx += 1
+            return True
+        except Exception as e:
+            print(f"Error opening file {file_path}: {str(e)}")
+            return False
+
+    def __iter__(self):
+        """Iterate over the dataset in a streaming fashion."""
+        print(f"Rank {self.rank}: Starting iteration")
+
+        # First try to open a file
+        if not self._open_next_file():
+            print(f"Rank {self.rank}: Failed to open first file")
+            return
+
+        while True:
+            try:
+                # Get next batch from current file
+                batch = next(self.current_batches)
+                batch_dict = batch.to_pydict()
+                processed = self._process_batch(batch_dict)
+
+                # Update sample count
+                batch_size = len(processed["latents"])
+                self.total_samples += batch_size
+
+                # Print progress
+                if self.total_samples % 1000 == 0:
+                    print(
+                        f"Rank {self.rank}: Processed {self.total_samples} samples"
+                    )
+
+                # Yield each item in the batch
+                for lat, emb, mask, info in zip(processed["latents"],
+                                                processed["embeddings"],
+                                                processed["masks"],
+                                                processed["info"]):
+                    if lat.numel() == 0:  # Split is validation
+                        yield lat, emb, mask, info
+                    else:
+                        yield lat[:, -self.num_latent_t:], emb, mask, info
+
+            except StopIteration:
+                # Current file is exhausted, try next file
+                print(
+                    f"Rank {self.rank}: Current file exhausted, trying next file"
+                )
+                self.current_batches = None
+                if not self._open_next_file():
+                    print(
+                        f"Rank {self.rank}: No more files to process. Total samples: {self.total_samples}"
+                    )
+                    break
+            except Exception as e:
+                print(f"Error processing batch: {str(e)}")
+                self.current_batches = None
+                if not self._open_next_file():
+                    print(
+                        f"Rank {self.rank}: Failed to open next file after error"
+                    )
+                    break
+
+        # Clean up
+        if self.current_reader is not None:
+            self.current_reader.close()
+
+    def _process_batch(self, batch):
+        """Process a PyArrow batch into tensors."""
+        out = {"lat": [], "emb": [], "msk": [], "info": []}
+
+        for i in range(len(batch["vae_latent_bytes"])):
+            vae_latent_bytes = batch["vae_latent_bytes"][i]
+            vae_latent_shape = batch["vae_latent_shape"][i]
+            text_embedding_bytes = batch["text_embedding_bytes"][i]
+            text_embedding_shape = batch["text_embedding_shape"][i]
+            text_attention_mask_bytes = batch["text_attention_mask_bytes"][i]
+            text_attention_mask_shape = batch["text_attention_mask_shape"][i]
+
+            # Process latent
+            if not vae_latent_shape:  # No VAE latent is stored. Split is validation
+                lat = np.array([])
+            else:
+                lat = np.frombuffer(vae_latent_bytes,
+                                    dtype=np.float32).reshape(vae_latent_shape)
+                # Make array writable
+                lat = np.copy(lat)
+
+            if random.random() < self.cfg_rate:
+                emb = np.zeros((512, 4096), dtype=np.float32)
+            else:
+                emb = np.frombuffer(
+                    text_embedding_bytes,
+                    dtype=np.float32).reshape(text_embedding_shape)
+                # Make array writable
+                emb = np.copy(emb)
+            if emb.shape[0] < 512:
+                padded_emb = np.zeros((512, emb.shape[1]), dtype=np.float32)
+                padded_emb[:emb.shape[0], :] = emb
+                emb = padded_emb
+            elif emb.shape[0] > 512:
+                emb = emb[:512, :]
+
+            # Process mask
+            if len(text_attention_mask_bytes) > 0 and len(
+                    text_attention_mask_shape) > 0:
+                msk = np.frombuffer(text_attention_mask_bytes,
+                                    dtype=np.uint8).astype(np.bool_)
+                msk = msk.reshape(1, -1)
+                # Make array writable
+                msk = np.copy(msk)
+                if msk.shape[1] < 512:
+                    padded_msk = np.zeros((1, 512), dtype=np.bool_)
+                    padded_msk[:, :msk.shape[1]] = msk
+                    msk = padded_msk
+                elif msk.shape[1] > 512:
+                    msk = msk[:, :512]
+            else:
+                msk = np.ones((1, 512), dtype=np.bool_)
+            # to string
+            file_name = str(batch["file_name"][i])
+            # Collect metadata
+            info = {
+                "width": batch["width"][i],
+                "height": batch["height"][i],
+                "num_frames": batch["num_frames"][i],
+                "duration_sec": batch["duration_sec"][i],
+                "fps": batch["fps"][i],
+                "file_name": batch["file_name"][i],
+                "caption": batch["caption"][i],
+            }
+
+            out["lat"].append(torch.from_numpy(lat))
+            out["emb"].append(torch.from_numpy(emb))
+            out["msk"].append(torch.from_numpy(msk))
+            out["info"].append(info)
+
+        return {
+            "latents": torch.stack(out["lat"]) if out["lat"] else None,
+            "embeddings": torch.stack(out["emb"]) if out["emb"] else None,
+            "masks": torch.stack(out["msk"]) if out["msk"] else None,
+            "info": out["info"]
+        }
+
+
+def bind_cpu_cores(local_rank, cpu_per_process=16):
+    """根据local_rank绑定固定cpu核。"""
+    start = local_rank * cpu_per_process
+    end = start + cpu_per_process
+    cores = list(range(start, end))
+    print(f"[Rank {local_rank}] Binding to CPU cores: {cores}")
+    os.sched_setaffinity(0, cores)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description='Benchmark Parquet dataset loading speed')
+    parser.add_argument('--path',
+                        type=str,
+                        default=dataset_path,
+                        help='Path to Parquet dataset')
+    parser.add_argument('--batch_size',
+                        type=int,
+                        default=4,
+                        help='Batch size for DataLoader')
+    parser.add_argument('--num_batches',
+                        type=int,
+                        default=100,
+                        help='Number of batches to benchmark')
+    parser.add_argument('--vae_debug', action="store_true")
+    args = parser.parse_args()
+
+    # Initialize distributed training
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    rank = int(os.environ.get("RANK", 0))
+
+    # Initialize CUDA device first
+    if torch.cuda.is_available():
+        torch.cuda.set_device(local_rank)
+        device = torch.device(f"cuda:{local_rank}")
+    else:
+        device = torch.device("cpu")
+
+    # Initialize distributed training
+    if world_size > 1:
+        dist.init_process_group(backend="nccl",
+                                init_method="env://",
+                                world_size=world_size,
+                                rank=rank)
+        print(
+            f"Initialized process: rank={rank}, local_rank={local_rank}, world_size={world_size}, device={device}"
+        )
+
+    # Bind CPU cores after distributed initialization
+    # bind_cpu_cores(local_rank, cpu_per_process=16)
+
+    # Create dataset
+    dataset = ParquetVideoTextDataset(
+        args.path,
+        batch_size=args.batch_size,
+        rank=rank,
+        world_size=world_size,
+    )
+
+    # Create DataLoader with proper settings
+    dataloader = StatefulDataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        num_workers=1,  # Reduce number of workers to avoid memory issues
+        prefetch_factor=2,
+        shuffle=False,
+        pin_memory=True,
+        drop_last=True)
+
+    # Example of how to load dataloader state
+    # if os.path.exists("/workspace/FastVideo/dataloader_state.pt"):
+    #     dataloader_state = torch.load("/workspace/FastVideo/dataloader_state.pt")
+    #     dataloader.load_state_dict(dataloader_state[rank])
+
+    # Warm-up with synchronization
+    if rank == 0:
+        print("Warming up...")
+    for i, (latents, embeddings, masks, infos) in enumerate(dataloader):
+        # Example of how to save dataloader state
+        # if i == 30:
+        #     dist.barrier()
+        #     local_data = {rank: dataloader.state_dict()}
+        #     gathered_data = [None] * world_size
+        #     dist.all_gather_object(gathered_data, local_data)
+        #     if rank == 0:
+        #         global_state_dict = {}
+        #         for d in gathered_data:
+        #             global_state_dict.update(d)
+        #         torch.save(global_state_dict, "dataloader_state.pt")
+        assert torch.sum(masks[0]).item() == torch.count_nonzero(
+            embeddings[0]).item() // 4096
+        if args.vae_debug:
+            from diffusers.utils import export_to_video
+            from diffusers.video_processor import VideoProcessor
+
+            from fastvideo.v1.configs.models.vaes import WanVAEConfig
+            from fastvideo.v1.fastvideo_args import FastVideoArgs
+            from fastvideo.v1.models.loader.component_loader import VAELoader
+            VAE_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers/vae"
+            fastvideo_args = FastVideoArgs(
+                model_path=VAE_PATH,
+                vae_config=WanVAEConfig(load_encoder=False),
+                vae_precision="fp32")
+            fastvideo_args.device = device
+            vae_loader = VAELoader()
+            vae = vae_loader.load(model_path=VAE_PATH,
+                                  architecture="",
+                                  fastvideo_args=fastvideo_args)
+
+            videoprocessor = VideoProcessor(vae_scale_factor=8)
+
+            with torch.inference_mode():
+                video = vae.decode(latents[0].unsqueeze(0).to(device))
+                video = videoprocessor.postprocess_video(video)
+                video_path = os.path.join("/workspace/FastVideo/debug_videos",
+                                          infos["caption"][0][:50] + ".mp4")
+                export_to_video(video[0], video_path, fps=16)
+
+        # Move data to device
+        # latents = latents.to(device)
+        # embeddings = embeddings.to(device)
+
+    if world_size > 1:
+        dist.barrier()
+
+    # Benchmark
+    if rank == 0:
+        print(f"Benchmarking with batch_size={args.batch_size}")
+    start_time = time.time()
+    total_samples = 0
+    for i, (latents, embeddings, masks,
+            infos) in enumerate(tqdm.tqdm(dataloader, total=args.num_batches)):
+        if i >= args.num_batches:
+            break
+
+        # Move data to device
+        latents = latents.to(device)
+        embeddings = embeddings.to(device)
+
+        # Calculate actual batch size
+        batch_size = latents.size(0)
+        total_samples += batch_size
+
+        # Print progress only from rank 0
+        if rank == 0 and (i + 1) % 10 == 0:
+            elapsed = time.time() - start_time
+            samples_per_sec = total_samples / elapsed
+            print(
+                f"Batch {i+1}/{args.num_batches}, Speed: {samples_per_sec:.2f} samples/sec"
+            )
+
+    # Final statistics
+    if world_size > 1:
+        dist.barrier()
+
+    if rank == 0:
+        elapsed = time.time() - start_time
+        samples_per_sec = total_samples / elapsed
+
+        print("\nBenchmark Results:")
+        print(f"Total time: {elapsed:.2f} seconds")
+        print(f"Total samples: {total_samples}")
+        print(f"Average speed: {samples_per_sec:.2f} samples/sec")
+        print(f"Time per batch: {elapsed/args.num_batches*1000:.2f} ms")
+
+    if world_size > 1:
+        dist.destroy_process_group()

fastvideo/v1/dataset/t2v_datasets.py

@@ -46,7 +46,8 @@ class DataSetProg(metaclass=SingletonMeta):
 
         for i in range(self.num_workers):
             self.n_used_elements[i] = 0
-            per_worker = int(math.ceil(len(self.elements) / float(self.num_workers)))
+            per_worker = int(
+                math.ceil(len(self.elements) / float(self.num_workers)))
             start = i * per_worker
             end = min(start + per_worker, len(self.elements))
             self.worker_elements[i] = self.elements[start:end]
@@ -57,7 +58,9 @@ class DataSetProg(metaclass=SingletonMeta):
         else:
             worker_id = work_info.id
 
-        idx = self.worker_elements[worker_id][self.n_used_elements[worker_id] % len(self.worker_elements[worker_id])]
+        idx = self.worker_elements[worker_id][
+            self.n_used_elements[worker_id] %
+            len(self.worker_elements[worker_id])]
         self.n_used_elements[worker_id] += 1
         return idx
 
@@ -65,7 +68,10 @@ class DataSetProg(metaclass=SingletonMeta):
 dataset_prog = DataSetProg()
 
 
-def filter_resolution(h, w, max_h_div_w_ratio=17 / 16, min_h_div_w_ratio=8 / 16):
+def filter_resolution(h,
+                      w,
+                      max_h_div_w_ratio=17 / 16,
+                      min_h_div_w_ratio=8 / 16):
     if h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio:
         return True
     return False
@@ -73,7 +79,14 @@ def filter_resolution(h, w, max_h_div_w_ratio=17 / 16, min_h_div_w_ratio=8 / 16)
 
 class T2V_dataset(Dataset):
 
-    def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop):
+    def __init__(self,
+                 args,
+                 transform,
+                 temporal_sample,
+                 tokenizer,
+                 transform_topcrop,
+                 start_idx=0):
+        self.start_idx = start_idx
         self.data = args.data_merge_path
         self.num_frames = args.num_frames
         self.train_fps = args.train_fps
@@ -102,7 +115,8 @@ class T2V_dataset(Dataset):
         self.lengths = self.sample_num_frames
 
         n_elements = len(cap_list)
-        dataset_prog.set_cap_list(args.dataloader_num_workers, cap_list, n_elements)
+        dataset_prog.set_cap_list(args.dataloader_num_workers, cap_list,
+                                  n_elements)
 
         print(f"video length: {len(dataset_prog.cap_list)}", flush=True)
 
@@ -129,7 +143,8 @@ class T2V_dataset(Dataset):
         video_path = dataset_prog.cap_list[idx]["path"]
         assert os.path.exists(video_path), f"file {video_path} do not exist!"
         frame_indices = dataset_prog.cap_list[idx]["sample_frame_index"]
-        torchvision_video, _, metadata = torchvision.io.read_video(video_path, output_format="TCHW")
+        torchvision_video, _, metadata = torchvision.io.read_video(
+            video_path, output_format="TCHW")
         video = torchvision_video[frame_indices]
         video = self.transform(video)
         video = rearrange(video, "t c h w -> c t h w")
@@ -160,16 +175,17 @@ class T2V_dataset(Dataset):
         )
         input_ids = text_tokens_and_mask["input_ids"]
         cond_mask = text_tokens_and_mask["attention_mask"]
-        return dict(
-            pixel_values=video,
-            text=text,
-            input_ids=input_ids,
-            cond_mask=cond_mask,
-            path=video_path,
-        )
+        return dict(pixel_values=video,
+                    text=text,
+                    input_ids=input_ids,
+                    cond_mask=cond_mask,
+                    path=video_path,
+                    fps=dataset_prog.cap_list[idx]["fps"],
+                    duration=dataset_prog.cap_list[idx]["duration"])
 
     def get_image(self, idx):
-        image_data = dataset_prog.cap_list[idx]  # [{'path': path, 'cap': cap}, ...]
+        image_data = dataset_prog.cap_list[
+            idx]  # [{'path': path, 'cap': cap}, ...]
 
         image = Image.open(image_data["path"]).convert("RGB")  # [h, w, c]
         image = torch.from_numpy(np.array(image))  # [h, w, c]
@@ -178,13 +194,15 @@ class T2V_dataset(Dataset):
         #     h, w = i.shape[-2:]
         #     assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only image with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But found ratio is {round(h / w, 2)} with the shape of {i.shape}'
 
-        image = (self.transform_topcrop(image) if "human_images" in image_data["path"] else self.transform(image)
+        image = (self.transform_topcrop(image) if "human_images"
+                 in image_data["path"] else self.transform(image)
                  )  #  [1 C H W] -> num_img [1 C H W]
         image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
 
         image = image.float() / 127.5 - 1.0
 
-        caps = (image_data["cap"] if isinstance(image_data["cap"], list) else [image_data["cap"]])
+        caps = (image_data["cap"]
+                if isinstance(image_data["cap"], list) else [image_data["cap"]])
         caps = [random.choice(caps)]
         text = caps
         input_ids, cond_mask = [], []
@@ -238,10 +256,12 @@ class T2V_dataset(Dataset):
                     cnt_no_resolution += 1
                     continue
                 else:
-                    if (resolution.get("height", None) is None or resolution.get("width", None) is None):
+                    if (resolution.get("height", None) is None
+                            or resolution.get("width", None) is None):
                         cnt_no_resolution += 1
                         continue
-                    height, width = i["resolution"]["height"], i["resolution"]["width"]
+                    height, width = i["resolution"]["height"], i["resolution"][
+                        "width"]
                     aspect = self.max_height / self.max_width
                     hw_aspect_thr = 1.5
                     is_pick = filter_resolution(
@@ -259,29 +279,34 @@ class T2V_dataset(Dataset):
                 i["num_frames"] = math.ceil(fps * duration)
                 # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration.
                 if i["num_frames"] / fps > self.video_length_tolerance_range * (
-                        self.num_frames / self.train_fps *
-                        self.speed_factor):  # too long video is not suitable for this training stage (self.num_frames)
+                        self.num_frames / self.train_fps * self.speed_factor
+                ):  # too long video is not suitable for this training stage (self.num_frames)
                     cnt_too_long += 1
                     continue
 
                 # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24)
                 frame_interval = fps / self.train_fps
                 start_frame_idx = 0
-                frame_indices = np.arange(start_frame_idx, i["num_frames"], frame_interval).astype(int)
+                frame_indices = np.arange(start_frame_idx, i["num_frames"],
+                                          frame_interval).astype(int)
 
                 # comment out it to enable dynamic frames training
-                if (len(frame_indices) < self.num_frames and random.random() < self.drop_short_ratio):
+                if (len(frame_indices) < self.num_frames
+                        and random.random() < self.drop_short_ratio):
                     cnt_too_short += 1
                     continue
 
                 #  too long video will be temporal-crop randomly
                 if len(frame_indices) > self.num_frames:
-                    begin_index, end_index = self.temporal_sample(len(frame_indices))
+                    begin_index, end_index = self.temporal_sample(
+                        len(frame_indices))
                     frame_indices = frame_indices[begin_index:end_index]
                     # frame_indices = frame_indices[:self.num_frames]  # head crop
                 i["sample_frame_index"] = frame_indices.tolist()
                 new_cap_list.append(i)
-                i["sample_num_frames"] = len(i["sample_frame_index"])  # will use in dataloader(group sampler)
+                i["sample_num_frames"] = len(
+                    i["sample_frame_index"]
+                )  # will use in dataloader(group sampler)
                 sample_num_frames.append(i["sample_num_frames"])
             elif path.endswith(".jpg"):  # image
                 cnt_img += 1
@@ -290,13 +315,15 @@ class T2V_dataset(Dataset):
                 sample_num_frames.append(i["sample_num_frames"])
             else:
                 raise NameError(
-                    f"Unknown file extension {path.split('.')[-1]}, only support .mp4 for video and .jpg for image")
+                    f"Unknown file extension {path.split('.')[-1]}, only support .mp4 for video and .jpg for image"
+                )
         # import ipdb;ipdb.set_trace()
         main_print(
             f"no_cap: {cnt_no_cap}, too_long: {cnt_too_long}, too_short: {cnt_too_short}, "
             f"no_resolution: {cnt_no_resolution}, resolution_mismatch: {cnt_resolution_mismatch}, "
             f"Counter(sample_num_frames): {Counter(sample_num_frames)}, cnt_movie: {cnt_movie}, cnt_img: {cnt_img}, "
-            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}")
+            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}"
+        )
         return new_cap_list, sample_num_frames
 
     def decord_read(self, path, frame_indices):
@@ -308,11 +335,14 @@ class T2V_dataset(Dataset):
 
     def read_jsons(self, data):
         cap_lists = []
-        with open(data, "r") as f:
-            folder_anno = [i.strip().split(",") for i in f.readlines() if len(i.strip()) > 0]
+        with open(data) as f:
+            folder_anno = [
+                i.strip().split(",") for i in f.readlines()
+                if len(i.strip()) > 0
+            ]
         print(folder_anno)
         for folder, anno in folder_anno:
-            with open(anno, "r") as f:
+            with open(anno) as f:
                 sub_list = json.load(f)
             for i in range(len(sub_list)):
                 sub_list[i]["path"] = opj(folder, sub_list[i]["path"])
@@ -320,5 +350,5 @@ class T2V_dataset(Dataset):
         return cap_lists
 
     def get_cap_list(self):
-        cap_lists = self.read_jsons(self.data)
+        cap_lists = self.read_jsons(self.data)[self.start_idx:]
         return cap_lists

fastvideo/v1/dataset/transform.py

@@ -21,15 +21,19 @@ def center_crop_arr(pil_image, image_size):
     https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
     """
     while min(*pil_image.size) >= 2 * image_size:
-        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), resample=Image.BOX)
+        pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size),
+                                     resample=Image.BOX)
 
     scale = image_size / min(*pil_image.size)
-    pil_image = pil_image.resize(tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC)
+    pil_image = pil_image.resize(tuple(
+        round(x * scale) for x in pil_image.size),
+                                 resample=Image.BICUBIC)
 
     arr = np.array(pil_image)
     crop_y = (arr.shape[0] - image_size) // 2
     crop_x = (arr.shape[1] - image_size) // 2
-    return Image.fromarray(arr[crop_y:crop_y + image_size, crop_x:crop_x + image_size])
+    return Image.fromarray(arr[crop_y:crop_y + image_size,
+                               crop_x:crop_x + image_size])
 
 
 def crop(clip, i, j, h, w):
@@ -44,7 +48,9 @@ def crop(clip, i, j, h, w):
 
 def resize(clip, target_size, interpolation_mode):
     if len(target_size) != 2:
-        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
+        raise ValueError(
+            f"target size should be tuple (height, width), instead got {target_size}"
+        )
     return torch.nn.functional.interpolate(
         clip,
         size=target_size,
@@ -56,7 +62,9 @@ def resize(clip, target_size, interpolation_mode):
 
 def resize_scale(clip, target_size, interpolation_mode):
     if len(target_size) != 2:
-        raise ValueError(f"target size should be tuple (height, width), instead got {target_size}")
+        raise ValueError(
+            f"target size should be tuple (height, width), instead got {target_size}"
+        )
     H, W = clip.size(-2), clip.size(-1)
     scale_ = target_size[0] / min(H, W)
     return torch.nn.functional.interpolate(
@@ -166,7 +174,8 @@ def normalize_video(clip):
     """
     _is_tensor_video_clip(clip)
     if not clip.dtype == torch.uint8:
-        raise TypeError("clip tensor should have data type uint8. Got %s" % str(clip.dtype))
+        raise TypeError("clip tensor should have data type uint8. Got %s" %
+                        str(clip.dtype))
     # return clip.float().permute(3, 0, 1, 2) / 255.0
     return clip.float() / 255.0
 
@@ -227,7 +236,9 @@ class RandomCropVideo:
         th, tw = self.size
 
         if h < th or w < tw:
-            raise ValueError(f"Required crop size {(th, tw)} is larger than input image size {(h, w)}")
+            raise ValueError(
+                f"Required crop size {(th, tw)} is larger than input image size {(h, w)}"
+            )
 
         if w == tw and h == th:
             return 0, 0, h, w
@@ -301,7 +312,9 @@ class LongSideResizeVideo:
         else:
             h = int(h * self.size / w)
             w = self.size
-        resize_clip = resize(clip, target_size=(h, w), interpolation_mode=self.interpolation_mode)
+        resize_clip = resize(clip,
+                             target_size=(h, w),
+                             interpolation_mode=self.interpolation_mode)
         return resize_clip
 
     def __repr__(self) -> str:
@@ -321,7 +334,8 @@ class CenterCropResizeVideo:
         interpolation_mode="bilinear",
     ):
         if len(size) != 2:
-            raise ValueError(f"size should be tuple (height, width), instead got {size}")
+            raise ValueError(
+                f"size should be tuple (height, width), instead got {size}")
         self.size = size
         self.top_crop = top_crop
         self.interpolation_mode = interpolation_mode
@@ -335,7 +349,10 @@ class CenterCropResizeVideo:
                 size is (T, C, crop_size, crop_size)
         """
         # clip_center_crop = center_crop_using_short_edge(clip)
-        clip_center_crop = center_crop_th_tw(clip, self.size[0], self.size[1], top_crop=self.top_crop)
+        clip_center_crop = center_crop_th_tw(clip,
+                                             self.size[0],
+                                             self.size[1],
+                                             top_crop=self.top_crop)
         # import ipdb;ipdb.set_trace()
         clip_center_crop_resize = resize(
             clip_center_crop,
@@ -361,7 +378,8 @@ class UCFCenterCropVideo:
     ):
         if isinstance(size, tuple):
             if len(size) != 2:
-                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+                raise ValueError(
+                    f"size should be tuple (height, width), instead got {size}")
             self.size = size
         else:
             self.size = (size, size)
@@ -376,7 +394,9 @@ class UCFCenterCropVideo:
             torch.tensor: scale resized / center cropped video clip.
                 size is (T, C, crop_size, crop_size)
         """
-        clip_resize = resize_scale(clip=clip, target_size=self.size, interpolation_mode=self.interpolation_mode)
+        clip_resize = resize_scale(clip=clip,
+                                   target_size=self.size,
+                                   interpolation_mode=self.interpolation_mode)
         clip_center_crop = center_crop(clip_resize, self.size)
         return clip_center_crop
 
@@ -396,7 +416,8 @@ class KineticsRandomCropResizeVideo:
     ):
         if isinstance(size, tuple):
             if len(size) != 2:
-                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+                raise ValueError(
+                    f"size should be tuple (height, width), instead got {size}")
             self.size = size
         else:
             self.size = (size, size)
@@ -405,7 +426,8 @@ class KineticsRandomCropResizeVideo:
 
     def __call__(self, clip):
         clip_random_crop = random_shift_crop(clip)
-        clip_resize = resize(clip_random_crop, self.size, self.interpolation_mode)
+        clip_resize = resize(clip_random_crop, self.size,
+                             self.interpolation_mode)
         return clip_resize
 
 
@@ -418,7 +440,8 @@ class CenterCropVideo:
     ):
         if isinstance(size, tuple):
             if len(size) != 2:
-                raise ValueError(f"size should be tuple (height, width), instead got {size}")
+                raise ValueError(
+                    f"size should be tuple (height, width), instead got {size}")
             self.size = size
         else:
             self.size = (size, size)
@@ -514,7 +537,7 @@ class RandomHorizontalFlipVideo:
 #  ------------------------------------------------------------
 #  ---------------------  Sampling  ---------------------------
 #  ------------------------------------------------------------
-class TemporalRandomCrop(object):
+class TemporalRandomCrop:
     """Temporally crop the given frame indices at a random location.
 
     Args:
@@ -531,7 +554,7 @@ class TemporalRandomCrop(object):
         return begin_index, end_index
 
 
-class DynamicSampleDuration(object):
+class DynamicSampleDuration:
     """Temporally crop the given frame indices at a random location.
 
     Args:
@@ -545,7 +568,8 @@ class DynamicSampleDuration(object):
     def __call__(self, t, h, w):
         if self.extra_1:
             t = t - 1
-        truncate_t_list = list(range(t + 1))[t // 2:][::self.t_stride]  # need half at least
+        truncate_t_list = list(
+            range(t + 1))[t // 2:][::self.t_stride]  # need half at least
         truncate_t = random.choice(truncate_t_list)
         if self.extra_1:
             truncate_t = truncate_t + 1
@@ -560,14 +584,18 @@ if __name__ == "__main__":
     from torchvision import transforms
     from torchvision.utils import save_image
 
-    vframes, aframes, info = io.read_video(filename="./v_Archery_g01_c03.avi", pts_unit="sec", output_format="TCHW")
+    vframes, aframes, info = io.read_video(filename="./v_Archery_g01_c03.avi",
+                                           pts_unit="sec",
+                                           output_format="TCHW")
 
     trans = transforms.Compose([
         Normalize255(),
         RandomHorizontalFlipVideo(),
         UCFCenterCropVideo(512),
         # NormalizeVideo(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
-        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
+        transforms.Normalize(mean=[0.5, 0.5, 0.5],
+                             std=[0.5, 0.5, 0.5],
+                             inplace=True),
     ])
 
     target_video_len = 32
@@ -582,7 +610,10 @@ if __name__ == "__main__":
     # print(start_frame_ind)
     # print(end_frame_ind)
     assert end_frame_ind - start_frame_ind >= target_video_len
-    frame_indice = np.linspace(start_frame_ind, end_frame_ind - 1, target_video_len, dtype=int)
+    frame_indice = np.linspace(start_frame_ind,
+                               end_frame_ind - 1,
+                               target_video_len,
+                               dtype=int)
     print(frame_indice)
 
     select_vframes = vframes[frame_indice]
@@ -593,11 +624,14 @@ if __name__ == "__main__":
     print(select_vframes_trans.shape)
     print(select_vframes_trans.dtype)
 
-    select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) * 255).to(dtype=torch.uint8)
+    select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) *
+                                255).to(dtype=torch.uint8)
     print(select_vframes_trans_int.dtype)
     print(select_vframes_trans_int.permute(0, 2, 3, 1).shape)
 
-    io.write_video("./test.avi", select_vframes_trans_int.permute(0, 2, 3, 1), fps=8)
+    io.write_video("./test.avi",
+                   select_vframes_trans_int.permute(0, 2, 3, 1),
+                   fps=8)
 
     for i in range(target_video_len):
         save_image(