add pyproject and docker to ssim test trigger

trigger tests
[bugfix] [training] use separate generator for validation (#610 )
2025-07-15 13:26:12 -07:00 · 2025-07-15 13:26:12 -07:00 · 2025-07-15 13:19:38 -07:00 · 2025-07-15 15:06:44 -05:00 · 2025-07-15 11:22:56 -07:00 · 2025-07-09 17:44:24 -05:00
290 changed files with 10744 additions and 3553 deletions
@@ -0,0 +1,155 @@
 env:
  IMAGE_VERSION: "py3.12-latest"
  BUILDKITE_CLEAN_CHECKOUT: true
 steps:
  - label: "pre-commit"
    command: ".buildkite/scripts/pre_commit.sh"
    agents:
      queue: "default"
  - wait
  - label: "Trigger Tests"
    plugins:
      - monorepo-diff#v1.4.0:
          diff: "git diff --name-only $BUILDKITE_PULL_REQUEST_BASE_BRANCH...HEAD"
          watch:
            - path: 
                - "fastvideo/v1/models/encoders/**"
                - "fastvideo/v1/models/loader/**"
                - "fastvideo/v1/tests/encoders/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Encoder Tests"
                env:
                  - TEST_TYPE=encoder
                agents:
                  queue: "default"
            - path:
                - "fastvideo/v1/models/vaes/**"
                - "fastvideo/v1/models/loader/**"
                - "fastvideo/v1/tests/vaes/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "VAE Tests"
                env:
                  - TEST_TYPE=vae
                agents:
                  queue: "default"
            - path:
                - "fastvideo/v1/models/dits/**"
                - "fastvideo/v1/models/loader/**"
                - "fastvideo/v1/tests/transformers/**"
                - "fastvideo/v1/layers/**"
                - "fastvideo/v1/attention/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Transformer Tests"
                env:
                  - TEST_TYPE=transformer
                agents:
                  queue: "default"
            - path: 
                - "fastvideo/v1/**/*.py"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 30m .buildkite/scripts/pr_test.sh"
                label: "SSIM Tests"
                env:
                  - TEST_TYPE=ssim
                agents:
                  queue: "default"
            - path:
                - "fastvideo/v1/tests/lora/**"
                - "fastvideo/v1/models/loader/**"
                - "fastvideo/v1/tests/transformers/**"
                - "fastvideo/v1/pipelines/**"
                - "fastvideo/v1/layers/lora/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "LoRA Inference Tests"
                env:
                  - TEST_TYPE=inference_lora
                agents:
                  queue: "default"
            - path:
                - "fastvideo/v1/**"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Training Tests"
                env:
                  - TEST_TYPE=training
                agents:
                  queue: "default"
            - path:
                - "fastvideo/v1/**"
                - "csrc/attn/vsa/**"
                - "csrc/attn/tk/**"
                - "csrc/attn/setup_vsa.py"
                - "csrc/attn/config_vsa.py"
                - "csrc/attn/vsa.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Training Tests VSA"
                env:
                  - TEST_TYPE=training_vsa
                agents:
                  queue: "default"
            - path:
                - "fastvideo/v1/**"
                - "csrc/attn/st_attn/**"
                - "csrc/attn/setup_sta.py"
                - "csrc/attn/config_sta.py"
                - "csrc/attn/st_attn.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Inference Tests STA"
                env:
                  - TEST_TYPE=inference_sta
                agents:
                  queue: "default"
            - path:
                - "csrc/attn/st_attn/**"
                - "csrc/attn/setup_sta.py"
                - "csrc/attn/config_sta.py"
                - "csrc/attn/st_attn.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Precision Tests STA"
                env:
                  - TEST_TYPE=precision_sta
                agents:
                  queue: "default"
            - path:
                - "csrc/attn/vsa/**"
                - "csrc/attn/tk/**"
                - "csrc/attn/setup_vsa.py"
                - "csrc/attn/config_vsa.py"
                - "csrc/attn/vsa.cpp"
                - "pyproject.toml"
                - "docker/Dockerfile.python3.12"
              config:
                command: "timeout 15m .buildkite/scripts/pr_test.sh"
                label: "Precision Tests VSA"
                env:
                  - TEST_TYPE=precision_vsa
                agents:
                  queue: "default"
@@ -0,0 +1,121 @@
 #!/bin/bash
 set -uo pipefail
 log() {
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
 }
 log "=== Starting Modal test execution ==="
 # Change to the project directory
 cd "$(dirname "$0")/../.."
 PROJECT_ROOT=$(pwd)
 log "Project root: $PROJECT_ROOT"
 # Install Modal if not available
 if ! python3 -m modal --version &> /dev/null; then
    log "Modal not found, installing..."
    python3 -m pip install modal
    # Verify installation
    if ! python3 -m modal --version &> /dev/null; then
        log "Error: Failed to install modal. Please install it manually."
        exit 1
    fi
 fi
 log "modal version: $(python3 -m modal --version)"
 # Set up Modal authentication using Buildkite secrets
 log "Setting up Modal authentication from Buildkite secrets..."
 MODAL_TOKEN_ID=$(buildkite-agent secret get modal_token_id)
 MODAL_TOKEN_SECRET=$(buildkite-agent secret get modal_token_secret)
 WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
 WANDB_API_KEY=$(buildkite-agent secret get wandb_api_key)
 if [ -n "$MODAL_TOKEN_ID" ] && [ -n "$MODAL_TOKEN_SECRET" ]; then
    log "Retrieved Modal credentials from Buildkite secrets"
    python3 -m modal token set --token-id "$MODAL_TOKEN_ID" --token-secret "$MODAL_TOKEN_SECRET" --profile buildkite-ci --activate --verify
    if [ $? -eq 0 ]; then
        log "Modal authentication successful"
    else
        log "Error: Failed to set Modal credentials"
        exit 1
    fi
 else
    log "Error: Could not retrieve Modal credentials from Buildkite secrets."
    log "Please ensure 'modal_token_id' and 'modal_token_secret' secrets are set in Buildkite."
    exit 1
 fi
 MODAL_TEST_FILE="fastvideo/v1/tests/modal/pr_test.py"
 if [ -z "${TEST_TYPE:-}" ]; then
    log "Error: TEST_TYPE environment variable is not set"
    exit 1
 fi
 log "Test type: $TEST_TYPE"
 MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$BUILDKITE_PULL_REQUEST IMAGE_VERSION=$IMAGE_VERSION"
 case "$TEST_TYPE" in
    "encoder")
        log "Running encoder tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
        ;;
    "vae")
        log "Running VAE tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
        ;;
    "transformer")
        log "Running transformer tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
        ;;
    "ssim")
        log "Running SSIM tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_ssim_tests"
        ;;
    "training")
        log "Running training tests..."
        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests"
        ;;
    "training_vsa")
        log "Running training VSA tests..."
        MODAL_COMMAND="$MODAL_ENV WANDB_API_KEY=$WANDB_API_KEY python3 -m modal run $MODAL_TEST_FILE::run_training_tests_VSA"
        ;;
    "inference_sta")
        log "Running inference STA tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_tests_STA"
        ;;
    "precision_sta")
        log "Running precision STA tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_STA"
        ;;
    "precision_vsa")
        log "Running precision VSA tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_precision_tests_VSA"
        ;;
    "inference_lora")
        log "Running LoRA tests..."
        MODAL_COMMAND="$MODAL_ENV python3 -m modal run $MODAL_TEST_FILE::run_inference_lora_tests"
        ;;
    *)
        log "Error: Unknown test type: $TEST_TYPE"
        exit 1
        ;;
 esac
 log "Executing: $MODAL_COMMAND"
 eval "$MODAL_COMMAND"
 TEST_EXIT_CODE=$?
 if [ $TEST_EXIT_CODE -eq 0 ]; then
    log "Modal test completed successfully"
 else
    log "Error: Modal test failed with exit code: $TEST_EXIT_CODE"
 fi
 log "=== Test execution completed with exit code: $TEST_EXIT_CODE ==="
 exit $TEST_EXIT_CODE
@@ -0,0 +1,40 @@
 #!/bin/bash
 set -uo pipefail
 log() {
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
 }
 log "=== Starting pre-commit checks ==="
 cd "$(dirname "$0")/../.."
 PROJECT_ROOT=$(pwd)
 log "Project root: $PROJECT_ROOT"
 if ! python3 -m pre_commit --version &> /dev/null; then
    log "pre-commit not found, installing..."
    python3 -m pip install --user pre-commit==4.0.1
    if ! python3 -m pre_commit --version &> /dev/null; then
        log "Error: Failed to install pre-commit."
        exit 1
    fi
 fi
 log "Pre-commit version: $(python3 -m pre_commit --version)"
 log "Installing/updating pre-commit hooks..."
 python3 -m pre_commit install --install-hooks
 log "Running pre-commit checks on all files..."
 python3 -m pre_commit run --all-files
 PRE_COMMIT_EXIT_CODE=$?
 if [ $PRE_COMMIT_EXIT_CODE -eq 0 ]; then
    log "Pre-commit checks completed successfully"
 else
    log "Error: Pre-commit checks failed with exit code: $PRE_COMMIT_EXIT_CODE"
 fi
 log "=== Pre-commit checks completed with exit code: $PRE_COMMIT_EXIT_CODE ==="
 exit $PRE_COMMIT_EXIT_CODE 
@@ -160,8 +160,7 @@ def execute_command(pod_id):
    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",
+        "source $HOME/.local/bin/env && source /opt/venv/bin/activate",
        "conda activate fastvideo-dev",
        args.test_command
    ]
@@ -13,4 +13,4 @@
      ]
    }
  ]
-}
+}
@@ -12,13 +12,11 @@ on:
    paths:
      - "fastvideo/**/*.py"
      - ".github/workflows/pr-test.yml"
      - "pyproject.toml"
      - "docker/Dockerfile.python3.12"
      - "csrc/**"
  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
@@ -44,6 +42,26 @@ on:
        required: false
        default: false
        type: boolean
      run_training_test_VSA:
        description: "Run training-test-VSA"
        required: false
        default: false
        type: boolean
      run_inference_test_STA:
        description: "Run inference-test-STA"
        required: false
        default: false
        type: boolean
      run_precision_test_STA:
        description: "Run precision-test-STA"
        required: false
        default: false
        type: boolean
      run_precision_test_VSA:
        description: "Run precision-test-VSA"
        required: false
        default: false
        type: boolean
      run_nightly_test:
        description: "Run nightly-test"
        required: false
@@ -53,6 +71,7 @@ on:
 env:
  PYTHONUNBUFFERED: "1"
 concurrency:
  group: pr-test-${{ github.ref }}
  cancel-in-progress: true
@@ -70,28 +89,71 @@ jobs:
      vae-test: ${{ steps.filter.outputs.vae-test }}
      transformer-test: ${{ steps.filter.outputs.transformer-test }}
      training-test: ${{ steps.filter.outputs.training-test }}
      training-test-VSA: ${{ steps.filter.outputs.training-test-VSA }}
      inference-test-STA: ${{ steps.filter.outputs.inference-test-STA }}
      precision-test-STA: ${{ steps.filter.outputs.precision-test-STA }}
      precision-test-VSA: ${{ steps.filter.outputs.precision-test-VSA }}
    steps:
      - uses: actions/checkout@v4
      - uses: dorny/paths-filter@v3
        id: filter
        with:
          filters: |
            # Define reusable path patterns
            common-paths: &common-paths
              - 'pyproject.toml'
              - 'docker/Dockerfile.python3.12'
            sta-kernel-paths: &sta-kernel-paths
              - 'csrc/attn/st_attn/**'
              - 'csrc/attn/setup_sta.py'
              - 'csrc/attn/config_sta.py'
              - 'csrc/attn/st_attn.cpp'
            vsa-kernel-paths: &vsa-kernel-paths
              - 'csrc/attn/vsa/**'
              - 'csrc/attn/tk/**'
              - 'csrc/attn/setup_vsa.py'
              - 'csrc/attn/config_vsa.py'
              - 'csrc/attn/vsa.cpp'
            vsa-paths: &vsa-paths
              - 'fastvideo/v1/**'
              - *common-paths
              - *vsa-kernel-paths
            # Actual tests
            encoder-test:
              - 'fastvideo/v1/models/encoders/**'
-              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/models/loader/**'
              - 'fastvideo/v1/tests/encoders/**'
              - *common-paths
            vae-test:
              - 'fastvideo/v1/models/vaes/**'
-              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/models/loader/**'
              - 'fastvideo/v1/tests/vaes/**'
              - *common-paths
            transformer-test:
              - 'fastvideo/v1/models/dits/**'
-              - 'fastvideo/v1/models/loaders/**'
+              - 'fastvideo/v1/models/loader/**'
              - 'fastvideo/v1/tests/transformers/**'
              - 'fastvideo/v1/layers/**'
              - 'fastvideo/v1/attention/**'
              - *common-paths
            training-test:
              - 'fastvideo/v1/**'
              - *common-paths
            training-test-VSA:
              - 'fastvideo/v1/**'
              - *common-paths
              - *vsa-kernel-paths
            inference-test-STA:
              - 'fastvideo/v1/**'
              - *common-paths
              - *sta-kernel-paths
            precision-test-STA:
              - *common-paths
              - *sta-kernel-paths
            precision-test-VSA:
              - *common-paths
              - *vsa-kernel-paths
  encoder-test:
    needs: change-filter
@@ -104,8 +166,8 @@ jobs:
      gpu_type: "NVIDIA A40"
      gpu_count: 1
      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -122,8 +184,8 @@ jobs:
      gpu_type: "NVIDIA A40"
      gpu_count: 1
      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -140,8 +202,8 @@ jobs:
      gpu_type: "NVIDIA L40S"
      gpu_count: 1
      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -150,8 +212,7 @@ jobs:
  ssim-test:
    needs: change-filter
    if: >-
-      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
+      github.event_name != 'workflow_dispatch' || (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
    strategy:
      fail-fast: false
      matrix:
@@ -168,7 +229,7 @@ jobs:
      volume_size: 200
      disk_size: 200
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
+      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
      timeout_minutes: 60
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -177,7 +238,7 @@ jobs:
  training-test:
    needs: change-filter
    if: >-
-      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
+      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.training-test == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
@@ -186,8 +247,86 @@ jobs:
      gpu_count: 4
      volume_size: 100
      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/training -srP"
+      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/Vanilla -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
  training-test-VSA:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.training-test-VSA == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test_VSA == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "training-test-VSA"
      gpu_type: "NVIDIA H100 NVL"
      gpu_count: 2
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/training/VSA -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
  inference-test-STA:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.inference-test-STA == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_inference_test_STA == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "inference-test-STA"
      gpu_type: "NVIDIA H100 NVL"
      gpu_count: 2
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/inference/STA -srP"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
  precision-test-STA:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-STA == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_STA == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "precision-test-STA"
      gpu_type: "NVIDIA H100 NVL"
      gpu_count: 1
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_sta.py"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
  precision-test-VSA:
    needs: change-filter
    if: >-
      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-VSA == 'true') || 
      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_VSA == 'true')
    uses: ./.github/workflows/runpod-test.yml
    with:
      job_id: "precision-test-VSA"
      gpu_type: "NVIDIA H100 NVL"
      gpu_count: 1
      volume_size: 100
      disk_size: 100
      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_block_sparse.py"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -203,8 +342,8 @@ jobs:
      gpu_count: 4
      volume_size: 100
      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs"
+      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/v1/tests/nightly/test_e2e_overfit_single_sample.py -vs"
      timeout_minutes: 30
    secrets:
      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -212,7 +351,8 @@ jobs:
      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
  runpod-cleanup:
-    needs: [encoder-test, vae-test, transformer-test, ssim-test] # Add other jobs to this list as you create them
+    # Add other jobs to this list as you create them
    needs: [encoder-test, vae-test, transformer-test, ssim-test, training-test, training-test-VSA, inference-test-STA, precision-test-STA, precision-test-VSA] 
    if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }}
    runs-on: ubuntu-latest
    steps:
@@ -229,7 +369,7 @@ jobs:
      - name: Cleanup all RunPod instances
        env:
-          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12"]'
+          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12", "training-test", "training-test-VSA", "inference-test-STA", "precision-test-STA", "precision-test-VSA"]'
          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
          GITHUB_RUN_ID: ${{ github.run_id }}
-        run: python .github/scripts/runpod_cleanup.py
+        run: python .github/scripts/runpod_cleanup.py
@@ -59,3 +59,5 @@ docs/source/inference/examples/
 # Static images
 !docs/source/_static/images/**/*.png
 !comfyui/assets/**/*.png
 !comfyui/assets/**/*.gif
@@ -60,7 +60,7 @@ repos:
  rev: v1.15.0
  hooks:
  - id: mypy
-    args: [--python-version, '3.10', --follow-imports, "skip", ]
+    args: [--python-version, '3.10', --follow-imports, "skip" ]
    additional_dependencies: [types-cachetools, types-setuptools, types-PyYAML, types-requests]
 - repo: local
  hooks:
@@ -69,7 +69,7 @@ repos:
    entry: bash
    args:
      - -c
-      - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
+      - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/" | grep -v "^fastvideo/v1/tests/inference/lora/L40S_reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0'
    language: system
    always_run: true
    pass_filenames: false
@@ -8,13 +8,18 @@ It features a clean, consistent API that works across popular video models, maki
 With FastVideo's optimizations, you can achieve more than 3x inference improvement compared to other systems.
 <p align="center">
