fix bugs

.github/scripts/runpod_api.py

@@ -1,250 +0,0 @@
-import argparse
-import json
-import os
-import subprocess
-import sys
-import time
-
-import requests
-
-
-def parse_arguments():
-    """Parse command line arguments"""
-    parser = argparse.ArgumentParser(description='Run tests on RunPod GPU')
-    parser.add_argument('--gpu-type', type=str, help='GPU type to use')
-    parser.add_argument('--gpu-count',
-                        type=int,
-                        help='Number of GPUs to use',
-                        default=1)
-    parser.add_argument('--test-command', type=str, help='Test command to run')
-    parser.add_argument('--disk-size',
-                        type=int,
-                        default=20,
-                        help='Container disk size in GB (default: 20)')
-    parser.add_argument('--volume-size',
-                        type=int,
-                        default=20,
-                        help='Persistent volume size in GB (default: 20)')
-    parser.add_argument(
-        '--image',
-        type=str,
-        required=True,
-        help='Docker image to use')
-    return parser.parse_args()
-
-
-args = parse_arguments()
-API_KEY = os.environ['RUNPOD_API_KEY']
-RUN_ID = os.environ['GITHUB_RUN_ID']
-JOB_ID = os.environ['JOB_ID']
-PODS_API = "https://rest.runpod.io/v1/pods"
-HEADERS = {
-    "Content-Type": "application/json",
-    "Authorization": f"Bearer {API_KEY}"
-}
-
-
-def create_pod():
-    """Create a RunPod instance"""
-    # Ensure image name is lowercase (Docker requirement)
-    image_name = args.image.lower()
-    print(f"Using specified image: {image_name}")
-
-    docker_start_cmd = [
-        "bash", 
-        "-c", 
-        "apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-server -y;mkdir -p ~/.ssh;cd $_;chmod 700 ~/.ssh;echo \"$PUBLIC_KEY\" >> authorized_keys;chmod 700 authorized_keys;service ssh start;sleep infinity"
-    ]
-    
-    print(f"Creating RunPod instance with GPU: {args.gpu_type}...")
-    payload = {
-        "name": f"fastvideo-{JOB_ID}-{RUN_ID}",
-        "containerDiskInGb": args.disk_size,
-        "volumeInGb": args.volume_size,
-        "gpuTypeIds": [args.gpu_type],
-        "gpuCount": args.gpu_count,
-        "imageName": image_name,
-        "allowedCudaVersions": ["12.4"],
-        "dockerStartCmd": docker_start_cmd
-    }
-
-    response = requests.post(PODS_API, headers=HEADERS, json=payload)
-    response_data = response.json()
-    print(f"Response: {json.dumps(response_data, indent=2)}")
-
-    return response_data["id"]
-
-
-def wait_for_pod(pod_id):
-    """Wait for pod to be in RUNNING state and fully ready with SSH access"""
-    print("Waiting for RunPod to be ready...")
-
-    # First wait for RUNNING status
-    max_attempts = 10
-    attempts = 0
-    while attempts < max_attempts:
-        response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS)
-        pod_data = response.json()
-        status = pod_data["desiredStatus"]
-
-        if status == "RUNNING":
-            print("RunPod is running! Now waiting for ports to be assigned...")
-            break
-
-        print(
-            f"Current status: {status}, waiting... (attempt {attempts+1}/{max_attempts})"
-        )
-        time.sleep(2)
-        attempts += 1
-
-    if attempts >= max_attempts:
-        raise TimeoutError(
-            "Timed out waiting for RunPod to reach RUNNING state")
-
-    # Wait for ports to be assigned
-    max_attempts = 50
-    attempts = 0
-    while attempts < max_attempts:
-        response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS)
-        pod_data = response.json()
-        port_mappings = pod_data.get("portMappings")
-
-        if (port_mappings is not None and "22" in port_mappings
-                and pod_data.get("publicIp", "") != ""):
-            print("RunPod is ready with SSH access!")
-            print(f"SSH IP: {pod_data['publicIp']}")
-            print(f"SSH Port: {port_mappings['22']}")
-            break
-
-        print(
-            f"Waiting for SSH port and public IP to be available... (attempt {attempts+1}/{max_attempts})"
-        )
-        time.sleep(20)
-        attempts += 1
-
-    if attempts >= max_attempts:
-        raise TimeoutError("Timed out waiting for RunPod SSH access")
-
-
-def execute_command(pod_id):
-    """Execute command on the pod via SSH using system SSH client"""
-    print(f"Running command: {args.test_command}")
-
-    response = requests.get(f"{PODS_API}/{pod_id}", headers=HEADERS)
-    pod_data = response.json()
-    ssh_ip = pod_data["publicIp"]
-    ssh_port = pod_data["portMappings"]["22"]
-
-    # Copy the repository to the pod using scp
-    repo_dir = os.path.abspath(os.getcwd())
-    repo_name = os.path.basename(repo_dir)
-
-    print(f"Copying repository from {repo_dir} to RunPod...")
-
-    tar_command = [
-        "tar", "-czf", "/tmp/repo.tar.gz", "-C",
-        os.path.dirname(repo_dir), repo_name
-    ]
-    subprocess.run(tar_command, check=True)
-
-    # Copy the tarball to the pod
-    scp_command = [
-        "scp", "-o", "StrictHostKeyChecking=no", "-o",
-        "UserKnownHostsFile=/dev/null", "-o", "ServerAliveInterval=60", "-o",
-        "ServerAliveCountMax=10", "-P",
-        str(ssh_port), "/tmp/repo.tar.gz", f"root@{ssh_ip}:/tmp/"
-    ]
-    subprocess.run(scp_command, check=True)
-
-    # For custom image, we can use the pre-configured environment
-    setup_steps = [
-        "tar -xzf /tmp/repo.tar.gz --no-same-owner -C /workspace/",
-        f"cd /workspace/{repo_name}",
-        "source /opt/conda/etc/profile.d/conda.sh",
-        "conda activate fastvideo-dev",
-        args.test_command
-    ]
-    
-    remote_command = " && ".join(setup_steps)
-
-    ssh_command = [
-        "ssh", "-o", "StrictHostKeyChecking=no", "-o",
-        "UserKnownHostsFile=/dev/null", "-o", "ServerAliveInterval=60", "-o",
-        "ServerAliveCountMax=10", "-p",
-        str(ssh_port), f"root@{ssh_ip}", remote_command
-    ]
-
-    print(f"Connecting to {ssh_ip}:{ssh_port}...")
-
-    try:
-        process = subprocess.Popen(ssh_command,
-                                   stdout=subprocess.PIPE,
-                                   stderr=subprocess.STDOUT,
-                                   universal_newlines=True,
-                                   bufsize=0)
-
-        stdout_lines = []
-
-        print("Command output:")
-
-        for line in iter(process.stdout.readline, ''):
-            print(line.strip())
-            stdout_lines.append(line)
-
-        process.wait()
-
-        return_code = process.returncode
-        success = return_code == 0
-
-        stdout_str = "".join(stdout_lines)
-
-        if success:
-            print("Command executed successfully")
-        else:
-            print(f"Command failed with exit code {return_code}")
-
-        result = {
-            "success": success,
-            "return_code": return_code,
-            "stdout": stdout_str,
-            "stderr": ""
-        }
-        return result
-
-    except Exception as e:
-        print(f"Error executing SSH command: {str(e)}")
-        result = {"success": False, "error": str(e), "stdout": "", "stderr": ""}
-        return result
-
-
-def terminate_pod(pod_id):
-    """Terminate the pod"""
-    print("Terminating RunPod...")
-    requests.delete(f"{PODS_API}/{pod_id}", headers=HEADERS)
-    print(f"Terminated pod {pod_id}")
-
-
-def main():
-    pod_id = None
-    try:
-        pod_id = create_pod()
-        wait_for_pod(pod_id)
-        result = execute_command(pod_id)
-
-        if result.get("error") is not None:
-            print(f"Error executing command: {result['error']}")
-            sys.exit(1)
-
-        if not result.get("success", False):
-            print(
-                "Tests failed - check the output above for details on which tests failed"
-            )
-            sys.exit(1)
-
-    finally:
-        if pod_id:
-            terminate_pod(pod_id)
-
-
-if __name__ == "__main__":
-    main()

.github/scripts/runpod_cleanup.py

@@ -1,90 +0,0 @@
-import json
-import os
-import sys
-import uuid
-
-import requests
-
-API_KEY = os.environ['RUNPOD_API_KEY']
-RUN_ID = os.environ.get('GITHUB_RUN_ID', str(uuid.uuid4()))
-PODS_API = "https://rest.runpod.io/v1/pods"
-HEADERS = {
-    "Content-Type": "application/json",
-    "Authorization": f"Bearer {API_KEY}"
-}
-
-
-def get_job_ids():
-    """Parse job IDs from environment variable"""
-    job_ids_str = os.environ.get('JOB_IDS')
-    try:
-        job_ids = json.loads(job_ids_str)
-        if not isinstance(job_ids, list):
-            print("Error: JOB_IDS is not a list.")
-            sys.exit(1)
-        return job_ids
-    except json.JSONDecodeError as e:
-        print(f"Error parsing JOB_IDS: {e}")
-        sys.exit(1)
-
-
-def cleanup_pods():
-    """Find and terminate RunPod instances"""
-    print(f"Run ID: {RUN_ID}")
-
-    single_job_id = os.environ.get('JOB_ID')
-    
-    if single_job_id:
-        job_ids = [single_job_id]
-        print(f"Job ID: {single_job_id}")
-    else:
-        job_ids = get_job_ids()
-        print(f"Job IDs: {job_ids}")
-
-    # Get all pods associated with RunPod API_KEY
-    try:
-        response = requests.get(PODS_API, headers=HEADERS)
-        response.raise_for_status()
-        pods = response.json()
-    except requests.exceptions.RequestException as e:
-        print(f"Error getting pods: {e}")
-        sys.exit(1)
-
-    # Find and terminate pods created by this workflow run
-    terminated_pods = []
-    for pod in pods:
-        pod_name = pod.get("name", "")
-        pod_id = pod.get("id")
-
-        # Check if this pod was created by one of our jobs
-        if any(f"{job_id}-{RUN_ID}" in pod_name for job_id in job_ids):
-            print(f"Found pod: {pod_id} ({pod_name})")
-            try:
-                print(f"Terminating pod {pod_id}...")
-                term_response = requests.delete(f"{PODS_API}/{pod_id}",
-                                                headers=HEADERS)
-                term_response.raise_for_status()
-                terminated_pods.append(pod_id)
-                print(f"Successfully terminated pod {pod_id}")
-            except requests.exceptions.RequestException as e:
-                print(f"Error terminating pod {pod_id}: {e}")
-                sys.exit(1)
-    
-    if terminated_pods:
-        if single_job_id:
-            print(f"Terminated pod: {terminated_pods[0]}")
-        else:
-            print(f"Terminated {len(terminated_pods)} pods: {terminated_pods}")
-    else:
-        if single_job_id:
-            print(f"No pod found matching pattern: {single_job_id}-{RUN_ID}")
-        else:
-            print("No pods found to terminate.")
-
-
-def main():
-    cleanup_pods()
-
-
-if __name__ == "__main__":
-    main()

.github/workflows/build-image.yml

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

.github/workflows/codespell.yml

@@ -0,0 +1,45 @@
+name: codespell
+
+on:
+  # Trigger the workflow on push or pull request,
+  # but only for the main branch
+  push:
+    branches:
+      - main
+    paths:
+      - "**/*.py"
+      - "**/*.md"
+      - "**/*.rst"
+      - pyproject.toml
+      - requirements-lint.txt
+      - .github/workflows/codespell.yml
+  pull_request:
+    branches:
+      - main
+    paths:
+      - "**/*.py"
+      - "**/*.md"
+      - "**/*.rst"
+      - pyproject.toml
+      - requirements-lint.txt
+      - .github/workflows/codespell.yml
+
+jobs:
+  codespell:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'  # or any version you need
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements-lint.txt
+      - name: Spelling check with codespell
+        run: |
+          # Refer to the above environment variable here
+          codespell --toml pyproject.toml $CODESPELL_EXCLUDES

.github/workflows/docs.yml

@@ -1,83 +0,0 @@
-# Sample workflow for building and deploying a Hugo site to GitHub Pages
-name: Deploy FastVideo Docs to Pages
-
-on:
-  # Runs on pushes targeting the default branch
-  push:
-    branches:
-      - main
-    paths:
-      - "docs/**/*.md"
-      - "fastvideo/v1/examples/**/*.py"
-  pull_request:
-    branches: 
-      - main
-    types: [opened, ready_for_review, synchronize, reopened]
-    paths:
-      - "docs/**/*.md"
-      - "fastvideo/v1/examples/**/*.py"
-
-  # Allows you to run this workflow manually from the Actions tab
-  workflow_dispatch:
-
-# Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages
-permissions:
-  contents: read
-  pages: write
-  id-token: write
-
-# Allow only one concurrent deployment, skipping runs queued between the run in-progress and latest queued.
-# However, do NOT cancel in-progress runs as we want to allow these production deployments to complete.
-concurrency:
-  group: "pages"
-  cancel-in-progress: false
-
-# Default to bash
-defaults:
-  run:
-    shell: bash
-
-jobs:
-  pre-commit:
-    uses: ./.github/workflows/pre-commit.yml
-
-  # Build job
-  build:
-    runs-on: ubuntu-latest
-    needs: pre-commit
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-      - name: Setup Pages
-        id: pages
-        uses: actions/configure-pages@v5
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-      - name: Install dependencies
-        run: |
-          cd docs
-          pip install -r requirements-docs.txt
-      - name: Build docs
-        run: |
-          cd docs
-          make clean
-          make html
-      - name: Upload artifact
-        uses: actions/upload-pages-artifact@v3
-        with:
-          path: ./docs/build/html
-
-  # Deployment job
-  deploy:
-    environment:
-      name: github-pages
-      url: ${{ steps.deployment.outputs.page_url }}
-    if: ${{ github.event_name == 'push' }}
-    runs-on: ubuntu-latest
-    needs: build
-    steps:
-      - name: Deploy to GitHub Pages
-        id: deployment
-        uses: actions/deploy-pages@v4

.github/workflows/fastvideo-publish.yml

@@ -6,7 +6,6 @@ on:
       - main
     paths:
       - 'pyproject.toml'  # Trigger when pyproject.toml changes
-  workflow_dispatch:
 
 jobs:
   check-version-change:
@@ -16,7 +15,7 @@ jobs:
       new-version: ${{ steps.check-version.outputs.new-version }}
     steps:
       - name: Checkout code
-        uses: actions/checkout@v4
+        uses: actions/checkout@v3
         with:
           fetch-depth: 2
 
@@ -24,11 +23,11 @@ jobs:
         id: check-version
         run: |
           # Get current commit's version
-          NEW_VERSION=$(grep -oP "version\\s*=\\s*\"\\K[^\"]+\"" pyproject.toml)
+          NEW_VERSION=$(grep -oP 'version\s*=\s*"\K[^"]+' pyproject.toml)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./pyproject.toml | grep -oP "version\\s*=\\s*\"\\K[^\"]+\"" || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./pyproject.toml | grep -oP 'version\s*=\s*"\K[^"]+' || echo "0.0.0")
           echo "Old version: $OLD_VERSION"
           
           if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -42,17 +41,17 @@ jobs:
 
   build-publish-main:
     needs: check-version-change
-    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
+    if: needs.check-version-change.outputs.version-changed == 'true'
     runs-on: ubuntu-latest
     permissions:
       id-token: write  # Needed for OIDC Trusted Publishing
 
     steps:
       - name: Checkout code
-        uses: actions/checkout@v4
+        uses: actions/checkout@v3
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v4
         with:
           python-version: '3.10'
 

.github/workflows/matchers/actionlint.json

@@ -1,17 +0,0 @@
-{
-  "problemMatcher": [
-    {
-      "owner": "actionlint",
-      "pattern": [
-        {
-          "regexp": "^(?:\\x1b\\[\\d+m)?(.+?)(?:\\x1b\\[\\d+m)*:(?:\\x1b\\[\\d+m)*(\\d+)(?:\\x1b\\[\\d+m)*:(?:\\x1b\\[\\d+m)*(\\d+)(?:\\x1b\\[\\d+m)*: (?:\\x1b\\[\\d+m)*(.+?)(?:\\x1b\\[\\d+m)* \\[(.+?)\\]$",
-          "file": 1,
-          "line": 2,
-          "column": 3,
-          "message": 4,
-          "code": 5
-        }
-      ]
-    }
-  ]
-}

.github/workflows/matchers/mypy.json

@@ -1,16 +0,0 @@
-{
-  "problemMatcher": [
-    {
-      "owner": "mypy",
-      "pattern": [
-        {
-          "regexp": "^(.+):(\\d+):\\s(error|warning):\\s(.+)$",
-          "file": 1,
-          "line": 2,
-          "severity": 3,
-          "message": 4
-        }
-      ]
-    }
-  ]
-}

.github/workflows/pr-test.yml

@@ -1,295 +0,0 @@
-name: PR Test
-
-on:
-  push:
-    branches: [main]
-    paths:
-      - "fastvideo/**/*.py"
-      - ".github/workflows/pr-test.yml"
-  pull_request:
-    branches: [main]
-    types: [opened, ready_for_review, synchronize, reopened]
-    paths:
-      - "fastvideo/**/*.py"
-      - ".github/workflows/pr-test.yml"
-  workflow_dispatch:
-    inputs:
-      custom_image:
-        description: "Custom image from this repository (default: fastvideo-dev:latest)"
-        required: false
-        default: "fastvideo-dev:latest"
-        type: string
-      run_encoder_test:
-        description: "Run encoder-test"
-        required: false
-        default: false
-        type: boolean
-      run_vae_test:
-        description: "Run vae-test"
-        required: false
-        default: false
-        type: boolean
-      run_transformer_test:
-        description: "Run transformer-test"
-        required: false
-        default: false
-        type: boolean
-      run_ssim_test:
-        description: "Run ssim-test"
-        required: false
-        default: false
-        type: boolean
-
-env:
-  PYTHONUNBUFFERED: "1"
-
-concurrency:
-  group: pr-test-${{ github.ref }}
-  cancel-in-progress: true
-
-jobs:
-  pre-commit:
-    uses: ./.github/workflows/pre-commit.yml
-
-  change-filter:
-    runs-on: ubuntu-latest
-    needs: pre-commit
-    if: ${{ github.event.pull_request.draft == false || github.event_name == 'workflow_dispatch' }}
-    outputs:
-      encoder-test: ${{ steps.filter.outputs.encoder-test }}
-      vae-test: ${{ steps.filter.outputs.vae-test }}
-      transformer-test: ${{ steps.filter.outputs.transformer-test }}
-    steps:
-      - uses: actions/checkout@v4
-      - uses: dorny/paths-filter@v3
-        id: filter
-        with:
-          filters: |
-            encoder-test:
-              - 'fastvideo/v1/models/encoders/**'
-              - 'fastvideo/v1/models/loaders/**'
-              - 'fastvideo/v1/tests/encoders/**'
-            vae-test:
-              - 'fastvideo/v1/models/vaes/**'
-              - 'fastvideo/v1/models/loaders/**'
-              - 'fastvideo/v1/tests/vaes/**'
-            transformer-test:
-              - 'fastvideo/v1/models/dits/**'
-              - 'fastvideo/v1/models/loaders/**'
-              - 'fastvideo/v1/tests/transformers/**'
-
-  encoder-test:
-    needs: change-filter
-    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
-
-  vae-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.vae-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_vae_test == 'true')
-    runs-on: ubuntu-latest
-    environment: runpod-runners
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-
-      - name: Set up SSH key
-        run: |
-          mkdir -p ~/.ssh
-          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
-          chmod 600 ~/.ssh/id_rsa
-          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
-
-      - name: Install dependencies
-        run: pip install requests
-
-      - name: Run tests on RunPod
-        env:
-          JOB_ID: "vae-test"
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-        timeout-minutes: 30
-        run: >-
-          python .github/scripts/runpod_api.py 
-          --gpu-type "NVIDIA A40"
-          --gpu-count 1
-          --volume-size 100
-          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
-
-      - name: Terminate RunPod Instances
-        if: ${{ always() }}
-        env:
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          JOB_ID: "vae-test"
-        run: python .github/scripts/runpod_cleanup.py
-
-  transformer-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.transformer-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_transformer_test == 'true')
-    runs-on: ubuntu-latest
-    environment: runpod-runners
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-
-      - name: Set up SSH key
-        run: |
-          mkdir -p ~/.ssh
-          echo "${{ secrets.RUNPOD_PRIVATE_KEY }}" > ~/.ssh/id_rsa
-          chmod 600 ~/.ssh/id_rsa
-          ssh-keygen -y -f ~/.ssh/id_rsa > ~/.ssh/id_rsa.pub
-
-      - name: Install dependencies
-        run: pip install requests
-
-      - name: Run tests on RunPod
-        env:
-          JOB_ID: "transformer-test"
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-        timeout-minutes: 30
-        run: >-
-          python .github/scripts/runpod_api.py 
-          --gpu-type "NVIDIA L40S"
-          --gpu-count 1
-          --volume-size 100
-          --image "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
-          --test-command "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
-
-      - name: Terminate RunPod Instances
-        if: ${{ always() }}
-        env:
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-          JOB_ID: "transformer-test"
-        run: python .github/scripts/runpod_cleanup.py
-      
-  ssim-test:
-    needs: change-filter
-    if: >-
-      (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
-
-  runpod-cleanup:
-    needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
-    if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }}
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-
-      - name: Install dependencies
-        run: pip install requests
-
-      - name: Cleanup all RunPod instances
-        env:
-          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test"]' # JSON array of job IDs
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-        run: python .github/scripts/runpod_cleanup.py

.github/workflows/pre-commit.yml

@@ -1,18 +0,0 @@
-name: pre-commit
-
-on:
-    workflow_call:
-
-jobs:
-  pre-commit:
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v4
-    - uses: actions/setup-python@v5
-      with:
-        python-version: "3.10"
-    - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json"
-    - run: echo "::add-matcher::.github/workflows/matchers/mypy.json"
-    - uses: pre-commit/action@v3.0.1
-      with:
-        extra_args: --all-files --hook-stage manual

.github/workflows/ruff.yml

@@ -0,0 +1,50 @@
+name: ruff
+
+on:
+  # Trigger the workflow on push or pull request,
+  # but only for the main branch
+  push:
+    branches:
+      - main
+    paths:
+      - "**/*.py"
+      - pyproject.toml
+      - requirements-lint.txt
+      - .github/workflows/matchers/ruff.json
+      - .github/workflows/ruff.yml
+  pull_request:
+    branches:
+      - main
+    # This workflow is only relevant when one of the following files changes.
+    # However, we have github configured to expect and require this workflow
+    # to run and pass before github with auto-merge a pull request. Until github
+    # allows more flexible auto-merge policy, we can just run this on every PR.
+    # It doesn't take that long to run, anyway.
+    #paths:
+    #  - "**/*.py"
+    #  - pyproject.toml
+    #  - requirements-lint.txt
+    #  - .github/workflows/matchers/ruff.json
+    #  - .github/workflows/ruff.yml
+
+jobs:
+  ruff:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'  # or any version you need
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements-lint.txt
+      - name: Analysing the code with ruff
+        run: |
+          ruff check .
+      - name: Run isort
+        run: |
+          isort . --check-only

