fix denoising stage init

.buildkite/pipeline.yml

@@ -1,66 +0,0 @@
-env:
-  IMAGE_VERSION: "py3.12-latest"
-
-steps:
-  - block: "Start Build"
-    blocked_state: "running"
-    prompt: "Approve build?"
-
-  - label: "Trigger Tests"
-    command: |
-      echo "Current working directory: $(pwd)"
-      echo "Current branch:"
-      git branch --show-current
-      echo "Full diff:"
-      git diff --name-only $BUILDKITE_PULL_REQUEST_BASE_BRANCH...HEAD
-    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/loaders/**"
-                - "fastvideo/v1/tests/encoders/**"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: "Encoder Tests"
-                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
-                  - TEST_TYPE=encoder
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/v1/models/vaes/**"
-                - "fastvideo/v1/models/loaders/**"
-                - "fastvideo/v1/tests/vaes/**"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: "VAE Tests"
-                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
-                  - TEST_TYPE=vae
-                agents:
-                  queue: "default"
-            - path:
-                - "fastvideo/v1/models/dits/**"
-                - "fastvideo/v1/models/loaders/**"
-                - "fastvideo/v1/tests/transformers/**"
-                - "fastvideo/v1/layers/**"
-                - "fastvideo/v1/attention/**"
-              config:
-                command: "timeout 30m .buildkite/scripts/pr_test.sh"
-                label: "Transformer Tests"
-                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
-                  - TEST_TYPE=transformer
-                agents:
-                  queue: "default"
-            - path: "fastvideo/v1/**/*.py"
-              config:
-                command: "timeout 60m .buildkite/scripts/pr_test.sh"
-                label: "SSIM Tests"
-                env:
-                  - BUILDKITE_CLEAN_CHECKOUT=true
-                  - TEST_TYPE=ssim
-                agents:
-                  queue: "default"

.buildkite/scripts/pr_test.sh

@@ -1,91 +0,0 @@
-#!/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)
-
-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"
-
-case "$TEST_TYPE" in
-    "encoder")
-        log "Running encoder tests..."
-        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_encoder_tests"
-        ;;
-    "vae")
-        log "Running VAE tests..."
-        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_vae_tests"
-        ;;
-    "transformer")
-        log "Running transformer tests..."
-        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_transformer_tests"
-        ;;
-    "ssim")
-        log "Running SSIM tests..."
-        MODAL_COMMAND="python3 -m modal run $MODAL_TEST_FILE::run_ssim_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

.github/scripts/runpod_api.py

@@ -160,7 +160,8 @@ def execute_command(pod_id):
     setup_steps = [
         "tar -xzf /tmp/repo.tar.gz --no-same-owner -C /workspace/",
         f"cd /workspace/{repo_name}",
-        "source $HOME/.local/bin/env && source /opt/venv/bin/activate",
+        "source /opt/conda/etc/profile.d/conda.sh",
+        "conda activate fastvideo-dev",
         args.test_command
     ]
     

.github/workflows/pr-test.yml

@@ -12,14 +12,12 @@ on:
     paths:
       - "fastvideo/**/*.py"
       - ".github/workflows/pr-test.yml"
-      - "pyproject.toml"
-      - "docker/Dockerfile.python3.12"
   workflow_dispatch:
     inputs:
       custom_image:
-        description: "Custom image from this repository (default: fastvideo-dev:py3.12-latest)"
+        description: "Custom image from this repository (default: fastvideo-dev:latest)"
         required: false
-        default: "fastvideo-dev:py3.12-latest"
+        default: "fastvideo-dev:latest"
         type: string
       run_encoder_test:
         description: "Run encoder-test"
@@ -41,26 +39,6 @@ on:
         required: false
         default: false
         type: boolean
-      run_training_test:
-        description: "Run training-test"
-        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_nightly_test:
-        description: "Run nightly-test"
-        required: false
-        default: false
-        type: boolean
 
 env:
   PYTHONUNBUFFERED: "1"
@@ -81,9 +59,6 @@ jobs:
       encoder-test: ${{ steps.filter.outputs.encoder-test }}
       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 }}
     steps:
       - uses: actions/checkout@v4
       - uses: dorny/paths-filter@v3
@@ -94,34 +69,16 @@ jobs:
               - 'fastvideo/v1/models/encoders/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/encoders/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
             vae-test:
               - 'fastvideo/v1/models/vaes/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/vaes/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
             transformer-test:
               - 'fastvideo/v1/models/dits/**'
               - 'fastvideo/v1/models/loaders/**'
               - 'fastvideo/v1/tests/transformers/**'
               - 'fastvideo/v1/layers/**'
               - 'fastvideo/v1/attention/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
-            training-test:
-              - 'fastvideo/v1/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
-            training-test-VSA:
-              - 'fastvideo/v1/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
-            inference-test-STA:
-              - 'fastvideo/v1/**'
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.12'
 
   encoder-test:
     needs: change-filter
@@ -134,8 +91,8 @@ jobs:
       gpu_type: "NVIDIA A40"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/encoders -s"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -152,8 +109,8 @@ jobs:
       gpu_type: "NVIDIA A40"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/vaes -s"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -170,8 +127,8 @@ jobs:
       gpu_type: "NVIDIA L40S"
       gpu_count: 1
       volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
+      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:latest' }}"
+      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/transformers -s"
       timeout_minutes: 30
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
@@ -198,88 +155,11 @@ jobs:
       volume_size: 200
       disk_size: 200
       image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
+      test_command: "pip install -e .[test] && pytest ./fastvideo/v1/tests/ssim -vs"
       timeout_minutes: 60
     secrets:
       RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
       RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-  
-  training-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "training-test"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 4
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
-      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' && github.event.pull_request.draft == false) || 
-      (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: 1
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || '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' && github.event.pull_request.draft == false) || 
-      (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: 1
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || '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 }}
-
-  nightly-test:
-    if: >-
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_nightly_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "nightly-test"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 4
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/${{ github.event.inputs.custom_image || 'fastvideo-dev:py3.12-latest' }}"
-      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 }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-      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

.github/workflows/runpod-test.yml

@@ -43,8 +43,6 @@ on:
         required: true
       RUNPOD_PRIVATE_KEY:
         required: true
-      WANDB_API_KEY:
-        required: false
 
 jobs:
   run-test:
@@ -57,7 +55,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.12"
+          python-version: "3.10"
 
       - name: Set up SSH key
         run: |
@@ -74,7 +72,6 @@ jobs:
           JOB_ID: ${{ inputs.job_id }}
           RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
           GITHUB_RUN_ID: ${{ github.run_id }}
-          WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
         timeout-minutes: ${{ inputs.timeout_minutes }}
         run: >-
           python .github/scripts/runpod_api.py 

.github/workflows/sta-publish.yml

@@ -5,7 +5,7 @@ on:
     branches:
       - main
     paths:
-      - "csrc/attn/setup_sta.py"
+      - "csrc/sliding_tile_attention/setup.py"
   workflow_dispatch:
 
 jobs:
@@ -23,13 +23,13 @@ jobs:
       - name: Check if version changed
         id: check-version
         run: |
-          cd csrc/attn
+          cd csrc/sliding_tile_attention
           # Get current commit's version
-          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup_sta.py)
+          NEW_VERSION=$(grep -oP 'VERSION\s*=\s*"\K[^"]+' setup.py)
           echo "New version: $NEW_VERSION"
           
           # Get previous version from git history