-    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg" target="_blank"> <b>Slack</b> </a> |
+    | <a href="https://hao-ai-lab.github.io/FastVideo"><b>Documentation</b></a> | <a href="https://hao-ai-lab.github.io/FastVideo/inference/inference_quick_start.html"><b> Quick Start</b></a> | 🤗 <a href="https://huggingface.co/FastVideo/FastHunyuan"  target="_blank"><b>FastHunyuan</b></a>  | 🤗 <a href="https://huggingface.co/FastVideo/FastMochi-diffusers" target="_blank"><b>FastMochi</b></a> | 🟣💬 <a href="https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ" target="_blank"> <b>Slack</b> </a> |
 </p>
 <div align="center">
 <img src=assets/perf.png width="90%"/>
 </div>
 ## NEWS
 - ```2025/06/14```: Release finetuning and inference code for [VSA](https://arxiv.org/pdf/2505.13389)
 - ```2025/04/24```: [FastVideo V1](https://hao-ai-lab.github.io/blogs/fastvideo/) is released!
 - ```2025/02/18```: Release the inference code for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/).
 ## Key Features
 FastVideo has the following features:
@@ -91,7 +96,7 @@ For a more detailed guide, please see our [inference quick start](https://hao-ai
 ## Distillation and Finetuning
 - [Distillation Guide](https://hao-ai-lab.github.io/FastVideo/training/distillation.html)
- [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetuning.html)
+- [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetune.html)
 ## 📑 Development Plan
@@ -111,7 +116,7 @@ For a more detailed guide, please see our [inference quick start](https://hao-ai
 ## 🤝 Contributing
-We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/developer_guide/overview.html)
+We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/contributing/overview.html)
 ## Acknowledgement
 We learned and reused code from the following projects:
@@ -128,6 +133,15 @@ We thank MBZUAI and [Anyscale](https://www.anyscale.com/) for their support thro
 If you use FastVideo for your research, please cite our paper:
 ```bibtex
@misc{zhang2025vsafastervideodiffusion,
      title={VSA: Faster Video Diffusion with Trainable Sparse Attention}, 
      author={Peiyuan Zhang and Haofeng Huang and Yongqi Chen and Will Lin and Zhengzhong Liu and Ion Stoica and Eric Xing and Hao Zhang},
      year={2025},
      eprint={2505.13389},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2505.13389}, 
 }
@misc{zhang2025fastvideogenerationsliding,
      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},
@@ -0,0 +1,138 @@
 # ComfyUI-FastVideo
 A custom node suite for ComfyUI that provides accelerated video generation using [FastVideo](https://github.com/hao-ai-labs/FastVideo). See the [blog post](https://hao-ai-lab.github.io/blogs/fastvideo/) about FastVideo V1 to learn more.
 ## Multi-GPU Parallel Inference
 One of the key features ComfyUI-FastVideo brings to ComfyUI is its ability to distribute the generation workload across multiple GPUs, resulting in significantly faster inference times.
 ![Wan2.1-I2V-14B-480P-Diffusers](./assets/wani2v.gif)
 Example of Wan2.1-I2V-14B-480P-Diffusers model running on 4 GPUs.
 ## Features
 - Generate high-quality videos from text prompts and images
 - Configurable video parameters (prompt, resolution, frame count, FPS)
 - Support for multiple GPUs with tensor and sequence parallelism
 - Advanced configuration options for VAE, Text Encoder, and DIT components
 - Interruption/cancellation support for long-running generations
 ## Installation
 ### Requirements
 - [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
 - CUDA-capable GPU(s) with sufficient VRAM
 ### Install using ComfyUI Manager
 Coming soon!
 ### Manual Installation
 #### Copy the FastVideo `comfyui` directory into your ComfyUI custom_nodes directory:
 ```bash
 cp -r /path/to/FastVideo/comfyui /path/to/ComfyUI/custom_nodes/FastVideo
 ```
 #### Install dependencies:
 Currently, the only dependency is `fastvideo`, which can be installed using pip.
 ```bash
 pip install fastvideo
 ```
 #### Install missing custom nodes:
 `ComfyUI-VideoHelperSuite`:
 ```bash
 cd /path/to/ComfyUI/custom_nodes
 git clone https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite.git
 ```
 If you're seeing `ImportError: libGL.so.1: cannot open shared object file: No such file or directory`,
 you may need to install ffmpeg
 ```bash
 apt-get update && apt-get install ffmpeg
 ```
 ## Usage
 After installation, the following nodes will be available in the ComfyUI interface under the "fastvideo" category:
 - **Video Generator**: The main node for generating videos from prompts
 - **Inference Args**: Configure video generation parameters
 - **VAE Config**
 - **Text Encoder Config**
 - **DIT Config**
 - **Load Image Path**: Load images for potential conditioning
 You may have noticed many arguments on the nodes have 'auto' as the default value. This is because FastVideo will automatically detect the best values for these parameters based on the model and the hardware. However, you can also manually configure these parameters to get the best performance for your specific use case. We plan on releasing more optimized workflow files for different models and hardware configurations in the future.
 You can see what some of the default configurations are by looking at the FastVideo repo:
 - [Wan2.1-I2V-14B-480P-Diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/v1/configs/wan_14B_i2v_480p_pipeline.json)
 - [FastHunyuan-diffusers](https://github.com/hao-ai-lab/FastVideo/blob/main/fastvideo/v1/configs/fasthunyuan_t2v.json)
 ### Node Configuration
 #### Video Generator
 - **prompt**: Text description of the video to generate
 - **output_path**: Directory where generated videos will be saved
 - **num_gpus**: Number of GPUs to use for generation
 - **model_path**: Path to the FastVideo model
 - **embedded_cfg_scale**: Classifier-free guidance scale
 - **sp_size**: Sequence parallelism size (usually should match num_gpus)
 - **tp_size**: Tensor parallelism size (usually should match num_gpus)
 - **precision**: Model precision (fp16 or bf16)
 `model_path takes either a model id from huggingface or a local path to a model. Models by default will be downloaded to ~/.cache/huggingface/hub/ and cached for subsequent runs.`
 #### Inference Args
 - **height/width**: Resolution of the output video
 - **num_frames**: Number of frames to generate
 - **num_inference_steps**: Number of diffusion steps per frame
 - **guidance_scale**: Classifier-free guidance scale
 - **flow_shift**: Frame flow shift parameter
 - **seed**: Random seed for reproducible generation
 - **fps**: Frames per second of the output video
 - **image_path**: Optional path to input image for conditioning (for i2v models)
 ## Memory Management
 Models will remain loaded in GPU memory between runs when you only change inference arguments (such as prompt, resolution, frame count, FPS, guidance scale, etc.) or the prompt text. This allows for faster subsequent generations since the model doesn't need to be reloaded.
 However, if you need to change the following parameters, you will need to restart the ComfyUI server:
 - **Number of GPUs** (`num_gpus`)
 - **Model path** (`model_path`)
 - **Tensor parallelism size** (`tp_size`)
 - **Sequence parallelism size** (`sp_size`)
 These parameters affect the model's distribution across GPUs and require a complete reinitialization of the model pipeline.
 ## Example workflows
 ### Text to Video
 FastVideo-FastHunyuan-diffusers
 ![FastVideo-FastHunyuan-diffusers](./assets/fasthunyuan.png)
 - [FastHunyuan-diffusers.json](./examples/FastHunyuan-diffusers.json)
 ### Image to Video
 Wan2.1-I2V-14B-480P-Diffusers
 ![Wan2.1-I2V-14B-480P-Diffusers](./assets/wani2v.png)
 - [Wan2.1-I2V-14B-480P-Diffusers.json](./examples/Wan2.1-I2V-14B-480P-Diffusers.json)
 ## License
 This project is licensed under Apache 2.0.
@@ -0,0 +1,5 @@
 from .video_generator.nodes import (NODE_CLASS_MAPPINGS,
                                    NODE_DISPLAY_NAME_MAPPINGS)
 WEB_DIRECTORY = "./web"
 __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY']
@@ -0,0 +1,645 @@
 {
  "id": "23a1f065-bbba-4a8f-b144-944e1318fcbf",
  "revision": 0,
  "last_node_id": 8,
  "last_link_id": 7,
  "nodes": [
    {
      "id": 4,
      "type": "VAEConfig",
      "pos": [
        374.2159423828125,
        554.85888671875
      ],
      "size": [
        334.080078125,
        322
      ],
      "flags": {},
      "order": 0,
      "mode": 0,
      "inputs": [],
      "outputs": [
        {
          "name": "vae_config",
          "type": "VAE_CONFIG",
          "links": [
            2
          ]
        }
      ],
      "properties": {
        "Node name for S&R": "VAEConfig"
      },
      "widgets_values": [
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999
      ],
      "auto_widget_states": {
        "load_encoder": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": true
        },
        "load_decoder": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": true
        },
        "tile_sample_min_height": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 256
        },
        "tile_sample_min_width": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 256
        },
        "tile_sample_min_num_frames": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 16
        },
        "tile_sample_stride_height": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 192
        },
        "tile_sample_stride_width": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 192
        },
        "tile_sample_stride_num_frames": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 12
        },
        "blend_num_frames": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 0
        },
        "use_tiling": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": true
        },
        "use_temporal_tiling": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": true
        },
        "use_parallel_tiling": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": true
        }
      }
    },
    {
      "id": 5,
      "type": "TextEncoderConfig",
      "pos": [
        416.4937744140625,
        953.6171875
      ],
      "size": [
        270,
        106
      ],
      "flags": {},
      "order": 1,
      "mode": 0,
      "inputs": [],
      "outputs": [
        {
          "name": "text_encoder_config",
          "type": "TEXT_ENCODER_CONFIG",
          "links": [
            7
          ]
        }
      ],
      "properties": {
        "Node name for S&R": "TextEncoderConfig"
      },
      "widgets_values": [
        -99999,
        -99999,
        -99999
      ],
      "auto_widget_states": {
        "prefix": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": ""
        },
        "quant_config": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": ""
        },
        "lora_config": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": ""
        }
      }
    },
    {
      "id": 6,
      "type": "DITConfig",
      "pos": [
        415.1928405761719,
        1154.1573486328125
      ],
      "size": [
        270,
        82
      ],
      "flags": {},
      "order": 2,
      "mode": 0,
      "inputs": [],
      "outputs": [
        {
          "name": "dit_config",
          "type": "DIT_CONFIG",
          "links": [
            6
          ]
        }
      ],
      "properties": {
        "Node name for S&R": "DITConfig"
      },
      "widgets_values": [
        -99999,
        -99999
      ],
      "auto_widget_states": {
        "prefix": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": ""
        },
        "quant_config": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": ""
        }
      }
    },
    {
      "id": 1,
      "type": "VideoGenerator",
      "pos": [
        818.804931640625,
        348.9299621582031
      ],
      "size": [
        400,
        436
      ],
      "flags": {},
      "order": 4,
      "mode": 0,
      "inputs": [
        {
          "name": "inference_args",
          "shape": 7,
          "type": "INFERENCE_ARGS",
          "link": 3
        },
        {
          "name": "vae_config",
          "shape": 7,
          "type": "VAE_CONFIG",
          "link": 2
        },
        {
          "name": "text_encoder_config",
          "shape": 7,
          "type": "TEXT_ENCODER_CONFIG",
          "link": 7
        },
        {
          "name": "dit_config",
          "shape": 7,
          "type": "DIT_CONFIG",
          "link": 6
        }
      ],
      "outputs": [
        {
          "name": "video_path",
          "type": "STRING",
          "links": [
            4
          ]
        }
      ],
      "properties": {
        "Node name for S&R": "VideoGenerator"
      },
      "widgets_values": [
        "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes wide with interest. The playful yet serene atmosphere is complemented by soft natural light filtering through the petals. Mid-shot, warm and cheerful tones.",
        "/workspace/ComfyUI/outputs_video/",
        2,
        "FastVideo/FastHunyuan-diffusers",
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999,
        -99999
      ],
      "auto_widget_states": {
        "embedded_cfg_scale": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 6
        },
        "sp_size": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 2
        },
        "tp_size": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 2
        },
        "vae_precision": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": "fp16"
        },
        "vae_tiling": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": true
        },
        "vae_sp": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": true
        },
        "text_encoder_precision": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": "fp16"
        },
        "precision": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": "fp16"
        },
        "use_cpu_offload": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": true
        }
      }
    },
    {
      "id": 3,
      "type": "VHS_LoadVideoPath",
      "pos": [
        1350.136962890625,
        331.20361328125
      ],
      "size": [
        231.8896484375,
        286
      ],
      "flags": {},
      "order": 5,
      "mode": 0,
      "inputs": [
        {
          "name": "meta_batch",
          "shape": 7,
          "type": "VHS_BatchManager",
          "link": null
        },
        {
          "name": "vae",
          "shape": 7,
          "type": "VAE",
          "link": null
        },
        {
          "name": "video",
          "type": "STRING",
          "widget": {
            "name": "video"
          },
          "link": 4
        }
      ],
      "outputs": [
        {
          "name": "IMAGE",
          "type": "IMAGE",
          "links": [
            5
          ]
        },
        {
          "name": "frame_count",
          "type": "INT",
          "links": null
        },
        {
          "name": "audio",
          "type": "AUDIO",
          "links": null
        },
        {
          "name": "video_info",
          "type": "VHS_VIDEOINFO",
          "links": null
        }
      ],
      "properties": {
        "Node name for S&R": "VHS_LoadVideoPath"
      },
      "widgets_values": {
        "video": "",
        "force_rate": 0,
        "custom_width": 0,
        "custom_height": 0,
        "frame_load_cap": 0,
        "skip_first_frames": 0,
        "select_every_nth": 1,
        "format": "Wan",
        "videopreview": {
          "hidden": false,
          "paused": false,
          "params": {
            "filename": "",
            "type": "path",
            "format": "video/",
            "force_rate": 0,
            "custom_width": 0,
            "custom_height": 0,
            "frame_load_cap": 0,
            "skip_first_frames": 0,
            "select_every_nth": 1
          }
        }
      }
    },
    {
      "id": 2,
      "type": "InferenceArgs",
      "pos": [
        411.46307373046875,
        178.18182373046875
      ],
      "size": [
        278.73828125,
        298
      ],
      "flags": {},
      "order": 3,
      "mode": 0,
      "inputs": [],
      "outputs": [
        {
          "name": "inference_args",
          "type": "INFERENCE_ARGS",
          "links": [
            3
          ]
        }
      ],
      "properties": {
        "Node name for S&R": "InferenceArgs"
      },
      "widgets_values": [
        720,
        1280,
        45,
        6,
        -99999,
        -99999,
        1025,
        "fixed",
        24,
        -99999,
        -99999
      ],
      "auto_widget_states": {
        "height": {
          "isAuto": false,
          "value": 720,
          "cachedValue": 720
        },
        "width": {
          "isAuto": false,
          "value": 1280,
          "cachedValue": 1280
        },
        "num_frames": {
          "isAuto": false,
          "value": 45,
          "cachedValue": 45
        },
        "num_inference_steps": {
          "isAuto": false,
          "value": 6,
          "cachedValue": 6
        },
        "guidance_scale": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 1
        },
        "flow_shift": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": 17
        },
        "seed": {
          "isAuto": false,
          "value": 1025,
          "cachedValue": 1024
        },
        "fps": {
          "isAuto": false,
          "value": 24,
          "cachedValue": 24
        },
        "image_path": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": "X://insert/path/here.mp4"
        },
        "enable_teacache": {
          "isAuto": true,
          "value": -99999,
          "cachedValue": true
        }
      }
    },
    {
      "id": 8,
      "type": "VHS_VideoCombine",
      "pos": [
        1668.3499755859375,
        328.22625732421875
      ],
      "size": [
        507.507080078125,
        622.2227172851562
      ],
      "flags": {},
      "order": 6,
      "mode": 0,
      "inputs": [
        {
          "name": "images",
          "type": "IMAGE",
          "link": 5
        },
        {
          "name": "audio",
          "shape": 7,
          "type": "AUDIO",
          "link": null
        },
        {
          "name": "meta_batch",
          "shape": 7,
          "type": "VHS_BatchManager",
          "link": null
        },
        {
          "name": "vae",
          "shape": 7,
          "type": "VAE",
          "link": null
        }
      ],
      "outputs": [
        {
          "name": "Filenames",
          "type": "VHS_FILENAMES",
          "links": null
        }
      ],
      "properties": {
        "Node name for S&R": "VHS_VideoCombine"
      },
      "widgets_values": {
        "frame_rate": 24,
        "loop_count": 0,
        "filename_prefix": "",
        "format": "video/h264-mp4",
        "pix_fmt": "yuv420p",
        "crf": 19,
        "save_metadata": true,
        "trim_to_audio": false,
        "pingpong": false,
        "save_output": false,
        "videopreview": {
          "hidden": false,
          "paused": false,
          "params": {
            "filename": "._00003.mp4",
            "subfolder": "",
            "type": "temp",
            "format": "video/h264-mp4",
            "frame_rate": 24,
            "workflow": "._00003.png",
            "fullpath": "/workspace/ComfyUI/temp/._00003.mp4"
          }
        }
      }
    }
  ],
  "links": [
    [
      2,
      4,
      0,
      1,
      1,
      "VAE_CONFIG"
    ],
    [
      3,
      2,
      0,
      1,
      0,
      "INFERENCE_ARGS"
    ],
    [
      4,
      1,
      0,
      3,
      2,
      "STRING"
    ],
    [
      5,
      3,
      0,
      8,
      0,
      "IMAGE"
    ],
    [
      6,
      6,
      0,
      1,
      3,
      "DIT_CONFIG"
    ],
    [
      7,
      5,
      0,
      1,
      2,
      "TEXT_ENCODER_CONFIG"
    ]
  ],
  "groups": [],
  "config": {},
  "extra": {
    "ds": {
      "scale": 0.9090909090909091,
      "offset": [
        112.86678372727341,
        -71.45635903989245
      ]
    },
    "frontendVersion": "1.20.4",
    "VHS_latentpreview": false,
    "VHS_latentpreviewrate": 0,
    "VHS_MetadataImage": true,
    "VHS_KeepIntermediate": true
  },
  "version": 0.4
 }
@@ -0,0 +1,697 @@
 {
  "id": "23a1f065-bbba-4a8f-b144-944e1318fcbf",
  "revision": 0,
  "last_node_id": 8,
  "last_link_id": 7,
  "nodes": [
    {
      "id": 7,
      "type": "LoadImagePath",
      "pos": [
        33.15385437011719,
        191.2037353515625
      ],
      "size": [
        270,
        334
      ],
      "flags": {},
      "order": 0,
      "mode": 0,
      "inputs": [],
      "outputs": [
        {
          "name": "image_path",
          "type": "STRING",
          "links": [
            1
          ]
        },
        {
          "name": "IMAGE",
          "type": "IMAGE",
          "links": null
        },
        {
          "name": "MASK",
          "type": "MASK",
          "links": null
        }
      ],
      "properties": {
        "Node name for S&R": "LoadImagePath"
      },
      "widgets_values": [
        "woman.jpg",
        "image"
      ]
    },
    {
      "id": 3,
      "type": "VHS_LoadVideoPath",
      "pos": [
        1350.136962890625,
        331.20361328125
      ],
      "size": [
        231.8896484375,
        286
      ],
      "flags": {},
      "order": 6,
      "mode": 0,
      "inputs": [
        {
          "name": "meta_batch",
          "shape": 7,
          "type": "VHS_BatchManager",
          "link": null
        },
        {
          "name": "vae",
          "shape": 7,
          "type": "VAE",
          "link": null
        },
        {
          "name": "video",
          "type": "STRING",
          "widget": {
            "name": "video"
          },
          "link": 4
        }
      ],
      "outputs": [
        {
          "name": "IMAGE",
          "type": "IMAGE",
          "links": [
            5
          ]
        },
        {
          "name": "frame_count",
          "type": "INT",
          "links": null
        },
        {
          "name": "audio",
          "type": "AUDIO",
          "links": null
        },
        {
          "name": "video_info",
          "type": "VHS_VIDEOINFO",
          "links": null
        }
      ],
      "properties": {
        "Node name for S&R": "VHS_LoadVideoPath"
      },
      "widgets_values": {
        "video": "",
        "force_rate": 0,
        "custom_width": 0,
        "custom_height": 0,
        "frame_load_cap": 0,
        "skip_first_frames": 0,
        "select_every_nth": 1,
        "format": "Wan",
        "videopreview": {
          "hidden": false,
          "paused": false,
          "params": {
            "filename": "",
            "type": "path",
            "format": "video/",
            "force_rate": 0,
            "custom_width": 0,
            "custom_height": 0,
            "frame_load_cap": 0,
            "skip_first_frames": 0,
            "select_every_nth": 1
          }
        }
      }
    },
    {
      "id": 8,
      "type": "VHS_VideoCombine",
      "pos": [
        1668.3499755859375,
        328.22625732421875
      ],
      "size": [
        214.7587890625,
        334
      ],
      "flags": {},
      "order": 7,
      "mode": 0,
      "inputs": [
        {
          "name": "images",
          "type": "IMAGE",
          "link": 5
        },
        {
          "name": "audio",
          "shape": 7,
          "type": "AUDIO",
          "link": null
        },
        {
          "name": "meta_batch",
          "shape": 7,
          "type": "VHS_BatchManager",
          "link": null
        },
        {
          "name": "vae",
          "shape": 7,
          "type": "VAE",
          "link": null
        }
      ],
      "outputs": [
        {
          "name": "Filenames",
          "type": "VHS_FILENAMES",
          "links": null
        }
      ],
      "properties": {
        "Node name for S&R": "VHS_VideoCombine"
      },
      "widgets_values": {
        "frame_rate": 24,
        "loop_count": 0,
        "filename_prefix": "",
        "format": "video/h264-mp4",
        "pix_fmt": "yuv420p",
        "crf": 19,
        "save_metadata": true,
        "trim_to_audio": false,
        "pingpong": false,
        "save_output": false,
        "videopreview": {
          "hidden": false,
          "paused": false,
          "params": {}
        }
      }
    },
    {
      "id": 4,
      "type": "VAEConfig",
      "pos": [
        374.2159423828125,
        554.85888671875
      ],
      "size": [
        334.080078125,
        322
      ],
      "flags": {},
      "order": 1,
      "mode": 0,
      "inputs": [],
      "outputs": [
        {
          "name": "vae_config",
          "type": "VAE_CONFIG",
          "links": [
            2
          ]
        }
      ],
      "properties": {
        "Node name for S&R": "VAEConfig"
      },
      "widgets_values": [
        true,
        true,
        256,
        256,
        16,
        192,
        192,
        12,
        0,
        true,
        true,
        true
      ],
      "auto_widget_states": {
        "load_encoder": {
          "isAuto": true,
          "value": true,
          "cachedValue": true
        },
        "load_decoder": {
          "isAuto": true,
          "value": true,
          "cachedValue": true
        },
        "tile_sample_min_height": {
          "isAuto": true,
          "value": 256,
          "cachedValue": 256
        },
        "tile_sample_min_width": {
          "isAuto": true,
          "value": 256,
          "cachedValue": 256
        },
        "tile_sample_min_num_frames": {
          "isAuto": true,
          "value": 16,
          "cachedValue": 16
        },
        "tile_sample_stride_height": {
          "isAuto": true,
          "value": 192,
          "cachedValue": 192
        },
        "tile_sample_stride_width": {
          "isAuto": true,
          "value": 192,
          "cachedValue": 192
        },
        "tile_sample_stride_num_frames": {
          "isAuto": true,
          "value": 12,
          "cachedValue": 12
        },
        "blend_num_frames": {
          "isAuto": true,
          "value": 0,
          "cachedValue": 0
        },
        "use_tiling": {
          "isAuto": true,
          "value": true,
          "cachedValue": true
        },
        "use_temporal_tiling": {
          "isAuto": true,
          "value": true,
          "cachedValue": true
        },
        "use_parallel_tiling": {
          "isAuto": true,
          "value": true,
          "cachedValue": true
        }
      }
    },
    {
      "id": 2,
      "type": "InferenceArgs",
      "pos": [
        411.46307373046875,
        178.18182373046875
      ],
      "size": [
        278.73828125,
        298
      ],
      "flags": {},
      "order": 4,
      "mode": 0,
      "inputs": [
        {
          "name": "image_path",
          "shape": 7,
          "type": "STRING",
          "widget": {
            "name": "image_path"
          },
          "link": 1
        }
      ],
      "outputs": [
        {
          "name": "inference_args",
          "type": "INFERENCE_ARGS",
          "links": [
            3
          ]
        }
      ],
      "properties": {
        "Node name for S&R": "InferenceArgs"
      },
      "widgets_values": [
        832,
        480,
        45,
        20,
        1,
        17,
        1024,
        "fixed",
        24,
        "X://insert/path/here.mp4",
        true
      ],
      "auto_widget_states": {
        "height": {
          "isAuto": false,
          "value": 832,
          "cachedValue": 720
        },
        "width": {
          "isAuto": false,
          "value": 480,
          "cachedValue": 1280
        },
        "num_frames": {
          "isAuto": false,
          "value": 45,
          "cachedValue": 45
        },
        "num_inference_steps": {
          "isAuto": false,
          "value": 20,
          "cachedValue": 6
        },
        "guidance_scale": {
          "isAuto": true,
          "value": 1,
          "cachedValue": 1
        },
        "flow_shift": {
          "isAuto": true,
          "value": 17,
          "cachedValue": 17
        },
        "seed": {
          "isAuto": false,
          "value": 1024,
          "cachedValue": 1024
        },
        "fps": {
          "isAuto": false,
          "value": 24,
          "cachedValue": 24
        },
        "image_path": {
          "isAuto": true,
          "value": "X://insert/path/here.mp4",
          "cachedValue": "X://insert/path/here.mp4"
        },
        "enable_teacache": {
          "isAuto": true,
          "value": true,
          "cachedValue": true
        }
      }
    },
    {
      "id": 1,
      "type": "VideoGenerator",
      "pos": [
        818.804931640625,
        348.9299621582031
      ],
      "size": [
        400,
        436
      ],
      "flags": {},
      "order": 5,
      "mode": 0,
      "inputs": [
        {
          "name": "inference_args",
          "shape": 7,
          "type": "INFERENCE_ARGS",
          "link": 3
        },
        {
          "name": "vae_config",
          "shape": 7,
          "type": "VAE_CONFIG",
          "link": 2
        },
        {
          "name": "text_encoder_config",
          "shape": 7,
          "type": "TEXT_ENCODER_CONFIG",
          "link": 7
        },
        {
          "name": "dit_config",
          "shape": 7,
          "type": "DIT_CONFIG",
          "link": 6
        }
      ],
      "outputs": [
        {
          "name": "video_path",
          "type": "STRING",
          "links": [
            4
          ]
        }
      ],
      "properties": {
        "Node name for S&R": "VideoGenerator"
      },
      "widgets_values": [
        "A woman crying from laughter.",
        "/workspace/ComfyUI/outputs_video/",
        4,
        "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
        6,
        2,
        2,
        "fp16",
        true,
        true,
        "fp16",
        "fp16",
        true
      ],
      "auto_widget_states": {
        "embedded_cfg_scale": {
          "isAuto": true,
          "value": 6,
          "cachedValue": 6
        },
        "sp_size": {
          "isAuto": true,
          "value": 2,
          "cachedValue": 2
        },
        "tp_size": {
          "isAuto": true,
          "value": 2,
          "cachedValue": 2
        },
        "vae_precision": {
          "isAuto": true,
          "value": "fp16",
          "cachedValue": "fp16"
        },
        "vae_tiling": {
          "isAuto": true,
          "value": true,
          "cachedValue": true
        },
        "vae_sp": {
          "isAuto": true,
          "value": true,
          "cachedValue": true
        },
        "text_encoder_precision": {
          "isAuto": true,
          "value": "fp16",
          "cachedValue": "fp16"
        },
        "precision": {
          "isAuto": true,
          "value": "fp16",
          "cachedValue": "fp16"
        },
        "use_cpu_offload": {
          "isAuto": true,
          "value": true,
          "cachedValue": true
        }
      }
    },
    {
      "id": 5,
      "type": "TextEncoderConfig",
      "pos": [
        416.4937744140625,
        953.6171875
      ],
      "size": [
        270,
        106
      ],
      "flags": {},
      "order": 2,
      "mode": 0,
      "inputs": [],
      "outputs": [
        {
          "name": "text_encoder_config",
          "type": "TEXT_ENCODER_CONFIG",
          "links": [
            7
          ]
        }
      ],
      "properties": {
        "Node name for S&R": "TextEncoderConfig"
      },
      "widgets_values": [
        "",
        "",
        ""
      ],
      "auto_widget_states": {
        "prefix": {
          "isAuto": true,
          "value": "",
          "cachedValue": ""
        },
        "quant_config": {
          "isAuto": true,
          "value": "",
          "cachedValue": ""
        },
        "lora_config": {
          "isAuto": true,
          "value": "",
          "cachedValue": ""
        }
      }
    },
    {
      "id": 6,
      "type": "DITConfig",
      "pos": [
        415.1928405761719,
        1154.1573486328125
      ],
      "size": [
        270,
        82
      ],
      "flags": {},
      "order": 3,
      "mode": 0,
      "inputs": [],
      "outputs": [
        {
          "name": "dit_config",
          "type": "DIT_CONFIG",
          "links": [
            6
          ]
        }
      ],
      "properties": {
        "Node name for S&R": "DITConfig"
      },
      "widgets_values": [
        "",
        ""
      ],
      "auto_widget_states": {
        "prefix": {
          "isAuto": true,
          "value": "",
          "cachedValue": ""
        },
        "quant_config": {
          "isAuto": true,
          "value": "",
          "cachedValue": ""
        }
      }
    }
  ],
  "links": [
    [
      1,
      7,
      0,
      2,
      0,
      "STRING"
    ],
    [
      2,
      4,
      0,
      1,
      1,
      "VAE_CONFIG"
    ],
    [
      3,
      2,
      0,
      1,
      0,
      "INFERENCE_ARGS"
    ],
    [
      4,
      1,
      0,
      3,
      2,
      "STRING"
    ],
    [
      5,
      3,
      0,
      8,
      0,
      "IMAGE"
    ],
    [
      6,
      6,
      0,
      1,
      3,
      "DIT_CONFIG"
    ],
    [
      7,
      5,
      0,
      1,
      2,
      "TEXT_ENCODER_CONFIG"
    ]
  ],
  "groups": [],
  "config": {},
  "extra": {
    "ds": {
      "scale": 0.8264462809917354,
      "offset": [
        646.7950212991898,
        66.17259910028655
      ]
    },
    "frontendVersion": "1.20.4",
    "VHS_latentpreview": false,
    "VHS_latentpreviewrate": 0,
    "VHS_MetadataImage": true,
    "VHS_KeepIntermediate": true
  },
  "version": 0.4
 }
@@ -0,0 +1,31 @@
 class DITConfig:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "optional": {
                "prefix": ("STRING", {
                    "default": ""
                }),
                "quant_config": ("STRING", {
                    "default": ""
                }),
            }
        }
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        return True
    RETURN_TYPES = ("DIT_CONFIG", )
    RETURN_NAMES = ("dit_config", )
    FUNCTION = "set_args"
    CATEGORY = "fastvideo"
    def set_args(self, prefix, quant_config):
        raw_args = {"prefix": prefix, "quant_config": quant_config}
        # Filter out keys where value is -99999
        args = {k: v for k, v in raw_args.items() if str(int(v)) != str(-99999)}
        return (args, )
@@ -0,0 +1,89 @@
 class InferenceArgs:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "optional": {
                "height": ("INT", {
                    "default": 720
                }),
                "width": ("INT", {
                    "default": 1280
                }),
                "num_frames": ("INT", {
                    "default": 45
                }),
                "num_inference_steps": ("INT", {
                    "default": 6
                }),
                "guidance_scale": ("FLOAT", {
                    "default": 1.0
                }),
                "flow_shift": ("INT", {
                    "default": 17
                }),
                "seed": ("INT", {
                    "default": 1024
                }),
                "fps": ("INT", {
                    "default": 24
                }),
                "image_path": ("STRING", {
                    "default": "X://insert/path/here.mp4"
                }),
                "enable_teacache": ([True, False], {
                    "default": False
                }),
            }
        }
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        return True
    RETURN_TYPES = ("INFERENCE_ARGS", )
    RETURN_NAMES = ("inference_args", )
    FUNCTION = "set_args"
    CATEGORY = "fastvideo"
    def set_args(
        self,
        height,
        width,
        num_frames,
        num_inference_steps,
        guidance_scale,
        flow_shift,
        seed,
        fps,
        image_path,
        enable_teacache,
    ):
        raw_args = {
            "height": height,
            "width": width,
            "num_frames": num_frames,
            "num_inference_steps": num_inference_steps,
            "guidance_scale": guidance_scale,
            "flow_shift": flow_shift,
            "seed": seed,
            "fps": fps,
            "image_path": image_path,
            "enable_teacache": enable_teacache,
        }
        # Filter out keys where value is -99999, handling different types properly
        args = {}
        for k, v in raw_args.items():
            try:
                if isinstance(v, str):
                    if v != "-99999":
                        args[k] = v
                elif v != -99999:
                    # If it's not a string, compare directly
                    args[k] = v
            except (ValueError, TypeError):
                # Include any value that causes an error in comparison
                args[k] = v
        return (args, )
@@ -0,0 +1,103 @@
 import hashlib
 import os
 import folder_paths
 import numpy as np
 import torch
 from PIL import Image, ImageOps, ImageSequence
 from .node_helpers import pillow
 class LoadImagePath:
    @classmethod
    def INPUT_TYPES(s):
        input_dir = folder_paths.get_input_directory()
        files = [
            f for f in os.listdir(input_dir)
            if os.path.isfile(os.path.join(input_dir, f))
        ]
        files = folder_paths.filter_files_content_types(files, ["image"])
        return {
            "required": {
                "image": (sorted(files), {
                    "image_upload": True
                })
            },
        }
    CATEGORY = "fastvideo"
    RETURN_TYPES = ("STRING", "IMAGE", "MASK")
    RETURN_NAMES = ("image_path", "IMAGE", "MASK")
    FUNCTION = "load_image"
    def load_image(self, image):
        image_path = folder_paths.get_annotated_filepath(image)
        img = pillow(Image.open, image_path)
        output_images: list[torch.Tensor] = []
        output_masks: list[torch.Tensor] = []
        w, h = None, None
        excluded_formats = ['MPO']
        for i in ImageSequence.Iterator(img):
            processed_image = pillow(ImageOps.exif_transpose, i)
            if processed_image is None:
                continue
            if processed_image.mode == 'I':
                processed_image = processed_image.point(lambda i: i * (1 / 255))
            image = processed_image.convert("RGB")
            if len(output_images) == 0:
                w = image.size[0]
                h = image.size[1]
            if image.size[0] != w or image.size[1] != h:
                continue
            image = np.array(image).astype(np.float32) / 255.0
            image = torch.from_numpy(image)[
                None,
            ]
            if 'A' in processed_image.getbands():
                mask = np.array(processed_image.getchannel('A')).astype(
                    np.float32) / 255.0
                mask = 1. - torch.from_numpy(mask)
            elif processed_image.mode == 'P' and 'transparency' in processed_image.info:
                mask = np.array(
                    processed_image.convert('RGBA').getchannel('A')).astype(
                        np.float32) / 255.0
                mask = 1. - torch.from_numpy(mask)
            else:
                mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
            output_images.append(image)
            output_masks.append(mask.unsqueeze(0))
        if len(output_images) > 1 and img.format not in excluded_formats:
            output_image = torch.cat(output_images, dim=0)
            output_mask = torch.cat(output_masks, dim=0)
        else:
            output_image = output_images[0]
            output_mask = output_masks[0]
        return (image_path, output_image, output_mask)
    @classmethod
    def IS_CHANGED(s, image):
        image_path = folder_paths.get_annotated_filepath(image)
        m = hashlib.sha256()
        with open(image_path, 'rb') as f:
            m.update(f.read())
        return m.digest().hex()
    @classmethod
    def VALIDATE_INPUTS(s, image):
        if not folder_paths.exists_annotated_filepath(image):
            return "Invalid image file: {}".format(image)
        return True
@@ -0,0 +1,68 @@
 import hashlib
 from collections.abc import Callable
 from typing import Any, TypeVar
 import torch
 from comfy.cli_args import args
 from PIL import ImageFile, UnidentifiedImageError
 T = TypeVar('T')
 def conditioning_set_values(conditioning: list[Any],
                            values: dict[str, Any] | None = None) -> list[Any]:
    if values is None:
        values = {}
    c = []
    for t in conditioning:
        n = [t[0], t[1].copy()]
        for k in values:
            n[1][k] = values[k]
        c.append(n)
    return c
 def pillow(fn: Callable[[Any], T], arg: Any) -> T:
    prev_value = None
    try:
        x = fn(arg)
    except (OSError, UnidentifiedImageError, ValueError
            ):  #PIL issues #4472 and #2445, also fixes ComfyUI issue #3416
        prev_value = ImageFile.LOAD_TRUNCATED_IMAGES
        ImageFile.LOAD_TRUNCATED_IMAGES = True
        x = fn(arg)
    finally:
        if prev_value is not None:
            ImageFile.LOAD_TRUNCATED_IMAGES = prev_value
    return x
 def hasher() -> Callable[[], Any]:
    hashfuncs = {
        "md5": hashlib.md5,
        "sha1": hashlib.sha1,
        "sha256": hashlib.sha256,
        "sha512": hashlib.sha512
    }
    return hashfuncs[args.default_hashing_function]
 def string_to_torch_dtype(string: str) -> torch.dtype | None:
    if string == "fp32":
        return torch.float32
    if string == "fp16":
        return torch.float16
    if string == "bf16":
        return torch.bfloat16
    return None
 def image_alpha_fix(destination: torch.Tensor,
                    source: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    if destination.shape[-1] < source.shape[-1]:
        source = source[..., :destination.shape[-1]]
    elif destination.shape[-1] > source.shape[-1]:
        destination = torch.nn.functional.pad(destination, (0, 1))
        destination[..., -1] = 1.0
    return destination, source
@@ -0,0 +1,24 @@
 from .dit_config import DITConfig
 from .inference_args import InferenceArgs
 from .load_image import LoadImagePath
 from .text_encoder_config import TextEncoderConfig
 from .vae_config import VAEConfig
 from .video_generator import VideoGenerator
 NODE_CLASS_MAPPINGS = {
    "VideoGenerator": VideoGenerator,
    "InferenceArgs": InferenceArgs,
    "VAEConfig": VAEConfig,
    "TextEncoderConfig": TextEncoderConfig,
    "DITConfig": DITConfig,
    "LoadImagePath": LoadImagePath
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
    "VideoGenerator": "Video Generator",
    "InferenceArgs": "Inference Args",
    "VAEConfig": "VAE Config",
    "TextEncoderConfig": "Text Encoder Config",
    "DITConfig": "DIT Config",
    "LoadImagePath": "Load Image Path"
 }
@@ -0,0 +1,38 @@
 class TextEncoderConfig:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "optional": {
                "prefix": ("STRING", {
                    "default": ""
                }),
                "quant_config": ("STRING", {
                    "default": ""
                }),
                "lora_config": ("STRING", {
                    "default": ""
                }),
            }
        }
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        return True
    RETURN_TYPES = ("TEXT_ENCODER_CONFIG", )
    RETURN_NAMES = ("text_encoder_config", )
    FUNCTION = "set_args"
    CATEGORY = "fastvideo"
    def set_args(self, prefix, quant_config, lora_config):
        raw_args = {
            "prefix": prefix,
            "quant_config": quant_config,
            "lora_config": lora_config
        }
        # Filter out keys where value is -99999
        args = {k: v for k, v in raw_args.items() if str(int(v)) != str(-99999)}
        return (args, )
@@ -0,0 +1,88 @@
 class VAEConfig:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "optional": {
                "load_encoder": ([True, False], {
                    "default": True
                }),
                "load_decoder": ([True, False], {
                    "default": True
                }),
                "tile_sample_min_height": ("INT", {
                    "default": 256
                }),
                "tile_sample_min_width": ("INT", {
                    "default": 256
                }),
                "tile_sample_min_num_frames": ("INT", {
                    "default": 16
                }),
                "tile_sample_stride_height": ("INT", {
                    "default": 192
                }),
                "tile_sample_stride_width": ("INT", {
                    "default": 192
                }),
                "tile_sample_stride_num_frames": ("INT", {
                    "default": 12
                }),
                "blend_num_frames": ("INT", {
                    "default": 0
                }),
                "use_tiling": ([True, False], {
                    "default": True
                }),
                "use_temporal_tiling": ([True, False], {
                    "default": True
                }),
                "use_parallel_tiling": ([True, False], {
                    "default": True
                }),
            }
        }
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        return True
    RETURN_TYPES = ("VAE_CONFIG", )
    RETURN_NAMES = ("vae_config", )
    FUNCTION = "set_args"
    CATEGORY = "fastvideo"
    def set_args(
        self,
        load_encoder,
        load_decoder,
        tile_sample_min_height,
        tile_sample_min_width,
        tile_sample_min_num_frames,
        tile_sample_stride_height,
        tile_sample_stride_width,
        tile_sample_stride_num_frames,
        blend_num_frames,
        use_tiling,
        use_temporal_tiling,
        use_parallel_tiling,
    ):
        raw_args = {
            "load_encoder": load_encoder,
            "load_decoder": load_decoder,
            "tile_sample_min_height": tile_sample_min_height,
            "tile_sample_min_width": tile_sample_min_width,
            "tile_sample_min_num_frames": tile_sample_min_num_frames,
            "tile_sample_stride_height": tile_sample_stride_height,
            "tile_sample_stride_width": tile_sample_stride_width,
            "tile_sample_stride_num_frames": tile_sample_stride_num_frames,
            "blend_num_frames": blend_num_frames,
            "use_tiling": use_tiling,
            "use_temporal_tiling": use_temporal_tiling,
            "use_parallel_tiling": use_parallel_tiling,
        }
        # Filter out any value explicitly set to -99999
        args = {k: v for k, v in raw_args.items() if str(int(v)) != str(-99999)}
        return (args, )
@@ -0,0 +1,315 @@
 from __future__ import annotations
 import glob
 import os
 import signal
 import sys
 import threading
 import time
 from typing import Any
 from comfy.model_management import processing_interrupted
 from fastvideo import PipelineConfig
 from fastvideo import VideoGenerator as FastVideoGenerator
 sys.path.insert(
    0,
    os.path.dirname(
        os.path.dirname(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
 # Custom exception for interruption
 class GenerationInterruptedException(Exception):
    pass
 # Custom exception for interruption that ComfyUI will recognize
 class GenerationCancelledException(Exception):
    def __init__(self,
                 message: str = "Generation was cancelled by user") -> None:
        self.message = message
        super().__init__(self.message)
 def update_config_from_args(config: Any, args_dict: dict[str, Any]) -> None:
    """
    Update configuration object from arguments dictionary.
    Args:
        config: The configuration object to update
        args_dict: Dictionary containing arguments
    """
    for key, value in args_dict.items():
        if hasattr(config, key) and value is not None:
            if key == "text_encoder_precisions" and isinstance(value, list):
                setattr(config, key, tuple(value))
            else:
                setattr(config, key, value)
 class VideoGenerator:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "prompt": ("STRING", {
                    "multiline":
                    True,
                    "default":
                    "A ripe orange tumbles gently from a tree and lands on the head of a lounging capybara, "
                    "who blinks slowly in response. The moment is quietly humorous and oddly serene, framed by "
                    "lush green foliage and dappled sunlight. Mid-shot, warm and whimsical tones."
                }),
                "output_path": ("STRING", {
                    "default": "/workspace/ComfyUI/outputs_video/"
                }),
                "num_gpus": ("INT", {
                    "default": 2,
                    "min": 1,
                    "max": 16
                }),
                "model_path": ("STRING", {
                    "default": "FastVideo/FastHunyuan-diffusers"
                })
            },
            "optional": {
                "inference_args": ("INFERENCE_ARGS", ),
                "embedded_cfg_scale": ("FLOAT", {
                    "default": 6.0
                }),
                "sp_size": ("INT", {
                    "default": 2
                }),
                "tp_size": ("INT", {
                    "default": 2
                }),
                "vae_config": ("VAE_CONFIG", ),
                "vae_precision": (["fp16", "bf16"], {
                    "default": "fp16"
                }),
                "vae_tiling": ([True, False], {
                    "default": True
                }),
                "vae_sp": ([True, False], {
                    "default": False
                }),
                "text_encoder_config": ("TEXT_ENCODER_CONFIG", ),
                "text_encoder_precision": (["fp16", "bf16"], {
                    "default": "fp16"
                }),
                "dit_config": ("DIT_CONFIG", ),
                "precision": (["fp16", "bf16"], {
                    "default": "fp16"
                }),
                "use_cpu_offload": ([True, False], {
                    "default": False
                }),
            }
        }
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        return True
    RETURN_TYPES = ("STRING", )
    RETURN_NAMES = ("video_path", )
    FUNCTION = "launch_inference"
    CATEGORY = "fastvideo"
    generator: FastVideoGenerator | None = None
    _interrupt_thread: threading.Thread | None = None
    _generation_active: bool = False
    _generation_interrupted: bool = False
    _interrupt_event: threading.Event = threading.Event()
    _generation_thread: threading.Thread | None = None
    _generation_result: str | None = None
    _generation_exception: Exception | None = None
    def _monitor_for_interruption(self):
        """Background thread that monitors for interruption requests"""
        time.sleep(2)  # Give the generation thread time to send execute_forward
        while self._generation_active and not self._interrupt_event.is_set():
            if processing_interrupted():
                print("Video generation interrupted by user")
                self._generation_interrupted = True
                # Try to send interrupt signal to worker processes
                if self.generator is not None and hasattr(
                        self.generator, 'executor'):
                    try:
                        # The MultiprocExecutor has a workers attribute
                        if hasattr(self.generator.executor, 'workers'):
                            for worker in self.generator.executor.workers:
                                if worker.is_alive():
                                    os.kill(worker.pid, signal.SIGINT)
                            print("Interrupt signal sent to worker processes")
                    except Exception as e:
                        print(f"Error sending interrupt signal: {e}")
                # Set the interrupt event to notify other threads
                self._interrupt_event.set()
                break
            time.sleep(0.5)
    def _run_generation(self, prompt: str, output_path: str,
                        inference_args: dict[str, Any]) -> None:
        """Thread function to run the generation"""
        try:
            if self.generator is not None:
                self.generator.generate_video(prompt=prompt,
                                              output_path=output_path,
                                              **inference_args)
                self._generation_result = os.path.join(output_path,
                                                       f"{prompt[:100]}.mp4")
            else:
                raise RuntimeError("Generator is not initialized")
        except Exception as e:
            self._generation_exception = e
            self._interrupt_event.set()
    def load_output_video(self, output_dir):
        video_extensions = ["*.mp4", "*.avi", "*.mov", "*.mkv"]
        video_files = []
        for ext in video_extensions:
            video_files.extend(glob.glob(os.path.join(output_dir, ext)))
        if not video_files:
            print("No video files found in output directory: %s", output_dir)
            return ""
        video_files.sort()
        return video_files[0]
    def launch_inference(
        self,
        prompt,
        output_path,
        num_gpus,
        model_path,
        embedded_cfg_scale,
        sp_size,
        tp_size,
        vae_precision,
        vae_tiling,
        vae_sp,
        text_encoder_precision,
        precision,
        inference_args=None,
        vae_config=None,
        text_encoder_config=None,
        dit_config=None,
        use_cpu_offload=None,
    ):
        print('Running FastVideo inference')
        # Reset interruption flag and event
        self._generation_interrupted = False
        self._interrupt_event.clear()
        self._generation_result = None
        self._generation_exception = None
        # Load pipeline config from model path
        pipeline_config = PipelineConfig.from_pretrained(model_path)
        print('pipeline_config', pipeline_config)
        # Update configs with provided config dictionaries
        if dit_config is not None:
            update_config_from_args(pipeline_config.dit_config, dit_config)
        if vae_config is not None:
            update_config_from_args(pipeline_config.vae_config, vae_config)
        if text_encoder_config is not None:
            update_config_from_args(pipeline_config.text_encoder_configs,
                                    text_encoder_config)
        # Update top-level pipeline config with remaining arguments
        raw_pipeline_args = {}
        if embedded_cfg_scale is not None:
            raw_pipeline_args['embedded_cfg_scale'] = embedded_cfg_scale
        if precision is not None:
            raw_pipeline_args['precision'] = precision
        if vae_precision is not None:
            raw_pipeline_args['vae_precision'] = vae_precision
        if vae_tiling is not None:
            raw_pipeline_args['vae_tiling'] = vae_tiling
        if vae_sp is not None:
            raw_pipeline_args['vae_sp'] = vae_sp
        if text_encoder_precision is not None:
            raw_pipeline_args['text_encoder_precision'] = text_encoder_precision
        # Filter out any value explicitly set to -99999 (auto values)
        pipeline_args = {
            k: v
            for k, v in raw_pipeline_args.items() if str(int(v)) != str(-99999)
        }
        update_config_from_args(pipeline_config, pipeline_args)
        raw_generation_args = {}
        if num_gpus is not None:
            raw_generation_args['num_gpus'] = num_gpus
        if tp_size is not None:
            raw_generation_args['tp_size'] = tp_size
        if sp_size is not None:
            raw_generation_args['sp_size'] = sp_size
        if use_cpu_offload is not None:
            raw_generation_args['use_cpu_offload'] = use_cpu_offload
        generation_args = {
            k: v
            for k, v in raw_generation_args.items()
            if str(int(v)) != str(-99999)
        }
        if self.generator is None:
            print('generation_args', generation_args)
            print('pipeline_config', pipeline_config)
            self.generator = FastVideoGenerator.from_pretrained(
                model_path=model_path,
                **generation_args,
                pipeline_config=pipeline_config)
        print('inference_args', inference_args)
        # Start a thread to run the generation
        self._generation_thread = threading.Thread(target=self._run_generation,
                                                   args=(prompt, output_path,
                                                         inference_args),
                                                   daemon=True)
        self._generation_thread.start()
        # Start a background thread to monitor for interruptions
        self._generation_active = True
        self._interrupt_thread = threading.Thread(
            target=self._monitor_for_interruption, daemon=True)
        self._interrupt_thread.start()
        # Wait for either completion or interruption
        while self._generation_thread.is_alive(
        ) and not self._interrupt_event.is_set():
            self._generation_thread.join(timeout=0.5)
        self._generation_active = False
        if self._interrupt_thread:
            self._interrupt_thread.join(timeout=1.0)
            self._interrupt_thread = None
        if self._generation_interrupted:
            print("Video generation was cancelled by user")
            raise GenerationCancelledException()
        elif self._generation_exception:
            # Re-raise the exception from the generation thread
            raise self._generation_exception
        elif self._generation_result:
            return (self._generation_result, )
        else:
            # This shouldn't happen, but just in case
            print("Generation completed but no result was produced")
            raise Exception("Generation failed to produce a result")
@@ -0,0 +1,593 @@
 import { app } from '../../../scripts/app.js'
 function chainCallback(object, property, callback) {
    if (object == undefined) {
        console.error("Tried to add callback to non-existent object");
        return;
    }
    if (property in object && object[property]) {
        const callback_orig = object[property];
        object[property] = function () {
            const r = callback_orig.apply(this, arguments);
            return callback.apply(this, arguments) ?? r;
        };
    } else {
        object[property] = callback;
    }
 }
 function drawAutoAnnotated(ctx, node, widget_width, y, H) {
    const litegraph_base = LiteGraph;
    const show_text = app.canvas.ds.scale >= 0.5;
    const margin = 15;
    const autoTextWidth = 30;
    const autoTextRightMargin = 5;
    ctx.textAlign = 'left';
    ctx.strokeStyle = litegraph_base.WIDGET_OUTLINE_COLOR;
    ctx.fillStyle = litegraph_base.WIDGET_BGCOLOR;
    ctx.beginPath();
    if (show_text && ctx.roundRect) {
        ctx.roundRect(margin, y, widget_width - margin * 2, H, [H * 0.5]);
    } else {
        ctx.rect(margin, y, widget_width - margin * 2, H);
    }
    ctx.fill();
    if (show_text) {
        if (!this.disabled) ctx.stroke();
        const isAuto = this.isAuto === true;
        ctx.save();
        if (isAuto) {
            ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR;
            ctx.strokeStyle = litegraph_base.WIDGET_TEXT_COLOR;
        } else {
            ctx.fillStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
            ctx.strokeStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
        }
        // Position for the cog
        const cogX = widget_width - autoTextRightMargin - autoTextWidth - 6;
        const cogY = y + H * 0.5;
        const cogRadius = 6;
        const toothLength = 2;
        const numTeeth = 8;
        const holeRadius = 2; // Radius of the center hole
        // Draw the cog
        ctx.beginPath();
        ctx.arc(cogX, cogY, cogRadius - toothLength, 0, Math.PI * 2);
        ctx.fill();
        // Draw the center hole (by clearing it)
        ctx.beginPath();
        ctx.arc(cogX, cogY, holeRadius, 0, Math.PI * 2);
        ctx.fillStyle = litegraph_base.WIDGET_BGCOLOR;
        ctx.fill();
        // Reset fill style for the teeth
        if (isAuto) {
            ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR;
        } else {
            ctx.fillStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
        }
        // Draw teeth
        ctx.beginPath();
        for (let i = 0; i < numTeeth; i++) {
            const angle = (i / numTeeth) * Math.PI * 2;
            const innerX = cogX + (cogRadius - toothLength) * Math.cos(angle);
            const innerY = cogY + (cogRadius - toothLength) * Math.sin(angle);
            const outerX = cogX + cogRadius * Math.cos(angle);
            const outerY = cogY + cogRadius * Math.sin(angle);
            ctx.moveTo(innerX, innerY);
            ctx.lineTo(outerX, outerY);
        }
        ctx.lineWidth = 2;
        ctx.stroke();
        ctx.restore();
        // Draw label
        ctx.fillStyle = litegraph_base.WIDGET_SECONDARY_TEXT_COLOR;
        const label = this.label || this.name;
        if (label != null) {
            ctx.fillText(label, margin * 2 + 5, y + H * 0.7);
        }
        // Draw value
        ctx.textAlign = 'right';
        const text = isAuto ? "auto" : this.displayValue();
        ctx.fillStyle = isAuto ? litegraph_base.WIDGET_SECONDARY_TEXT_COLOR : litegraph_base.WIDGET_TEXT_COLOR;
        ctx.fillText(text, widget_width - autoTextRightMargin - autoTextWidth - 15, y + H * 0.7);
        // Draw increment/decrement buttons if not in AUTO mode and not a string widget
        if (!isAuto && !this.disabled && this.config[0] !== "FVAUTOSTRING") {
            // Draw decrement button (left triangle)
            ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR;
            ctx.beginPath();
            ctx.moveTo(margin + 16, y + 5);
            ctx.lineTo(margin + 6, y + H * 0.5);
            ctx.lineTo(margin + 16, y + H - 5);
            ctx.fill();
            // Draw increment button (right triangle)
            ctx.beginPath();
            ctx.moveTo(widget_width - margin - 16, y + 5);
            ctx.lineTo(widget_width - margin - 6, y + H * 0.5);
            ctx.lineTo(widget_width - margin - 16, y + H - 5);
            ctx.fill();
        }
    }
 }
 function mouseAutoAnnotated(event, [x, y], node) {
    const widget_width = node.size[0];
    const margin = 15;
    const H = 20; // Widget height
    const autoTextWidth = 30;
    const autoTextRightMargin = 5;
    const cogRadius = 6;
    if (this.isAuto) {
        if (event.type === "pointerup" || event.type === "mouseup") {
            const cogX = widget_width - autoTextRightMargin - autoTextWidth - 6;
            const cogLeftEdge = cogX - cogRadius;
            const cogRightEdge = cogX + cogRadius;
            if (x > cogLeftEdge && x < cogRightEdge) {
                this.isAuto = false;
                this.value = this.cachedValue !== undefined ? this.cachedValue : (this.options.default || 0);
                if (this.callback) {
                    this.callback(this.value);
                }
                node.graph.setDirtyCanvas(true, false);
            }
        }
        // Block ALL events in auto mode except cog clicks
        event.preventDefault?.();
        event.stopPropagation?.();
        event.stopImmediatePropagation?.();
        return true; // Always return true to indicate event was handled
    }
    // Determine if clicking on increment/decrement buttons
    const delta = this.config[0] === "FVAUTOSTRING" ? 0 :
        (x < 40 ? -1 : x > widget_width - 48 ? 1 : 0);
    if (event.type === "pointerdown" || event.type === "mousedown") {
        // ComfyUI appears to intercept pointerdown events, so this code path is never reached
        console.log("pointerdown received (unexpected)");
        return false;
    } else if (event.type === "pointerup" || event.type === "mouseup") {
        // Stop event propagation to prevent double handling
        event.preventDefault?.();
        event.stopPropagation?.();
        event.stopImmediatePropagation?.();
        const cogX = widget_width - autoTextRightMargin - autoTextWidth - 6;
        const cogLeftEdge = widget_width - autoTextRightMargin - autoTextWidth - 6 - cogRadius;
        const cogRightEdge = widget_width - autoTextRightMargin - autoTextWidth - 6 + cogRadius;
        if (x > cogLeftEdge && x < cogRightEdge) {
            this.isAuto = !this.isAuto;
            if (this.isAuto) {
                this.cachedValue = this.value;
                this.value = -99999;
            } else {
                this.value = this.cachedValue !== undefined ? this.cachedValue : (this.options.default || 0);
            }
            if (this.callback) {
                this.callback(this.value);
            }
            node.graph.setDirtyCanvas(true, false);
            return true;
        }
        // If in auto mode and NOT clicking the cog, block all other interactions
        if (this.isAuto) {
            return true;
        }
        // Handle increment/decrement buttons if not in auto mode
        if (delta !== 0 && !this.isAuto) {
            if (this.config[0] === "FVAUTOCOMBO") {
                const options = this.options.values || [];
                if (options.length === 0) return true;
                let currentIndex = -1;
                for (let i = 0; i < options.length; i++) {
                    const optValue = typeof options[i] === 'object' ? options[i].value : options[i];
                    if (optValue == this.value || String(optValue) === String(this.value)) {
                        currentIndex = i;
                        break;
                    }
                }
                if (currentIndex === -1) {
                    currentIndex = 0;
                }
                let newIndex = currentIndex + delta;
                if (newIndex < 0) {
                    newIndex = options.length - 1;
                } else if (newIndex >= options.length) {
                    newIndex = 0;
                }
                const newOption = options[newIndex];
                this.value = typeof newOption === 'object' ? newOption.value : newOption;
                if (this.callback) {
                    this.callback(this.value);
                }
                node.graph.setDirtyCanvas(true, false);
                return true;
            } else {
                let v = parseFloat(this.value);
                const increment = delta * 0.1 * (this.options.step || 1);
                v += increment;
                // Apply min/max constraints
                if (this.options.min != null) {
                    v = Math.max(this.options.min, v);
                }
                if (this.options.max != null) {
                    v = Math.min(this.options.max, v);
                }
                // Round to precision or to integer
                if (this.config[0] === "FVAUTOINT") {
                    v = Math.round(v);
                } else if (this.options.precision !== undefined) {
                    const precision = Math.pow(10, this.options.precision);
                    v = Math.round(v * precision) / precision;
                }
                this.value = v;
                if (this.callback) {
                    this.callback(this.value);
                }
                node.graph.setDirtyCanvas(true, false);
                return true;
            }
        }
        if (delta === 0 && !this.isAuto) {
            if (this.config[0] === "FVAUTOCOMBO") {
                const options = this.options.values || [];
                // Create menu items
                const menuItems = options.map(opt => {
                    const value = typeof opt === 'object' ? opt.value : opt;
                    const label = typeof opt === 'object' ? opt.label : opt.toString();
                    return {
                        content: label,
                        callback: () => {
                            this.value = value;
                            if (this.callback) {
                                this.callback(this.value);
                            }
                            node.graph.setDirtyCanvas(true, false);
                        }
                    };
                });
                new LiteGraph.ContextMenu(menuItems, {
                    event: event,
                    title: null,
                    callback: null,
                    extra: node
                });
                return true;
            } else if (this.config[0] === "FVAUTOSTRING") {
                const d_callback = (v) => {
                    this.value = v;
                    if (this.callback) {
                        this.callback(this.value);
                    }
                    node.graph.setDirtyCanvas(true, false);
                };
                const dialog = app.canvas.prompt(
                    'Value',
                    this.value,
                    d_callback,
                    event
                );
                return true;
            } else {
                // For numeric widgets, show input dialog
                const d_callback = (v) => {
                    this.value = this.parseValue?.(v) ?? Number(v);
                    // Apply min/max constraints
                    if (this.options.min != null) {
                        this.value = Math.max(this.options.min, this.value);
                    }
                    if (this.options.max != null) {
                        this.value = Math.min(this.options.max, this.value);
                    }
                    // Round to precision or to integer
                    if (this.config[0] === "FVAUTOINT") {
                        this.value = Math.round(this.value);
                    } else if (this.options.precision !== undefined) {
                        const precision = Math.pow(10, this.options.precision);
                        this.value = Math.round(this.value * precision) / precision;
                    }
                    if (this.callback) {
                        this.callback(this.value);
                    }
                    node.graph.setDirtyCanvas(true, false);
                };
                const dialog = app.canvas.prompt(
                    'Value',
                    this.value,
                    d_callback,
                    event
                );
                return true;
            }
        }
        return true;
    }
    return false;
 }
 function makeAutoAnnotated(widget, inputData) {
    const original = {
        callback: widget.callback,
        type: widget.type,
        value: widget.value
    };
    // Add AUTO properties to the widget
    Object.assign(widget, {
        type: "BOOLEAN",
        draw: drawAutoAnnotated,
        mouse: mouseAutoAnnotated,
        onMouse: null, // Explicitly disable original onMouse handler
        isAuto: true,
        cachedValue: widget.value,
        config: inputData,
        options: Object.assign({}, inputData[1], widget.options),
        original: original,  // Store original properties for reference
        // Disable other potential mouse handlers with no-op functions
        onClick: function () {
            return false;
        },
        onPointerUp: function () {
            return false;
        },
        onPointerDown: function () {
            return false;
        },
        onMouseUp: function () {
            return false;
        },
        onMouseDown: function () {
            return false;
        },
        computeSize(width) {
            return [width, 20];
        },
        displayValue: function () {
            if (this.config[0] === "FVAUTOINT") {
                return Math.round(this.value).toString();
            }
            if (this.config[0] === "FVAUTOCOMBO") {
                return this.value;
            }
            if (this.config[0] === "FVAUTOSTRING") {
                return this.value;
            }
            // For FLOAT values, check if it's actually an integer
            if (Number.isInteger(this.value)) {
                return this.value.toString();
            }
            return this.value.toFixed(this.options.precision || 2);
        },
        parseValue: function (v) {
            if (this.config[0] === "FVAUTOSTRING") {
                return v;
            }
            if (typeof v === "string") {
                return parseFloat(v);
            }
            return v;
        },
        serializeValue: function () {
            // Return special value for AUTO mode
            return this.isAuto ? -99999 : this.value;
        },
        deserializeValue: function (data) {
            if (data === -99999) {
                this.isAuto = true;
                this.value = -99999;
            } else {
                this.isAuto = false;
                this.value = data;
                this.cachedValue = data;
            }
        }
    });
    // Override callback to handle AUTO mode
    widget.callback = function (v) {
        if (this.isAuto) {
            return; // Don't call the original callback in AUTO mode
        }
        const result = original.callback?.call(this, v);
        return result;
    };
    // Override any potential click handlers
    const originalOnClick = widget.onClick;
    if (originalOnClick) {
        widget.onClick = function (...args) {
            if (this.isAuto) {
                return false;
            }
            return originalOnClick.call(this, ...args);
        };
    }
    return widget;
 }
 app.registerExtension({
    name: "FastVideo.AutoWidgets",
    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        if (nodeData?.name == "VideoGenerator" || nodeData?.name === "InferenceArgs" || nodeData?.name === "VAEConfig" ||
            nodeData?.name === "TextEncoderConfig" || nodeData?.name === "DITConfig") {
            // Add serialization support
            chainCallback(nodeType.prototype, "onSerialize", function (info) {
                if (!this.widgets) {
                    return;
                }
                // Ensure widgets_values exists
                if (!info.widgets_values) {
                    info.widgets_values = {};
                }
                // Store AUTO widget states in a separate property
                if (!info.auto_widget_states) {
                    info.auto_widget_states = {};
                }
                // Handle AUTO widgets specially
                for (const w of this.widgets) {
                    if (w.type === "BOOLEAN" && w.isAuto !== undefined) {
                        // Store the serialized value (for Python node)
                        info.widgets_values[w.name] = w.serializeValue();
                        // Store the full state (for UI restoration)
                        info.auto_widget_states[w.name] = {
                            isAuto: w.isAuto,
                            value: w.value,
                            cachedValue: w.cachedValue
                        };
                    }
                }
            });
            // Add deserialization support
            chainCallback(nodeType.prototype, "onConfigure", function (info) {
                if (!this.widgets) {
                    return;
                }
                // First, restore from widgets_values (for backward compatibility)
                if (info.widgets_values && Array.isArray(info.widgets_values)) {
                    for (let i = 0; i < this.widgets.length && i < info.widgets_values.length; i++) {
                        const w = this.widgets[i];
                        const value = info.widgets_values[i];
                        if (w.type === "BOOLEAN" && w.isAuto !== undefined) {
                            w.deserializeValue(value);
                        }
                    }
                }
                // Then, restore full state if available
                if (info.auto_widget_states) {
                    for (const w of this.widgets) {
                        if (w.type === "BOOLEAN" && w.isAuto !== undefined && w.name in info.auto_widget_states) {
                            const state = info.auto_widget_states[w.name];
                            w.isAuto = state.isAuto;
                            w.cachedValue = state.cachedValue;
                            w.value = state.isAuto ? -99999 : state.value;
                            w.callback?.(w.value);
                        }
                    }
                }
                // Force a redraw
                this.graph?.setDirtyCanvas(true, true);
            });
            // Override addInput to handle AUTO widgets
            chainCallback(nodeType.prototype, "onNodeCreated", function () {
                // Convert any existing widgets to AUTO widgets if needed
                let new_widgets = [];
                const intWidgetNames = ["sp_size", "tp_size", "height", "width", "num_frames", "num_inference_steps", "flow_shift", "seed", "fps", "scale_factor",
                    "tile_sample_min_height", "tile_sample_min_width", "tile_sample_min_num_frames", "tile_sample_stride_height", "tile_sample_stride_width",
                    "tile_sample_stride_num_frames", "blend_num_frames"
                ]
                const floatWidgetNames = ["embedded_cfg_scale", "guidance_scale"]
                const comboWidgetNames = ["vae_tiling", "vae_precision", "vae_sp", "text_encoder_precision", "precision",
                    "load_encoder", "load_decoder", "use_tiling", "use_temporal_tiling", "use_parallel_tiling", "use_cpu_offload", "enable_teacache"
                ]
                const stringWidgetNames = ["prefix", "quant_config", "lora_config", "image_path"]
                if (this.widgets) {
                    for (let w of this.widgets) {
                        if (intWidgetNames.includes(w.name)) {
                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOINT", { "default": 0 }]));
                        } else if (floatWidgetNames.includes(w.name)) {
                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOFLOAT", { "default": 0 }]));
                        } else if (comboWidgetNames.includes(w.name)) {
                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOCOMBO", { "default": 0 }]));
                        } else if (stringWidgetNames.includes(w.name)) {
                            new_widgets.push(makeAutoAnnotated(w, ["FVAUTOSTRING", { "default": "" }]));
                        } else {
                            new_widgets.push(w);
                        }
                    }
                    this.widgets = new_widgets;
                    const autoWidgets = this.widgets.filter(w => w.type === "BOOLEAN" && w.isAuto !== undefined);
                }
                this.graph?.setDirtyCanvas(true, true);
            });
        }
    },
    async init() {
        // Force a redraw of all nodes when the extension initializes
        if (app.graph) {
            setTimeout(() => {
                app.graph.setDirtyCanvas(true, true);
            }, 1000);
        }
    }
 });
 console.log("FastVideo.core.js loaded");
@@ -4,7 +4,7 @@
 ## Installation
-We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
+We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only support H100/H200, because ThunderKittens uses TMA but doesn't support Blackwell yet.
 First, install C++20 for ThunderKittens:
 ```bash
 sudo apt update
@@ -53,8 +53,14 @@ out = sliding_tile_attention(q, k, v, window_size, 0, False)
 ## Test
 ```bash
-python test/test_sta.py
+python tests/test_sta.py # test STA
 python tests/test_block_sparse.py # test VSA
 ```
 ## Benchmark
 ```bash
 python benchmarks/bench_sta.py
 ```
 ## How Does STA Work?
 We give a demo for 2D STA with window size (6,6) operating on a (10, 10) image. 
@@ -5,6 +5,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 import torch
 from st_attn import sliding_tile_attention
 from triton.testing import do_bench
 def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"):
@@ -13,16 +14,16 @@ def flops(batch, seqlen, nheads, headdim, causal, mode="fwd"):
    return f if mode == "fwd" else (2.5 * f if mode == "bwd" else 3.5 * f)
-def efficiency(flop, time):
+def compute_TFLOPS(flops, ms):
-    flop = flop / 1e12
+    flops = flops / 1e12
-    time = time / 1e6
+    ms = ms / 1e3
-    return flop / time
+    return flops / ms
 def benchmark_attention(configurations):
    results = {'fwd': defaultdict(list), 'bwd': defaultdict(list)}
-    for B, H, N, D, causal in configurations:
+    for B, H, N, D, causal, dit_seq_shape, window_size in configurations:
        print("=" * 60)
        print(f"Timing forward and backward pass for B={B}, H={H}, N={N}, D={D}, causal={causal}")
@@ -30,38 +31,31 @@ def benchmark_attention(configurations):
        k = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
        v = torch.randn(B, H, N, D, dtype=torch.bfloat16, device='cuda', requires_grad=False).contiguous()
-        grad_output = torch.randn_like(q, requires_grad=False).contiguous()
+        # grad_output = torch.randn_like(q, requires_grad=False).contiguous()
        # qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous()
        # kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous()
        # vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous()
        qg = torch.zeros_like(q, requires_grad=False, dtype=torch.float).contiguous()
        kg = torch.zeros_like(k, requires_grad=False, dtype=torch.float).contiguous()
        vg = torch.zeros_like(v, requires_grad=False, dtype=torch.float).contiguous()
-        # Prepare for timing forward pass
+        # # Warmup for forward pass
-        start_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+        # for _ in range(10):
-        end_events_fwd = [torch.cuda.Event(enable_timing=True) for _ in range(10)]
+        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
-        torch.cuda.empty_cache()
+        # # Time the forward pass
-        torch.cuda.synchronize()
+        # for i in range(10):
        #     start_events_fwd[i].record()
        #     o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, dit_seq_shape)
        #     end_events_fwd[i].record()
        ms = do_bench(lambda: sliding_tile_attention(q, k, v, [window_size] * 24, 0, False, dit_seq_shape))
-        # Warmup for forward pass
+        # times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)]
-        for _ in range(10):
+        # time_us_fwd = np.mean(times_fwd) * 1000
            o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, '18x48x80')
-        # Time the forward pass
+        tflops_fwd = compute_TFLOPS(flops(B, N, H, D, causal, 'fwd'), ms)
        for i in range(10):
            start_events_fwd[i].record()
            o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, '18x48x80')
            end_events_fwd[i].record()
        torch.cuda.synchronize()
        times_fwd = [s.elapsed_time(e) for s, e in zip(start_events_fwd, end_events_fwd)]
        time_us_fwd = np.mean(times_fwd) * 1000
        tflops_fwd = efficiency(flops(B, N, H, D, causal, 'fwd'), time_us_fwd)
        results['fwd'][(D, causal)].append((N, tflops_fwd))
-        print(f"Average time for forward pass in us: {time_us_fwd:.2f}")
+        print(f"Average time for forward pass (ms): {ms:.2f}")
-        print(f"Average efficiency for forward pass in TFLOPS: {tflops_fwd}")
+        print(f"Average TFLOPS: {tflops_fwd}")
        print("-" * 60)
        # torch.cuda.empty_cache()
@@ -85,15 +79,14 @@ def benchmark_attention(configurations):
        # times_bwd = [s.elapsed_time(e) for s, e in zip(start_events_bwd, end_events_bwd)]
        # time_us_bwd = np.mean(times_bwd) * 1000
-        # tflops_bwd = efficiency(flops(B, N, H, D, causal, 'bwd'), time_us_bwd)
+        # tflops_bwd = compute_TFLOPS(flops(B, N, H, D, causal, 'bwd'), ms)
        # results['bwd'][(D, causal)].append((N, tflops_bwd))
-        # print(f"Average time for backward pass in us: {time_us_bwd:.2f}")
+        # print(f"Average time for backward pass(ms): {ms:.2f}")
-        # print(f"Average efficiency for backward pass in TFLOPS: {tflops_bwd}")
+        # print(f"Average TFLOPS: {tflops_bwd}")
-        print("=" * 60)
+        # print("=" * 60)
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
    return results
@@ -124,7 +117,10 @@ def plot_results(results):
 # Example list of configurations to test
 configurations = [
-    (2, 24, 69120, 128, False),
+    (2, 24, 69120, 128, False, '18x48x80', [3, 6, 10]),
    (2, 24, 69120, 128, True, '18x48x80', [3, 6, 10]),
    (2, 24, 82944, 128, False, '36x48x48', [3, 3, 6]), # Stepvideo
    (2, 24, 82944, 128, True, '36x48x48', [3, 3, 6]),
    # (16, 16, 768*16,  128, False),
    # (16, 16, 768*2,  128, False),
    # (16, 16, 768*4,  128, False),
@@ -4,9 +4,17 @@
 #include <cooperative_groups.h>
 #include <iostream>
 #include <stdio.h>
 #include <c10/cuda/CUDAGuard.h>
 // #define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value)))
 __device__ __forceinline__ int clamp_int(int value, int min, int max) {
    return (value < min) ? min : ((value > max) ? max : value);
 }
 // #define ABS(x) ((x) < 0 ? -(x) : (x))
 __device__ __forceinline__ int abs_int(int value) {
    return (value < 0) ? -value : value;
 }
 #define CLAMP(value, min, max) ((value) < (min) ? (min) : ((value) > (max) ? (max) : (value)))
 #define ABS(x) ((x) < 0 ? -(x) : (x))
 constexpr int CONSUMER_WARPGROUPS = (3); 
 constexpr int PRODUCER_WARPGROUPS = (1); 
@@ -117,16 +125,16 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
            int qt = seq_idx / 6 / (CH * CW);
            int qh = (seq_idx / 6) % (CH * CW) / CW;
            int qw = (seq_idx / 6) % CW;
-            qt = CLAMP(qt, DT, CT-DT-1);
+            qt = clamp_int(qt, DT, CT-DT-1);
-            qh = CLAMP(qh, DH, CH-DH-1);
+            qh = clamp_int(qh, DH, CH-DH-1);
-            qw = CLAMP(qw, DW, CW-DW-1);
+            qw = clamp_int(qw, DW, CW-DW-1);
            int count = 0;
            int j = 0;
            while (count < K::stages - 1) {
                int kt = j / 3 / (CH * CW);
                int kh = (j / 3) % (CH * CW) / CW;
                int kw = (j / 3) % CW;
-                bool mask = (ABS(qt - kt) <= DT) && (ABS(qh - kh) <= DH) && (ABS(qw - kw) <= DW);
+                bool mask = (abs_int(qt - kt) <= DT) && (abs_int(qh - kh) <= DH) && (abs_int(qw - kw) <= DW);
                if (mask){
                    coord<k_tile> kv_tile_idx = {blockIdx.z, kv_head_idx, j, 0};
                    tma::expect_bytes(k_smem_arrived[count], sizeof(k_tile));
@@ -167,15 +175,15 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
                int qt = seq_idx / 6 / (CH * CW);
                int qh = (seq_idx / 6) % (CH * CW) / CW;
                int qw = (seq_idx / 6) % CW;
-                qt = CLAMP(qt, DT, CT-DT-1);
+                qt = clamp_int(qt, DT, CT-DT-1);
-                qh = CLAMP(qh, DH, CH-DH-1);
+                qh = clamp_int(qh, DH, CH-DH-1);
-                qw = CLAMP(qw, DW, CW-DW-1);
+                qw = clamp_int(qw, DW, CW-DW-1);
-                int k_t_min = CLAMP(qt-DT, 0, CT-1);
+                int k_t_min = clamp_int(qt-DT, 0, CT-1);
-                int k_t_max = CLAMP(qt+DT, 0, CT-1);
+                int k_t_max = clamp_int(qt+DT, 0, CT-1);
-                int k_h_min = CLAMP(qh-DH, 0, CH-1);
+                int k_h_min = clamp_int(qh-DH, 0, CH-1);
-                int k_h_max = CLAMP(qh+DH, 0, CH-1);
+                int k_h_max = clamp_int(qh+DH, 0, CH-1);
-                int k_w_min = CLAMP(qw-DW, 0, CW-1);
+                int k_w_min = clamp_int(qw-DW, 0, CW-1);
-                int k_w_max = CLAMP(qw+DW, 0, CW-1);
+                int k_w_max = clamp_int(qw+DW, 0, CW-1);
                int count = 0;
                for (int kt = k_t_min; kt <= k_t_max; kt++) {
                    for (int kh = k_h_min; kh <= k_h_max; kh++) {
@@ -234,7 +242,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals<D> g) {
            // the last three kv blocks are for text, we process them separately
            kv_iters = img_kv_blocks - 1;
        } else {
-            kv_iters = CLAMP(DT*2+1, 1, CT) * CLAMP(DH*2+1, 1, CH) * CLAMP(DW*2+1, 1, CW) * 3 - 1 ; 
+            kv_iters = clamp_int(DT*2+1, 1, CT) * clamp_int(DH*2+1, 1, CH) * clamp_int(DW*2+1, 1, CW) * 3 - 1 ; 
        }
        kittens::wait(qsmem_semaphore, 0);
@@ -415,8 +423,9 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
    float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
    float* d_l   = reinterpret_cast<float*>(l_ptr);
-    cudaDeviceSynchronize();
+    //cudadevicesynchronize();
-    auto stream = at::cuda::getCurrentCUDAStream().stream(); 
+    const c10::cuda::OptionalCUDAGuard device_guard(q.device());
    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
    if (head_dim == 128) {
@@ -442,8 +451,8 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
        globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len),  static_cast<int>(text_length), static_cast<int>(hr)};
-        auto mem_size = kittens::MAX_SHARED_MEMORY;
+        constexpr int mem_size = kittens::MAX_SHARED_MEMORY;
-        auto threads  = NUM_WORKERS * kittens::WARP_THREADS;
+        int 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);
@@ -823,10 +832,10 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
        }
        CHECK_CUDA_ERROR(cudaGetLastError());
-        cudaStreamSynchronize(stream);
+        // cudaStreamSynchronize(stream);
    }
    return o;
-    cudaDeviceSynchronize();
+    //cudadevicesynchronize();
 }