.github/workflows/sta-publish.yml

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

.github/workflows/test.yml

@@ -11,10 +11,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Check out repository
-        uses: actions/checkout@v4
+        uses: actions/checkout@v3
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v4
         with:
           python-version: '3.12'  # or any version you need
 
@@ -23,9 +23,11 @@ jobs:
           python -m pip install --upgrade pip setuptools wheel
           pip install torch
           pip install packaging ninja
+          # remove st-attn dependency because no cuda environment
+          sed -i '/st_attn/d' pyproject.toml
           pip install -e .
           pip install pytest
 
       - name: Run Pytest
         run: |
-          pytest --ignore csrc/sliding_tile_attention/test
+          pytest --ignore csrc/sliding_tile_attention/test

.github/workflows/yapf.yml

@@ -0,0 +1,38 @@
+name: yapf
+
+on:
+  # Trigger the workflow on push or pull request,
+  # but only for the main branch
+  push:
+    branches:
+      - main
+    paths:
+      - "**/*.py"
+      - .github/workflows/yapf.yml
+  pull_request:
+    branches:
+      - main
+    paths:
+      - "**/*.py"
+      - .github/workflows/yapf.yml
+
+jobs:
+  yapf:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'  # or any version you need
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install yapf==0.32.0
+          pip install toml==0.10.2
+      - name: Running yapf
+        run: |
+          yapf --diff --recursive .

.gitignore

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

.pre-commit-config.yaml

@@ -1,81 +0,0 @@
-default_stages:
-  - pre-commit # Run locally
-  - manual # Run in CI
-exclude: |
-  (?x)(
-    fastvideo/v1/third_party/.*|
-    csrc/.*|
-    assets/.*|
-    tests/.*|
-    demo/.*|
-    predict\.py|
-    scripts/.*|
-    fastvideo/data_preprocess/.*|
-    fastvideo/dataset/.*|
-    fastvideo/distill/.*|
-    fastvideo/distill\.py|
-    fastvideo/distill_adv\.py|
-    fastvideo/models/.*|
-    fastvideo/sample/.*|
-    fastvideo/train\.py|
-    fastvideo/utils/.*|
-    fastvideo/v1/examples/.*|
-    .github/workflows/fastvideo-publish.yml|
-    .github/workflows/sta-publish.yml
-  )
-repos:
-- repo: https://github.com/google/yapf
-  rev: v0.43.0
-  hooks:
-  - id: yapf
-    args: [--in-place, --verbose]
-    additional_dependencies: [toml] # TODO: Remove when yapf is upgraded
-- repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.11.4
-  hooks:
-  - id: ruff
-    args: [--output-format, github, --fix]
-- repo: https://github.com/codespell-project/codespell
-  rev: v2.4.1
-  hooks:
-  - id: codespell
-    additional_dependencies: ['tomli']
-    args: ['--toml', 'pyproject.toml']
-- repo: https://github.com/PyCQA/isort
-  rev: 6.0.1
-  hooks:
-  - id: isort
-- repo: https://github.com/jackdewinter/pymarkdown
-  rev: v0.9.29
-  hooks:
-  - id: pymarkdown
-    args: [fix]
-- repo: https://github.com/rhysd/actionlint
-  rev: v1.7.7
-  hooks:
-  - id: actionlint
-- repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.15.0
-  hooks:
-  - id: mypy
-    args: [--python-version, '3.10', --follow-imports, "skip", ]
-    additional_dependencies: [types-cachetools, types-setuptools, types-PyYAML, types-requests]
-- repo: local
-  hooks:
-  - id: check-filenames
-    name: Check for spaces in all filenames
-    entry: bash
-    args:
-      - -c
-      - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
-    language: system
-    always_run: true
-    pass_filenames: false
-  # Keep `suggestion` last
-  - id: suggestion
-    name: Suggestion
-    entry: bash -c 'echo "To bypass pre-commit hooks, add --no-verify to git commit."'
-    language: system
-    verbose: true
-    pass_filenames: false
-  # Insert new entries above the `suggestion` entry

Dockerfile

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

README.md

@@ -4,12 +4,19 @@
 
 FastVideo is a lightweight framework for accelerating large video diffusion models.
 
+
 <p align="center">
-    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | 🤗 <a href="https://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>
+    🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank">FastHunyuan</a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank">FastMochi</a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> Slack </a> 
+</p> 
+
+
+
+
 
 https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1
 
+
+
 FastVideo currently offers: (with more to come)
 
 - [NEW!] [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
@@ -21,6 +28,8 @@ FastVideo currently offers: (with more to come)
 
 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/).
@@ -28,38 +37,21 @@ Dev in progress and highly experimental.
 - ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment.
 - ```2024/12/17```: `FastVideo` v1.0 is released.
 
-## 🔧 Installation from source
-The code is tested on Python 3.10-3.12, CUDA 12.4 and H100.
 
+## 🔧 Installation
+The code is tested on Python 3.10.0, CUDA 12.4 and H100.
 ```
-# Clone FastVideo
-git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
-
-# Install FastVideo
-pip install -e .
-
-# Install Flash Attention (optional)
-pip install flash-attn==2.7.0.post2
+./env_setup.sh fastvideo
 ```
-
 To try Sliding Tile Attention (optional), please follow the instruction in [csrc/sliding_tile_attention/README.md](csrc/sliding_tile_attention/README.md) to install STA.
 
-You can also install the Sliding Tile Attention package using
-
-```
-pip install st_attn==0.0.4
-```
-
 ## 🚀 Inference
-### Inference StepVideo with Sliding Tile Attention
+### 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
+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
@@ -67,49 +59,44 @@ sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo
 
 ### Inference HunyuanVideo with Sliding Tile Attention
 First, download the model:
-
 ```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model 
 ```
-
 We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only.
-
 ```bash
 sh scripts/inference/inference_hunyuan_STA.sh
 ```
-
 ### Video Demos using STA + Teacache
 Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100.
 
 ### Inference FastHunyuan on single RTX4090
 We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM.
-
 ```bash
 # Download the model weight
 python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model
 # CLI inference
 bash scripts/inference/inference_hunyuan_hf_quantization.sh
 ```
-
 For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU):
 
+
 | Configuration                  | Memory to Init Transformer | Peak Memory After Init Pipeline (Denoise) | Diffusion Time | End-to-End Time |
 |--------------------------------|----------------------------|--------------------------------------------|----------------|-----------------|
 | BF16 + Pipeline CPU Offload    | 23.883G                   | 33.744G                                    | 81s            | 121.5s          |
 | INT8 + Pipeline CPU Offload    | 13.911G                   | 27.979G                                    | 88s            | 116.7s          |
 | NF4 + Pipeline CPU Offload     | 9.453G                    | 19.26G                                     | 78s            | 114.5s          |
+           
+
 
 For improved quality in generated videos, we recommend using a GPU with 80GB of memory to run the BF16 model with the original Hunyuan pipeline. To execute the inference, use the following section:
 
 ### FastHunyuan
-
 ```bash
 # Download the model weight
 python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan --local_dir=data/FastHunyuan --repo_type=model
 # CLI inference
 bash scripts/inference/inference_hunyuan.sh
 ```
-
 You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo).
 
 ### FastMochi
@@ -121,99 +108,79 @@ python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffuser
 bash scripts/inference/inference_mochi_sp.sh
 ```
 
+
 ## 🎯 Distill
 Our distillation recipe is based on [Phased Consistency Model](https://github.com/G-U-N/Phased-Consistency-Model). We did not find significant improvement using multi-phase distillation, so we keep the one phase setup similar to the original latent consistency model's recipe.
 We use the [MixKit](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/all_mixkit) dataset for distillation. To avoid running the text encoder and VAE during training, we preprocess all data to generate text embeddings and VAE latents.
 Preprocessing instructions can be found [data_preprocess.md](docs/data_preprocess.md). For convenience, we also provide preprocessed data that can be downloaded directly using the following command:
-
 ```bash
 python scripts/huggingface/download_hf.py --repo_id=FastVideo/HD-Mixkit-Finetune-Hunyuan --local_dir=data/HD-Mixkit-Finetune-Hunyuan --repo_type=dataset
 ```
-
 Next, download the original model weights with:
-
 ```bash
 python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model # original hunyuan
 python scripts/huggingface/download_hf.py --repo_id=genmo/mochi-1-preview --local_dir=data/mochi --repo_type=model # original mochi
 ```
-
 To launch the distillation process, use the following commands:
-
 ```
 bash scripts/distill/distill_hunyuan.sh # for hunyuan
 bash scripts/distill/distill_mochi.sh # for mochi
 ```
-
 We also provide an optional script for distillation with adversarial loss, located at `fastvideo/distill_adv.py`. Although we tried adversarial loss, we did not observe significant improvements.
 ## Finetune
 ### ⚡ Full Finetune
 Ensure your data is prepared and preprocessed in the format specified in [data_preprocess.md](docs/data_preprocess.md). For convenience, we also provide a mochi preprocessed Black Myth Wukong data that can be downloaded directly:
-
 ```bash
 python scripts/huggingface/download_hf.py --repo_id=FastVideo/Mochi-Black-Myth --local_dir=data/Mochi-Black-Myth --repo_type=dataset
 ```
-
 Download the original model weights as specified in [Distill Section](#-distill):
 
 Then you can run the finetune with:
-
 ```
 bash scripts/finetune/finetune_mochi.sh # for mochi
 ```
-
 **Note that for finetuning, we did not tune the hyperparameters in the provided script.**
-### ⚡ Lora Finetune
+### ⚡ Lora Finetune 
 
 Hunyuan supports Lora fine-tuning of videos up to 720p. Demos and prompts of Black-Myth-Wukong can be found in [here](https://huggingface.co/FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight). You can download the Lora weight through:
-
 ```bash
 python scripts/huggingface/download_hf.py --repo_id=FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight --local_dir=data/Hunyuan-Black-Myth-Wukong-lora-weight --repo_type=model
 ```
-
 #### Minimum Hardware Requirement
 - 40 GB GPU memory each for 2 GPUs with lora.
-- 30 GB GPU memory each for 2 GPUs with CPU offload and lora.
+- 30 GB GPU memory each for 2 GPUs with CPU offload and lora.  
+
 
 Currently, both Mochi and Hunyuan models support Lora finetuning through diffusers. To generate personalized videos from your own dataset, you'll need to follow three main steps: dataset preparation, finetuning, and inference.
 
 #### Dataset Preparation
-We provide scripts to better help you get started to train on your own characters!
+We provide scripts to better help you get started to train on your own characters!  
 You can run this to organize your dataset to get the videos2caption.json before preprocess. Specify your video folder and corresponding caption folder (caption files should be .txt files and have the same name with its video):
-
 ```
 python scripts/dataset_preparation/prepare_json_file.py --video_dir data/input_videos/ --prompt_dir data/captions/ --output_path data/output_folder/videos2caption.json --verbose
 ```
-
 Also, we provide script to resize your videos:
-
 ```
-python scripts/data_preprocess/resize_videos.py
+python scripts/data_preprocess/resize_videos.py 
 ```
-
 #### Finetuning
 After basic dataset preparation and preprocess, you can start to finetune your model using Lora:
-
 ```
 bash scripts/finetune/finetune_hunyuan_hf_lora.sh
 ```
-
 #### Inference
 For inference with Lora checkpoint, you can run the following scripts with additional parameter `--lora_checkpoint_dir`:
-
 ```