-          OLD_VERSION=$(git show HEAD~1:./setup_sta.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
+          OLD_VERSION=$(git show HEAD~1:./setup.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
           echo "Old version: $OLD_VERSION"
           
           if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
@@ -136,21 +136,19 @@ jobs:
 
       - name: Build wheel
         run: |
-          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
-          
           # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
           # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
           # However this still fails so I'm using a newer version of setuptools
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn  # Move into the correct folder
-          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_sta.py bdist_wheel --dist-dir=dist
+          cd csrc/sliding_tile_attention  # Move into the correct folder
+          git submodule update --init --recursive tk # Ensure ThunderKittens submodule is initialized
+          python setup.py bdist_wheel --dist-dir=dist
 
       - name: Rename wheel file
         run: |
-          cd csrc/attn
+          cd csrc/sliding_tile_attention
 
           CUDA_SHORT_VERSION=$(echo ${{ matrix.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
           TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-version }} | cut -d. -f1,2)
@@ -165,7 +163,7 @@ jobs:
         uses: actions/upload-artifact@v4
         with:
           name: ${{ env.wheel_name }}
-          path: csrc/attn/dist/*.whl
+          path: csrc/sliding_tile_attention/dist/*.whl
           retention-days: 90
 
   publish_package:
@@ -231,19 +229,17 @@ jobs:
 
       - name: Build source distribution
         run: |
-          export PYTHONPATH=$GITHUB_WORKSPACE:$PYTHONPATH
-          
           # We want setuptools >= 49.6.0 otherwise we can't compile the extension if system CUDA version is 11.7 and pytorch cuda version is 11.6
           # https://github.com/pytorch/pytorch/blob/664058fa83f1d8eede5d66418abff6e20bd76ca8/torch/utils/cpp_extension.py#L810
           # However this still fails so I'm using a newer version of setuptools
           pip install setuptools
           pip install ninja packaging wheel
 
-          cd csrc/attn  # Move into the correct folder
-          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup_sta.py sdist --dist-dir=dist
+          cd csrc/sliding_tile_attention  # Move into the correct folder
+          git submodule update --init --recursive tk # Ensure ThunderKittens submodule is initialized
+          python setup.py sdist --dist-dir=dist
 
       - name: Publish release distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: csrc/attn/dist/
+          packages-dir: csrc/sliding_tile_attention/dist/

.github/workflows/test.yml

@@ -28,4 +28,4 @@ jobs:
 
       - name: Run Pytest
         run: |
-          pytest --ignore csrc/attn/test
+          pytest --ignore csrc/sliding_tile_attention/test

.gitmodules

@@ -1,3 +1,3 @@
-[submodule "csrc/attn/tk"]
-	path = csrc/attn/tk
+[submodule "csrc/sliding_tile_attention/tk"]
+	path = csrc/sliding_tile_attention/tk
 	url = https://github.com/HazyResearch/ThunderKittens.git

README.md

@@ -91,7 +91,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/finetune.html)
+- [Finetuning Guide](https://hao-ai-lab.github.io/FastVideo/training/finetuning.html)
 
 ## 📑 Development Plan
 
@@ -111,7 +111,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/contributing/overview.html)
+We welcome all contributions. Please check out our guide [here](https://hao-ai-lab.github.io/FastVideo/developer_guide/overview.html)
 
 ## Acknowledgement
 We learned and reused code from the following projects:

csrc/attn/bench/bench_vsa.py

@@ -1,225 +0,0 @@
-import torch
-import argparse
-from flash_attn.utils.benchmark import benchmark_forward
-from vsa import block_sparse_attention_fwd, block_sparse_attention_backward
-from vsa import BLOCK_M, BLOCK_N
-
-import numpy as np
-import random
-
-def set_seed(seed: int = 42):
-    # Python random module
-    random.seed(seed)
-
-    # NumPy
-    np.random.seed(seed)
-
-    # PyTorch
-    torch.manual_seed(seed)
-    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=1, help='Batch size')
-    parser.add_argument('--num_heads', type=int, default=12, help='Number of heads')
-    parser.add_argument('--head_dim', type=int, default=64, help='Head dimension')
-    parser.add_argument('--topk', type=int, default=None, help='Number of kv blocks each q block attends to')
-    parser.add_argument('--seq_lengths', type=int, nargs='+', default=[49152], help='Sequence lengths to benchmark')
-    return parser.parse_args()
-
-def create_input_tensors(batch, head, seq_len, headdim):
-    """Create random input tensors for attention."""
-    q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    return q, k, v
-
-def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"):
-    """
-    Generate a block sparse pattern where each q block attends to exactly k kv blocks.
-    
-    Args:
-        bs: batch size
-        h: number of heads
-        num_q_blocks: number of query blocks
-        num_kv_blocks: number of key-value blocks
-        k: number of kv blocks each q block attends to
-        device: device to create tensors on
-        
-    Returns:
-        q2k_block_sparse_index: [bs, h, num_q_blocks, k]
-            Contains the indices of kv blocks that each q block attends to.
-        q2k_block_sparse_num: [bs, h, num_q_blocks]
-            Contains the number of kv blocks that each q block attends to (all equal to k).
-        k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks]
-            Contains the indices of q blocks that attend to each kv block.
-        k2q_block_sparse_num: [bs, h, num_kv_blocks]
-            Contains the number of q blocks that attend to each kv block.
-        block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]
-            Binary mask where 1 indicates attention connection.
-    """
-    # Ensure k is not larger than num_kv_blocks
-    k = min(k, num_kv_blocks)
-    
-    # Create random scores for sampling
-    scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device)
-    
-    # Get top-k indices for each q block
-    _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1)
-    q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32)
-    
-    # sort q2k_block_sparse_index
-    q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1)
-
-    # All q blocks attend to exactly k kv blocks
-    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device)
-    
-    # Create the corresponding mask
-    block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)
-    
-    # Fill in the mask based on the indices
-    for b in range(bs):
-        for head in range(h):
-            for q_idx in range(num_q_blocks):
-                kv_indices = q2k_block_sparse_index[b, head, q_idx]
-                block_sparse_mask[b, head, q_idx, kv_indices] = True
-    
-    # Create the reverse mapping (k2q)
-    # First, initialize lists to collect q indices for each kv block
-    k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)]
-    
-    # Populate the lists based on q2k mapping
-    for b in range(bs):
-        for head in range(h):
-            flat_idx = b * h + head
-            for q_idx in range(num_q_blocks):
-                kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist()
-                for kv_idx in kv_indices:
-                    k2q_indices_list[flat_idx][kv_idx].append(q_idx)
-    
-    # Find the maximum number of q blocks that attend to any kv block
-    max_q_per_kv = 0
-    for flat_idx in range(bs * h):
-        for kv_idx in range(num_kv_blocks):
-            max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx]))
-    
-    # Create tensors for k2q mapping
-    k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, 
-                                        dtype=torch.int32, device=device)
-    k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), 
-                                       dtype=torch.int32, device=device)
-    
-    # Fill the tensors
-    for b in range(bs):
-        for head in range(h):
-            flat_idx = b * h + head
-            for kv_idx in range(num_kv_blocks):
-                q_indices = k2q_indices_list[flat_idx][kv_idx]
-                num_q = len(q_indices)
-                k2q_block_sparse_num[b, head, kv_idx] = num_q
-                if num_q > 0:
-                    k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor(
-                        q_indices, dtype=torch.int32, device=device)
-                
-    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
-
-def benchmark_block_sparse_attention(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops):
-    """Benchmark block sparse attention forward and backward passes."""
-    print("\n=== BLOCK SPARSE ATTENTION BENCHMARK ===")
-    
-    # Forward pass
-    # Warm-up run
-    o, l_vec = block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
-    torch.cuda.synchronize()
-    
-    # Benchmark forward
-    _, fwd_time = benchmark_forward(
-        block_sparse_attention_fwd,
-        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num,
-        repeats=20,
-        verbose=False,
-        desc='Block Sparse Forward'
-    )
-    
-    sparse_tflops = flops / fwd_time.mean * 1e-12
-    print(f"Block Sparse Forward - TFLOPS: {sparse_tflops:.2f}")
-    
-    # Backward pass
-    grad_output = torch.randn_like(o)
-    
-    # Warm-up runs
-    for _ in range(5):
-        block_sparse_attention_backward(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num)
-    torch.cuda.synchronize()
-    
-    # Benchmark backward
-    _, bwd_time = benchmark_forward(
-        block_sparse_attention_backward,
-        q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num,
-        repeats=20,
-        verbose=False,
-        desc='Block Sparse Backward'
-    )
-    bwd_flops = 2.5 * flops  # Approximation
-
-    sparse_bwd_tflops = bwd_flops / bwd_time.mean * 1e-12
-    print(f"Block Sparse Backward - TFLOPS: {sparse_bwd_tflops:.2f}")
-    
-    return sparse_tflops, sparse_bwd_tflops
-
-def main():
-    args = parse_arguments()
-
-    set_seed(42)
-    
-    # Extract parameters
-    batch = args.batch_size
-    head = args.num_heads
-    headdim = args.head_dim
-    
-    print(f"Block Sparse Attention Benchmark")
-    print(f"batch: {batch}, head: {head}, headdim: {headdim}")
-    
-    # Test with different sequence lengths
-    for seq_len in args.seq_lengths:
-        # Skip very long sequences if they might cause OOM
-        if seq_len > 16384 and batch > 1:
-            continue
-            
-        print("="*100)
-        print(f"\nSequence length: {seq_len}")
-        
-        # Calculate theoretical FLOPs for attention
-        flops = 4 * batch * head * headdim * seq_len * seq_len
-        
-        # Create input tensors
-        q, k, v = create_input_tensors(batch, head, seq_len, headdim)
-        
-        # Setup block sparse parameters
-        num_q_blocks = seq_len // BLOCK_M
-        num_kv_blocks = seq_len // BLOCK_N
-        
-        # Determine k value (number of kv blocks per q block)
-        topk = args.topk
-        if topk is None:
-            topk = num_kv_blocks // 10  # Default to ~90% sparsity if k is not specified
-        topk = max(1, topk)       
-        print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)")
-        
-        # Generate block sparse pattern
-        q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, _ = generate_block_sparse_pattern(
-            batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda")
-        
-        # Benchmark block sparse attention
-        sparse_fwd, sparse_bwd = benchmark_block_sparse_attention(
-            q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, flops
-        )
-        
-        # Print results
-        print("\n=== PERFORMANCE RESULTS ===")
-        print(f"Block Sparse Forward - TFLOPS: {sparse_fwd:.2f}")
-        print(f"Block Sparse Backward - TFLOPS: {sparse_bwd:.2f}")
-
-if __name__ == "__main__":
-    main()

csrc/attn/config_vsa.py

@@ -1,15 +0,0 @@
-### ADD TO THIS TO REGISTER NEW KERNELS
-sources = {
-    'block_sparse': {
-        'source_files': {
-            'h100': 'vsa/block_sparse_h100.cu'
-        }
-    }
-}
-
-### WHICH KERNELS DO WE WANT TO BUILD?
-# (oftentimes during development work you don't need to redefine them all.)
-kernels = ['block_sparse']
-
-### WHICH GPU TARGET DO WE WANT TO BUILD FOR?
-target = 'h100'

csrc/attn/setup_vsa.py

@@ -1,76 +0,0 @@
-import os
-import subprocess
-
-from csrc.attn.config_vsa import kernels, sources, target
-from setuptools import find_packages, setup
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
-
-target = target.lower()
-
-# Package metadata
-PACKAGE_NAME = "vsa"
-VERSION = "0.0.1"
-AUTHOR = "Hao AI Lab"
-DESCRIPTION = "Video Sparse Attention Kernel Used in FastVideo"
-URL = "https://github.com/hao-ai-lab/FastVideo/tree/main/csrc/attn"
-
-# Set environment variables
-tk_root = os.getenv('THUNDERKITTENS_ROOT', os.path.abspath(os.path.join(os.getcwd(), 'tk/')))
-python_include = subprocess.check_output(['python', '-c',
-                                          "import sysconfig; print(sysconfig.get_path('include'))"]).decode().strip()
-torch_include = subprocess.check_output([
-    'python', '-c',
-    "import torch; from torch.utils.cpp_extension import include_paths; print(' '.join(['-I' + p for p in include_paths()]))"
-]).decode().strip()
-print('vsa root:', tk_root)
-print('Python include:', python_include)
-print('Torch include directories:', torch_include)
-
-# CUDA flags
-cuda_flags = [
-    '-DNDEBUG', '-Xcompiler=-Wno-psabi', '-Xcompiler=-fno-strict-aliasing', '--expt-extended-lambda',
-    '--expt-relaxed-constexpr', '-forward-unknown-to-host-compiler', '--use_fast_math', '-std=c++20', '-O3',
-    '-Xnvlink=--verbose', '-Xptxas=--verbose', '-Xptxas=--warn-on-spills', f'-I{tk_root}/include',
-    f'-I{tk_root}/prototype', f'-I{python_include}', '-DTORCH_COMPILE'
-] + torch_include.split()
-cpp_flags = ['-std=c++20', '-O3']
-
-if target == 'h100':
-    cuda_flags.append('-DKITTENS_HOPPER')
-    cuda_flags.append('-arch=sm_90a')
-else:
-    raise ValueError(f'Target {target} not supported')
-
-source_files = ['vsa.cpp']
-for k in kernels:
-    if target not in sources[k]['source_files']:
-        raise KeyError(f'Target {target} not found in source files for kernel {k}')
-    if isinstance(sources[k]['source_files'][target], list):
-        source_files.extend(sources[k]['source_files'][target])
-    else:
-        source_files.append(sources[k]['source_files'][target])
-    cpp_flags.append(f'-DTK_COMPILE_{k.replace(" ", "_").upper()}')
-
-setup(name=PACKAGE_NAME,
-      version=VERSION,
-      author=AUTHOR,
-      description=DESCRIPTION,
-      url=URL,
-      packages=find_packages(),
-      ext_modules=[
-          CUDAExtension('vsa_cuda',
-                        sources=source_files,
-                        extra_compile_args={
-                            'cxx': cpp_flags,
-                            'nvcc': cuda_flags
-                        },
-                        libraries=['cuda'])
-      ],
-      cmdclass={'build_ext': BuildExtension},
-      classifiers=[
-          "Programming Language :: Python :: 3",
-          "Environment :: GPU :: NVIDIA CUDA :: 12",
-          "License :: OSI Approved :: Apache Software License",
-      ],
-      python_requires='>=3.10',
-      install_requires=["torch>=2.5.0"])

csrc/attn/test/test_block_sparse.py

@@ -1,266 +0,0 @@
-import torch
-import argparse
-from flash_attn.utils.benchmark import benchmark_forward
-from flash_attn import flash_attn_func
-from vsa import block_sparse_attention_fwd, block_sparse_attention_backward, BlockSparseAttentionFunction
-from vsa import BLOCK_M, BLOCK_N
-
-import numpy as np
-import random
-
-def set_seed(seed: int = 42):
-    # Python random module
-    random.seed(seed)
-
-    # NumPy
-    np.random.seed(seed)
-
-    # PyTorch
-    torch.manual_seed(seed)
-    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): 
-    x, xx = quant_o.float(), fa2_o.float() 
-    sim = torch.nn.functional.cosine_similarity(x.reshape(1, -1), xx.reshape(1, -1)).item()
-    l1 =   ((x - xx).abs().sum() / xx.abs().sum() ).item()
-    rmse = torch.sqrt(torch.mean((x -xx) ** 2)).item()
-
-    return sim, l1, rmse
-
-def create_input_tensors(batch, head, seq_len, headdim):
-    """Create random input tensors for attention."""
-    q = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    k = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    v = torch.randn(batch, head, seq_len, headdim, dtype=torch.bfloat16, device="cuda")
-    return q, k, v
-
-def generate_block_sparse_pattern(bs, h, num_q_blocks, num_kv_blocks, k, device="cuda"):
-    """
-    Generate a block sparse pattern where each q block attends to exactly k kv blocks.
-    
-    Args:
-        bs: batch size
-        h: number of heads
-        num_q_blocks: number of query blocks
-        num_kv_blocks: number of key-value blocks
-        k: number of kv blocks each q block attends to
-        device: device to create tensors on
-        
-    Returns:
-        q2k_block_sparse_index: [bs, h, num_q_blocks, k]
-            Contains the indices of kv blocks that each q block attends to.
-        q2k_block_sparse_num: [bs, h, num_q_blocks]
-            Contains the number of kv blocks that each q block attends to (all equal to k).
-        k2q_block_sparse_index: [bs, h, num_kv_blocks, num_q_blocks]
-            Contains the indices of q blocks that attend to each kv block.
-        k2q_block_sparse_num: [bs, h, num_kv_blocks]
-            Contains the number of q blocks that attend to each kv block.
-        block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]
-            Binary mask where 1 indicates attention connection.
-    """
-    # Ensure k is not larger than num_kv_blocks
-    k = min(k, num_kv_blocks)
-    
-    # Create random scores for sampling
-    scores = torch.rand(bs, h, num_q_blocks, num_kv_blocks, device=device)
-    
-    # Get top-k indices for each q block
-    _, q2k_block_sparse_index = torch.topk(scores, k, dim=-1)
-    q2k_block_sparse_index = q2k_block_sparse_index.to(torch.int32)
-
-    # sort q2k_block_sparse_index
-    q2k_block_sparse_index, _ = torch.sort(q2k_block_sparse_index, dim=-1)
-
-    # All q blocks attend to exactly k kv blocks
-    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks), k, dtype=torch.int32, device=device)
-    
-    # Create the corresponding mask
-    block_sparse_mask = torch.zeros(bs, h, num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)
-    
-    # Fill in the mask based on the indices
-    for b in range(bs):
-        for head in range(h):
-            for q_idx in range(num_q_blocks):
-                kv_indices = q2k_block_sparse_index[b, head, q_idx]
-                block_sparse_mask[b, head, q_idx, kv_indices] = True
-    
-    # Create the reverse mapping (k2q)
-    # First, initialize lists to collect q indices for each kv block
-    k2q_indices_list = [[[] for _ in range(num_kv_blocks)] for _ in range(bs * h)]
-    
-    # Populate the lists based on q2k mapping
-    for b in range(bs):
-        for head in range(h):
-            flat_idx = b * h + head
-            for q_idx in range(num_q_blocks):
-                kv_indices = q2k_block_sparse_index[b, head, q_idx].tolist()
-                for kv_idx in kv_indices:
-                    k2q_indices_list[flat_idx][kv_idx].append(q_idx)
-    
-    # Find the maximum number of q blocks that attend to any kv block
-    max_q_per_kv = 0
-    for flat_idx in range(bs * h):
-        for kv_idx in range(num_kv_blocks):
-            max_q_per_kv = max(max_q_per_kv, len(k2q_indices_list[flat_idx][kv_idx]))
-    
-    # Create tensors for k2q mapping
-    k2q_block_sparse_index = torch.full((bs, h, num_kv_blocks, max_q_per_kv), -1, 
-                                        dtype=torch.int32, device=device)
-    k2q_block_sparse_num = torch.zeros((bs, h, num_kv_blocks), 
-                                       dtype=torch.int32, device=device)
-    
-    # Fill the tensors
-    for b in range(bs):
-        for head in range(h):
-            flat_idx = b * h + head
-            for kv_idx in range(num_kv_blocks):
-                q_indices = k2q_indices_list[flat_idx][kv_idx]
-                num_q = len(q_indices)
-                k2q_block_sparse_num[b, head, kv_idx] = num_q
-                if num_q > 0:
-                    k2q_block_sparse_index[b, head, kv_idx, :num_q] = torch.tensor(
-                        q_indices, dtype=torch.int32, device=device)
-                
-    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask
-
-def main():
-    args = parse_arguments()
-
-    set_seed(42)
-    
-    # Extract parameters
-    batch = args.batch_size
-    head = args.num_heads
-    headdim = args.head_dim
-    num_iterations = args.num_iterations
-    
-    print(f"Block Sparse Attention Benchmark")
-    print(f"batch: {batch}, head: {head}, headdim: {headdim}, iterations: {num_iterations}")
-    
-    # Test with different sequence lengths
-    for seq_len in args.seq_lengths:
-        # Skip very long sequences if they might cause OOM
-        # if seq_len > 16384 and batch > 1:
-        #     continue
-            
-        print("="*100)
-        print(f"\nSequence length: {seq_len}")
-        
-        # Collect metrics across iterations
-        forward_metrics = {'sim': [], 'l1': [], 'rmse': []}
-        grad_q_metrics = {'sim': [], 'l1': [], 'rmse': []}
-        grad_k_metrics = {'sim': [], 'l1': [], 'rmse': []}
-        grad_v_metrics = {'sim': [], 'l1': [], 'rmse': []}
-        
-        for iter_idx in range(num_iterations):
-            if num_iterations > 1:
-                print(f"\nIteration {iter_idx+1}/{num_iterations}")
-            
-            # Create input tensors
-            q, k, v = create_input_tensors(batch, head, seq_len, headdim)
-            
-            # Setup block sparse parameters
-            num_q_blocks = seq_len // BLOCK_M
-            num_kv_blocks = seq_len // BLOCK_N
-            
-            # Determine k value (number of kv blocks per q block)
-            topk = args.topk
-            if topk is None:
-                topk = num_kv_blocks // 10  # Default to ~90% sparsity if k is not specified
-            topk = max(1, topk)
-            if iter_idx == 0:  # Only print this once
-                print(f"Using topk={topk} kv blocks per q block (out of {num_kv_blocks} total kv blocks)")
-
-            # Generate block sparse pattern
-            q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num, block_sparse_mask = generate_block_sparse_pattern(
-                batch, head, num_q_blocks, num_kv_blocks, topk, device="cuda")
-
-            # expand block_sparse_mask to full mask
-            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)
-            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)
-            forward_metrics['sim'].append(sim)
-            forward_metrics['l1'].append(l1)
-            forward_metrics['rmse'].append(rmse)
-
-            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)
-            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)
-            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)
-            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}")
-        
-        # 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"L1 error: mean={np.mean(forward_metrics['l1']):.6f}, std={np.std(forward_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(forward_metrics['rmse']):.6f}, std={np.std(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"L1 error: mean={np.mean(grad_q_metrics['l1']):.6f}, std={np.std(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("\nGradient K metrics:")
-            print(f"Similarity: mean={np.mean(grad_k_metrics['sim']):.6f}, std={np.std(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"RMSE: mean={np.mean(grad_k_metrics['rmse']):.6f}, std={np.std(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"L1 error: mean={np.mean(grad_v_metrics['l1']):.6f}, std={np.std(grad_v_metrics['l1']):.6f}")
-            print(f"RMSE: mean={np.mean(grad_v_metrics['rmse']):.6f}, std={np.std(grad_v_metrics['rmse']):.6f}")
-
-if __name__ == "__main__":
-    main()

csrc/attn/test/test_fa2.py

@@ -1,136 +0,0 @@
-import torch
-from tqdm import tqdm
-import matplotlib.pyplot as plt
-import numpy as np
-
-def pytorch_test(Q, K, V, dO):
-    q_ = Q.to(torch.float64).requires_grad_()
-    k_ = K.to(torch.float64).requires_grad_()
-    v_ = V.to(torch.float64).requires_grad_()
-    dO_ = dO.to(torch.float64)
-    
-    # manual pytorch implementation of scaled dot product attention
-    QK = torch.matmul(q_, k_.transpose(-2, -1))
-    QK /= (q_.size(-1) ** 0.5)
-    
-    # Causal mask removed since causal is always false
-    
-    QK = torch.nn.functional.softmax(QK, dim=-1)
-    output = torch.matmul(QK, v_)
-    
-    output.backward(dO_)
-    
-    q_grad = q_.grad
-    k_grad = k_.grad
-    v_grad = v_.grad
-    
-    return output, q_grad, k_grad, v_grad
-
-def fa2_test(Q, K, V, dO):
-    Q.requires_grad = True
-    K.requires_grad = True
-    V.requires_grad = True
-    output = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=False)
-    output.backward(dO)
-    
-    return output, Q.grad, K.grad, V.grad
-
-def generate_tensor(shape, mean, std, dtype, device):
-    tensor = torch.randn(shape, dtype=dtype, device=device)
-
-    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
-    scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
-    
-    return scaled_tensor.contiguous()
-
-def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all', test_mode='forward_backward'):
-    results = {
-        'FA2 vs PT': {'sum_diff': 0, 'sum_abs': 0, 'max_diff': 0},
-    }
-
-    for _ in range(num_iterations):
-        torch.manual_seed(0)
-        
-        Q  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
-        K  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
-        V  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
-        dO = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
-        
-        pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, dO)
-        fa2_o, fa2_qg, fa2_kg, fa2_vg = fa2_test(Q, K, V, dO)
-        
-        if test_mode == 'forward_only':
-            tensors_fa2_pt = [(pt_o, fa2_o)]
-        else:  # 'forward_backward'
-            if error_mode == 'output':
-                tensors_fa2_pt = [(pt_o, fa2_o)]
-            elif error_mode == 'backward':
-                tensors_fa2_pt = [(pt_qg, fa2_qg),
-                               (pt_kg, fa2_kg),
-                               (pt_vg, fa2_vg)]
-            else:  # 'all'
-                tensors_fa2_pt = [(pt_o, fa2_o),
-                               (pt_qg, fa2_qg),
-                               (pt_kg, fa2_kg),
-                               (pt_vg, fa2_vg)]
-                
-        for pt, fa2 in tensors_fa2_pt:
-            diff = pt - fa2
-            abs_diff = torch.abs(diff)
-            results['FA2 vs PT']['sum_diff'] += torch.sum(abs_diff).item()
-            results['FA2 vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
-            results['FA2 vs PT']['max_diff'] = max(results['FA2 vs PT']['max_diff'], torch.max(abs_diff).item())
-                
-        torch.cuda.empty_cache()
-
-    # Calculate total elements based on test mode and error mode
-    if test_mode == 'forward_only':
-        total_elements = b * h * n * d * num_iterations
-    else:  # 'forward_backward'
-        total_elements = b * h * n * d * num_iterations * (1 if error_mode == 'output' else 3 if error_mode == 'backward' else 4)
-    
-    for name, data in results.items():
-        avg_diff = data['sum_diff'] / total_elements
-        max_diff = data['max_diff']
-        results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff}
-
-    return results
-
-def generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward'):
-    seq_lengths = [768 * (2**i) for i in range(1)]
-
-    print(f"\n{'='*80}")
-    print(f"ATTENTION ERROR COMPARISON TABLE (b={b}, h={h}, d={d}, mean={mean}, std={std})")
-    print(f"Mode: {error_mode}, Test: {test_mode}")
-    print(f"{'='*80}")
-    
-    # Print header
-    print(f"{'Seq Length':<12} | {'FA2 vs PT Avg':<15} | {'FA2 vs PT Max':<15}")
-    print(f"{'-'*12} | {'-'*15} | {'-'*15}")
-    
-    for n in seq_lengths:
-        results = check_correctness(b, h, n, d, mean, std, error_mode=error_mode, test_mode=test_mode)
-        
-        fa2_pt_avg = results['FA2 vs PT']['avg_diff']
-        fa2_pt_max = results['FA2 vs PT']['max_diff']
-        
-        # Print row
-        print(f"{n:<12} | {fa2_pt_avg:<15.6e} | {fa2_pt_max:<15.6e}")
-    
-    print(f"{'='*80}\n")
-
-# fix random seed
-torch.manual_seed(0)
-
-# Example usage
-b, h, d = 2, 2, 64
-mean = 1e-1
-std = 10
-
-# Test forward only
-generate_error_tables(b, h, d, mean, std, error_mode='output', test_mode='forward_only')
-
-# Test forward and backward
-generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward')
-
-print("Attention error comparison completed.")

csrc/attn/test/test_full_attn.py

@@ -1,175 +0,0 @@
-import torch
-from flash_attn_interface import flash_attn_func
-from st_attn import mha_forward, mha_backward
-import random
-from tqdm import tqdm
-import matplotlib.pyplot as plt
-import numpy as np
-
-def pytorch_test(Q, K, V, dO):
-    q_ = Q.to(torch.float64).requires_grad_()
-    k_ = K.to(torch.float64).requires_grad_()
-    v_ = V.to(torch.float64).requires_grad_()
-    dO_ = dO.to(torch.float64)
-    
-    # manual pytorch implementation of scaled dot product attention
-    QK = torch.matmul(q_, k_.transpose(-2, -1))
-    QK /= (q_.size(-1) ** 0.5)
-    
-    # Causal mask removed since causal is always false
-    
-    QK = torch.nn.functional.softmax(QK, dim=-1)
-    output = torch.matmul(QK, v_)
-    
-    output.backward(dO_)
-    
-    q_grad = q_.grad
-    k_grad = k_.grad
-    v_grad = v_.grad
-    
-    return output, q_grad, k_grad, v_grad
-
-def fa2_test(Q, K, V, dO):
-    Q.requires_grad = True
-    K.requires_grad = True
-    V.requires_grad = True
-    output = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=False)
-    output.backward(dO)
-    
-    return output, Q.grad, K.grad, V.grad
-
-
-def mha_kernel_test(Q, K, V, dO, mode): 
-    Q.requires_grad = True
-    K.requires_grad = True
-    V.requires_grad = True
-    
-    o, l_vec = mha_forward(Q, K, V)
-    
-    if mode == 'forward_only':
-        return o, None, None, None
-    else:  # 'forward_backward'
-        qg, kg, vg = mha_backward(Q, K, V, o, l_vec, dO)
-        return o, qg, kg, vg
-
-def generate_tensor(shape, mean, std, dtype, device):
-    tensor = torch.randn(shape, dtype=dtype, device=device)
-
-    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
-    scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
-    
-    return scaled_tensor.contiguous()
-
-def check_correctness(b, h, n, d, mean, std, num_iterations=100, error_mode='all', test_mode='forward_backward'):
-    results = {
-        'MHA vs PT': {'sum_diff': 0, 'sum_abs': 0, 'max_diff': 0},
-        'FA2 vs PT': {'sum_diff': 0, 'sum_abs': 0, 'max_diff': 0},
-    }
-
-    for _ in range(num_iterations):
-        torch.manual_seed(0)
-        
-        Q  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
-        K  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
-        V  = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
-        dO = generate_tensor((b, h, n, d), mean, std, torch.bfloat16, 'cuda')
-        
-        pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, dO)
-        fa2_o, fa2_qg, fa2_kg, fa2_vg = fa2_test(Q, K, V, dO)
-        
-        if test_mode == 'forward_only':
-            mha_o, _, _, _ = mha_kernel_test(Q, K, V, dO, 'forward_only')
-            tensors_mha_pt = [(pt_o, mha_o)]
-            tensors_fa2_pt = [(pt_o, fa2_o)]
-        else:  # 'forward_backward'
-            mha_o, mha_qg, mha_kg, mha_vg = mha_kernel_test(Q, K, V, dO, 'forward_backward')
-            
-            if error_mode == 'output':
-                tensors_mha_pt = [(pt_o, mha_o)]
-                tensors_fa2_pt = [(pt_o, fa2_o)]
-            elif error_mode == 'backward':
-                tensors_mha_pt = [(pt_qg, mha_qg),
-                               (pt_kg, mha_kg),
-                               (pt_vg, mha_vg)]
-                tensors_fa2_pt = [(pt_qg, fa2_qg),
-                               (pt_kg, fa2_kg),
-                               (pt_vg, fa2_vg)]
-            else:  # 'all'
-                tensors_mha_pt = [(pt_o, mha_o),
-                               (pt_qg, mha_qg),
-                               (pt_kg, mha_kg),
-                               (pt_vg, mha_vg)]
-                tensors_fa2_pt = [(pt_o, fa2_o),
-                               (pt_qg, fa2_qg),
-                               (pt_kg, fa2_kg),
-                               (pt_vg, fa2_vg)]
-                
-        for pt, mha in tensors_mha_pt:
-            diff = pt - mha
-            abs_diff = torch.abs(diff)
-            results['MHA vs PT']['sum_diff'] += torch.sum(abs_diff).item()
-            results['MHA vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
-            results['MHA vs PT']['max_diff'] = max(results['MHA vs PT']['max_diff'], torch.max(abs_diff).item())
-                
-        for pt, fa2 in tensors_fa2_pt:
-            diff = pt - fa2
-            abs_diff = torch.abs(diff)
-            results['FA2 vs PT']['sum_diff'] += torch.sum(abs_diff).item()
-            results['FA2 vs PT']['sum_abs'] += torch.sum(torch.abs(pt)).item()
-            results['FA2 vs PT']['max_diff'] = max(results['FA2 vs PT']['max_diff'], torch.max(abs_diff).item())
-                
-        torch.cuda.empty_cache()
-
-    # Calculate total elements based on test mode and error mode
-    if test_mode == 'forward_only':
-        total_elements = b * h * n * d * num_iterations
-    else:  # 'forward_backward'
-        total_elements = b * h * n * d * num_iterations * (1 if error_mode == 'output' else 3 if error_mode == 'backward' else 4)
-    
-    for name, data in results.items():
-        avg_diff = data['sum_diff'] / total_elements
-        max_diff = data['max_diff']
-        results[name] = {'avg_diff': avg_diff, 'max_diff': max_diff}
-
-    return results
-
-def generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward'):
-    seq_lengths = [768 * (2**i) for i in range(1)]
-
-    print(f"\n{'='*80}")
-    print(f"MHA ERROR COMPARISON TABLE (b={b}, h={h}, d={d}, mean={mean}, std={std})")
-    print(f"Mode: {error_mode}, Test: {test_mode}")
-    print(f"{'='*80}")
-    
-    # Print header
-    print(f"{'Seq Length':<12} | {'MHA vs PT Avg':<15} | {'MHA vs PT Max':<15} | {'FA2 vs PT Avg':<15} | {'FA2 vs PT Max':<15}")
-    print(f"{'-'*12} | {'-'*15} | {'-'*15} | {'-'*15} | {'-'*15}")
-    
-    for n in seq_lengths:
-        results = check_correctness(b, h, n, d, mean, std, error_mode=error_mode, test_mode=test_mode)
-        
-        mha_pt_avg = results['MHA vs PT']['avg_diff']
-        mha_pt_max = results['MHA vs PT']['max_diff']
-        fa2_pt_avg = results['FA2 vs PT']['avg_diff']
-        fa2_pt_max = results['FA2 vs PT']['max_diff']
-        
-        # Print row
-        print(f"{n:<12} | {mha_pt_avg:<15.6e} | {mha_pt_max:<15.6e} | {fa2_pt_avg:<15.6e} | {fa2_pt_max:<15.6e}")
-    
-    print(f"{'='*80}\n")
-
-# fix random seed
-torch.manual_seed(0)
-
-# Example usage
-b, h, d = 2, 2, 64
-mean = 1e-1
-std = 10
-
-# Test forward only
-generate_error_tables(b, h, d, mean, std, error_mode='output', test_mode='forward_only')
-
-# Test forward and backward
-generate_error_tables(b, h, d, mean, std, error_mode='all', test_mode='forward_backward')
-
-print("MHA attention error comparison completed.")

csrc/attn/vsa.cpp

@@ -1,27 +0,0 @@
-#include <torch/extension.h>
-#include <ATen/ATen.h>
-
-#include <vector>
-#include <cuda_fp16.h>
-#include <cuda_bf16.h>
-
-#include <cuda_runtime.h>
-
-#ifdef TK_COMPILE_BLOCK_SPARSE
-extern std::vector<torch::Tensor> block_sparse_attention_forward(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v,  torch::Tensor q2k_block_sparse_index, torch::Tensor q2k_block_sparse_num
-); 
-extern std::vector<torch::Tensor> block_sparse_attention_backward(
-    torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, torch::Tensor l_vec, torch::Tensor og, torch::Tensor k2q_block_sparse_index, torch::Tensor k2q_block_sparse_num
-);
-#endif
-
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.doc() = "Video Sparse Attention Kernels"; // optional module docstring
-
-#ifdef TK_COMPILE_BLOCK_SPARSE
-    m.def("block_sparse_fwd",  torch::wrap_pybind_function(block_sparse_attention_forward), "block sparse attention");
-    m.def("block_sparse_bwd",  torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward");
-#endif
-}

csrc/attn/vsa/__init__.py

@@ -1,470 +0,0 @@
-import math
-
-import torch
-from torch.utils.checkpoint import detach_variable
-from typing import Tuple
-try:
-    from vsa_cuda import block_sparse_fwd, block_sparse_bwd
-except ImportError:
-    block_sparse_fwd = None
-    block_sparse_bwd = None
-
-
-BLOCK_M = 64
-BLOCK_N = 64
-
-def video_sparse_attn(q, k, v, topk, block_size, compress_attn_weight=None):
-    """
-    q: [batch_size, num_heads, seq_len, head_dim]
-    k: [batch_size, num_heads, seq_len, head_dim]
-    v: [batch_size, num_heads, seq_len, head_dim]
-    topk: int
-    block_size: int or tuple of 3 ints
-    video_shape: tuple of (T, H, W)
-    compress_attn_weight: [batch_size, num_heads, seq_len, head_dim]
-    select_attn_weight: [batch_size, num_heads, seq_len, head_dim]
-    
-    V1 of sparse attention. Include compress attn and sparse attn branch, use average pooling to compress. 
-    Assume q, k, v is flattened in this way: [batch_size, num_heads, T//block_size[0], H//block_size[1], W//block_size[2], block_size[0], block_size[1], block_size[2]]
-    """
-
-    if isinstance(block_size, int):
-        block_size = (block_size, block_size, block_size)
-
-    block_elements = block_size[0] * block_size[1] * block_size[2]
-    assert block_elements % 64 == 0 and block_elements >= 64
-    assert q.shape[2] % block_elements == 0
-    batch_size, num_heads, seq_len, head_dim = q.shape
-    # compress attn
-    q_compress = q.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).mean(dim=3)
-    k_compress = k.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).mean(dim=3)
-    v_compress = v.view(batch_size, num_heads, seq_len // block_elements,
-                        block_elements, head_dim).mean(dim=3)
-
-    output_compress, block_attn_score = torch_attention(q_compress, k_compress,
-                                                        v_compress)
-
-    output_compress = output_compress.view(batch_size, num_heads,
-                                           seq_len // block_elements, 1,
-                                           head_dim)
-    output_compress = output_compress.repeat(1, 1, 1, block_elements,
-                                             1).view(batch_size, num_heads,
-                                                     seq_len, head_dim)
-
-    q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num = generate_topk_block_sparse_pattern(
-        block_attn_score, topk)
-
-    output_select = block_sparse_attn(q, k, v, q2k_block_sparse_index,
-                                      q2k_block_sparse_num,
-                                      k2q_block_sparse_index,
-                                      k2q_block_sparse_num)
-
-    if compress_attn_weight is not None:
-        final_output = output_compress * compress_attn_weight + output_select
-    else:
-        final_output = output_compress + output_select
-    return final_output
-
-def torch_attention(q, k, v) -> Tuple[torch.Tensor, torch.Tensor]:
-    QK = torch.matmul(q, k.transpose(-2, -1))
-    QK /= (q.size(-1)**0.5)
-
-    # Causal mask removed since causal is always false
-
-    QK = torch.nn.functional.softmax(QK, dim=-1)
-    output = torch.matmul(QK, v)
-    return output, QK
-
-def generate_topk_block_sparse_pattern(block_attn_score: torch.Tensor,
-                                       topk: int):
-    """
-    Generate a block sparse pattern where each q block attends to exactly topk kv blocks,
-    based on the provided attention scores.
-    
-    Args:
-        block_attn_score: [bs, h, num_q_blocks, num_kv_blocks]
-            Attention scores between query and key blocks
-        topk: int
-            Number of kv blocks each q block attends to
-        
-    Returns:
-        q2k_block_sparse_index: [bs, h, num_q_blocks, topk]
-            Contains the indices of kv blocks that each q block attends to.
-        q2k_block_sparse_num: [bs, h, num_q_blocks]
-            Contains the number of kv blocks that each q block attends to (all equal to topk).
-        k2q_block_sparse_index: [bs, h, num_kv_blocks, max_q_per_kv]
-            Contains the indices of q blocks that attend to each kv block.
-        k2q_block_sparse_num: [bs, h, num_kv_blocks]
-            Contains the number of q blocks that attend to each kv block.
-    """
-    device = block_attn_score.device
-    # Extract dimensions from block_attn_score
-    bs, h, num_q_blocks, num_kv_blocks = block_attn_score.shape
-
-    sorted_result = torch.sort(block_attn_score, dim=-1, descending=True)
-
-    sorted_indice = sorted_result.indices
-
-    q2k_block_sparse_index, _ = torch.sort(sorted_indice[:, :, :, :topk],
-                                           dim=-1)
-    q2k_block_sparse_index = q2k_block_sparse_index.to(dtype=torch.int32)
-    q2k_block_sparse_num = torch.full((bs, h, num_q_blocks),
-                                      topk,
-                                      device=device,
-                                      dtype=torch.int32)
-
-    block_map = topk_index_to_map(q2k_block_sparse_index,
-                                  num_kv_blocks,
-                                  transpose_map=True)
-    k2q_block_sparse_index, k2q_block_sparse_num = map_to_index(
-        block_map.transpose(2, 3))
-
-    return q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
-
-@torch._dynamo.disable
-def block_sparse_attn(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
-    """
-    Differentiable block sparse attention function.
-    
-    Args:
-        q: Query tensor [batch_size, num_heads, seq_len_q, head_dim]
-        k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim]
-        v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim]
-        q2k_block_sparse_index: Indices for query-to-key sparse blocks
-        q2k_block_sparse_num: Number of sparse blocks for each query block
-        k2q_block_sparse_index: Indices for key-to-query sparse blocks (for backward pass)
-        k2q_block_sparse_num: Number of sparse blocks for each key block (for backward pass)
-    
-    Returns:
-        output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim]
-    """
-    return BlockSparseAttentionFunction.apply(
-        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num
-    )
-
-def block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num):
-    """
-    block_sparse_mask: [bs, h, num_q_blocks, num_kv_blocks]. 
-        [*, *, i, j] = 1 means the i-th q block should attend to the j-th kv block.
-    """
-    # assert all elements in q2k_block_sparse_num can be devisible by 2
-    o, lse = block_sparse_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
-    return o, lse
-
-def block_sparse_attention_backward(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num):
-    grad_output = grad_output.contiguous()
-    grad_q, grad_k, grad_v = block_sparse_bwd(q, k, v, o, l_vec, grad_output, k2q_block_sparse_index, k2q_block_sparse_num)
-    return grad_q, grad_k, grad_v
-
-## pytorch sdpa version of block sparse ##
-import triton
-import triton.language as tl
-
-@triton.jit
-def index_to_mask_kernel(
-    q2k_block_sparse_index_ptr,
-    q2k_block_sparse_num_ptr,
-    mask_ptr,
-    batch_size: tl.constexpr,
-    num_heads: tl.constexpr,
-    num_q_blocks: tl.constexpr,
-    num_k_blocks: tl.constexpr,
-    max_kv_blocks: tl.constexpr,
-    BLOCK_Q: tl.constexpr,
-    BLOCK_K: tl.constexpr,
-):
-    bh, q, id = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64), tl.program_id(2).to(tl.int64)
-    b = bh // num_heads
-    h = bh % num_heads
-
-    num_valid_blocks = tl.load(q2k_block_sparse_num_ptr + b * num_heads * num_q_blocks + h * num_q_blocks + q)
-
-    if num_valid_blocks <= id:
-        return
-    k = tl.load(q2k_block_sparse_index_ptr + b * num_heads * num_q_blocks * max_kv_blocks + h * num_q_blocks * max_kv_blocks + q * max_kv_blocks + id)
-
-    full_mask = (tl.arange(0, BLOCK_Q)[:, None] < BLOCK_Q) & (tl.arange(0, BLOCK_K)[None, :] < BLOCK_K)
-
-    q_lengths = num_q_blocks * BLOCK_Q
-    k_lengths = num_k_blocks * BLOCK_K
-    mask_ptr_base = mask_ptr + b * num_heads * q_lengths * k_lengths + h * q_lengths * k_lengths + q * BLOCK_Q * k_lengths + k * BLOCK_K
-
-    tl.store(mask_ptr_base + tl.arange(0, BLOCK_Q)[:, None] * k_lengths + tl.arange(0, BLOCK_K)[None, :], full_mask)
-
-def index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, BLOCK_Q, BLOCK_K, num_k_blocks):
-    """
-    Convert block sparse indices to a mask.
-    
-    Args:
-        q2k_block_sparse_index: Indices for query-to-key sparse blocks
-        q2k_block_sparse_num: Number of sparse blocks for each query block
-    
-    Returns:
-        mask: Block sparse mask tensor
-    """
-    batch_size, num_heads, num_q_blocks, max_kv_blocks = q2k_block_sparse_index.shape
-    assert q2k_block_sparse_num.shape == (batch_size, num_heads, num_q_blocks)
-
-    mask = torch.zeros((batch_size, num_heads, num_q_blocks * BLOCK_Q, num_k_blocks * BLOCK_K), dtype=torch.bool, device=q2k_block_sparse_index.device)
-
-    grid = (batch_size * num_heads, num_q_blocks, max_kv_blocks)
-    index_to_mask_kernel[grid](
-        q2k_block_sparse_index,
-        q2k_block_sparse_num,
-        mask,
-        batch_size,
-        num_heads,
-        num_q_blocks,
-        num_k_blocks,
-        max_kv_blocks,
-        BLOCK_Q=BLOCK_Q,
-        BLOCK_K=BLOCK_K,
-    )
-
-    return mask
-
-@triton.jit
-def topk_index_to_map_kernel(
-    map_ptr,
-    index_ptr,
-    map_bs_stride,
-    map_h_stride,
-    map_q_stride,
-    map_kv_stride,
-    index_bs_stride,
-    index_h_stride,
-    index_q_stride,
-    index_kv_stride,
-    topk: tl.constexpr,
-):
-    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
-    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
-
-    for i in tl.static_range(topk):
-        index = tl.load(index_ptr_base + i * index_kv_stride)
-        tl.store(map_ptr_base + index * map_kv_stride, 1.0)
-
-@triton.jit
-def map_to_index_kernel(
-    map_ptr,
-    index_ptr,
-    index_num_ptr,
-    map_bs_stride,
-    map_h_stride,
-    map_q_stride,
-    map_kv_stride,
-    index_bs_stride,
-    index_h_stride,
-    index_q_stride,
-    index_kv_stride,
-    index_num_bs_stride,
-    index_num_h_stride,
-    index_num_q_stride,
-    num_kv_blocks: tl.constexpr,
-):
-    b, h, q = tl.program_id(0), tl.program_id(1), tl.program_id(2)
-    index_ptr_base = index_ptr + b * index_bs_stride + h * index_h_stride + q * index_q_stride
-    map_ptr_base = map_ptr + b * map_bs_stride + h * map_h_stride + q * map_q_stride
-
-    num = 0
-    for i in tl.static_range(num_kv_blocks):
-        map_entry = tl.load(map_ptr_base + i * map_kv_stride)
-        if map_entry:
-            tl.store(index_ptr_base + num * index_kv_stride, i)
-            num += 1
-
-    tl.store(
-        index_num_ptr + b * index_num_bs_stride + h * index_num_h_stride +
-        q * index_num_q_stride, num)
-
-def topk_index_to_map(index: torch.Tensor,
-                      num_kv_blocks: int,
-                      transpose_map: bool = False):
-    """
-    Convert topk indices to a map.
-    
-    Args:
-        index: [bs, h, num_q_blocks, topk]
-            The topk indices tensor.
-        num_kv_blocks: int
-            The number of key-value blocks in the block_map returned
-        transpose_map: bool
-            If True, the block_map will be transposed on the final two dimensions.
-    
-    Returns:
-        block_map: [bs, h, num_q_blocks, num_kv_blocks]
-            A binary map where 1 indicates that the q block attends to the kv block.
-    """
-    bs, h, num_q_blocks, topk = index.shape
-
-    if transpose_map is False:
-        block_map = torch.zeros((bs, h, num_q_blocks, num_kv_blocks),
-                                dtype=torch.bool,
-                                device=index.device)
-    else:
-        block_map = torch.zeros((bs, h, num_kv_blocks, num_q_blocks),
-                                dtype=torch.bool,
-                                device=index.device)
-        block_map = block_map.transpose(2, 3)
-
-    grid = (bs, h, num_q_blocks)
-    topk_index_to_map_kernel[grid](
-        block_map,
-        index,
-        block_map.stride(0),
-        block_map.stride(1),
-        block_map.stride(2),
-        block_map.stride(3),
-        index.stride(0),
-        index.stride(1),
-        index.stride(2),
-        index.stride(3),
-        topk=topk,
-    )
-
-    return block_map
-
-def map_to_index(block_map: torch.Tensor):
-    """
-    Convert a block map to indices and counts.
-    
-    Args:
-        block_map: [bs, h, num_q_blocks, num_kv_blocks]
-            The block map tensor.
-    
-    Returns:
-        index: [bs, h, num_q_blocks, num_kv_blocks]
-            The indices of the blocks.
-        index_num: [bs, h, num_q_blocks]
-            The number of blocks for each q block.
-    """
-    bs, h, num_q_blocks, num_kv_blocks = block_map.shape
-
-    index = torch.full((block_map.shape),
-                       -1,
-                       dtype=torch.int32,
-                       device=block_map.device)
-    index_num = torch.empty((bs, h, num_q_blocks),
-                            dtype=torch.int32,
-                            device=block_map.device)
-
-    grid = (bs, h, num_q_blocks)
-    map_to_index_kernel[grid](
-        block_map,
-        index,
-        index_num,
-        block_map.stride(0),
-        block_map.stride(1),
-        block_map.stride(2),
-        block_map.stride(3),
-        index.stride(0),
-        index.stride(1),
-        index.stride(2),
-        index.stride(3),
-        index_num.stride(0),
-        index_num.stride(1),
-        index_num.stride(2),
-        num_kv_blocks=num_kv_blocks,
-    )
-
-    return index, index_num
-
-class BlockSparseAttentionFunction(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, k2q_block_sparse_index, k2q_block_sparse_num):
-        o, lse = block_sparse_attention_fwd(q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)
-        ctx.save_for_backward(q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num)
-        return o
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        q, k, v, o, lse, k2q_block_sparse_index, k2q_block_sparse_num = ctx.saved_tensors
-        grad_q, grad_k, grad_v = block_sparse_attention_backward(
-            q, k, v, o, lse, grad_output, k2q_block_sparse_index, k2q_block_sparse_num
-        )
-        return grad_q, grad_k, grad_v, None, None, None, None
-
-
-class DummyOperator(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x):
-        return x
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        return grad_output
-
-class CheckpointSDPA(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, obj, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k):
-        """Forward pass."""
-        with torch.no_grad():
-            mask = index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k, k.shape[2] // block_k)
-            outputs = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
-        ctx.save_for_backward(*detach_variable((q, k, v, q2k_block_sparse_index, q2k_block_sparse_num)))
-        ctx.block_q = block_q
-        ctx.block_k = block_k
-        # the obj is passed in, then it can access the saved input
-        # tensors later for recomputation
-        obj.ctx = ctx
-        return outputs
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        """Backward pass."""
-        inputs = ctx.saved_tensors
-        output = ctx.output
-        torch.autograd.backward(output, grad_output)
-        ctx.output = None
-        grads = tuple(inp.grad for inp in inputs)
-        return (None, ) + grads + (None, None)
-
-
-class BlockSparseAttnTorch:
-    def __init__(self):
-        self.ctx = None
-    
-    def recompute_mask(self, _):
-        recomputed_mask = index_to_mask(self.q2k_block_sparse_index, self.q2k_block_sparse_num, self.block_q, self.block_k, self.num_kv_blocks)
-        mask_size = recomputed_mask.untyped_storage().size()
-        self.mask.untyped_storage().resize_(mask_size)
-        self.mask.untyped_storage().copy_(recomputed_mask.untyped_storage()) 
-    
-    def recompute(self, _):
-        q, k, v, q2k_block_sparse_index, q2k_block_sparse_num = self.ctx.saved_tensors
-        block_q = self.ctx.block_q
-        block_k = self.ctx.block_k
-        mask = index_to_mask(q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k, k.shape[2] // block_k)
-        with torch.enable_grad():
-            output = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask)
-        self.ctx.output = output
-        self.ctx = None
-
-    @torch._dynamo.disable
-    def forward(self, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k):
-        """
-        Differentiable block sparse attention function using PyTorch.
-        
-        Args:
-            q: Query tensor [batch_size, num_heads, seq_len_q, head_dim]
-            k: Key tensor [batch_size, num_heads, seq_len_kv, head_dim]
-            v: Value tensor [batch_size, num_heads, seq_len_kv, head_dim]
-            q2k_block_sparse_index: Indices for query-to-key sparse blocks
-            q2k_block_sparse_num: Number of sparse blocks for each query block
-            block_q: Block size for query
-            block_k: Block size for key-value
-
-        Returns:
-            output: Attention output tensor [batch_size, num_heads, seq_len_q, head_dim]
-        """
-
-        output = CheckpointSDPA.apply(
-            self, q, k, v, q2k_block_sparse_index, q2k_block_sparse_num, block_q, block_k
-        )
-
-        o = DummyOperator.apply(output)
-        o.register_hook(self.recompute)
-        return o

csrc/sliding_tile_attention/README.md

@@ -6,6 +6,7 @@
 ## Installation
 We test our code on Pytorch 2.5.0 and CUDA>=12.4. Currently we only have implementation on H100.
 First, install C++20 for ThunderKittens:
+
 ```bash
 sudo apt update
 sudo apt install gcc-11 g++-11
@@ -15,27 +16,17 @@ sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave
 sudo apt update
 sudo apt install clang-11
 ```
-
-## Environment Setup
-First, set up your CUDA environment:
+Install STA:
 ```bash
 export CUDA_HOME=/usr/local/cuda-12.4
 export PATH=${CUDA_HOME}/bin:${PATH} 
 export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
 git submodule update --init --recursive
-```
-
-## Install Sliding Tile Attention (STA)
-```bash
-python setup_sta.py install
-```
-
-## Install Video Sparse Attention (VSA)
-```bash
-python setup_vsa.py install
+python setup.py install
 ```
 
 ## Usage
+
 ```python
 from st_attn import sliding_tile_attention
 # assuming video size (T, H, W) = (30, 48, 80), text tokens = 256 with padding. 

csrc/sliding_tile_attention/config.py

@@ -1,6 +1,6 @@
 ### ADD TO THIS TO REGISTER NEW KERNELS
 sources = {
-    'st_attn': {
+    'attn': {
         'source_files': {
             'h100': 'st_attn/st_attn_h100.cu'  # define these source files for each GPU target desired.
         }
@@ -9,7 +9,7 @@ sources = {
 
 ### WHICH KERNELS DO WE WANT TO BUILD?
 # (oftentimes during development work you don't need to redefine them all.)
-kernels = ['st_attn']
+kernels = ['attn']
 
 ### WHICH GPU TARGET DO WE WANT TO BUILD FOR?
 target = 'h100'

csrc/sliding_tile_attention/setup.py

@@ -1,7 +1,7 @@
 import os
 import subprocess
 
-from csrc.attn.config_sta import kernels, sources, target
+from config import kernels, sources, target
 from setuptools import find_packages, setup
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 

csrc/sliding_tile_attention/st_attn.cpp

@@ -7,7 +7,8 @@
 
 #include <cuda_runtime.h>
 
-#ifdef TK_COMPILE_ST_ATTN
+
+#ifdef TK_COMPILE_ATTN
 extern torch::Tensor sta_forward(
     torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, int kernel_t_size, int kernel_w_size, int kernel_h_size, int text_length, bool process_text, bool has_text, int kernel_aspect_ratio_flag
 ); 
@@ -16,8 +17,8 @@ extern torch::Tensor sta_forward(
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.doc() = "Sliding Block Attention Kernels"; // optional module docstring
-    
-#ifdef TK_COMPILE_ST_ATTN
+
+#ifdef TK_COMPILE_ATTN
     m.def("sta_fwd",  torch::wrap_pybind_function(sta_forward), "sliding tile attention, assuming tile size is (6,8,8)");
 #endif
-}
+}

csrc/sliding_tile_attention/st_attn/__init__.py

@@ -1,22 +1,19 @@
 import math
 
 import torch
-from torch.utils.checkpoint import detach_variable
-try:
-    from st_attn_cuda import sta_fwd
-except ImportError:
-    sta_fwd = None
+from st_attn_cuda import sta_fwd
 
-def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True, dit_seq_shape='30x48x80'):
+
+def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_text=True, img_latent_shape='30*48*80'):
     seq_length = q_all.shape[2]
-    dit_seq_shape_mapping = {
+    img_latent_shape_mapping = {
         '30x48x80':1,
         '36x48x48':2,
         '18x48x80':3, 
     }
     if has_text:
         assert q_all.shape[
-            2] >= 115200 and q_all.shape[2] <= 115456, f"Unsupported {dit_seq_shape}, current shape is {q_all.shape}, only support '30x48x80' for HunyuanVideo"
+            2] >= 115200, "STA currently only supports video with latent size (30, 48, 80), which is 117 frames x 768 x 1280 pixels"
         assert q_all.shape[1] == len(window_size), "Number of heads must match the number of window sizes"
         target_size = math.ceil(seq_length / 384) * 384
         pad_size = target_size - seq_length
@@ -25,14 +22,14 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te
             k_all = torch.cat([k_all, k_all[:, :, -pad_size:]], dim=2)
             v_all = torch.cat([v_all, v_all[:, :, -pad_size:]], dim=2)
     else:
-        if dit_seq_shape == '36x48x48': # Stepvideo 204x768x68
+        if img_latent_shape == '36x48x48': # Stepvideo 204x768x68
             assert q_all.shape[2] == 82944
-        elif dit_seq_shape == '18x48x80': # Wan 69x768x1280
+        elif img_latent_shape == '18x48x80': # Wan 69x768x1280
             assert q_all.shape[2] == 69120
         else:
-            raise ValueError(f"Unsupported {dit_seq_shape}, current shape is {q_all.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan")
+            raise ValueError(f"Unsupported {img_latent_shape}, current shape is {q_all.shape}, only support '36x48x48' for Stepvideo and '18x48x80' for Wan")
 
-    kernel_aspect_ratio_flag = dit_seq_shape_mapping[dit_seq_shape]
+    kernel_aspect_ratio_flag = img_latent_shape_mapping[img_latent_shape]
     hidden_states = torch.empty_like(q_all)
     # This for loop is ugly. but it is actually quite efficient. The sequence dimension alone can already oversubscribe SMs
     for head_index, (t_kernel, h_kernel, w_kernel) in enumerate(window_size):
@@ -46,4 +43,4 @@ def sliding_tile_attention(q_all, k_all, v_all, window_size, text_length, has_te
             _ = sta_fwd(q_head, k_head, v_head, o_head, t_kernel, h_kernel, w_kernel, text_length, False, has_text, kernel_aspect_ratio_flag)
     if has_text:
         _ = sta_fwd(q_all, k_all, v_all, hidden_states, 3, 3, 3, text_length, True, True, kernel_aspect_ratio_flag)
-    return hidden_states[:, :, :seq_length]
+    return hidden_states[:, :, :seq_length]

csrc/sliding_tile_attention/st_attn/st_attn_h100.cu

@@ -829,4 +829,3 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o,
     return o;
     cudaDeviceSynchronize();
 }
-

csrc/sliding_tile_attention/test/bench.py

@@ -45,12 +45,12 @@ def benchmark_attention(configurations):
 
         # Warmup for forward pass
         for _ in range(10):
-            o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, '18x48x80')
+            o = sliding_tile_attention(q, k, v, [[6, 6, 6]] * 24, 0, False)
 
         # Time the forward pass
         for i in range(10):
             start_events_fwd[i].record()
-            o = sliding_tile_attention(q, k, v, [[3, 6, 10]] * 24, 0, False, '18x48x80')
+            o = sliding_tile_attention(q, k, v, [[6, 6, 6]] * 24, 0, False)
             end_events_fwd[i].record()
 
         torch.cuda.synchronize()
@@ -124,7 +124,7 @@ def plot_results(results):
 
 # Example list of configurations to test
 configurations = [
-    (2, 24, 69120, 128, False),
+    (2, 24, 82944, 128, False),
     # (16, 16, 768*16,  128, False),
     # (16, 16, 768*2,  128, False),
     # (16, 16, 768*4,  128, False),

csrc/sliding_tile_attention/test/test_sta.py

@@ -2,28 +2,27 @@ import torch
 from flex_sta_ref import get_sliding_tile_attention_mask
 from st_attn import sliding_tile_attention
 from torch.nn.attention.flex_attention import flex_attention
-# from flash_attn_interface import flash_attn_func
 from tqdm import tqdm
 
 flex_attention = torch.compile(flex_attention, dynamic=False)
 
 
 def flex_test(Q, K, V, kernel_size):
-    mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), (18, 48, 80), 0, 'cuda', 0)
+    mask = get_sliding_tile_attention_mask(kernel_size, (6, 8, 8), (36, 48, 48), 39, 'cuda', 0)
     output = flex_attention(Q, K, V, block_mask=mask)
 
     return output
 
 
 def h100_fwd_kernel_test(Q, K, V, kernel_size):
-    o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 0, False, '18x48x80')
+    o = sliding_tile_attention(Q, K, V, [kernel_size] * 24, 39, False)
     return o
 
 
 def generate_tensor(shape, mean, std, dtype, device):
     tensor = torch.randn(shape, dtype=dtype, device=device)
 
-    magnitude = torch.norm(tensor, dim=-1, keepdim=True)
+    magnitude = torch.linalg.norm(tensor, dim=-1, keepdim=True)
     scaled_tensor = tensor * (torch.randn(magnitude.shape, dtype=dtype, device=device) * std + mean) / magnitude
 
     return scaled_tensor.contiguous()
@@ -37,7 +36,7 @@ def check_correctness(b, h, n, d, causal, mean, std, num_iterations=50, error_mo
             'max_diff': 0
         },
     }
-    kernel_size_ls = [(3, 3, 5), (3, 1, 10)]
+    kernel_size_ls = [(6, 1, 6), (6, 6, 1)]
     from tqdm import tqdm
     for kernel_size in tqdm(kernel_size_ls):
         for _ in range(num_iterations):
@@ -72,14 +71,25 @@ def check_correctness(b, h, n, d, causal, mean, std, num_iterations=50, error_mo
     return results
 
 
+def generate_error_graphs(b, h, d, causal, mean, std, error_mode='all'):
+    seq_lengths = [82944]
+
+    tk_avg_errors, tk_max_errors = [], []
+
+    for n in tqdm(seq_lengths, desc="Generating error data"):
+        results = check_correctness(b, h, n, d, causal, mean, std, error_mode=error_mode)
+
+        tk_avg_errors.append(results['TK vs FLEX']['avg_diff'])
+        tk_max_errors.append(results['TK vs FLEX']['max_diff'])
+
+
 # Example usage
 b, h, d = 2, 24, 128
-n = 69120  # Sequence length
 causal = False
 mean = 1e-1
 std = 10
 
-# Run correctness check directly
-results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output')
-print(f"Average difference: {results['TK vs FLEX']['avg_diff']}")
-print(f"Maximum difference: {results['TK vs FLEX']['max_diff']}")
+for mode in ['output']:
+    generate_error_graphs(b, h, d, causal, mean, std, error_mode=mode)
+
+print("Error graphs generated and saved for all modes.")

docker/Dockerfile.python3.12

@@ -1,9 +1,7 @@
-FROM nvidia/cuda:12.8.0-cudnn-devel-ubuntu22.04
+FROM nvidia/cuda:12.4.1-devel-ubuntu20.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
-SHELL ["/bin/bash", "-c"]
-
 WORKDIR /FastVideo
 
 RUN apt-get update && apt-get install -y --no-install-recommends \
@@ -11,25 +9,17 @@ 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/*
 
-# Set up C++20 compilers for ThunderKittens
-RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
+    bash Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda && \
+    rm Miniconda3-latest-Linux-x86_64.sh
 
-# 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
+ENV PATH=/opt/conda/bin:$PATH
 
-# Install uv and source its environment
-RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
-    echo 'source $HOME/.local/bin/env' >> /root/.bashrc
+RUN conda create --name fastvideo-dev python=3.12.9 -y
+
+SHELL ["/bin/bash", "-c"]
 
 # Copy just the pyproject.toml first to leverage Docker cache
 COPY pyproject.toml ./
@@ -37,36 +27,22 @@ COPY pyproject.toml ./
 # Create a dummy README to satisfy the installation
 RUN echo "# Placeholder" > README.md
 
-# Create and activate virtual environment with specific Python version and seed
-RUN source $HOME/.local/bin/env && \
-    uv venv --python 3.12 --seed /opt/venv && \
-    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
+RUN conda run -n fastvideo-dev pip install --no-cache-dir --upgrade pip && \
+    conda run -n fastvideo-dev pip install --no-cache-dir .[dev] && \
+    conda run -n fastvideo-dev pip install --no-cache-dir flash-attn==2.7.4.post1 --no-build-isolation && \
+    conda clean -afy
 
 COPY . .
 
-# 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
+RUN conda run -n fastvideo-dev pip install --no-cache-dir -e .[dev]
 
-# Install STA (Sliding Tile Attention)
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn && \
-    git submodule update --init --recursive && \
-    python setup_sta.py install
+# Remove authentication headers
+RUN git config --unset-all http.https://github.com/.extraheader || true
 
-# Install VSA 
-RUN source $HOME/.local/bin/env && \
-    source /opt/venv/bin/activate && \
-    cd csrc/attn && \
-    git submodule update --init --recursive && \
-    python setup_vsa.py install
+# Set up automatic conda environment activation for all shells
+RUN echo 'source /opt/conda/etc/profile.d/conda.sh' >> /root/.bashrc && \
+    echo 'conda activate fastvideo-dev' >> /root/.bashrc && \
+    # Ensure .bashrc is sourced for SSH login shells
+    echo 'if [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
 EXPOSE 22

docs/source/sliding_tile_attention/demo.md

@@ -1,7 +1,7 @@
 (sta-demo)=
 
 # 🔍 Demo
-This is is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
+There is a demo for 2D STA with window size (6,6) operating on a (10, 10) image.
 
 <div style="text-align: center;">
   <video controls width="800">
@@ -9,9 +9,3 @@ This is is a demo for 2D STA with window size (6,6) operating on a (10, 10) imag
     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
-```

docs/source/training/data_preprocess.md

@@ -7,40 +7,70 @@ To save GPU memory, we precompute text embeddings and VAE latents to eliminate t
 We provide a sample dataset to help you get started. Download the source media using the following command:
 
 ```bash
-python scripts/huggingface/download_hf.py --repo_id=FastVideo/mini_i2v_dataset --local_dir=FastVideo/mini_i2v_dataset --repo_type=dataset
+python scripts/huggingface/download_hf.py --repo_id=FastVideo/Image-Vid-Finetune-Src --local_dir=data/Image-Vid-Finetune-Src --repo_type=dataset
 ```
 
-The folder `crush-smol_raw/` contains raw videos and captions for testing preprocessing, while `crush-smol_preprocessed/` contains latents prepared for testing training.
-
 To preprocess the dataset for fine-tuning or distillation, run:
 
 ```
-bash scripts/preprocess/v1_preprocess_wan_data_t2v # for wan
+bash scripts/preprocess/preprocess_mochi_data.sh # for mochi
+bash scripts/preprocess/preprocess_hunyuan_data.sh # for hunyuan
 ```
 
+The preprocessed dataset will be stored in `Image-Vid-Finetune-Mochi` or `Image-Vid-Finetune-HunYuan` correspondingly.
+
 ## Process your own dataset
 
-If you wish to create your own dataset for finetuning or distillation, please refer `mini_i2v_dataset/crush-smol_raw/` to structure you video dataset in the following format:
+If you wish to create your own dataset for finetuning or distillation, please structure you video dataset in the following format:
 
 ```
-path_to_your_dataset_folder/
-├── videos/
-│   ├── 0.mp4
+path_to_dataset_folder/
+├── media/
+│   ├── 0.jpg
 │   ├── 1.mp4
-├── videos.txt
-└── prompt.txt
+│   ├── 2.jpg
+├── video2caption.json
+└── merge.txt
 ```
 
-To geranate the `videos2caption.json` and `merge.txt`, run
+Format the JSON file as a list, where each item represents a media source:
 
-``` python
-python scripts/dataset_preparation/prepare_json_file.py --data_folder mini_i2v_dataset/crush-smol_raw/ --output your_output_folder
-```
-
-Adjust the `DATA_MERGE_PATH` and `OUTPUT_DIR` in `scripts/preprocess/v1_preprocess_****.sh` accordingly and run:
+For image media,
 
 ```
-bash scripts/preprocess/v1_preprocess_****.sh
+{
+    "path": "0.jpg",
+    "cap": ["captions"]
+}
+```
+
+For video media,
+
+```
+{
+    "path": "1.mp4",
+    "resolution": {
+      "width": 848,
+      "height": 480
+    },
+    "fps": 30.0,
+    "duration": 6.033333333333333,
+    "cap": [
+      "caption"
+    ]
+  }
+```
+
+Use a txt file (merge.txt) to contain the source folder for media and the JSON file for meta information:
+
+```
+path_to_media_source_foder,path_to_json_file
+```
+
+Adjust the `DATA_MERGE_PATH` and `OUTPUT_DIR` in `scripts/preprocess/preprocess_****_data.sh` accordingly and run:
+
+```
+bash scripts/preprocess/preprocess_****_data.sh
 ```
 
 The preprocessed data will be put into the `OUTPUT_DIR` and the `videos2caption.json` can be used in finetune and distill scripts.

docs/source/training/finetune.md

@@ -16,13 +16,6 @@ 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:

examples/inference/sta_mask_search/inference_wan_sta.sh

@@ -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=1
+num_gpu=$(nvidia-smi --query-gpu=gpu_name --format=csv,noheader | wc -l)
 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_${i}.txt \
+        --prompt_path ./assets/prompt_extend_${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_${i}.txt \
+        --prompt_path ./assets/prompt_extend_${i}.txt \
         --output_path $output_path \
         --STA_mode $STA_mode \
         --skip_time_steps $skip_time_steps &

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.sh

@@ -1,88 +0,0 @@
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/"
-VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/"
-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 5000
-  --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
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS
-  --sp_size 8
-  --tp_size 8
-  --hsdp_replicate_dim 1
-  --hsdp_shard_dim 8
-)
-
-# Model arguments
-model_args=(
-  --model_path "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-  --pretrained_model_name_or_path "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"
-)
-
-# Dataset arguments
-dataset_args=(
-  --data_path "$DATA_DIR"
-  --dataloader_num_workers 1
-)
-
-# Validation arguments
-validation_args=(
-  --log_validation
-  --validation_preprocessed_path "$VALIDATION_DIR"
-  --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 6000
-  --weight_decay 1e-4
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --allow_tf32
-  --checkpoints_total_limit 3
-  --cfg 0.0
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-)
-
-# 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[@]}"

examples/training/finetune/wan_i2v_14b_480p/crush_smol/finetune_i2v.slurm

@@ -1,97 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=FV_2N_14B
-#SBATCH --partition=main
-#SBATCH --qos=hao
-#SBATCH --nodes=4
-#SBATCH --ntasks=4
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --nodelist=fs-mbz-gpu-[400-550]
-#SBATCH --mem=1440G
-#SBATCH --output=4n_i2v/4n_i2v_%j.out
-#SBATCH --error=4n_i2v/4n_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_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=8
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA
-
-DATA_DIR=data/crush-smol_processed_i2v/combined_parquet_dataset
-VALIDATION_DIR=data/crush-smol_processed_i2v/validation_parquet_dataset
-NUM_GPUS=8
-# 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
-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\
-    --model_path Wan-AI/Wan2.1-I2V-14B-480P-Diffusers \
-    --inference_mode False\
-    --pretrained_model_name_or_path Wan-AI/Wan2.1-I2V-14B-480P-Diffusers \
-    --cache_dir "/home/ray/.cache"\
-    --data_path "$DATA_DIR"\
-    --validation_preprocessed_path "$VALIDATION_DIR"\
-    --train_batch_size=1\
-    --num_latent_t 16 \
-    --num_gpus $NUM_GPUS \
-    --sp_size $NUM_GPUS \
-    --tp_size $NUM_GPUS \
-    --hsdp_replicate_dim $SLURM_JOB_NUM_NODES \
-    --hsdp_shard_dim $NUM_GPUS \
-    --train_sp_batch_size 1\
-    --dataloader_num_workers 10\
-    --gradient_accumulation_steps=2\
-    --max_train_steps=10000 \
-    --learning_rate=5e-5\
-    --mixed_precision="bf16"\
-    --checkpointing_steps=11000 \
-    --validation_steps 100\
-    --validation_sampling_steps "40" \
-    --log_validation \
-    --checkpoints_total_limit 3\
-    --allow_tf32\
-    --ema_start_step 0\
-    --cfg 0.0\
-    --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"\
-    --tracker_project_name wan_i2v_finetune \
-    --num_height 480 \
-    --num_width 832 \
-    --num_frames 77 \
-    --validation_guidance_scale "1.0" \
-    --num_euler_timesteps 50 \
-    --multi_phased_distill_schedule "4000-1" \
-    --weight_decay 1e-4 \
-    --not_apply_cfg_solver \
-    --dit_precision "fp32" \
-    --max_grad_norm 1.0

examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json

@@ -1,31 +0,0 @@
-{
-  "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
-    }
-  ]
-}

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.sh

@@ -1,88 +0,0 @@
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-# 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_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/"
-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 1
-  --num_latent_t 8 
-  --num_height 480
-  --num_width 832
-  --num_frames 77
-)
-
-# Parallel arguments
-parallel_args=(
-  --num_gpus $NUM_GPUS \
-  --sp_size $NUM_GPUS \
-  --tp_size $NUM_GPUS \
-  --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_preprocessed_path $VALIDATION_DIR
-  --validation_steps 100
-  --validation_sampling_steps "50" 
-  --validation_guidance_scale "1.0"
-)
-
-# Optimizer arguments
-optimizer_args=(
-  --learning_rate 1e-5
-  --mixed_precision "bf16"
-  --checkpointing_steps 6000
-  --weight_decay 0.01
-  --max_grad_norm 1.0
-)
-
-# Miscellaneous arguments
-miscellaneous_args=(
-  --inference_mode False
-  --allow_tf32
-  --checkpoints_total_limit 3
-  --cfg 0.0
-  --multi_phased_distill_schedule "4000-1"
-  --not_apply_cfg_solver
-  --dit_precision "fp32"
-  --num_euler_timesteps 50
-  --ema_start_step 0
-)
-
-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[@]}"

examples/training/finetune/wan_t2v_1_3b/crush_smol/finetune_t2v.slurm

@@ -1,98 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=FV_2N_14B
-#SBATCH --partition=main
-#SBATCH --qos=hao
-#SBATCH --nodes=1
-#SBATCH --ntasks=1
-#SBATCH --ntasks-per-node=1
-#SBATCH --gres=gpu:8
-#SBATCH --cpus-per-task=128
-#SBATCH --nodelist=fs-mbz-gpu-[400-550]
-#SBATCH --mem=1440G
-#SBATCH --output=4n_i2v/4n_i2v_%j.out
-#SBATCH --error=4n_i2v/4n_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_API_KEY='8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde'
-
-echo "MASTER_ADDR: $MASTER_ADDR"
-echo "NODE_RANK: $NODE_RANK"
-
-# Configs
-NUM_GPUS=8
-
-export WANDB_BASE_URL="https://api.wandb.ai"
-export WANDB_MODE=online
-# 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_DIR="data/crush-smol_processed_t2v/validation_parquet_dataset/"
-NUM_GPUS=8
-# 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
-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\
-    --model_path $MODEL_PATH \
-    --inference_mode False\
-    --pretrained_model_name_or_path $MODEL_PATH \
-    --cache_dir "/home/ray/.cache"\
-    --data_path "$DATA_DIR"\
-    --validation_preprocessed_path "$VALIDATION_DIR"\
-    --train_batch_size=1\
-    --num_latent_t 8 \
-    --num_gpus $NUM_GPUS \
-    --sp_size $NUM_GPUS \
-    --tp_size $NUM_GPUS \
-    --hsdp_replicate_dim $SLURM_JOB_NUM_NODES \
-    --hsdp_shard_dim $NUM_GPUS \
-    --train_sp_batch_size 1\
-    --dataloader_num_workers 10\
-    --gradient_accumulation_steps=1\
-    --max_train_steps=10000 \
-    --learning_rate=5e-5\
-    --mixed_precision="bf16"\
-    --checkpointing_steps=11000 \
-    --validation_steps 100\
-    --validation_sampling_steps "40" \
-    --log_validation \
-    --checkpoints_total_limit 3\
-    --allow_tf32\
-    --ema_start_step 0\
-    --cfg 0.0\
-    --output_dir="$DATA_DIR/outputs/wan_i2v_finetune_2n"\
-    --tracker_project_name wan_i2v_finetune \
-    --num_height 480 \
-    --num_width 832 \
-    --num_frames 77 \
-    --validation_guidance_scale "1.0" \
-    --num_euler_timesteps 50 \
-    --multi_phased_distill_schedule "4000-1" \
-    --weight_decay 1e-4 \
-    --not_apply_cfg_solver \
-    --dit_precision "fp32" \
-    --max_grad_norm 1.0

examples/training/finetune/wan_t2v_1_3b/crush_smol/overfit.json

@@ -1,13 +0,0 @@
-{
-  "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": "examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation_dataset/yYcK4nANZz4-Scene-034.mp4",
-      "num_inference_steps": 40,
-      "height": 480,
-      "width": 832,
-      "num_frames": 77
-    }
-  ]
-}

examples/training/finetune/wan_t2v_1_3b/crush_smol/v1_preprocess_wan_data_t2v.sh

@@ -1,24 +0,0 @@
-# export WANDB_MODE="offline"
-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 \
-    --preprocess_task "t2v" 

examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json

@@ -1,31 +0,0 @@
-{
-  "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
-    }
-  ]
-}

fastvideo/data_preprocess/preprocess.py

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

fastvideo/data_preprocess/preprocess_vae_latents.py

@@ -104,7 +104,13 @@ if __name__ == "__main__":
         default=None,
         help="The output directory where the model predictions and checkpoints will be written.",
     )
-
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
 
     args = parser.parse_args()
     main(args)

fastvideo/distill.py

@@ -671,6 +671,13 @@ if __name__ == "__main__":
         help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
               ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
     )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
 
     # optimizer & scheduler & Training
     parser.add_argument("--num_train_epochs", type=int, default=100)

fastvideo/distill_adv.py

@@ -693,6 +693,13 @@ if __name__ == "__main__":
         help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
               ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
     )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
 
     # optimizer & scheduler & Training
     parser.add_argument("--num_train_epochs", type=int, default=100)

fastvideo/train.py

@@ -520,7 +520,13 @@ if __name__ == "__main__":
         help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
               ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
     )
-
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
 
     # optimizer & scheduler & Training
     parser.add_argument("--num_train_epochs", type=int, default=100)

fastvideo/v1/STA_configuration.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 import json
 import os
 from collections import defaultdict
@@ -6,8 +5,6 @@ from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 
-from fastvideo.v1.utils import dict_to_3d_list
-
 
 def configure_sta(mode: str = 'STA_searching',
                   layer_num: int = 40,
@@ -351,6 +348,21 @@ 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]]],
         prompt: str,

fastvideo/v1/attention/__init__.py

@@ -3,15 +3,12 @@
 from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
                                                       AttentionMetadata,
                                                       AttentionMetadataBuilder)
-from fastvideo.v1.attention.layer import (DistributedAttention,
-                                          DistributedAttention_VSA,
-                                          LocalAttention)
+from fastvideo.v1.attention.layer import DistributedAttention, LocalAttention
 from fastvideo.v1.attention.selector import get_attn_backend
 
 __all__ = [
     "DistributedAttention",
     "LocalAttention",
-    "DistributedAttention_VSA",
     "AttentionBackend",
     "AttentionMetadata",
     "AttentionMetadataBuilder",

fastvideo/v1/attention/backends/sage_attn.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from typing import List, Optional, Type
 
 import torch

fastvideo/v1/attention/backends/sdpa.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from typing import List, Optional, Type
 
 import torch

fastvideo/v1/attention/backends/sliding_tile_attn.py

@@ -1,7 +1,6 @@
-# SPDX-License-Identifier: Apache-2.0
 import json
 from dataclasses import dataclass
-from typing import Any, List, Optional, Type
+from typing import Any, Dict, List, Optional, Type
 
 import torch
 from einops import rearrange
@@ -17,11 +16,33 @@ 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:
@@ -118,7 +139,7 @@ class SlidingTileAttentionImpl(AttentionImpl):
         self.sp_size = sp_group.world_size
         # STA config
         self.STA_base_tile_size = [6, 8, 8]
-        self.dit_seq_shape_mapping = RangeDict({
+        self.img_latent_shape_mapping = RangeDict({
             (115200, 115456): '30x48x80',
             82944: '36x48x48',
             69120: '18x48x80',
@@ -133,9 +154,9 @@ class SlidingTileAttentionImpl(AttentionImpl):
         x = rearrange(x,
                       "b (sp t h w) head d -> b (t sp h w) head d",
                       sp=self.sp_size,
-                      t=self.dit_seq_shape_int[0] // self.sp_size,
-                      h=self.dit_seq_shape_int[1],
-                      w=self.dit_seq_shape_int[2])
+                      t=self.img_latent_shape_int[0] // self.sp_size,
+                      h=self.img_latent_shape_int[1],
+                      w=self.img_latent_shape_int[2])
         return rearrange(
             x,
             "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
@@ -159,9 +180,9 @@ class SlidingTileAttentionImpl(AttentionImpl):
         return rearrange(x,
                          "b (t sp h w) head d -> b (sp t h w) head d",
                          sp=self.sp_size,
-                         t=self.dit_seq_shape_int[0] // self.sp_size,
-                         h=self.dit_seq_shape_int[1],
-                         w=self.dit_seq_shape_int[2])
+                         t=self.img_latent_shape_int[0] // self.sp_size,
+                         h=self.img_latent_shape_int[1],
+                         w=self.img_latent_shape_int[2])
 
     def preprocess_qkv(
         self,
@@ -169,12 +190,14 @@ class SlidingTileAttentionImpl(AttentionImpl):
         attn_metadata: AttentionMetadata,
     ) -> torch.Tensor:
         img_sequence_length = qkv.shape[1]
-        self.dit_seq_shape_str = self.dit_seq_shape_mapping[img_sequence_length]
-        self.full_window_size = self.full_window_mapping[self.dit_seq_shape_str]
-        self.dit_seq_shape_int = list(
-            map(int, self.dit_seq_shape_str.split('x')))
-        self.img_seq_length = self.dit_seq_shape_int[
-            0] * self.dit_seq_shape_int[1] * self.dit_seq_shape_int[2]
+        self.img_latent_shape_str = self.img_latent_shape_mapping[
+            img_sequence_length]
+        self.full_window_size = self.full_window_mapping[
+            self.img_latent_shape_str]
+        self.img_latent_shape_int = list(
+            map(int, self.img_latent_shape_str.split('x')))
+        self.img_seq_length = self.img_latent_shape_int[
+            0] * self.img_latent_shape_int[1] * self.img_latent_shape_int[2]
         return self.tile(qkv)
 
     def postprocess_output(
@@ -229,12 +252,12 @@ class SlidingTileAttentionImpl(AttentionImpl):
             for window_size in STA_param[:-1]:
                 sparse_hidden_states = sliding_tile_attention(
                     query, key, value, [window_size] * head_num, text_length,
-                    has_text, self.dit_seq_shape_str).transpose(1, 2)
+                    has_text, self.img_latent_shape_str).transpose(1, 2)
                 sparse_attn_hidden_states_all.append(sparse_hidden_states)
 
             hidden_states = sliding_tile_attention(
                 query, key, value, [full_mask_window] * head_num, text_length,
-                has_text, self.dit_seq_shape_str).transpose(1, 2)
+                has_text, self.img_latent_shape_str).transpose(1, 2)
 
             attn_L2_loss = []
             attn_L1_loss = []
@@ -265,12 +288,18 @@ class SlidingTileAttentionImpl(AttentionImpl):
                     forward_batch.mask_search_final_result_pos[timestep].append(
                         layer_loss_save)
         else:
+            # windows = [
+            #     self.mask_strategy[timestep][layer_idx][head_idx + start_head]
+            #     for head_idx in range(head_num)
+            # ]
             windows = [
                 STA_param[head_idx + start_head] for head_idx in range(head_num)
             ]
-
+            # if has_text is False:
+            #     from IPython import embed
+            #     embed()
             hidden_states = sliding_tile_attention(
                 query, key, value, windows, text_length, has_text,
-                self.dit_seq_shape_str).transpose(1, 2)
+                self.img_latent_shape_str).transpose(1, 2)
 
         return hidden_states

fastvideo/v1/attention/backends/video_sparse_attn.py

@@ -1,198 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import math
-from dataclasses import dataclass
-from typing import List, Optional, Type
-
-import torch
-from einops import rearrange
-
-try:
-    from vsa import video_sparse_attn
-except ImportError:
-    video_sparse_attn = None
-
-from fastvideo.v1.attention.backends.abstract import (AttentionBackend,
-                                                      AttentionImpl,
-                                                      AttentionMetadata,
-                                                      AttentionMetadataBuilder)
-from fastvideo.v1.distributed import get_sp_group
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
-
-logger = init_logger(__name__)
-
-
-class VideoSparseAttentionBackend(AttentionBackend):
-
-    accept_output_buffer: bool = True
-
-    @staticmethod
-    def get_supported_head_sizes() -> List[int]:
-        return [64, 128]
-
-    @staticmethod
-    def get_name() -> str:
-        return "VIDEO_SPARSE_ATTN"
-
-    @staticmethod
-    def get_impl_cls() -> Type["VideoSparseAttentionImpl"]:
-        return VideoSparseAttentionImpl
-
-    @staticmethod
-    def get_metadata_cls() -> Type["VideoSparseAttentionMetadata"]:
-        return VideoSparseAttentionMetadata
-
-    @staticmethod
-    def get_builder_cls() -> Type["VideoSparseAttentionMetadataBuilder"]:
-        return VideoSparseAttentionMetadataBuilder
-
-
-@dataclass
-class VideoSparseAttentionMetadata(AttentionMetadata):
-    current_timestep: int
-    dit_seq_shape: List[int]
-    VSA_sparsity: float
-
-
-class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
-
-    def __init__(self):
-        pass
-
-    def prepare(self):
-        pass
-
-    def build(
-        self,
-        current_timestep: int,
-        forward_batch: ForwardBatch,
-        fastvideo_args: FastVideoArgs,
-    ) -> VideoSparseAttentionMetadata:
-        if forward_batch.latents is None:
-            raise ValueError("latents cannot be None")
-
-        raw_latent_shape = forward_batch.raw_latent_shape
-        if raw_latent_shape is None:
-            raise ValueError("raw_latent_shape cannot be None")
-
-        patch_size = fastvideo_args.pipeline_config.dit_config.patch_size
-        dit_seq_shape = [
-            raw_latent_shape[2] // patch_size[0],
-            raw_latent_shape[3] // patch_size[1],
-            raw_latent_shape[4] // patch_size[2]
-        ]
-        VSA_sparsity = forward_batch.VSA_sparsity
-
-        return VideoSparseAttentionMetadata(current_timestep=current_timestep,
-                                            dit_seq_shape=dit_seq_shape,
-                                            VSA_sparsity=VSA_sparsity)
-
-
-class VideoSparseAttentionImpl(AttentionImpl):
-
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        causal: bool,
-        softmax_scale: float,
-        num_kv_heads: Optional[int] = None,
-        prefix: str = "",
-        **extra_impl_args,
-    ) -> None:
-        self.prefix = prefix
-        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.full_window_size: List[int]
-        self.img_seq_length: int
-
-    def tile(self, x: torch.Tensor) -> torch.Tensor:
-        x = rearrange(x,
-                      "b (sp t h w) head d -> b (t sp h w) head d",
-                      sp=self.sp_size,
-                      t=self.dit_seq_shape[0] // self.sp_size,
-                      h=self.dit_seq_shape[1],
-                      w=self.dit_seq_shape[2])
-
-        return rearrange(
-            x,
-            "b (n_t ts_t n_h ts_h n_w ts_w) h d -> b (n_t n_h n_w ts_t ts_h ts_w) h d",
-            n_t=self.full_window_size[0],
-            n_h=self.full_window_size[1],
-            n_w=self.full_window_size[2],
-            ts_t=self.VSA_base_tile_size[0],
-            ts_h=self.VSA_base_tile_size[1],
-            ts_w=self.VSA_base_tile_size[2])
-
-    def untile(self, x: torch.Tensor) -> torch.Tensor:
-        x = rearrange(
-            x,
-            "b (n_t n_h n_w ts_t ts_h ts_w) h d -> b (n_t ts_t n_h ts_h n_w ts_w) h d",
-            n_t=self.full_window_size[0],
-            n_h=self.full_window_size[1],
-            n_w=self.full_window_size[2],
-            ts_t=self.VSA_base_tile_size[0],
-            ts_h=self.VSA_base_tile_size[1],
-            ts_w=self.VSA_base_tile_size[2])
-        return rearrange(x,
-                         "b (t sp h w) head d -> b (sp t h w) head d",
-                         sp=self.sp_size,
-                         t=self.dit_seq_shape[0] // self.sp_size,
-                         h=self.dit_seq_shape[1],
-                         w=self.dit_seq_shape[2])
-
-    def preprocess_qkv(
-        self,
-        qkv: torch.Tensor,
-        attn_metadata: VideoSparseAttentionMetadata,
-    ) -> torch.Tensor:
-        self.dit_seq_shape = attn_metadata.dit_seq_shape
-        self.full_window_size = [
-            self.dit_seq_shape[0] // self.VSA_base_tile_size[0],
-            self.dit_seq_shape[1] // self.VSA_base_tile_size[1],
-            self.dit_seq_shape[2] // self.VSA_base_tile_size[2]
-        ]
-        self.img_seq_length = math.prod(self.dit_seq_shape)
-        return self.tile(qkv)
-
-    def postprocess_output(
-        self,
-        output: torch.Tensor,
-        attn_metadata: VideoSparseAttentionMetadata,
-    ) -> torch.Tensor:
-        return self.untile(output)
-
-    def forward(  # type: ignore[override]
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        gate_compress: torch.Tensor,
-        attn_metadata: VideoSparseAttentionMetadata,
-    ) -> torch.Tensor:
-        query = query.transpose(1, 2).contiguous()
-        key = key.transpose(1, 2).contiguous()
-        value = value.transpose(1, 2).contiguous()
-        gate_compress = gate_compress.transpose(1, 2).contiguous()
-
-        VSA_sparsity = attn_metadata.VSA_sparsity
-
-        cur_topk = math.ceil(
-            (1 - VSA_sparsity) *
-            (self.img_seq_length / math.prod(self.VSA_base_tile_size)))
-
-        if video_sparse_attn is None:
-            raise NotImplementedError("video_sparse_attn is not installed")
-
-        hidden_states = video_sparse_attn(
-            query,
-            key,
-            value,
-            topk=cur_topk,
-            block_size=(4, 4, 4),
-            compress_attn_weight=gate_compress).transpose(1, 2)
-
-        return hidden_states

fastvideo/v1/attention/layer.py

@@ -9,10 +9,10 @@ from fastvideo.v1.attention.selector import (backend_name_to_enum,
                                              get_attn_backend)
 from fastvideo.v1.distributed.communication_op import (
     sequence_model_parallel_all_gather, sequence_model_parallel_all_to_all_4D)
-from fastvideo.v1.distributed.parallel_state import (get_sp_parallel_rank,
-                                                     get_sp_world_size)
+from fastvideo.v1.distributed.parallel_state import (
+    get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size)
 from fastvideo.v1.forward_context import ForwardContext, get_forward_context
-from fastvideo.v1.platforms import AttentionBackendEnum
+from fastvideo.v1.platforms import _Backend
 from fastvideo.v1.utils import get_compute_dtype
 
 
@@ -26,8 +26,8 @@ class DistributedAttention(nn.Module):
                  num_kv_heads: Optional[int] = None,
                  softmax_scale: Optional[float] = None,
                  causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[
-                     AttentionBackendEnum, ...]] = None,
+                 supported_attention_backends: Optional[Tuple[_Backend,
+                                                              ...]] = None,
                  prefix: str = "",
                  **extra_impl_args) -> None:
         super().__init__()
@@ -45,13 +45,13 @@ class DistributedAttention(nn.Module):
             dtype,
             supported_attention_backends=supported_attention_backends)
         impl_cls = attn_backend.get_impl_cls()
-        self.attn_impl = impl_cls(num_heads=num_heads,
-                                  head_size=head_size,
-                                  causal=causal,
-                                  softmax_scale=self.softmax_scale,
-                                  num_kv_heads=num_kv_heads,
-                                  prefix=f"{prefix}.impl",
-                                  **extra_impl_args)
+        self.impl = impl_cls(num_heads=num_heads,
+                             head_size=head_size,
+                             causal=causal,
+                             softmax_scale=self.softmax_scale,
+                             num_kv_heads=num_kv_heads,
+                             prefix=f"{prefix}.impl",
+                             **extra_impl_args)
         self.num_heads = num_heads
         self.head_size = head_size
         self.num_kv_heads = num_kv_heads
@@ -86,8 +86,8 @@ class DistributedAttention(nn.Module):
         assert q.dim() == 4 and k.dim() == 4 and v.dim(
         ) == 4, "Expected 4D tensors"
         batch_size, seq_len, num_heads, head_dim = q.shape
-        local_rank = get_sp_parallel_rank()
-        world_size = get_sp_world_size()
+        local_rank = get_sequence_model_parallel_rank()
+        world_size = get_sequence_model_parallel_world_size()
 
         forward_context: ForwardContext = get_forward_context()
         ctx_attn_metadata = forward_context.attn_metadata
@@ -100,7 +100,7 @@ class DistributedAttention(nn.Module):
                                                     scatter_dim=2,
                                                     gather_dim=1)
         # Apply backend-specific preprocess_qkv
-        qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata)
+        qkv = self.impl.preprocess_qkv(qkv, ctx_attn_metadata)
 
         # Concatenate with replicated QKV if provided
         if replicated_q is not None:
@@ -116,7 +116,7 @@ class DistributedAttention(nn.Module):
 
         q, k, v = qkv.chunk(3, dim=0)
 
-        output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
+        output = self.impl.forward(q, k, v, ctx_attn_metadata)
 
         # Redistribute back if using sequence parallelism
         replicated_output = None
@@ -127,73 +127,7 @@ class DistributedAttention(nn.Module):
             replicated_output = sequence_model_parallel_all_gather(
                 replicated_output.contiguous(), dim=2)
         # Apply backend-specific postprocess_output
-        output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
-
-        output = sequence_model_parallel_all_to_all_4D(output,
-                                                       scatter_dim=1,
-                                                       gather_dim=2)
-        return output, replicated_output
-
-
-class DistributedAttention_VSA(DistributedAttention):
-    """Distributed attention layer with VSA support.
-    """
-
-    def forward(
-        self,
-        q: torch.Tensor,
-        k: torch.Tensor,
-        v: torch.Tensor,
-        replicated_q: Optional[torch.Tensor] = None,
-        replicated_k: Optional[torch.Tensor] = None,
-        replicated_v: Optional[torch.Tensor] = None,
-        gate_compress: Optional[torch.Tensor] = None,
-    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-        """Forward pass for distributed attention.
-        
-        Args:
-            q (torch.Tensor): Query tensor [batch_size, seq_len, num_heads, head_dim]
-            k (torch.Tensor): Key tensor [batch_size, seq_len, num_heads, head_dim]
-            v (torch.Tensor): Value tensor [batch_size, seq_len, num_heads, head_dim]
-            gate_compress (torch.Tensor): Gate compress tensor [batch_size, seq_len, num_heads, head_dim]
-            replicated_q (Optional[torch.Tensor]): Replicated query tensor, typically for text tokens
-            replicated_k (Optional[torch.Tensor]): Replicated key tensor
-            replicated_v (Optional[torch.Tensor]): Replicated value tensor
-            
-        Returns:
-            Tuple[torch.Tensor, Optional[torch.Tensor]]: A tuple containing:
-                - o (torch.Tensor): Output tensor after attention for the main sequence
-                - replicated_o (Optional[torch.Tensor]): Output tensor for replicated tokens, if provided
-        """
-        # Check text tokens are not supported for VSA now
-        assert replicated_q is None and replicated_k is None and replicated_v is None, "Replicated QKV is not supported for VSA now"
-        # Check input shapes
-        assert q.dim() == 4 and k.dim() == 4 and v.dim(
-        ) == 4, "Expected 4D tensors"
-
-        forward_context: ForwardContext = get_forward_context()
-        ctx_attn_metadata = forward_context.attn_metadata
-
-        # Stack QKV
-        qkvg = torch.cat([q, k, v, gate_compress],
-                         dim=0)  # [3, seq_len, num_heads, head_dim]
-
-        # Redistribute heads across sequence dimension
-        qkvg = sequence_model_parallel_all_to_all_4D(qkvg,
-                                                     scatter_dim=2,
-                                                     gather_dim=1)
-
-        qkvg = self.attn_impl.preprocess_qkv(qkvg, ctx_attn_metadata)
-
-        q, k, v, gate_compress = qkvg.chunk(4, dim=0)
-        output = self.attn_impl.forward(
-            q, k, v, gate_compress, ctx_attn_metadata)  # type: ignore[call-arg]
-
-        # Redistribute back if using sequence parallelism
-        replicated_output = None
-
-        # Apply backend-specific postprocess_output
-        output = self.attn_impl.postprocess_output(output, ctx_attn_metadata)
+        output = self.impl.postprocess_output(output, ctx_attn_metadata)
 
         output = sequence_model_parallel_all_to_all_4D(output,
                                                        scatter_dim=1,
@@ -211,8 +145,8 @@ class LocalAttention(nn.Module):
                  num_kv_heads: Optional[int] = None,
                  softmax_scale: Optional[float] = None,
                  causal: bool = False,
-                 supported_attention_backends: Optional[Tuple[
-                     AttentionBackendEnum, ...]] = None,
+                 supported_attention_backends: Optional[Tuple[_Backend,
+                                                              ...]] = None,
                  **extra_impl_args) -> None:
         super().__init__()
         if softmax_scale is None:
@@ -228,12 +162,12 @@ class LocalAttention(nn.Module):
             dtype,
             supported_attention_backends=supported_attention_backends)
         impl_cls = attn_backend.get_impl_cls()
-        self.attn_impl = impl_cls(num_heads=num_heads,
-                                  head_size=head_size,
-                                  softmax_scale=self.softmax_scale,
-                                  num_kv_heads=num_kv_heads,
-                                  causal=causal,
-                                  **extra_impl_args)
+        self.impl = impl_cls(num_heads=num_heads,
+                             head_size=head_size,
+                             softmax_scale=self.softmax_scale,
+                             num_kv_heads=num_kv_heads,
+                             causal=causal,
+                             **extra_impl_args)
         self.num_heads = num_heads
         self.head_size = head_size
         self.num_kv_heads = num_kv_heads
@@ -264,5 +198,5 @@ class LocalAttention(nn.Module):
         forward_context: ForwardContext = get_forward_context()
         ctx_attn_metadata = forward_context.attn_metadata
 
-        output = self.attn_impl.forward(q, k, v, ctx_attn_metadata)
+        output = self.impl.forward(q, k, v, ctx_attn_metadata)
         return output

fastvideo/v1/attention/selector.py

@@ -11,13 +11,13 @@ import torch
 import fastvideo.v1.envs as envs
 from fastvideo.v1.attention.backends.abstract import AttentionBackend
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.platforms import AttentionBackendEnum, current_platform
+from fastvideo.v1.platforms import _Backend, current_platform
 from fastvideo.v1.utils import STR_BACKEND_ENV_VAR, resolve_obj_by_qualname
 
 logger = init_logger(__name__)
 
 
-def backend_name_to_enum(backend_name: str) -> Optional[AttentionBackendEnum]:
+def backend_name_to_enum(backend_name: str) -> Optional[_Backend]:
     """
     Convert a string backend name to a _Backend enum value.
 
@@ -27,11 +27,11 @@ def backend_name_to_enum(backend_name: str) -> Optional[AttentionBackendEnum]:
             loaded.
     """
     assert backend_name is not None
-    return AttentionBackendEnum[backend_name] if backend_name in AttentionBackendEnum.__members__ else \
+    return _Backend[backend_name] if backend_name in _Backend.__members__ else \
           None
 
 
-def get_env_variable_attn_backend() -> Optional[AttentionBackendEnum]:
+def get_env_variable_attn_backend() -> Optional[_Backend]:
     '''
     Get the backend override specified by the FastVideo attention
     backend environment variable, if one is specified.
@@ -53,11 +53,10 @@ 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: Optional[_Backend] = None
 
 
-def global_force_attn_backend(
-        attn_backend: Optional[AttentionBackendEnum]) -> None:
+def global_force_attn_backend(attn_backend: Optional[_Backend]) -> None:
     '''
     Force all attention operations to use a specified backend.
 
@@ -72,7 +71,7 @@ def global_force_attn_backend(
     forced_attn_backend = attn_backend
 
 
-def get_global_forced_attn_backend() -> Optional[AttentionBackendEnum]:
+def get_global_forced_attn_backend() -> Optional[_Backend]:
     '''
     Get the currently-forced choice of attention backend,
     or None if auto-selection is currently enabled.
@@ -83,8 +82,7 @@ def get_global_forced_attn_backend() -> Optional[AttentionBackendEnum]:
 def get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                 ...]] = None,
+    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
 ) -> Type[AttentionBackend]:
     return _cached_get_attn_backend(head_size, dtype,
                                     supported_attention_backends)
@@ -94,8 +92,7 @@ def get_attn_backend(
 def _cached_get_attn_backend(
     head_size: int,
     dtype: torch.dtype,
-    supported_attention_backends: Optional[Tuple[AttentionBackendEnum,
-                                                 ...]] = None,
+    supported_attention_backends: Optional[Tuple[_Backend, ...]] = None,
 ) -> Type[AttentionBackend]:
     # Check whether a particular choice of backend was
     # previously forced.
@@ -105,7 +102,7 @@ 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: Optional[_Backend] = (
         get_global_forced_attn_backend())
     if backend_by_global_setting is not None:
         selected_backend = backend_by_global_setting
@@ -128,7 +125,7 @@ def _cached_get_attn_backend(
 
 @contextmanager
 def global_force_attn_backend_context_manager(
-        attn_backend: AttentionBackendEnum) -> Generator[None, None, None]:
+        attn_backend: _Backend) -> Generator[None, None, None]:
     '''
     Globally force a FastVideo attention backend override within a
     context manager, reverting the global attention backend

fastvideo/v1/configs/models/base.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field, fields
 from typing import Any, Dict
 

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

@@ -1,10 +1,9 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Any, List, Optional, Tuple
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
 from fastvideo.v1.layers.quantization import QuantizationConfig
-from fastvideo.v1.platforms import AttentionBackendEnum
+from fastvideo.v1.platforms import _Backend
 
 
 @dataclass
@@ -13,10 +12,11 @@ class DiTArchConfig(ArchConfig):
     _compile_conditions: list = field(default_factory=list)
     _param_names_mapping: dict = field(default_factory=dict)
     _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)
+    _supported_attention_backends: Tuple[_Backend,
+                                         ...] = (_Backend.SLIDING_TILE_ATTN,
+                                                 _Backend.SAGE_ATTN,
+                                                 _Backend.FLASH_ATTN,
+                                                 _Backend.TORCH_SDPA)
 
     hidden_size: int = 0
     num_attention_heads: int = 0

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

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import List, Optional, Tuple
 

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

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import List, Optional, Tuple, Union
 

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

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import List, Optional, Tuple
 

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

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Any, Dict, List, Optional, Tuple
 
@@ -6,14 +5,14 @@ import torch
 
 from fastvideo.v1.configs.models.base import ArchConfig, ModelConfig
 from fastvideo.v1.layers.quantization import QuantizationConfig
-from fastvideo.v1.platforms import AttentionBackendEnum
+from fastvideo.v1.platforms import _Backend
 
 
 @dataclass
 class EncoderArchConfig(ArchConfig):
     architectures: List[str] = field(default_factory=lambda: [])
-    _supported_attention_backends: Tuple[AttentionBackendEnum, ...] = (
-        AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA)
+    _supported_attention_backends: Tuple[_Backend, ...] = (_Backend.FLASH_ATTN,
+                                                           _Backend.TORCH_SDPA)
     output_hidden_states: bool = False
     use_return_dict: bool = True
 

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

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Optional
 

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

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Optional
 

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

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Optional
 

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

@@ -1,6 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
-import argparse
-import dataclasses
 from dataclasses import dataclass, field
 from typing import Any, Union
 
@@ -131,12 +128,3 @@ class VAEConfig(ModelConfig):
         )
 
         return parser
-
-    @classmethod
-    def from_cli_args(cls, args: argparse.Namespace) -> "VAEConfig":
-        kwargs = {}
-        for attr in dataclasses.fields(cls):
-            value = getattr(args, attr.name, None)
-            if value is not None:
-                kwargs[attr.name] = value
-        return cls(**kwargs)

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

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Tuple
 

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

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 
 from fastvideo.v1.configs.models.vaes.base import VAEArchConfig, VAEConfig

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

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Tuple
 

fastvideo/v1/configs/pipelines/__init__.py

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

fastvideo/v1/configs/pipelines/base.py

@@ -1,31 +1,18 @@
-# SPDX-License-Identifier: Apache-2.0
 import json
 from dataclasses import asdict, dataclass, field, fields
-from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing import Any, Callable, Dict, Optional, Tuple, cast
 
 import torch
 
 from fastvideo.v1.configs.models import (DiTConfig, EncoderConfig, ModelConfig,
                                          VAEConfig)
 from fastvideo.v1.configs.models.encoders import BaseEncoderOutput
-from fastvideo.v1.configs.utils import update_config_from_args
 from fastvideo.v1.logger import init_logger
-from fastvideo.v1.utils import (FlexibleArgumentParser, StoreBoolean,
-                                shallow_asdict)
+from fastvideo.v1.utils import shallow_asdict
 
 logger = init_logger(__name__)
 
 
-class STA_Mode(str, Enum):
-    """STA (Sliding Tile Attention) modes."""
-    STA_INFERENCE = "STA_inference"
-    STA_SEARCHING = "STA_searching"
-    STA_TUNING = "STA_tuning"
-    STA_TUNING_CFG = "STA_tuning_cfg"
-    NONE = None
-
-
 def preprocess_text(prompt: str) -> str:
     return prompt
 
@@ -34,282 +21,59 @@ def postprocess_text(output: BaseEncoderOutput) -> torch.tensor:
     raise NotImplementedError
 
 
-# config for a single pipeline
 @dataclass
 class PipelineConfig:
     """Base configuration for all pipeline architectures."""
-    model_path: str = ""
-    pipeline_config_path: Optional[str] = None
-
     # Video generation parameters
     embedded_cfg_scale: float = 6.0
     flow_shift: Optional[float] = None
     disable_autocast: bool = False
 
     # Model configuration
-    dit_config: DiTConfig = field(default_factory=DiTConfig)
-    dit_precision: str = "bf16"
+    precision: str = "bf16"
 
     # VAE configuration
-    vae_config: VAEConfig = field(default_factory=VAEConfig)
     vae_precision: str = "fp16"
     vae_tiling: bool = True
     vae_sp: bool = True
+    vae_config: VAEConfig = field(default_factory=VAEConfig)
 
-    # Image encoder configuration
-    image_encoder_config: EncoderConfig = field(default_factory=EncoderConfig)
-    image_encoder_precision: str = "fp32"
+    # DiT configuration
+    dit_config: DiTConfig = field(default_factory=DiTConfig)
 
     # Text encoder configuration
-    DEFAULT_TEXT_ENCODER_PRECISIONS = ("fp16", )
-    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
-        default_factory=lambda: (EncoderConfig(), ))
     text_encoder_precisions: Tuple[str, ...] = field(
         default_factory=lambda: ("fp16", ))
+    text_encoder_configs: Tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (EncoderConfig(), ))
     preprocess_text_funcs: Tuple[Callable[[str], str], ...] = field(
         default_factory=lambda: (preprocess_text, ))
     postprocess_text_funcs: Tuple[Callable[[BaseEncoderOutput], torch.tensor],
                                   ...] = field(default_factory=lambda:
                                                (postprocess_text, ))
 
-    # LoRA parameters
-    lora_path: Optional[str] = None
-    lora_nickname: Optional[
-        str] = "default"  # for swapping adapters in the pipeline
-    lora_target_names: Optional[List[
-        str]] = None  # can restrict list of layers to adapt, e.g. ["q_proj"]
-
-    # StepVideo specific parameters
-    pos_magic: Optional[str] = None
-    neg_magic: Optional[str] = None
-    timesteps_scale: Optional[bool] = None
-
-    # STA (Sliding Tile Attention) parameters
+    # STA (Spatial-Temporal Attention) parameters
     mask_strategy_file_path: Optional[str] = None
-    STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
+    STA_mode: str = "STA_inference"
     skip_time_steps: int = 15
 
     # Compilation
-    # enable_torch_compile: bool = False
-
-    @staticmethod
-    def add_cli_args(parser: FlexibleArgumentParser,
-                     prefix: str = "") -> FlexibleArgumentParser:
-        prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else ""
-
-        # model_path will be conflicting with the model_path in FastVideoArgs,
-        # so we add it separately if prefix is not empty
-        if prefix_with_dot != "":
-            parser.add_argument(
-                f"--{prefix_with_dot}model-path",
-                type=str,
-                dest=f"{prefix_with_dot.replace('-', '_')}model_path",
-                default=PipelineConfig.model_path,
-                help="Path to the pretrained model",
-            )
-
-        parser.add_argument(
-            f"--{prefix_with_dot}pipeline-config-path",
-            type=str,
-            dest=f"{prefix_with_dot.replace('-', '_')}pipeline_config_path",
-            default=PipelineConfig.pipeline_config_path,
-            help="Path to the pipeline config",
-        )
-        parser.add_argument(
-            f"--{prefix_with_dot}embedded-cfg-scale",
-            type=float,
-            dest=f"{prefix_with_dot.replace('-', '_')}embedded_cfg_scale",
-            default=PipelineConfig.embedded_cfg_scale,
-            help="Embedded CFG scale",
-        )
-        parser.add_argument(
-            f"--{prefix_with_dot}flow-shift",
-            type=float,
-            dest=f"{prefix_with_dot.replace('-', '_')}flow_shift",
-            default=PipelineConfig.flow_shift,
-            help="Flow shift parameter",
-        )
-
-        # DiT configuration
-        parser.add_argument(
-            f"--{prefix_with_dot}dit-precision",
-            type=str,
-            dest=f"{prefix_with_dot.replace('-', '_')}dit_precision",
-            default=PipelineConfig.dit_precision,
-            choices=["fp32", "fp16", "bf16"],
-            help="Precision for the DiT model",
-        )
-
-        # VAE configuration
-        parser.add_argument(
-            f"--{prefix_with_dot}vae-precision",
-            type=str,
-            dest=f"{prefix_with_dot.replace('-', '_')}vae_precision",
-            default=PipelineConfig.vae_precision,
-            choices=["fp32", "fp16", "bf16"],
-            help="Precision for VAE",
-        )
-        parser.add_argument(
-            f"--{prefix_with_dot}vae-tiling",
-            action=StoreBoolean,
-            dest=f"{prefix_with_dot.replace('-', '_')}vae_tiling",
-            default=PipelineConfig.vae_tiling,
-            help="Enable VAE tiling",
-        )
-        parser.add_argument(
-            f"--{prefix_with_dot}vae-sp",
-            action=StoreBoolean,
-            dest=f"{prefix_with_dot.replace('-', '_')}vae_sp",
-            help="Enable VAE spatial parallelism",
-        )
-
-        # Text encoder configuration
-        parser.add_argument(
-            f"--{prefix_with_dot}text-encoder-precisions",
-            nargs="+",
-            type=str,
-            dest=f"{prefix_with_dot.replace('-', '_')}text_encoder_precisions",
-            default=PipelineConfig.DEFAULT_TEXT_ENCODER_PRECISIONS,
-            choices=["fp32", "fp16", "bf16"],
-            help="Precision for each text encoder",
-        )
-
-        # Image encoder configuration
-        parser.add_argument(
-            f"--{prefix_with_dot}image-encoder-precision",
-            type=str,
-            dest=f"{prefix_with_dot.replace('-', '_')}image_encoder_precision",
-            default=PipelineConfig.image_encoder_precision,
-            choices=["fp32", "fp16", "bf16"],
-            help="Precision for image encoder",
-        )
-        parser.add_argument(
-            f"--{prefix_with_dot}pos_magic",
-            type=str,
-            dest=f"{prefix_with_dot.replace('-', '_')}pos_magic",
-            default=PipelineConfig.pos_magic,
-            help="Positive magic prompt for sampling, used in stepvideo",
-        )
-        parser.add_argument(
-            f"--{prefix_with_dot}neg_magic",
-            type=str,
-            dest=f"{prefix_with_dot.replace('-', '_')}neg_magic",
-            default=PipelineConfig.neg_magic,
-            help="Negative magic prompt for sampling, used in stepvideo",
-        )
-        parser.add_argument(
-            f"--{prefix_with_dot}timesteps_scale",
-            type=bool,
-            dest=f"{prefix_with_dot.replace('-', '_')}timesteps_scale",
-            default=PipelineConfig.timesteps_scale,
-            help=
-            "Bool for applying scheduler scale in set_timesteps, used in stepvideo",
-        )
-
-        # Add VAE configuration arguments
-        from fastvideo.v1.configs.models.vaes.base import VAEConfig
-        VAEConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}vae-config")
-
-        # Add DiT configuration arguments
-        from fastvideo.v1.configs.models.dits.base import DiTConfig
-        DiTConfig.add_cli_args(parser, prefix=f"{prefix_with_dot}dit-config")
-
-        return parser
-
-    def update_config_from_dict(self,
-                                args: Dict[str, Any],
-                                prefix: str = "") -> None:
-        prefix_with_dot = f"{prefix}." if (prefix.strip() != "") else ""
-        update_config_from_args(self, args, prefix, pop_args=True)
-        update_config_from_args(self.vae_config,
-                                args,
-                                f"{prefix_with_dot}vae_config",
-                                pop_args=True)
-        update_config_from_args(self.dit_config,
-                                args,
-                                f"{prefix_with_dot}dit_config",
-                                pop_args=True)
+    enable_torch_compile: bool = False
 
     @classmethod
     def from_pretrained(cls, model_path: str) -> "PipelineConfig":
-        """
-        use the pipeline class setting from model_path to match the pipeline config
-        """
         from fastvideo.v1.configs.pipelines.registry import (
-            get_pipeline_config_cls_from_name)
-        pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)
-
-        return cast(PipelineConfig, pipeline_config_cls(model_path=model_path))
-
-    @classmethod
-    def from_kwargs(cls,
-                    kwargs: Dict[str, Any],
-                    config_cli_prefix: str = "") -> "PipelineConfig":
-        """
-        Load PipelineConfig from kwargs Dictionary.
-        kwargs: dictionary of kwargs
-        config_cli_prefix: prefix of CLI arguments for this PipelineConfig instance
-        """
-        from fastvideo.v1.configs.pipelines.registry import (
-            get_pipeline_config_cls_from_name)
-
-        prefix_with_dot = f"{config_cli_prefix}." if (config_cli_prefix.strip()
-                                                      != "") else ""
-        model_path: Optional[str] = kwargs.get(prefix_with_dot + 'model_path',
-                                               None) or kwargs.get('model_path')
-        pipeline_config_or_path: Optional[Union[str, PipelineConfig, Dict[
-            str, Any]]] = kwargs.get(prefix_with_dot + 'pipeline_config',
-                                     None) or kwargs.get('pipeline_config')
-        if model_path is None:
-            raise ValueError("model_path is required in kwargs")
-
-        # 1. Get the pipeline config class from the registry
-        pipeline_config_cls = get_pipeline_config_cls_from_name(model_path)
-
-        # 2. Instantiate PipelineConfig
-        if pipeline_config_cls is None:
+            get_pipeline_config_cls_for_name)
+        pipeline_config_cls = get_pipeline_config_cls_for_name(model_path)
+        if pipeline_config_cls is not None:
+            pipeline_config = pipeline_config_cls()
+        else:
             logger.warning(
-                "Couldn't find pipeline config for %s. Using the default pipeline config.",
+                "Couldn't find an optimal sampling param for %s. Using the default sampling param.",
                 model_path)
             pipeline_config = cls()
-        else:
-            pipeline_config = pipeline_config_cls()
 
-        # 3. Load PipelineConfig from a json file or a PipelineConfig object if provided
-        if isinstance(pipeline_config_or_path, str):
-            pipeline_config.load_from_json(pipeline_config_or_path)
-            kwargs[prefix_with_dot +
-                   'pipeline_config_path'] = pipeline_config_or_path
-        elif isinstance(pipeline_config_or_path, PipelineConfig):
-            pipeline_config = pipeline_config_or_path
-        elif isinstance(pipeline_config_or_path, dict):
-            pipeline_config.update_pipeline_config(pipeline_config_or_path)
-
-        # 4. Update PipelineConfig from CLI arguments if provided
-        kwargs[prefix_with_dot + 'model_path'] = model_path
-        pipeline_config.update_config_from_dict(kwargs, config_cli_prefix)
-        return pipeline_config
-
-    def check_pipeline_config(self) -> None:
-        if self.vae_sp and not self.vae_tiling:
-            raise ValueError(
-                "Currently enabling vae_sp requires enabling vae_tiling, please set --vae-tiling to True."
-            )
-
-        if len(self.text_encoder_configs) != len(self.text_encoder_precisions):
-            raise ValueError(
-                f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text encoder precisions ({len(self.text_encoder_precisions)})"
-            )
-
-        if len(self.text_encoder_configs) != len(self.preprocess_text_funcs):
-            raise ValueError(
-                f"Length of text encoder configs ({len(self.text_encoder_configs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
-            )
-
-        if len(self.preprocess_text_funcs) != len(self.postprocess_text_funcs):
-            raise ValueError(
-                f"Length of text postprocess functions ({len(self.postprocess_text_funcs)}) must be equal to length of text preprocessing functions ({len(self.preprocess_text_funcs)})"
-            )
+        return cast(PipelineConfig, pipeline_config)
 
     def dump_to_json(self, file_path: str):
         output_dict = shallow_asdict(self)

fastvideo/v1/configs/pipelines/hunyuan.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Callable, Tuple, TypedDict
 
@@ -80,7 +79,7 @@ class HunyuanConfig(PipelineConfig):
                      (llama_postprocess_text, clip_postprocess_text))
 
     # Precision for each component
-    dit_precision: str = "bf16"
+    precision: str = "bf16"
     vae_precision: str = "fp16"
     text_encoder_precisions: Tuple[str, ...] = field(
         default_factory=lambda: ("fp16", "fp16"))

fastvideo/v1/configs/pipelines/registry.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 """Registry for pipeline weight-specific configurations."""
 
 import os
@@ -19,7 +18,7 @@ from fastvideo.v1.utils import (maybe_download_model_index,
 logger = init_logger(__name__)
 
 # Registry maps specific model weights to their config classes
-PIPE_NAME_TO_CONFIG: Dict[str, Type[PipelineConfig]] = {
+WEIGHT_CONFIG_REGISTRY: Dict[str, Type[PipelineConfig]] = {
     "FastVideo/FastHunyuan-diffusers": FastHunyuanConfig,
     "hunyuanvideo-community/HunyuanVideo": HunyuanConfig,
     "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V480PConfig,
@@ -51,74 +50,37 @@ PIPELINE_FALLBACK_CONFIG: Dict[str, Type[PipelineConfig]] = {
 }
 
 
-def get_pipeline_config_cls_from_name(
-        pipeline_name_or_path: str) -> Type[PipelineConfig]:
-    """Get the appropriate configuration class for a given pipeline name or path.
+def get_pipeline_config_cls_for_name(
+        pipeline_name_or_path: str) -> Optional[type[PipelineConfig]]:
+    """Get the appropriate config class for specific pretrained weights."""
 
-    This function implements a multi-step lookup process to find the most suitable
-    configuration class for a given pipeline. It follows this order:
-    1. Exact match in the PIPE_NAME_TO_CONFIG
-    2. Partial match in the PIPE_NAME_TO_CONFIG
-    3. Fallback to class name in the model_index.json
-    4. else raise an error
+    if os.path.exists(pipeline_name_or_path):
+        config = verify_model_config_and_directory(pipeline_name_or_path)
+        logger.warning(
+            "FastVideo may not correctly identify the optimal config for this model, as the local directory may have been renamed."
+        )
+    else:
+        config = maybe_download_model_index(pipeline_name_or_path)
 
-    Args:
-        pipeline_name_or_path (str): The name or path of the pipeline. This can be:
-            - A registered model ID (e.g., "FastVideo/FastHunyuan-diffusers")
-            - A local path to a model directory
-            - A model ID that will be downloaded
-
-    Returns:
-        Type[PipelineConfig]: The configuration class that best matches the pipeline.
-            This will be one of:
-            - A specific weight configuration class if an exact match is found
-            - A fallback configuration class based on the pipeline architecture
-            - The base PipelineConfig class if no matches are found
-
-    Note:
-        - For local paths, the function will verify the model configuration
-        - For remote models, it will attempt to download the model index
-        - Warning messages are logged when falling back to less specific configurations
-    """
-
-    pipeline_config_cls: Optional[Type[PipelineConfig]] = None
+    pipeline_name = config["_class_name"]
 
     # First try exact match for specific weights
-    if pipeline_name_or_path in PIPE_NAME_TO_CONFIG:
-        pipeline_config_cls = PIPE_NAME_TO_CONFIG[pipeline_name_or_path]
+    if pipeline_name_or_path in WEIGHT_CONFIG_REGISTRY:
+        return WEIGHT_CONFIG_REGISTRY[pipeline_name_or_path]
 
     # Try partial matches (for local paths that might include the weight ID)
-    for registered_id, config_class in PIPE_NAME_TO_CONFIG.items():
+    for registered_id, config_class in WEIGHT_CONFIG_REGISTRY.items():
         if registered_id in pipeline_name_or_path:
-            pipeline_config_cls = config_class
-            break
+            return config_class
 
     # If no match, try to use the fallback config
-    if pipeline_config_cls is None:
-        if os.path.exists(pipeline_name_or_path):
-            config = verify_model_config_and_directory(pipeline_name_or_path)
-        else:
-            config = maybe_download_model_index(pipeline_name_or_path)
-        logger.warning(
-            "Trying to use the config from the model_index.json. FastVideo may not correctly identify the optimal config for this model in this situation."
-        )
+    fallback_config = None
+    # Try to determine pipeline architecture for fallback
+    for pipeline_type, detector in PIPELINE_DETECTOR.items():
+        if detector(pipeline_name.lower()):
+            fallback_config = PIPELINE_FALLBACK_CONFIG.get(pipeline_type)
+            break
 
-        pipeline_name = config["_class_name"]
-        # Try to determine pipeline architecture for fallback
-        for pipeline_type, detector in PIPELINE_DETECTOR.items():
-            if detector(pipeline_name.lower()):
-                pipeline_config_cls = PIPELINE_FALLBACK_CONFIG.get(
-                    pipeline_type)
-                break
-
-        if pipeline_config_cls is not None:
-            logger.warning(
-                "No match found for pipeline %s, using fallback config %s.",
-                pipeline_name_or_path, pipeline_config_cls)
-
-    if pipeline_config_cls is None:
-        raise ValueError(
-            f"No match found for pipeline {pipeline_name_or_path}, please check the pipeline name or path."
-        )
-
-    return pipeline_config_cls
+    logger.warning("No match found for pipeline %s, using fallback config %s.",
+                   pipeline_name_or_path, fallback_config)
+    return fallback_config

fastvideo/v1/configs/pipelines/stepvideo.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 
 from fastvideo.v1.configs.models import DiTConfig, VAEConfig

fastvideo/v1/configs/pipelines/wan.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 from typing import Callable, Tuple
 

fastvideo/v1/configs/sample/base.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Union
 

fastvideo/v1/configs/sample/hunyuan.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 
 from fastvideo.v1.configs.sample.base import SamplingParam

fastvideo/v1/configs/sample/registry.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 import os
 from typing import Any, Callable, Dict, Optional
 

fastvideo/v1/configs/sample/stepvideo.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass
 
 from fastvideo.v1.configs.sample.base import SamplingParam

fastvideo/v1/configs/sample/teacache.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 
 from fastvideo.v1.configs.sample.base import CacheParams

fastvideo/v1/configs/sample/wan.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 from dataclasses import dataclass, field
 
 from fastvideo.v1.configs.sample.base import SamplingParam

fastvideo/v1/configs/utils.py

@@ -1,45 +0,0 @@
-from typing import Any, Dict
-
-
-def update_config_from_args(config: Any,
-                            args_dict: Dict[str, Any],
-                            prefix: str = "",
-                            pop_args: bool = False) -> None:
-    """
-    Update configuration object from arguments dictionary.
-    
-    Args:
-        config: The configuration object to update
-        args_dict: Dictionary containing arguments
-        prefix: Prefix for the configuration parameters in the args_dict.
-               If None, assumes direct attribute mapping without prefix.
-    """
-    # Handle top-level attributes (no prefix)
-    args_not_to_remove = [
-        'model_path',
-    ]
-    args_to_remove = []
-    if prefix.strip() == "":
-        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)
-                if pop_args:
-                    args_to_remove.append(key)
-    else:
-        # Handle nested attributes with prefix
-        prefix_with_dot = f"{prefix}."
-        for key, value in args_dict.items():
-            if key.startswith(prefix_with_dot) and value is not None:
-                attr_name = key[len(prefix_with_dot):]
-                if hasattr(config, attr_name):
-                    setattr(config, attr_name, value)
-                if pop_args:
-                    args_to_remove.append(key)
-
-    if pop_args:
-        for key in args_to_remove:
-            if key not in args_not_to_remove:
-                args_dict.pop(key)

fastvideo/v1/dataset/__init__.py

@@ -1,17 +1,15 @@
-# SPDX-License-Identifier: Apache-2.0
+import os
+
 from torchvision import transforms
 from torchvision.transforms import Lambda
+from transformers import AutoTokenizer
 
-from fastvideo.v1.dataset.parquet_dataset_map_style import (
-    build_parquet_map_style_dataloader)
-from fastvideo.v1.dataset.preprocessing_datasets import (
-    VideoCaptionMergedDataset)
+from fastvideo.v1.dataset.t2v_datasets import T2V_dataset
 from fastvideo.v1.dataset.transform import (CenterCropResizeVideo, Normalize255,
                                             TemporalRandomCrop)
-from fastvideo.v1.dataset.validation_dataset import ValidationDataset
 
 
-def getdataset(args) -> VideoCaptionMergedDataset:
+def getdataset(args, start_idx=0) -> T2V_dataset:
     temporal_sample = TemporalRandomCrop(args.num_frames)  # 16 x
     norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
     resize_topcrop = [
@@ -29,14 +27,15 @@ def getdataset(args) -> VideoCaptionMergedDataset:
         *resize_topcrop,
         norm_fun,
     ])
-    return VideoCaptionMergedDataset(data_merge_path=args.data_merge_path,
-                                     args=args,
-                                     transform=transform,
-                                     temporal_sample=temporal_sample,
-                                     transform_topcrop=transform_topcrop)
+    tokenizer_path = os.path.join(args.model_path, "tokenizer")
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
+                                              cache_dir=args.cache_dir)
+    if args.dataset == "t2v":
+        return T2V_dataset(args,
+                           transform=transform,
+                           temporal_sample=temporal_sample,
+                           tokenizer=tokenizer,
+                           transform_topcrop=transform_topcrop,
+                           start_idx=start_idx)
 
-
-__all__ = [
-    "build_parquet_map_style_dataloader", "ValidationDataset",
-    "VideoCaptionMergedDataset"
-]
+    raise NotImplementedError(args.dataset)

fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py

@@ -1,185 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import argparse
-import os
-import pathlib
-import time
-
-import torch.distributed as dist
-import torch.distributed.checkpoint as dist_cp
-
-from fastvideo.v1.dataset.parquet_dataset_iterable_style import (
-    build_parquet_iterable_style_dataloader)
-from fastvideo.v1.distributed import get_world_rank
-from fastvideo.v1.distributed.parallel_state import (
-    cleanup_dist_env_and_memory, get_torch_device,
-    maybe_init_distributed_environment_and_model_parallel)
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-def main() -> None:
-    parser = argparse.ArgumentParser(
-        description="Benchmark parquet iterable style dataset loading speed")
-    parser.add_argument(
-        "--path",
-        type=str,
-        help="Path to parquet dataset",
-    )
-    parser.add_argument("--batch_size",
-                        type=int,
-                        default=4,
-                        help="Batch size for DataLoader")
-    parser.add_argument("--num_data_workers",
-                        type=int,
-                        help="Number of DataLoader workers")
-    parser.add_argument("--num_epoch",
-                        type=int,
-                        default=2,
-                        help="Number of epoches to benchmark")
-    parser.add_argument("--verify_resume",
-                        action="store_true",
-                        help="Verify resume")
-    parser.add_argument(
-        "--num_batches_per_epoch",
-        type=int,
-        default=1000,
-        help="Number of batches to benchmark",
-    )
-    parser.add_argument('--checkpoint_path',
-                        type=str,
-                        default='dataloader_checkpoint',
-                        help='Path to save/load checkpoint')
-    '''
-    example launch command:
-    torchrun --nproc_per_node=1 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 2 --num_data_workers 2 --num_epoch 2 --num_batches_per_epoch 2 --verify_resume
-    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 2 --num_data_workers 1 --num_epoch 2 --num_batches_per_epoch 5 --verify_resume
-    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_iterable_style.py --path /mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents/ --batch_size 2 --num_data_workers 4 --num_epoch 2 --num_batches_per_epoch 100 
-    '''
-    args = parser.parse_args()
-    world_size = int(os.environ.get("WORLD_SIZE", 1))
-    maybe_init_distributed_environment_and_model_parallel(
-        tp_size=(world_size + 1) // 2, sp_size=(world_size + 1) // 2)
-    logger.info("Initialized distributed environment with world_size=%d",
-                world_size)
-
-    # Create DataLoader with proper settings
-    dataset, dataloader = build_parquet_iterable_style_dataloader(
-        args.path, args.batch_size, args.num_data_workers)
-    logger.info("Initialized dataloader")
-
-    if args.verify_resume:
-        # First pass - record latent sums
-        first_pass_sums = []
-        for i, (latents, embeddings, masks,
-                caption_text) in enumerate(dataloader):
-            latent_sum = latents.sum().item()
-            first_pass_sums.append(latent_sum)
-            logger.info("Batch %d latent sum: %f", i, latent_sum)
-            if i >= args.num_batches_per_epoch - 1:
-                break
-
-        # Save dataloader state using distributed checkpoint
-        checkpoint_dir = pathlib.Path(args.checkpoint_path)
-        logger.info("Rank %d: Saving dataloader state to %s", get_world_rank(),
-                    checkpoint_dir)
-        states = {"dataloader": dataloader}
-
-        begin_time = time.monotonic()
-        dist_cp.save(states, checkpoint_id=checkpoint_dir.as_posix())
-        end_time = time.monotonic()
-
-        logger.info("Rank %d: Saved checkpoint in %.2f seconds",
-                    get_world_rank(), end_time - begin_time)
-
-        # Make sure all processes wait for checkpoint to be saved
-        if world_size > 1:
-            dist.barrier()
-
-        # Recreate dataloader and load state
-        dataset, dataloader = build_parquet_iterable_style_dataloader(
-            args.path, args.batch_size, args.num_data_workers)
-        load_states = {"dataloader": dataloader}
-        dist_cp.load(load_states, checkpoint_id=checkpoint_dir.as_posix())
-        logger.info("Rank %d: Loaded dataloader state from %s",
-                    get_world_rank(), checkpoint_dir)
-
-        # Second pass - verify latent sums match
-        for i, (latents, embeddings, masks) in enumerate(dataloader):
-            latent_sum = latents.sum().item()
-            first_pass_sums.append(latent_sum)
-            logger.info("Batch %d latent sum: %f",
-                        i + args.num_batches_per_epoch, latent_sum)
-            if i >= args.num_batches_per_epoch - 1:
-                break
-
-        dataset, dataloader = build_parquet_iterable_style_dataloader(
-            args.path, args.batch_size, args.num_data_workers)
-        # Second pass - verify latent sums match
-        second_pass_sums = []
-        for i, (latents, embeddings, masks,
-                caption_text) in enumerate(dataloader):
-            latent_sum = latents.sum().item()
-            second_pass_sums.append(latent_sum)
-            logger.info("Batch %d latent sum: %f (should match first pass: %f)",
-                        i, latent_sum, first_pass_sums[i])
-            if i >= args.num_batches_per_epoch * 2 - 1:
-                break
-
-        # Verify all sums match
-        if all(
-                abs(a - b) < 1e-6
-                for a, b in zip(first_pass_sums, second_pass_sums)):
-            logger.info(
-                "All latent sums match between passes - resume verification successful!"
-            )
-        else:
-            raise ValueError(
-                "Latent sums do not match between passes - resume verification failed!"
-            )
-
-    start_time = time.time()
-    total_samples = 0
-    total_batches = 0
-    for _ in range(args.num_epoch):
-        for i, (latents, embeddings, masks,
-                caption_text) in enumerate(dataloader):
-            if i >= args.num_batches_per_epoch:
-                break
-
-            # Move data to device
-            latents = latents.to(get_torch_device())
-            embeddings = embeddings.to(get_torch_device())
-
-            # Calculate actual batch size
-            batch_size = latents.size(0)
-            total_samples += batch_size
-            total_batches += 1
-
-            # Print progress only from rank 0
-            if get_world_rank() == 0 and (i + 1) % 10 == 0:
-                elapsed = time.time() - start_time
-                samples_per_sec = total_samples / elapsed
-                logger.info("Batch %d/%d, Speed: %.2f samples/sec", i + 1,
-                            args.num_batches_per_epoch, samples_per_sec)
-
-    # Final statistics
-    if world_size > 1:
-        dist.barrier()
-
-    if get_world_rank() == 0:
-        elapsed = time.time() - start_time
-        samples_per_sec = total_samples / elapsed
-
-        logger.info("\nBenchmark Results:")
-        logger.info("Total time: %.2f seconds", elapsed)
-        logger.info("Total samples: %d", total_samples)
-        logger.info("Average speed: %.2f samples/sec", samples_per_sec)
-        logger.info("Time per batch: %.2f ms", elapsed / total_batches * 1000)
-
-
-if __name__ == "__main__":
-    try:
-        main()
-    finally:
-        cleanup_dist_env_and_memory()

fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_map_style.py

@@ -1,187 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import argparse
-import os
-import pathlib
-import time
-
-import torch
-import torch.distributed as dist
-import torch.distributed.checkpoint as dist_cp
-
-from fastvideo.v1.dataset.parquet_dataset_map_style import (
-    build_parquet_map_style_dataloader)
-from fastvideo.v1.distributed import get_world_rank
-from fastvideo.v1.distributed.parallel_state import (
-    cleanup_dist_env_and_memory, get_torch_device,
-    maybe_init_distributed_environment_and_model_parallel)
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-def main() -> None:
-    torch.multiprocessing.set_start_method("spawn", force=True)
-    parser = argparse.ArgumentParser(
-        description="Benchmark parquet map style dataset loading speed")
-    parser.add_argument(
-        "--path",
-        type=str,
-        help="Path to parquet dataset",
-    )
-    parser.add_argument("--batch_size",
-                        type=int,
-                        default=4,
-                        help="Batch size for DataLoader")
-    parser.add_argument("--num_data_workers",
-                        type=int,
-                        help="Number of DataLoader workers")
-    parser.add_argument("--num_epoch",
-                        type=int,
-                        default=2,
-                        help="Number of epoches to benchmark")
-    parser.add_argument("--verify_resume",
-                        action="store_true",
-                        help="Verify resume")
-    parser.add_argument(
-        "--num_batches_per_epoch",
-        type=int,
-        default=1000,
-        help="Number of batches to benchmark",
-    )
-    parser.add_argument('--checkpoint_path',
-                        type=str,
-                        default='dataloader_checkpoint',
-                        help='Path to save/load checkpoint')
-    '''
-    example launch command:
-    torchrun --nproc_per_node=1 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_map_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 4 --num_data_workers 1 --num_epoch 2 --num_batches_per_epoch 3 --verify_resume
-    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_map_style.py --path data/crush-smol/latents/combined_parquet_dataset --batch_size 2 --num_data_workers 1 --num_epoch 2 --num_batches_per_epoch 5 --verify_resume
-    torchrun --nproc_per_node=8 --master_port=12358 fastvideo/v1/dataset/benchmarks/benchmark_parquet_dataset_map_style.py --path /mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn/latents/ --batch_size 2 --num_data_workers 4 --num_epoch 2 --num_batches_per_epoch 100 
-    '''
-    args = parser.parse_args()
-    world_size = int(os.environ.get("WORLD_SIZE", 1))
-    maybe_init_distributed_environment_and_model_parallel(
-        tp_size=(world_size + 1) // 2, sp_size=(world_size + 1) // 2)
-    logger.info("Initialized distributed environment with world_size=%d",
-                world_size)
-
-    # Create DataLoader with proper settings
-    dataset, dataloader = build_parquet_map_style_dataloader(
-        args.path, args.batch_size, args.num_data_workers)
-    logger.info("Initialized dataloader with %d batches", len(dataloader))
-
-    if args.verify_resume:
-        # First pass - record latent sums
-        first_pass_sums = []
-        for i, (latents, embeddings, masks,
-                caption_text) in enumerate(dataloader):
-            latent_sum = latents.sum().item()
-            first_pass_sums.append(latent_sum)
-            logger.info("Batch %d latent sum: %f", i, latent_sum)
-            if i >= args.num_batches_per_epoch - 1:
-                break
-
-        # Save dataloader state using distributed checkpoint
-        checkpoint_dir = pathlib.Path(args.checkpoint_path)
-        logger.info("Rank %d: Saving dataloader state to %s", get_world_rank(),
-                    checkpoint_dir)
-        states = {"dataloader": dataloader}
-
-        begin_time = time.monotonic()
-        dist_cp.save(states, checkpoint_id=checkpoint_dir.as_posix())
-        end_time = time.monotonic()
-
-        logger.info("Rank %d: Saved checkpoint in %.2f seconds",
-                    get_world_rank(), end_time - begin_time)
-
-        # Make sure all processes wait for checkpoint to be saved
-        if world_size > 1:
-            dist.barrier()
-
-        # Recreate dataloader and load state
-        dataset, dataloader = build_parquet_map_style_dataloader(
-            args.path, args.batch_size, args.num_data_workers)
-        load_states = {"dataloader": dataloader}
-        dist_cp.load(load_states, checkpoint_id=checkpoint_dir.as_posix())
-        logger.info("Rank %d: Loaded dataloader state from %s",
-                    get_world_rank(), checkpoint_dir)
-
-        for i, (latents, embeddings, masks,
-                caption_text) in enumerate(dataloader):
-            latent_sum = latents.sum().item()
-            first_pass_sums.append(latent_sum)
-            logger.info("Batch %d latent sum: %f",
-                        i + args.num_batches_per_epoch, latent_sum)
-            if i >= args.num_batches_per_epoch - 1:
-                break
-
-        dataset, dataloader = build_parquet_map_style_dataloader(
-            args.path, args.batch_size, args.num_data_workers)
-
-        # Second pass - verify latent sums match
-        second_pass_sums = []
-        for i, (latents, embeddings, masks) in enumerate(dataloader):
-            latent_sum = latents.sum().item()
-            second_pass_sums.append(latent_sum)
-            logger.info("Batch %d latent sum: %f (should match first pass: %f)",
-                        i, latent_sum, first_pass_sums[i])
-            if i >= args.num_batches_per_epoch * 2 - 1:
-                break
-
-        # Verify all sums match
-        if all(
-                abs(a - b) < 1e-6
-                for a, b in zip(first_pass_sums, second_pass_sums)):
-            logger.info(
-                "All latent sums match between passes - resume verification successful!"
-            )
-        else:
-            raise ValueError(
-                "Latent sums do not match between passes - resume verification failed!"
-            )
-
-    start_time = time.time()
-    total_samples = 0
-    total_batches = 0
-    for _ in range(args.num_epoch):
-        for i, (latents, embeddings, masks,
-                caption_text) in enumerate(dataloader):
-            if i >= args.num_batches_per_epoch:
-                break
-
-            # Move data to device
-            latents = latents.to(get_torch_device())
-            embeddings = embeddings.to(get_torch_device())
-
-            # Calculate actual batch size
-            batch_size = latents.size(0)
-            total_samples += batch_size
-            total_batches += 1
-
-            # Print progress only from rank 0
-            if get_world_rank() == 0 and (i + 1) % 10 == 0:
-                elapsed = time.time() - start_time
-                samples_per_sec = total_samples / elapsed
-                logger.info("Batch %d/%d, Speed: %.2f samples/sec", i + 1,
-                            args.num_batches_per_epoch, samples_per_sec)
-
-    # Final statistics
-    if world_size > 1:
-        dist.barrier()
-
-    if get_world_rank() == 0:
-        elapsed = time.time() - start_time
-        samples_per_sec = total_samples / elapsed
-
-        logger.info("\nBenchmark Results:")
-        logger.info("Total time: %.2f seconds", elapsed)
-        logger.info("Total samples: %d", total_samples)
-        logger.info("Average speed: %.2f samples/sec", samples_per_sec)
-        logger.info("Time per batch: %.2f ms", elapsed / total_batches * 1000)
-
-
-if __name__ == "__main__":
-    try:
-        main()
-    finally:
-        cleanup_dist_env_and_memory()

fastvideo/v1/dataset/dataloader/schema.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 # schema.py
 """
 Unified data schema and format for saving and loading image/video data after
@@ -26,47 +25,15 @@ pyarrow_schema_i2v = pa.schema([
     pa.field("text_embedding_shape", pa.list_(pa.int64())),
     # e.g., 'bfloat16' or 'float32'
     pa.field("text_embedding_dtype", pa.string()),
+    pa.field("text_attention_mask_bytes", pa.binary()),
+    # e.g., [SeqLen]
+    pa.field("text_attention_mask_shape", pa.list_(pa.int64())),
+    # e.g., 'bool' or 'int8'
+    pa.field("text_attention_mask_dtype", pa.string()),
     #I2V
     pa.field("clip_feature_bytes", pa.binary()),
     pa.field("clip_feature_shape", pa.list_(pa.int64())),
     pa.field("clip_feature_dtype", pa.string()),
-    pa.field("first_frame_latent_bytes", pa.binary()),
-    pa.field("first_frame_latent_shape", pa.list_(pa.int64())),
-    pa.field("first_frame_latent_dtype", pa.string()),
-    # I2V Validation
-    pa.field("pil_image_bytes", pa.binary()),
-    pa.field("pil_image_shape", pa.list_(pa.int64())),
-    pa.field("pil_image_dtype", pa.string()),
-    # --- Metadata ---
-    pa.field("file_name", pa.string()),
-    pa.field("caption", pa.string()),
-    pa.field("media_type", pa.string()),  # 'image' or 'video'
-    pa.field("width", pa.int64()),
-    pa.field("height", pa.int64()),
-    # -- Video-specific (can be null/default for images) ---
-    # Number of frames processed (e.g., 1 for image, N for video)
-    pa.field("num_frames", pa.int64()),
-    pa.field("duration_sec", pa.float64()),
-    pa.field("fps", pa.float64()),
-])
-
-pyarrow_schema_i2v_validation = pa.schema([
-    pa.field("id", pa.string()),
-    # --- Text encoder output tensor ---
-    # Tensors are stored as raw bytes with shape and dtype info for loading
-    pa.field("text_embedding_bytes", pa.binary()),
-    # e.g., [SeqLen, Dim]
-    pa.field("text_embedding_shape", pa.list_(pa.int64())),
-    # e.g., 'bfloat16' or 'float32'
-    pa.field("text_embedding_dtype", pa.string()),
-    #I2V
-    pa.field("clip_feature_bytes", pa.binary()),
-    pa.field("clip_feature_shape", pa.list_(pa.int64())),
-    pa.field("clip_feature_dtype", pa.string()),
-    # I2V Validation
-    pa.field("pil_image_bytes", pa.binary()),
-    pa.field("pil_image_shape", pa.list_(pa.int64())),
-    pa.field("pil_image_dtype", pa.string()),
     # --- Metadata ---
     pa.field("file_name", pa.string()),
     pa.field("caption", pa.string()),
@@ -96,6 +63,11 @@ pyarrow_schema_t2v = pa.schema([
     pa.field("text_embedding_shape", pa.list_(pa.int64())),
     # e.g., 'bfloat16' or 'float32'
     pa.field("text_embedding_dtype", pa.string()),
+    pa.field("text_attention_mask_bytes", pa.binary()),
+    # e.g., [SeqLen]
+    pa.field("text_attention_mask_shape", pa.list_(pa.int64())),
+    # e.g., 'bool' or 'int8'
+    pa.field("text_attention_mask_dtype", pa.string()),
     # --- Metadata ---
     pa.field("file_name", pa.string()),
     pa.field("caption", pa.string()),
@@ -107,26 +79,4 @@ pyarrow_schema_t2v = pa.schema([
     pa.field("num_frames", pa.int64()),
     pa.field("duration_sec", pa.float64()),
     pa.field("fps", pa.float64()),
-])
-
-pyarrow_schema_t2v_validation = pa.schema([
-    pa.field("id", pa.string()),
-    # --- Text encoder output tensor ---
-    # Tensors are stored as raw bytes with shape and dtype info for loading
-    pa.field("text_embedding_bytes", pa.binary()),
-    # e.g., [SeqLen, Dim]
-    pa.field("text_embedding_shape", pa.list_(pa.int64())),
-    # e.g., 'bfloat16' or 'float32'
-    pa.field("text_embedding_dtype", pa.string()),
-    # --- Metadata ---
-    pa.field("file_name", pa.string()),
-    pa.field("caption", pa.string()),
-    pa.field("media_type", pa.string()),  # 'image' or 'video'
-    pa.field("width", pa.int64()),
-    pa.field("height", pa.int64()),
-    # -- Video-specific (can be null/default for images) ---
-    # Number of frames processed (e.g., 1 for image, N for video)
-    pa.field("num_frames", pa.int64()),
-    pa.field("duration_sec", pa.float64()),
-    pa.field("fps", pa.float64()),
-])
+])

fastvideo/v1/dataset/json_prepare.py

@@ -0,0 +1,136 @@
+import argparse
+import json
+import os
+import time
+from multiprocessing import Pool, cpu_count
+from pathlib import Path
+
+import torchvision
+from tqdm import tqdm
+
+
+def get_video_info(video_path):
+    """Get video information using torchvision."""
+    # Read video tensor (T, C, H, W)
+    video_tensor, _, info = torchvision.io.read_video(str(video_path),
+                                                      output_format="TCHW",
+                                                      pts_unit="sec")
+
+    num_frames = video_tensor.shape[0]
+    height = video_tensor.shape[2]
+    width = video_tensor.shape[3]
+    fps = info.get("video_fps", 0)
+    duration = num_frames / fps if fps > 0 else 0
+
+    # Extract name
+    _, _, videos_dir, video_name = str(video_path).split("/")
+
+    return {
+        "path": str(video_name),
+        "resolution": {
+            "width": width,
+            "height": height
+        },
+        "size": os.path.getsize(video_path),
+        "fps": fps,
+        "duration": duration,
+        "num_frames": num_frames
+    }
+
+
+def prepare_dataset_json(folder_path,
+                         output_name="videos2caption.json",
+                         num_workers=None) -> None:
+    """Prepare dataset information from a folder containing videos and prompt.txt."""
+    folder_path = Path(folder_path)
+
+    # Read prompt file
+    prompt_file = folder_path / "prompt.txt"
+    if not prompt_file.exists():
+        raise FileNotFoundError(f"prompt.txt not found in {folder_path}")
+
+    with open(prompt_file) as f:
+        prompts = [line.strip() for line in f.readlines() if line.strip()]
+
+    # Read videos file
+    videos_file = folder_path / "videos.txt"
+    if not videos_file.exists():
+        raise FileNotFoundError(f"videos.txt not found in {folder_path}")
+
+    with open(videos_file) as f:
+        video_paths = [line.strip() for line in f.readlines() if line.strip()]
+
+    if len(prompts) != len(video_paths):
+        raise ValueError(
+            f"Number of prompts ({len(prompts)}) does not match number of videos ({len(video_paths)})"
+        )
+
+    # Prepare arguments for multiprocessing
+    process_args = [folder_path / video_path for video_path in video_paths]
+
+    # Determine number of workers
+    if num_workers is None:
+        num_workers = max(1, cpu_count() - 1)  # Leave one CPU free
+
+    # Process videos in parallel
+    start_time = time.time()
+    with Pool(num_workers) as pool:
+        results = list(
+            tqdm(pool.imap(get_video_info, process_args),
+                 total=len(process_args),
+                 desc="Processing videos",
+                 unit="video"))
+
+    # Combine results with prompts
+    dataset_info = []
+    for result, prompt in zip(results, prompts):
+        result["cap"] = [prompt]
+        dataset_info.append(result)
+
+    # Calculate total processing time
+    total_time = time.time() - start_time
+    total_videos = len(dataset_info)
+    avg_time_per_video = total_time / total_videos if total_videos > 0 else 0
+
+    print("\nProcessing completed:")
+    print(f"Total videos processed: {total_videos}")
+    print(f"Total time: {total_time:.2f} seconds")
+    print(f"Average time per video: {avg_time_per_video:.2f} seconds")
+
+    # Save to JSON file
+    output_file = folder_path / output_name
+    with open(output_file, 'w') as f:
+        json.dump(dataset_info, f, indent=2)
+
+    # Create merge.txt
+    merge_file = folder_path / "merge.txt"
+    with open(merge_file, 'w') as f:
+        f.write(f"{folder_path}/videos,{output_file}\n")
+
+    print(f"Dataset information saved to {output_file}")
+    print(f"Merge file created at {merge_file}")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description='Prepare video dataset information in JSON format')
+    parser.add_argument(
+        '--folder',
+        type=str,
+        required=True,
+        help='Path to the folder containing videos and prompt.txt')
+    parser.add_argument(
+        '--output',
+        type=str,
+        default='videos2caption.json',
+        help='Name of the output JSON file (default: videos2caption.json)')
+    parser.add_argument('--workers',
+                        type=int,
+                        default=32,
+                        help='Number of worker processes (default: 16)')
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    prepare_dataset_json(args.folder, args.output, args.workers)

fastvideo/v1/dataset/latent_datasets.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 import json
 import os
 import random

fastvideo/v1/dataset/parquet_dataset_iterable_style.py

@@ -1,278 +0,0 @@
-import os
-import pickle
-import random
-from typing import Dict, List, Tuple
-
-import numpy as np
-import pyarrow as pa
-import pyarrow.parquet as pq
-import torch
-import tqdm
-from torch.utils.data import IterableDataset, get_worker_info
-from torchdata.stateful_dataloader import StatefulDataLoader
-
-from fastvideo.v1.dataset.utils import collate_latents_embs_masks
-from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
-                                      get_world_size)
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class BatchIterator:
-    # TODO: Implement state_dict and load_state_dict to support resume.
-    def __init__(self, files, batch_size, text_padding_length, keys,
-                 worker_num_samples, read_batch_size):
-        self.files = files
-        self.batch_size = batch_size
-        self.text_padding_length = text_padding_length
-        self.keys = keys
-        self.worker_num_samples = worker_num_samples
-        self.processed_samples = 0
-        self.buffer = []
-        self.read_batch_size = read_batch_size
-
-    def __iter__(self):
-        for file in self.files:
-            if self.processed_samples >= self.worker_num_samples:
-                return
-
-            reader = pq.ParquetFile(file)
-            for batch in reader.iter_batches(batch_size=self.read_batch_size):
-                if self.processed_samples >= self.worker_num_samples:
-                    return
-
-                self.buffer.extend(batch.to_pylist())
-
-                while len(self.buffer) >= self.batch_size:
-                    if self.processed_samples >= self.worker_num_samples:
-                        return
-
-                    batch_to_process = self.buffer[:self.batch_size]
-                    self.buffer = self.buffer[self.batch_size:]
-
-                    all_latents, all_embs, all_masks, caption_text = collate_latents_embs_masks(
-                        batch_to_process, self.text_padding_length, self.keys)
-                    self.processed_samples += self.batch_size
-                    yield all_latents, all_embs, all_masks, caption_text
-
-
-class LatentsParquetIterStyleDataset(IterableDataset):
-    """Efficient loader for video-text data from a directory of Parquet files."""
-
-    # Modify this in the future if we want to add more keys, for example, in image to video.
-    keys = [("vae_latent", "latent"), ("text_embedding")]
-
-    def __init__(self,
-                 path: str,
-                 batch_size: int = 1024,
-                 cfg_rate: float = 0.1,
-                 num_workers: int = 1,
-                 drop_last: bool = True,
-                 text_padding_length: int = 512,
-                 seed: int = 42,
-                 read_batch_size: int = 32,
-                 parquet_schema: pa.Schema = None):
-        super().__init__()
-        self.path = str(path)
-        self.batch_size = batch_size
-        self.parquet_schema = parquet_schema
-        self.cfg_rate = cfg_rate
-        self.text_padding_length = text_padding_length
-        self.seed = seed
-        self.read_batch_size = read_batch_size
-        # Get distributed training info
-        self.global_rank = get_world_rank()
-        self.world_size = get_world_size()
-        self.sp_world_size = get_sp_world_size()
-        self.num_sp_groups = self.world_size // self.sp_world_size
-        num_workers = 1 if num_workers == 0 else num_workers
-        # Get sharding info
-        shard_parquet_files, shard_total_samples, shard_parquet_lengths = shard_parquet_files_across_sp_groups_and_workers(
-            self.path, self.num_sp_groups, num_workers, seed)
-
-        if drop_last:
-            self.worker_num_samples = min(
-                shard_total_samples) // batch_size * batch_size
-            # Assign files to current rank's SP group
-            ith_sp_group = self.global_rank // self.sp_world_size
-            self.sp_group_parquet_files = shard_parquet_files[ith_sp_group::self
-                                                              .num_sp_groups]
-            self.sp_group_parquet_lengths = shard_parquet_lengths[
-                ith_sp_group::self.num_sp_groups]
-            self.sp_group_num_samples = shard_total_samples[ith_sp_group::self.
-                                                            num_sp_groups]
-            logger.info(
-                "In total %d parquet files, %d samples, after sharding we retain %d samples due to drop_last",
-                sum([len(shard) for shard in shard_parquet_files]),
-                sum(shard_total_samples),
-                self.worker_num_samples * self.num_sp_groups * num_workers)
-        else:
-            raise ValueError("drop_last must be True")
-        logger.info("Each dataloader worker will load %d samples",
-                    self.worker_num_samples)
-
-    def __iter__(self):
-        worker_info = get_worker_info()
-        worker_id = worker_info.id if worker_info is not None else 1
-
-        worker_files = self.sp_group_parquet_files[worker_id]
-
-        batch_iterator = BatchIterator(
-            files=worker_files,
-            batch_size=self.batch_size,
-            text_padding_length=self.text_padding_length,
-            keys=self.keys,
-            worker_num_samples=self.worker_num_samples,
-            read_batch_size=self.read_batch_size)  # type: ignore
-
-        yield from batch_iterator
-
-        if batch_iterator.processed_samples != self.worker_num_samples:
-            raise ValueError(
-                "Rank %d, Worker %d: Not enough samples to process, this should not happen",
-                self.global_rank, worker_id)
-
-
-def shard_parquet_files_across_sp_groups_and_workers(
-    path: str,
-    num_sp_groups: int,
-    num_workers: int,
-    seed: int = 42,
-) -> Tuple[List[List[str]], List[int], List[Dict[str, int]]]:
-    """
-    Shard parquet files across SP groups and workers in a balanced way.
-    
-    Args:
-        path: Directory containing parquet files
-        num_sp_groups: Number of SP groups to shard across
-        num_workers: Number of workers per SP group
-        seed: Random seed for shuffling
-        
-    Returns:
-        Tuple containing:
-        - List of lists of parquet files for each shard
-        - List of total samples per shard
-        - List of dictionaries mapping file paths to their lengths
-    """
-    # Check if sharding plan already exists
-    sharding_info_dir = os.path.join(
-        path, f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers")
-    if os.path.exists(sharding_info_dir):
-        logger.info("Sharding plan already exists")
-        logger.info("Loading sharding plan from %s", sharding_info_dir)
-        try:
-            with open(
-                    os.path.join(sharding_info_dir, "shard_parquet_files.pkl"),
-                    "rb") as f:
-                shard_parquet_files = pickle.load(f)
-            with open(
-                    os.path.join(sharding_info_dir, "shard_total_samples.pkl"),
-                    "rb") as f:
-                shard_total_samples = pickle.load(f)
-            with open(
-                    os.path.join(sharding_info_dir,
-                                 "shard_parquet_lengths.pkl"), "rb") as f:
-                shard_parquet_lengths = pickle.load(f)
-            return shard_parquet_files, shard_total_samples, shard_parquet_lengths
-        except Exception as e:
-            logger.error("Error loading sharding plan: %s", str(e))
-            logger.info("Falling back to creating new sharding plan")
-
-    if get_world_rank() == 0:
-        logger.info("Scanning for parquet files in %s", path)
-
-        # Find all parquet files
-        parquet_files = []
-
-        for root, _, files in os.walk(path):
-            for file in files:
-                if file.endswith('.parquet'):
-                    parquet_files.append(os.path.join(root, file))
-
-        if not parquet_files:
-            raise ValueError("No parquet files found in %s", path)
-
-        # Calculate file lengths efficiently using a single pass
-        logger.info("Calculating file lengths...")
-        lengths = []
-        for file in tqdm.tqdm(parquet_files, desc="Reading parquet files"):
-            lengths.append(pq.ParquetFile(file).metadata.num_rows)
-
-        total_samples = sum(lengths)
-        logger.info("Found %d files with %d total samples", len(parquet_files),
-                    total_samples)
-
-        # Sort files by length for better balancing
-        sorted_indices = np.argsort(lengths)
-        sorted_files = [parquet_files[i] for i in sorted_indices]
-        sorted_lengths = [lengths[i] for i in sorted_indices]
-
-        # Create shards
-        num_shards = num_sp_groups * num_workers
-        shard_parquet_files = [[] for _ in range(num_shards)]
-        shard_total_samples = [0] * num_shards
-        shard_parquet_lengths = [{} for _ in range(num_shards)]
-
-        # Distribute files to shards using a greedy approach
-        logger.info("Distributing files to shards...")
-        for file, length in zip(reversed(sorted_files),
-                                reversed(sorted_lengths)):
-            # Find shard with minimum current length
-            target_shard = np.argmin(shard_total_samples)
-            shard_parquet_files[target_shard].append(file)
-            shard_total_samples[target_shard] += length
-            shard_parquet_lengths[target_shard][file] = length
-        #randomize each shard
-        for shard in shard_parquet_files:
-            random.seed(seed)
-            random.shuffle(shard)
-
-        save_dir = os.path.join(
-            path,
-            f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers")
-        os.makedirs(save_dir, exist_ok=True)
-        with open(os.path.join(save_dir, "shard_parquet_files.pkl"), "wb") as f:
-            pickle.dump(shard_parquet_files, f)
-        with open(os.path.join(save_dir, "shard_total_samples.pkl"), "wb") as f:
-            pickle.dump(shard_total_samples, f)
-        with open(os.path.join(save_dir, "shard_parquet_lengths.pkl"),
-                  "wb") as f:
-            pickle.dump(shard_parquet_lengths, f)
-        logger.info("Saved sharding info to %s", save_dir)
-
-    # wait for all ranks to finish
-    torch.distributed.barrier()
-    # recursive call
-    return shard_parquet_files_across_sp_groups_and_workers(
-        path, num_sp_groups, num_workers, seed)
-
-
-def build_parquet_iterable_style_dataloader(
-    path: str,
-    batch_size: int,
-    num_data_workers: int,
-    cfg_rate: float = 0.0,
-    drop_last: bool = True,
-    text_padding_length: int = 512,
-    seed: int = 42,
-    read_batch_size: int = 32
-) -> Tuple[LatentsParquetIterStyleDataset, StatefulDataLoader]:
-    """Build a dataloader for the LatentsParquetIterStyleDataset."""
-    dataset = LatentsParquetIterStyleDataset(
-        path=path,
-        batch_size=batch_size,
-        cfg_rate=cfg_rate,
-        num_workers=num_data_workers,
-        drop_last=drop_last,
-        text_padding_length=text_padding_length,
-        seed=seed,
-        read_batch_size=read_batch_size)
-
-    loader = StatefulDataLoader(
-        dataset,
-        batch_size=1,
-        num_workers=num_data_workers,
-        pin_memory=True,
-    )
-    return dataset, loader

fastvideo/v1/dataset/parquet_dataset_map_style.py

@@ -1,314 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import os
-import pickle
-from typing import Any, Dict, List, Tuple
-
-import pyarrow as pa
-import pyarrow.parquet as pq
-# Torch in general
-import torch
-import tqdm
-# Dataset
-from torch.utils.data import Dataset, Sampler
-from torchdata.stateful_dataloader import StatefulDataLoader
-
-from fastvideo.v1.dataset.utils import collate_rows_from_parquet_schema
-from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
-                                      get_world_size)
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-class DP_SP_BatchSampler(Sampler[List[int]]):
-    """
-    A simple sequential batch sampler that yields batches of indices.
-    """
-
-    def __init__(
-        self,
-        batch_size: int,
-        dataset_size: int,
-        num_sp_groups: int,
-        sp_world_size: int,
-        global_rank: int,
-        drop_last: bool = True,
-        drop_first_row: bool = False,
-        seed: int = 0,
-    ):
-        self.batch_size = batch_size
-        self.dataset_size = dataset_size
-        self.drop_last = drop_last
-        self.seed = seed
-        self.num_sp_groups = num_sp_groups
-        self.global_rank = global_rank
-        self.sp_world_size = sp_world_size
-
-        # ── epoch-level RNG ────────────────────────────────────────────────
-        rng = torch.Generator().manual_seed(self.seed)
-        # Create a random permutation of all indices
-        global_indices = torch.randperm(self.dataset_size, generator=rng)
-
-        if drop_first_row:
-            # drop 0 in global_indices
-            global_indices = global_indices[global_indices != 0]
-            self.dataset_size = self.dataset_size - 1
-
-        if self.drop_last:
-            # For drop_last=True, we:
-            # 1. Ensure total samples is divisible by (batch_size * num_sp_groups)
-            # 2. This guarantees each SP group gets same number of complete batches
-            # 3. Prevents uneven batch sizes across SP groups at end of epoch
-            num_batches = self.dataset_size // self.batch_size
-            num_global_batches = num_batches // self.num_sp_groups
-            global_indices = global_indices[:num_global_batches *
-                                            self.num_sp_groups *
-                                            self.batch_size]
-        else:
-            if self.dataset_size % (self.num_sp_groups * self.batch_size) != 0:
-                # add more indices to make it divisible by (batch_size * num_sp_groups)
-                padding_size = self.num_sp_groups * self.batch_size - (
-                    self.dataset_size % (self.num_sp_groups * self.batch_size))
-                logger.info("Padding the dataset from %d to %d",
-                            self.dataset_size, self.dataset_size + padding_size)
-                global_indices = torch.cat(
-                    [global_indices, global_indices[:padding_size]])
-
-        # shard the indices to each sp group
-        ith_sp_group = self.global_rank // self.sp_world_size
-        sp_group_local_indices = global_indices[ith_sp_group::self.
-                                                num_sp_groups]
-        self.sp_group_local_indices = sp_group_local_indices
-        logger.info("Dataset size for each sp group: %d",
-                    len(sp_group_local_indices))
-
-    def __iter__(self):
-        indices = self.sp_group_local_indices
-        for i in range(0, len(indices), self.batch_size):
-            batch_indices = indices[i:i + self.batch_size]
-            yield batch_indices.tolist()
-
-    def __len__(self):
-        return len(self.sp_group_local_indices) // self.batch_size
-
-
-def get_parquet_files_and_length(path: str):
-    # Check if cached info exists
-    cache_dir = os.path.join(path, "map_style_cache")
-    cache_file = os.path.join(cache_dir, "file_info.pkl")
-
-    if os.path.exists(cache_file):
-        logger.info("Loading cached file info from %s", cache_file)
-        try:
-            with open(cache_file, "rb") as f:
-                file_names_sorted, lengths_sorted = pickle.load(f)
-            return file_names_sorted, lengths_sorted
-        except Exception as e:
-            logger.error("Error loading cached file info: %s", str(e))
-            logger.info("Falling back to scanning files")
-
-    # If no cache exists or loading failed, scan files
-    if get_world_rank() == 0:
-        lengths = []
-        file_names = []
-        for root, _, files in os.walk(path):
-            for file in sorted(files):
-                if file.endswith('.parquet'):
-                    file_path = os.path.join(root, file)
-                    file_names.append(file_path)
-        for file_path in tqdm.tqdm(file_names,
-                                   desc="Reading parquet files to get lengths"):
-            num_rows = pq.ParquetFile(file_path).metadata.num_rows
-            lengths.append(num_rows)
-        # sort according to file name to ensure all rank has the same order (in case os.walk is not sorted)
-        file_names_sorted, lengths_sorted = zip(
-            *sorted(zip(file_names, lengths), key=lambda x: x[0]))
-        assert len(
-            file_names_sorted) != 0, "No parquet files found in the dataset"
-
-        os.makedirs(cache_dir, exist_ok=True)
-        with open(cache_file, "wb") as f:
-            pickle.dump((file_names_sorted, lengths_sorted), f)
-        logger.info("Saved file info to %s", cache_file)
-
-    # Wait for rank 0 to finish saving
-    if get_world_size() > 1:
-        torch.distributed.barrier()
-
-    return get_parquet_files_and_length(path)
-
-
-def read_row_from_parquet_file(parquet_files: List[str], global_row_idx: int,
-                               lengths: List[int]) -> Dict[str, Any]:
-    '''
-    Read a row from a parquet file.
-    Args:
-        parquet_files: List[str]
-        global_row_idx: int
-        lengths: List[int]
-    Returns:
-    '''
-    # find the parquet file and local row index
-    cumulative = 0
-    for file_index in range(len(lengths)):
-        if cumulative + lengths[file_index] > global_row_idx:
-            local_row_idx = global_row_idx - cumulative
-            break
-        cumulative += lengths[file_index]
-
-    parquet_file = pq.ParquetFile(parquet_files[file_index])
-
-    # Calculate the row group to read into memory and the local idx
-    # This way we can avoid reading in the entire parquet file
-    cumulative = 0
-    for i in range(parquet_file.num_row_groups):
-        num_rows = parquet_file.metadata.row_group(i).num_rows
-        if cumulative + num_rows > local_row_idx:
-            row_group_index = i
-            local_index = local_row_idx - cumulative
-            break
-        cumulative += num_rows
-
-    row_group = parquet_file.read_row_group(row_group_index).to_pydict()
-    row_dict = {k: v[local_index] for k, v in row_group.items()}
-    del row_group
-
-    return row_dict
-
-
-# ────────────────────────────────────────────────────────────────────────────
-# 2.  Dataset with batched __getitems__
-# ────────────────────────────────────────────────────────────────────────────
-class LatentsParquetMapStyleDataset(Dataset):
-    """
-    Return latents[B,C,T,H,W] and embeddings[B,L,D] in pinned CPU memory.
-    Note: 
-    Using parquet for map style dataset is not efficient, we mainly keep it for backward compatibility and debugging.
-    """
-    # Modify this in the future if we want to add more keys, for example, in image to video.
-    keys = [("vae_latent", "latent"), "text_embedding", "clip_feature",
-            "first_frame_latent", "pil_image"]
-
-    def __init__(
-        self,
-        path: str,
-        batch_size: int,
-        parquet_schema: pa.Schema,
-        cfg_rate: float = 0.0,
-        seed: int = 42,
-        drop_last: bool = True,
-        drop_first_row: bool = False,
-        text_padding_length: int = 512,
-    ):
-        super().__init__()
-        self.path = path
-        self.cfg_rate = cfg_rate
-        self.parquet_schema = parquet_schema
-        if cfg_rate > 0.0:
-            raise ValueError(
-                "cfg_rate > 0.0 is not supported for now because it will trigger bug when num_data_workers > 0"
-            )
-        logger.info("Initializing LatentsParquetMapStyleDataset with path: %s",
-                    path)
-        self.parquet_files, self.lengths = get_parquet_files_and_length(path)
-        self.batch = batch_size
-        self.text_padding_length = text_padding_length
-        self.sampler = DP_SP_BatchSampler(
-            batch_size=batch_size,
-            dataset_size=sum(self.lengths),
-            num_sp_groups=get_world_size() // get_sp_world_size(),
-            sp_world_size=get_sp_world_size(),
-            global_rank=get_world_rank(),
-            drop_last=drop_last,
-            drop_first_row=drop_first_row,
-            seed=seed,
-        )
-        logger.info("Dataset initialized with %d parquet files and %d rows",
-                    len(self.parquet_files), sum(self.lengths))
-
-    def get_validation_negative_prompt(
-            self) -> tuple[torch.Tensor, torch.Tensor, str]:
-        """
-        Get the negative prompt for validation. 
-        This method ensures the negative prompt is loaded and cached properly.
-        Returns the processed negative prompt data (latents, embeddings, masks, info).
-        """
-
-        # Read first row from first parquet file
-        file_path = self.parquet_files[0]
-        row_idx = 0
-        # Read the negative prompt data
-        row_dict = read_row_from_parquet_file([file_path], row_idx,
-                                              [self.lengths[0]])
-
-        batch = collate_rows_from_parquet_schema([row_dict],
-                                                 self.parquet_schema,
-                                                 self.text_padding_length)
-        negative_prompt = batch['info_list'][0]['prompt']
-        negative_prompt_embedding = batch['text_embedding']
-        negative_prompt_attention_mask = batch['text_attention_mask']
-        if len(negative_prompt_embedding.shape) == 2:
-            negative_prompt_embedding = negative_prompt_embedding.unsqueeze(0)
-        if len(negative_prompt_attention_mask.shape) == 1:
-            negative_prompt_attention_mask = negative_prompt_attention_mask.unsqueeze(
-                0).unsqueeze(0)
-
-        return negative_prompt_embedding, negative_prompt_attention_mask, negative_prompt
-
-    # PyTorch calls this ONLY because the batch_sampler yields a list
-    def __getitems__(self, indices: List[int]) -> Dict[str, Any]:
-        """
-        Batch fetch using read_row_from_parquet_file for each index.
-        """
-        rows = [
-            read_row_from_parquet_file(self.parquet_files, idx, self.lengths)
-            for idx in indices
-        ]
-
-        # all_latents, all_embs, all_masks, caption_text, all_extra_latents, all_infos = collate_latents_embs_masks(
-        #     rows, self.text_padding_length, self.keys)
-        # return all_latents, all_embs, all_masks, caption_text, all_extra_latents, all_infos
-        batch = collate_rows_from_parquet_schema(rows, self.parquet_schema,
-                                                 self.text_padding_length)
-        return batch
-
-    def __len__(self):
-        return sum(self.lengths)
-
-
-# ────────────────────────────────────────────────────────────────────────────
-# 3.  Loader helper – everything else stays just like your original trainer
-# ────────────────────────────────────────────────────────────────────────────
-def passthrough(batch):
-    return batch
-
-
-def build_parquet_map_style_dataloader(
-        path,
-        batch_size,
-        num_data_workers,
-        parquet_schema,
-        cfg_rate=0.0,
-        drop_last=True,
-        drop_first_row=False,
-        text_padding_length=512,
-        seed=42) -> Tuple[LatentsParquetMapStyleDataset, StatefulDataLoader]:
-    dataset = LatentsParquetMapStyleDataset(
-        path,
-        batch_size,
-        cfg_rate=cfg_rate,
-        drop_last=drop_last,
-        drop_first_row=drop_first_row,
-        text_padding_length=text_padding_length,
-        parquet_schema=parquet_schema,
-        seed=seed)
-
-    loader = StatefulDataLoader(
-        dataset,
-        batch_sampler=dataset.sampler,
-        collate_fn=passthrough,
-        num_workers=num_data_workers,
-        pin_memory=True,
-        persistent_workers=num_data_workers > 0,
-    )
-    return dataset, loader

fastvideo/v1/dataset/parquet_datasets.py

@@ -0,0 +1,470 @@
+import argparse
+import json
+import os
+import random
+import time
+from collections import defaultdict
+from typing import Any, Dict, List
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+import tqdm
+from einops import rearrange
+from torch import distributed as dist
+from torch.utils.data import Dataset
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+from fastvideo.v1.distributed import (get_dp_group,
+                                      get_sequence_model_parallel_rank,
+                                      get_sp_group)
+from fastvideo.v1.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class ParquetVideoTextDataset(Dataset):
+    """Efficient loader for video-text data from a directory of Parquet files."""
+
+    def __init__(self,
+                 path: str,
+                 batch_size: int = 1024,
+                 rank: int = 0,
+                 world_size: int = 1,
+                 cfg_rate: float = 0.0,
+                 num_latent_t: int = 2,
+                 seed: int = 0,
+                 validation: bool = False):
+        super().__init__()
+        self.path = str(path)
+        self.batch_size = batch_size
+        self.rank = rank
+        self.local_rank = get_sequence_model_parallel_rank()
+        self.sp_group = get_sp_group()
+        self.dp_group = get_dp_group()
+        self.dp_world_size = self.dp_group.world_size
+        self.sp_world_size = self.sp_group.world_size
+        self.world_size = int(os.getenv("WORLD_SIZE", 1))
+        self.cfg_rate = cfg_rate
+        self.num_latent_t = num_latent_t
+        self.local_indices = None
+        self.validation = validation
+
+        # Negative prompt caching
+        self.neg_metadata = None
+        self.cached_neg_prompt: Dict[str, Any] | None = None
+
+        self.plan_output_dir = os.path.join(
+            self.path,
+            f"data_plan_{self.world_size}_{self.sp_world_size}_{self.dp_world_size}.json"
+        )
+
+        ranks = get_sp_group().ranks
+        group_ranks: List[List] = [[] for _ in range(self.world_size)]
+        torch.distributed.all_gather_object(group_ranks, ranks)
+
+        if rank == 0:
+            # If a plan already exists, then skip creating a new plan
+            # This will be useful when resume training
+            if os.path.exists(self.plan_output_dir):
+                print(f"Using existing plan from {self.plan_output_dir}")
+            else:
+                print(f"Creating new plan for {self.plan_output_dir}")
+                # Find all parquet files recursively, and record num_rows for each file
+                print(f"Scanning for parquet files in {self.path}")
+                metadatas = []
+                for root, _, files in os.walk(self.path):
+                    for file in sorted(files):
+                        if file.endswith('.parquet'):
+                            file_path = os.path.join(root, file)
+                            num_rows = pq.ParquetFile(
+                                file_path).metadata.num_rows
+                            for row_idx in range(num_rows):
+                                metadatas.append((file_path, row_idx))
+
+                # the negative prompt is always the first row in the first
+                # parquet file
+                if validation:
+                    self.neg_metadata = metadatas[0]
+                    metadatas = metadatas[1:]
+
+                # Generate the plan that distribute rows among workers
+                random.seed(seed)
+                random.shuffle(metadatas)
+
+                # Get all sp groups
+                # e.g. if num_gpus = 4, sp_size = 2
+                # group_ranks = [(0, 1), (2, 3)]
+                # We will assign the same batches of data to ranks in the same sp group, and we'll assign different batches to ranks in different sp groups
+                # e.g. plan = {0: [row 1, row 4], 1: [row 1, row 4], 2: [row 2, row 3], 3: [row 2, row 3]}
+                group_ranks_list: List[Any] = list(
+                    set(tuple(r) for r in group_ranks))
+                num_sp_groups = len(group_ranks_list)
+                plan = defaultdict(list)
+                for idx, metadata in enumerate(metadatas):
+                    sp_group_idx = idx % num_sp_groups
+                    for global_rank in group_ranks_list[sp_group_idx]:
+                        plan[global_rank].append(metadata)
+
+                if validation:
+                    assert self.neg_metadata is not None
+                    plan["negative_prompt"] = [self.neg_metadata]
+                with open(self.plan_output_dir, "w") as f:
+                    json.dump(plan, f)
+        else:
+            pass
+
+        dist.barrier()
+        if validation:
+            with open(self.plan_output_dir) as f:
+                plan = json.load(f)
+            self.neg_metadata = plan["negative_prompt"][0]
+
+    def _load_and_cache_negative_prompt(self) -> None:
+        """Load and cache the negative prompt. Only rank 0 in each SP group should call this."""
+        if not self.validation or self.neg_metadata is None:
+            return
+
+        if self.cached_neg_prompt is not None:
+            return
+
+        # Only rank 0 in each SP group should read the negative prompt
+        try:
+            file_path, row_idx = self.neg_metadata
+            parquet_file = pq.ParquetFile(file_path)
+
+            # Since negative prompt is always the first row (row_idx = 0),
+            # it's always in the first row group
+            row_group_index = 0
+            local_index = row_idx  # This will be 0 for the negative prompt
+
+            row_group = parquet_file.read_row_group(row_group_index).to_pydict()
+            row_dict = {k: v[local_index] for k, v in row_group.items()}
+            del row_group
+
+            # Process the negative prompt row
+            self.cached_neg_prompt = self._process_row(row_dict)
+
+        except Exception as e:
+            logger.error("Failed to load negative prompt: %s", e)
+            self.cached_neg_prompt = None
+
+    def get_validation_negative_prompt(
+        self
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Dict[str, Any]]:
+        """
+        Get the negative prompt for validation. 
+        This method ensures the negative prompt is loaded and cached properly.
+        Returns the processed negative prompt data (latents, embeddings, masks, info).
+        """
+        if not self.validation:
+            raise ValueError(
+                "get_validation_negative_prompt() can only be called in validation mode"
+            )
+
+        # Load and cache if needed (only rank 0 in SP group will actually load)
+        if self.cached_neg_prompt is None:
+            self._load_and_cache_negative_prompt()
+
+        if self.cached_neg_prompt is None:
+            raise RuntimeError(
+                f"Rank {self.rank} (SP rank {self.local_rank}): Could not retrieve negative prompt data"
+            )
+
+        # Extract the components
+        lat, emb, mask, info = (self.cached_neg_prompt["latents"],
+                                self.cached_neg_prompt["embeddings"],
+                                self.cached_neg_prompt["masks"],
+                                self.cached_neg_prompt["info"])
+
+        # Apply the same processing as in __getitem__
+        if lat.numel() == 0:  # Validation parquet
+            return lat, emb, mask, info
+        else:
+            lat = lat[:, -self.num_latent_t:]
+            if self.sp_world_size > 1:
+                lat = rearrange(lat,
+                                "t (n s) h w -> t n s h w",
+                                n=self.sp_world_size).contiguous()
+                lat = lat[:, self.local_rank, :, :, :]
+            return lat, emb, mask, info
+
+    def __len__(self):
+        if self.local_indices is None:
+            try:
+                with open(self.plan_output_dir) as f:
+                    plan = json.load(f)
+                self.local_indices = plan[str(self.rank)]
+            except Exception as err:
+                raise Exception(
+                    "The data plan hasn't been created yet") from err
+        assert self.local_indices is not None
+        return len(self.local_indices)
+
+    def __getitem__(self, idx):
+        if self.local_indices is None:
+            try:
+                with open(self.plan_output_dir) as f:
+                    plan = json.load(f)
+                self.local_indices = plan[self.rank]
+            except Exception as err:
+                raise Exception(
+                    "The data plan hasn't been created yet") from err
+        assert self.local_indices is not None
+        file_path, row_idx = self.local_indices[idx]
+        parquet_file = pq.ParquetFile(file_path)
+
+        # Calculate the row group to read into memory and the local idx
+        # This way we can avoid reading in the entire parquet file
+        cumulative = 0
+        for i in range(parquet_file.num_row_groups):
+            num_rows = parquet_file.metadata.row_group(i).num_rows
+            if cumulative + num_rows > row_idx:
+                row_group_index = i
+                local_index = row_idx - cumulative
+                break
+            cumulative += num_rows
+
+        row_group = parquet_file.read_row_group(row_group_index).to_pydict()
+        row_dict = {k: v[local_index] for k, v in row_group.items()}
+        del row_group
+
+        processed = self._process_row(row_dict)
+        lat, emb, mask, info = processed["latents"], processed[
+            "embeddings"], processed["masks"], processed["info"]
+        if lat.numel() == 0:  # Validation parquet
+            return lat, emb, mask, info
+        else:
+            lat = lat[:, -self.num_latent_t:]
+            if self.sp_world_size > 1:
+                lat = rearrange(lat,
+                                "t (n s) h w -> t n s h w",
+                                n=self.sp_world_size).contiguous()
+                lat = lat[:, self.local_rank, :, :, :]
+            return lat, emb, mask, info
+
+    def _process_row(self, row) -> Dict[str, Any]:
+        """Process a PyArrow batch into tensors."""
+
+        vae_latent_bytes = row["vae_latent_bytes"]
+        vae_latent_shape = row["vae_latent_shape"]
+        text_embedding_bytes = row["text_embedding_bytes"]
+        text_embedding_shape = row["text_embedding_shape"]
+        text_attention_mask_bytes = row["text_attention_mask_bytes"]
+        text_attention_mask_shape = row["text_attention_mask_shape"]
+
+        # Process latent
+        if not vae_latent_shape:  # No VAE latent is stored. Split is validation
+            lat = np.array([])
+        else:
+            lat = np.frombuffer(vae_latent_bytes,
+                                dtype=np.float32).reshape(vae_latent_shape)
+            # Make array writable
+            lat = np.copy(lat)
+
+        if random.random() < self.cfg_rate:
+            emb = np.zeros((512, 4096), dtype=np.float32)
+        else:
+            emb = np.frombuffer(text_embedding_bytes,
+                                dtype=np.float32).reshape(text_embedding_shape)
+            # Make array writable
+            emb = np.copy(emb)
+        if emb.shape[0] < 512:
+            padded_emb = np.zeros((512, emb.shape[1]), dtype=np.float32)
+            padded_emb[:emb.shape[0], :] = emb
+            emb = padded_emb
+        elif emb.shape[0] > 512:
+            emb = emb[:512, :]
+
+        # Process mask
+        if len(text_attention_mask_bytes) > 0 and len(
+                text_attention_mask_shape) > 0:
+            msk = np.frombuffer(text_attention_mask_bytes,
+                                dtype=np.uint8).astype(np.bool_)
+            msk = msk.reshape(1, -1)
+            # Make array writable
+            msk = np.copy(msk)
+            if msk.shape[1] < 512:
+                padded_msk = np.zeros((1, 512), dtype=np.bool_)
+                padded_msk[:, :msk.shape[1]] = msk
+                msk = padded_msk
+            elif msk.shape[1] > 512:
+                msk = msk[:, :512]
+        else:
+            msk = np.ones((1, 512), dtype=np.bool_)
+
+        # Collect metadata
+        info = {
+            "width": row["width"],
+            "height": row["height"],
+            "num_frames": row["num_frames"],
+            "duration_sec": row["duration_sec"],
+            "fps": row["fps"],
+            "file_name": row["file_name"],
+            "caption": row["caption"],
+        }
+
+        return {
+            "latents": torch.from_numpy(lat),
+            "embeddings": torch.from_numpy(emb),
+            "masks": torch.from_numpy(msk),
+            "info": info
+        }
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description='Benchmark Parquet dataset loading speed')
+    parser.add_argument('--path',
+                        type=str,
+                        default="your/dataset/path",
+                        help='Path to Parquet dataset')
+    parser.add_argument('--batch_size',
+                        type=int,
+                        default=4,
+                        help='Batch size for DataLoader')
+    parser.add_argument('--num_batches',
+                        type=int,
+                        default=100,
+                        help='Number of batches to benchmark')
+    parser.add_argument('--vae_debug', action="store_true")
+    args = parser.parse_args()
+
+    # Initialize distributed training
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    rank = int(os.environ.get("RANK", 0))
+
+    # Initialize CUDA device first
+    if torch.cuda.is_available():
+        torch.cuda.set_device(local_rank)
+        device = torch.device(f"cuda:{local_rank}")
+    else:
+        device = torch.device("cpu")
+
+    # Initialize distributed training
+    if world_size > 1:
+        dist.init_process_group(backend="nccl",
+                                init_method="env://",
+                                world_size=world_size,
+                                rank=rank)
+        print(
+            f"Initialized process: rank={rank}, local_rank={local_rank}, world_size={world_size}, device={device}"
+        )
+
+    # Create dataset
+    dataset = ParquetVideoTextDataset(
+        args.path,
+        batch_size=args.batch_size,
+        rank=rank,
+        world_size=world_size,
+    )
+
+    # Create DataLoader with proper settings
+    dataloader = StatefulDataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        num_workers=1,  # Reduce number of workers to avoid memory issues
+        prefetch_factor=2,
+        shuffle=False,
+        pin_memory=True,
+        drop_last=True)
+
+    # Example of how to load dataloader state
+    # if os.path.exists("/workspace/FastVideo/dataloader_state.pt"):
+    #     dataloader_state = torch.load("/workspace/FastVideo/dataloader_state.pt")
+    #     dataloader.load_state_dict(dataloader_state[rank])
+
+    # Warm-up with synchronization
+    if rank == 0:
+        print("Warming up...")
+    for i, (latents, embeddings, masks, infos) in enumerate(dataloader):
+        # Example of how to save dataloader state
+        # if i == 30:
+        #     dist.barrier()
+        #     local_data = {rank: dataloader.state_dict()}
+        #     gathered_data = [None] * world_size
+        #     dist.all_gather_object(gathered_data, local_data)
+        #     if rank == 0:
+        #         global_state_dict = {}
+        #         for d in gathered_data:
+        #             global_state_dict.update(d)
+        #         torch.save(global_state_dict, "dataloader_state.pt")
+        assert torch.sum(masks[0]).item() == torch.count_nonzero(
+            embeddings[0]).item() // 4096
+        if args.vae_debug:
+            from diffusers.utils import export_to_video
+            from diffusers.video_processor import VideoProcessor
+
+            from fastvideo.v1.configs.models.vaes import WanVAEConfig
+            from fastvideo.v1.fastvideo_args import FastVideoArgs
+            from fastvideo.v1.models.loader.component_loader import VAELoader
+            VAE_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers/vae"
+            fastvideo_args = FastVideoArgs(
+                model_path=VAE_PATH,
+                vae_config=WanVAEConfig(load_encoder=False),
+                vae_precision="fp32")
+            fastvideo_args.device = device
+            vae_loader = VAELoader()
+            vae = vae_loader.load(model_path=VAE_PATH,
+                                  architecture="",
+                                  fastvideo_args=fastvideo_args)
+
+            videoprocessor = VideoProcessor(vae_scale_factor=8)
+
+            with torch.inference_mode():
+                video = vae.decode(latents[0].unsqueeze(0).to(device))
+                video = videoprocessor.postprocess_video(video)
+                video_path = os.path.join("/workspace/FastVideo/debug_videos",
+                                          infos["caption"][0][:50] + ".mp4")
+                export_to_video(video[0], video_path, fps=16)
+
+        # Move data to device
+        # latents = latents.to(device)
+        # embeddings = embeddings.to(device)
+
+    if world_size > 1:
+        dist.barrier()
+
+    # Benchmark
+    if rank == 0:
+        print(f"Benchmarking with batch_size={args.batch_size}")
+    start_time = time.time()
+    total_samples = 0
+    for i, (latents, embeddings, masks,
+            infos) in enumerate(tqdm.tqdm(dataloader, total=args.num_batches)):
+        if i >= args.num_batches:
+            break
+
+        # Move data to device
+        latents = latents.to(device)
+        embeddings = embeddings.to(device)
+
+        # Calculate actual batch size
+        batch_size = latents.size(0)
+        total_samples += batch_size
+
+        # Print progress only from rank 0
+        if rank == 0 and (i + 1) % 10 == 0:
+            elapsed = time.time() - start_time
+            samples_per_sec = total_samples / elapsed
+            print(
+                f"Batch {i+1}/{args.num_batches}, Speed: {samples_per_sec:.2f} samples/sec"
+            )
+
+    # Final statistics
+    if world_size > 1:
+        dist.barrier()
+
+    if rank == 0:
+        elapsed = time.time() - start_time
+        samples_per_sec = total_samples / elapsed
+
+        print("\nBenchmark Results:")
+        print(f"Total time: {elapsed:.2f} seconds")
+        print(f"Total samples: {total_samples}")
+        print(f"Average speed: {samples_per_sec:.2f} samples/sec")
+        print(f"Time per batch: {elapsed/args.num_batches*1000:.2f} ms")
+
+    if world_size > 1:
+        dist.destroy_process_group()

fastvideo/v1/dataset/preprocessing_datasets.py

@@ -1,592 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import json
-import math
-import os
-import random
-from abc import ABC, abstractmethod
-from collections import Counter
-from dataclasses import dataclass
-from os.path import join as opj
-from typing import Any, Dict, List, Optional, Union
-
-import numpy as np
-import torch
-import torchvision
-from einops import rearrange
-from PIL import Image
-from transformers import AutoTokenizer
-
-from fastvideo.v1.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-@dataclass
-class PreprocessBatch:
-    """
-    Batch information for dataset processing stages.
-    
-    This class holds all the information about a video-caption or image-caption pair
-    as it moves through the processing pipeline. Fields are populated by different stages.
-    """
-    # Raw metadata
-    path: str
-    cap: Union[str, List[str]]
-    resolution: Optional[Dict] = None
-    fps: Optional[float] = None
-    duration: Optional[float] = None
-
-    # Processed metadata
-    num_frames: Optional[int] = None
-    sample_frame_index: Optional[List[int]] = None
-    sample_num_frames: Optional[int] = None
-
-    # Processed data
-    pixel_values: Optional[torch.Tensor] = None
-    text: Optional[str] = None
-    input_ids: Optional[torch.Tensor] = None
-    cond_mask: Optional[torch.Tensor] = None
-
-    @property
-    def is_video(self) -> bool:
-        """Check if this is a video item."""
-        return self.path.endswith(".mp4")
-
-    @property
-    def is_image(self) -> bool:
-        """Check if this is an image item."""
-        return self.path.endswith(".jpg")
-
-
-class DatasetStage(ABC):
-    """
-    Abstract base class for dataset processing stages.
-    
-    Similar to PipelineStage but designed for dataset preprocessing operations.
-    """
-
-    @abstractmethod
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """
-        Process the dataset batch.
-        
-        Args:
-            batch: Dataset batch to process
-            **kwargs: Additional processing parameters
-            
-        Returns:
-            Processed batch
-        """
-        raise NotImplementedError
-
-
-class DatasetFilterStage(ABC):
-    """
-    Abstract base class for dataset filtering stages.
-    
-    These stages can filter out items during metadata processing.
-    """
-
-    @abstractmethod
-    def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
-        """
-        Check if batch should be kept.
-        
-        Args:
-            batch: Dataset batch to check
-            **kwargs: Additional parameters
-            
-        Returns:
-            True if batch should be kept, False otherwise
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """
-        Process the dataset batch (for non-filtering operations).
-        
-        Args:
-            batch: Dataset batch to process
-            **kwargs: Additional processing parameters
-            
-        Returns:
-            Processed batch
-        """
-        raise NotImplementedError
-
-
-class DataValidationStage(DatasetFilterStage):
-    """Stage for validating data items."""
-
-    def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
-        """
-        Validate data item.
-        
-        Args:
-            batch: Dataset batch to validate
-            
-        Returns:
-            True if valid, False if invalid
-        """
-        # Check for caption
-        if batch.cap is None:
-            return False
-
-        if batch.is_video:
-            # Validate video-specific fields
-            if batch.duration is None or batch.fps is None:
-                return False
-        elif not batch.is_image:
-            return False
-
-        return True
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """Process does nothing for validation - filtering is handled by should_keep."""
-        return batch
-
-
-class ResolutionFilterStage(DatasetFilterStage):
-    """Stage for filtering data items based on resolution constraints."""
-
-    def __init__(self,
-                 max_h_div_w_ratio: float = 17 / 16,
-                 min_h_div_w_ratio: float = 8 / 16,
-                 max_height: int = 1024,
-                 max_width: int = 1024):
-        self.max_h_div_w_ratio = max_h_div_w_ratio
-        self.min_h_div_w_ratio = min_h_div_w_ratio
-        self.max_height = max_height
-        self.max_width = max_width
-
-    def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
-        """
-        Check if data item passes resolution filtering.
-        
-        Args:
-            batch: Dataset batch with resolution information
-            
-        Returns:
-            True if passes filter, False otherwise
-        """
-        # Only apply to videos
-        if not batch.is_video:
-            return True
-
-        if batch.resolution is None:
-            return False
-
-        height = batch.resolution.get("height", None)
-        width = batch.resolution.get("width", None)
-        if height is None or width is None:
-            return False
-
-        # Check aspect ratio
-        aspect = self.max_height / self.max_width
-        hw_aspect_thr = 1.5
-
-        return self.filter_resolution(
-            height,
-            width,
-            max_h_div_w_ratio=hw_aspect_thr * aspect,
-            min_h_div_w_ratio=1 / hw_aspect_thr * aspect,
-        )
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """Process does nothing for resolution filtering - filtering is handled by should_keep."""
-        return batch
-
-    def filter_resolution(self, h: int, w: int, max_h_div_w_ratio: float,
-                          min_h_div_w_ratio: float) -> bool:
-        """Filter based on height/width ratio."""
-        return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio
-
-
-class FrameSamplingStage(DatasetFilterStage):
-    """Stage for temporal frame sampling and indexing."""
-
-    def __init__(self,
-                 num_frames: int,
-                 train_fps: int,
-                 speed_factor: int = 1,
-                 video_length_tolerance_range: float = 5.0,
-                 drop_short_ratio: float = 0.0):
-        self.num_frames = num_frames
-        self.train_fps = train_fps
-        self.speed_factor = speed_factor
-        self.video_length_tolerance_range = video_length_tolerance_range
-        self.drop_short_ratio = drop_short_ratio
-
-    def should_keep(self, batch: PreprocessBatch, **kwargs) -> bool:
-        """
-        Check if video should be kept based on length constraints.
-        
-        Args:
-            batch: Dataset batch
-            
-        Returns:
-            True if should be kept, False otherwise
-        """
-        if batch.is_image:
-            return True
-
-        if batch.duration is None or batch.fps is None:
-            return False
-
-        num_frames = math.ceil(batch.fps * batch.duration)
-
-        # Check if video is too long
-        if (num_frames / batch.fps > self.video_length_tolerance_range *
-            (self.num_frames / self.train_fps * self.speed_factor)):
-            return False
-
-        # Resample frame indices to check length
-        frame_interval = batch.fps / self.train_fps
-        start_frame_idx = 0
-        frame_indices = np.arange(start_frame_idx, num_frames,
-                                  frame_interval).astype(int)
-
-        # Filter short videos
-        return not (len(frame_indices) < self.num_frames
-                    and random.random() < self.drop_short_ratio)
-
-    def process(self,
-                batch: PreprocessBatch,
-                temporal_sample_fn=None,
-                **kwargs) -> PreprocessBatch:
-        """
-        Process frame sampling for video data items.
-        
-        Args:
-            batch: Dataset batch
-            temporal_sample_fn: Function for temporal sampling
-            
-        Returns:
-            Updated batch with frame sampling info
-        """
-        if batch.is_image:
-            # For images, just add sample info
-            batch.sample_frame_index = [0]
-            batch.sample_num_frames = 1
-            return batch
-
-        assert batch.duration is not None and batch.fps is not None
-        batch.num_frames = math.ceil(batch.fps * batch.duration)
-
-        # Resample frame indices
-        frame_interval = batch.fps / self.train_fps
-        start_frame_idx = 0
-        frame_indices = np.arange(start_frame_idx, batch.num_frames,
-                                  frame_interval).astype(int)
-
-        # Temporal crop if too long
-        if len(frame_indices
-               ) > self.num_frames and temporal_sample_fn is not None:
-            begin_index, end_index = temporal_sample_fn(len(frame_indices))
-            frame_indices = frame_indices[begin_index:end_index]
-
-        batch.sample_frame_index = frame_indices.tolist()
-        batch.sample_num_frames = len(frame_indices)
-
-        return batch
-
-
-class VideoTransformStage(DatasetStage):
-    """Stage for video data transformation."""
-
-    def __init__(self, transform) -> None:
-        self.transform = transform
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """
-        Transform video data.
-        
-        Args:
-            batch: Dataset batch with video information
-            
-        Returns:
-            Batch with transformed video tensor
-        """
-        if not batch.is_video:
-            return batch
-
-        assert os.path.exists(batch.path), f"file {batch.path} do not exist!"
-        assert batch.sample_frame_index is not None, "Frame indices must be set before transformation"
-
-        torchvision_video, _, metadata = torchvision.io.read_video(
-            batch.path, output_format="TCHW")
-        video = torchvision_video[batch.sample_frame_index]
-        if self.transform is not None:
-            video = self.transform(video)
-        video = rearrange(video, "t c h w -> c t h w")
-        video = video.to(torch.uint8)
-
-        h, w = video.shape[-2:]
-        assert (
-            h / w <= 17 / 16 and h / w >= 8 / 16
-        ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({batch.path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}"
-
-        video = video.float() / 127.5 - 1.0
-        batch.pixel_values = video
-        return batch
-
-
-class ImageTransformStage(DatasetStage):
-    """Stage for image data transformation."""
-
-    def __init__(self, transform, transform_topcrop) -> None:
-        self.transform = transform
-        self.transform_topcrop = transform_topcrop
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """
-        Transform image data.
-        
-        Args:
-            batch: Dataset batch with image information
-            
-        Returns:
-            Batch with transformed image tensor
-        """
-        if not batch.is_image:
-            return batch
-
-        image = Image.open(batch.path).convert("RGB")
-        image = torch.from_numpy(np.array(image))
-        image = rearrange(image, "h w c -> c h w").unsqueeze(0)
-
-        if self.transform_topcrop is not None:
-            image = self.transform_topcrop(image)
-        elif self.transform is not None:
-            image = self.transform(image)
-
-        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
-        image = image.float() / 127.5 - 1.0
-        batch.pixel_values = image
-        return batch
-
-
-class TextEncodingStage(DatasetStage):
-    """Stage for text tokenization and encoding."""
-
-    def __init__(self, tokenizer, text_max_length: int, cfg_rate: float = 0.0):
-        self.tokenizer = tokenizer
-        self.text_max_length = text_max_length
-        self.cfg_rate = cfg_rate
-
-    def process(self, batch: PreprocessBatch, **kwargs) -> PreprocessBatch:
-        """
-        Process text data.
-        
-        Args:
-            batch: Dataset batch with caption information
-            
-        Returns:
-            Batch with encoded text information
-        """
-        text = batch.cap
-        if not isinstance(text, list):
-            text = [text]
-        text = [random.choice(text)]
-
-        text = text[0] if random.random() > self.cfg_rate else ""
-        text_tokens_and_mask = self.tokenizer(
-            text,
-            max_length=self.text_max_length,
-            padding="max_length",
-            truncation=True,
-            return_attention_mask=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-
-        batch.text = text
-        batch.input_ids = text_tokens_and_mask["input_ids"]
-        batch.cond_mask = text_tokens_and_mask["attention_mask"]
-        return batch
-
-
-class VideoCaptionMergedDataset(torch.utils.data.IterableDataset,
-                                torch.distributed.checkpoint.stateful.Stateful):
-    """
-    Merged dataset for video and caption data with stage-based processing.
-    
-    This dataset processes video and image data through a series of stages:
-    - Data validation
-    - Resolution filtering  
-    - Frame sampling
-    - Transformation
-    - Text encoding
-    """
-
-    def __init__(self,
-                 data_merge_path: str,
-                 args,
-                 transform,
-                 temporal_sample,
-                 transform_topcrop,
-                 start_idx: int = 0):
-        self.data_merge_path = data_merge_path
-        self.start_idx = start_idx
-        self.args = args
-        self.temporal_sample = temporal_sample
-
-        # Initialize tokenizer
-        tokenizer_path = os.path.join(args.model_path, "tokenizer")
-        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path,
-                                                  cache_dir=args.cache_dir)
-
-        # Initialize processing stages
-        self._init_stages(args, transform, transform_topcrop, tokenizer)
-
-        # Process metadata
-        self.processed_batches = self._process_metadata()
-
-    def _init_stages(self, args, transform, transform_topcrop,
-                     tokenizer) -> None:
-        """Initialize all processing stages."""
-        self.validation_stage = DataValidationStage()
-        self.resolution_filter_stage = ResolutionFilterStage(
-            max_height=args.max_height, max_width=args.max_width)
-        self.frame_sampling_stage = FrameSamplingStage(
-            num_frames=args.num_frames,
-            train_fps=args.train_fps,
-            speed_factor=args.speed_factor,
-            video_length_tolerance_range=args.video_length_tolerance_range,
-            drop_short_ratio=args.drop_short_ratio)
-        self.video_transform_stage = VideoTransformStage(transform)
-        self.image_transform_stage = ImageTransformStage(
-            transform, transform_topcrop)
-        self.text_encoding_stage = TextEncodingStage(
-            tokenizer=tokenizer,
-            text_max_length=args.text_max_length,
-            cfg_rate=args.cfg)
-
-    def _load_raw_data(self) -> List[Dict]:
-        """Load raw data from JSON files."""
-        all_data = []
-
-        # Read folder-annotation pairs
-        with open(self.data_merge_path) as f:
-            folder_anno_pairs = [
-                line.strip().split(",") for line in f if line.strip()
-            ]
-
-        # Process each folder-annotation pair
-        for folder, annotation_file in folder_anno_pairs:
-            with open(annotation_file) as f:
-                data_items = json.load(f)
-
-            # Update paths with folder prefix
-            for item in data_items:
-                item["path"] = opj(folder, item["path"])
-
-            all_data.extend(data_items)
-
-        return all_data[self.start_idx:]
-
-    def _process_metadata(self) -> List[PreprocessBatch]:
-        """Process the raw metadata through all filtering stages."""
-        raw_data = self._load_raw_data()
-        processed_batches = []
-
-        # Initialize counters
-        filter_counts = {
-            "validation_failed": 0,
-            "resolution_failed": 0,
-            "frame_sampling_failed": 0
-        }
-        sample_num_frames: List[int] = []
-
-        for item in raw_data:
-            batch = PreprocessBatch(path=item["path"],
-                                    cap=item["cap"],
-                                    resolution=item.get("resolution"),
-                                    fps=item.get("fps"),
-                                    duration=item.get("duration"))
-
-            # Apply filtering stages
-            if not self._apply_filter_stages(batch, filter_counts):
-                continue
-
-            # Apply frame sampling processing
-            batch = self.frame_sampling_stage.process(
-                batch, temporal_sample_fn=self.temporal_sample)
-
-            processed_batches.append(batch)
-            assert batch.sample_num_frames is not None
-            sample_num_frames.append(batch.sample_num_frames)
-
-        self._log_filtering_stats(filter_counts, sample_num_frames,
-                                  len(raw_data), len(processed_batches))
-        return processed_batches
-
-    def _apply_filter_stages(self, batch: PreprocessBatch,
-                             filter_counts: Dict[str, int]) -> bool:
-        """Apply all filter stages and update counters. Returns True if batch should be kept."""
-        if not self.validation_stage.should_keep(batch):
-            filter_counts["validation_failed"] += 1
-            return False
-
-        if not self.resolution_filter_stage.should_keep(batch):
-            filter_counts["resolution_failed"] += 1
-            return False
-
-        if not self.frame_sampling_stage.should_keep(batch):
-            filter_counts["frame_sampling_failed"] += 1
-            return False
-
-        return True
-
-    def _log_filtering_stats(self, filter_counts: Dict[str, int],
-                             sample_num_frames: List[int], before_count: int,
-                             after_count: int):
-        """Log filtering statistics."""
-        logger.info(
-            "validation_failed: %d, resolution_failed: %d, frame_sampling_failed: %d, "
-            "Counter(sample_num_frames): %s, before filter: %d, after filter: %d",
-            filter_counts['validation_failed'],
-            filter_counts['resolution_failed'],
-            filter_counts['frame_sampling_failed'], Counter(sample_num_frames),
-            before_count, after_count)
-
-    def __iter__(self):
-        """Iterate through processed data items."""
-        for idx in range(len(self.processed_batches)):
-            yield self._get_item(idx)
-
-    def __len__(self):
-        return len(self.processed_batches)
-
-    def _get_item(self, idx: int) -> Dict:
-        """Get a single processed data item."""
-        batch = self.processed_batches[idx]
-
-        # Apply transformation stages
-        batch = self.video_transform_stage.process(batch)
-        batch = self.image_transform_stage.process(batch)
-        batch = self.text_encoding_stage.process(batch)
-
-        # Build result dictionary
-        result = {
-            "pixel_values": batch.pixel_values,
-            "text": batch.text,
-            "input_ids": batch.input_ids,
-            "cond_mask": batch.cond_mask,
-            "path": batch.path,
-        }
-
-        # Add video-specific fields
-        if batch.is_video:
-            result.update({"fps": batch.fps, "duration": batch.duration})
-
-        return result
-
-    def state_dict(self) -> Dict[str, Any]:
-        """Return state dict for checkpointing."""
-        return {"processed_batches": self.processed_batches}
-
-    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        """Load state dict from checkpoint."""
-        self.processed_batches = state_dict["processed_batches"]

fastvideo/v1/dataset/t2v_datasets.py

@@ -0,0 +1,351 @@
+import json
+import math
+import os
+import random
+from collections import Counter
+from os.path import join as opj
+
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+from PIL import Image
+from torch.utils.data import Dataset
+
+from fastvideo.utils.dataset_utils import DecordInit
+from fastvideo.utils.logging_ import main_print
+
+
+class SingletonMeta(type):
+    _instances: dict[type, 'SingletonMeta'] = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            instance = super().__call__(*args, **kwargs)
+            cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+class DataSetProg(metaclass=SingletonMeta):
+
+    def __init__(self) -> None:
+        self.cap_list: list[dict] = []
+        self.elements: list[int] = []
+        self.num_workers = 1
+        self.n_elements = 0
+        self.worker_elements: dict[int, list[int]] = {}
+        self.n_used_elements: dict[int, int] = {}
+
+    def set_cap_list(self, num_workers, cap_list, n_elements) -> None:
+        self.num_workers = num_workers
+        self.cap_list = cap_list
+        self.n_elements = n_elements
+        self.elements = list(range(n_elements))
+        random.shuffle(self.elements)
+        print(f"n_elements: {len(self.elements)}", flush=True)
+
+        for i in range(self.num_workers):
+            self.n_used_elements[i] = 0
+            per_worker = int(
+                math.ceil(len(self.elements) / float(self.num_workers)))
+            start = i * per_worker
+            end = min(start + per_worker, len(self.elements))
+            self.worker_elements[i] = self.elements[start:end]
+
+    def get_item(self, work_info) -> int:
+        worker_id = 0 if work_info is None else work_info.id
+
+        idx = self.worker_elements[worker_id][
+            self.n_used_elements[worker_id] %
+            len(self.worker_elements[worker_id])]
+        self.n_used_elements[worker_id] += 1
+        return idx
+
+
+dataset_prog = DataSetProg()
+
+
+def filter_resolution(h: int,
+                      w: int,
+                      max_h_div_w_ratio: float = 17 / 16,
+                      min_h_div_w_ratio: float = 8 / 16) -> bool:
+    return h / w <= max_h_div_w_ratio and h / w >= min_h_div_w_ratio
+
+
+class T2V_dataset(Dataset):
+
+    def __init__(self,
+                 args,
+                 transform,
+                 temporal_sample,
+                 tokenizer,
+                 transform_topcrop,
+                 start_idx=0) -> None:
+        self.start_idx = start_idx
+        self.data = args.data_merge_path
+        self.num_frames = args.num_frames
+        self.train_fps = args.train_fps
+        self.use_image_num = args.use_image_num
+        self.transform = transform
+        self.transform_topcrop = transform_topcrop
+        self.temporal_sample = temporal_sample
+        self.tokenizer = tokenizer
+        self.text_max_length = args.text_max_length
+        self.cfg = args.cfg
+        self.speed_factor = args.speed_factor
+        self.max_height = args.max_height
+        self.max_width = args.max_width
+        self.drop_short_ratio = args.drop_short_ratio
+        assert self.speed_factor >= 1
+        self.v_decoder = DecordInit()
+        self.video_length_tolerance_range = args.video_length_tolerance_range
+        self.support_Chinese = True
+        if "mt5" not in args.text_encoder_name:
+            self.support_Chinese = False
+
+        cap_list = self.get_cap_list()
+
+        assert len(cap_list) > 0
+        cap_list, self.sample_num_frames = self.define_frame_index(cap_list)
+        self.lengths = self.sample_num_frames
+
+        n_elements = len(cap_list)
+        dataset_prog.set_cap_list(args.dataloader_num_workers, cap_list,
+                                  n_elements)
+
+        print(f"video length: {len(dataset_prog.cap_list)}", flush=True)
+
+    def set_checkpoint(self, n_used_elements):
+        for i in range(len(dataset_prog.n_used_elements)):
+            dataset_prog.n_used_elements[i] = n_used_elements
+
+    def __len__(self):
+        return dataset_prog.n_elements
+
+    def __getitem__(self, idx):
+
+        data = self.get_data(idx)
+        return data
+
+    def get_data(self, idx) -> dict:
+        path = dataset_prog.cap_list[idx]["path"]
+        if path.endswith(".mp4"):
+            return self.get_video(idx)
+        else:
+            return self.get_image(idx)
+
+    def get_video(self, idx) -> dict:
+        video_path = dataset_prog.cap_list[idx]["path"]
+        assert os.path.exists(video_path), f"file {video_path} do not exist!"
+        frame_indices = dataset_prog.cap_list[idx]["sample_frame_index"]
+
+        torchvision_video, _, metadata = torchvision.io.read_video(
+            video_path, output_format="TCHW")
+        video = torchvision_video[frame_indices]
+        video = self.transform(video)
+        video = rearrange(video, "t c h w -> c t h w")
+        video = video.to(torch.uint8)
+        assert video.dtype == torch.uint8
+
+        h, w = video.shape[-2:]
+        assert (
+            h / w <= 17 / 16 and h / w >= 8 / 16
+        ), f"Only videos with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But video ({video_path}) found ratio is {round(h / w, 2)} with the shape of {video.shape}"
+
+        video = video.float() / 127.5 - 1.0
+
+        text = dataset_prog.cap_list[idx]["cap"]
+        if not isinstance(text, list):
+            text = [text]
+        text = [random.choice(text)]
+
+        text = text[0] if random.random() > self.cfg else ""
+        text_tokens_and_mask = self.tokenizer(
+            text,
+            max_length=self.text_max_length,
+            padding="max_length",
+            truncation=True,
+            return_attention_mask=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        input_ids = text_tokens_and_mask["input_ids"]
+        cond_mask = text_tokens_and_mask["attention_mask"]
+        return dict(pixel_values=video,
+                    text=text,
+                    input_ids=input_ids,
+                    cond_mask=cond_mask,
+                    path=video_path,
+                    fps=dataset_prog.cap_list[idx]["fps"],
+                    duration=dataset_prog.cap_list[idx]["duration"])
+
+    def get_image(self, idx) -> dict:
+        image_data = dataset_prog.cap_list[
+            idx]  # [{'path': path, 'cap': cap}, ...]
+
+        image = Image.open(image_data["path"]).convert("RGB")  # [h, w, c]
+        image = torch.from_numpy(np.array(image))  # [h, w, c]
+        image = rearrange(image, "h w c -> c h w").unsqueeze(0)  #  [1 c h w]
+        # for i in image:
+        #     h, w = i.shape[-2:]
+        #     assert h / w <= 17 / 16 and h / w >= 8 / 16, f'Only image with a ratio (h/w) less than 17/16 and more than 8/16 are supported. But found ratio is {round(h / w, 2)} with the shape of {i.shape}'
+
+        image = (self.transform_topcrop(image) if "human_images"
+                 in image_data["path"] else self.transform(image)
+                 )  #  [1 C H W] -> num_img [1 C H W]
+        image = image.transpose(0, 1)  # [1 C H W] -> [C 1 H W]
+
+        image = image.float() / 127.5 - 1.0
+
+        caps: list[str] = (image_data["cap"] if isinstance(
+            image_data["cap"], list) else [image_data["cap"]])
+        caps = [random.choice(caps)]
+        text = caps
+        input_ids, cond_mask = [], []
+        single_text = text[0] if random.random() > self.cfg else ""
+        text_tokens_and_mask = self.tokenizer(
+            single_text,
+            max_length=self.text_max_length,
+            padding="max_length",
+            truncation=True,
+            return_attention_mask=True,
+            add_special_tokens=True,
+            return_tensors="pt",
+        )
+        input_ids = text_tokens_and_mask["input_ids"]  # 1, l
+        cond_mask = text_tokens_and_mask["attention_mask"]  # 1, l
+        return dict(
+            pixel_values=image,
+            text=text,
+            input_ids=input_ids,
+            cond_mask=cond_mask,
+            path=image_data["path"],
+        )
+
+    def define_frame_index(self, cap_list) -> tuple[list[dict], list[int]]:
+        new_cap_list = []
+        sample_num_frames = []
+        cnt_too_long = 0
+        cnt_too_short = 0
+        cnt_no_cap = 0
+        cnt_no_resolution = 0
+        cnt_resolution_mismatch = 0
+        cnt_movie = 0
+        cnt_img = 0
+        for i in cap_list:
+            path = i["path"]
+            cap = i.get("cap", None)
+            # ======no caption=====
+            if cap is None:
+                cnt_no_cap += 1
+                continue
+            if path.endswith(".mp4"):
+                # ======no fps and duration=====
+                duration = i.get("duration", None)
+                fps = i.get("fps", None)
+                if fps is None or duration is None:
+                    continue
+
+                # ======resolution mismatch=====
+                resolution = i.get("resolution", None)
+                if resolution is None:
+                    cnt_no_resolution += 1
+                    continue
+                else:
+                    if (resolution.get("height", None) is None
+                            or resolution.get("width", None) is None):
+                        cnt_no_resolution += 1
+                        continue
+                    height, width = i["resolution"]["height"], i["resolution"][
+                        "width"]
+                    aspect = self.max_height / self.max_width
+                    hw_aspect_thr = 1.5
+                    is_pick = filter_resolution(
+                        height,
+                        width,
+                        max_h_div_w_ratio=hw_aspect_thr * aspect,
+                        min_h_div_w_ratio=1 / hw_aspect_thr * aspect,
+                    )
+                    if not is_pick:
+                        print("resolution mismatch")
+                        cnt_resolution_mismatch += 1
+                        continue
+
+                # if path == 'finetrainers/3dgs-dissolve/videos/1.mp4':
+                #     from IPython import embed; embed()
+                i["num_frames"] = math.ceil(fps * duration)
+                # max 5.0 and min 1.0 are just thresholds to filter some videos which have suitable duration.
+                if i["num_frames"] / fps > self.video_length_tolerance_range * (
+                        self.num_frames / self.train_fps * self.speed_factor
+                ):  # too long video is not suitable for this training stage (self.num_frames)
+                    cnt_too_long += 1
+                    continue
+
+                # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24)
+                frame_interval = fps / self.train_fps
+                start_frame_idx = 0
+                frame_indices = np.arange(start_frame_idx, i["num_frames"],
+                                          frame_interval).astype(int)
+
+                # comment out it to enable dynamic frames training
+                if (len(frame_indices) < self.num_frames
+                        and random.random() < self.drop_short_ratio):
+                    cnt_too_short += 1
+                    continue
+
+                #  too long video will be temporal-crop randomly
+                if len(frame_indices) > self.num_frames:
+                    begin_index, end_index = self.temporal_sample(
+                        len(frame_indices))
+                    frame_indices = frame_indices[begin_index:end_index]
+                    # frame_indices = frame_indices[:self.num_frames]  # head crop
+                i["sample_frame_index"] = frame_indices.tolist()
+                new_cap_list.append(i)
+                i["sample_num_frames"] = len(
+                    i["sample_frame_index"]
+                )  # will use in dataloader(group sampler)
+                sample_num_frames.append(i["sample_num_frames"])
+            elif path.endswith(".jpg"):  # image
+                cnt_img += 1
+                new_cap_list.append(i)
+                i["sample_num_frames"] = 1
+                sample_num_frames.append(i["sample_num_frames"])
+            else:
+                raise NameError(
+                    f"Unknown file extension {path.split('.')[-1]}, only support .mp4 for video and .jpg for image"
+                )
+        # import ipdb;ipdb.set_trace()
+        main_print(
+            f"no_cap: {cnt_no_cap}, too_long: {cnt_too_long}, too_short: {cnt_too_short}, "
+            f"no_resolution: {cnt_no_resolution}, resolution_mismatch: {cnt_resolution_mismatch}, "
+            f"Counter(sample_num_frames): {Counter(sample_num_frames)}, cnt_movie: {cnt_movie}, cnt_img: {cnt_img}, "
+            f"before filter: {len(cap_list)}, after filter: {len(new_cap_list)}"
+        )
+        return new_cap_list, sample_num_frames
+
+    def decord_read(self, path, frame_indices) -> torch.Tensor:
+        decord_vr = self.v_decoder(path)
+        video_data = decord_vr.get_batch(frame_indices).asnumpy()
+        video_data = torch.from_numpy(video_data)
+        video_data = video_data.permute(0, 3, 1, 2)  # (T, H, W, C) -> (T C H W)
+        return video_data
+
+    def read_jsons(self, data) -> list[dict]:
+        cap_lists = []
+        with open(data) as f:
+            folder_anno = [
+                i.strip().split(",") for i in f.readlines()
+                if len(i.strip()) > 0
+            ]
+        print(folder_anno)
+        for folder, anno in folder_anno:
+            with open(anno) as f:
+                sub_list = json.load(f)
+            for i in range(len(sub_list)):
+                sub_list[i]["path"] = opj(folder, sub_list[i]["path"])
+            cap_lists += sub_list
+        return cap_lists
+
+    def get_cap_list(self) -> list:
+        cap_lists = self.read_jsons(self.data)[self.start_idx:]
+        return cap_lists

fastvideo/v1/dataset/transform.py

@@ -1,4 +1,3 @@
-# SPDX-License-Identifier: Apache-2.0
 import random
 
 import torch