@@ -7,6 +7,7 @@ from vsa import BLOCK_M, BLOCK_N
 import numpy as np
 import random
 import gc
 def set_seed(seed: int = 42):
    # Python random module
@@ -20,15 +21,6 @@ def set_seed(seed: int = 42):
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if using multi-GPU
 def parse_arguments():
    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
    parser.add_argument('--num_heads', type=int, default=6, help='Number of heads')
    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
    parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to')
    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark')
    parser.add_argument('--num_iterations', type=int, default=100, help='Number of test iterations to run')
    return parser.parse_args()
@torch.no_grad
 def precision_metric(quant_o, fa2_o): 
@@ -135,9 +127,7 @@ def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device=
    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
-def main():
+def main(args):
    args = parse_arguments()
    set_seed(42)
    # Extract parameters
@@ -191,23 +181,36 @@ def main():
            block_mask_expanded = block_sparse_mask.unsqueeze(-1).unsqueeze(-2)  # [b, h, num_q_blocks, num_kv_blocks, 1, 1]
            block_mask_expanded = block_mask_expanded.expand(-1, -1, -1, -1, BLOCK_M, BLOCK_N)  # [b, h, num_q_blocks, num_kv_blocks, BLOCK_M, BLOCK_N]
            full_mask = block_mask_expanded.permute(0, 1, 2, 4, 3, 5).reshape(batch, head, seq_len, seq_len)
            q_sdpa = q.clone()
            k_sdpa = k.clone()
            v_sdpa = v.clone()
            q.requires_grad = True
            k.requires_grad = True
            v.requires_grad = True