-bash scripts/inference/inference_hunyuan_hf.sh
+bash scripts/inference/inference_hunyuan_hf.sh 
 ```
-
 **We also provide scripts for Mochi in the same directory.**
 
 #### Finetune with Both Image and Video
-Our codebase support finetuning with both image and video.
-
+Our codebase support finetuning with both image and video. 
 ```bash
 bash scripts/finetune/finetune_hunyuan.sh
 bash scripts/finetune/finetune_mochi_lora_mix.sh
 ```
-
 For Image-Video Mixture Fine-tuning, make sure to enable the `--group_frame` option in your script.
 
 ## 📑 Development Plan
@@ -238,26 +205,26 @@ We learned and reused code from the following projects: [PCM](https://github.com
 
 We thank MBZUAI and Anyscale for their support throughout this project.
 
-## Citation
+## Citation 
 If you use FastVideo for your research, please cite our paper:
 
 ```bibtex
 @misc{zhang2025fastvideogenerationsliding,
-      title={Fast Video Generation with Sliding Tile Attention},
+      title={Fast Video Generation with Sliding Tile Attention}, 
       author={Peiyuan Zhang and Yongqi Chen and Runlong Su and Hangliang Ding and Ion Stoica and Zhenghong Liu and Hao Zhang},
       year={2025},
       eprint={2502.04507},
       archivePrefix={arXiv},
       primaryClass={cs.CV},
-      url={https://arxiv.org/abs/2502.04507},
+      url={https://arxiv.org/abs/2502.04507}, 
 }
 @misc{ding2025efficientvditefficientvideodiffusion,
-      title={Efficient-vDiT: Efficient Video Diffusion Transformers With Attention Tile},
+      title={Efficient-vDiT: Efficient Video Diffusion Transformers With Attention Tile}, 
       author={Hangliang Ding and Dacheng Li and Runlong Su and Peiyuan Zhang and Zhijie Deng and Ion Stoica and Hao Zhang},
       year={2025},
       eprint={2502.06155},
       archivePrefix={arXiv},
       primaryClass={cs.CV},
-      url={https://arxiv.org/abs/2502.06155},
+      url={https://arxiv.org/abs/2502.06155}, 
 }
 ```

csrc/sliding_tile_attention/setup.py

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

csrc/sliding_tile_attention/st_attn.cpp

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

csrc/sliding_tile_attention/st_attn/__init__.py

@@ -1,46 +1,212 @@
 import math
-
+import subprocess
 import torch
-from st_attn_cuda import sta_fwd
+from torch.nn.attention.flex_attention import flex_attention
+from functools import lru_cache
+from typing import Tuple
+from torch import BoolTensor, IntTensor
+from torch.nn.attention.flex_attention import create_block_mask
+
+# Peiyuan: This is neccesay. Dont know why. see https://github.com/pytorch/pytorch/issues/135028
+torch._inductor.config.realize_opcount_threshold = 100
 
 
-def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True, img_latent_shape='30*48*80'):
-    seq_length = q_all.shape[2]
-    img_latent_shape_mapping = {
-        '30x48x80':1,
-        '36x48x48':2,
-        '18x48x80':3, 
-    }
-    if has_text:
-        assert q_all.shape[
-            2] >= 115200, "STA currently only supports video with latent size (30, 48, 80), which is 117 frames x 768 x 1280 pixels"
-        assert q_all.shape[1] == len(window_size), "Number of heads must match the number of window sizes"
-        target_size = math.ceil(seq_length / 384) * 384
-        pad_size = target_size - seq_length
-        if pad_size > 0:
-            q_all = torch.cat([q_all, q_all[:, :, -pad_size:]], dim=2)
-            k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2)
-            v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2)
-    else:
-        if img_latent_shape == '36x48x48': # Stepvideo 204x768x68
-            assert q_all.shape[2] == 82944
-        elif img_latent_shape == '18x48x80': # Wan 69x768x1280
-            assert q_all.shape[2] == 69120
+def generate_sta_mask(canvas_twh, kernel_twh, tile_twh, text_length):
+    """Generates a 3D NATTEN attention mask with a given kernel size.
+    
+    Args:
+        canvas_t: The time dimension of the canvas.
+        canvas_h: The height of the canvas.
+        canvas_w: The width of the canvas.
+        kernel_t: The time dimension of the kernel.
+        kernel_h: The height of the kernel.
+        kernel_w: The width of the kernel.
+    """
+    canvas_t, canvas_h, canvas_w = canvas_twh
+    kernel_t, kernel_h, kernel_w = kernel_twh
+    tile_t_size, tile_h_size, tile_w_size = tile_twh
+    total_tile_size = tile_t_size * tile_h_size * tile_w_size
+    canvas_tile_t, canvas_tile_h, canvas_tile_w = canvas_t // tile_t_size, canvas_h // tile_h_size, canvas_w // tile_w_size
+    img_seq_len = canvas_t * canvas_h * canvas_w
+
+    def get_tile_t_x_y(idx: IntTensor) -> Tuple[IntTensor, IntTensor, IntTensor]:
+        tile_id = idx // total_tile_size
+        tile_t = tile_id // (canvas_tile_h * canvas_tile_w)
+        tile_h = (tile_id % (canvas_tile_h * canvas_tile_w)) // canvas_tile_w
+        tile_w = tile_id % canvas_tile_w
+        return tile_t, tile_h, tile_w
+
+    def sta_mask_3d(
+        b: IntTensor,
+        h: IntTensor,
+        q_idx: IntTensor,
+        kv_idx: IntTensor,
+    ) -> BoolTensor:
+        q_t_tile, q_x_tile, q_y_tile = get_tile_t_x_y(q_idx)
+        kv_t_tile, kv_x_tile, kv_y_tile = get_tile_t_x_y(kv_idx)
+        # kernel nominally attempts to center itself on the query, but kernel center
+        # is clamped to a fixed distance (kernel half-length) from the canvas edge
+        kernel_center_t = q_t_tile.clamp(kernel_t // 2, (canvas_tile_t - 1) - kernel_t // 2)
+        kernel_center_x = q_x_tile.clamp(kernel_h // 2, (canvas_tile_h - 1) - kernel_h // 2)
+        kernel_center_y = q_y_tile.clamp(kernel_w // 2, (canvas_tile_w - 1) - kernel_w // 2)
+        time_mask = (kernel_center_t - kv_t_tile).abs() <= kernel_t // 2
+        hori_mask = (kernel_center_x - kv_x_tile).abs() <= kernel_h // 2
+        vert_mask = (kernel_center_y - kv_y_tile).abs() <= kernel_w // 2
+        image_mask = (q_idx < img_seq_len) & (kv_idx < img_seq_len)
+        image_to_text_mask = (q_idx < img_seq_len) & (kv_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length)
+        text_to_all_mask = (q_idx >= img_seq_len) & (kv_idx < img_seq_len + text_length)
+        return (image_mask & time_mask & hori_mask & vert_mask) | image_to_text_mask | text_to_all_mask
+
+    sta_mask_3d.__name__ = f"natten_3d_c{canvas_t}x{canvas_w}x{canvas_h}_k{kernel_t}x{kernel_w}x{kernel_h}"
+    return sta_mask_3d
+
+
+def get_sliding_tile_attention_mask(kernel_size, tile_size, img_size, text_length, device, text_max_len=256):
+    img_seq_len = img_size[0] * img_size[1] * img_size[2]
+    image_mask = generate_sta_mask(img_size, kernel_size, tile_size, text_length)
+    mask = create_block_mask(image_mask,
+                             B=None,
+                             H=None,
+                             Q_LEN=img_seq_len + text_max_len,
+                             KV_LEN=img_seq_len + text_max_len,
+                             device=device,
+                             _compile=True)
+    return mask
+
+
+def get_gpu_type():
+    try:
+        # Run nvidia-smi to get GPU information
+        result = subprocess.check_output(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader']).decode('utf-8')
+        
+        # Check if H100 is in any of the GPU names
+        gpus = [gpu.strip() for gpu in result.split('\n') if gpu.strip()]
+        
+        for gpu in gpus:
+            if 'H100' in gpu:
+                return 'H100'
+            if '4090' in gpu:
+                return '4090'
+                
+        return None
+    except Exception as e:
+        return None
+
+gpu_type = get_gpu_type()
+
+if gpu_type == 'H100':
+    from st_attn_cuda import sta_fwd
+
+    
+@lru_cache(maxsize=32)
+def get_compiled_flex_attention(strategy, tile_size, image_size, text_length, device):
+    """
+    Create and compile flex attention with a specific sliding block mask.
+    This function is cached to avoid recompiling for the same parameters.
+    
+    Args:
+        strategy (tuple): A tuple (t, h, w) defining the strategy
+        tile_size (tuple): A tuple (ts_t, ts_h, ts_w) defining the tile size
+        image_size (tuple): A tuple (n_t, n_h, n_w) defining the image size
+        text_length (int): The text length
+        device (str): The device to use
+        
+    Returns:
+        function: A compiled flex attention function with the specified mask
+    """
+    # Convert strategy to the required format (ceil(t*3/2), h*2, w)
+    adjusted_strategy = strategy
+    
+    # Get the sliding block attention mask
+    mask = get_sliding_tile_attention_mask(
+        adjusted_strategy, 
+        tile_size, 
+        image_size, 
+        text_length, 
+        device
+    )
+    
+    def flex_attn_with_mask(q, k, v, scale=None):
+        return flex_attention(q, k, v, block_mask=mask, scale=scale)
+    
+    # Compile the wrapper function
+    compiled_flex_attn = torch.compile(flex_attn_with_mask)
+    
+    return compiled_flex_attn
+
+def flex_sliding_tile_attention(q_all, k_all, v_all, strategy, tile_size, 
+                           image_size, text_length, scale=None):
+    device = q_all.device
+    
+    # Get the compiled flex attention function (cached if called with same parameters)
+    compiled_flex_attn = get_compiled_flex_attention(
+        strategy, 
+        tile_size,
+        image_size, 
+        text_length, 
+        device
+    )
+    
+    
+    # Apply the compiled flex attention
+    output = compiled_flex_attn(q_all, k_all, v_all, scale=scale)
+
+        
+    return output
+
+def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True):
+    if gpu_type == 'H100':
+        seq_length = q_all.shape[2]
+        if has_text:
+            assert q_all.shape[
+                2] == 115456, "STA currently only supports video with latent size (30, 48, 80), which is 117 frames x 768 x 1280 pixels"
+            assert q_all.shape[1] == len(window_size), "Number of heads must match the number of window sizes"
+            target_size = math.ceil(seq_length / 384) * 384
+            pad_size = target_size - seq_length
+            if pad_size > 0:
+                q_all = torch.cat([q_all, q_all[:, :, -pad_size:]], dim=2)
+                k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2)
+                v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2)
         else:
-            raise ValueError(f"Unsupported {img_latent_shape}, current shape is {q_all.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan")
+            assert q_all.shape[2] == 82944
 
-    kernel_aspect_ratio_flag = img_latent_shape_mapping[img_latent_shape]
-    hidden_states = torch.empty_like(q_all)
-    # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
-    for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size):
-        for batch in range(q_all.shape[0]):
-            q_head, k_head, v_head, o_head = (q_all[batch:batch + 1, head_index:head_index + 1],
-                                              k_all[batch:batch + 1,
-                                                    head_index:head_index + 1], v_all[batch:batch + 1,
-                                                                                      head_index:head_index + 1],
-                                              hidden_states[batch:batch + 1, head_index:head_index + 1])
+        hidden_states = torch.empty_like(q_all)
+        # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
+        for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size):
+            for batch in range(q_all.shape[0]):
+                q_head, k_head, v_head, o_head = (q_all[batch:batch + 1, head_index:head_index + 1],
+                                                k_all[batch:batch + 1,
+                                                        head_index:head_index + 1], v_all[batch:batch + 1,
+                                                                                        head_index:head_index + 1],
+                                                hidden_states[batch:batch + 1, head_index:head_index + 1])
 
-            _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text, kernel_aspect_ratio_flag)
-    if has_text:
-        _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True, kernel_aspect_ratio_flag)
-    return hidden_states[:, :, :seq_length]
+                _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text)
+        if has_text:
+            _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True)
+        return hidden_states[:, :, :seq_length]
+    else:
+        assert q_all.shape[
+                2] == 46336, "Flex STA currently only supports video with latent size (12, 48, 80), which is 45 frames x 768 x 1280 pixels"
+        head_num = q_all.size(1)
+        hidden_states = torch.empty_like(q_all)
+        strategy_to_heads = {}
+        for head_index in range(head_num):
+            strategy = tuple(window_size[head_index])  # Convert list to tuple for dict key
+            if strategy not in strategy_to_heads:
+                strategy_to_heads[strategy] = []
+            strategy_to_heads[strategy].append(head_index)
+        for strategy, heads in strategy_to_heads.items():                
+            # Gather all heads with this strategy
+            query_heads = torch.cat([q_all[:, head_idx:head_idx + 1, :, :] for head_idx in heads], dim=1)
+            key_heads = torch.cat([k_all[:, head_idx:head_idx + 1, :, :] for head_idx in heads], dim=1)
+            value_heads = torch.cat([v_all[:, head_idx:head_idx + 1, :, :] for head_idx in heads], dim=1)
+            
+            # Process all heads with this strategy at once
+            # processed_heads = selected_attn_processor[processor_idx](query_heads, key_heads, value_heads)
+            processed_heads = flex_sliding_tile_attention(query_heads, key_heads, value_heads, strategy, (6, 8, 8), (12, 48, 80), text_length)
+            
+            # Distribute results back to the correct positions
+            for i, head_idx in enumerate(heads):
+                hidden_states[:, head_idx:head_idx + 1, :, :] = processed_heads[:, i:i + 1, :, :]
+            
+        return hidden_states

csrc/sliding_tile_attention/st_attn/st_attn_h100.cu

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

docs/Makefile

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

docs/README.md

@@ -1,20 +0,0 @@
-# FastVideo documents
-
-## Build the docs
-
-```bash
-# Install dependencies.
-pip install -r requirements-docs.txt
-
-# Build the docs.
-make clean
-make html
-```
-
-## Open the docs with your browser
-
-```bash
-python -m http.server -d build/html/
-```
-
-Launch your browser and open localhost:8000.

docs/data_preprocess.md

@@ -1,16 +1,16 @@
 
+
+
 ## 🧱 Data Preprocess
 
 To save GPU memory, we precompute text embeddings and VAE latents to eliminate the need to load the text encoder and VAE during training.
 
-We provide a sample dataset to help you get started. Download the source media using the following command:
 
+We provide a sample dataset to help you get started. Download the source media using the following command:
 ```bash
 python scripts/huggingface/download_hf.py --repo_id=FastVideo/Image-Vid-Finetune-Src --local_dir=data/Image-Vid-Finetune-Src --repo_type=dataset
 ```
-
 To preprocess the dataset for fine-tuning or distillation, run:
-
 ```
 bash scripts/preprocess/preprocess_mochi_data.sh # for mochi
 bash scripts/preprocess/preprocess_hunyuan_data.sh # for hunyuan
@@ -33,16 +33,13 @@ path_to_dataset_folder/
 Format the JSON file as a list, where each item represents a media source:
 
 For image media,
-
 ```
 {
     "path": "0.jpg",
     "cap": ["captions"]
 }
 ```
-
-For video media,
-
+For video media, 
 ```
 {
     "path": "1.mp4",
@@ -65,9 +62,7 @@ path_to_media_source_foder,path_to_json_file
 ```
 
 Adjust the `DATA_MERGE_PATH` and `OUTPUT_DIR` in `scripts/preprocess/preprocess_****_data.sh` accordingly and run:
-
 ```
 bash scripts/preprocess/preprocess_****_data.sh
 ```
-
 The preprocessed data will be put into the `OUTPUT_DIR` and the `videos2caption.json` can be used in finetune and distill scripts.

docs/make.bat

@@ -1,35 +0,0 @@
-@ECHO OFF
-
-pushd %~dp0
-
-REM Command file for Sphinx documentation
-
-if "%SPHINXBUILD%" == "" (
-	set SPHINXBUILD=sphinx-build
-)
-set SOURCEDIR=source
-set BUILDDIR=build
-
-%SPHINXBUILD% >NUL 2>NUL
-if errorlevel 9009 (
-	echo.
-	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
-	echo.installed, then set the SPHINXBUILD environment variable to point
-	echo.to the full path of the 'sphinx-build' executable. Alternatively you
-	echo.may add the Sphinx directory to PATH.
-	echo.
-	echo.If you don't have Sphinx installed, grab it from
-	echo.https://www.sphinx-doc.org/
-	exit /b 1
-)
-
-if "%1" == "" goto help
-
-%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
-goto end
-
-:help
-%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
-
-:end
-popd

docs/requirements-docs.txt

@@ -1,25 +0,0 @@
-sphinx==6.2.1
-sphinx-argparse==0.4.0
-sphinx-book-theme==1.0.1
-sphinx-copybutton==0.5.2
-sphinx-design==0.6.1
-sphinx-togglebutton==0.3.2
-myst-parser==3.0.1
-msgspec
-cloudpickle
-
-# 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

docs/seed_parameter_behavior.md

@@ -1,51 +0,0 @@
-# Seed Parameter Behavior in vLLM
-
-## Overview
-
-The `seed` parameter in vLLM is used to control the random states for various random number generators. This parameter can affect the behavior of random operations in user code, especially when working with models in vLLM.
-
-## Default Behavior
-
-By default, the `seed` parameter is set to `None`. When the `seed` parameter is `None`, the global random states for `random`, `np.random`, and `torch.manual_seed` are not set. This means that the random operations will behave as expected, without any fixed random states.
-
-## Specifying a Seed
-
-If a specific seed value is provided, the global random states for `random`, `np.random`, and `torch.manual_seed` will be set accordingly. This can be useful for reproducibility, as it ensures that the random operations produce the same results across multiple runs.
-
-## Example Usage
-
-### Without Specifying a Seed
-
-```python
-import random
-from vllm import LLM
-
-# Initialize a vLLM model without specifying a seed
-model = LLM(model="Qwen/Qwen2.5-0.5B-Instruct")
-
-# Try generating random numbers
-print(random.randint(0, 100))  # Outputs different numbers across runs
-```
-
-### Specifying a Seed
-
-```python
-import random
-from vllm import LLM
-
-# Initialize a vLLM model with a specific seed
-model = LLM(model="Qwen/Qwen2.5-0.5B-Instruct", seed=42)
-
-# Try generating random numbers
-print(random.randint(0, 100))  # Outputs the same number across runs
-```
-
-## Important Notes
-
-- If the `seed` parameter is not specified, the behavior of global random states remains unaffected.
-- If a specific seed value is provided, the global random states for `random`, `np.random`, and `torch.manual_seed` will be set to that value.
-- This behavior can be useful for reproducibility but may lead to non-intuitive behavior if the user is not explicitly aware of it.
-
-## Conclusion
-
-Understanding the behavior of the `seed` parameter in vLLM is crucial for ensuring the expected behavior of random operations in your code. By default, the `seed` parameter is set to `None`, which means that the global random states are not affected. However, specifying a seed value can help achieve reproducibility in your experiments.

docs/source/_static/custom.css

@@ -1,8 +0,0 @@
-.vertical-table-header th.head:not(.stub) {
-    writing-mode: sideways-lr;
-    white-space: nowrap;
-    max-width: 0;
-    p {
-       margin: 0;
-    }
-}

docs/source/_static/custom.js

@@ -1,18 +0,0 @@
-// Update URL search params when tab is clicked
-  document.addEventListener("DOMContentLoaded", function () {
-    const tabs = document.querySelectorAll(".sd-tab-label");
-
-    function updateURL(tab) {
-      const syncGroup = tab.getAttribute("data-sync-group");
-      const syncId = tab.getAttribute("data-sync-id");
-      if (syncGroup && syncId) {
-          const url = new URL(window.location);
-          url.searchParams.set(syncGroup, syncId);
-          window.history.replaceState(null, "", url);
-      }
-    }
-
-    tabs.forEach(tab => {
-        tab.addEventListener("click", () => updateURL(tab));
-    });
-});

docs/source/_templates/sections/header.html

@@ -1,39 +0,0 @@
-<style>
-  .notification-bar {
-    width: 100vw;
-    display: flex;
-    justify-content: center;
-    align-items: center;
-    font-size: 16px;
-    padding: 0 6px 0 6px;
-  }
-  .notification-bar p {
-    margin: 0;
-  }
-  .notification-bar a {
-    font-weight: bold;
-    text-decoration: none;
-  }
-
-  /* Light mode styles (default) */
-  .notification-bar {
-    background-color: #fff3cd;
-    color: #856404;
-  }
-  .notification-bar a {
-    color: #d97706;
-  }
-
-  /* Dark mode styles */
-  html[data-theme=dark] .notification-bar {
-    background-color: #333;
-    color: #ddd;
-  }
-  html[data-theme=dark] .notification-bar a {
-    color: #ffa500; /* Brighter color for visibility */
-  }
-</style>
-
-<div class="notification-bar">
-  <p>You are viewing the latest developer preview docs. <a href="https://docs.vllm.ai/en/stable/">Click here</a> to view docs for the latest stable release.</p>
-</div>

docs/source/conf.py

@@ -1,260 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-# Configuration file for the Sphinx documentation builder.
-#
-# This file only contains a selection of the most common options. For a full
-# list see the documentation:
-# https://www.sphinx-doc.org/en/master/usage/configuration.html
-
-# -- Path setup --------------------------------------------------------------
-
-# If extensions (or modules to document with autodoc) are in another directory,
-# add these directories to sys.path here. If the directory is relative to the
-# documentation root, use os.path.abspath to make it absolute, like shown here.
-
-import datetime
-import inspect
-import logging
-import os
-import sys
-from typing import Optional
-
-import requests
-from sphinx.ext import autodoc
-
-logger = logging.getLogger(__name__)
-sys.path.append(os.path.abspath("../.."))
-
-# -- Project information -----------------------------------------------------
-
-project = 'FastVideo'
-copyright = f'{datetime.datetime.now().year}, FastVideo Team'
-author = 'the FastVideo Team'
-
-# -- General configuration ---------------------------------------------------
-
-# Add any Sphinx extension module names here, as strings. They can be
-# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
-# ones.
-extensions = [
-    "sphinx.ext.napoleon",
-    "sphinx.ext.linkcode",
-    "sphinx.ext.intersphinx",
-    "sphinx_copybutton",
-    "sphinx.ext.autodoc",
-    "sphinx.ext.autosummary",
-    "myst_parser",
-    "sphinxarg.ext",
-    "sphinx_design",
-    "sphinx_togglebutton",
-]
-myst_enable_extensions = [
-    "colon_fence",
-]
-
-# Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
-
-# List of patterns, relative to source directory, that match files and
-# directories to ignore when looking for source files.
-# This pattern also affects html_static_path and html_extra_path.
-exclude_patterns: list[str] = ["**/*.template.md", "**/*.inc.md"]
-
-# Exclude the prompt "$" when copying code
-copybutton_prompt_text = r"\$ "
-copybutton_prompt_is_regexp = True
-
-# -- Options for HTML output -------------------------------------------------
-
-# The theme to use for HTML and HTML Help pages.  See the documentation for
-# a list of builtin themes.
-#
-html_title = project
-html_theme = 'sphinx_book_theme'
-html_logo = '../../assets/logo.jpg'
-#html_favicon = 'assets/logos/vllm-logo-only-light.ico'
-html_theme_options = {
-    'path_to_docs': 'docs/source',
-    'repository_url': 'https://github.com/hao-ai-lab/FastVideo/',
-    'use_repository_button': True,
-    'use_edit_page_button': True,
-}
-# Add any paths that contain custom static files (such as style sheets) here,
-# relative to this directory. They are copied after the builtin static files,
-# so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ["_static"]
-html_js_files = ["custom.js"]
-html_css_files = ["custom.css"]
-
-myst_url_schemes = {
-    'http': None,
-    'https': None,
-    'mailto': None,
-    'ftp': None,
-    "gh-issue": {
-        "url":
-        "https://github.com/hao-ai-lab/FastVideo/issues/{{path}}#{{fragment}}",
-        "title": "Issue #{{path}}",
-        "classes": ["github"],
-    },
-    "gh-pr": {
-        "url":
-        "https://github.com/hao-ai-lab/FastVideo/pull/{{path}}#{{fragment}}",
-        "title": "Pull Request #{{path}}",
-        "classes": ["github"],
-    },
-    "gh-dir": {
-        "url": "https://github.com/hao-ai-lab/FastVideo/tree/main/{{path}}",
-        "title": "{{path}}",
-        "classes": ["github"],
-    },
-    "gh-file": {
-        "url": "https://github.com/hao-ai-lab/FastVideo/blob/main/{{path}}",
-        "title": "{{path}}",
-        "classes": ["github"],
-    },
-}
-
-# see https://docs.readthedocs.io/en/stable/reference/environment-variables.html # noqa
-READTHEDOCS_VERSION_TYPE = os.environ.get('READTHEDOCS_VERSION_TYPE')
-if READTHEDOCS_VERSION_TYPE == "tag":
-    # remove the warning banner if the version is a tagged release
-    header_file = os.path.join(os.path.dirname(__file__),
-                               "_templates/sections/header.html")
-    # The file might be removed already if the build is triggered multiple times
-    # (readthedocs build both HTML and PDF versions separately)
-    if os.path.exists(header_file):
-        os.remove(header_file)
-
-
-# Generate additional rst documentation here.
-def setup(app):
-    from docs.source.generate_examples import generate_examples
-    generate_examples()
-
-
-_cached_base: str = ""
-_cached_branch: str = ""
-
-
-def get_repo_base_and_branch(
-        pr_number: str) -> tuple[Optional[str], Optional[str]]:
-    global _cached_base, _cached_branch
-    if _cached_base and _cached_branch:
-        return _cached_base, _cached_branch
-
-    url = f"https://api.github.com/repos/hao-ai-lab/FastVideo/pulls/{pr_number}"
-    response = requests.get(url)
-    if response.status_code == 200:
-        data = response.json()
-        _cached_base = data['head']['repo']['full_name']
-        _cached_branch = data['head']['ref']
-        return _cached_base, _cached_branch
-    else:
-        logger.error("Failed to fetch PR details: %s", response)
-        return None, None
-
-
-def linkcode_resolve(domain, info):
-    if domain != 'py':
-        return None
-    if not info['module']:
-        return None
-    module = info['module']
-
-    # try to determine the correct file and line number to link to
-    obj = sys.modules[module]
-
-    # get as specific as we can
-    lineno: int = 0
-    filename: str = ""
-    try:
-        for part in info['fullname'].split('.'):
-            obj = getattr(obj, part)
-
-            if not (inspect.isclass(obj) or inspect.isfunction(obj)
-                    or inspect.ismethod(obj)):
-                obj = obj.__class__  # type: ignore[assignment]
-
-            lineno = inspect.getsourcelines(obj)[1]
-            filename = (inspect.getsourcefile(obj)
-                        or f"{filename}.py").split("FastVideo/", 1)[1]
-    except Exception:
-        # For some things, like a class member, won't work, so
-        # we'll use the line number of the parent (the class)
-        pass
-
-    if filename.startswith("checkouts/"):
-        # a PR build on readthedocs
-        pr_number = filename.split("/")[1]
-        filename = filename.split("/", 2)[2]
-        base, branch = get_repo_base_and_branch(pr_number)
-        if base and branch:
-            return f"https://github.com/{base}/blob/{branch}/{filename}#L{lineno}"
-
-    # Otherwise, link to the source file on the main branch
-    return f"https://github.com/hao-ai-lab/FastVideo/blob/main/{filename}#L{lineno}"
-
-
-# Mock out external dependencies here, otherwise the autodoc pages may be blank.
-autodoc_mock_imports = [
-    "blake3",
-    "compressed_tensors",
-    "cpuinfo",
-    "cv2",
-    "torch",
-    "transformers",
-    "psutil",
-    "prometheus_client",
-    "sentencepiece",
-    "vllm._C",
-    "PIL",
-    "numpy",
-    'triton',
-    "tqdm",
-    "tensorizer",
-    "pynvml",
-    "outlines",
-    "xgrammar",
-    "librosa",
-    "soundfile",
-    "gguf",
-    "lark",
-    "decord",
-]
-
-for mock_target in autodoc_mock_imports:
-    if mock_target in sys.modules:
-        logger.info(
-            "Potentially problematic mock target (%s) found; "
-            "autodoc_mock_imports cannot mock modules that have already "
-            "been loaded into sys.modules when the sphinx build starts.",
-            mock_target)
-
-
-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":
-    ("https://typing-extensions.readthedocs.io/en/latest", None),
-    "aiohttp": ("https://docs.aiohttp.org/en/stable", None),
-    "pillow": ("https://pillow.readthedocs.io/en/stable", None),
-    "numpy": ("https://numpy.org/doc/stable", None),
-    "torch": ("https://pytorch.org/docs/stable", None),
-    "psutil": ("https://psutil.readthedocs.io/en/stable", None),
-}
-
-autodoc_preserve_defaults = True
-autodoc_warningiserror = True
-
-navigation_with_keys = False

docs/source/developer_guide/overview.md

@@ -1,138 +0,0 @@
-(developer-guide)=
-
-# Contributing to FastVideo
-
-Thank you for your interest in contributing to FastVideo. We want to make the process as smooth for you as possible and this is a guide to help get you started!
-
-Our community is open to everyone and welcomes any contributions no matter how large or small.
-
-# Developer Environment:
-Do make sure you have CUDA 12.4 installed and supported. FastVideo currently only support Linux and CUDA GPUs, but we hope to support other platforms in the future.
-
-We recommend using a fresh Python 3.10 Conda environment to develop FastVideo:
-
-Install Miniconda:
-
-```
-wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
-bash Miniconda3-latest-Linux-x86_64.sh
-source ~/.bashrc
-```
-
-Create and activate a Conda environment for FastVideo:
-
-```
-conda create -n fastvideo python=3.10 -y
-conda activate fastvideo
-```
-
-Clone the FastVideo repository and go to the FastVideo directory:
-
-```
-git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
-
-```
-
-Now you can install FastVideo and setup git hooks for running linting. By using `pre-commit`, the linters will run and have to pass before you'll be able to make a commit.
-
-```bash
-pip install -e .[dev]
-
-# Can also install flash-attn (optional)
-pip install flash-attn==2.7.0.post2 --no-build-isolation 
-
-# Linting, formatting and static type checking
-pre-commit install --hook-type pre-commit --hook-type commit-msg
-
-# You can manually run pre-commit with
-pre-commit run --all-files
-
-# Unit tests
-pytest tests/
-```
-
----
-## 🐳 Using the FastVideo Docker Image
-
-If you prefer a containerized development environment or want to avoid managing dependencies manually, you can use our prebuilt Docker image:
-
-**Image:** [`ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest`](https://ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev)
-
-### Starting the container
-
-```bash
-docker run --gpus all -it ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
-```
-
-This will:
-
-- Start the container with GPU access  
-- Drop you into a shell with the `fastvideo-dev` Conda environment preconfigured
-
-### Using the container
-
-```bash
-# Conda environment should already be active
-# FastVideo package installed in editable mode
-
-# Pull the latest changes from remote
-cd /FastVideo
-git pull
-
-# Run linters and tests
-pre-commit run --all-files
-pytest tests/
-```
-
----
-
-## 📦 Developing FastVideo on RunPod
-
-You can easily use the FastVideo Docker image as a custom container on [RunPod](https://www.runpod.io) for development or experimentation.
-
-### Creating a new pod
-
-Choose a GPU that supports CUDA 12.4
-
-![RunPod CUDA selection](../_static/images/runpod_cuda.png)
-
-When creating your pod template, use this image:
-
-```
-ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:latest
-```
-
-Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.io/pods/configuration/use-ssh)):
-
-```bash
-bash -c "apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-server -y;mkdir -p ~/.ssh;cd $_;chmod 700 ~/.ssh;echo \"$PUBLIC_KEY\" >> authorized_keys;chmod 700 authorized_keys;service ssh start;sleep infinity"
-```
-
-![RunPod template configuration](../_static/images/runpod_template.png)
-
-After deploying, the pod will take a few minutes to pull the image and start the SSH service.
-
-![RunPod ssh](../_static/images/runpod_ssh.png)
-
-### Working with the pod
-
-After SSH'ing into your pod, you'll find the `fastvideo-dev` Conda environment already activated.
-
-To pull in the latest changes from the GitHub repo:
-
-```bash
-cd /FastVideo
-git pull
-```
-
-`If you have a persistent volume and want to keep your code changes, you can move /FastVideo to /workspace/FastVideo, or simply clone the repository there.`
-
-Run your development workflows as usual:
-
-```bash
-# Run linters
-pre-commit run --all-files
-
-# Run tests
-pytest tests/
-```

docs/source/generate_examples.py

@@ -1,249 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/docs/source/generate_examples.py
-
-import itertools
-import re
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Optional
-
-ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
-ROOT_DIR_RELATIVE = '../../../..'
-EXAMPLE_DIR = ROOT_DIR / "fastvideo/v1/examples"
-EXAMPLE_DOC_DIR = ROOT_DIR / "docs/source/getting_started/examples"
-
-
-def fix_case(text: str) -> str:
-    subs = {
-        "api": "API",
-        "cli": "CLI",
-        "cpu": "CPU",
-        "llm": "LLM",
-        "tpu": "TPU",
-        "aqlm": "AQLM",
-        "gguf": "GGUF",
-        "lora": "LoRA",
-        "rlhf": "RLHF",
-        "vllm": "vLLM",
-        "openai": "OpenAI",
-        "multilora": "MultiLoRA",
-        "mlpspeculator": "MLPSpeculator",
-        r"fp\d+": lambda x: x.group(0).upper(),  # e.g. fp16, fp32
-        r"int\d+": lambda x: x.group(0).upper(),  # e.g. int8, int16
-    }
-    for pattern, repl in subs.items():
-        text = re.sub(rf'\b{pattern}\b', repl, text,
-                      flags=re.IGNORECASE)  # type: ignore[call-overload]
-    return text
-
-
-@dataclass
-class Index:
-    """
-    Index class to generate a structured document index.
-
-    Attributes:
-        path (Path): The path save the index file to.
-        title (str): The title of the index.
-        description (str): A brief description of the index.
-        caption (str): An optional caption for the table of contents.
-        maxdepth (int): The maximum depth of the table of contents. Defaults to 1.
-        documents (list[str]): A list of document paths to include in the index. Defaults to an empty list.
-
-    Methods:
-        generate() -> str:
-            Generates the index content as a string in the specified format.
-    """ # noqa: E501
-    path: Path
-    title: str
-    description: str
-    caption: str
-    maxdepth: int = 1
-    documents: list[str] = field(default_factory=list)
-
-    def generate(self) -> str:
-        content = f"# {self.title}\n\n{self.description}\n\n"
-        content += ":::{toctree}\n"
-        content += f":caption: {self.caption}\n:maxdepth: {self.maxdepth}\n"
-        content += "\n".join(self.documents) + "\n:::\n"
-        return content
-
-
-@dataclass
-class Example:
-    """
-    Example class for generating documentation content from a given path.
-
-    Attributes:
-        path (Path): The path to the main directory or file.
-        category (str): The category of the document.
-        main_file (Path): The main file in the directory.
-        other_files (list[Path]): list of other files in the directory.
-        title (str): The title of the document.
-
-    Methods:
-        __post_init__(): Initializes the main_file, other_files, and title attributes.
-        determine_main_file() -> Path: Determines the main file in the given path.
-        determine_other_files() -> list[Path]: Determines other files in the directory excluding the main file.
-        determine_title() -> str: Determines the title of the document.
-        generate() -> str: Generates the documentation content.
-    """ # noqa: E501
-    path: Path
-    category: Optional[str] = None
-    main_file: Path = field(init=False)
-    other_files: list[Path] = field(init=False)
-    title: str = field(init=False)
-
-    def __post_init__(self):
-        self.main_file = self.determine_main_file()
-        self.other_files = self.determine_other_files()
-        self.title = self.determine_title()
-
-    def determine_main_file(self) -> Path:
-        """
-        Determines the main file in the given path.
-        If the path is a file, it returns the path itself. Otherwise, it searches
-        for Markdown files (*.md) in the directory and returns the first one found.
-        Returns:
-            Path: The main file path, either the original path if it's a file or the first
-            Markdown file found in the directory.
-        Raises:
-            IndexError: If no Markdown files are found in the directory.
-        """ # noqa: E501
-        return self.path if self.path.is_file() else list(
-            self.path.glob("*.md")).pop()
-
-    def determine_other_files(self) -> list[Path]:
-        """
-        Determine other files in the directory excluding the main file.
-
-        This method checks if the given path is a file. If it is, it returns an empty list.
-        Otherwise, it recursively searches through the directory and returns a list of all
-        files that are not the main file.
-
-        Returns:
-            list[Path]: A list of Path objects representing the other files in the directory.
-        """ # noqa: E501
-        if self.path.is_file():
-            return []
-        is_other_file = lambda file: file.is_file() and file != self.main_file
-        return [file for file in self.path.rglob("*")
-                if is_other_file(file)]  # type: ignore[no-untyped-call]
-
-    def determine_title(self) -> str:
-        return fix_case(self.path.stem.replace("_", " ").title())
-
-    def generate(self) -> str:
-        # Convert the path to a relative path from __file__
-        make_relative = lambda path: ROOT_DIR_RELATIVE / path.relative_to(
-            ROOT_DIR)
-
-        content = f"Source <gh-file:{self.path.relative_to(ROOT_DIR)}>.\n\n"
-        include = "include" if self.main_file.suffix == ".md" else \
-            "literalinclude"
-        if include == "literalinclude":
-            content += f"# {self.title}\n\n"
-        content += f":::{{{include}}} {make_relative(self.main_file)}\n"  # type: ignore[no-untyped-call]
-        if include == "literalinclude":
-            content += f":language: {self.main_file.suffix[1:]}\n"
-        content += ":::\n\n"
-
-        if not self.other_files:
-            return content
-
-        content += "## Example materials\n\n"
-        for file in sorted(self.other_files):
-            include = "include" if file.suffix == ".md" else "literalinclude"
-            content += f":::{{admonition}} {file.relative_to(self.path)}\n"
-            content += ":class: dropdown\n\n"
-            content += f":::{{{include}}} {make_relative(file)}\n:::\n"  # type: ignore[no-untyped-call]
-            content += ":::\n\n"
-
-        return content
-
-
-def generate_examples():
-    # Create the EXAMPLE_DOC_DIR if it doesn't exist
-    if not EXAMPLE_DOC_DIR.exists():
-        EXAMPLE_DOC_DIR.mkdir(parents=True)
-
-    # Create empty indices
-    examples_index = Index(
-        path=EXAMPLE_DOC_DIR / "examples_index.md",
-        title="Examples",
-        description=
-        "A collection of examples demonstrating usage of FastVideo.\nAll documented examples are autogenerated using <gh-file:docs/source/generate_examples.py> from examples found in <gh-file:examples>.",  # noqa: E501
-        caption="Examples",
-        maxdepth=2)
-    # Category indices stored in reverse order because they are inserted into
-    # examples_index.documents at index 0 in order
-    category_indices = {
-        # "other":
-        # Index(
-        #     path=EXAMPLE_DOC_DIR / "examples_other_index.md",
-        #     title="Other",
-        #     description=
-        #     "Other examples that don't strongly fit into the online or offline serving categories.",  # noqa: E501
-        #     caption="Examples",
-        # ),
-        # "online_serving":
-        # Index(
-        #     path=EXAMPLE_DOC_DIR / "examples_online_serving_index.md",
-        #     title="Online Serving",
-        #     description=
-        #     "Online serving examples demonstrate how to use FastVideo in an online setting, where the model is queried for predictions in real-time.",  # noqa: E501
-        #     caption="Examples",
-        # ),
-        "inference":
-        Index(
-            path=EXAMPLE_DOC_DIR / "examples_inference_index.md",
-            title="Inference",
-            description=
-            "Inference examples demonstrate how to use FastVideo in an offline setting, where the model is queried for predictions in batches. We recommend starting with <project:basic.md>.",  # noqa: E501
-            caption="Examples",
-        ),
-    }
-
-    examples = []
-    glob_patterns = ["*.py", "*.md", "*.sh"]
-    # Find categorised examples
-    for category in category_indices:
-        print(category)
-        category_dir = EXAMPLE_DIR / category
-        globs = [category_dir.glob(pattern) for pattern in glob_patterns]
-        for path in itertools.chain(*globs):
-            examples.append(Example(path, category))
-        # Find examples in subdirectories
-        for path in category_dir.glob("*/*.md"):
-            examples.append(Example(path.parent, category))
-    # Find uncategorised examples
-    globs = [EXAMPLE_DIR.glob(pattern) for pattern in glob_patterns]
-    for path in itertools.chain(*globs):
-        examples.append(Example(path))
-    # Find examples in subdirectories
-    for path in EXAMPLE_DIR.glob("*/*.md"):
-        # Skip categorised examples
-        if path.parent.name in category_indices:
-            continue
-        examples.append(Example(path.parent))
-
-    # Generate the example documentation
-    for example in sorted(examples, key=lambda e: e.path.stem):
-        print(example)
-        doc_path = EXAMPLE_DOC_DIR / f"{example.path.stem}.md"
-        with open(doc_path, "w+") as f:
-            f.write(example.generate())
-        # Add the example to the appropriate index
-        assert example.category is not None
-        index = category_indices.get(example.category, examples_index)
-        index.documents.append(example.path.stem)
-
-    # Generate the index files
-    for category_index in category_indices.values():
-        if category_index.documents:
-            examples_index.documents.insert(0, category_index.path.name)
-            with open(category_index.path, "w+") as f:
-                f.write(category_index.generate())
-
-    with open(examples_index.path, "w+") as f:
-        f.write(examples_index.generate())

docs/source/getting_started/installation.md

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

docs/source/index.md

@@ -1,89 +0,0 @@
-# Welcome to FastVideo
-
-:::{figure} ../../assets/logo.jpg
-:align: center
-:alt: FastVideo
-:class: no-scaled-link
-:width: 60%
-:::
-
-:::{raw} html
-<p style="text-align:center">
-<strong>FastVideo is a lightweight framework for accelerating large video diffusion models.
-</strong>
-</p>
-
-<p style="text-align:center">
-<script async defer src="https://buttons.github.io/buttons.js"></script>
-<a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/" data-show-count="true" data-size="large" aria-label="Star">Star</a>
-<a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/subscription" data-icon="octicon-eye" data-size="large" aria-label="Watch">Watch</a>
-<a class="github-button" href="https://github.com/hao-ai-lab/FastVideo/fork" data-icon="octicon-repo-forked" data-size="large" aria-label="Fork">Fork</a>
-</p>
-:::
-
-FastVideo is a lightweight framework for accelerating large video diffusion models developed by the [Hao AI Lab](https://hao-ai-lab.github.io/).
-
-<div style="text-align: center;">
-  <video controls width="800">
-    <source src="https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1" type="video/mp4">
-    Your browser does not support the video tag.
-  </video>
-</div>
-
-FastVideo currently offers: (with more to come)
-
-- [NEW!] [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
-- FastHunyuan and FastMochi: consistency distilled video diffusion models for 8x inference speedup.
-- First open distillation recipes for video DiT, based on [PCM](https://github.com/G-U-N/Phased-Consistency-Model).
-- Support distilling/finetuning/inferencing state-of-the-art open video DiTs: 1. Mochi 2. Hunyuan.
-- Scalable training with FSDP, sequence parallelism, and selective activation checkpointing, with near linear scaling to 64 GPUs.
-- Memory efficient finetuning with LoRA, precomputed latent, and precomputed text embeddings.
-
-Dev in progress and highly experimental.
-
-## Documentation
-
-% How to start using vLLM?
-
-:::{toctree}
-:caption: Getting Started
-:maxdepth: 1
-
-getting_started/installation
-getting_started/examples/examples_index
-:::
-
-% What is STA Kernel?
-
-:::{toctree}
-:caption: Sliding Tile Attention
-:maxdepth: 1
-
-sliding_tile_attention/installation
-sliding_tile_attention/usage
-sliding_tile_attention/test
-sliding_tile_attention/demo
-:::
-
-:::{toctree}
-:caption: Inference
-:maxdepth: 1
-
-inference/wanvideo
-inference/stepvideo
-inference/hunyuanvideo
-inference/fasthunyuan
-inference/fastmochi
-:::
-
-:::{toctree}
-:caption: Developer Guide
-:maxdepth: 1
-
-developer_guide/overview
-:::
-
-## Indices and tables
-
-- {ref}`genindex`
-- {ref}`modindex`

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/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/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.

docs/source/sliding_tile_attention/demo.md

@@ -1,11 +0,0 @@
-(sta-demo)=
-
-# Demo
-There is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
-
-<div style="text-align: center;">
-  <video controls width="800">
-    <source src="https://github.com/user-attachments/assets/f3b6dd79-7b43-4b60-a0fa-3d6495ec5747" type="video/mp4">
-    Your browser does not support the video tag.
-  </video>
-</div>

docs/source/sliding_tile_attention/installation.md

@@ -1,25 +0,0 @@
-(sta-installation)=
-
-# Installation
-We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
-First, install C++20 for ThunderKittens:
-
-```bash
-sudo apt update
-sudo apt install gcc-11 g++-11
-
-sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-sudo apt update
-sudo apt install clang-11
-```
-
-Install STA:
-
-```bash
-export CUDA_HOME=/usr/local/cuda-12.4
-export PATH=${CUDA_HOME}/bin:${PATH} 
-export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-git submodule update --init --recursive
-python setup.py install
-```

docs/source/sliding_tile_attention/test.md

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

docs/source/sliding_tile_attention/usage.md

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

env_setup.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+# install torch
+pip install torch==2.5.0 torchvision --index-url https://download.pytorch.org/whl/cu124
+
+# install FA2 and diffusers
+pip install packaging ninja && pip install flash-attn==2.7.0.post2 --no-build-isolation 
+
+pip install -r requirements-lint.txt
+
+# install fastvideo
+pip install -e .

examples/inference/basic/README.md

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

examples/inference/basic/basic.py

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

fastvideo/__init__.py

@@ -1,3 +0,0 @@
-from fastvideo.v1.entrypoints.video_generator import VideoGenerator
-
-__all__ = ["VideoGenerator"]

fastvideo/models/hunyuan/diffusion/pipelines/pipeline_hunyuan_video.py

@@ -546,6 +546,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                                              MultiPipelineCallbacks, ]] = None,
         callback_on_step_end_tensor_inputs: List[str] = ["latents"],
         vae_ver: str = "88-4c-sd",
+        use_cpu_offload: bool = False,
         enable_tiling: bool = False,
         enable_vae_sp: bool = False,
         n_tokens: Optional[int] = None,
@@ -904,14 +905,22 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                 latents = (latents / self.vae.config.scaling_factor + self.vae.config.shift_factor)
             else:
                 latents = latents / self.vae.config.scaling_factor
+            
+            if use_cpu_offload:
+                print("cpu offloaded")
+                self.transformer = self.transformer.to('cpu')
 
             with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
                 if enable_tiling:
+                    print("tiling enabled")
                     self.vae.enable_tiling()
                 if enable_vae_sp:
                     self.vae.enable_parallel()
                 image = self.vae.decode(latents, return_dict=False, generator=generator)[0]
-
+            
+            if use_cpu_offload:
+                self.transformer = self.transformer.to(device)
+            
             if expand_temporal_dim or image.shape[2] == 1:
                 image = image.squeeze(2)
 

fastvideo/models/hunyuan/inference.py

@@ -15,7 +15,7 @@ from fastvideo.models.hunyuan.text_encoder import TextEncoder
 from fastvideo.models.hunyuan.utils.data_utils import align_to
 from fastvideo.models.hunyuan.vae import load_vae
 from fastvideo.utils.parallel_states import nccl_info
-
+from fastvideo.models.hunyuan.modules.fp8 import convert_fp8_linear
 
 class Inference(object):
 
@@ -76,7 +76,10 @@ class Inference(object):
 
         # =========================== Build main model ===========================
         logger.info("Building model...")
-        factor_kwargs = {"device": device, "dtype": PRECISION_TO_TYPE[args.precision]}
+        if args.use_cpu_offload:
+            factor_kwargs = {"device": 'cpu', "dtype": PRECISION_TO_TYPE[args.precision]}
+        else:
+            factor_kwargs = {"device": device, "dtype": PRECISION_TO_TYPE[args.precision]}
         in_channels = args.latent_channels
         out_channels = args.latent_channels
 
@@ -86,6 +89,11 @@ class Inference(object):
             out_channels=out_channels,
             factor_kwargs=factor_kwargs,
         )
+        
+        if args.use_fp8:
+            print("loading fp8 model")
+            convert_fp8_linear(model, args.dit_weight, original_dtype=PRECISION_TO_TYPE[args.precision])
+        
         model = model.to(device)
         model = Inference.load_state_dict(args, model, pretrained_model_path)
         if args.enable_torch_compile:
@@ -453,6 +461,7 @@ class HunyuanVideoSampler(Inference):
         # Pipeline inference
         # ========================================================================
         start_time = time.time()
+        torch._dynamo.config.cache_size_limit = 125
         samples = self.pipeline(
             prompt=prompt,
             height=target_height,
@@ -469,6 +478,7 @@ class HunyuanVideoSampler(Inference):
             data_type="video" if target_video_length > 1 else "image",
             is_progress_bar=True,
             vae_ver=self.args.vae,
+            use_cpu_offload=self.args.use_cpu_offload,
             enable_tiling=self.args.vae_tiling,
             enable_vae_sp=self.args.vae_sp,
             mask_strategy=mask_strategy,

fastvideo/models/hunyuan/modules/attenion.py

@@ -34,11 +34,11 @@ def attention(
     return out
 
 
-def tile(x, sp_size):
-    x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80)
+def tile(x, sp_size, t_size):
+    x = rearrange(x, "b (sp t h w) head d -> b (t sp h w) head d", sp=sp_size, t=(t_size // sp_size), h=48, w=80)
     return rearrange(x,
                      "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
-                     n_t=5,
+                     n_t= (t_size // 6),
                      n_h=6,
                      n_w=10,
                      ts_t=6,
@@ -46,16 +46,16 @@ def tile(x, sp_size):
                      ts_w=8)
 
 
-def untile(x, sp_size):
+def untile(x, sp_size, t_size):
     x = rearrange(x,
                   "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
-                  n_t=5,
+                  n_t=(t_size // 6),
                   n_h=6,
                   n_w=10,
                   ts_t=6,
                   ts_h=8,
                   ts_w=8)
-    return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=30 // sp_size, h=48, w=80)
+    return rearrange(x, "b (t sp h w) head d -> b (sp t h w) head d", sp=sp_size, t=(t_size // sp_size), h=48, w=80)
 
 
 def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=None):
@@ -83,9 +83,10 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=
     encoder_sequence_length = encoder_query.size(1)
 
     if mask_strategy[0] is not None:
-        query = torch.cat([tile(query, nccl_info.sp_size), encoder_query], dim=1).transpose(1, 2)
-        key = torch.cat([tile(key, nccl_info.sp_size), encoder_key], dim=1).transpose(1, 2)
-        value = torch.cat([tile(value, nccl_info.sp_size), encoder_value], dim=1).transpose(1, 2)
+        t_size = int(query.shape[1] / (8 * 6 * 8 * 10)) # ts_h n_h ts_w n_w
+        query = torch.cat([tile(query, nccl_info.sp_size, t_size), encoder_query], dim=1).transpose(1, 2)
+        key = torch.cat([tile(key, nccl_info.sp_size, t_size), encoder_key], dim=1).transpose(1, 2)
+        value = torch.cat([tile(value, nccl_info.sp_size, t_size), encoder_value], dim=1).transpose(1, 2)
 
         head_num = query.size(1)
         current_rank = nccl_info.rank_within_group
@@ -107,7 +108,7 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=
                                                                           dim=1)
 
     if mask_strategy[0] is not None:
-        hidden_states = untile(hidden_states, nccl_info.sp_size)
+        hidden_states = untile(hidden_states, nccl_info.sp_size, t_size)
 
     if get_sequence_parallel_state():
         hidden_states = all_to_all_4D(hidden_states, scatter_dim=1, gather_dim=2)

fastvideo/models/hunyuan/modules/fp8.py

@@ -0,0 +1,100 @@
+import os
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+def get_fp_maxval(bits=8, mantissa_bit=3, sign_bits=1):
+    _bits = torch.tensor(bits)
+    _mantissa_bit = torch.tensor(mantissa_bit)
+    _sign_bits = torch.tensor(sign_bits)
+    M = torch.clamp(torch.round(_mantissa_bit), 1, _bits - _sign_bits)
+    E = _bits - _sign_bits - M
+    bias = 2 ** (E - 1) - 1
+    mantissa = 1
+    for i in range(mantissa_bit - 1):
+        mantissa += 1 / (2 ** (i+1))
+    maxval = mantissa * 2 ** (2**E - 1 - bias)
+    return maxval
+
+def quantize_to_fp8(x, bits=8, mantissa_bit=3, sign_bits=1):
+    """
+    Default is E4M3.
+    """
+    bits = torch.tensor(bits)
+    mantissa_bit = torch.tensor(mantissa_bit)
+    sign_bits = torch.tensor(sign_bits)
+    M = torch.clamp(torch.round(mantissa_bit), 1, bits - sign_bits)
+    E = bits - sign_bits - M
+    bias = 2 ** (E - 1) - 1
+    mantissa = 1
+    for i in range(mantissa_bit - 1):
+        mantissa += 1 / (2 ** (i+1))
+    maxval = mantissa * 2 ** (2**E - 1 - bias)
+    minval = - maxval
+    minval = - maxval if sign_bits == 1 else torch.zeros_like(maxval)
+    input_clamp = torch.min(torch.max(x, minval), maxval)
+    log_scales = torch.clamp((torch.floor(torch.log2(torch.abs(input_clamp)) + bias)).detach(), 1.0)
+    log_scales = 2.0 ** (log_scales - M - bias.type(x.dtype))
+    # dequant
+    qdq_out = torch.round(input_clamp / log_scales) * log_scales
+    return qdq_out, log_scales
+
+def fp8_tensor_quant(x, scale, bits=8, mantissa_bit=3, sign_bits=1):
+    for i in range(len(x.shape) - 1):
+        scale = scale.unsqueeze(-1)
+    new_x = x / scale
+    quant_dequant_x, log_scales = quantize_to_fp8(new_x, bits=bits, mantissa_bit=mantissa_bit, sign_bits=sign_bits)
+    return quant_dequant_x, scale, log_scales
+
+def fp8_activation_dequant(qdq_out, scale, dtype):
+    qdq_out = qdq_out.type(dtype)
+    quant_dequant_x = qdq_out * scale.to(dtype)
+    return quant_dequant_x
+
+def fp8_linear_forward(cls, original_dtype, input):
+    weight_dtype = cls.weight.dtype
+    #####
+    if cls.weight.dtype != torch.float8_e4m3fn:
+        maxval = get_fp_maxval()
+        scale = torch.max(torch.abs(cls.weight.flatten())) / maxval
+        linear_weight, scale, log_scales = fp8_tensor_quant(cls.weight, scale)
+        linear_weight = linear_weight.to(torch.float8_e4m3fn)
+        weight_dtype = linear_weight.dtype
+    else:
+        scale = cls.fp8_scale.to(cls.weight.device)
+        linear_weight = cls.weight
+    #####
+
+    if weight_dtype == torch.float8_e4m3fn and cls.weight.sum() != 0:
+        if True or len(input.shape) == 3:
+            cls_dequant = fp8_activation_dequant(linear_weight, scale, original_dtype)
+            if cls.bias != None:
+                output = F.linear(input, cls_dequant, cls.bias)
+            else:
+                output = F.linear(input, cls_dequant)
+            return output
+        else:
+            return cls.original_forward(input.to(original_dtype))
+    else:
+        return cls.original_forward(input)
+
+def convert_fp8_linear(module, dit_weight_path, original_dtype, params_to_keep={}):
+    setattr(module, "fp8_matmul_enabled", True)
+
+    # loading fp8 mapping file
+    fp8_map_path = dit_weight_path.replace('.pt', '_map.pt')
+    if os.path.exists(fp8_map_path):
+        fp8_map = torch.load(fp8_map_path, map_location=lambda storage, loc: storage)
+    else:
+        raise ValueError(f"Invalid fp8_map path: {fp8_map_path}.")
+
+    fp8_layers = []
+    for key, layer in module.named_modules():
+        if isinstance(layer, nn.Linear) and ('double_blocks' in key or 'single_blocks' in key):
+            fp8_layers.append(key)
+            original_forward = layer.forward
+            layer.weight = torch.nn.Parameter(layer.weight.to(torch.float8_e4m3fn))
+            setattr(layer, "fp8_scale", fp8_map[key].to(dtype=original_dtype))
+            setattr(layer, "original_forward", original_forward)
+            setattr(layer, "forward", lambda input, m=layer: fp8_linear_forward(m, original_dtype, input))

fastvideo/sample/sample_t2v_hunyuan_STA.py

@@ -85,7 +85,9 @@ def teacache_forward(
             img_mod2_gate,
         ) = self.double_blocks[0].img_mod(vec_).chunk(6, dim=-1)
         normed_inp = self.double_blocks[0].img_norm1(inp)
-        modulated_inp = modulate(normed_inp, shift=img_mod1_shift, scale=img_mod1_scale)
+        modulated_inp = modulate(normed_inp, shift=img_mod1_shift, scale=img_mod1_scale).to("cpu")
+        del inp, vec_, img_mod1_shift, img_mod1_scale, normed_inp
+
         if self.cnt == 0 or self.cnt == self.num_steps - 1:
             should_calc = True
             self.accumulated_rel_l1_distance = 0
@@ -106,9 +108,10 @@ def teacache_forward(
             self.cnt = 0
     if self.enable_teacache:
         if not should_calc:
-            img += self.previous_residual
+            img += self.previous_residual.to(img.device)
+            self.previous_residual = self.previous_residual.to(img.device)
         else:
-            ori_img = img.clone()
+            ori_img = img.clone().to("cpu")
             # --------------------- Pass through DiT blocks ------------------------
             for index, block in enumerate(self.double_blocks):
                 double_block_args = [img, txt, vec, freqs_cis, text_mask, mask_strategy[index]]
@@ -133,7 +136,8 @@ def teacache_forward(
                         features_list.append(x[:, :img_seq_len, ...])
 
             img = x[:, :img_seq_len, ...]
-            self.previous_residual = img - ori_img
+            self.previous_residual = (img.clone().to("cpu") - ori_img).to("cpu")
+            del ori_img
     else:
         # --------------------- Pass through DiT blocks ------------------------
         for index, block in enumerate(self.double_blocks):
@@ -301,6 +305,11 @@ if __name__ == "__main__":
         action="store_true",
         help="Use CPU offload for the model load.",
     )
+    parser.add_argument(
+        "--use-fp8",
+        action="store_true",
+        help="Use FP8 Quantization for the model load.",
+    )
     parser.add_argument(
         "--dit-weight",
         type=str,
@@ -373,6 +382,7 @@ if __name__ == "__main__":
     parser.add_argument("--text-states-dim-2", type=int, default=768)
     parser.add_argument("--tokenizer-2", type=str, default="clipL")
     parser.add_argument("--text-len-2", type=int, default=77)
+    parser.add_argument("--vae_tiling", action='store_true')
     parser.add_argument("--skip_time_steps", type=int, default=10)
     parser.add_argument(
         "--mask_strategy_selected",

fastvideo/sample/sample_t2v_stepvideo_STA.py

@@ -341,7 +341,7 @@ if __name__ == "__main__":
     )
 
     # TeaCache
-    pipeline.transformer.__class__.enable_teacache = True
+    pipeline.transformer.__class__.enable_teacache = args.enable_teacache
     pipeline.transformer.__class__.cnt = 0
     pipeline.transformer.__class__.num_steps = args.infer_steps
     pipeline.transformer.__class__.rel_l1_thresh = args.rel_l1_thresh  # 0.1 for 1.6x speedup, 0.15 for 2.1x speedup

fastvideo/v1/attention/__init__.py

@@ -1,17 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
-                                                      AttentionMetadata,
-                                                      AttentionMetadataBuilder)
-from fastvideo.v1.attention.layer import DistributedAttention, LocalAttention
-from fastvideo.v1.attention.selector import get_attn_backend
-
-__all__ = [
-    "DistributedAttention",
-    "LocalAttention",
-    "AttentionBackend",
-    "AttentionMetadata",
-    "AttentionMetadataBuilder",
-    # "AttentionState",
-    "get_attn_backend",
-]

fastvideo/v1/attention/backends/abstract.py

@@ -1,246 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/backends/abstract.py
-
-from abc import ABC, abstractmethod
-from dataclasses import dataclass, fields
-from typing import (TYPE_CHECKING, Any, Dict, Generic, Optional, Protocol, Set,
-                    Type, TypeVar)
-
-if TYPE_CHECKING:
-    from fastvideo.v1.fastvideo_args import FastVideoArgs
-    from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-
-import torch
-
-
-class AttentionBackend(ABC):
-    """Abstract class for attention backends."""
-    # For some attention backends, we allocate an output tensor before
-    # calling the custom op. When piecewise cudagraph is enabled, this
-    # makes sure the output tensor is allocated inside the cudagraph.
-    accept_output_buffer: bool = False
-
-    @staticmethod
-    @abstractmethod
-    def get_name() -> str:
-        raise NotImplementedError
-
-    @staticmethod
-    @abstractmethod
-    def get_impl_cls() -> Type["AttentionImpl"]:
-        raise NotImplementedError
-
-    @staticmethod
-    @abstractmethod
-    def get_metadata_cls() -> Type["AttentionMetadata"]:
-        raise NotImplementedError
-
-    # @staticmethod
-    # @abstractmethod
-    # def get_state_cls() -> Type["AttentionState"]:
-    #     raise NotImplementedError
-
-    # @classmethod
-    # def make_metadata(cls, *args, **kwargs) -> "AttentionMetadata":
-    #     return cls.get_metadata_cls()(*args, **kwargs)
-
-    @staticmethod
-    @abstractmethod
-    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
-        raise NotImplementedError
-
-
-@dataclass
-class AttentionMetadata:
-    """Attention metadata for prefill and decode batched together."""
-    # 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]:
-        """Similar to dataclasses.asdict, but avoids deepcopying."""
-        if skip_fields is None:
-            skip_fields = set()
-        # Note that if we add dataclasses as fields, they will need
-        # similar handling.
-        return {
-            field.name: getattr(self, field.name)
-            for field in fields(self) if field.name not in skip_fields
-        }
-
-
-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."""
-
-    @abstractmethod
-    def __init__(self) -> None:
-        """Create the builder, remember some configuration and parameters."""
-        raise NotImplementedError
-
-    @abstractmethod
-    def prepare(self) -> None:
-        """Prepare for one batch."""
-        raise NotImplementedError
-
-    @abstractmethod
-    def build(
-        self,
-        current_timestep: int,
-        forward_batch: "ForwardBatch",
-        fastvideo_args: "FastVideoArgs",
-    ) -> T:
-        """Build attention metadata with on-device tensors."""
-        raise NotImplementedError
-
-
-class AttentionLayer(Protocol):
-
-    _k_scale: torch.Tensor
-    _v_scale: torch.Tensor
-    _k_scale_float: float
-    _v_scale_float: float
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        kv_cache: torch.Tensor,
-        attn_metadata: AttentionMetadata,
-    ) -> torch.Tensor:
-        ...
-
-
-class AttentionImpl(ABC, Generic[T]):
-
-    @abstractmethod
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        softmax_scale: float,
-        causal: bool = False,
-        num_kv_heads: Optional[int] = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        raise NotImplementedError
-
-    def preprocess_qkv(self, qkv: torch.Tensor,
-                       attn_metadata: T) -> torch.Tensor:
-        """Preprocess QKV tensor before performing attention operation.
-
-        Default implementation returns the tensor unchanged.
-        Subclasses can override this to implement custom preprocessing
-        like reshaping, tiling, scaling, or other transformations.
-
-        Called AFTER all_to_all for distributed attention
-        
-        Args:
-            qkv: The query-key-value tensor
-            attn_metadata: Metadata for the attention operation
-            
-        Returns:
-            Processed QKV tensor
-        """
-        return qkv
-
-    def postprocess_output(
-        self,
-        output: torch.Tensor,
-        attn_metadata: T,
-    ) -> torch.Tensor:
-        """Postprocess the output tensor after the attention operation.
-
-        Default implementation returns the tensor unchanged.
-        Subclasses can override this to implement custom postprocessing
-        like untiling, scaling, or other transformations.
-
-        Called BEFORE all_to_all for distributed attention
-
-        Args:
-            output: The output tensor from the attention operation
-            attn_metadata: Metadata for the attention operation
-
-        Returns:
-            Postprocessed output tensor
-        """
-
-        return output
-
-    @abstractmethod
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: T,
-    ) -> torch.Tensor:
-        raise NotImplementedError

fastvideo/v1/attention/backends/flash_attn.py

@@ -1,79 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from typing import List, Optional, Type
-
-import torch
-from flash_attn import flash_attn_func as flash_attn_2_func
-
-try:
-    from flash_attn_interface import flash_attn_func as flash_attn_3_func
-
-    # flash_attn 3 has slightly different API: it returns lse by default
-    flash_attn_func = lambda q, k, v, softmax_scale, causal: flash_attn_3_func(
-        q, k, v, softmax_scale, causal)[0]
-except ImportError:
-    flash_attn_func = flash_attn_2_func
-
-from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
-                                                      AttentionImpl,
-                                                      AttentionMetadata,
-                                                      AttentionMetadataBuilder)
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class FlashAttentionBackend(AttentionBackend):
-
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
-        return [32, 64, 96, 128, 160, 192, 224, 256]
-
-    @staticmethod
-    def get_name() -> str:
-        return "FLASH_ATTN"
-
-    @staticmethod
-    def get_impl_cls() -> Type["FlashAttentionImpl"]:
-        return FlashAttentionImpl
-
-    @staticmethod
-    def get_metadata_cls() -> Type["AttentionMetadata"]:
-        raise NotImplementedError
-
-    @staticmethod
-    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
-        raise NotImplementedError
-
-
-class FlashAttentionImpl(AttentionImpl):
-
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        causal: bool,
-        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata,
-    ):
-        output = flash_attn_func(
-            query,  # type: ignore[no-untyped-call]
-            key,
-            value,
-            softmax_scale=self.softmax_scale,
-            causal=self.causal)
-        return output

fastvideo/v1/attention/backends/sdpa.py

@@ -1,73 +0,0 @@
-from typing import List, Optional, Type
-
-import torch
-
-from fastvideo.v1.attention.backends.abstract import (
-    AttentionBackend)  # FlashAttentionMetadata,
-from fastvideo.v1.attention.backends.abstract import (AttentionImpl,
-                                                      AttentionMetadata)
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class SDPABackend(AttentionBackend):
-
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
-        return [32, 64, 96, 128, 160, 192, 224, 256]
-
-    @staticmethod
-    def get_name() -> str:
-        return "SDPA"
-
-    @staticmethod
-    def get_impl_cls() -> Type["SDPAImpl"]:
-        return SDPAImpl
-
-    # @staticmethod
-    # def get_metadata_cls() -> Type["AttentionMetadata"]:
-    #     return FlashAttentionMetadata
-
-
-class SDPAImpl(AttentionImpl):
-
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        causal: bool,
-        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        self.causal = causal
-        self.softmax_scale = softmax_scale
-        self.dropout = extra_impl_args.get("dropout_p", 0.0)
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        attn_metadata: AttentionMetadata,
-    ) -> torch.Tensor:
-        # transpose to bs, heads, seq_len, head_dim
-        query = query.transpose(1, 2)
-        key = key.transpose(1, 2)
-        value = value.transpose(1, 2)
-        attn_kwargs = {
-            "attn_mask": None,
-            "dropout_p": self.dropout,
-            "is_causal": self.causal,
-            "scale": self.softmax_scale
-        }
-        if query.shape[1] != key.shape[1]:
-            attn_kwargs["enable_gqa"] = True
-        output = torch.nn.functional.scaled_dot_product_attention(
-            query, key, value, **attn_kwargs)
-        output = output.transpose(1, 2)
-        return output

fastvideo/v1/attention/backends/sliding_tile_attn.py

@@ -1,241 +0,0 @@
-import json
-from dataclasses import dataclass
-from typing import List, Optional, Type
-
-import torch
-from einops import rearrange
-from st_attn import sliding_tile_attention
-
-import fastvideo.v1.envs as envs
-from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
-                                                      AttentionImpl,
-                                                      AttentionMetadata,
-                                                      AttentionMetadataBuilder)
-from fastvideo.v1.distributed import get_sp_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-
-logger = init_logger(__name__)
-
-
-# TODO(will-refactor): move this to a utils file
-def dict_to_3d_list(mask_strategy) -> List[List[List[Optional[torch.Tensor]]]]:
-    indices = [tuple(map(int, key.split('_'))) for key in mask_strategy]
-
-    max_timesteps_idx = max(
-        timesteps_idx for timesteps_idx, layer_idx, head_idx in indices) + 1
-    max_layer_idx = max(layer_idx
-                        for timesteps_idx, layer_idx, head_idx in indices) + 1
-    max_head_idx = max(head_idx
-                       for timesteps_idx, layer_idx, head_idx in indices) + 1
-
-    result = [[[None for _ in range(max_head_idx)]
-               for _ in range(max_layer_idx)] for _ in range(max_timesteps_idx)]
-
-    for key, value in mask_strategy.items():
-        timesteps_idx, layer_idx, head_idx = map(int, key.split('_'))
-        result[timesteps_idx][layer_idx][head_idx] = value
-
-    return result
-
-
-class RangeDict(dict):
-
-    def __getitem__(self, item):
-        for key in self.keys():
-            if isinstance(key, tuple):
-                low, high = key
-                if low <= item <= high:
-                    return super().__getitem__(key)
-            elif key == item:
-                return super().__getitem__(key)
-        raise KeyError(f"seq_len {item} not supported for STA")
-
-
-class SlidingTileAttentionBackend(AttentionBackend):
-
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
-        # TODO(will-refactor): check this
-        return [32, 64, 96, 128, 160, 192, 224, 256]
-
-    @staticmethod
-    def get_name() -> str:
-        return "SLIDING_TILE_ATTN"
-
-    @staticmethod
-    def get_impl_cls() -> Type["SlidingTileAttentionImpl"]:
-        return SlidingTileAttentionImpl
-
-    @staticmethod
-    def get_metadata_cls() -> Type["SlidingTileAttentionMetadata"]:
-        return SlidingTileAttentionMetadata
-
-    @staticmethod
-    def get_builder_cls() -> Type["SlidingTileAttentionMetadataBuilder"]:
-        return SlidingTileAttentionMetadataBuilder
-
-
-@dataclass
-class SlidingTileAttentionMetadata(AttentionMetadata):
-    current_timestep: int
-
-
-class SlidingTileAttentionMetadataBuilder(AttentionMetadataBuilder):
-
-    def __init__(self):
-        pass
-
-    def prepare(self):
-        pass
-
-    def build(
-        self,
-        current_timestep: int,
-        forward_batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> SlidingTileAttentionMetadata:
-
-        return SlidingTileAttentionMetadata(current_timestep=current_timestep, )
-
-
-class SlidingTileAttentionImpl(AttentionImpl):
-
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        causal: bool,
-        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        # TODO(will-refactor): for now this is the mask strategy, but maybe we should
-        # have a more general config for STA?
-        config_file = envs.FASTVIDEO_ATTENTION_CONFIG
-        if config_file is None:
-            raise ValueError("FASTVIDEO_ATTENTION_CONFIG is not set")
-
-        with open(config_file) as f:
-            mask_strategy = json.load(f)
-        mask_strategy = dict_to_3d_list(mask_strategy)
-
-        self.prefix = prefix
-        self.mask_strategy = mask_strategy
-        sp_group = get_sp_group()
-        self.sp_size = sp_group.world_size
-        # STA config
-        self.STA_base_tile_size = [6, 8, 8]
-        self.img_latent_shape_mapping = RangeDict({
-            (115200, 115456): '30x48x80',
-            82944: '36x48x48',
-            69120: '18x48x80',
-        })
-        self.full_window_mapping = {
-            '30x48x80': [5, 6, 10],
-            '36x48x48': [6, 6, 6],
-            '18x48x80': [3, 6, 10]
-        }
-
-    def tile(self, x: torch.Tensor) -> torch.Tensor:
-        x = rearrange(x,
-                      "b (sp t h w) head d -> b (t sp h w) head d",
-                      sp=self.sp_size,
-                      t=self.img_latent_shape_int[0] // self.sp_size,
-                      h=self.img_latent_shape_int[1],
-                      w=self.img_latent_shape_int[2])
-        return rearrange(
-            x,
-            "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
-            n_t=self.full_window_size[0],
-            n_h=self.full_window_size[1],
-            n_w=self.full_window_size[2],
-            ts_t=self.STA_base_tile_size[0],
-            ts_h=self.STA_base_tile_size[1],
-            ts_w=self.STA_base_tile_size[2])
-
-    def untile(self, x: torch.Tensor) -> torch.Tensor:
-        x = rearrange(
-            x,
-            "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
-            n_t=self.full_window_size[0],
-            n_h=self.full_window_size[1],
-            n_w=self.full_window_size[2],
-            ts_t=self.STA_base_tile_size[0],
-            ts_h=self.STA_base_tile_size[1],
-            ts_w=self.STA_base_tile_size[2])
-        return rearrange(x,
-                         "b (t sp h w) head d -> b (sp t h w) head d",
-                         sp=self.sp_size,
-                         t=self.img_latent_shape_int[0] // self.sp_size,
-                         h=self.img_latent_shape_int[1],
-                         w=self.img_latent_shape_int[2])
-
-    def preprocess_qkv(
-        self,
-        qkv: torch.Tensor,
-        attn_metadata: AttentionMetadata,
-    ) -> torch.Tensor:
-        img_sequence_length = qkv.shape[1]
-        self.img_latent_shape_str = self.img_latent_shape_mapping[
-            img_sequence_length]
-        self.full_window_size = self.full_window_mapping[
-            self.img_latent_shape_str]
-        self.img_latent_shape_int = list(
-            map(int, self.img_latent_shape_str.split('x')))
-        self.img_seq_length = self.img_latent_shape_int[
-            0] * self.img_latent_shape_int[1] * self.img_latent_shape_int[2]
-        return self.tile(qkv)
-
-    def postprocess_output(
-        self,
-        output: torch.Tensor,
-        attn_metadata: SlidingTileAttentionMetadata,
-    ) -> torch.Tensor:
-        return self.untile(output)
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        attn_metadata: SlidingTileAttentionMetadata,
-    ) -> torch.Tensor:
-
-        assert self.mask_strategy is not None, "mask_strategy cannot be None for SlidingTileAttention"
-        assert self.mask_strategy[
-            0] is not None, "mask_strategy[0] cannot be None for SlidingTileAttention"
-
-        timestep = attn_metadata.current_timestep
-        # pattern:'.double_blocks.0.attn.impl' or '.single_blocks.0.attn.impl'
-        layer_idx = int(self.prefix.split('.')[-3])
-
-        # TODO: remove hardcode
-
-        text_length = q.shape[1] - self.img_seq_length
-        has_text = text_length > 0
-
-        query = q.transpose(1, 2).contiguous()
-        key = k.transpose(1, 2).contiguous()
-        value = v.transpose(1, 2).contiguous()
-
-        head_num = query.size(1)
-        sp_group = get_sp_group()
-        current_rank = sp_group.rank_in_group
-        start_head = current_rank * head_num
-        windows = [
-            self.mask_strategy[timestep][layer_idx][head_idx + start_head]
-            for head_idx in range(head_num)
-        ]
-        # if has_text is False:
-        #     from IPython import embed
-        #     embed()
-        hidden_states = sliding_tile_attention(
-            query, key, value, windows, text_length, has_text,
-            self.img_latent_shape_str).transpose(1, 2)
-
-        return hidden_states

fastvideo/v1/attention/layer.py

@@ -1,204 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-from typing import Optional, Tuple
-
-import torch
-import torch.nn as nn
-
-from fastvideo.v1.attention.selector import (backend_name_to_enum,
-                                             get_attn_backend)
-from fastvideo.v1.distributed.communication_op import (
-    sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D)
-from fastvideo.v1.distributed.parallel_state import (
-    get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size)
-from fastvideo.v1.forward_context import ForwardContext, get_forward_context
-from fastvideo.v1.platforms import _Backend
-
-
-class DistributedAttention(nn.Module):
-    """Distributed attention layer.
-    """
-
-    def __init__(self,
-                 num_heads: int,
-                 head_size: int,
-                 num_kv_heads: Optional[int] = None,
-                 softmax_scale: Optional[float] = None,
-                 causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[_Backend,
-                                                              ...]] = None,
-                 prefix: str = "",
-                 **extra_impl_args) -> None:
-        super().__init__()
-        if softmax_scale is None:
-            self.softmax_scale = head_size**-0.5
-        else:
-            self.softmax_scale = softmax_scale
-
-        if num_kv_heads is None:
-            num_kv_heads = num_heads
-
-        dtype = torch.get_default_dtype()
-        attn_backend = get_attn_backend(
-            head_size,
-            dtype,
-            supported_attention_backends=supported_attention_backends)
-        impl_cls = attn_backend.get_impl_cls()
-        self.impl = impl_cls(num_heads=num_heads,
-                             head_size=head_size,
-                             causal=causal,
-                             softmax_scale=self.softmax_scale,
-                             num_kv_heads=num_kv_heads,
-                             prefix=f"{prefix}.impl",
-                             **extra_impl_args)
-        self.num_heads = num_heads
-        self.head_size = head_size
-        self.num_kv_heads = num_kv_heads
-        self.backend = backend_name_to_enum(attn_backend.get_name())
-        self.dtype = dtype
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        replicated_q: Optional[torch.Tensor] = None,
-        replicated_k: Optional[torch.Tensor] = None,
-        replicated_v: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-        """Forward pass for distributed attention.
-        
-        Args:
-            q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim]
-            k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim]
-            v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim]
-            replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens
-            replicated_k (Optional[torch.Tensor]): Replicated key tensor
-            replicated_v (Optional[torch.Tensor]): Replicated value tensor
-            
-        Returns:
-            Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing:
-                - o (torch.Tensor): Output tensor after attention for the main sequence
-                - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided
-        """
-        # 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()
-
-        forward_context: ForwardContext = get_forward_context()
-        ctx_attn_metadata = forward_context.attn_metadata
-
-        # Stack QKV
-        qkv = torch.cat([q, k, v], dim=0)  # [3, seq_len, num_heads, head_dim]
-
-        # Redistribute heads across sequence dimension
-        qkv = sequence_model_parallel_all_to_all_4D(qkv,
-                                                    scatter_dim=2,
-                                                    gather_dim=1)
-        # Apply backend-specific preprocess_qkv
-        qkv = self.impl.preprocess_qkv(qkv, ctx_attn_metadata)
-
-        # Concatenate with replicated QKV if provided
-        if replicated_q is not None:
-            assert replicated_k is not None and replicated_v is not None
-            replicated_qkv = torch.cat(
-                [replicated_q, replicated_k, replicated_v],
-                dim=0)  # [3, seq_len, num_heads, head_dim]
-            heads_per_rank = num_heads // world_size
-            replicated_qkv = replicated_qkv[:, :, local_rank *
-                                            heads_per_rank:(local_rank + 1) *
-                                            heads_per_rank]
-            qkv = torch.cat([qkv, replicated_qkv], dim=1)
-
-        q, k, v = qkv.chunk(3, dim=0)
-
-        output = self.impl.forward(q, k, v, ctx_attn_metadata)
-
-        # Redistribute back if using sequence parallelism
-        replicated_output = None
-        if replicated_q is not None:
-            replicated_output = output[:, seq_len * world_size:]
-            output = output[:, :seq_len * world_size]
-            # TODO: make this asynchronous
-            replicated_output = sequence_model_parallel_all_gather(
-                replicated_output.contiguous(), dim=2)
-        # Apply backend-specific postprocess_output
-        output = self.impl.postprocess_output(output, ctx_attn_metadata)
-
-        output = sequence_model_parallel_all_to_all_4D(output,
-                                                       scatter_dim=1,
-                                                       gather_dim=2)
-        return output, replicated_output
-
-
-class LocalAttention(nn.Module):
-    """Attention layer.
-    """
-
-    def __init__(self,
-                 num_heads: int,
-                 head_size: int,
-                 num_kv_heads: Optional[int] = None,
-                 softmax_scale: Optional[float] = None,
-                 causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[_Backend,
-                                                              ...]] = None,
-                 **extra_impl_args) -> None:
-        super().__init__()
-        if softmax_scale is None:
-            self.softmax_scale = head_size**-0.5
-        else:
-            self.softmax_scale = softmax_scale
-        if num_kv_heads is None:
-            num_kv_heads = num_heads
-
-        dtype = torch.get_default_dtype()
-        attn_backend = get_attn_backend(
-            head_size,
-            dtype,
-            supported_attention_backends=supported_attention_backends)
-        impl_cls = attn_backend.get_impl_cls()
-        self.impl = impl_cls(num_heads=num_heads,
-                             head_size=head_size,
-                             softmax_scale=self.softmax_scale,
-                             num_kv_heads=num_kv_heads,
-                             causal=causal,
-                             **extra_impl_args)
-        self.num_heads = num_heads
-        self.head_size = head_size
-        self.num_kv_heads = num_kv_heads
-        self.backend = backend_name_to_enum(attn_backend.get_name())
-        self.dtype = dtype
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-    ) -> torch.Tensor:
-        """
-        Apply local attention between query, key and value tensors.
-        
-        Args:
-            q (torch.Tensor): Query tensor of shape [batch_size, seq_len, num_heads, head_dim]
-            k (torch.Tensor): Key tensor of shape [batch_size, seq_len, num_heads, head_dim] 
-            v (torch.Tensor): Value tensor of shape [batch_size, seq_len, num_heads, head_dim]
-            
-        Returns:
-            torch.Tensor: Output tensor after local attention
-        """
-        # Check input shapes
-        assert q.dim() == 4 and k.dim() == 4 and v.dim(
-        ) == 4, "Expected 4D tensors"
-
-        forward_context: ForwardContext = get_forward_context()
-        ctx_attn_metadata = forward_context.attn_metadata
-
-        output = self.impl.forward(q, k, v, ctx_attn_metadata)
-        return output

fastvideo/v1/attention/selector.py

@@ -1,155 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py
-
-import os
-from contextlib import contextmanager
-from functools import cache
-from typing import Generator, Optional, Tuple, Type, cast
-
-import torch
-
-import fastvideo.v1.envs as envs
-from fastvideo.v1.attention.backends.abstract import AttentionBackend
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.platforms import _Backend, current_platform
-from fastvideo.v1.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
-
-logger = init_logger(__name__)
-
-
-def backend_name_to_enum(backend_name: str) -> Optional[_Backend]:
-    """
-    Convert a string backend name to a _Backend enum value.
-
-    Returns:
-    * _Backend: enum value if backend_name is a valid in-tree type
-    * None: otherwise it's an invalid in-tree type or an out-of-tree platform is
-            loaded.
-    """
-    assert backend_name is not None
-    return _Backend[backend_name] if backend_name in _Backend.__members__ else \
-          None
-
-
-def get_env_variable_attn_backend() -> Optional[_Backend]:
-    '''
-    Get the backend override specified by the FastVideo attention
-    backend environment variable, if one is specified.
-
-    Returns:
-
-    * _Backend enum value if an override is specified
-    * None otherwise
-    '''
-    backend_name = os.environ.get(STR_BACKEND_ENV_VAR)
-    return (None
-            if backend_name is None else backend_name_to_enum(backend_name))
-
-
-# Global state allows a particular choice of backend
-# to be forced, overriding the logic which auto-selects
-# a backend based on system & workload configuration
-# (default behavior if this variable is None)
-#
-# THIS SELECTION TAKES PRECEDENCE OVER THE
-# FASTVIDEO ATTENTION BACKEND ENVIRONMENT VARIABLE
-forced_attn_backend: Optional[_Backend] = None
-
-
-def global_force_attn_backend(attn_backend: Optional[_Backend]) -> None:
-    '''
-    Force all attention operations to use a specified backend.
-
-    Passing `None` for the argument re-enables automatic
-    backend selection.,
-
-    Arguments:
-
-    * attn_backend: backend selection (None to revert to auto)
-    '''
-    global forced_attn_backend
-    forced_attn_backend = attn_backend
-
-
-def get_global_forced_attn_backend() -> Optional[_Backend]:
-    '''
-    Get the currently-forced choice of attention backend,
-    or None if auto-selection is currently enabled.
-    '''
-    return forced_attn_backend
-
-
-def get_attn_backend(
-    head_size: int,
-    dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
-) -> Type[AttentionBackend]:
-    return _cached_get_attn_backend(head_size, dtype,
-                                    supported_attention_backends)
-
-
-@cache
-def _cached_get_attn_backend(
-    head_size: int,
-    dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
-) -> Type[AttentionBackend]:
-    # Check whether a particular choice of backend was
-    # previously forced.
-    #
-    # THIS SELECTION OVERRIDES THE FASTVIDEO_ATTENTION_BACKEND
-    # ENVIRONMENT VARIABLE.
-    if not supported_attention_backends:
-        raise ValueError("supported_attention_backends is empty")
-    selected_backend = None
-    backend_by_global_setting: Optional[_Backend] = (
-        get_global_forced_attn_backend())
-    if backend_by_global_setting is not None:
-        selected_backend = backend_by_global_setting
-    else:
-        # Check the environment variable and override if specified
-        backend_by_env_var: Optional[str] = envs.FASTVIDEO_ATTENTION_BACKEND
-        if backend_by_env_var is not None:
-            selected_backend = backend_name_to_enum(backend_by_env_var)
-
-    # get device-specific attn_backend
-    if selected_backend not in supported_attention_backends:
-        selected_backend = None
-    attention_cls = current_platform.get_attn_backend_cls(
-        selected_backend, head_size, dtype)
-    if not attention_cls:
-        raise ValueError(
-            f"Invalid attention backend for {current_platform.device_name}")
-    return cast(Type[AttentionBackend], resolve_obj_by_qualname(attention_cls))
-
-
-@contextmanager
-def global_force_attn_backend_context_manager(
-        attn_backend: _Backend) -> Generator[None, None, None]:
-    '''
-    Globally force a FastVideo attention backend override within a
-    context manager, reverting the global attention backend
-    override to its prior state upon exiting the context
-    manager.
-
-    Arguments:
-
-    * attn_backend: attention backend to force
-
-    Returns:
-
-    * Generator
-    '''
-
-    # Save the current state of the global backend override (if any)
-    original_value = get_global_forced_attn_backend()
-
-    # Globally force the new backend override
-    global_force_attn_backend(attn_backend)
-
-    # Yield control back to the enclosed code block
-    try:
-        yield
-    finally:
-        # Revert the original global backend override, if any
-        global_force_attn_backend(original_value)

fastvideo/v1/configs/models/__init__.py

@@ -1,7 +0,0 @@
-from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
-from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
-
-__all__ = [
-    "ArchConfig", "ModelConfig",
-    "VAEArchConfig", "VAEConfig"
-]

fastvideo/v1/configs/models/base.py

@@ -1,47 +0,0 @@
-from dataclasses import dataclass, fields
-from typing import Dict, Any
-
-# 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 & overriden by each model arch_config
-# 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
-@dataclass
-class ArchConfig:
-    pass
-
-@dataclass
-class ModelConfig:
-    # Every model config parameter can be categorized into either ArchConfig or everything else
-    # Diffuser/Transformer parameters
-    arch_config: ArchConfig = ArchConfig()
-
-    # FastVideo-specific parameters here
-    # i.e. STA, quantization, teacache
-
-    # 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}'")
-
-    def update_model_config(
-        self,
-        source_model_dict: Dict[str, Any]    
-    ) -> None:
-        assert "arch_config" not in source_model_dict.keys(), "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:
-                print(f"{type(self).__name__} does not contain field '{key}'!")
-                raise AttributeError(f"Invalid field: {key}")

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

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

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

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

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

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

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

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

fastvideo/v1/configs/pipelines/__init__.py

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

fastvideo/v1/configs/pipelines/base.py

@@ -1,103 +0,0 @@
-from dataclasses import dataclass, asdict, fields
-from typing import Optional, Dict, Any
-import json
-
-from fastvideo.v1.configs.models import ModelConfig, VAEConfig
-from fastvideo.v1.utils import shallow_asdict
-
-@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 # Deprecated
-    vae_config: VAEConfig = VAEConfig()
-
-    # DiT configuration
-    num_channels_latents: Optional[int] = None # Deprecated
-
-    # Image encoder configuration
-    image_encoder_precision: str = "fp32"
-
-    # Text encoder configuration
-    text_encoder_precision: str = "fp16"
-    text_len: int = -1 # Deprecated
-    hidden_state_skip_layer: int = 0 # Deprecated
-
-    # STA (Spatial-Temporal Attention) parameters
-    mask_strategy_file_path: Optional[str] = None
-    enable_torch_compile: bool = False
-
-    neg_prompt: Optional[str] = None
-
-    def dump_to_json(self, file_path: str):
-        output_dict = shallow_asdict(self)
-        for key, value in output_dict.items():
-            if isinstance(value, ModelConfig):
-                model_dict = asdict(value)
-                # Model Arch Config should be hidden away from the users
-                model_dict.pop("arch_config")
-                output_dict[key] = model_dict
-        
-        with open(file_path, "w") as f:
-            json.dump(output_dict, f, indent=2)
-
-    def load_from_json(self, file_path: str):
-        with open(file_path, "r") as f:
-            input_pipeline_dict = json.load(f)
-        self.update_pipeline_config(input_pipeline_dict)
-
-    def update_pipeline_config(
-        self,
-        source_pipeline_dict: Dict[str, Any]    
-    ) -> None:
-        for f in fields(self):
-            key = f.name
-            if key in source_pipeline_dict:
-                current_value = getattr(self, key)
-                new_value = source_pipeline_dict[key]
-
-                # If it's a nested ModelConfig, update it recursively
-                if isinstance(current_value, ModelConfig):
-                    current_value.update_model_config(new_value)
-                else:
-                    setattr(self, key, new_value)
-
-@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/pipelines/hunyuan.py

@@ -1,47 +0,0 @@
-from dataclasses import dataclass
-from fastvideo.v1.configs.pipelines.base import BaseConfig
-
-from fastvideo.v1.configs.models import VAEConfig
-from fastvideo.v1.configs.models.vaes import HunyuanVAEConfig
-
-@dataclass
-class HunyuanConfig(BaseConfig):
-    """Base configuration for HunYuan pipeline architecture."""
-
-    # HunyuanConfig-specific parameters with defaults
-    # VAE
-    vae_config: VAEConfig = HunyuanVAEConfig()
-    # 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
-
-    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
-    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/pipelines/registry.py

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

fastvideo/v1/configs/pipelines/wan.py

@@ -1,58 +0,0 @@
-from dataclasses import dataclass
-from fastvideo.v1.configs.pipelines.base import BaseConfig
-
-from fastvideo.v1.configs.models import VAEConfig
-from fastvideo.v1.configs.models.vaes import WanVAEConfig
-
-@dataclass
-class WanT2V480PConfig(BaseConfig):
-    """Base configuration for Wan T2V 1.3B pipeline architecture."""
-
-    # WanConfig-specific parameters with defaults
-    # VAE
-    vae_config: VAEConfig = WanVAEConfig()
-    vae_tiling: bool = False
-    vae_sp: bool = False
-
-    # 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
-
-    def __post_init__(self):
-        self.vae_config.load_encoder = False
-        self.vae_config.load_decoder = True
-
-@dataclass
-class WanI2V480PConfig(WanT2V480PConfig):
-    """Base configuration for Wan I2V 14B 480P pipeline architecture."""
-
-    # WanConfig-specific parameters with defaults
-    # Denoising stage
-    guidance_scale: float = 5.0
-    num_inference_steps: int = 40
-
-    # Precision for each component
-    image_encoder_precision: str = "fp32"
-
-    def __post_init__(self):
-        self.vae_config.load_encoder = True
-        self.vae_config.load_decoder = True

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

@@ -1,41 +0,0 @@
-{
-  "height": 480,
-  "width": 832,
-  "num_frames": 81,
-  "fps": 16,
-  "num_inference_steps": 50,
-  "guidance_scale": 3.0,
-  "seed": 1024,
-  "guidance_rescale": 0.0,
-  "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
-  },
-  "num_channels_latents": null,
-  "image_encoder_precision": "fp32",
-  "text_encoder_precision": "fp32",
-  "text_len": 512,
-  "hidden_state_skip_layer": 0,
-  "mask_strategy_file_path": null,
-  "enable_torch_compile": false,
-  "neg_prompt": "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"
-}

fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json

@@ -1,41 +0,0 @@
-{
-  "height": 480,
-  "width": 832,
-  "num_frames": 81,
-  "fps": 16,
-  "num_inference_steps": 40,
-  "guidance_scale": 5.0,
-  "seed": 1024,
-  "guidance_rescale": 0.0,
-  "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
-  },
-  "num_channels_latents": null,
-  "image_encoder_precision": "fp32",
-  "text_encoder_precision": "fp32",
-  "text_len": 512,
-  "hidden_state_skip_layer": 0,
-  "mask_strategy_file_path": null,
-  "enable_torch_compile": false,
-  "neg_prompt": "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"
-}

fastvideo/v1/default_configs/v1_inference_hunyuan_config.yaml

@@ -1,16 +0,0 @@
-num_gpus: 4
-model_path: FastVideo/FastHunyuan-diffusers
-master_port: 29503
-sp_size: 4 
-tp_size: 4 
-height: 720 
-width: 1280 
-num_frames: 125 
-num_inference_steps: 6 
-guidance_scale: 1 
-embedded_cfg_scale: 6 
-flow_shift: 17 
-prompt_path: ./assets/prompt.txt 
-seed: 1024 
-output_path: outputs_video/ 
-vae-sp: True

fastvideo/v1/distributed/__init__.py

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

fastvideo/v1/distributed/communication_op.py

@@ -1,32 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/communication_op.py
-
-import torch
-import torch.distributed
-
-from fastvideo.v1.distributed.parallel_state import get_sp_group, get_tp_group
-
-
-def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
-    """All-reduce the input tensor across model parallel group."""
-    return get_tp_group().all_reduce(input_)
-
-
-def tensor_model_parallel_all_gather(input_: torch.Tensor,
-                                     dim: int = -1) -> torch.Tensor:
-    """All-gather the input tensor across model parallel group."""
-    return get_tp_group().all_gather(input_, dim)
-
-
-# TODO: remove model, make it sequence_parallel
-def sequence_model_parallel_all_to_all_4D(input_: torch.Tensor,
-                                          scatter_dim: int = 2,
-                                          gather_dim: int = 1) -> torch.Tensor:
-    """All-to-all communication of 4D tensors (e.g. QKV matrices) across sequence parallel group."""
-    return get_sp_group().all_to_all_4D(input_, scatter_dim, gather_dim)
-
-
-def sequence_model_parallel_all_gather(input_: torch.Tensor,
-                                       dim: int = -1) -> torch.Tensor:
-    """All-gather the input tensor across model parallel group."""
-    return get_sp_group().all_gather(input_, dim)

fastvideo/v1/distributed/device_communicators/base_device_communicator.py

@@ -1,194 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/base_device_communicator.py
-
-from typing import Optional
-
-import torch
-import torch.distributed as dist
-from torch.distributed import ProcessGroup
-
-
-class DeviceCommunicatorBase:
-    """
-    Base class for device-specific communicator.
-    It can use the `cpu_group` to initialize the communicator.
-    If the device has PyTorch integration (PyTorch can recognize its
-    communication backend), the `device_group` will also be given.
-    """
-
-    def __init__(self,
-                 cpu_group: ProcessGroup,
-                 device: Optional[torch.device] = None,
-                 device_group: Optional[ProcessGroup] = None,
-                 unique_name: str = ""):
-        self.device = device or torch.device("cpu")
-        self.cpu_group = cpu_group
-        self.device_group = device_group
-        self.unique_name = unique_name
-        self.rank = dist.get_rank(cpu_group)
-        self.world_size = dist.get_world_size(cpu_group)
-        self.ranks = dist.get_process_group_ranks(cpu_group)
-        self.global_rank = dist.get_rank()
-        self.global_world_size = dist.get_world_size()
-        self.rank_in_group = dist.get_group_rank(self.cpu_group,
-                                                 self.global_rank)
-
-    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
-        dist.all_reduce(input_, group=self.device_group)
-        return input_
-
-    def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
-        if dim < 0:
-            # Convert negative dim to positive.
-            dim += input_.dim()
-        input_size = input_.size()
-        # NOTE: we have to use concat-style all-gather here,
-        # stack-style all-gather has compatibility issues with
-        # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
-        output_size = (input_size[0] * self.world_size, ) + input_size[1:]
-        # Allocate output tensor.
-        output_tensor = torch.empty(output_size,
-                                    dtype=input_.dtype,
-                                    device=input_.device)
-        # All-gather.
-        dist.all_gather_into_tensor(output_tensor,
-                                    input_,
-                                    group=self.device_group)
-        # Reshape
-        output_tensor = output_tensor.reshape((self.world_size, ) + input_size)
-        output_tensor = output_tensor.movedim(0, dim)
-        output_tensor = output_tensor.reshape(input_size[:dim] +
-                                              (self.world_size *
-                                               input_size[dim], ) +
-                                              input_size[dim + 1:])
-        return output_tensor
-
-    def gather(self,
-               input_: torch.Tensor,
-               dst: int = 0,
-               dim: int = -1) -> Optional[torch.Tensor]:
-        """
-        NOTE: We assume that the input tensor is on the same device across
-        all the ranks.
-        NOTE: `dst` is the local rank of the destination rank.
-        """
-        world_size = self.world_size
-        assert -input_.dim() <= dim < input_.dim(), (
-            f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
-        if dim < 0:
-            # Convert negative dim to positive.
-            dim += input_.dim()
-
-        # Allocate output tensor.
-        if self.rank_in_group == dst:
-            gather_list = [torch.empty_like(input_) for _ in range(world_size)]
-        else:
-            gather_list = None
-        # Gather.
-        torch.distributed.gather(input_,
-                                 gather_list,
-                                 dst=self.ranks[dst],
-                                 group=self.device_group)
-        if self.rank_in_group == dst:
-            output_tensor = torch.cat(gather_list, dim=dim)
-        else:
-            output_tensor = None
-        return output_tensor
-
-    def all_to_all_4D(self,
-                      input_: torch.Tensor,
-                      scatter_dim: int = 2,
-                      gather_dim: int = 1) -> torch.Tensor:
-        """Specialized all-to-all operation for 4D tensors (e.g., for QKV matrices).
-        
-        Args:
-            input_ (torch.Tensor): 4D input tensor to be scattered and gathered.
-            scatter_dim (int, optional): Dimension along which to scatter. Defaults to 2.
-            gather_dim (int, optional): Dimension along which to gather. Defaults to 1.
-            
-        Returns:
-            torch.Tensor: Output tensor after all-to-all operation.
-        """
-        # Bypass the function if we are using only 1 GPU.
-        if self.world_size == 1:
-            return input_
-
-        assert input_.dim(
-        ) == 4, f"input must be 4D tensor, got {input_.dim()} and shape {input_.shape}"
-
-        if scatter_dim == 2 and gather_dim == 1:
-            # input: (bs, seqlen/P, hc, hs) output: (bs, seqlen, hc/P, hs)
-            bs, shard_seqlen, hc, hs = input_.shape
-            seqlen = shard_seqlen * self.world_size
-            shard_hc = hc // self.world_size
-
-            # Reshape and transpose for scattering
-            input_t = (input_.reshape(bs, shard_seqlen, self.world_size,
-                                      shard_hc, hs).transpose(0,
-                                                              2).contiguous())
-
-            output = torch.empty_like(input_t)
-
-            torch.distributed.all_to_all_single(output,
-                                                input_t,
-                                                group=self.device_group)
-            torch.cuda.synchronize()
-
-            # Reshape and transpose back
-            output = output.reshape(seqlen, bs, shard_hc,
-                                    hs).transpose(0, 1).contiguous().reshape(
-                                        bs, seqlen, shard_hc, hs)
-
-            return output
-
-        elif scatter_dim == 1 and gather_dim == 2:
-            # input: (bs, seqlen, hc/P, hs) output: (bs, seqlen/P, hc, hs)
-            bs, seqlen, shard_hc, hs = input_.shape
-            hc = shard_hc * self.world_size
-            shard_seqlen = seqlen // self.world_size
-
-            # Reshape and transpose for scattering
-            input_t = (input_.reshape(bs, self.world_size, shard_seqlen,
-                                      shard_hc, hs).transpose(0, 3).transpose(
-                                          0, 1).contiguous().reshape(
-                                              self.world_size, shard_hc,
-                                              shard_seqlen, bs, hs))
-            output = torch.empty_like(input_t)
-
-            torch.distributed.all_to_all_single(output,
-                                                input_t,
-                                                group=self.device_group)
-            torch.cuda.synchronize()
-
-            # Reshape and transpose back
-            output = output.reshape(hc, shard_seqlen, bs,
-                                    hs).transpose(0, 2).contiguous().reshape(
-                                        bs, shard_seqlen, hc, hs)
-
-            return output
-        else:
-            raise RuntimeError(
-                "scatter_dim must be 1 or 2 and gather_dim must be 1 or 2")
-
-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
-        """Sends a tensor to the destination rank in a non-blocking way"""
-        """NOTE: `dst` is the local rank of the destination rank."""
-        if dst is None:
-            dst = (self.rank_in_group + 1) % self.world_size
-        torch.distributed.send(tensor, self.ranks[dst], self.device_group)
-
-    def recv(self,
-             size: torch.Size,
-             dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
-        """Receives a tensor from the source rank."""
-        """NOTE: `src` is the local rank of the source rank."""
-        if src is None:
-            src = (self.rank_in_group - 1) % self.world_size
-
-        tensor = torch.empty(size, dtype=dtype, device=self.device)
-        torch.distributed.recv(tensor, self.ranks[src], self.device_group)
-        return tensor
-
-    def destroy(self) -> None:
-        pass

fastvideo/v1/distributed/device_communicators/cuda_communicator.py

@@ -1,76 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/cuda_communicator.py
-
-from typing import Optional
-
-import torch
-from torch.distributed import ProcessGroup
-
-from fastvideo.v1.distributed.device_communicators.base_device_communicator import (
-    DeviceCommunicatorBase)
-
-
-class CudaCommunicator(DeviceCommunicatorBase):
-
-    def __init__(self,
-                 cpu_group: ProcessGroup,
-                 device: Optional[torch.device] = None,
-                 device_group: Optional[ProcessGroup] = None,
-                 unique_name: str = ""):
-        super().__init__(cpu_group, device, device_group, unique_name)
-
-        from fastvideo.v1.distributed.device_communicators.pynccl import (
-            PyNcclCommunicator)
-
-        self.pynccl_comm: Optional[PyNcclCommunicator] = None
-        if self.world_size > 1:
-            self.pynccl_comm = PyNcclCommunicator(
-                group=self.cpu_group,
-                device=self.device,
-            )
-
-    def all_reduce(self, input_):
-        pynccl_comm = self.pynccl_comm
-        assert pynccl_comm is not None
-        out = pynccl_comm.all_reduce(input_)
-        if out is None:
-            # fall back to the default all-reduce using PyTorch.
-            # this usually happens during testing.
-            # when we run the model, allreduce only happens for the TP
-            # group, where we always have either custom allreduce or pynccl.
-            out = input_.clone()
-            torch.distributed.all_reduce(out, group=self.device_group)
-        return out
-
-    def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
-        """Sends a tensor to the destination rank in a non-blocking way"""
-        """NOTE: `dst` is the local rank of the destination rank."""
-        if dst is None:
-            dst = (self.rank_in_group + 1) % self.world_size
-
-        pynccl_comm = self.pynccl_comm
-        if pynccl_comm is not None and not pynccl_comm.disabled:
-            pynccl_comm.send(tensor, dst)
-        else:
-            torch.distributed.send(tensor, self.ranks[dst], self.device_group)
-
-    def recv(self,
-             size: torch.Size,
-             dtype: torch.dtype,
-             src: Optional[int] = None) -> torch.Tensor:
-        """Receives a tensor from the source rank."""
-        """NOTE: `src` is the local rank of the source rank."""
-        if src is None:
-            src = (self.rank_in_group - 1) % self.world_size
-
-        tensor = torch.empty(size, dtype=dtype, device=self.device)
-        pynccl_comm = self.pynccl_comm
-        if pynccl_comm is not None and not pynccl_comm.disabled:
-            pynccl_comm.recv(tensor, src)
-        else:
-            torch.distributed.recv(tensor, self.ranks[src], self.device_group)
-        return tensor
-
-    def destroy(self) -> None:
-        if self.pynccl_comm is not None:
-            self.pynccl_comm = None

fastvideo/v1/distributed/device_communicators/pynccl.py

@@ -1,218 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl.py
-
-from typing import Optional, Union
-
-# ===================== import region =====================
-import torch
-import torch.distributed as dist
-from torch.distributed import ProcessGroup, ReduceOp
-
-from fastvideo.v1.distributed.device_communicators.pynccl_wrapper import (
-    NCCLLibrary, buffer_type, cudaStream_t, ncclComm_t, ncclDataTypeEnum,
-    ncclRedOpTypeEnum, ncclUniqueId)
-from fastvideo.v1.distributed.utils import StatelessProcessGroup
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import current_stream
-
-logger = init_logger(__name__)
-
-
-class PyNcclCommunicator:
-
-    def __init__(
-        self,
-        group: Union[ProcessGroup, StatelessProcessGroup],
-        device: Union[int, str, torch.device],
-        library_path: Optional[str] = None,
-    ):
-        """
-        Args:
-            group: the process group to work on. If None, it will use the
-                default process group.
-            device: the device to bind the PyNcclCommunicator to. If None,
-                it will be bind to f"cuda:{local_rank}".
-            library_path: the path to the NCCL library. If None, it will
-                use the default library path.
-        It is the caller's responsibility to make sure each communicator
-        is bind to a unique device.
-        """
-        if not isinstance(group, StatelessProcessGroup):
-            assert dist.is_initialized()
-            assert dist.get_backend(group) != dist.Backend.NCCL, (
-                "PyNcclCommunicator should be attached to a non-NCCL group.")
-            # note: this rank is the rank in the group
-            self.rank = dist.get_rank(group)
-            self.world_size = dist.get_world_size(group)
-        else:
-            self.rank = group.rank
-            self.world_size = group.world_size
-
-        self.group = group
-
-        # if world_size == 1, no need to create communicator
-        if self.world_size == 1:
-            self.available = False
-            self.disabled = True
-            return
-        try:
-            self.nccl = NCCLLibrary(library_path)
-        except Exception:
-            # disable because of missing NCCL library
-            # e.g. in a non-GPU environment
-            self.available = False
-            self.disabled = True
-            return
-
-        self.available = True
-        self.disabled = False
-
-        logger.info("FastVideo is using nccl==%s", self.nccl.ncclGetVersion())
-
-        if self.rank == 0:
-            # get the unique id from NCCL
-            self.unique_id = self.nccl.ncclGetUniqueId()
-        else:
-            # construct an empty unique id
-            self.unique_id = ncclUniqueId()
-
-        if not isinstance(group, StatelessProcessGroup):
-            tensor = torch.ByteTensor(list(self.unique_id.internal))
-            ranks = dist.get_process_group_ranks(group)
-            # arg `src` in `broadcast` is the global rank
-            dist.broadcast(tensor, src=ranks[0], group=group)
-            byte_list = tensor.tolist()
-            for i, byte in enumerate(byte_list):
-                self.unique_id.internal[i] = byte
-        else:
-            self.unique_id = group.broadcast_obj(self.unique_id, src=0)
-        if isinstance(device, int):
-            device = torch.device(f"cuda:{device}")
-        elif isinstance(device, str):
-            device = torch.device(device)
-        # now `device` is a `torch.device` object
-        assert isinstance(device, torch.device)
-        self.device = device
-        # nccl communicator and stream will use this device
-        # `torch.cuda.device` is a context manager that changes the
-        # current cuda device to the specified one
-        with torch.cuda.device(device):
-            self.comm: ncclComm_t = self.nccl.ncclCommInitRank(
-                self.world_size, self.unique_id, self.rank)
-
-            stream = current_stream()
-            # A small all_reduce for warmup.
-            data = torch.zeros(1, device=device)
-            self.all_reduce(data)
-            stream.synchronize()
-            del data
-
-    def all_reduce(self,
-                   in_tensor: torch.Tensor,
-                   op: ReduceOp = ReduceOp.SUM,
-                   stream=None) -> torch.Tensor:
-        if self.disabled:
-            return None
-        # nccl communicator created on a specific device
-        # will only work on tensors on the same device
-        # otherwise it will cause "illegal memory access"
-        assert in_tensor.device == self.device, (
-            f"this nccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {in_tensor.device}")
-
-        out_tensor = torch.empty_like(in_tensor)
-
-        if stream is None:
-            stream = current_stream()
-        self.nccl.ncclAllReduce(buffer_type(in_tensor.data_ptr()),
-                                buffer_type(out_tensor.data_ptr()),
-                                in_tensor.numel(),
-                                ncclDataTypeEnum.from_torch(in_tensor.dtype),
-                                ncclRedOpTypeEnum.from_torch(op), self.comm,
-                                cudaStream_t(stream.cuda_stream))
-        return out_tensor
-
-    def all_gather(self,
-                   output_tensor: torch.Tensor,
-                   input_tensor: torch.Tensor,
-                   stream=None):
-        if self.disabled:
-            return
-        # nccl communicator created on a specific device
-        # will only work on tensors on the same device
-        # otherwise it will cause "illegal memory access"
-        assert input_tensor.device == self.device, (
-            f"this nccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {input_tensor.device}")
-        if stream is None:
-            stream = current_stream()
-        self.nccl.ncclAllGather(buffer_type(input_tensor.data_ptr()),
-                                buffer_type(output_tensor.data_ptr()),
-                                input_tensor.numel(),
-                                ncclDataTypeEnum.from_torch(input_tensor.dtype),
-                                self.comm, cudaStream_t(stream.cuda_stream))
-
-    def reduce_scatter(self,
-                       output_tensor: torch.Tensor,
-                       input_tensor: torch.Tensor,
-                       op: ReduceOp = ReduceOp.SUM,
-                       stream=None):
-        if self.disabled:
-            return
-        # nccl communicator created on a specific device
-        # will only work on tensors on the same device
-        # otherwise it will cause "illegal memory access"
-        assert input_tensor.device == self.device, (
-            f"this nccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {input_tensor.device}")
-        if stream is None:
-            stream = current_stream()
-        self.nccl.ncclReduceScatter(
-            buffer_type(input_tensor.data_ptr()),
-            buffer_type(output_tensor.data_ptr()), output_tensor.numel(),
-            ncclDataTypeEnum.from_torch(input_tensor.dtype),
-            ncclRedOpTypeEnum.from_torch(op), self.comm,
-            cudaStream_t(stream.cuda_stream))
-
-    def send(self, tensor: torch.Tensor, dst: int, stream=None):
-        if self.disabled:
-            return
-        assert tensor.device == self.device, (
-            f"this nccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {tensor.device}")
-        if stream is None:
-            stream = current_stream()
-        self.nccl.ncclSend(buffer_type(tensor.data_ptr()), tensor.numel(),
-                           ncclDataTypeEnum.from_torch(tensor.dtype), dst,
-                           self.comm, cudaStream_t(stream.cuda_stream))
-
-    def recv(self, tensor: torch.Tensor, src: int, stream=None):
-        if self.disabled:
-            return
-        assert tensor.device == self.device, (
-            f"this nccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {tensor.device}")
-        if stream is None:
-            stream = current_stream()
-        self.nccl.ncclRecv(buffer_type(tensor.data_ptr()), tensor.numel(),
-                           ncclDataTypeEnum.from_torch(tensor.dtype), src,
-                           self.comm, cudaStream_t(stream.cuda_stream))
-
-    def broadcast(self, tensor: torch.Tensor, src: int, stream=None):
-        if self.disabled:
-            return
-        assert tensor.device == self.device, (
-            f"this nccl communicator is created to work on {self.device}, "
-            f"but the input tensor is on {tensor.device}")
-        if stream is None:
-            stream = current_stream()
-        if src == self.rank:
-            sendbuff = buffer_type(tensor.data_ptr())
-            # NCCL requires the sender also to have a receive buffer
-            recvbuff = buffer_type(tensor.data_ptr())
-        else:
-            sendbuff = buffer_type()
-            recvbuff = buffer_type(tensor.data_ptr())
-        self.nccl.ncclBroadcast(sendbuff, recvbuff, tensor.numel(),
-                                ncclDataTypeEnum.from_torch(tensor.dtype), src,
-                                self.comm, cudaStream_t(stream.cuda_stream))

fastvideo/v1/distributed/device_communicators/pynccl_wrapper.py

@@ -1,341 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/device_communicators/pynccl_wrapper.py
-
-# This file is a pure Python wrapper for the NCCL library.
-# The main purpose is to use NCCL combined with CUDA graph.
-# Before writing this script, we tried the following approach:
-# 1. We tried to use `cupy`, it calls NCCL correctly, but `cupy` itself
-#  often gets stuck when initializing the NCCL communicator.
-# 2. We tried to use `torch.distributed`, but `torch.distributed.all_reduce`
-#  contains many other potential cuda APIs, that are not allowed during
-#  capturing the CUDA graph. For further details, please check
-# https://discuss.pytorch.org/t/pytorch-cudagraph-with-nccl-operation-failed/ .
-#
-# Another rejected idea is to write a C/C++ binding for NCCL. It is usually
-# doable, but we often encounter issues related with nccl versions, and need
-# to switch between different versions of NCCL. See
-# https://github.com/NVIDIA/nccl/issues/1234 for more details.
-# A C/C++ binding is not flexible enough to handle this. It requires
-# recompilation of the code every time we want to switch between different
-# versions. This current implementation, with a **pure** Python wrapper, is
-# more flexible. We can easily switch between different versions of NCCL by
-# changing the environment variable `FASTVIDEO_NCCL_SO_PATH`, or the `so_file`
-# variable in the code.
-
-#TODO(will): support FASTVIDEO_NCCL_SO_PATH
-
-import ctypes
-import platform
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
-
-import torch
-from torch.distributed import ReduceOp
-
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import find_nccl_library
-
-logger = init_logger(__name__)
-
-# === export types and functions from nccl to Python ===
-# for the original nccl definition, please check
-# https://github.com/NVIDIA/nccl/blob/master/src/nccl.h.in
-
-ncclResult_t = ctypes.c_int
-ncclComm_t = ctypes.c_void_p
-
-
-class ncclUniqueId(ctypes.Structure):
-    _fields_ = [("internal", ctypes.c_byte * 128)]
-
-
-cudaStream_t = ctypes.c_void_p
-buffer_type = ctypes.c_void_p
-
-ncclDataType_t = ctypes.c_int
-
-
-class ncclDataTypeEnum:
-    ncclInt8 = 0
-    ncclChar = 0
-    ncclUint8 = 1
-    ncclInt32 = 2
-    ncclInt = 2
-    ncclUint32 = 3
-    ncclInt64 = 4
-    ncclUint64 = 5
-    ncclFloat16 = 6
-    ncclHalf = 6
-    ncclFloat32 = 7
-    ncclFloat = 7
-    ncclFloat64 = 8
-    ncclDouble = 8
-    ncclBfloat16 = 9
-    ncclNumTypes = 10
-
-    @classmethod
-    def from_torch(cls, dtype: torch.dtype) -> int:
-        if dtype == torch.int8:
-            return cls.ncclInt8
-        if dtype == torch.uint8:
-            return cls.ncclUint8
-        if dtype == torch.int32:
-            return cls.ncclInt32
-        if dtype == torch.int64:
-            return cls.ncclInt64
-        if dtype == torch.float16:
-            return cls.ncclFloat16
-        if dtype == torch.float32:
-            return cls.ncclFloat32
-        if dtype == torch.float64:
-            return cls.ncclFloat64
-        if dtype == torch.bfloat16:
-            return cls.ncclBfloat16
-        raise ValueError(f"Unsupported dtype: {dtype}")
-
-
-ncclRedOp_t = ctypes.c_int
-
-
-class ncclRedOpTypeEnum:
-    ncclSum = 0
-    ncclProd = 1
-    ncclMax = 2
-    ncclMin = 3
-    ncclAvg = 4
-    ncclNumOps = 5
-
-    @classmethod
-    def from_torch(cls, op: ReduceOp) -> int:
-        if op == ReduceOp.SUM:
-            return cls.ncclSum
-        if op == ReduceOp.PRODUCT:
-            return cls.ncclProd
-        if op == ReduceOp.MAX:
-            return cls.ncclMax
-        if op == ReduceOp.MIN:
-            return cls.ncclMin
-        if op == ReduceOp.AVG:
-            return cls.ncclAvg
-        raise ValueError(f"Unsupported op: {op}")
-
-
-@dataclass
-class Function:
-    name: str
-    restype: Any
-    argtypes: List[Any]
-
-
-class NCCLLibrary:
-    exported_functions = [
-        # const char* ncclGetErrorString(ncclResult_t result)
-        Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]),
-        # ncclResult_t  ncclGetVersion(int *version);
-        Function("ncclGetVersion", ncclResult_t,
-                 [ctypes.POINTER(ctypes.c_int)]),
-        # ncclResult_t ncclGetUniqueId(ncclUniqueId* uniqueId);
-        Function("ncclGetUniqueId", ncclResult_t,
-                 [ctypes.POINTER(ncclUniqueId)]),
-        # ncclResult_t  ncclCommInitRank(
-        #   ncclComm_t* comm, int nranks, ncclUniqueId commId, int rank);
-        # note that ncclComm_t is a pointer type, so the first argument
-        # is a pointer to a pointer
-        Function("ncclCommInitRank", ncclResult_t, [
-            ctypes.POINTER(ncclComm_t), ctypes.c_int, ncclUniqueId, ctypes.c_int
-        ]),
-        # ncclResult_t  ncclAllReduce(
-        #   const void* sendbuff, void* recvbuff, size_t count,
-        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
-        #   cudaStream_t stream);
-        # note that cudaStream_t is a pointer type, so the last argument
-        # is a pointer
-        Function("ncclAllReduce", ncclResult_t, [
-            buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
-            ncclRedOp_t, ncclComm_t, cudaStream_t
-        ]),
-
-        # ncclResult_t  ncclAllGather(
-        #   const void* sendbuff, void* recvbuff, size_t count,
-        #   ncclDataType_t datatype, ncclComm_t comm,
-        #   cudaStream_t stream);
-        # note that cudaStream_t is a pointer type, so the last argument
-        # is a pointer
-        Function("ncclAllGather", ncclResult_t, [
-            buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
-            ncclComm_t, cudaStream_t
-        ]),
-
-        # ncclResult_t  ncclReduceScatter(
-        #   const void* sendbuff, void* recvbuff, size_t count,
-        #   ncclDataType_t datatype, ncclRedOp_t op, ncclComm_t comm,
-        #   cudaStream_t stream);
-        # note that cudaStream_t is a pointer type, so the last argument
-        # is a pointer
-        Function("ncclReduceScatter", ncclResult_t, [
-            buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
-            ncclRedOp_t, ncclComm_t, cudaStream_t
-        ]),
-
-        # ncclResult_t  ncclSend(
-        #   const void* sendbuff, size_t count, ncclDataType_t datatype,
-        #   int dest, ncclComm_t comm, cudaStream_t stream);
-        Function("ncclSend", ncclResult_t, [
-            buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int,
-            ncclComm_t, cudaStream_t
-        ]),
-
-        # ncclResult_t  ncclRecv(
-        #   void* recvbuff, size_t count, ncclDataType_t datatype,
-        #   int src, ncclComm_t comm, cudaStream_t stream);
-        Function("ncclRecv", ncclResult_t, [
-            buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int,
-            ncclComm_t, cudaStream_t
-        ]),
-
-        # ncclResult_t ncclBroadcast(
-        #   const void* sendbuff, void* recvbuff, size_t count,
-        #   ncclDataType_t datatype, int root, ncclComm_t comm,
-        #   cudaStream_t stream);
-        Function("ncclBroadcast", ncclResult_t, [
-            buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t,
-            ctypes.c_int, ncclComm_t, cudaStream_t
-        ]),
-
-        # be cautious! this is a collective call, it will block until all
-        # processes in the communicator have called this function.
-        # because Python object destruction can happen in random order,
-        # it is better not to call it at all.
-        # ncclResult_t  ncclCommDestroy(ncclComm_t comm);
-        Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]),
-    ]
-
-    # class attribute to store the mapping from the path to the library
-    # to avoid loading the same library multiple times
-    path_to_library_cache: Dict[str, Any] = {}
-
-    # class attribute to store the mapping from library path
-    #  to the corresponding dictionary
-    path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}
-
-    def __init__(self, so_file: Optional[str] = None):
-
-        so_file = so_file or find_nccl_library()
-
-        try:
-            if so_file not in NCCLLibrary.path_to_dict_mapping:
-                lib = ctypes.CDLL(so_file)
-                NCCLLibrary.path_to_library_cache[so_file] = lib
-            self.lib = NCCLLibrary.path_to_library_cache[so_file]
-        except Exception as e:
-            logger.error(
-                "Failed to load NCCL library from %s ."
-                "It is expected if you are not running on NVIDIA/AMD GPUs."
-                "Otherwise, the nccl library might not exist, be corrupted "
-                "or it does not support the current platform %s."
-                "If you already have the library, please set the "
-                "environment variable FASTVIDEO_NCCL_SO_PATH"
-                " to point to the correct nccl library path.", so_file,
-                platform.platform())
-            raise e
-
-        if so_file not in NCCLLibrary.path_to_dict_mapping:
-            _funcs: Dict[str, Any] = {}
-            for func in NCCLLibrary.exported_functions:
-                f = getattr(self.lib, func.name)
-                f.restype = func.restype
-                f.argtypes = func.argtypes
-                _funcs[func.name] = f
-            NCCLLibrary.path_to_dict_mapping[so_file] = _funcs
-        self._funcs = NCCLLibrary.path_to_dict_mapping[so_file]
-
-    def ncclGetErrorString(self, result: ncclResult_t) -> str:
-        return str(self._funcs["ncclGetErrorString"](result).decode("utf-8"))
-
-    def NCCL_CHECK(self, result: ncclResult_t) -> None:
-        if result != 0:
-            error_str = self.ncclGetErrorString(result)
-            raise RuntimeError(f"NCCL error: {error_str}")
-
-    def ncclGetVersion(self) -> str:
-        version = ctypes.c_int()
-        self.NCCL_CHECK(self._funcs["ncclGetVersion"](ctypes.byref(version)))
-        version_str = str(version.value)
-        # something like 21903 --> "2.19.3"
-        major = version_str[0].lstrip("0")
-        minor = version_str[1:3].lstrip("0")
-        patch = version_str[3:].lstrip("0")
-        return f"{major}.{minor}.{patch}"
-
-    def ncclGetUniqueId(self) -> ncclUniqueId:
-        unique_id = ncclUniqueId()
-        self.NCCL_CHECK(self._funcs["ncclGetUniqueId"](ctypes.byref(unique_id)))
-        return unique_id
-
-    def ncclCommInitRank(self, world_size: int, unique_id: ncclUniqueId,
-                         rank: int) -> ncclComm_t:
-        comm = ncclComm_t()
-        self.NCCL_CHECK(self._funcs["ncclCommInitRank"](ctypes.byref(comm),
-                                                        world_size, unique_id,
-                                                        rank))
-        return comm
-
-    def ncclAllReduce(self, sendbuff: buffer_type, recvbuff: buffer_type,
-                      count: int, datatype: int, op: int, comm: ncclComm_t,
-                      stream: cudaStream_t) -> None:
-        # `datatype` actually should be `ncclDataType_t`
-        # and `op` should be `ncclRedOp_t`
-        # both are aliases of `ctypes.c_int`
-        # when we pass int to a function, it will be converted to `ctypes.c_int`
-        # by ctypes automatically
-        self.NCCL_CHECK(self._funcs["ncclAllReduce"](sendbuff, recvbuff, count,
-                                                     datatype, op, comm,
-                                                     stream))
-
-    def ncclReduceScatter(self, sendbuff: buffer_type, recvbuff: buffer_type,
-                          count: int, datatype: int, op: int, comm: ncclComm_t,
-                          stream: cudaStream_t) -> None:
-        # `datatype` actually should be `ncclDataType_t`
-        # and `op` should be `ncclRedOp_t`
-        # both are aliases of `ctypes.c_int`
-        # when we pass int to a function, it will be converted to `ctypes.c_int`
-        # by ctypes automatically
-        self.NCCL_CHECK(self._funcs["ncclReduceScatter"](sendbuff, recvbuff,
-                                                         count, datatype, op,
-                                                         comm, stream))
-
-    def ncclAllGather(self, sendbuff: buffer_type, recvbuff: buffer_type,
-                      count: int, datatype: int, comm: ncclComm_t,
-                      stream: cudaStream_t) -> None:
-        # `datatype` actually should be `ncclDataType_t`
-        # which is an aliases of `ctypes.c_int`
-        # when we pass int to a function, it will be converted to `ctypes.c_int`
-        # by ctypes automatically
-        self.NCCL_CHECK(self._funcs["ncclAllGather"](sendbuff, recvbuff, count,
-                                                     datatype, comm, stream))
-
-    def ncclSend(self, sendbuff: buffer_type, count: int, datatype: int,
-                 dest: int, comm: ncclComm_t, stream: cudaStream_t) -> None:
-        self.NCCL_CHECK(self._funcs["ncclSend"](sendbuff, count, datatype, dest,
-                                                comm, stream))
-
-    def ncclRecv(self, recvbuff: buffer_type, count: int, datatype: int,
-                 src: int, comm: ncclComm_t, stream: cudaStream_t) -> None:
-        self.NCCL_CHECK(self._funcs["ncclRecv"](recvbuff, count, datatype, src,
-                                                comm, stream))
-
-    def ncclBroadcast(self, sendbuff: buffer_type, recvbuff: buffer_type,
-                      count: int, datatype: int, root: int, comm: ncclComm_t,
-                      stream: cudaStream_t) -> None:
-        self.NCCL_CHECK(self._funcs["ncclBroadcast"](sendbuff, recvbuff, count,
-                                                     datatype, root, comm,
-                                                     stream))
-
-    def ncclCommDestroy(self, comm: ncclComm_t) -> None:
-        self.NCCL_CHECK(self._funcs["ncclCommDestroy"](comm))
-
-
-__all__ = [
-    "NCCLLibrary", "ncclDataTypeEnum", "ncclRedOpTypeEnum", "ncclUniqueId",
-    "ncclComm_t", "cudaStream_t", "buffer_type"
-]

fastvideo/v1/distributed/utils.py

@@ -1,191 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/utils.py
-
-# Copyright 2023 The vLLM team.
-# Adapted from
-# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/tensor_parallel/utils.py
-# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
-import dataclasses
-import pickle
-import time
-from collections import deque
-from typing import Any, Deque, Dict, Optional, Sequence, Tuple
-
-import torch
-from torch.distributed import TCPStore
-
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-def ensure_divisibility(numerator, denominator) -> None:
-    """Ensure that numerator is divisible by the denominator."""
-    assert numerator % denominator == 0, "{} is not divisible by {}".format(
-        numerator, denominator)
-
-
-def divide(numerator: int, denominator: int) -> int:
-    """Ensure that numerator is divisible by the denominator and return
-    the division value."""
-    ensure_divisibility(numerator, denominator)
-    return numerator // denominator
-
-
-def split_tensor_along_last_dim(
-    tensor: torch.Tensor,
-    num_partitions: int,
-    contiguous_split_chunks: bool = False,
-) -> Sequence[torch.Tensor]:
-    """ Split a tensor along its last dimension.
-
-        Arguments:
-            tensor: input tensor.
-            num_partitions: number of partitions to split the tensor
-            contiguous_split_chunks: If True, make each chunk contiguous
-                                     in memory.
-
-        Returns:
-            A list of Tensors
-    """
-    # Get the size and dimension.
-    last_dim = tensor.dim() - 1
-    last_dim_size = divide(tensor.size()[last_dim], num_partitions)
-    # Split.
-    tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
-    # NOTE: torch.split does not create contiguous tensors by default.
-    if contiguous_split_chunks:
-        return tuple(chunk.contiguous() for chunk in tensor_list)
-
-    return tuple(tensor_list)
-
-
-@dataclasses.dataclass
-class StatelessProcessGroup:
-    """A dataclass to hold a metadata store, and the rank, world_size of the
-    group. Only use it to communicate metadata between processes.
-    For data-plane communication, create NCCL-related objects.
-    """
-    rank: int
-    world_size: int
-    store: torch._C._distributed_c10d.Store
-    data_expiration_seconds: int = 3600  # 1 hour
-
-    # dst rank -> counter
-    send_dst_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
-    # src rank -> counter
-    recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict)
-    broadcast_send_counter: int = 0
-    broadcast_recv_src_counter: Dict[int, int] = dataclasses.field(
-        default_factory=dict)
-
-    # A deque to store the data entries, with key and timestamp.
-    entries: Deque[Tuple[str, float]] = dataclasses.field(default_factory=deque)
-
-    def __post_init__(self):
-        assert self.rank < self.world_size
-        self.send_dst_counter = {i: 0 for i in range(self.world_size)}
-        self.recv_src_counter = {i: 0 for i in range(self.world_size)}
-        self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)}
-
-    def send_obj(self, obj: Any, dst: int):
-        """Send an object to a destination rank."""
-        self.expire_data()
-        key = f"send_to/{dst}/{self.send_dst_counter[dst]}"
-        self.store.set(key, pickle.dumps(obj))
-        self.send_dst_counter[dst] += 1
-        self.entries.append((key, time.time()))
-
-    def expire_data(self) -> None:
-        """Expire data that is older than `data_expiration_seconds` seconds."""
-        while self.entries:
-            # check the oldest entry
-            key, timestamp = self.entries[0]
-            if time.time() - timestamp > self.data_expiration_seconds:
-                self.store.delete_key(key)
-                self.entries.popleft()
-            else:
-                break
-
-    def recv_obj(self, src: int) -> Any:
-        """Receive an object from a source rank."""
-        obj = pickle.loads(
-            self.store.get(f"send_to/{self.rank}/{self.recv_src_counter[src]}"))
-        self.recv_src_counter[src] += 1
-        return obj
-
-    def broadcast_obj(self, obj: Optional[Any], src: int) -> Any:
-        """Broadcast an object from a source rank to all other ranks.
-        It does not clean up after all ranks have received the object.
-        Use it for limited times, e.g., for initialization.
-        """
-        if self.rank == src:
-            self.expire_data()
-            key = (f"broadcast_from/{src}/"
-                   f"{self.broadcast_send_counter}")
-            self.store.set(key, pickle.dumps(obj))
-            self.broadcast_send_counter += 1
-            self.entries.append((key, time.time()))
-            return obj
-        else:
-            key = (f"broadcast_from/{src}/"
-                   f"{self.broadcast_recv_src_counter[src]}")
-            recv_obj = pickle.loads(self.store.get(key))
-            self.broadcast_recv_src_counter[src] += 1
-            return recv_obj
-
-    def all_gather_obj(self, obj: Any) -> list[Any]:
-        """All gather an object from all ranks."""
-        gathered_objs = []
-        for i in range(self.world_size):
-            if i == self.rank:
-                gathered_objs.append(obj)
-                self.broadcast_obj(obj, src=self.rank)
-            else:
-                recv_obj = self.broadcast_obj(None, src=i)
-                gathered_objs.append(recv_obj)
-        return gathered_objs
-
-    def barrier(self):
-        """A barrier to synchronize all ranks."""
-        for i in range(self.world_size):
-            if i == self.rank:
-                self.broadcast_obj(None, src=self.rank)
-            else:
-                self.broadcast_obj(None, src=i)
-
-    @staticmethod
-    def create(
-        host: str,
-        port: int,
-        rank: int,
-        world_size: int,
-        data_expiration_seconds: int = 3600,
-    ) -> "StatelessProcessGroup":
-        """A replacement for `torch.distributed.init_process_group` that does not
-        pollute the global state.
-
-        If we have process A and process B called `torch.distributed.init_process_group`
-        to form a group, and then we want to form another group with process A, B, C,
-        D, it is not possible in PyTorch, because process A and process B have already
-        formed a group, and process C and process D cannot join that group. This
-        function is a workaround for this issue.
-
-        `torch.distributed.init_process_group` is a global call, while this function
-        is a stateless call. It will return a `StatelessProcessGroup` object that can be
-        used for exchanging metadata. With this function, process A and process B
-        can call `StatelessProcessGroup.create` to form a group, and then process A, B,
-        C, and D can call `StatelessProcessGroup.create` to form another group.
-        """ # noqa
-        store = TCPStore(
-            host_name=host,
-            port=port,
-            world_size=world_size,
-            is_master=(rank == 0),
-        )
-
-        return StatelessProcessGroup(
-            rank=rank,
-            world_size=world_size,
-            store=store,
-            data_expiration_seconds=data_expiration_seconds)

fastvideo/v1/entrypoints/cli/cli_types.py

@@ -1,26 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/types.py
-
-import argparse
-
-from fastvideo.v1.utils import FlexibleArgumentParser
-
-
-class CLISubcommand:
-    """Base class for CLI subcommands"""
-
-    name: str
-
-    def cmd(self, args: argparse.Namespace) -> None:
-        """Execute the command with the given arguments"""
-        raise NotImplementedError
-
-    def validate(self, args: argparse.Namespace) -> None:
-        """Validate the arguments for this command"""
-        pass
-
-    def subparser_init(
-            self,
-            subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
-        """Initialize the subparser for this command"""
-        raise NotImplementedError

fastvideo/v1/entrypoints/cli/generate.py

@@ -1,87 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/serve.py
-
-import argparse
-from typing import List, cast
-
-from fastvideo.v1.entrypoints.cli import utils
-from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.utils import FlexibleArgumentParser
-
-
-class GenerateSubcommand(CLISubcommand):
-    """The `generate` subcommand for the FastVideo CLI"""
-
-    def __init__(self) -> None:
-        self.name = "generate"
-        super().__init__()
-
-    def cmd(self, args: argparse.Namespace) -> None:
-        excluded_args = [
-            'subparser', 'config', 'num_gpus', 'master_port',
-            'dispatch_function'
-        ]
-
-        # Create a filtered dictionary of arguments
-        filtered_args = {
-            k: v
-            for k, v in vars(args).items()
-            if k not in excluded_args and v is not None
-        }
-
-        main_args = []
-
-        for key, value in filtered_args.items():
-            # Convert underscores to dashes in argument names
-            arg_name = f"--{key.replace('_', '-')}"
-
-            # Handle boolean flags
-            if isinstance(value, bool):
-                if value:
-                    main_args.append(arg_name)
-            else:
-                main_args.append(arg_name)
-                main_args.append(str(value))
-
-        utils.launch_distributed(args.num_gpus,
-                                 main_args,
-                                 master_port=args.master_port)
-
-    def validate(self, args: argparse.Namespace) -> None:
-        if args.num_gpus is not None and args.num_gpus <= 0:
-            raise ValueError("Number of gpus must be positive")
-
-        if args.master_port is not None and (args.master_port < 1024
-                                             or args.master_port > 65535):
-            raise ValueError("Master port must be between 1024 and 65535")
-
-    def subparser_init(
-            self,
-            subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
-        generate_parser = subparsers.add_parser(
-            "generate",
-            help="Run inference on a model",
-            usage=
-            "fastvideo generate --model-path MODEL_PATH_OR_ID --prompt PROMPT [OPTIONS]"
-        )
-
-        generate_parser.add_argument(
-            "--config",
-            type=str,
-            default='',
-            required=False,
-            help="Read CLI options from a config YAML file.")
-
-        generate_parser.add_argument("--master-port",
-                                     type=int,
-                                     default=None,
-                                     help="Port for the master process")
-
-        generate_parser = FastVideoArgs.add_cli_args(generate_parser)
-
-        return cast(FlexibleArgumentParser, generate_parser)
-
-
-def cmd_init() -> List[CLISubcommand]:
-    return [GenerateSubcommand()]

fastvideo/v1/entrypoints/cli/main.py

@@ -1,39 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/entrypoints/cli/main.py
-
-from typing import List
-
-from fastvideo.v1.entrypoints.cli.cli_types import CLISubcommand
-from fastvideo.v1.entrypoints.cli.generate import cmd_init as generate_cmd_init
-from fastvideo.v1.utils import FlexibleArgumentParser
-
-
-def cmd_init() -> List[CLISubcommand]:
-    """Initialize all commands from separate modules"""
-    commands = []
-    commands.extend(generate_cmd_init())
-    return commands
-
-
-def main() -> None:
-    parser = FlexibleArgumentParser(description="FastVideo CLI")
-    parser.add_argument('-v', '--version', action='version', version='0.1.0')
-
-    subparsers = parser.add_subparsers(required=False, dest="subparser")
-
-    cmds = {}
-    for cmd in cmd_init():
-        cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
-        cmds[cmd.name] = cmd
-    args = parser.parse_args()
-    if args.subparser in cmds:
-        cmds[args.subparser].validate(args)
-
-    if hasattr(args, "dispatch_function"):
-        args.dispatch_function(args)
-    else:
-        parser.print_help()
-
-
-if __name__ == "__main__":
-    main()

fastvideo/v1/entrypoints/cli/utils.py

@@ -1,60 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-
-import os
-import subprocess
-import sys
-from typing import List, Optional
-
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-def launch_distributed(num_gpus: int,
-                       args: List[str],
-                       master_port: Optional[int] = None) -> int:
-    """
-    Launch a distributed job with the given arguments
-    
-    Args:
-        num_gpus: Number of GPUs to use
-        args: Arguments to pass to v1_fastvideo_inference.py (defaults to sys.argv[1:])
-        master_port: Port for the master process (default: random)
-    """
-
-    current_env = os.environ.copy()
-    python_executable = sys.executable
-    project_root = os.path.abspath(
-        os.path.join(os.path.dirname(__file__), "../../../.."))
-    main_script = os.path.join(project_root,
-                               "fastvideo/v1/sample/v1_fastvideo_inference.py")
-
-    cmd = [
-        python_executable, "-m", "torch.distributed.run",
-        f"--nproc_per_node={num_gpus}"
-    ]
-
-    if master_port is not None:
-        cmd.append(f"--master_port={master_port}")
-
-    cmd.append(main_script)
-    cmd.extend(args)
-
-    logger.info("Running inference with %d GPU(s)", num_gpus)
-    logger.info("Launching command: %s", " ".join(cmd))
-
-    current_env["PYTHONIOENCODING"] = "utf-8"
-    process = subprocess.Popen(cmd,
-                               env=current_env,
-                               stdout=subprocess.PIPE,
-                               stderr=subprocess.STDOUT,
-                               universal_newlines=True,
-                               bufsize=1,
-                               encoding='utf-8',
-                               errors='replace')
-
-    if process.stdout:
-        for line in iter(process.stdout.readline, ''):
-            print(line.strip())
-
-    return process.wait()

fastvideo/v1/entrypoints/video_generator.py

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

fastvideo/v1/envs.py

@@ -1,213 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/envs.py
-
-import os
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
-
-if TYPE_CHECKING:
-    FASTVIDEO_RINGBUFFER_WARNING_INTERVAL: int = 60
-    FASTVIDEO_NCCL_SO_PATH: Optional[str] = None
-    LD_LIBRARY_PATH: Optional[str] = None
-    LOCAL_RANK: int = 0
-    CUDA_VISIBLE_DEVICES: Optional[str] = None
-    FASTVIDEO_CACHE_ROOT: str = os.path.expanduser("~/.cache/fastvideo")
-    FASTVIDEO_CONFIG_ROOT: str = os.path.expanduser("~/.config/fastvideo")
-    FASTVIDEO_CONFIGURE_LOGGING: int = 1
-    FASTVIDEO_LOGGING_LEVEL: str = "INFO"
-    FASTVIDEO_LOGGING_PREFIX: str = ""
-    FASTVIDEO_LOGGING_CONFIG_PATH: Optional[str] = None
-    FASTVIDEO_TRACE_FUNCTION: int = 0
-    FASTVIDEO_ATTENTION_BACKEND: Optional[str] = None
-    FASTVIDEO_ATTENTION_CONFIG: Optional[str] = None
-    FASTVIDEO_WORKER_MULTIPROC_METHOD: str = "fork"
-    FASTVIDEO_TARGET_DEVICE: str = "cuda"
-    MAX_JOBS: Optional[str] = None
-    NVCC_THREADS: Optional[str] = None
-    CMAKE_BUILD_TYPE: Optional[str] = None
-    VERBOSE: bool = False
-    FASTVIDEO_SERVER_DEV_MODE: bool = False
-
-
-def get_default_cache_root() -> str:
-    return os.getenv(
-        "XDG_CACHE_HOME",
-        os.path.join(os.path.expanduser("~"), ".cache"),
-    )
-
-
-def get_default_config_root() -> str:
-    return os.getenv(
-        "XDG_CONFIG_HOME",
-        os.path.join(os.path.expanduser("~"), ".config"),
-    )
-
-
-def maybe_convert_int(value: Optional[str]) -> Optional[int]:
-    if value is None:
-        return None
-    return int(value)
-
-
-# The begin-* and end* here are used by the documentation generator
-# to extract the used env vars.
-
-# begin-env-vars-definition
-
-environment_variables: Dict[str, Callable[[], Any]] = {
-
-    # ================== Installation Time Env Vars ==================
-
-    # Target device of FastVideo, supporting [cuda (by default),
-    # rocm, neuron, cpu, openvino]
-    "FASTVIDEO_TARGET_DEVICE":
-    lambda: os.getenv("FASTVIDEO_TARGET_DEVICE", "cuda"),
-
-    # Maximum number of compilation jobs to run in parallel.
-    # By default this is the number of CPUs
-    "MAX_JOBS":
-    lambda: os.getenv("MAX_JOBS", None),
-
-    # Number of threads to use for nvcc
-    # By default this is 1.
-    # If set, `MAX_JOBS` will be reduced to avoid oversubscribing the CPU.
-    "NVCC_THREADS":
-    lambda: os.getenv("NVCC_THREADS", None),
-
-    # If set, fastvideo will use precompiled binaries (*.so)
-    "FASTVIDEO_USE_PRECOMPILED":
-    lambda: bool(os.environ.get("FASTVIDEO_USE_PRECOMPILED")) or bool(
-        os.environ.get("FASTVIDEO_PRECOMPILED_WHEEL_LOCATION")),
-
-    # CMake build type
-    # If not set, defaults to "Debug" or "RelWithDebInfo"
-    # Available options: "Debug", "Release", "RelWithDebInfo"
-    "CMAKE_BUILD_TYPE":
-    lambda: os.getenv("CMAKE_BUILD_TYPE"),
-
-    # If set, fastvideo will print verbose logs during installation
-    "VERBOSE":
-    lambda: bool(int(os.getenv('VERBOSE', '0'))),
-
-    # Root directory for FASTVIDEO configuration files
-    # Defaults to `~/.config/fastvideo` unless `XDG_CONFIG_HOME` is set
-    # Note that this not only affects how fastvideo finds its configuration files
-    # during runtime, but also affects how fastvideo installs its configuration
-    # files during **installation**.
-    "FASTVIDEO_CONFIG_ROOT":
-    lambda: os.path.expanduser(
-        os.getenv(
-            "FASTVIDEO_CONFIG_ROOT",
-            os.path.join(get_default_config_root(), "fastvideo"),
-        )),
-
-    # ================== Runtime Env Vars ==================
-
-    # Root directory for FASTVIDEO cache files
-    # Defaults to `~/.cache/fastvideo` unless `XDG_CACHE_HOME` is set
-    "FASTVIDEO_CACHE_ROOT":
-    lambda: os.path.expanduser(
-        os.getenv(
-            "FASTVIDEO_CACHE_ROOT",
-            os.path.join(get_default_cache_root(), "fastvideo"),
-        )),
-
-    # Interval in seconds to log a warning message when the ring buffer is full
-    "FASTVIDEO_RINGBUFFER_WARNING_INTERVAL":
-    lambda: int(os.environ.get("FASTVIDEO_RINGBUFFER_WARNING_INTERVAL", "60")),
-
-    # Path to the NCCL library file. It is needed because nccl>=2.19 brought
-    # by PyTorch contains a bug: https://github.com/NVIDIA/nccl/issues/1234
-    "FASTVIDEO_NCCL_SO_PATH":
-    lambda: os.environ.get("FASTVIDEO_NCCL_SO_PATH", None),
-
-    # when `FASTVIDEO_NCCL_SO_PATH` is not set, fastvideo will try to find the nccl
-    # library file in the locations specified by `LD_LIBRARY_PATH`
-    "LD_LIBRARY_PATH":
-    lambda: os.environ.get("LD_LIBRARY_PATH", None),
-
-    # Internal flag to enable Dynamo fullgraph capture
-    "FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE":
-    lambda: bool(
-        os.environ.get("FASTVIDEO_TEST_DYNAMO_FULLGRAPH_CAPTURE", "1") != "0"),
-
-    # local rank of the process in the distributed setting, used to determine
-    # the GPU device id
-    "LOCAL_RANK":
-    lambda: int(os.environ.get("LOCAL_RANK", "0")),
-
-    # used to control the visible devices in the distributed setting
-    "CUDA_VISIBLE_DEVICES":
-    lambda: os.environ.get("CUDA_VISIBLE_DEVICES", None),
-
-    # timeout for each iteration in the engine
-    "FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S":
-    lambda: int(os.environ.get("FASTVIDEO_ENGINE_ITERATION_TIMEOUT_S", "60")),
-
-    # Logging configuration
-    # If set to 0, fastvideo will not configure logging
-    # If set to 1, fastvideo will configure logging using the default configuration
-    #    or the configuration file specified by FASTVIDEO_LOGGING_CONFIG_PATH
-    "FASTVIDEO_CONFIGURE_LOGGING":
-    lambda: int(os.getenv("FASTVIDEO_CONFIGURE_LOGGING", "1")),
-    "FASTVIDEO_LOGGING_CONFIG_PATH":
-    lambda: os.getenv("FASTVIDEO_LOGGING_CONFIG_PATH"),
-
-    # this is used for configuring the default logging level
-    "FASTVIDEO_LOGGING_LEVEL":
-    lambda: os.getenv("FASTVIDEO_LOGGING_LEVEL", "INFO"),
-
-    # if set, FASTVIDEO_LOGGING_PREFIX will be prepended to all log messages
-    "FASTVIDEO_LOGGING_PREFIX":
-    lambda: os.getenv("FASTVIDEO_LOGGING_PREFIX", ""),
-
-    # Trace function calls
-    # If set to 1, fastvideo will trace function calls
-    # Useful for debugging
-    "FASTVIDEO_TRACE_FUNCTION":
-    lambda: int(os.getenv("FASTVIDEO_TRACE_FUNCTION", "0")),
-
-    # Backend for attention computation
-    # Available options:
-    # - "TORCH_SDPA": use torch.nn.MultiheadAttention
-    # - "FLASH_ATTN": use FlashAttention
-    # - "STA" : use sliding tile attention
-    "FASTVIDEO_ATTENTION_BACKEND":
-    lambda: os.getenv("FASTVIDEO_ATTENTION_BACKEND", None),
-
-    # Path to the attention configuration file. Only used for sliding tile
-    # attention for now.
-    "FASTVIDEO_ATTENTION_CONFIG":
-    lambda: (None if os.getenv("FASTVIDEO_ATTENTION_CONFIG", None) is None else
-             os.path.expanduser(os.getenv("FASTVIDEO_ATTENTION_CONFIG", "."))),
-
-    # Use dedicated multiprocess context for workers.
-    # Both spawn and fork work
-    "FASTVIDEO_WORKER_MULTIPROC_METHOD":
-    lambda: os.getenv("FASTVIDEO_WORKER_MULTIPROC_METHOD", "fork"),
-
-    # Enables torch profiler if set. Path to the directory where torch profiler
-    # traces are saved. Note that it must be an absolute path.
-    "FASTVIDEO_TORCH_PROFILER_DIR":
-    lambda: (None
-             if os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", None) is None else os.
-             path.expanduser(os.getenv("FASTVIDEO_TORCH_PROFILER_DIR", "."))),
-
-    # If set, fastvideo will run in development mode, which will enable
-    # some additional endpoints for developing and debugging,
-    # e.g. `/reset_prefix_cache`
-    "FASTVIDEO_SERVER_DEV_MODE":
-    lambda: bool(int(os.getenv("FASTVIDEO_SERVER_DEV_MODE", "0"))),
-}
-
-# end-env-vars-definition
-
-
-def __getattr__(name: str):
-    # lazy evaluation of environment variables
-    if name in environment_variables:
-        return environment_variables[name]()
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
-
-
-def __dir__():
-    return list(environment_variables.keys())