-            q_sdpa.requires_grad = True
+
            k_sdpa.requires_grad = True
            v_sdpa.requires_grad = True
            # testing forward
            o = BlockSparseAttentionFunction.apply(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num)
            del q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask, block_mask_expanded
            grad_o = torch.randn_like(o)
            o.backward(grad_o)
            # clear memory
            q_sdpa = q.detach().clone()
            k_sdpa = k.detach().clone()
            v_sdpa = v.detach().clone()
            q_sdpa.requires_grad = True
            k_sdpa.requires_grad = True
            v_sdpa.requires_grad = True
            q.data = torch.empty(0, device=q.device)
            k.data = torch.empty(0, device=k.device)
            v.data = torch.empty(0, device=v.device)
            torch.cuda.empty_cache()
            o_sdpa = torch.nn.functional.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, attn_mask=full_mask)
            sim, l1, rmse = precision_metric(o, o_sdpa)
            assert sim > 0.9999, f"SSIM too low: {sim}"
            assert l1 < 8e-5, f"l1 too large: {l1}"
            assert rmse < 2e-5, f"RMSE too large: {rmse}"
            forward_metrics['sim'].append(sim)
            forward_metrics['l1'].append(l1)
            forward_metrics['rmse'].append(rmse)
@@ -215,52 +218,72 @@ def main():
            print(f"block_sparse_attention_fwd vs torch.nn.functional.scaled_dot_product_attention:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
            # test backward
            grad_o = torch.randn_like(o)
            o.backward(grad_o)
            o_sdpa.backward(grad_o)
            sim, l1, rmse = precision_metric(q.grad, q_sdpa.grad)
            # Error bounds collected on H100
            assert sim > 0.9999, f"SSIM too low: {sim}"
            assert l1 < 4e-3, f"l1 too large: {l1}"
            assert rmse < 3e-4, f"RMSE too large: {rmse}"
            grad_q_metrics['sim'].append(sim)
            grad_q_metrics['l1'].append(l1)
            grad_q_metrics['rmse'].append(rmse)
            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_q:\nsim: {sim}, l1: {l1}, rmse: {rmse}")        
            sim, l1, rmse = precision_metric(k.grad, k_sdpa.grad)
            assert sim > 0.9999, f"SSIM too low: {sim}"
            assert l1 < 4e-3, f"l1 too large: {l1}"
            assert rmse < 2e-4, f"RMSE too large: {rmse}"
            grad_k_metrics['sim'].append(sim)
            grad_k_metrics['l1'].append(l1)
            grad_k_metrics['rmse'].append(rmse)
            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_k:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
            sim, l1, rmse = precision_metric(v.grad, v_sdpa.grad)
            assert sim > 0.9999, f"SSIM too low: {sim}"
            assert l1 < 1e-4, f"l1 too large: {l1}"
            assert rmse < 2e-5, f"RMSE too large: {rmse}"
            grad_v_metrics['sim'].append(sim)
            grad_v_metrics['l1'].append(l1)
            grad_v_metrics['rmse'].append(rmse)
            print(f"block_sparse_attention_bwd vs torch.nn.functional.scaled_dot_product_attention grad_v:\nsim: {sim}, l1: {l1}, rmse: {rmse}")
-        
+            
            del o, o_sdpa, grad_o, q_sdpa, k_sdpa, v_sdpa
            gc.collect()
            torch.cuda.empty_cache()
        # Print summary statistics if multiple iterations were run
        if num_iterations > 1:
            print("\n" + "="*50)
            print(f"Summary Statistics (over {num_iterations} iterations):")
            print("\nForward metrics:")
-            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}")
+            print(f"Similarity: mean={np.mean(forward_metrics['sim']):.6f}, std={np.std(forward_metrics['sim']):.6f}, min={np.min(forward_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}")
+            print(f"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}, max={np.max(forward_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}")
+            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(forward_metrics['rmse']):.6f}, max={np.max(forward_metrics['rmse']):.6f}")
            print("\nGradient Q metrics:")
-            print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}")
+            print(f"Similarity: mean={np.mean(grad_q_metrics['sim']):.6f}, std={np.std(grad_q_metrics['sim']):.6f}, min={np.min(grad_q_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(grad_q_metrics['l1']):.6f}, max={np.max(grad_q_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_q_metrics['rmse']):.6f}, std={np.std(grad_q_metrics['rmse']):.6f}, max={np.max(grad_q_metrics['rmse']):.6f}")
            print("\nGradient K metrics:")
-            print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}")
+            print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(grad_k_metrics['sim']):.6f}, min={np.min(grad_k_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_k_metrics['l1']):.6f}, std={np.std(grad_k_metrics['l1']):.6f}, max={np.max(grad_k_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(grad_k_metrics['rmse']):.6f}, max={np.max(grad_k_metrics['rmse']):.6f}")
            print("\nGradient V metrics:")
-            print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}")
+            print(f"Similarity: mean={np.mean(grad_v_metrics['sim']):.6f}, std={np.std(grad_v_metrics['sim']):.6f}, min={np.min(grad_v_metrics['sim']):.6f}")
-            print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}")
+            print(f"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}, max={np.max(grad_v_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}")
+            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}, max={np.max(grad_v_metrics['rmse']):.6f}")
 if __name__ == "__main__":
-    main()
+    parser = argparse.ArgumentParser(description='Benchmark Block Sparse Attention')
    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
    parser.add_argument('--num_heads', type=int, default=6, help='Number of heads')
    parser.add_argument('--head_dim', type=int, default=128, help='Head dimension')
    parser.add_argument('--topk', type=int, default=64, help='Number of kv blocks each q block attends to')
    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[29120], help='Sequence lengths to benchmark')
    parser.add_argument('--num_iterations', type=int, default=50, help='Number of test iterations to run')
    args = parser.parse_args()
    main(args)
@@ -81,5 +81,7 @@ std = 10
 # Run correctness check directly
 results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output')
 assert results['TK vs FLEX']['avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large"
 assert results['TK vs FLEX']['max_diff'] < 4e-2, f"Maximum difference: {results['TK vs FLEX']['max_diff']} is too large"
 print(f"Average difference: {results['TK vs FLEX']['avg_diff']}")
 print(f"Maximum difference: {results['TK vs FLEX']['max_diff']}")
@@ -3,6 +3,8 @@
 #include "kittens.cuh"
 #include <cooperative_groups.h>
 #include <iostream>
 #include <c10/cuda/CUDAGuard.h>
 using namespace kittens;
 namespace cg = cooperative_groups;
@@ -940,8 +942,9 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
    float* l_ptr = reinterpret_cast<float*>(l_vec.data_ptr<float>());
    float* d_l   = reinterpret_cast<float*>(l_ptr);
-    cudaDeviceSynchronize();
+    //cudadevicesynchronize();
-    auto stream = at::cuda::getCurrentCUDAStream().stream(); 
+    const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
    if (head_dim == 64) {
        using q_tile    =         st_bf<fwd_attend_ker_tile_dims<64>::qo_height, fwd_attend_ker_tile_dims<64>::tile_width>;
@@ -966,7 +969,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
        globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_kv_blocks_per_q), reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr())};
-        auto mem_size = 54000;
+        constexpr int mem_size = 54000;
        dim3 grid(seq_len/(64), qo_heads, batch);
@@ -979,7 +982,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
        fwd_attend_ker<64><<<grid, (128), mem_size, stream>>>(g);
        CHECK_CUDA_ERROR(cudaGetLastError());
-        cudaStreamSynchronize(stream);
+        // cudaStreamSynchronize(stream);
    }
    if (head_dim == 128) {
@@ -1005,7 +1008,7 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
        globals g{qg_arg, kg_arg, vg_arg, lg_arg, og_arg, static_cast<int>(seq_len), static_cast<int>(hr), static_cast<int>(max_kv_blocks_per_q), reinterpret_cast<int32_t*>(q2k_block_sparse_index.data_ptr()), reinterpret_cast<int32_t*>(q2k_block_sparse_num.data_ptr())};
-        auto mem_size = 54000;
+        constexpr int mem_size = 54000;
        dim3 grid(seq_len/(64), qo_heads, batch);
@@ -1018,11 +1021,11 @@ block_sparse_attention_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v
        fwd_attend_ker<128><<<grid, (128), mem_size, stream>>>(g);
        CHECK_CUDA_ERROR(cudaGetLastError());
-        cudaStreamSynchronize(stream);
+        // cudaStreamSynchronize(stream);
    }
    return {o, l_vec};
-    cudaDeviceSynchronize();
+    //cudadevicesynchronize();
 }
 std::vector<torch::Tensor> 
@@ -1132,13 +1135,14 @@ block_sparse_attention_backward(torch::Tensor q,
    float* d_kg = reinterpret_cast<float*>(kg_ptr);
    float* d_vg = reinterpret_cast<float*>(vg_ptr);
-    auto mem_size = kittens::MAX_SHARED_MEMORY; 
+    constexpr int mem_size = kittens::MAX_SHARED_MEMORY; 
-    auto threads  = 4 * kittens::WARP_THREADS;
+    int threads  = 4 * kittens::WARP_THREADS;
-    cudaDeviceSynchronize();
+    //cudadevicesynchronize();
-    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    const c10::cuda::OptionalCUDAGuard device_guard(q.device());    
    const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); 
-    cudaStreamSynchronize(stream);
+    //  cudaStreamSynchronize(stream);
    // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256");
    dim3 grid_bwd(seq_len/(4*kittens::TILE_ROW_DIM<bf16>*4), qo_heads, batch);
@@ -1222,7 +1226,7 @@ block_sparse_attention_backward(torch::Tensor q,
        dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
        threads = 128;
-        cudaDeviceSynchronize();
+        //cudadevicesynchronize();
        {
            cudaFuncSetAttribute(
@@ -1240,8 +1244,8 @@ block_sparse_attention_backward(torch::Tensor q,
        }
        // CHECK_CUDA_ERROR(cudaGetLastError());
-        cudaStreamSynchronize(stream);
+        // cudaStreamSynchronize(stream);
-        cudaDeviceSynchronize();
+        //cudadevicesynchronize();
        // const auto kernel_end = std::chrono::high_resolution_clock::now();
        // std::cout << "Kernel Time: " << std::chrono::duration_cast<std::chrono::microseconds>(kernel_end - start).count() << "us" << std::endl;
        // std::cout << "---" << std::endl;
@@ -1326,7 +1330,7 @@ block_sparse_attention_backward(torch::Tensor q,
        dim3 grid_bwd_2(seq_len/64, qo_heads, batch);
        threads = 128;
-        cudaDeviceSynchronize();
+        //cudadevicesynchronize();
        {
            cudaFuncSetAttribute(
@@ -1338,10 +1342,10 @@ block_sparse_attention_backward(torch::Tensor q,
            bwd_attend_ker<128><<<grid_bwd_2, threads, 113000, stream>>>(bwd_global); 
        }
-        cudaStreamSynchronize(stream);
+        // cudaStreamSynchronize(stream);
-        cudaDeviceSynchronize();
+        //cudadevicesynchronize();
    }
    return {qg, kg, vg};
-    cudaDeviceSynchronize();
+    //cudadevicesynchronize();
 }
@@ -1,7 +1,9 @@
-FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
+FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04
 ENV DEBIAN_FRONTEND=noninteractive
 SHELL ["/bin/bash", "-c"]
 WORKDIR /FastVideo
 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -9,17 +11,25 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
    git \
    ca-certificates \
    openssh-server \
    zsh \
    vim \
    curl \
    gcc-11 \
    g++-11 \
    clang-11 \
    && rm -rf /var/lib/apt/lists/*
-RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+# Set up C++20 compilers for ThunderKittens
-    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
    rm Miniconda3-latest-Linux-x86_64.sh
-ENV PATH=/opt/conda/bin:$PATH
+# Set CUDA environment variables
 ENV CUDA_HOME=/usr/local/cuda-12.8
 ENV PATH=${CUDA_HOME}/bin:${PATH}
 ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-RUN conda create --name fastvideo-dev python=3.12.9 -y
+# Install uv and source its environment
-
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
-SHELL ["/bin/bash", "-c"]
+    echo 'source $HOME/.local/bin/env' >> /root/.bashrc
 # Copy just the pyproject.toml first to leverage Docker cache
 COPY pyproject.toml ./
@@ -27,22 +37,36 @@ 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 && \
+# Create and activate virtual environment with specific Python version and seed
-    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
+RUN source $HOME/.local/bin/env && \
-    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
+    uv venv --python 3.12 --seed /opt/venv && \
-    conda clean -afy
+    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir --upgrade pip && \
    uv pip install --no-cache-dir .[dev] && \
    uv pip install --no-cache-dir flash-attn==2.8.0.post2 --no-build-isolation
 COPY . .
-RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
+# Install dependencies using uv and set up shell configuration
 RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    uv pip install --no-cache-dir -e .[dev] && \
    git config --unset-all http.https://github.com/.extraheader || true && \
    echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
    echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
-# Remove authentication headers
+# Install STA (Sliding Tile Attention)
-RUN git config --unset-all http.https://github.com/.extraheader || true
+RUN source $HOME/.local/bin/env && \
    source /opt/venv/bin/activate && \
    cd csrc/attn && \
    git submodule update --init --recursive && \
    python setup_sta.py install
-# Set up automatic conda environment activation for all shells
+# Install VSA 
-RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
+RUN source $HOME/.local/bin/env && \
-    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
+    source /opt/venv/bin/activate && \
-    # Ensure .bashrc is sourced for SSH login shells
+    cd csrc/attn && \
-    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
+    git submodule update --init --recursive && \
    python setup_vsa.py install
 EXPOSE 22
@@ -18,7 +18,6 @@ import os
 import re
 import sys
 from pathlib import Path
 from typing import Optional
 import requests
@@ -168,8 +167,7 @@ _cached_base: str = ""
 _cached_branch: str = ""
-def get_repo_base_and_branch(
+def get_repo_base_and_branch(pr_number: str) -> tuple[str | None, str | None]:
        pr_number: str) -> tuple[Optional[str], Optional[str]]:
    global _cached_base, _cached_branch
    if _cached_base and _cached_branch:
        return _cached_base, _cached_branch
@@ -8,12 +8,14 @@ You can easily use the FastVideo Docker image as a custom container on [RunPod](
 Choose a GPU that supports CUDA 12.4
 Pick 1 or 2 L40S GPU(s)
 ![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
+ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest
 ```
 Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.io/pods/configuration/use-ssh)):
@@ -5,7 +5,6 @@ import itertools
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Optional
 ROOT_DIR = Path(__file__).parent.parent.parent.resolve()
 ROOT_DIR_RELATIVE = '../../../..'
@@ -89,7 +88,7 @@ class Example:
        generate() -> str: Generates the documentation content.
    """ # noqa: E501
    path: Path
-    category: Optional[str] = None
+    category: str | None = None
    main_file: Path = field(init=False)
    other_files: list[Path] = field(init=False)
    title: str = field(init=False)
@@ -117,4 +117,4 @@ If you're planning to contribute to FastVideo please see the following page:
 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.
+You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.
@@ -12,7 +12,7 @@ This guide explains how to implement a custom diffusion pipeline in FastVideo, l
 4. **Register Your Pipeline** - Make it discoverable by the framework
 5. **Configure Your Pipeline** - (Coming soon)
-Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).
+Need help? Join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
 ## Step 1: Pipeline Modules
@@ -27,7 +27,7 @@ fastvideo generate --help
 ### Hardware Configuration
 - `--num-gpus {NUM_GPUS}`: Number of GPUs to use
- `--tp-size {TP_SIZE}`: Tensor parallelism size (Typically should match the number of GPUs)
+- `--tp-size {TP_SIZE}`: Tensor parallelism size (only for the encoder, should not be larger than 1 if text encoder offload is enabled, as layerwise offload + prefetch is faster)
 - `--sp-size {SP_SIZE}`: Sequence parallelism size (Typically should match the number of GPUs)
 #### Video Configuration
@@ -68,7 +68,7 @@ Example configuration file (config.json):
    "output_path": "outputs/",
    "num_gpus": 2,
    "sp_size": 2,
-    "tp_size": 2,
+    "tp_size": 1,
    "num_frames": 45,
    "height": 720,
    "width": 1280,
@@ -102,7 +102,7 @@ prompt: "A beautiful woman in a red dress walking down a street"
 output_path: "outputs/"
 num_gpus: 2
 sp_size: 2
-tp_size: 2
+tp_size: 1
 num_frames: 45
 height: 720
 width: 1280
@@ -121,4 +121,4 @@ If the generated video doesn't match your prompt:
 - Learn about using [Optimizations](#inference-optimizations)
 - See [Examples](../examples/examples_inference_index.md) for more usage scenarios
 - Join our [Community Discord](https://discord.gg/JA7cksDz86).
- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-2zf6ru791-sRwI9lPIUJQq1mIeB_yjJg).
+- Join our [Community Slack](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ).
@@ -1,7 +1,7 @@
 (sta-demo)=
 # 🔍 Demo
-There is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
+This is 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">
@@ -9,3 +9,9 @@ There is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
    Your browser does not support the video tag.
  </video>
 </div>
 You can run STA using the following command:
 ```bash
 bash scripts/inference/v1_inference_wan_STA.sh
 ```
@@ -16,6 +16,13 @@ bash scripts/finetune/finetune_mochi.sh # for mochi
 ```
 **Note that for finetuning, we did not tune the hyperparameters in the provided script.**
 ## ⚡ Finetune with VSA
 Follow [data_preprocess.md](#v0-data-preprocess) to get parquet files for preproccessed latent, and then run:
 ```bash
 bash scripts/finetune/finetune_v1_VSA.sh
 ```
 ## ⚡ 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:
@@ -10,7 +10,7 @@ def main():
    # attempt to identify the optimal arguments.
    generator = VideoGenerator.from_pretrained(
        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        # if num_gpus > 1, FastVideo will automatically handle distributed setup
+        # FastVideo will automatically handle distributed setup
        num_gpus=2,
        use_fsdp_inference=True,
        use_cpu_offload=False
@@ -6,7 +6,7 @@ def main():
    # Initialize VideoGenerator with the Wan model
    generator = VideoGenerator.from_pretrained(
        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-        num_gpus=2,
+        num_gpus=1,
        lora_path="benjamin-paine/steamboat-willie-1.3b",
        lora_nickname="steamboat"
    )
@@ -16,6 +16,7 @@ def main():
        "num_frames": 81,
        "guidance_scale": 5.0,
        "num_inference_steps": 32,
        "seed": 42,
    }
    # Generate video with LoRA style
    prompt = "steamboat willie style, golden era animation, close-up of a short fluffy monster  kneeling beside a melting red candle. the mood is one of wonder and curiosity,  as the monster gazes at the flame with wide eyes and open mouth. Its pose and expression  convey a sense of innocence and playfulness, as if it is exploring the world around it for the first time.  The use of warm colors and dramatic lighting further enhances the cozy atmosphere of the image."
@@ -29,8 +30,17 @@ def main():
        negative_prompt=negative_prompt,
        **kwargs
    )
-
+    del generator
-    generator.set_lora_adapter(lora_nickname="flat_color", lora_path="motimalu/wan-flat-color-1.3b-v2")
+    
    # Until FSDP resharding bug is fixed, multi-lora requires reloading the model
    # see https://github.com/pytorch/pytorch/issues/157209
    generator = VideoGenerator.from_pretrained(
        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
        num_gpus=1,
        lora_path="motimalu/wan-flat-color-1.3b-v2",
        lora_nickname="flat_color"
    )
    # generator.set_lora_adapter(lora_nickname="flat_color", lora_path="motimalu/wan-flat-color-1.3b-v2")
    prompt = "flat color, no lineart, blending, negative space, artist:[john kafka|ponsuke kaikai|hara id 21|yoneyama mai|fuzichoco],  1girl, sakura miko, pink hair, cowboy shot, white shirt, floral print, off shoulder, outdoors, cherry blossom, tree shade, wariza, looking up, falling petals, half-closed eyes, white sky, clouds,  live2d animation, upper body, high quality cinematic video of a woman sitting under a sakura tree. Dreamy and lonely, the camera close-ups on the face of the woman as she turns towards the viewer. The Camera is steady, This is a cowboy shot. The animation is smooth and fluid."
    negative_prompt = "bad quality video,色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
    video = generator.generate_video(
@@ -5,7 +5,7 @@ export FASTVIDEO_ATTENTION_BACKEND=SLIDING_TILE_ATTN
 export MODEL_BASE=Wan-AI/Wan2.1-T2V-14B-Diffusers
 base_port=29503
-num_gpu=$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l)
+num_gpu=1
 gpu_ids=$(seq 0 $((num_gpu-1)))
 skip_time_steps=12
@@ -14,7 +14,7 @@ STA_mode="STA_searching"
 for i in $gpu_ids; do
    port=$((base_port+i))
    CUDA_VISIBLE_DEVICES=$i MASTER_PORT=$port python examples/inference/sta_mask_search/wan_example.py \
-        --prompt_path ./assets/prompt_extend_${i}.txt \
+        --prompt_path ./assets/prompt_${i}.txt \
        --output_path $output_path \
        --STA_mode $STA_mode &
    sleep 1
@@ -27,7 +27,7 @@ STA_mode="STA_tuning"
 for i in $gpu_ids; do
    port=$((base_port+i))
    CUDA_VISIBLE_DEVICES=$i MASTER_PORT=$port python examples/inference/sta_mask_search/wan_example.py \
-        --prompt_path ./assets/prompt_extend_${i}.txt \
+        --prompt_path ./assets/prompt_${i}.txt \
        --output_path $output_path \
        --STA_mode $STA_mode \
        --skip_time_steps $skip_time_steps &
@@ -0,0 +1,10 @@
 This directory contain e2e examples scripts for finetuning Wan2.1 I2V. 
 Execute the following commands from `FastVideo/` to run training:
 - Download crush-smol dataset:
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
 - Preprocess the videos and captions into latents:
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
 - Edit the following file and run finetuning:
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`
@@ -0,0 +1,3 @@
 #!/bin/bash
 python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -0,0 +1,94 @@
 #!/bin/bash
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export TOKENIZERS_PARALLELISM=false
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
 DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
 NUM_GPUS=8
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 # Training arguments
 training_args=(
  --tracker_project_name "wan_i2v_finetune"
  --output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
  --max_train_steps 2000
  --train_batch_size 4
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
  --num_latent_t 8
  --num_height 480
  --num_width 832
  --num_frames 77
  --enable_gradient_checkpointing_type "full"
 )
 # Parallel arguments
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size 4
  --tp_size 4
  --hsdp_replicate_dim 2
  --hsdp_shard_dim 4
 )
 # Model arguments
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
 )
 # Dataset arguments
 dataset_args=(
  --data_path "$DATA_DIR"
  --dataloader_num_workers 1
 )
 # Validation arguments
 validation_args=(
  --log_validation
  --validation_dataset_file "$VALIDATION_DATASET_FILE"
  --validation_steps 100
  --validation_sampling_steps "40"
  --validation_guidance_scale "1.0"
 )
 # Optimizer arguments
 optimizer_args=(
  --learning_rate 2e-5
  --mixed_precision "bf16"
  --checkpointing_steps 2000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
 # Miscellaneous arguments
 miscellaneous_args=(
  --inference_mode False
  --allow_tf32
  --checkpoints_total_limit 3
  --training_cfg_rate 0.1
  --multi_phased_distill_schedule "4000-1"
  --not_apply_cfg_solver
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
  --enable_gradient_checkpointing_type "full"
 )
 # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
 torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
    fastvideo/v1/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}"
@@ -0,0 +1,130 @@
 #!/bin/bash
 #SBATCH --job-name=i2v
 #SBATCH --partition=main
 #SBATCH --nodes=4
 #SBATCH --ntasks=4
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
 #SBATCH --mem=1440G
 #SBATCH --output=i2v_output/i2v_%j.out
 #SBATCH --error=i2v_output/i2v_%j.err
 #SBATCH --exclusive
 set -e -x
 # Environment Setup
 source ~/conda/miniconda/bin/activate
 conda activate will-fv
 # Basic Info
 export WANDB_MODE="online"
 export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
 # different cache dir for different processes
 export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
 export MASTER_PORT=29500
 export NODE_RANK=$SLURM_PROCID
 nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
 export MASTER_ADDR=${nodes[0]}
 export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
 export TOKENIZERS_PARALLELISM=false
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 echo "MASTER_ADDR: $MASTER_ADDR"
 echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
 DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
 # Training arguments
 training_args=(
  --tracker_project_name wan_i2v_finetune
  --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"
  --max_train_steps=2000
  --train_batch_size=2
  --train_sp_batch_size 1
  --gradient_accumulation_steps=1
  --num_latent_t 8
  --num_height 480
  --num_width 832
  --num_frames 77
  --enable_gradient_checkpointing_type "full"
 )
 # Parallel arguments
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size $NUM_GPUS
  --tp_size $NUM_GPUS
  --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
  --hsdp_shard_dim $NUM_GPUS
 )
 # Model arguments
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
 )
 # Dataset arguments
 dataset_args=(
  --data_path "$DATA_DIR"
  --dataloader_num_workers 10
 )
 # Validation arguments
 validation_args=(
  --log_validation
  --validation_dataset_file "$VALIDATION_DATASET_FILE"
  --validation_steps 100
  --validation_sampling_steps "40"
  --validation_guidance_scale "1.0"
 )
 # Optimizer arguments
 optimizer_args=(
  --learning_rate=1e-5
  --mixed_precision="bf16"
  --checkpointing_steps=1000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
 # Miscellaneous arguments
 miscellaneous_args=(
  --inference_mode False
  --allow_tf32
  --checkpoints_total_limit 3
  --training_cfg_rate 0.1
  --multi_phased_distill_schedule "4000-1"
  --not_apply_cfg_solver
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
  --enable_gradient_checkpointing_type "full"
 )
 srun torchrun \
 --nnodes $SLURM_JOB_NUM_NODES \
 --nproc_per_node $NUM_GPUS \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
    fastvideo/v1/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}"
@@ -0,0 +1,26 @@
 #!/bin/bash
 GPU_NUM=1 # 2,4,8
 MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_i2v_1_3b_inp/"
 VALIDATION_PATH="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json"
 torchrun --nproc_per_node=$GPU_NUM \
    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
    --model_path $MODEL_PATH \
    --data_merge_path $DATA_MERGE_PATH \
    --preprocess_video_batch_size 8 \
    --max_height 480 \
    --max_width 832 \
    --num_frames 77 \
    --dataloader_num_workers 0 \
    --output_dir=$OUTPUT_DIR \
    --model_type $MODEL_TYPE \
    --train_fps 16 \
    --validation_dataset_file $VALIDATION_PATH \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
    --preprocess_task "i2v" 
@@ -0,0 +1,31 @@
 {
  "data": [
    {
      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
      "image_path": null,
      "video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
      "num_inference_steps": 40,
      "height": 480,
      "width": 832,
      "num_frames": 77
    },
    {
      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
      "image_path": null,
      "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
      "num_inference_steps": 40,
      "height": 480,
      "width": 832,
      "num_frames": 77
    },
    {
      "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
      "image_path": null,
      "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
      "num_inference_steps": 40,
      "height": 480,
      "width": 832,
      "num_frames": 77
    }
  ]
 }
@@ -0,0 +1,10 @@
 This directory contain e2e examples scripts for finetuning Wan2.1 I2V. 
 Execute the following commands from `FastVideo/` to run training:
 - Download crush-smol dataset:
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/download_dataset.sh`
 - Preprocess the videos and captions into latents:
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/preprocess_wan_data_i2v.sh`
 - Edit the following file and run finetuning:
 `bash examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh`
@@ -0,0 +1,3 @@
 #!/bin/bash
 python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -0,0 +1,94 @@
 #!/bin/bash
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export TOKENIZERS_PARALLELISM=false
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
 NUM_GPUS=8
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 # Training arguments
 training_args=(
  --tracker_project_name "wan_i2v_finetune"
  --output_dir "$DATA_DIR/outputs/wan_i2v_finetune"
  --max_train_steps 2000
  --train_batch_size 1
  --train_sp_batch_size 1
  --gradient_accumulation_steps 1
  --num_latent_t 8
  --num_height 480
  --num_width 832
  --num_frames 77
  --enable_gradient_checkpointing_type "full"
 )
 # Parallel arguments
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size 8
  --tp_size 1
  --hsdp_replicate_dim 1
  --hsdp_shard_dim 8
 )
 # Model arguments
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
 )
 # Dataset arguments
 dataset_args=(
  --data_path "$DATA_DIR"
  --dataloader_num_workers 1
 )
 # Validation arguments
 validation_args=(
  --log_validation
  --validation_dataset_file "$VALIDATION_DATASET_FILE"
  --validation_steps 100
  --validation_sampling_steps "40"
  --validation_guidance_scale "1.0"
 )
 # Optimizer arguments
 optimizer_args=(
  --learning_rate 1e-5
  --mixed_precision "bf16"
  --checkpointing_steps 1000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
 # Miscellaneous arguments
 miscellaneous_args=(
  --inference_mode False
  --allow_tf32
  --checkpoints_total_limit 3
  --training_cfg_rate 0.1
  --multi_phased_distill_schedule "4000-1"
  --not_apply_cfg_solver
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
  --enable_gradient_checkpointing_type "full"
 )
 # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
 torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
    fastvideo/v1/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}"
@@ -0,0 +1,130 @@
 #!/bin/bash
 #SBATCH --job-name=i2v
 #SBATCH --partition=main
 #SBATCH --nodes=4
 #SBATCH --ntasks=4
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
 #SBATCH --mem=1440G
 #SBATCH --output=i2v_output/i2v_%j.out
 #SBATCH --error=i2v_output/i2v_%j.err
 #SBATCH --exclusive
 set -e -x
 # Environment Setup
 source ~/conda/miniconda/bin/activate
 conda activate will-fv
 # Basic Info
 export WANDB_MODE="online"
 export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
 # different cache dir for different processes
 export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
 export MASTER_PORT=29500
 export NODE_RANK=$SLURM_PROCID
 nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
 export MASTER_ADDR=${nodes[0]}
 export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
 export TOKENIZERS_PARALLELISM=false
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 echo "MASTER_ADDR: $MASTER_ADDR"
 echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
 # Training arguments
 training_args=(
  --tracker_project_name wan_i2v_finetune
  --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"
  --max_train_steps=2000
  --train_batch_size=2
  --train_sp_batch_size 1
  --gradient_accumulation_steps=1
  --num_latent_t 8
  --num_height 480
  --num_width 832
  --num_frames 77
  --enable_gradient_checkpointing_type "full"
 )
 # Parallel arguments
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size $NUM_GPUS
  --tp_size 1
  --hsdp_replicate_dim $SLURM_JOB_NUM_NODES
  --hsdp_shard_dim $NUM_GPUS
 )
 # Model arguments
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
 )
 # Dataset arguments
 dataset_args=(
  --data_path "$DATA_DIR"
  --dataloader_num_workers 10
 )
 # Validation arguments
 validation_args=(
  --log_validation
  --validation_dataset_file "$VALIDATION_DATASET_FILE"
  --validation_steps 100
  --validation_sampling_steps "40"
  --validation_guidance_scale "1.0"
 )
 # Optimizer arguments
 optimizer_args=(
  --learning_rate=1e-5
  --mixed_precision="bf16"
  --checkpointing_steps=1000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
 # Miscellaneous arguments
 miscellaneous_args=(
  --inference_mode False
  --allow_tf32
  --checkpoints_total_limit 3
  --training_cfg_rate 0.1
  --multi_phased_distill_schedule "4000-1"
  --not_apply_cfg_solver
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
  --enable_gradient_checkpointing_type "full"
 )
 srun torchrun \
 --nnodes $SLURM_JOB_NUM_NODES \
 --nproc_per_node $NUM_GPUS \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
    fastvideo/v1/training/wan_i2v_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}"
@@ -0,0 +1,26 @@
 #!/bin/bash
 GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_i2v/"
 VALIDATION_PATH="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json"
 torchrun --nproc_per_node=$GPU_NUM \
    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
    --model_path $MODEL_PATH \
    --data_merge_path $DATA_MERGE_PATH \
    --preprocess_video_batch_size 8 \
    --max_height 480 \
    --max_width 832 \
    --num_frames 77 \
    --dataloader_num_workers 0 \
    --output_dir=$OUTPUT_DIR \
    --model_type $MODEL_TYPE \
    --train_fps 16 \
    --validation_dataset_file $VALIDATION_PATH \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
    --preprocess_task "i2v" 
@@ -0,0 +1,31 @@
 {
  "data": [
    {
      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
      "image_path": null,
      "video_path": "validation_dataset/yYcK4nANZz4-Scene-034.mp4",
      "num_inference_steps": 40,
      "height": 480,
      "width": 832,
      "num_frames": 77
    },
    {
      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
      "image_path": null,
      "video_path": "validation_dataset/yYcK4nANZz4-Scene-027.mp4",
      "num_inference_steps": 40,
      "height": 480,
      "width": 832,
      "num_frames": 77
    },
    {
      "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
      "image_path": null,
      "video_path": "validation_dataset/yYcK4nANZz4-Scene-030.mp4",
      "num_inference_steps": 40,
      "height": 480,
      "width": 832,
      "num_frames": 77
    }
  ]
 }
@@ -0,0 +1,10 @@
 This directory contain e2e examples scripts for finetuning Wan2.1 T2v. 
 Execute the following commands from `FastVideo/` to run training:
 - Download crush-smol dataset:
 `bash examples/training/finetune/wan_t2v_1_3b/crush_smol/download_dataset.sh`
 - Preprocess the videos and captions into latents:
 `bash examples/training/finetune/wan_t2v_1_3b/crush_smol/preprocess_wan_data_t2v.sh`
 - Edit the following file and run finetuning:
 `bash examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh`
@@ -0,0 +1,3 @@
 #!/bin/bash
 python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"
@@ -0,0 +1,93 @@
 #!/bin/bash
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 export TOKENIZERS_PARALLELISM=false
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
 NUM_GPUS=4
 # export CUDA_VISIBLE_DEVICES=4,5
 # Training arguments
 training_args=(
  --tracker_project_name "wan_t2v_finetune"
  --output_dir "outputs/wan_t2v_finetune"
  --max_train_steps 5000
  --train_batch_size 1
  --train_sp_batch_size 1
  --gradient_accumulation_steps 8
  --num_latent_t 8 
  --num_height 480
  --num_width 832
  --num_frames 77
  --enable_gradient_checkpointing_type "full"
 )
 # Parallel arguments
 parallel_args=(
  --num_gpus $NUM_GPUS 
  --sp_size $NUM_GPUS 
  --tp_size 1
  --hsdp_replicate_dim 1
  --hsdp_shard_dim $NUM_GPUS
 )
 # Model arguments
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
 )
 # Dataset arguments
 dataset_args=(
  --data_path $DATA_DIR
  --dataloader_num_workers 1
 )
 # Validation arguments
 validation_args=(
  --log_validation 
  --validation_dataset_file $VALIDATION_DATASET_FILE
  --validation_steps 50
  --validation_sampling_steps "50" 
  --validation_guidance_scale "1.0"
 )
 # Optimizer arguments
 optimizer_args=(
  --learning_rate 5e-5
  --mixed_precision "bf16"
  --checkpointing_steps 6000
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
 # Miscellaneous arguments
 miscellaneous_args=(
  --inference_mode False
  --allow_tf32
  --checkpoints_total_limit 3
  --training_cfg_rate 0.1
  --multi_phased_distill_schedule "4000-1"
  --not_apply_cfg_solver
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
  --enable_gradient_checkpointing_type "full"
 )
 torchrun \
  --nnodes 1 \
  --nproc_per_node $NUM_GPUS \
    fastvideo/v1/training/wan_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}"
@@ -0,0 +1,127 @@
 #!/bin/bash
 #SBATCH --job-name=t2v
 #SBATCH --partition=main
 #SBATCH --nodes=1
 #SBATCH --ntasks=1
 #SBATCH --ntasks-per-node=1
 #SBATCH --gres=gpu:8
 #SBATCH --cpus-per-task=128
 #SBATCH --mem=1440G
 #SBATCH --output=t2v_output/t2v_%j.out
 #SBATCH --error=t2v_output/t2v_%j.err
 #SBATCH --exclusive
 set -e -x
 # Environment Setup
 source ~/conda/miniconda/bin/activate
 conda activate will-fv
 # Basic Info
 export WANDB_MODE="online"
 export NCCL_P2P_DISABLE=1
 export TORCH_NCCL_ENABLE_MONITORING=0
 # different cache dir for different processes
 export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
 export MASTER_PORT=29500
 export NODE_RANK=$SLURM_PROCID
 nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) )
 export MASTER_ADDR=${nodes[0]}
 export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
 export TOKENIZERS_PARALLELISM=false
 export WANDB_BASE_URL="https://api.wandb.ai"
 export WANDB_MODE=online
 # export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
 echo "MASTER_ADDR: $MASTER_ADDR"
 echo "NODE_RANK: $NODE_RANK"
 # Configs
 NUM_GPUS=8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
 VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
 # export CUDA_VISIBLE_DEVICES=4,5
 # IP=[MASTER NODE IP]
 # Training arguments
 training_args=(
  --tracker_project_name wan_t2v_finetune
  --output_dir="outputs/wan_t2v_finetune"
  --max_train_steps=1000
  --train_batch_size=4
  --train_sp_batch_size 1
  --gradient_accumulation_steps=1
  --num_latent_t 8
  --num_height 480
  --num_width 832
  --num_frames 77
  --enable_gradient_checkpointing_type "full"
 )
 # Parallel arguments
 parallel_args=(
  --num_gpus $NUM_GPUS
  --sp_size 4
  --tp_size 1
  --hsdp_replicate_dim 2
  --hsdp_shard_dim 4
 )
 # Model arguments
 model_args=(
  --model_path $MODEL_PATH
  --pretrained_model_name_or_path $MODEL_PATH
 )
 # Dataset arguments
 dataset_args=(
  --data_path "$DATA_DIR"
  --dataloader_num_workers 10
 )
 # Validation arguments
 validation_args=(
  --log_validation
  --validation_dataset_file "$VALIDATION_DATASET_FILE"
  --validation_steps 100
  --validation_sampling_steps "50"
  --validation_guidance_scale "1.0"
 )
 # Optimizer arguments
 optimizer_args=(
  --learning_rate=5e-5
  --mixed_precision="bf16"
  --checkpointing_steps=500
  --weight_decay 1e-4
  --max_grad_norm 1.0
 )
 # Miscellaneous arguments
 miscellaneous_args=(
  --inference_mode False
  --allow_tf32
  --checkpoints_total_limit 3
  --training_cfg_rate 0.1
  --multi_phased_distill_schedule "4000-1"
  --not_apply_cfg_solver
  --dit_precision "fp32"
  --num_euler_timesteps 50
  --ema_start_step 0
  --enable_gradient_checkpointing_type "full"
 )
 srun torchrun \
 --nnodes $SLURM_JOB_NUM_NODES \
 --nproc_per_node $NUM_GPUS \
 --node_rank $SLURM_PROCID \
 --rdzv_backend=c10d \
 --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
    fastvideo/v1/training/wan_training_pipeline.py \
    "${parallel_args[@]}" \
    "${model_args[@]}" \
    "${dataset_args[@]}" \
    "${training_args[@]}" \
    "${optimizer_args[@]}" \
    "${validation_args[@]}" \
    "${miscellaneous_args[@]}"
@@ -0,0 +1,26 @@
 #!/bin/bash
 GPU_NUM=1 # 2,4,8
 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 MODEL_TYPE="wan"
 DATA_MERGE_PATH="data/crush-smol/merge.txt"
 OUTPUT_DIR="data/crush-smol_processed_t2v/"
 VALIDATION_PATH="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json"
 torchrun --nproc_per_node=$GPU_NUM \
    fastvideo/v1/pipelines/preprocess/v1_preprocess.py \
    --model_path $MODEL_PATH \
    --data_merge_path $DATA_MERGE_PATH \
    --preprocess_video_batch_size 8 \
    --max_height 480 \
    --max_width 832 \
    --num_frames 77 \
    --dataloader_num_workers 0 \
    --output_dir=$OUTPUT_DIR \
    --model_type $MODEL_TYPE \
    --train_fps 16 \
    --validation_dataset_file $VALIDATION_PATH \
    --samples_per_file 8 \
    --flush_frequency 8 \
    --video_length_tolerance_range 5 \
    --preprocess_task "t2v" 
@@ -0,0 +1,31 @@
 {
  "data": [
    {
      "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.",
      "image_path": null,
      "video_path": null,
      "num_inference_steps": 50,
      "height": 480,
      "width": 832,
      "num_frames": 77
    },
    {
      "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.",
      "image_path": null,
      "video_path": null,
      "num_inference_steps": 50,
      "height": 480,
      "width": 832,
      "num_frames": 77
    },
    {
      "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.",
      "image_path": null,
      "video_path": null,
      "num_inference_steps": 50,
      "height": 480,
      "width": 832,
      "num_frames": 77
    }
  ]
 }
@@ -62,6 +62,7 @@ SystemEnv = namedtuple(
 DEFAULT_CONDA_PATTERNS = {
    "torch",
    "numpy",
    "mypy"
    "cudatoolkit",
    "soumith",
    "mkl",
@@ -80,7 +81,6 @@ DEFAULT_CONDA_PATTERNS = {
 DEFAULT_PIP_PATTERNS = {
    "torch",
    "numpy",
    "mypy",
    "flake8",
    "triton",
    "optree",
@@ -2,16 +2,18 @@
 import json
 import os
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any
 import numpy as np
 from fastvideo.v1.utils import dict_to_3d_list
 def configure_sta(mode: str = 'STA_searching',
                  layer_num: int = 40,
                  time_step_num: int = 50,
                  head_num: int = 40,
-                  **kwargs) -> List[List[List[Any]]]:
+                  **kwargs) -> list[list[list[Any]]]:
    """
    Configure Sliding Tile Attention (STA) parameters based on the specified mode.
@@ -51,22 +53,22 @@ def configure_sta(mode: str = 'STA_searching',
    if mode == 'STA_searching':
        # Get parameters with defaults
-        mask_candidates: Optional[List[str]] = kwargs.get('mask_candidates')
+        mask_candidates: list[str] | None = kwargs.get('mask_candidates')
        if mask_candidates is None:
            raise ValueError(
                "mask_candidates is required for STA_searching mode")
-        mask_selected: List[int] = kwargs.get('mask_selected',
+        mask_selected: list[int] = kwargs.get('mask_selected',
                                              list(range(len(mask_candidates))))
        # Parse selected masks
-        selected_masks: List[List[int]] = []
+        selected_masks: list[list[int]] = []
        for index in mask_selected:
            mask = mask_candidates[index]
            masks_list = [int(x) for x in mask.split(',')]
            selected_masks.append(masks_list)
        # Create 3D mask structure with fixed dimensions (t=50, l=60)
-        masks_3d: List[List[List[List[int]]]] = []
+        masks_3d: list[list[list[list[int]]]] = []
        for i in range(time_step_num):  # Fixed t dimension = 50
            row = []
            for j in range(layer_num):  # Fixed l dimension = 60
@@ -77,25 +79,23 @@ def configure_sta(mode: str = 'STA_searching',
    elif mode == 'STA_tuning':
        # Get required parameters
-        mask_search_files_path: Optional[str] = kwargs.get(
+        mask_search_files_path: str | None = kwargs.get(
            'mask_search_files_path')
        if not mask_search_files_path:
            raise ValueError(
                "mask_search_files_path is required for STA_tuning mode")
        # Get optional parameters with defaults
-        mask_candidates_tuning: Optional[List[str]] = kwargs.get(
+        mask_candidates_tuning: list[str] | None = kwargs.get('mask_candidates')
            'mask_candidates')
        if mask_candidates_tuning is None:
            raise ValueError("mask_candidates is required for STA_tuning mode")
-        mask_selected_tuning: List[int] = kwargs.get(
+        mask_selected_tuning: list[int] = kwargs.get(
            'mask_selected', list(range(len(mask_candidates_tuning))))
-        skip_time_steps_tuning: Optional[int] = kwargs.get('skip_time_steps')
+        skip_time_steps_tuning: int | None = kwargs.get('skip_time_steps')
-        save_dir_tuning: Optional[str] = kwargs.get('save_dir',
+        save_dir_tuning: str | None = kwargs.get('save_dir', "mask_candidates")
                                                    "mask_candidates")
        # Parse selected masks
-        selected_masks_tuning: List[List[int]] = []
+        selected_masks_tuning: list[list[int]] = []
        for index in mask_selected_tuning:
            mask = mask_candidates_tuning[index]
            masks_list = [int(x) for x in mask.split(',')]
@@ -106,7 +106,7 @@ def configure_sta(mode: str = 'STA_searching',
        averaged_results = average_head_losses(results, selected_masks_tuning)
        # Add full attention mask for specific cases
-        full_attention_mask_tuning: Optional[List[int]] = kwargs.get(
+        full_attention_mask_tuning: list[int] | None = kwargs.get(
            'full_attention_mask')
        if full_attention_mask_tuning is not None:
            selected_masks_tuning.append(full_attention_mask_tuning)
@@ -147,11 +147,11 @@ def configure_sta(mode: str = 'STA_searching',
        return mask_strategy_3d
    elif mode == 'STA_tuning_cfg':
        # Get required parameters for both positive and negative paths
-        mask_search_files_path_pos: Optional[str] = kwargs.get(
+        mask_search_files_path_pos: str | None = kwargs.get(
            'mask_search_files_path_pos')
-        mask_search_files_path_neg: Optional[str] = kwargs.get(
+        mask_search_files_path_neg: str | None = kwargs.get(
            'mask_search_files_path_neg')
-        save_dir_cfg: Optional[str] = kwargs.get('save_dir')
+        save_dir_cfg: str | None = kwargs.get('save_dir')
        if not mask_search_files_path_pos or not mask_search_files_path_neg or not save_dir_cfg:
            raise ValueError(
@@ -159,16 +159,16 @@ def configure_sta(mode: str = 'STA_searching',
            )
        # Get optional parameters with defaults
-        mask_candidates_cfg: Optional[List[str]] = kwargs.get('mask_candidates')
+        mask_candidates_cfg: list[str] | None = kwargs.get('mask_candidates')
        if mask_candidates_cfg is None:
            raise ValueError(
                "mask_candidates is required for STA_tuning_cfg mode")
-        mask_selected_cfg: List[int] = kwargs.get(
+        mask_selected_cfg: list[int] = kwargs.get(
            'mask_selected', list(range(len(mask_candidates_cfg))))
-        skip_time_steps_cfg: Optional[int] = kwargs.get('skip_time_steps')
+        skip_time_steps_cfg: int | None = kwargs.get('skip_time_steps')
        # Parse selected masks
-        selected_masks_cfg: List[List[int]] = []
+        selected_masks_cfg: list[list[int]] = []
        for index in mask_selected_cfg:
            mask = mask_candidates_cfg[index]
            masks_list = [int(x) for x in mask.split(',')]
@@ -185,7 +185,7 @@ def configure_sta(mode: str = 'STA_searching',
                                               selected_masks_cfg)
        # Add full attention mask for specific cases
-        full_attention_mask_cfg: Optional[List[int]] = kwargs.get(
+        full_attention_mask_cfg: list[int] | None = kwargs.get(
            'full_attention_mask')
        if full_attention_mask_cfg is not None:
            selected_masks_cfg.append(full_attention_mask_cfg)
@@ -225,7 +225,7 @@ def configure_sta(mode: str = 'STA_searching',
    else:  # STA_inference
        # Get parameters with defaults
-        load_path: Optional[str] = kwargs.get(
+        load_path: str | None = kwargs.get(
            'load_path', "mask_candidates/mask_strategy.json")
        if load_path is None:
            raise ValueError("load_path is required for STA_inference mode")
@@ -246,9 +246,9 @@ def configure_sta(mode: str = 'STA_searching',
 # Helper functions
-def read_specific_json_files(folder_path: str) -> List[Dict[str, Any]]:
+def read_specific_json_files(folder_path: str) -> list[dict[str, Any]]:
    """Read and parse JSON files containing mask search results."""
-    json_contents: List[Dict[str, Any]] = []
+    json_contents: list[dict[str, Any]] = []
    # List files only in the current directory (no walk)
    files = os.listdir(folder_path)
@@ -266,11 +266,11 @@ def read_specific_json_files(folder_path: str) -> List[Dict[str, Any]]:
 def average_head_losses(
-        results: List[Dict[str, Any]],
+        results: list[dict[str, Any]],
-        selected_masks: List[List[int]]) -> Dict[str, Dict[str, np.ndarray]]:
+        selected_masks: list[list[int]]) -> dict[str, dict[str, np.ndarray]]:
    """Average losses across all prompts for each mask strategy."""
    # Initialize a dictionary to store the averaged results
-    averaged_losses: Dict[str, Dict[str, np.ndarray]] = {}
+    averaged_losses: dict[str, dict[str, np.ndarray]] = {}
    loss_type = 'L2_loss'
    # Get all loss types (e.g., 'L2_loss')
    averaged_losses[loss_type] = {}
@@ -292,14 +292,14 @@ def average_head_losses(
 def select_best_mask_strategy(
-        averaged_results: Dict[str, Dict[str, np.ndarray]],
+        averaged_results: dict[str, dict[str, np.ndarray]],
-        selected_masks: List[List[int]],
+        selected_masks: list[list[int]],
        skip_time_steps: int = 12,
        timesteps: int = 50,
        head_num: int = 40
-) -> Tuple[Dict[str, List[int]], float, Dict[str, int]]:
+) -> tuple[dict[str, list[int]], float, dict[str, int]]:
    """Select the best mask strategy for each head based on loss minimization."""
-    best_mask_strategy: Dict[str, List[int]] = {}
+    best_mask_strategy: dict[str, list[int]] = {}
    loss_type = 'L2_loss'
    # Get the shape of time steps and layers
    layers = len(averaged_results[loss_type][str(selected_masks[0])][0])
@@ -308,7 +308,7 @@ def select_best_mask_strategy(
    total_tokens = 0  # total number of masked tokens
    total_length = 0  # total sequence length
-    strategy_counts: Dict[str, int] = {
+    strategy_counts: dict[str, int] = {
        str(strategy): 0
        for strategy in selected_masks
    }
@@ -349,38 +349,23 @@ def select_best_mask_strategy(
    return best_mask_strategy, overall_sparsity, strategy_counts
 def dict_to_3d_list(mask_strategy: Optional[Dict[str, List[int]]],
                    t_max: int = 50,
                    l_max: int = 60,
                    h_max: int = 24) -> List[List[List[Optional[List[int]]]]]:
    result: List[List[List[Optional[List[int]]]]] = [[[
        None for _ in range(h_max)
    ] for _ in range(l_max)] for _ in range(t_max)]
    if mask_strategy is None:
        return result
    for key, value in mask_strategy.items():
        t, layer_idx, h = map(int, key.split('_'))
        result[t][layer_idx][h] = value
    return result
 def save_mask_search_results(
-        mask_search_final_result: List[Dict[str, List[float]]],
+        mask_search_final_result: list[dict[str, list[float]]],
        prompt: str,
-        mask_strategies: List[str],
+        mask_strategies: list[str],
-        output_dir: str = 'output/mask_search_result/') -> Optional[str]:
+        output_dir: str = 'output/mask_search_result/') -> str | None:
    if not mask_search_final_result:
        print("No mask search results to save")
        return None
    # Create result dictionary with defaultdict for nested lists
-    mask_search_dict: Dict[str, Dict[str, List[List[float]]]] = {
+    mask_search_dict: dict[str, dict[str, list[list[float]]]] = {
        "L2_loss": defaultdict(list),
        "L1_loss": defaultdict(list)
    }
    mask_selected = list(range(len(mask_strategies)))
-    selected_masks: List[List[int]] = []
+    selected_masks: list[list[int]] = []
    for index in mask_selected:
        mask = mask_strategies[index]
        masks_list = [int(x) for x in mask.split(',')]
@@ -390,7 +375,7 @@ def save_mask_search_results(
    for i, mask_strategy in enumerate(selected_masks):
        mask_strategy_str = str(mask_strategy)
        # Process L2 loss
-        step_results: List[List[float]] = []
+        step_results: list[list[float]] = []
        for step_data in mask_search_final_result:
            if isinstance(step_data, dict) and "L2_loss" in step_data:
                layer_losses = [float(loss) for loss in step_data["L2_loss"]]
@@ -10,8 +10,8 @@ from fastvideo.v1.attention.selector import get_attn_backend
 __all__ = [
    "DistributedAttention",
    "DistributedAttention_VSA",
    "LocalAttention",
    "DistributedAttention_VSA",
    "AttentionBackend",
    "AttentionMetadata",
    "AttentionMetadataBuilder",
@@ -3,8 +3,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, fields
-from typing import (TYPE_CHECKING, Any, Dict, Generic, Optional, Protocol, Set,
+from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar
                    Type, TypeVar)
 if TYPE_CHECKING:
    from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -27,12 +26,12 @@ class AttentionBackend(ABC):
    @staticmethod
    @abstractmethod
-    def get_impl_cls() -> Type["AttentionImpl"]:
+    def get_impl_cls() -> type["AttentionImpl"]:
        raise NotImplementedError
    @staticmethod
    @abstractmethod
-    def get_metadata_cls() -> Type["AttentionMetadata"]:
+    def get_metadata_cls() -> type["AttentionMetadata"]:
        raise NotImplementedError
    # @staticmethod
@@ -46,7 +45,7 @@ class AttentionBackend(ABC):
    @staticmethod
    @abstractmethod
-    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
        raise NotImplementedError
@@ -57,8 +56,7 @@ class AttentionMetadata:
    current_timestep: int
    def asdict_zerocopy(self,
-                        skip_fields: Optional[Set[str]] = None
+                        skip_fields: set[str] | None = None) -> dict[str, Any]:
                        ) -> Dict[str, Any]:
        """Similar to dataclasses.asdict, but avoids deepcopying."""
        if skip_fields is None:
            skip_fields = set()
@@ -124,7 +122,7 @@ class AttentionImpl(ABC, Generic[T]):
        head_size: int,
        softmax_scale: float,
        causal: bool = False,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from typing import List, Optional, Type
 import torch
 from flash_attn import flash_attn_func as flash_attn_2_func
@@ -28,7 +26,7 @@ class FlashAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True
    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [32, 64, 96, 128, 160, 192, 224, 256]
    @staticmethod
@@ -36,15 +34,15 @@ class FlashAttentionBackend(AttentionBackend):
        return "FLASH_ATTN"
    @staticmethod
-    def get_impl_cls() -> Type["FlashAttentionImpl"]:
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
        return FlashAttentionImpl
    @staticmethod
-    def get_metadata_cls() -> Type["AttentionMetadata"]:
+    def get_metadata_cls() -> type["AttentionMetadata"]:
        raise NotImplementedError
    @staticmethod
-    def get_builder_cls() -> Type["AttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["AttentionMetadataBuilder"]:
        raise NotImplementedError
@@ -56,7 +54,7 @@ class FlashAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
 from typing import List, Optional, Type
 import torch
 from sageattention import sageattn
@@ -18,7 +17,7 @@ class SageAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True
    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [32, 64, 96, 128, 160, 192, 224, 256]
    @staticmethod
@@ -26,7 +25,7 @@ class SageAttentionBackend(AttentionBackend):
        return "SAGE_ATTN"
    @staticmethod
-    def get_impl_cls() -> Type["SageAttentionImpl"]:
+    def get_impl_cls() -> type["SageAttentionImpl"]:
        return SageAttentionImpl
    # @staticmethod
@@ -42,7 +41,7 @@ class SageAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
 from typing import List, Optional, Type
 import torch
@@ -17,7 +16,7 @@ class SDPABackend(AttentionBackend):
    accept_output_buffer: bool = True
    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [32, 64, 96, 128, 160, 192, 224, 256]
    @staticmethod
@@ -25,7 +24,7 @@ class SDPABackend(AttentionBackend):
        return "SDPA"
    @staticmethod
-    def get_impl_cls() -> Type["SDPAImpl"]:
+    def get_impl_cls() -> type["SDPAImpl"]:
        return SDPAImpl
    # @staticmethod
@@ -41,7 +40,7 @@ class SDPAImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import json
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Type
+from typing import Any
 import torch
 from einops import rearrange
@@ -17,33 +17,11 @@ from fastvideo.v1.fastvideo_args import FastVideoArgs
 from fastvideo.v1.forward_context import ForwardContext, get_forward_context
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.utils import dict_to_3d_list
 logger = init_logger(__name__)
 # TODO(will-refactor): move this to a utils file
 def dict_to_3d_list(
        mask_strategy: Dict[str,
                            Any]) -> 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: int) -> str:
@@ -62,7 +40,7 @@ class SlidingTileAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True
    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        # TODO(will-refactor): check this
        return [32, 64, 96, 128, 160, 192, 224, 256]
@@ -71,22 +49,22 @@ class SlidingTileAttentionBackend(AttentionBackend):
        return "SLIDING_TILE_ATTN"
    @staticmethod
-    def get_impl_cls() -> Type["SlidingTileAttentionImpl"]:
+    def get_impl_cls() -> type["SlidingTileAttentionImpl"]:
        return SlidingTileAttentionImpl
    @staticmethod
-    def get_metadata_cls() -> Type["SlidingTileAttentionMetadata"]:
+    def get_metadata_cls() -> type["SlidingTileAttentionMetadata"]:
        return SlidingTileAttentionMetadata
    @staticmethod
-    def get_builder_cls() -> Type["SlidingTileAttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["SlidingTileAttentionMetadataBuilder"]:
        return SlidingTileAttentionMetadataBuilder
@dataclass
 class SlidingTileAttentionMetadata(AttentionMetadata):
    current_timestep: int
-    STA_param: List[List[
+    STA_param: list[list[
        Any]]  # each timestep with one metadata, shape [num_layers, num_heads]
@@ -120,7 +98,7 @@ class SlidingTileAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 import math
 from dataclasses import dataclass
 from typing import List, Optional, Type
 import torch
 from einops import rearrange
@@ -28,7 +27,7 @@ class VideoSparseAttentionBackend(AttentionBackend):
    accept_output_buffer: bool = True
    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
+    def get_supported_head_sizes() -> list[int]:
        return [64, 128]
    @staticmethod
@@ -36,22 +35,22 @@ class VideoSparseAttentionBackend(AttentionBackend):
        return "VIDEO_SPARSE_ATTN"
    @staticmethod
-    def get_impl_cls() -> Type["VideoSparseAttentionImpl"]:
+    def get_impl_cls() -> type["VideoSparseAttentionImpl"]:
        return VideoSparseAttentionImpl
    @staticmethod
-    def get_metadata_cls() -> Type["VideoSparseAttentionMetadata"]:
+    def get_metadata_cls() -> type["VideoSparseAttentionMetadata"]:
        return VideoSparseAttentionMetadata
    @staticmethod
-    def get_builder_cls() -> Type["VideoSparseAttentionMetadataBuilder"]:
+    def get_builder_cls() -> type["VideoSparseAttentionMetadataBuilder"]:
        return VideoSparseAttentionMetadataBuilder
@dataclass
 class VideoSparseAttentionMetadata(AttentionMetadata):
    current_timestep: int
-    dit_seq_shape: List[int]
+    dit_seq_shape: list[int]
    VSA_sparsity: float
@@ -97,7 +96,7 @@ class VideoSparseAttentionImpl(AttentionImpl):
        head_size: int,
        causal: bool,
        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
+        num_kv_heads: int | None = None,
        prefix: str = "",
        **extra_impl_args,
    ) -> None:
@@ -105,8 +104,8 @@ class VideoSparseAttentionImpl(AttentionImpl):
        sp_group = get_sp_group()
        self.sp_size = sp_group.world_size
        self.VSA_base_tile_size = [4, 4, 4]
-        self.dit_seq_shape: List[int]
+        self.dit_seq_shape: list[int]
-        self.full_window_size: List[int]
+        self.full_window_size: list[int]
        self.img_seq_length: int
    def tile(self, x: torch.Tensor) -> torch.Tensor:
@@ -1,7 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from typing import Optional, Tuple
 import torch
 import torch.nn as nn
@@ -23,11 +21,11 @@ class DistributedAttention(nn.Module):
    def __init__(self,
                 num_heads: int,
                 head_size: int,
-                 num_kv_heads: Optional[int] = None,
+                 num_kv_heads: int | None = None,
-                 softmax_scale: Optional[float] = None,
+                 softmax_scale: float | None = None,
                 causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
-                     AttentionBackendEnum, ...]] = None,
+                 | None = None,
                 prefix: str = "",
                 **extra_impl_args) -> None:
        super().__init__()
@@ -63,10 +61,10 @@ class DistributedAttention(nn.Module):
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
-        replicated_q: Optional[torch.Tensor] = None,
+        replicated_q: torch.Tensor | None = None,
-        replicated_k: Optional[torch.Tensor] = None,
+        replicated_k: torch.Tensor | None = None,
-        replicated_v: Optional[torch.Tensor] = None,
+        replicated_v: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Forward pass for distributed attention.
        Args:
@@ -144,11 +142,11 @@ class DistributedAttention_VSA(DistributedAttention):
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
-        replicated_q: Optional[torch.Tensor] = None,
+        replicated_q: torch.Tensor | None = None,
-        replicated_k: Optional[torch.Tensor] = None,
+        replicated_k: torch.Tensor | None = None,
-        replicated_v: Optional[torch.Tensor] = None,
+        replicated_v: torch.Tensor | None = None,
-        gate_compress: Optional[torch.Tensor] = None,
+        gate_compress: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Forward pass for distributed attention.
        Args:
@@ -208,11 +206,11 @@ class LocalAttention(nn.Module):
    def __init__(self,
                 num_heads: int,
                 head_size: int,
-                 num_kv_heads: Optional[int] = None,
+                 num_kv_heads: int | None = None,
-                 softmax_scale: Optional[float] = None,
+                 softmax_scale: float | None = None,
                 causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[
+                 supported_attention_backends: tuple[AttentionBackendEnum, ...]
-                     AttentionBackendEnum, ...]] = None,
+                 | None = None,
                 **extra_impl_args) -> None:
        super().__init__()
        if softmax_scale is None:
@@ -2,9 +2,10 @@
 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/attention/selector.py
 import os
 from collections.abc import Generator
 from contextlib import contextmanager
 from functools import cache
-from typing import Generator, Optional, Tuple, Type, cast
+from typing import cast
 import torch
@@ -17,7 +18,7 @@ from fastvideo.v1.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
 logger = init_logger(__name__)
-def backend_name_to_enum(backend_name: str) -> Optional[AttentionBackendEnum]:
+def backend_name_to_enum(backend_name: str) -> AttentionBackendEnum | None:
    """
    Convert a string backend name to a _Backend enum value.
@@ -31,7 +32,7 @@ def backend_name_to_enum(backend_name: str) -> Optional[AttentionBackendEnum]:
          None
-def get_env_variable_attn_backend() -> Optional[AttentionBackendEnum]:
+def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
    '''
    Get the backend override specified by the FastVideo attention
    backend environment variable, if one is specified.
@@ -53,11 +54,11 @@ def get_env_variable_attn_backend() -> Optional[AttentionBackendEnum]:
 #
 # THIS SELECTION TAKES PRECEDENCE OVER THE
 # FASTVIDEO ATTENTION BACKEND ENVIRONMENT VARIABLE
-forced_attn_backend: Optional[AttentionBackendEnum] = None
+forced_attn_backend: AttentionBackendEnum | None = None
 def global_force_attn_backend(
-        attn_backend: Optional[AttentionBackendEnum]) -> None:
+        attn_backend: AttentionBackendEnum | None) -> None:
    '''
    Force all attention operations to use a specified backend.
@@ -72,7 +73,7 @@ def global_force_attn_backend(
    forced_attn_backend = attn_backend
-def get_global_forced_attn_backend() -> Optional[AttentionBackendEnum]:
+def get_global_forced_attn_backend() -> AttentionBackendEnum | None:
    '''
    Get the currently-forced choice of attention backend,
    or None if auto-selection is currently enabled.
@@ -83,9 +84,9 @@ def get_global_forced_attn_backend() -> Optional[AttentionBackendEnum]:
 def get_attn_backend(
    head_size: int,
    dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
+    supported_attention_backends: tuple[AttentionBackendEnum, ...]
-                                                 ...]] = None,
+    | None = None,
-) -> Type[AttentionBackend]:
+) -> type[AttentionBackend]:
    return _cached_get_attn_backend(head_size, dtype,
                                    supported_attention_backends)
@@ -94,9 +95,9 @@ def get_attn_backend(
 def _cached_get_attn_backend(
    head_size: int,
    dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
+    supported_attention_backends: tuple[AttentionBackendEnum, ...]
-                                                 ...]] = None,
+    | None = None,
-) -> Type[AttentionBackend]:
+) -> type[AttentionBackend]:
    # Check whether a particular choice of backend was
    # previously forced.
    #
@@ -105,13 +106,13 @@ def _cached_get_attn_backend(
    if not supported_attention_backends:
        raise ValueError("supported_attention_backends is empty")
    selected_backend = None
-    backend_by_global_setting: Optional[AttentionBackendEnum] = (
+    backend_by_global_setting: AttentionBackendEnum | None = (
        get_global_forced_attn_backend())
    if backend_by_global_setting is not None:
        selected_backend = backend_by_global_setting
    else:
        # Check the environment variable and override if specified
-        backend_by_env_var: Optional[str] = envs.FASTVIDEO_ATTENTION_BACKEND
+        backend_by_env_var: str | None = envs.FASTVIDEO_ATTENTION_BACKEND
        if backend_by_env_var is not None:
            selected_backend = backend_name_to_enum(backend_by_env_var)
@@ -123,7 +124,7 @@ def _cached_get_attn_backend(
    if not attention_cls:
        raise ValueError(
            f"Invalid attention backend for {current_platform.device_name}")
-    return cast(Type[AttentionBackend], resolve_obj_by_qualname(attention_cls))
+    return cast(type[AttentionBackend], resolve_obj_by_qualname(attention_cls))
@contextmanager
@@ -4,7 +4,7 @@
  "use_cpu_offload": false,
  "disable_autocast": false,
  "precision": "bf16",
-  "vae_precision": "fp16",
+  "vae_precision": "fp32",
  "vae_tiling": true,
  "vae_sp": true,
  "vae_config": {
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field, fields
-from typing import Any, Dict
+from typing import Any
 from fastvideo.v1.logger import init_logger
@@ -12,7 +12,9 @@ logger = init_logger(__name__)
 # 3. Any field in ArchConfig is fixed upon initialization, and should be hidden away from users
@dataclass
 class ArchConfig:
-    pass
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
        default_factory=list
    )  # mapping from huggingface weight names to custom names
@dataclass
@@ -42,7 +44,7 @@ class ModelConfig:
        self.__dict__.update(state)
    # This should be used only when loading from transformers/diffusers
-    def update_model_arch(self, source_model_dict: Dict[str, Any]) -> None:
+    def update_model_arch(self, source_model_dict: dict[str, Any]) -> None:
        arch_config = self.arch_config
        valid_fields = {f.name for f in fields(arch_config)}
@@ -56,7 +58,7 @@ class ModelConfig:
        if hasattr(arch_config, "__post_init__"):
            arch_config.__post_init__()
-    def update_model_config(self, source_model_dict: Dict[str, Any]) -> None:
+    def update_model_config(self, source_model_dict: dict[str, Any]) -> None:
        assert "arch_config" not in source_model_dict, "Source model config shouldn't contain arch_config."
        valid_fields = {f.name for f in fields(self)}
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Any, List, Optional, Tuple
+from typing import Any
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
 from fastvideo.v1.layers.quantization import QuantizationConfig
@@ -11,9 +11,10 @@ from fastvideo.v1.platforms import AttentionBackendEnum
 class DiTArchConfig(ArchConfig):
    _fsdp_shard_conditions: list = field(default_factory=list)
    _compile_conditions: list = field(default_factory=list)
-    _param_names_mapping: dict = field(default_factory=dict)
+    param_names_mapping: dict = field(default_factory=dict)
-    _lora_param_names_mapping: dict = field(default_factory=dict)
+    reverse_param_names_mapping: dict = field(default_factory=dict)
-    _supported_attention_backends: Tuple[AttentionBackendEnum, ...] = (
+    lora_param_names_mapping: dict = field(default_factory=dict)
    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
        AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN,
        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA,
        AttentionBackendEnum.VIDEO_SPARSE_ATTN)
@@ -21,7 +22,7 @@ class DiTArchConfig(ArchConfig):
    hidden_size: int = 0
    num_attention_heads: int = 0
    num_channels_latents: int = 0
-    exclude_lora_layers: List[str] = field(default_factory=list)
+    exclude_lora_layers: list[str] = field(default_factory=list)
    def __post_init__(self) -> None:
        if not self._compile_conditions:
@@ -34,7 +35,7 @@ class DiTConfig(ModelConfig):
    # FastVideoDiT-specific parameters
    prefix: str = ""
-    quant_config: Optional[QuantizationConfig] = None
+    quant_config: QuantizationConfig | None = None
    @staticmethod
    def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any:
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import List, Optional, Tuple
 import torch
@@ -32,7 +31,7 @@ class HunyuanVideoArchConfig(DiTArchConfig):
    _compile_conditions: list = field(
        default_factory=lambda: [is_double_block, is_single_block, is_txt_in])
-    _param_names_mapping: dict = field(
+    param_names_mapping: dict = field(
        default_factory=lambda: {
            # 1. context_embedder.time_text_embed submodules (specific rules, applied first):
            r"^context_embedder\.time_text_embed\.timestep_embedder\.linear_1\.(.*)$":
@@ -147,6 +146,9 @@ class HunyuanVideoArchConfig(DiTArchConfig):
            r"final_layer.linear.\1",
        })
    # Reverse mapping for saving checkpoints: custom -> hf
    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
    patch_size: int = 2
    patch_size_t: int = 1
    in_channels: int = 16
@@ -157,14 +159,14 @@ class HunyuanVideoArchConfig(DiTArchConfig):
    num_layers: int = 20
    num_single_layers: int = 40
    num_refiner_layers: int = 2
-    rope_axes_dim: Tuple[int, int, int] = (16, 56, 56)
+    rope_axes_dim: tuple[int, int, int] = (16, 56, 56)
    guidance_embeds: bool = False
-    dtype: Optional[torch.dtype] = None
+    dtype: torch.dtype | None = None
    text_embed_dim: int = 4096
    pooled_projection_dim: int = 768
    rope_theta: int = 256
    qk_norm: str = "rms_norm"
-    exclude_lora_layers: List[str] = field(
+    exclude_lora_layers: list[str] = field(
        default_factory=lambda: ["img_in", "txt_in", "time_in", "vector_in"])
    def __post_init__(self):
@@ -1,19 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import List, Optional, Tuple, Union
 from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
 def is_blocks(n: str, m) -> bool:
    return "blocks" in n and str.isdigit(n.split(".")[-1])
@dataclass
 class StepVideoArchConfig(DiTArchConfig):
-    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+    _fsdp_shard_conditions: list = field(
        default_factory=lambda:
        [lambda n, m: "transformer_blocks" in n and n.split(".")[-1].isdigit()])
-    _param_names_mapping: dict = field(
+    param_names_mapping: dict = field(
        default_factory=lambda: {
            # transformer block
            r"^transformer_blocks\.(\d+)\.norm1\.(weight|bias)$":
@@ -41,18 +38,18 @@ class StepVideoArchConfig(DiTArchConfig):
    num_attention_heads: int = 48
    attention_head_dim: int = 128
    in_channels: int = 64
-    out_channels: Optional[int] = 64
+    out_channels: int | None = 64
    num_layers: int = 48
    dropout: float = 0.0
    patch_size: int = 1
    norm_type: str = "ada_norm_single"
    norm_elementwise_affine: bool = False
    norm_eps: float = 1e-6
-    caption_channels: Optional[Union[int, List[int], Tuple[int, ...]]] = field(
+    caption_channels: int | list[int] | tuple[int, ...] | None = field(
        default_factory=lambda: [6144, 1024])
-    attention_type: Optional[str] = "torch"
+    attention_type: str | None = "torch"
-    use_additional_conditions: Optional[bool] = False
+    use_additional_conditions: bool | None = False
-    exclude_lora_layers: List[str] = field(default_factory=lambda: [])
+    exclude_lora_layers: list[str] = field(default_factory=lambda: [])
    def __post_init__(self):
        self.hidden_size = self.num_attention_heads * self.attention_head_dim
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import List, Optional, Tuple
 from fastvideo.v1.configs.models.dits.base import DiTArchConfig, DiTConfig
@@ -13,7 +12,7 @@ def is_blocks(n: str, m) -> bool:
 class WanVideoArchConfig(DiTArchConfig):
    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
-    _param_names_mapping: dict = field(
+    param_names_mapping: dict = field(
        default_factory=lambda: {
            r"^patch_embedding\.(.*)$":
            r"patch_embedding.proj.\1",
@@ -49,12 +48,16 @@ class WanVideoArchConfig(DiTArchConfig):
            r"blocks.\1.ffn.fc_in.\2",
            r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$":
            r"blocks.\1.ffn.fc_out.\2",
-            r"blocks\.(\d+)\.norm2\.(.*)$":
+            r"^blocks\.(\d+)\.norm2\.(.*)$":
            r"blocks.\1.self_attn_residual_norm.norm.\2",
        })
    # Reverse mapping for saving checkpoints: custom -> hf
    reverse_param_names_mapping: dict = field(default_factory=lambda: {})
    # Some LoRA adapters use the original official layer names instead of hf layer names,
    # so apply this before the param_names_mapping
-    _lora_param_names_mapping: dict = field(
+    lora_param_names_mapping: dict = field(
        default_factory=lambda: {
            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2",
            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2",
@@ -70,7 +73,7 @@ class WanVideoArchConfig(DiTArchConfig):
            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
        })
-    patch_size: Tuple[int, int, int] = (1, 2, 2)
+    patch_size: tuple[int, int, int] = (1, 2, 2)
    text_len = 512
    num_attention_heads: int = 40
    attention_head_dim: int = 128
@@ -83,10 +86,10 @@ class WanVideoArchConfig(DiTArchConfig):
    cross_attn_norm: bool = True
    qk_norm: str = "rms_norm_across_heads"
    eps: float = 1e-6
-    image_dim: Optional[int] = None
+    image_dim: int | None = None
-    added_kv_proj_dim: Optional[int] = None
+    added_kv_proj_dim: int | None = None
    rope_max_seq_len: int = 1024
-    exclude_lora_layers: List[str] = field(default_factory=lambda: ["embedder"])
+    exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"])
    def __post_init__(self):
        super().__post_init__()
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any
 import torch
@@ -11,8 +11,8 @@ from fastvideo.v1.platforms import AttentionBackendEnum
@dataclass
 class EncoderArchConfig(ArchConfig):
-    architectures: List[str] = field(default_factory=lambda: [])
+    architectures: list[str] = field(default_factory=lambda: [])
-    _supported_attention_backends: Tuple[AttentionBackendEnum, ...] = (
+    _supported_attention_backends: tuple[AttentionBackendEnum, ...] = (
        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA)
    output_hidden_states: bool = False
    use_return_dict: bool = True
@@ -32,8 +32,11 @@ class TextEncoderArchConfig(EncoderArchConfig):
    output_past: bool = True
    scalable_attention: bool = True
    tie_word_embeddings: bool = False
-
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
-    tokenizer_kwargs: Dict[str, Any] = field(default_factory=dict)
+        default_factory=list
    )  # mapping from huggingface weight names to custom names
    tokenizer_kwargs: dict[str, Any] = field(default_factory=dict)
    _fsdp_shard_conditions: list = field(default_factory=lambda: [])
    def __post_init__(self) -> None:
        self.tokenizer_kwargs = {
@@ -50,11 +53,11 @@ class ImageEncoderArchConfig(EncoderArchConfig):
@dataclass
 class BaseEncoderOutput:
-    last_hidden_state: Optional[torch.FloatTensor] = None
+    last_hidden_state: torch.FloatTensor | None = None
-    pooler_output: Optional[torch.FloatTensor] = None
+    pooler_output: torch.FloatTensor | None = None
-    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
-    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
-    attention_mask: Optional[torch.Tensor] = None
+    attention_mask: torch.Tensor | None = None
@dataclass
@@ -62,8 +65,8 @@ class EncoderConfig(ModelConfig):
    arch_config: ArchConfig = field(default_factory=EncoderArchConfig)
    prefix: str = ""
-    quant_config: Optional[QuantizationConfig] = None
+    quant_config: QuantizationConfig | None = None
-    lora_config: Optional[Any] = None
+    lora_config: Any | None = None
@dataclass
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Optional
 from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                       ImageEncoderConfig,
@@ -8,6 +7,14 @@ from fastvideo.v1.configs.models.encoders.base import (ImageEncoderArchConfig,
                                                       TextEncoderConfig)
 def _is_transformer_layer(n: str, m) -> bool:
    return "layers" in n and str.isdigit(n.split(".")[-1])
 def _is_embeddings(n: str, m) -> bool:
    return n.endswith("embeddings")
@dataclass
 class CLIPTextArchConfig(TextEncoderArchConfig):
    vocab_size: int = 49408
@@ -27,6 +34,15 @@ class CLIPTextArchConfig(TextEncoderArchConfig):
    bos_token_id: int = 49406
    eos_token_id: int = 49407
    text_len: int = 77
    stacked_params_mapping: list[tuple[str, str,
                                       str]] = field(default_factory=lambda: [
                                           # (param_name, shard_name, shard_id)
                                           ("qkv_proj", "q_proj", "q"),
                                           ("qkv_proj", "k_proj", "k"),
                                           ("qkv_proj", "v_proj", "v"),
                                       ])
    _fsdp_shard_conditions: list = field(
        default_factory=lambda: [_is_transformer_layer, _is_embeddings])
@dataclass
@@ -45,6 +61,13 @@ class CLIPVisionArchConfig(ImageEncoderArchConfig):
    attention_dropout: float = 0.0
    initializer_range: float = 0.02
    initializer_factor: float = 1.0
    stacked_params_mapping: list[tuple[str, str,
                                       str]] = field(default_factory=lambda: [
                                           # (param_name, shard_name, shard_id)
                                           ("qkv_proj", "q_proj", "q"),
                                           ("qkv_proj", "k_proj", "k"),
                                           ("qkv_proj", "v_proj", "v"),
                                       ])
@dataclass
@@ -52,8 +75,8 @@ class CLIPTextConfig(TextEncoderConfig):
    arch_config: TextEncoderArchConfig = field(
        default_factory=CLIPTextArchConfig)
-    num_hidden_layers_override: Optional[int] = None
+    num_hidden_layers_override: int | None = None
-    require_post_norm: Optional[bool] = None
+    require_post_norm: bool | None = None
    prefix: str = "clip"
@@ -62,6 +85,6 @@ class CLIPVisionConfig(ImageEncoderConfig):
    arch_config: ImageEncoderArchConfig = field(
        default_factory=CLIPVisionArchConfig)
-    num_hidden_layers_override: Optional[int] = None
+    num_hidden_layers_override: int | None = None
-    require_post_norm: Optional[bool] = None
+    require_post_norm: bool | None = None
    prefix: str = "clip"
@@ -1,11 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Optional
 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                       TextEncoderConfig)
 def _is_transformer_layer(n: str, m) -> bool:
    return "layers" in n and str.isdigit(n.split(".")[-1])
 def _is_embeddings(n: str, m) -> bool:
    return n.endswith("embed_tokens")
 def _is_final_norm(n: str, m) -> bool:
    return n.endswith("norm")
@dataclass
 class LlamaArchConfig(TextEncoderArchConfig):
    vocab_size: int = 32000
@@ -13,7 +24,7 @@ class LlamaArchConfig(TextEncoderArchConfig):
    intermediate_size: int = 11008
    num_hidden_layers: int = 32
    num_attention_heads: int = 32
-    num_key_value_heads: Optional[int] = None
+    num_key_value_heads: int | None = None
    hidden_act: str = "silu"
    max_position_embeddings: int = 2048
    initializer_range: float = 0.02
@@ -25,13 +36,25 @@ class LlamaArchConfig(TextEncoderArchConfig):
    pretraining_tp: int = 1
    tie_word_embeddings: bool = False
    rope_theta: float = 10000.0
-    rope_scaling: Optional[float] = None
+    rope_scaling: float | None = None
    attention_bias: bool = False
    attention_dropout: float = 0.0
    mlp_bias: bool = False
-    head_dim: Optional[int] = None
+    head_dim: int | None = None
    hidden_state_skip_layer: int = 2
    text_len: int = 256
    stacked_params_mapping: list[tuple[str, str, str]] = field(
        default_factory=lambda: [
            # (param_name, shard_name, shard_id)
            (".qkv_proj", ".q_proj", "q"),
            (".qkv_proj", ".k_proj", "k"),
            (".qkv_proj", ".v_proj", "v"),
            (".gate_up_proj", ".gate_proj", 0),  # type: ignore
            (".gate_up_proj", ".up_proj", 1),  # type: ignore
        ])
    _fsdp_shard_conditions: list = field(
        default_factory=lambda:
        [_is_transformer_layer, _is_embeddings, _is_final_norm])
@dataclass
@@ -1,11 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Optional
 from fastvideo.v1.configs.models.encoders.base import (TextEncoderArchConfig,
                                                       TextEncoderConfig)
 def _is_transformer_layer(n: str, m) -> bool:
    return "block" in n and str.isdigit(n.split(".")[-1])
 def _is_embeddings(n: str, m) -> bool:
    return n.endswith("shared")
 def _is_final_layernorm(n: str, m) -> bool:
    return n.endswith("final_layer_norm")
@dataclass
 class T5ArchConfig(TextEncoderArchConfig):
    vocab_size: int = 32128
@@ -13,7 +24,7 @@ class T5ArchConfig(TextEncoderArchConfig):
    d_kv: int = 64
    d_ff: int = 2048
    num_layers: int = 6
-    num_decoder_layers: Optional[int] = None
+    num_decoder_layers: int | None = None
    num_heads: int = 8
    relative_attention_num_buckets: int = 32
    relative_attention_max_distance: int = 128
@@ -29,6 +40,16 @@ class T5ArchConfig(TextEncoderArchConfig):
    eos_token_id: int = 1
    classifier_dropout: float = 0.0
    text_len: int = 512
    stacked_params_mapping: list[tuple[str, str,
                                       str]] = field(default_factory=lambda: [
                                           # (param_name, shard_name, shard_id)
                                           (".qkv_proj", ".q", "q"),
                                           (".qkv_proj", ".k", "k"),
                                           (".qkv_proj", ".v", "v"),
                                       ])
    _fsdp_shard_conditions: list = field(
        default_factory=lambda:
        [_is_transformer_layer, _is_embeddings, _is_final_layernorm])
    # Referenced from https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/configuration_t5.py
    def __post_init__(self):
@@ -2,7 +2,7 @@
 import argparse
 import dataclasses
 from dataclasses import dataclass, field
-from typing import Any, Union
+from typing import Any
 import torch
@@ -12,7 +12,7 @@ from fastvideo.v1.utils import StoreBoolean
@dataclass
 class VAEArchConfig(ArchConfig):
-    scaling_factor: Union[float, torch.tensor] = 0
+    scaling_factor: float | torch.Tensor = 0
    temporal_compression_ratio: int = 4
    spatial_compression_ratio: int = 8
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Tuple
 from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
@@ -10,19 +9,19 @@ class HunyuanVAEArchConfig(VAEArchConfig):
    in_channels: int = 3
    out_channels: int = 3
    latent_channels: int = 16
-    down_block_types: Tuple[str, ...] = (
+    down_block_types: tuple[str, ...] = (
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
        "HunyuanVideoDownBlock3D",
    )
-    up_block_types: Tuple[str, ...] = (
+    up_block_types: tuple[str, ...] = (
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
        "HunyuanVideoUpBlock3D",
    )
-    block_out_channels: Tuple[int, ...] = (128, 256, 512, 512)
+    block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
    layers_per_block: int = 2
    act_fn: str = "silu"
    norm_num_groups: int = 32
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Tuple
 import torch
@@ -11,12 +10,12 @@ from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig
 class WanVAEArchConfig(VAEArchConfig):
    base_dim: int = 96
    z_dim: int = 16
-    dim_mult: Tuple[int, ...] = (1, 2, 4, 4)
+    dim_mult: tuple[int, ...] = (1, 2, 4, 4)
    num_res_blocks: int = 2
-    attn_scales: Tuple[float, ...] = ()
+    attn_scales: tuple[float, ...] = ()
-    temperal_downsample: Tuple[bool, ...] = (False, True, True)
+    temperal_downsample: tuple[bool, ...] = (False, True, True)
    dropout: float = 0.0
-    latents_mean: Tuple[float, ...] = (
+    latents_mean: tuple[float, ...] = (
        -0.7571,
        -0.7089,
        -0.9113,
@@ -34,7 +33,7 @@ class WanVAEArchConfig(VAEArchConfig):
        0.2503,
        -0.2921,
    )
-    latents_std: Tuple[float, ...] = (
+    latents_std: tuple[float, ...] = (
        2.8184,
        1.4541,
        2.3275,
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Will Lin	5aee5320b6	add pyproject and docker to ssim test trigger	2025-07-15 13:26:12 -07:00
SolitaryThinker	67bbe56cd4	trigger tests	2025-07-15 13:26:12 -07:00
William Lin	616d43c1cf	[bugfix] [training] use separate generator for validation (#610 )	2025-07-15 13:19:38 -07:00
Wenxuan Tan	7244a4b27f	[CI] Add LoRA inference tests (#546 )	2025-07-15 15:06:44 -05:00
Yongqi Chen	7e5ebb4582	[Feature][Training]Update example fine-tuning scripts to enable gradient checkpointing (#618 )	2025-07-15 11:22:56 -07:00
Wenxuan Tan	65ed588570	Set encoder TP size to 1 by default (#569 )	2025-07-09 17:44:24 -05:00
Wenxuan Tan	14adfe2edc	Remove all unnecessary torch.cuda.empty_cache (#606 )	2025-07-09 16:45:20 -05:00
William Lin	6198c6a640	[docs] update dev guide runpod image to py3.12 (#602 )	2025-07-07 19:56:07 -05:00
William Lin	e6b71b531b	[docs] Update slack invite (#601 )	2025-07-07 14:20:42 -05:00
William Lin	ae1d112c6a	[bugfix] [training] fix deadlock in latent datasets and init error in multi-node training (#598 )	2025-07-06 01:34:15 -05:00
William Lin	bf4de1f38f	[chore] Upgrade min Python version from 3.8 to 3.10 (#597 )	2025-07-04 22:08:31 -05:00
William Lin	66fdcc8e76	[Training] Use inference pipeline for training validation (#585 )	2025-07-04 17:23:55 -05:00
Wenxuan Tan	ed1e8d6bad	[Feature] Offload all text encoders by default (#594 )	2025-07-03 19:30:14 -05:00
Kevin Lin	b9423ca3f8	Add ComfyUI custom node for inference (#596 )	2025-07-03 16:09:03 -05:00
Wenxuan Tan	ad16289871	[LoRA] Fix lora merge weights (#579 )	2025-07-02 00:12:52 -05:00
Wenxuan Tan	2a41da1e6b	Fix VAE precisions (#588 )	2025-07-01 14:49:04 -05:00
William Lin	32133171da	[chore] Release 0.1.1 (#592 )	2025-07-01 01:43:12 -05:00
Kevin Lin	19674c6f29	[CI] Fix fork builds (#590 )	2025-07-01 01:03:34 -05:00
Yongqi Chen	508afb7002	[docs] Update Readme (#591 )	2025-06-30 23:48:45 -05:00
Yongqi Chen	288ea88105	[Feat][Training] Rename weight conversion function and update gradient checkpoint in scripts (#589 )	2025-07-01 00:20:02 -04:00
Jinzhe Pan	eb0f1318f3	[Feat] activation checkpointing (#584 )	2025-06-30 15:24:29 -05:00
William Lin	ce9b5910cc	[Training] add caption to validation log (#582 )	2025-06-30 02:42:17 -05:00
William Lin	d0e5a6214a	[misc] [training] Add --video_length_tolerance_range 10 to preprocessing scripts (#581 )	2025-06-30 02:21:22 -05:00
Wenxuan Tan	834562b2db	[CI] Fix pre-commit CI (#578 )	2025-06-29 16:52:29 -05:00
Wei (Will) Feng	060cc7b9ba	fully_shard usage on RMSNorm (#577 )	2025-06-29 16:35:24 -05:00
Yongqi Chen	6c58a5ba62	[Bugfix]Fix VSA sp for training/inference (#574 )	2025-06-29 13:44:33 -05:00
William Lin	48d9f61f86	[ci] [misc] fix training test threshold (#573 )	2025-06-28 22:17:18 -05:00
William Lin	5f938b5844	[Revert] "[Feature] Load weights from distributed" (#571 )	2025-06-28 20:55:14 -05:00
Wenxuan Tan	74da2a7370	Fix CLIP config (#568 )	2025-06-28 19:01:23 -05:00
Kevin Lin	580d6dfe1f	[CI] Add tests to Modal (#562 )	2025-06-28 14:02:16 -05:00
Wenxuan Tan	344e43006a	[CI] Fix SSIM and transformers CI (#564 )	2025-06-28 00:26:20 -05:00
Wenxuan Tan	c5155b256e	[Feature] Load weights from distributed (#470 )	2025-06-27 22:52:40 -05:00
William Lin	e005c7f3ac	[Docs] [Training] add readme for example training (#563 )	2025-06-27 14:50:42 -05:00
Yongqi Chen	ff5a79ef60	[Feature][Inference] Add VSA inference script (#561 )	2025-06-27 02:19:23 -05:00
William Lin	ab01dc4ba5	[Feature] [Training] Add i2v training (#559 )	2025-06-27 01:56:50 -05:00
William Lin	285a950c1b	[CI] fix vae and ssim tests (#557 )	2025-06-26 23:53:01 -05:00
William Lin	46a0a85d85	[Training] Fixes SP for training; Improve Datasets and schema (#555 )	2025-06-26 21:13:28 -05:00
Yongqi Chen	4aeabbc629	[Feature][Training] Add cfg rate for dataset loader (#556 )	2025-06-26 18:22:37 -04:00
Wenxuan Tan	949bb5c835	[CI] Fix CI checks (#553 )	2025-06-25 14:07:51 -05:00
Wenxuan Tan	aab74c1271	[Kernel] Remove all syncs from STA & VSA kernels (#517 )	2025-06-23 13:13:09 -07:00
Yongqi Chen	f89d86944f	[Feature][Training]Add diffusers format checkpoint saving for inference (#542 )	2025-06-22 01:23:41 -04:00
William Lin	8741d204a5	[Training] Refactor and improve validation datasets (#539 )	2025-06-21 17:58:35 -07:00
Wenxuan Tan	cdc85f58a8	[chore] Bump torch to 2.7.1 to support Blackwell (#483 )	2025-06-20 22:10:56 -07:00
William Lin	0262d2f089	[misc] [training] Reorganize training pipeline (#533 )	2025-06-20 20:42:25 -07:00
William Lin	62c0343465	[bugfix] [VSA] Fix layernorm type for VSA Wan2.1 TransformerBlock (#534 )	2025-06-20 00:24:51 -07:00
William Lin	1e1a023fb0	[bugfix] Fix stage validator for multi text encoder models (#535 )	2025-06-19 22:49:16 -07:00
William Lin	1d2517ad8e	[misc] Remove gradient checking code (#532 )	2025-06-18 23:29:25 -07:00
William Lin	d41186cb4a	[Feat] Add Stage input and output verification (#523 )	2025-06-18 23:29:11 -07:00
Kevin Lin Edenzzzz Wenxuan Tan	78e0c7eec9	Specify cu128 Pytorch installation (#530 ) Co-authored-by: Edenzzzz <wtan45@wisc.edu> Co-authored-by: Wenxuan Tan <wenxuan.tan@wisc.edu>	2025-06-18 20:02:50 -05:00
Wenxuan Tan	1c41a94b62	[Refactor] Move dict_to_3d_list under utils (#507 )	2025-06-18 13:34:37 -07:00
Yongqi Chen	2e66aafe20	[Bugfix][Readme]Fix readme website bugs and add VSA finetune docs (#531 )	2025-06-17 22:48:29 -07:00
Yongqi ChenandWill Lin	55074bda76	[CI] Add STA-inference/VSA-training test (#527 ) Co-authored-by: Will Lin <wlsaidhi@gmail.com>	2025-06-17 21:13:06 -07:00
William Lin	de65bec2b7	[Ci] add sta and vsa install to docker image (#528 )	2025-06-17 18:09:48 -07:00
Yongqi Chen	7664dd0de3	[Bugfix][Inference]Fix envs.attn_backend (#525 )	2025-06-17 18:38:06 -05:00
William Linandkevin314	019a88ced4	[CI][bugfix] Use new 3.12 docker image (#526 ) Co-authored-by: kevin314 <kevin.lin.cs1@gmail.com>	2025-06-17 15:37:08 -07:00
Kevin Lin	72de11abcc	[CI] Add current PR test workflow to Buildkite/Modal (#512 )	2025-06-17 13:29:22 -07:00
Kevin Lin	d71a4ebffc	[CI] Update Docker image to flash-attn 2.8.0 / CUDA 12.8 (#524 )	2025-06-16 17:48:23 -07:00
William Lin	1089ab43bf	[bugfix] [Training] use diffusers fp32layernorm for wan2.1 (#490 )	2025-06-15 22:45:48 -07:00
`@@ -117,4 +117,4 @@ If you're planning to contribute to FastVideo please see the following page:`

	`If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo).`	`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.`	`You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support.`
		`@@ -0,0 +1,3 @@`
							`#!/bin/bash`

							`python scripts/huggingface/download_hf.py --repo_id "wlsaidhi/crush-smol-merged" --local_dir "data/crush-smol" --repo_type "dataset"`