[ci] Add /test pre-commit slash command

- /test pre-commit triggers pre-commit via workflow_call
- Posts commit status to PR head SHA so check appears on PR page
- Mergify check-success~=pre-commit condition can now be satisfied

.buildkite/pipeline.yml

@@ -2,17 +2,45 @@ env:
   IMAGE_VERSION: "py3.12-latest"
   BUILDKITE_CLEAN_CHECKOUT: true
 
+notify:
+  - github_commit_status:
+      context: "fastcheck-passed"
+    if: build.env("TEST_SCOPE") == "fastcheck" || build.env("TEST_SCOPE") == null
+  - github_commit_status:
+      context: "full-suite-passed"
+    if: build.env("TEST_SCOPE") == "full"
+
 steps:
+  # ============================================================
+    - label: ":dart: Direct Test (${TEST_TYPE})"
+      if: build.env("TEST_SCOPE") == "direct"
+      command: "timeout 90m .buildkite/scripts/pr_test.sh"
+      retry:
+        automatic:
+          - exit_status: 128
+            limit: 3
+          - exit_status: -1
+            limit: 2
+      agents:
+        queue: "default"
+
   # ============================================================
   # Fastcheck: Runs on every PR (~10-15 min parallel)
   # Core component validation: encoders, VAEs, transformers,
   # CUDA kernels, and unit tests.
   # ============================================================
-  - label: "Trigger Fastcheck"
-    plugins:
-      - monorepo-diff#v1.4.0:
-          diff: 'git fetch origin "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" && git diff --name-only origin/"$BUILDKITE_PULL_REQUEST_BASE_BRANCH"...HEAD'
-          watch:
+    - label: "Trigger Fastcheck"
+      if: build.env("TEST_SCOPE") == "fastcheck" || build.env("TEST_SCOPE") == null
+      retry:
+        automatic:
+          - exit_status: 128
+            limit: 3
+          - exit_status: -1
+            limit: 2
+      plugins:
+        - monorepo-diff#v1.4.0:
+            diff: 'git fetch origin "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" && git diff --name-only "origin/${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}...HEAD"'
+            watch:
             - path:
                 - "fastvideo/models/encoders/**"
                 - "fastvideo/models/loader/**"
@@ -80,17 +108,24 @@ steps:
                   queue: "default"
 
   # ============================================================
-  # Full Suite: Runs only when FULL_SUITE=true
-  # Triggered by adding the 'ready' label to a PR (via GitHub
-  # Actions → Buildkite API). Includes integration tests, SSIM
-  # regression, training pipelines, and performance benchmarks.
+  # Full Suite: Runs when TEST_SCOPE=full
+  # Triggered by adding the 'ready' label (via ci-trigger-full-suite.yml)
+  # or on-demand via /test full slash command.
+  # Includes integration tests, SSIM regression, training pipelines,
+  # and performance benchmarks.
   # ============================================================
-  - label: "Trigger Full Suite"
-    if: build.env("FULL_SUITE") == "true"
-    plugins:
-      - monorepo-diff#v1.4.0:
-          diff: 'git fetch origin "$BUILDKITE_PULL_REQUEST_BASE_BRANCH" && git diff --name-only origin/"$BUILDKITE_PULL_REQUEST_BASE_BRANCH"...HEAD'
-          watch:
+    - label: "Trigger Full Suite"
+      if: build.env("TEST_SCOPE") == "full"
+      retry:
+        automatic:
+          - exit_status: 128
+            limit: 3
+          - exit_status: -1
+            limit: 2
+      plugins:
+        - monorepo-diff#v1.4.0:
+            diff: 'git fetch origin "${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}" && git diff --name-only "origin/${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}...HEAD"'
+            watch:
             - path:
                 - "fastvideo/**/*.py"
                 - "pyproject.toml"

.buildkite/scripts/pr_test.sh

@@ -59,7 +59,11 @@ if [ -z "${TEST_TYPE:-}" ]; then
 fi
 log "Test type: $TEST_TYPE"
 
-MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$BUILDKITE_PULL_REQUEST IMAGE_VERSION=$IMAGE_VERSION"
+EFFECTIVE_PR=${BUILDKITE_PULL_REQUEST:-false}
+if [ "$EFFECTIVE_PR" = "false" ] && [ -n "${PR_NUMBER:-}" ]; then
+    EFFECTIVE_PR=$PR_NUMBER
+fi
+MODAL_ENV="BUILDKITE_REPO=$BUILDKITE_REPO BUILDKITE_COMMIT=$BUILDKITE_COMMIT BUILDKITE_PULL_REQUEST=$EFFECTIVE_PR IMAGE_VERSION=$IMAGE_VERSION"
 
 case "$TEST_TYPE" in
     "encoder")

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,3 +1,22 @@
+<!--
+PR TITLE: Must start with a type tag, e.g.:
+  [feat] Add new model       [bugfix] Fix VAE tiling      [refactor] Restructure pipeline
+  [perf] Optimize kernel     [ci] Update tests             [docs] Add guide
+  [misc] Cleanup configs     [new-model] Port Flux2
+
+MERGE WORKFLOW:
+  1. Ensure pre-commit passes and you have at least 1 approval
+  2. Comment /merge (or add the "ready" label) to enter the Merge Queue
+  3. Full Test Suite runs automatically on a staging branch → auto-merge on success
+
+ON-DEMAND TESTING (write access required):
+  /test full       — Full Test Suite        /test ssim        — SSIM regression
+  /test training   — Training pipeline      /test encoder     — Encoder tests
+  /test transformer — Transformer tests     /test vae         — VAE tests
+  /test kernel     — CUDA kernel tests      /test unit        — Unit tests
+  See docs/contributing/pull_requests.md for all 17 test commands
+-->
+
 ## Purpose
 
 <!-- What does this PR do? Link the related issue if applicable. -->

.github/mergify.yml

@@ -1,168 +1,191 @@
+merge_protections:
+  - name: PR merge requirements
+    if:
+      - base = main
+    success_conditions:
+      - "title~=(?i)^\\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model)\\]"
+      - check-success~=pre-commit
+      - check-success=fastcheck-passed
+
 pull_request_rules:
 
   # ============================================================
-  # Auto-labeling by file path
+  # Type labels (from PR title prefix)
   # ============================================================
 
-  - name: label attention changes
+  - name: "label type: feat"
+    conditions:
+      - "title~=(?i)^\\[(feat|feature)\\]"
+      - -closed
+    actions:
+      label:
+        add: ["type: feat"]
+
+  - name: "label type: bugfix"
+    conditions:
+      - "title~=(?i)^\\[(bug)?fix\\]"
+      - -closed
+    actions:
+      label:
+        add: ["type: bugfix"]
+
+  - name: "label type: refactor"
+    conditions:
+      - "title~=(?i)^\\[refactor\\]"
+      - -closed
+    actions:
+      label:
+        add: ["type: refactor"]
+
+  - name: "label type: perf"
+    conditions:
+      - "title~=(?i)^\\[perf\\]"
+      - -closed
+    actions:
+      label:
+        add: ["type: perf"]
+
+  - name: "label type: ci"
+    conditions:
+      - "title~=(?i)^\\[ci\\]"
+      - -closed
+    actions:
+      label:
+        add: ["type: ci"]
+
+  - name: "label type: docs"
+    conditions:
+      - "title~=(?i)^\\[(doc|docs)\\]"
+      - -closed
+    actions:
+      label:
+        add: ["type: docs"]
+
+  - name: "label type: misc"
+    conditions:
+      - "title~=(?i)^\\[(misc|chore)\\]"
+      - -closed
+    actions:
+      label:
+        add: ["type: misc"]
+
+  - name: "label type: new-model"
+    conditions:
+      - "title~=(?i)^\\[new.?model\\]"
+      - -closed
+    actions:
+      label:
+        add: ["type: new-model"]
+
+  # ============================================================
+  # Scope labels (from changed files)
+  # ============================================================
+
+  - name: "label scope: training"
     conditions:
       - or:
-        - files~=^fastvideo/attention/
-        - files~=^csrc/attn/
+          - files~=^fastvideo/train/
+          - files~=^fastvideo/training/
+          - files~=^fastvideo/distillation/
+          - files~=^examples/train/
+          - files~=^examples/training/
+          - files~=^examples/distill/
       - -closed
-      - label!=stale
     actions:
       label:
-        add: [attention]
+        add: ["scope: training"]
 
-  - name: label encoder changes
+  - name: "label scope: inference"
     conditions:
-      - files~=^fastvideo/models/encoders/
+      - or:
+          - files~=^fastvideo/pipelines/basic/
+          - files~=^fastvideo/pipelines/stages/
+          - files~=^fastvideo/pipelines/samplers/
+          - files~=^fastvideo/entrypoints/
+          - files~=^fastvideo/worker/
+          - files~=^fastvideo/configs/sample/
+          - files~=^fastvideo/configs/pipelines/
+          - files~=^examples/inference/
       - -closed
-      - label!=stale
     actions:
       label:
-        add: [encoder]
+        add: ["scope: inference"]
 
-  - name: label vae changes
+  - name: "label scope: attention"
     conditions:
-      - files~=^fastvideo/models/vaes/
+      - files~=^fastvideo/attention/
       - -closed
-      - label!=stale
     actions:
       label:
-        add: [vae]
+        add: ["scope: attention"]
 
-  - name: label transformer/DiT changes
+  - name: "label scope: kernel"
     conditions:
-      - files~=^fastvideo/models/dits/
+      - or:
+          - files~=^fastvideo-kernel/
+          - files~=^csrc/
       - -closed
-      - label!=stale
     actions:
       label:
-        add: [transformer]
+        add: ["scope: kernel"]
 
-  - name: label training changes
+  - name: "label scope: data"
     conditions:
-      - files~=^fastvideo/training/
+      - or:
+          - files~=^fastvideo/dataset/
+          - files~=^fastvideo/pipelines/preprocess/
+          - files~=^examples/preprocessing/
       - -closed
-      - label!=stale
     actions:
       label:
-        add: [training]
+        add: ["scope: data"]
 
-  - name: label pipeline changes
+  - name: "label scope: infra"
     conditions:
-      - files~=^fastvideo/pipelines/
+      - or:
+          - files~=^\.github/
+          - files~=^\.buildkite/
+          - files~=^fastvideo/tests/
+          - files~=^docker/
       - -closed
-      - label!=stale
     actions:
       label:
-        add: [pipeline]
+        add: ["scope: infra"]
 
-  - name: label distributed changes
+  - name: "label scope: distributed"
     conditions:
       - files~=^fastvideo/distributed/
       - -closed
-      - label!=stale
     actions:
       label:
-        add: [distributed]
+        add: ["scope: distributed"]
 
-  - name: label kernel changes
+  - name: "label scope: docs"
+    conditions:
+      - files~=^docs/
+      - -closed
+    actions:
+      label:
+        add: ["scope: docs"]
+
+  - name: "label scope: ui"
+    conditions:
+      - files~=^ui/
+      - -closed
+    actions:
+      label:
+        add: ["scope: ui"]
+
+  - name: "label scope: model"
     conditions:
       - or:
-        - files~=^csrc/
-        - files~=^fastvideo-kernel/
+          - files~=^fastvideo/models/
+          - files~=^fastvideo/layers/
+          - files~=^fastvideo/configs/models/
       - -closed
-      - label!=stale
     actions:
       label:
-        add: [kernel]
-
-  - name: label data/dataset changes
-    conditions:
-      - or:
-        - files~=^fastvideo/dataset/
-        - files~=^fastvideo/data_preprocess/
-      - -closed
-      - label!=stale
-    actions:
-      label:
-        add: [data]
-
-  - name: label CLI/entrypoint changes
-    conditions:
-      - files~=^fastvideo/entrypoints/
-      - -closed
-      - label!=stale
-    actions:
-      label:
-        add: [cli]
-
-  - name: label documentation changes
-    conditions:
-      - or:
-        - files~=^docs/
-        - files~=\.md$
-      - -closed
-      - label!=stale
-    actions:
-      label:
-        add: [documentation]
-
-  - name: label CI changes
-    conditions:
-      - or:
-        - files~=^\.github/
-        - files~=^\.buildkite/
-        - files~=^docker/
-      - -closed
-      - label!=stale
-    actions:
-      label:
-        add: [ci]
-
-  - name: label test changes
-    conditions:
-      - or:
-        - files~=^fastvideo/tests/
-        - files~=^tests/
-      - -closed
-      - label!=stale
-    actions:
-      label:
-        add: [tests]
-
-  - name: label ComfyUI changes
-    conditions:
-      - files~=^comfyui/
-      - -closed
-      - label!=stale
-    actions:
-      label:
-        add: [comfyui]
-
-  - name: label layer changes
-    conditions:
-      - files~=^fastvideo/layers/
-      - -closed
-      - label!=stale
-    actions:
-      label:
-        add: [layers]
-
-  - name: label LoRA changes
-    conditions:
-      - or:
-        - files~=^fastvideo/layers/lora/
-        - files~=^fastvideo/tests/lora/
-        - title~=(?i)lora
-      - -closed
-      - label!=stale
-    actions:
-      label:
-        add: [lora]
+        add: ["scope: model"]
 
   # ============================================================
   # Pre-commit failure help comment
@@ -197,6 +220,7 @@ pull_request_rules:
 
           For future commits, `pre-commit` will run automatically on changed files before each commit.
 
+
   # ============================================================
   # Merge conflict detection
   # ============================================================
@@ -230,14 +254,17 @@ pull_request_rules:
         remove: [needs-rebase]
 
   # ============================================================
-  # Auto-rebase stale PRs
+  # Auto-merge and auto-rebase
   # ============================================================
 
-  - name: auto-merge when ready and approved
+  - name: auto-merge when ready and all checks pass
     conditions:
       - label=ready
+      - "title~=(?i)^\\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model)\\]"
       - "#approved-reviews-by>=1"
       - check-success~=pre-commit
+      - check-success=fastcheck-passed
+      - check-success=full-suite-passed
       - -conflict
       - -closed
       - -draft
@@ -245,11 +272,50 @@ pull_request_rules:
       merge:
         method: squash
 
-  - name: auto-rebase on ready label
+  - name: auto-rebase when ready and Full Suite passed
     conditions:
       - label=ready
       - "#approved-reviews-by>=1"
+      - check-success=full-suite-passed
       - -conflict
       - -closed
+      - -draft
     actions:
       rebase: {}
+
+  - name: remove ready label on Full Suite failure
+    conditions:
+      - label=ready
+      - check-failure=full-suite-passed
+    actions:
+      label:
+        remove: [ready]
+
+  # ============================================================
+  # PR title format help
+  # ============================================================
+
+  - name: comment on invalid PR title format
+    conditions:
+      - -closed
+      - -draft
+      - "-title~=(?i)^\\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model)\\]"
+    actions:
+      comment:
+        message: |
+          ## ⚠️ PR title format required
+
+          Your PR title must start with a type tag in brackets. Examples:
+          - `[feat] Add new model support`
+          - `[bugfix] Fix VAE tiling corruption`
+          - `[refactor] Restructure training pipeline`
+          - `[perf] Optimize attention kernel`
+          - `[ci] Update test infrastructure`
+          - `[docs] Add inference guide`
+          - `[misc] Clean up configs`
+          - `[new-model] Port Flux2 to FastVideo`
+
+          Valid tags: `feat`, `feature`, `bugfix`, `fix`, `refactor`, `perf`, `ci`, `doc`, `docs`, `misc`, `chore`, `kernel`, `new-model`
+
+          Please update your PR title and the merge protection check will pass automatically.
+

.github/scripts/runpod_api.py

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

.github/scripts/runpod_cleanup.py

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

.github/workflows/ci-precommit.yml

@@ -3,8 +3,6 @@ name: pre-commit
 on:
   pull_request:
     branches: [main]
-  push:
-    branches: [main]
   workflow_call:
 
 concurrency:
@@ -16,6 +14,7 @@ permissions:
 
 jobs:
   pre-commit:
+    if: github.event_name == 'workflow_call' || github.event.pull_request.draft != true
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4

.github/workflows/ci-slash-commands.yml

@@ -0,0 +1,230 @@
+name: Slash Commands
+
+on:
+  issue_comment:
+    types: [created]
+
+permissions:
+  contents: read
+  pull-requests: write
+
+jobs:
+  handle-merge:
+    if: >-
+      github.event.issue.pull_request != null
+      && startsWith(github.event.comment.body, '/merge')
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check write permission
+        id: perm
+        uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
+        with:
+          script: |
+            const { data: perm } = await github.rest.repos.getCollaboratorPermissionLevel({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              username: context.payload.comment.user.login,
+            });
+            const hasWrite = ['admin', 'write'].includes(perm.permission);
+            if (!hasWrite) {
+              core.setFailed(`User ${context.payload.comment.user.login} lacks write permission (has: ${perm.permission}).`);
+            }
+            core.setOutput('has_write', String(hasWrite));
+
+      - name: Add ready label and react
+        if: steps.perm.outputs.has_write == 'true'
+        uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
+        with:
+          script: |
+            const owner = context.repo.owner;
+            const repo = context.repo.repo;
+            const prNumber = context.payload.issue.number;
+            // Remove ready first to allow re-trigger (labeled event fires on add, not if already present)
+            try { await github.rest.issues.removeLabel({ owner, repo, issue_number: prNumber, name: 'ready' }); } catch {}
+            await github.rest.issues.addLabels({ owner, repo, issue_number: prNumber, labels: ['ready'] });
+            await github.rest.reactions.createForIssueComment({
+              owner, repo,
+              comment_id: context.payload.comment.id,
+              content: 'rocket',
+            });
+  parse-command:
+    if: >-
+      github.event.issue.pull_request != null
+      && startsWith(github.event.comment.body, '/test')
+    runs-on: ubuntu-latest
+    outputs:
+      test_type: ${{ steps.parse.outputs.test_type }}
+      test_scope: ${{ steps.parse.outputs.test_scope }}
+      full_suite: ${{ steps.parse.outputs.full_suite }}
+      pr_sha: ${{ steps.pr.outputs.sha }}
+      pr_branch: ${{ steps.pr.outputs.branch }}
+      has_write: ${{ steps.perm.outputs.has_write }}
+    steps:
+      - name: Check write permission
+        id: perm
+        uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
+        with:
+          script: |
+            const { data: perm } = await github.rest.repos.getCollaboratorPermissionLevel({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              username: context.payload.comment.user.login,
+            });
+            const hasWrite = ['admin', 'write'].includes(perm.permission);
+            core.setOutput('has_write', String(hasWrite));
+            if (!hasWrite) {
+              core.info(`User ${context.payload.comment.user.login} lacks write permission — ignoring.`);
+            }
+
+      - name: Parse /test command
+        id: parse
+        if: steps.perm.outputs.has_write == 'true'
+        shell: bash
+        env:
+          COMMENT: ${{ github.event.comment.body }}
+        run: |
+          set -euo pipefail
+          TEST_NAME=$(echo "$COMMENT" | grep -oP '(?<=/test\s)\S+' | head -1 || true)
+
+          VALID="encoder vae transformer kernel unit ssim training lora-inference lora-training distillation self-forcing vsa vmoba performance api full fastcheck pre-commit"
+          if [ -z "$TEST_NAME" ] || ! echo "$VALID" | grep -qw "$TEST_NAME"; then
+            echo "Unknown test: '$TEST_NAME'. Valid: $VALID"
+            exit 1
+          fi
+
+          declare -A MAP=(
+            [encoder]=encoder [vae]=vae [transformer]=transformer
+            [kernel]=kernel_tests [unit]=unit_test
+            [ssim]=ssim [training]=training
+            [lora-inference]=inference_lora [lora-training]=training_lora
+            [distillation]=distillation_dmd [self-forcing]=self_forcing
+            [vsa]=training_vsa [vmoba]=inference_vmoba
+            [performance]=performance [api]=api_server
+          )
+
+          if [ "$TEST_NAME" = "full" ]; then
+            {
+              echo "test_type=all"
+              echo "test_scope=full"
+              echo "full_suite=true"
+            } >> "$GITHUB_OUTPUT"
+          elif [ "$TEST_NAME" = "fastcheck" ]; then
+            {
+              echo "test_type=fastcheck"
+              echo "test_scope=fastcheck"
+              echo "full_suite=false"
+            } >> "$GITHUB_OUTPUT"
+          elif [ "$TEST_NAME" = "pre-commit" ]; then
+            {
+              echo "test_type="
+              echo "test_scope=precommit"
+              echo "full_suite=false"
+            } >> "$GITHUB_OUTPUT"
+          else
+            {
+              echo "test_type=${MAP[$TEST_NAME]}"
+              echo "test_scope=direct"
+              echo "full_suite=false"
+            } >> "$GITHUB_OUTPUT"
+          fi
+
+      - name: Get PR details
+        id: pr
+        if: steps.perm.outputs.has_write == 'true'
+        uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
+        with:
+          script: |
+            const { data: pr } = await github.rest.pulls.get({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              pull_number: context.payload.issue.number,
+            });
+            core.setOutput('sha', pr.head.sha);
+            core.setOutput('branch', pr.head.ref);
+
+  pre-commit:
+    needs: parse-command
+    if: >-
+      needs.parse-command.outputs.has_write == 'true'
+      && needs.parse-command.outputs.test_scope == 'precommit'
+    uses: ./.github/workflows/ci-precommit.yml
+
+  post-precommit-status:
+    needs: [parse-command, pre-commit]
+    if: always() && needs.parse-command.outputs.test_scope == 'precommit'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Post commit status
+        uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
+        env:
+          PR_SHA: ${{ needs.parse-command.outputs.pr_sha }}
+          RESULT: ${{ needs.pre-commit.result }}
+        with:
+          script: |
+            const state = process.env.RESULT === 'success' ? 'success' : 'failure';
+            await github.rest.repos.createCommitStatus({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              sha: process.env.PR_SHA,
+              state,
+              context: 'pre-commit',
+              description: `Triggered via /test pre-commit (${state})`,
+            });
+
+  trigger-buildkite:
+    needs: parse-command
+    if: >-
+      needs.parse-command.outputs.has_write == 'true'
+      && needs.parse-command.outputs.test_type != ''
+    runs-on: ubuntu-latest
+    steps:
+      - name: React to comment
+        uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
+        with:
+          script: |
+            await github.rest.reactions.createForIssueComment({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              comment_id: context.payload.comment.id,
+              content: 'rocket',
+            });
+
+      - name: Trigger Buildkite
+        env:
+          BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
+          PR_SHA: ${{ needs.parse-command.outputs.pr_sha }}
+          PR_BRANCH: ${{ needs.parse-command.outputs.pr_branch }}
+          PR_NUMBER: ${{ github.event.issue.number }}
+          TEST_SCOPE: ${{ needs.parse-command.outputs.test_scope }}
+          FULL_SUITE: ${{ needs.parse-command.outputs.full_suite }}
+          TEST_TYPE: ${{ needs.parse-command.outputs.test_type }}
+          BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
+          BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
+        run: |
+          curl -sS --fail-with-body -X POST \
+            "https://api.buildkite.com/v2/organizations/${BK_ORG}/pipelines/${BK_PIPELINE}/builds" \
+            -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
+            -H "Content-Type: application/json" \
+            --data-raw "$(jq -n \
+              --arg commit "$PR_SHA" \
+              --arg branch "$PR_BRANCH" \
+              --arg message "/test ${TEST_TYPE} on PR #${PR_NUMBER}" \
+              --argjson pr_id "$PR_NUMBER" \
+              --arg test_scope "$TEST_SCOPE" \
+              --arg full_suite "$FULL_SUITE" \
+              --arg test_type "$TEST_TYPE" \
+              --arg pr_number "$PR_NUMBER" \
+              '{
+                commit: $commit,
+                branch: $branch,
+                message: $message,
+                ignore_pipeline_branch_filters: true,
+                pull_request_id: $pr_id,
+                pull_request_base_branch: "main",
+                env: {
+                  TEST_SCOPE: $test_scope,
+                  FULL_SUITE: $full_suite,
+                  TEST_TYPE: $test_type,
+                  PR_NUMBER: $pr_number
+                }
+              }')"

.github/workflows/ci-trigger-full-suite.yml

@@ -0,0 +1,83 @@
+name: Trigger Full Suite
+
+on:
+  pull_request:
+    types: [labeled, synchronize]
+
+permissions:
+  contents: read
+  pull-requests: read
+
+concurrency:
+  group: full-suite-${{ github.event.pull_request.number }}
+  cancel-in-progress: true
+
+jobs:
+  trigger:
+    if: >-
+      (github.event.action == 'labeled' && github.event.label.name == 'ready')
+      || github.event.action == 'synchronize'
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check ready label
+        id: check
+        uses: actions/github-script@60a0d83039c74a4aee543508d2ffcb1c3799cdea # v7.0.1
+        with:
+          script: |
+            const { data: pr } = await github.rest.pulls.get({
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              pull_number: context.payload.pull_request.number,
+            });
+            const hasReady = pr.labels.some(l => l.name === 'ready');
+            core.setOutput('has_ready', String(hasReady));
+            if (!hasReady) core.info('No ready label — skipping Full Suite trigger.');
+
+      - name: Cancel previous Buildkite builds
+        if: steps.check.outputs.has_ready == 'true'
+        env:
+          BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
+          PR_BRANCH: ${{ github.event.pull_request.head.ref }}
+        run: |
+          # Find running builds for this branch with TEST_SCOPE=full and cancel them
+          builds=$(curl -sS -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
+            "https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds?branch=${PR_BRANCH}&state=running,scheduled" \
+            | jq -r '.[] | select(.env.TEST_SCOPE == "full") | .number')
+          for build_num in $builds; do
+            echo "Cancelling Buildkite build #$build_num"
+            curl -sS -X PUT -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
+              "https://api.buildkite.com/v2/organizations/${{ vars.BUILDKITE_ORG_SLUG }}/pipelines/${{ vars.BUILDKITE_PIPELINE_SLUG }}/builds/${build_num}/cancel"
+          done
+
+      - name: Trigger Buildkite Full Suite
+        if: steps.check.outputs.has_ready == 'true'
+        env:
+          BUILDKITE_API_TOKEN: ${{ secrets.BUILDKITE_API_TOKEN }}
+          PR_SHA: ${{ github.event.pull_request.head.sha }}
+          PR_BRANCH: ${{ github.event.pull_request.head.ref }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          BK_ORG: ${{ vars.BUILDKITE_ORG_SLUG }}
+          BK_PIPELINE: ${{ vars.BUILDKITE_PIPELINE_SLUG }}
+        run: |
+          curl -sS --fail-with-body -X POST \
+            "https://api.buildkite.com/v2/organizations/${BK_ORG}/pipelines/${BK_PIPELINE}/builds" \
+            -H "Authorization: Bearer $BUILDKITE_API_TOKEN" \
+            -H "Content-Type: application/json" \
+            --data-raw "$(jq -n \
+              --arg commit "$PR_SHA" \
+              --arg branch "$PR_BRANCH" \
+              --arg message "Full Suite for PR #${PR_NUMBER}" \
+              --argjson pr_id "$PR_NUMBER" \
+              '{
+                commit: $commit,
+                branch: $branch,
+                message: $message,
+                ignore_pipeline_branch_filters: true,
+                pull_request_id: $pr_id,
+                pull_request_base_branch: "main",
+                env: {
+                  TEST_SCOPE: "full",
+                  FULL_SUITE: "true",
+                  PR_NUMBER: ($pr_id | tostring)
+                }
+              }')"

.github/workflows/community-issue-labeler.yml

@@ -22,16 +22,28 @@ jobs:
             const labels = [];
 
             const rules = [
-              { keywords: ['training', 'finetune', 'fine-tune', 'lora', 'fsdp', 'distill'], label: 'training' },
-              { keywords: ['inference', 'generate', 'pipeline', 'slow', 'latency'], label: 'inference' },
-              { keywords: ['attention', 'vsa', 'flash', 'sta', 'vmoba', 'sparse attn'], label: 'attention' },
+              // scope labels (shared with PR labeling via Mergify)
+              // Mapping: label → repo directories
+              //   scope: training    → fastvideo/train/, fastvideo/training/, fastvideo/distillation/
+              //   scope: inference   → fastvideo/pipelines/, fastvideo/entrypoints/, fastvideo/worker/
+              //   scope: attention   → fastvideo/attention/
+              //   scope: kernel      → fastvideo-kernel/, csrc/
+              //   scope: model       → fastvideo/models/, fastvideo/layers/, fastvideo/configs/models/
+              //   scope: data        → fastvideo/dataset/, fastvideo/pipelines/preprocess/
+              //   scope: distributed → fastvideo/distributed/
+              //   scope: docs        → docs/
+              { keywords: ['training', 'finetune', 'fine-tune', 'lora', 'fsdp', 'distill'], label: 'scope: training' },
+              { keywords: ['inference', 'generate', 'pipeline', 'slow', 'latency'], label: 'scope: inference' },
+              { keywords: ['attention', 'vsa', 'flash', 'sta', 'vmoba', 'sparse attn'], label: 'scope: attention' },
+              { keywords: ['kernel', 'csrc', 'cuda kernel', 'thunderkittens'], label: 'scope: kernel' },
+              { keywords: ['wan', 'hunyuan', 'mochi', 'ltx', 'cogvideo', 'flux', 'sd3', 'cosmos'], label: 'scope: model' },
+              { keywords: ['dataset', 'dataloader', 'preprocessing', 'preprocess'], label: 'scope: data' },
+              { keywords: ['distributed', 'sequence parallel', 'fsdp', 'tensor parallel', 'multi-node', 'multi-gpu'], label: 'scope: distributed' },
+              { keywords: ['docs', 'documentation', 'tutorial', 'example'], label: 'scope: docs' },
+              // issue-only labels (cross-module, no single repo directory)
               { keywords: ['install', 'setup', 'pip', 'cuda', 'uv ', 'import error', 'modulenotfound'], label: 'installation' },
-              { keywords: ['docs', 'documentation', 'tutorial', 'example'], label: 'documentation' },
               { keywords: ['memory', 'oom', 'out of memory', 'gpu memory', 'vram'], label: 'performance' },
-              { keywords: ['wan', 'hunyuan', 'mochi', 'ltx', 'cogvideo'], label: 'model-specific' },
-              { keywords: ['windows', 'macos', 'mac os', 'apple', 'mps'], label: 'platform' },
-              { keywords: ['comfyui', 'comfy'], label: 'comfyui' },
-              { keywords: ['kernel', 'csrc', 'cuda kernel', 'thunderkittens'], label: 'kernel' },
+              { keywords: ['windows', 'macos', 'mac os', 'apple', 'mps', 'rocm', 'amd', 'npu'], label: 'platform' },
             ];
 
             for (const rule of rules) {

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

@@ -32,7 +32,7 @@ permissions:
 jobs:
   build-python-3-10:
     if: ${{ github.event.inputs.python_3_10 == 'true' }}
-    uses: ./.github/workflows/build-image-template.yml
+    uses: ./.github/workflows/_template-build-image.yml
     with:
       python_version: '3.10'
       dockerfile_path: docker/Dockerfile.python3.10
@@ -41,7 +41,7 @@ jobs:
 
   build-python-3-11:
     if: ${{ github.event.inputs.python_3_11 == 'true' }}
-    uses: ./.github/workflows/build-image-template.yml
+    uses: ./.github/workflows/_template-build-image.yml
     with:
       python_version: '3.11'
       dockerfile_path: docker/Dockerfile.python3.11
@@ -50,7 +50,7 @@ jobs:
 
   build-python-3-12:
     if: ${{ github.event.inputs.python_3_12 == 'true' }}
-    uses: ./.github/workflows/build-image-template.yml
+    uses: ./.github/workflows/_template-build-image.yml
     with:
       python_version: '3.12'
       dockerfile_path: docker/Dockerfile.python3.12
@@ -59,9 +59,9 @@ jobs:
 
   build-python-3-12-cuda-12-9:
     if: ${{ github.event.inputs.python_3_12_cuda_12_9 == 'true' }}
-    uses: ./.github/workflows/build-image-template.yml
+    uses: ./.github/workflows/_template-build-image.yml
     with:
       python_version: '3.12'
       dockerfile_path: docker/Dockerfile.python3.12.cuda12.9.1
       tag_suffix: py3.12-cuda12.9.1
-    secrets: inherit
+    secrets: inherit

.github/workflows/infra-docs.yml

@@ -7,14 +7,14 @@ on:
       - 'docs/**'
       - 'mkdocs.yml'
       - 'requirements-mkdocs.txt'
-      - '.github/workflows/docs.yml'
+      - '.github/workflows/infra-docs.yml'
   pull_request:
     branches: [ main ]
     paths:
       - 'docs/**'
       - 'mkdocs.yml'
       - 'requirements-mkdocs.txt'
-      - '.github/workflows/docs.yml'
+      - '.github/workflows/infra-docs.yml'
 
 permissions:
   contents: read

.github/workflows/pr-test.yml

@@ -1,334 +0,0 @@
-name: PR Test
-
-on:
-  push:
-    branches: [main]
-    paths:
-      - "fastvideo/**/*.py"
-      - ".github/workflows/pr-test.yml"
-  pull_request:
-    branches: [main]
-    types: [opened, ready_for_review, synchronize, reopened]
-    paths:
-      - "fastvideo/**/*.py"
-      - ".github/workflows/pr-test.yml"
-      - "pyproject.toml"
-      - "docker/Dockerfile.python3.12"
-      - "csrc/**"
-  workflow_dispatch:
-    inputs:
-      run_encoder_test:
-        description: "Run encoder-test"
-        required: false
-        default: false
-        type: boolean
-      run_vae_test:
-        description: "Run vae-test"
-        required: false
-        default: false
-        type: boolean
-      run_transformer_test:
-        description: "Run transformer-test"
-        required: false
-        default: false
-        type: boolean
-      run_ssim_test:
-        description: "Run ssim-test"
-        required: false
-        default: false
-        type: boolean
-      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_precision_test_VSA:
-        description: "Run precision-test-VSA"
-        required: false
-        default: false
-        type: boolean
-      run_unit_test:
-        description: "Run unit-test"
-        required: false
-        default: false
-        type: boolean
-
-env:
-  PYTHONUNBUFFERED: "1"
-
-
-concurrency:
-  group: pr-test-${{ github.ref }}
-  cancel-in-progress: true
-
-jobs:
-  change-filter:
-    runs-on: ubuntu-latest
-    if: ${{ github.event.pull_request.draft == false || github.event_name == 'workflow_dispatch' }}
-    outputs:
-      encoder-test: ${{ steps.filter.outputs.encoder-test }}
-      vae-test: ${{ steps.filter.outputs.vae-test }}
-      transformer-test: ${{ steps.filter.outputs.transformer-test }}
-      training-test: ${{ steps.filter.outputs.training-test }}
-      training-test-VSA: ${{ steps.filter.outputs.training-test-VSA }}
-      precision-test-VSA: ${{ steps.filter.outputs.precision-test-VSA }}
-      unit-test: ${{ steps.filter.outputs.unit-test }}
-    steps:
-      - uses: actions/checkout@v4
-      - uses: dorny/paths-filter@v3
-        id: filter
-        with:
-          filters: |
-            # Define reusable path patterns
-            common-paths: &common-paths
-              - 'pyproject.toml'
-              - 'docker/Dockerfile.python3.10'
-              - 'docker/Dockerfile.python3.11'
-              - 'docker/Dockerfile.python3.12'
-            vsa-kernel-paths: &vsa-kernel-paths
-              - 'csrc/attn/video_sparse_attn/**'
-              - 'csrc/attn/video_sparse_attn/tk/**'
-              - 'csrc/attn/video_sparse_attn/setup.py'
-              - 'csrc/attn/video_sparse_attn/config_vsa.py'
-              - 'csrc/attn/video_sparse_attn/vsa.cpp'
-            vsa-paths: &vsa-paths
-              - 'fastvideo/**'
-              - *common-paths
-              - *vsa-kernel-paths
-            
-            # Actual tests
-            encoder-test:
-              - 'fastvideo/models/encoders/**'
-              - 'fastvideo/models/loader/**'
-              - 'fastvideo/tests/encoders/**'
-              - *common-paths
-            vae-test:
-              - 'fastvideo/models/vaes/**'
-              - 'fastvideo/models/loader/**'
-              - 'fastvideo/tests/vaes/**'
-              - *common-paths
-            transformer-test:
-              - 'fastvideo/models/dits/**'
-              - 'fastvideo/models/loader/**'
-              - 'fastvideo/tests/transformers/**'
-              - 'fastvideo/layers/**'
-              - 'fastvideo/attention/**'
-              - *common-paths
-            training-test:
-              - 'fastvideo/**'
-              - *common-paths
-            training-test-VSA:
-              - 'fastvideo/**'
-              - *common-paths
-              - *vsa-kernel-paths
-            precision-test-VSA:
-              - *common-paths
-              - *vsa-kernel-paths
-            unit-test:
-              - 'fastvideo/**'
-              - *common-paths
-
-  encoder-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.encoder-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_encoder_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "encoder-test"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 1
-      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/encoders -s"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  vae-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.vae-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_vae_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "vae-test"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 1
-      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/vaes -s"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  transformer-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.transformer-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_transformer_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "transformer-test"
-      gpu_type: "NVIDIA L40S"
-      gpu_count: 1
-      volume_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/tests/transformers -s"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-      
-  ssim-test:
-    needs: change-filter
-    if: >-
-      github.event_name != 'workflow_dispatch' || (github.event_name == 'workflow_dispatch' && github.event.inputs.run_ssim_test == 'true')
-    strategy:
-      fail-fast: false
-      matrix:
-        python-version: [
-          # {version: "3.10", tag: "latest"},
-          # {version: "3.11", tag: "py3.11-latest"},
-          {version: "3.12", tag: "py3.12-latest"}
-        ]
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "ssim-test-py${{ matrix.python-version.version }}"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 2
-      volume_size: 200
-      disk_size: 200
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:${{ matrix.python-version.tag }}"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/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' && needs.change-filter.outputs.training-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "training-test"
-      gpu_type: "NVIDIA A40"
-      gpu_count: 4
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/tests/training/Vanilla -srP"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-      WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
-
-  training-test-VSA:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.training-test-VSA == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_training_test_VSA == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "training-test-VSA"
-      gpu_type: "NVIDIA H100 NVL"
-      gpu_count: 2
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/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 }}
-
-  precision-test-VSA:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.precision-test-VSA == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_precision_test_VSA == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "precision-test-VSA"
-      gpu_type: "NVIDIA H100 NVL"
-      gpu_count: 1
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && python csrc/attn/tests/test_vsa.py"
-      timeout_minutes: 30
-    secrets:
-      RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-      RUNPOD_PRIVATE_KEY: ${{ secrets.RUNPOD_PRIVATE_KEY }}
-
-  unit-test:
-    needs: change-filter
-    if: >-
-      (github.event_name != 'workflow_dispatch' && needs.change-filter.outputs.unit-test == 'true') || 
-      (github.event_name == 'workflow_dispatch' && github.event.inputs.run_unit_test == 'true')
-    uses: ./.github/workflows/runpod-test.yml
-    with:
-      job_id: "unit-test"
-      gpu_type: "NVIDIA L40S"
-      gpu_count: 1
-      volume_size: 100
-      disk_size: 100
-      image: "ghcr.io/${{ github.repository }}/fastvideo-dev:py3.12-latest"
-      test_command: "uv pip install -e .[test] && pytest ./fastvideo/dataset/ -vs && pytest ./fastvideo/workflow/ -vs && pytest ./fastvideo/entrypoints/ -vs"
-      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 }}/fastvideo-dev:py3.12-latest"
-  #     test_command: "wandb login $WANDB_API_KEY && uv pip install -e .[test] && pytest ./fastvideo/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:
-    # Add other jobs to this list as you create them
-    needs: [encoder-test, vae-test, transformer-test, ssim-test, training-test, training-test-VSA, precision-test-VSA] 
-    if: ${{ always() && ((github.event_name != 'workflow_dispatch' && github.event.pull_request.draft == false) || github.event_name == 'workflow_dispatch') }}
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: "3.10"
-
-      - name: Install dependencies
-        run: pip install requests
-
-      - name: Cleanup all RunPod instances
-        env:
-          JOB_IDS: '["encoder-test", "vae-test", "transformer-test", "ssim-test-py3.10", "ssim-test-py3.11", "ssim-test-py3.12", "training-test", "training-test-VSA", "precision-test-VSA"]'
-          RUNPOD_API_KEY: ${{ secrets.RUNPOD_API_KEY }}
-          GITHUB_RUN_ID: ${{ github.run_id }}
-        run: python .github/scripts/runpod_cleanup.py

.github/workflows/runpod-test.yml

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

.github/workflows/test.yml

@@ -1,31 +0,0 @@
-name: Run Tests
-
-on:
-  push:
-    branches: [ main ]
-  pull_request:
-    branches: [ main ]
-
-jobs:
-  test:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check out repository
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.12'  # or any version you need
-
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip setuptools wheel
-          pip install torch
-          pip install packaging ninja
-          pip install -e .
-          pip install pytest
-
-      - name: Run Pytest
-        run: |
-          pytest --ignore csrc/attn/test

.github/workflows/trigger-full-suite.yml

@@ -1,26 +0,0 @@
-name: Trigger Full Test Suite
-
-on:
-  pull_request:
-    types: [labeled]
-
-jobs:
-  trigger-buildkite:
-    if: >-
-      github.repository == 'hao-ai-lab/FastVideo'
-      && github.event.label.name == 'ready'
-    runs-on: ubuntu-latest
-    permissions:
-      pull-requests: read
-    steps:
-      - name: Trigger Buildkite build with FULL_SUITE
-        uses: buildkite/trigger-pipeline-action@v2.4.1
-        with:
-          buildkite_api_access_token: ${{ secrets.BUILDKITE_API_TOKEN }}
-          pipeline: "${{ vars.BUILDKITE_ORG_SLUG }}/${{ vars.BUILDKITE_PIPELINE_SLUG }}"
-          branch: "${{ github.event.pull_request.head.ref }}"
-          commit: "${{ github.event.pull_request.head.sha }}"
-          message: "Full suite for PR #${{ github.event.pull_request.number }}"
-          send_pull_request: "true"
-          pull_request_base_branch: "${{ github.event.pull_request.base.ref }}"
-          build_env_vars: '{"FULL_SUITE": "true"}'

.github/workflows/vsa-publish.yml

@@ -1,257 +0,0 @@
-name: Publish Video Sparse Attention Kernel to PyPI on Version Change
-
-on:
-  push:
-    branches:
-      - main
-    paths:
-      - "csrc/attn/video_sparse_attn/setup.py"
-  workflow_dispatch:
-
-jobs:
-  check-version-change:
-    runs-on: ubuntu-latest
-    outputs:
-      version-changed: ${{ steps.check-version.outputs.changed }}
-      new-version: ${{ steps.check-version.outputs.new-version }}
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-        with:
-          fetch-depth: 2
-
-      - name: Check if version changed
-        id: check-version
-        run: |
-          cd csrc/attn/video_sparse_attn
-          # Get current commit's version
-          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.py | grep -oP 'VERSION\s*=\s*"\K[^"]+' || echo "0.0.0")
-          echo "Old version: $OLD_VERSION"
-          
-          if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then
-            echo "Version changed from $OLD_VERSION to $NEW_VERSION"
-            echo "changed=true" >> $GITHUB_OUTPUT
-            echo "new-version=$NEW_VERSION" >> $GITHUB_OUTPUT
-          else
-            echo "Version did not change"
-            echo "changed=false" >> $GITHUB_OUTPUT
-          fi
-
-  build_wheels:
-    name: Build Wheel
-    needs: check-version-change
-    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
-    runs-on: ${{ matrix.os }}
-
-    strategy:
-      fail-fast: false
-      matrix:
-          # Using ubuntu-20.04 instead of 22.04 for more compatibility (glibc). Ideally we'd use the
-          # manylinux docker image, but I haven't figured out how to install CUDA on manylinux.
-          os: [ubuntu-22.04]
-          python-version: ['3.10', '3.11', '3.12', '3.13']
-          # For version reference https://pytorch.org/get-started/previous-versions/
-          torch-cuda:
-            - torch-version: '2.5.1'
-              cuda-version: '12.4.1'
-              torch-cuda-short: 'cu124'
-            - torch-version: '2.6.0'
-              cuda-version: '12.6.3'
-              torch-cuda-short: 'cu126'
-            - torch-version: '2.7.1'
-              cuda-version: '12.8.0'
-              torch-cuda-short: 'cu128'
-
-    steps:
-      - name: Free up disk space
-        run: |
-          echo "Initial disk space:"
-          df -h
-          
-          # Remove large directories
-          sudo rm -rf /usr/share/dotnet
-          sudo rm -rf /usr/local/lib/android
-          sudo rm -rf /opt/ghc
-          sudo rm -rf /usr/local/share/boost
-          sudo rm -rf /usr/share/swift
-          sudo rm -rf /usr/local/lib/node_modules
-          sudo rm -rf /usr/local/share/powershell
-          sudo rm -rf /usr/share/rust
-          sudo rm -rf /usr/local/.ghcup
-          
-          # Remove cached files
-          sudo rm -rf /var/lib/apt/lists/*
-          sudo rm -rf /var/cache/apt/archives/*
-          
-          echo "Disk space after cleanup:"
-          df -h
-
-      - name: Checkout
-        uses: actions/checkout@v4
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python-version }}
-
-      - name: Install CUDA ${{ matrix.torch-cuda.cuda-version }}
-        uses: Jimver/cuda-toolkit@v0.2.21
-        id: cuda-toolkit
-        with:
-          cuda: ${{ matrix.torch-cuda.cuda-version }}
-          linux-local-args: '["--toolkit"]'
-          method: 'network'
-
-      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
-        run: |
-          sudo apt update
-          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
-          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-          # Allow Git to Access Safe Directory
-          git config --global --add safe.directory /__w/FastVideo/FastVideo
-
-          # Set CUDA environment variables
-          export CUDA_HOME=/usr/local/cuda-${{ matrix.torch-cuda.cuda-version }}
-          export PATH=${CUDA_HOME}/bin:${PATH}
-          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-          
-          # Verify installation
-          gcc --version
-          g++ --version
-          clang-11 --version
-          nvcc --version
-
-      - name: Install PyTorch ${{ matrix.torch-cuda.torch-version }}+cu${{ matrix.torch-cuda.cuda-version }}
-        run: |
-          pip install --upgrade pip
-          # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
-          # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
-          pip install typing-extensions==4.12.2
-          # We want to figure out the CUDA version to download pytorch
-          # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
-          # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
-          pip install --no-cache-dir torch==${{ matrix.torch-cuda.torch-version }} --index-url https://download.pytorch.org/whl/${{matrix.torch-cuda.torch-cuda-short}}
-          nvcc --version
-          python --version
-          python -c "import torch; print('PyTorch:', torch.__version__)"
-          python -c "import torch; print('CUDA:', torch.version.cuda)"
-          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
-
-      - 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/video_sparse_attn  # Move into the correct folder
-          git submodule update --init --recursive # Ensure ThunderKittens submodule is initialized
-          python setup.py bdist_wheel --dist-dir=dist
-
-      - name: Rename wheel file
-        run: |
-          cd csrc/attn/video_sparse_attn
-
-          CUDA_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.cuda-version }} | cut -d. -f1,2 | sed 's/\.//g')
-          TORCH_SHORT_VERSION=$(echo ${{ matrix.torch-cuda.torch-version }} | cut -d. -f1,2)
-          # Get the correct version format
-          tmpname=cu${CUDA_SHORT_VERSION}torch${TORCH_SHORT_VERSION}
-          wheel_name=$(ls dist/*whl | xargs -n 1 basename | sed "s/-/+$tmpname-/2")
-          # Rename with version information
-          ls dist/*whl |xargs -I {} mv {} dist/${wheel_name}
-          echo "wheel_name=${wheel_name}" >> $GITHUB_ENV
-
-      - name: Upload wheel artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: ${{ env.wheel_name }}
-          path: csrc/attn/video_sparse_attn/dist/*.whl
-          retention-days: 90
-
-  publish_package:
-    name: Publish package
-    needs: [build_wheels, check-version-change]
-    if: ${{ needs.check-version-change.outputs.version-changed == 'true' || github.event_name == 'workflow_dispatch' }}
-    runs-on: ubuntu-22.04
-    permissions:
-      id-token: write  # Needed for OIDC Trusted Publishing
-
-    steps:
-      - uses: actions/checkout@v4
-
-      - uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-
-      - name: Install CUDA 12.4.1
-        uses: Jimver/cuda-toolkit@v0.2.21
-        id: cuda-toolkit
-        with:
-          cuda: 12.4.1
-          linux-local-args: '["--toolkit"]'
-          method: 'network'
-          sub-packages: '["nvcc"]'
-
-      - name: Install dependencies (GCC, Clang, CUDA Paths, Git)
-        run: |
-          sudo apt update
-          sudo apt install -y git patchelf gcc-11 g++-11 clang-11
-          sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
-
-          # Allow Git to Access Safe Directory
-          git config --global --add safe.directory /__w/FastVideo/FastVideo
-
-          # Set CUDA environment variables
-          export CUDA_HOME=/usr/local/cuda-12.4.1
-          export PATH=${CUDA_HOME}/bin:${PATH}
-          export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
-          
-          # Verify installation
-          gcc --version
-          g++ --version
-          clang-11 --version
-          nvcc --version
-
-      - name: Install PyTorch 2.5.1+cu12.4.1
-        run: |
-          pip install --upgrade pip
-          # With python 3.13 and torch 2.5.1, unless we update typing-extensions, we get error
-          # AttributeError: attribute '__default__' of 'typing.ParamSpec' objects is not writable
-          pip install typing-extensions==4.12.2
-          # We want to figure out the CUDA version to download pytorch
-          # e.g. we can have system CUDA version being 11.7 but if torch==1.12 then we need to download the wheel from cu116
-          # see https://github.com/pytorch/pytorch/blob/main/RELEASE.md#release-compatibility-matrix
-          export TORCH_CUDA_VERSION=124
-          pip install --no-cache-dir torch==2.5.1 --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION}
-          nvcc --version
-          python --version
-          python -c "import torch; print('PyTorch:', torch.__version__)"
-          python -c "import torch; print('CUDA:', torch.version.cuda)"
-          python -c "from torch.utils import cpp_extension; print (cpp_extension.CUDA_HOME)"
-
-      - 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/video_sparse_attn  # Move into the correct folder
-          git submodule update --init --recursive # 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/video_sparse_attn/dist/

.gitignore

@@ -85,8 +85,9 @@ docs/distillation/examples/
 dmd_t2v_output/
 preprocess_output_text/
 
-# Next.js / Node artifacts under ui/: see ui/.gitignore
 
 .claude/
 .codex/
+.sisyphus/
 openspec/
+fastvideo/tests/ssim/reference_videos/**

.pre-commit-config.yaml

@@ -18,10 +18,8 @@ exclude: |
     fastvideo/train\.py|
     fastvideo/utils/.*|
     examples/.*|
-    .github/workflows/fastvideo-publish.yml|
-    .github/workflows/sta-publish.yml|
-    .github/workflows/vsa-publish.yml|
-    .github/workflows/build-image-template.yml|
+    .github/workflows/publish-fastvideo.yml|
+    .github/workflows/_template-build-image.yml|
     docs/source/inference/support_matrix.md
   )
 repos:

.python-version

@@ -1 +0,0 @@
-3.12

docs/README.md

@@ -36,4 +36,4 @@ Then open your browser to: http://localhost:8000
 
 ## Automatic Deployment
 
-Documentation is automatically built and deployed to GitHub Pages when changes are pushed to the `main` branch via the `.github/workflows/docs.yml` workflow.
+Documentation is automatically built and deployed to GitHub Pages when changes are pushed to the `main` branch via the `.github/workflows/infra-docs.yml` workflow.

docs/contributing/ci_architecture.md

@@ -0,0 +1,306 @@
+# CI Architecture
+
+## Overview
+
+FastVideo uses a three-tier CI pipeline designed to keep feedback fast on every push while
+protecting `main` through a full GPU regression suite before any merge.
+
+```
+PR push
+  │
+  ├─► Tier 1: Pre-commit (~2 min)
+  │     GitHub Actions / ubuntu-latest
+  │     yapf, ruff, mypy, codespell, pymarkdown, actionlint, check-filenames
+  │
+  └─► Tier 2: Fastcheck (~10-20 min, path-filtered)
+        Buildkite / Modal GPU instances
+        Only runs tests for paths you changed
+
+              │ (developer comments /merge or maintainer adds 'ready' label)
+              ▼
+        Tier 3: Full Suite (~60-90 min)
+          Buildkite / Modal GPU instances
+          All integration, SSIM, training, and performance tests
+          Runs on the PR branch directly
+              │
+          pass ──► Mergify auto-squash-merges to main, branch deleted
+          fail ──► Mergify removes 'ready' label; fix and /merge again
+```
+
+---
+
+## CI Tiers
+
+### Tier 1: Pre-commit (every PR push)
+
+| Attribute | Value |
+|-----------|-------|
+| Triggered by | Every push to any PR branch, plus pushes to `main` |
+| Runs on | GitHub Actions, `ubuntu-latest` |
+| Duration | ~2 minutes |
+
+**Checks run** (from `.pre-commit-config.yaml`, stage: `manual`):
+
+| Hook | What it checks |
+|------|---------------|
+| `yapf` | Python code formatting |
+| `ruff` | Python linting and auto-fixable style issues |
+| `codespell` | Spelling errors in code and docs |
+| `pymarkdown` | Markdown formatting |
+| `actionlint` | GitHub Actions workflow syntax |
+| `mypy` | Static type checking (Python 3.10 target) |
+| `check-filenames` | No spaces in tracked filenames |
+
+A failure here means code style or type issues. Run `pre-commit run --all-files` locally to
+replicate CI results before pushing.
+
+---
+
+### Tier 2: Fastcheck (path-filtered, every PR push)
+
+| Attribute | Value |
+|-----------|-------|
+| Triggered by | Every push; the monorepo-diff plugin skips tests for unchanged paths |
+| Runs on | Buildkite, Modal GPU instances |
+| Duration | ~10-20 minutes per test (run in parallel) |
+
+**Tests and their path triggers:**
+
+| Buildkite label | `TEST_TYPE` | Triggers when you change |
+|-----------------|-------------|--------------------------|
+| Encoder Tests | `encoder` | `fastvideo/models/encoders/**`, `fastvideo/models/loader/**`, `fastvideo/tests/encoders/**`, `pyproject.toml`, `docker/Dockerfile.python3.12` |
+| VAE Tests | `vae` | `fastvideo/models/vaes/**`, `fastvideo/models/loader/**`, `fastvideo/tests/vaes/**`, `pyproject.toml`, `docker/Dockerfile.python3.12` |
+| Transformer Tests | `transformer` | `fastvideo/models/dits/**`, `fastvideo/models/loader/**`, `fastvideo/tests/transformers/**`, `fastvideo/layers/**`, `fastvideo/attention/**`, `pyproject.toml`, `docker/Dockerfile.python3.12` |
+| Kernel Tests | `kernel_tests` | `fastvideo-kernel/**`, `pyproject.toml`, `docker/Dockerfile.python3.12` |
+| Unit Tests | `unit_test` | `fastvideo/**`, `.buildkite/**`, `.github/**`, `pyproject.toml`, `docker/Dockerfile.python3.12` |
+
+A Fastcheck failure means a component-level regression. Check the Buildkite build log for the
+failing test's output.
+
+---
+
+### Tier 3: Full Test Suite (triggered by `ready` label)
+
+| Attribute | Value |
+|-----------|-------|
+| Triggered by | Adding the `ready` label to the PR (via `/merge` command), or a `/test full` command |
+| Runs on | Buildkite, Modal GPU instances |
+| Duration | 60-90 minutes total (tests run in parallel, path-filtered) |
+
+**Tests:**
+
+| Buildkite label | `TEST_TYPE` | Timeout |
+|-----------------|-------------|---------|
+| SSIM Tests | `ssim` | 90 min |
+| LoRA Inference Tests | `inference_lora` | 20 min |
+| Training Tests | `training` | 15 min |
+| Distillation DMD Tests | `distillation_dmd` | 15 min |
+| Self-Forcing Tests | `self_forcing` | 15 min |
+| LoRA Training Tests | `training_lora` | 15 min |
+| Training Tests VSA | `training_vsa` | 15 min |
+| Inference Tests VMoBA | `inference_vmoba` | 15 min |
+| Performance Tests | `performance` | 30 min |
+| API Server Tests | `api_server` | 30 min |
+
+A Full Suite failure removes the `ready` label automatically. A Mergify comment links to
+the Buildkite build. Fix the regression, push, and comment `/merge` again.
+
+---
+
+## Auto-merge Flow
+
+Mergify prevents untested code from landing on `main` by gating squash-merge on the Full
+Suite passing directly on the PR branch.
+
+**How it works:**
+
+1. A developer comments `/merge` on an approved PR (or a maintainer adds the `ready` label).
+2. The `ready` label triggers `ci-trigger-full-suite.yml`, which calls the Buildkite API to
+   run the Full Suite on the PR branch itself.
+3. While the Full Suite runs, Mergify also auto-rebases the PR branch against `main` if it
+   is behind and has no conflicts.
+4. Once the Full Suite posts `full-suite-passed`, Mergify checks all **merge conditions**:
+   - `pre-commit` check is green
+   - `fastcheck-passed` check is green
+   - `full-suite-passed` check is green
+   - At least 1 approved review (`#approved-reviews-by>=1`)
+   - PR title starts with a valid `[type]` tag
+   - PR is not a draft
+   - No merge conflicts
+5. If all conditions pass, Mergify squash-merges to `main` automatically. The branch is
+   deleted after merge.
+6. If the Full Suite fails, Mergify removes the `ready` label and posts a comment linking to
+   the Buildkite build. The developer fixes the issue, pushes, and comments `/merge` again.
+
+**Merge conditions summary:**
+
+| Condition | Meaning |
+|-----------|---------|
+| `check-success~=pre-commit` | Tier 1 pre-commit must be green |
+| `check-success=fastcheck-passed` | Tier 2 Fastcheck must be green |
+| `check-success=full-suite-passed` | Tier 3 Full Suite must be green |
+| `#approved-reviews-by>=1` | At least one approved review |
+| `title~=(?i)^\[(feat|bugfix|...)` | PR title has a valid type tag |
+| `-draft` | PR is not in draft state |
+| `-conflict` | No merge conflicts with base branch |
+| `-closed` | PR is still open |
+
+---
+
+## Label System
+
+Labels are applied automatically. You don't need to set them manually.
+
+### Type Labels (from PR title prefix)
+
+Applied by Mergify based on the `[tag]` at the start of the PR title.
+
+| Label | Matched title prefix | Meaning |
+|-------|---------------------|---------|
+| `type: feat` | `[feat]` or `[feature]` | New feature or capability |
+| `type: bugfix` | `[bugfix]` or `[fix]` | Bug fix |
+| `type: refactor` | `[refactor]` | Code restructuring, no behavior change |
+| `type: perf` | `[perf]` | Performance improvement |
+| `type: ci` | `[ci]` | CI/CD or tooling changes |
+| `type: docs` | `[doc]` or `[docs]` | Documentation only |
+| `type: misc` | `[misc]` or `[chore]` | Housekeeping, dependency bumps |
+| `type: new-model` | `[new-model]` | Adding a new model |
+
+### Scope Labels (from changed files)
+
+Applied by Mergify based on which paths you modified. Multiple scope labels can be added.
+
+| Label | File paths that trigger it |
+|-------|---------------------------|
+| `scope: training` | `fastvideo/train/`, `fastvideo/training/`, `fastvideo/distillation/`, `examples/train/`, `examples/training/`, `examples/distill/` |
+| `scope: inference` | `fastvideo/pipelines/basic/`, `fastvideo/pipelines/stages/`, `fastvideo/pipelines/samplers/`, `fastvideo/entrypoints/`, `fastvideo/worker/`, `fastvideo/configs/sample/`, `fastvideo/configs/pipelines/`, `examples/inference/` |
+| `scope: attention` | `fastvideo/attention/` |
+| `scope: kernel` | `fastvideo-kernel/`, `csrc/` |
+| `scope: data` | `fastvideo/dataset/`, `fastvideo/pipelines/preprocess/`, `examples/preprocessing/` |
+| `scope: infra` | `.github/`, `.buildkite/`, `fastvideo/tests/`, `docker/` |
+| `scope: distributed` | `fastvideo/distributed/` |
+| `scope: docs` | `docs/` |
+| `scope: ui` | `ui/` |
+| `scope: model` | `fastvideo/models/`, `fastvideo/layers/`, `fastvideo/configs/models/` |
+
+### Process Labels
+
+| Label | Who sets it | Meaning |
+|-------|-------------|---------|
+| `ready` | Developer (`/merge` command) or maintainer | Triggers Full Suite and enables auto-merge |
+| `needs-rebase` | Mergify (automatic) | PR has merge conflicts; rebase needed |
+| `do-not-merge` | Maintainer | Blocks queue entry regardless of other conditions |
+
+---
+
+## PR Title Format
+
+All PR titles targeting `main` must start with a bracketed type tag. This is enforced by a
+Mergify merge protection rule and is required before a PR can be squash-merged.
+
+**Format:**
+
+```
+[type] Short description
+```
+
+**Valid type tags:**
+
+`feat`, `feature`, `bugfix`, `fix`, `refactor`, `perf`, `ci`, `doc`, `docs`, `misc`, `chore`,
+`kernel`, `new-model`
+
+**Valid examples:**
+
+```
+[feat] Add causal Wan pipeline
+[bugfix] Fix VAE temporal tiling corruption
+[refactor] Restructure training framework
+[perf] Optimize FlashAttention kernel dispatch
+[docs] Add inference guide for LoRA
+[new-model] Port HunyuanVideo 1.5 to FastVideo
+```
+
+**Invalid examples (will block merge):**
+
+```
+Add causal Wan pipeline          ← missing type tag
+FEAT: Add pipeline               ← wrong format
+feat: Add pipeline               ← square brackets required
+```
+
+If your title is invalid, Mergify posts a comment explaining the required format and the
+merge protection check will remain failed until you update the title.
+
+---
+
+## Slash Commands
+
+Slash commands let contributors and maintainers trigger CI actions directly from PR comments.
+**Write permission to the repository is required.**
+
+The command is recognized within a few seconds. The workflow reacts with a 🚀 emoji to confirm.
+
+### `/merge`
+
+```
+/merge
+```
+
+Adds the `ready` label to the PR, which triggers the Full Suite on your PR branch and
+enables Mergify to auto-squash-merge once all conditions pass.
+
+The command first removes the `ready` label if it is already present, then re-adds it. This
+ensures the `labeled` event fires and a fresh Full Suite build is started even on a re-try.
+
+### `/test <name>`
+
+```
+/test <name>
+```
+
+Triggers a specific Buildkite test or suite on the current PR branch.
+
+| Command | Runs | Maps to `TEST_TYPE` |
+|---------|------|---------------------|
+| `/test encoder` | Encoder Tests (Fastcheck) | `encoder` |
+| `/test vae` | VAE Tests (Fastcheck) | `vae` |
+| `/test transformer` | Transformer Tests (Fastcheck) | `transformer` |
+| `/test kernel` | Kernel Tests (Fastcheck) | `kernel_tests` |
+| `/test unit` | Unit Tests (Fastcheck) | `unit_test` |
+| `/test ssim` | SSIM regression tests | `ssim` |
+| `/test training` | Training pipeline tests | `training` |
+| `/test lora-inference` | LoRA inference tests | `inference_lora` |
+| `/test lora-training` | LoRA training tests | `training_lora` |
+| `/test distillation` | DMD distillation tests | `distillation_dmd` |
+| `/test self-forcing` | Self-Forcing tests | `self_forcing` |
+| `/test vsa` | VSA training tests | `training_vsa` |
+| `/test vmoba` | VMoBA inference tests | `inference_vmoba` |
+| `/test performance` | Performance benchmarks | `performance` |
+| `/test api` | API server integration tests | `api_server` |
+| `/test full` | Entire Full Suite | all (with `TEST_SCOPE=full`) |
+| `/test fastcheck` | Entire Fastcheck suite | fastcheck (with `TEST_SCOPE=fastcheck`) |
+
+---
+
+## Auto Branch Cleanup
+
+After a PR is squash-merged to `main`, Mergify automatically deletes the head branch.
+Protected branches (`main`, `master`, `release/*`) are never deleted.
+
+---
+
+## Workflow File Reference
+
+| Filename | Trigger | What it does |
+|----------|---------|-------------|
+| `ci-precommit.yml` | Every push / PR against `main` | Runs pre-commit hooks (yapf, ruff, mypy, codespell, pymarkdown, actionlint, check-filenames) |
+| `ci-trigger-full-suite.yml` | `ready` label added to a PR | Calls Buildkite API to run Full Suite on the PR branch |
+| `ci-slash-commands.yml` | PR comment starting with `/merge` or `/test` | Handles slash commands; adds `ready` label or triggers Buildkite |
+| `community-issue-labeler.yml` | Issue opened or edited | Auto-labels issues by keyword matching against title and body |
+| `community-welcome.yml` | First contribution | Posts a welcome comment for first-time contributors |
+| `community-stale.yml` | Scheduled | Marks and closes stale issues and PRs |
+| `infra-build-image.yml` | Manual (`workflow_dispatch`) | Builds Docker images for CI |
+| `infra-docs.yml` | Changes to `docs/` merged to `main` | Builds and deploys documentation to GitHub Pages |
+| `publish-fastvideo.yml` | Version bump | Publishes `fastvideo` package to PyPI |
+| `publish-kernel.yml` | Version bump | Publishes `fastvideo-kernel` package to PyPI |
+| `publish-comfyui.yml` | Version bump | Publishes ComfyUI node package to PyPI |

docs/contributing/developer_env/runpod.md

@@ -1,65 +1,195 @@
+# RunPod Development Environment
 
-# 📦 Developing FastVideo on RunPod
+RunPod gives you on-demand cloud GPUs for FastVideo development. It's useful when you need a beefy GPU to test training runs, benchmark inference, or reproduce results without waiting for shared cluster time.
 
-You can easily use the FastVideo Pod Template on [RunPod](https://www.runpod.io) for development or experimentation.
+## Prerequisites
 
-## Creating a new pod
+- A [RunPod](https://www.runpod.io) account with billing configured
+- An SSH key pair. If you don't have one, generate it with `ssh-keygen -t ed25519`
+- Your public key (`~/.ssh/id_ed25519.pub`) ready to paste into RunPod
+
+## Step 1: Create a Pod
+
+**1. Verify your account**
+
+Make sure you're logged into the right RunPod account before spending credits.
 
-- Make sure you are using the correct RunPod account.
 ![RunPod Account Selection](../../assets/images/runpod_account.png)
 
-- Use "Additional Filters" to select CUDA 12.8.
+**2. Filter by CUDA version**
+
+Use "Additional Filters" to select CUDA 12.8.
+
 ![RunPod CUDA selection](../../assets/images/runpod_cuda.png)
 
-- Click "Deploy" and Pick a single A40 or RTX 4090 GPU.
+**3. Select a GPU**
+
+Click "Deploy" and pick a GPU. See [GPU Recommendations](#gpu-recommendations) below for guidance on which GPU to choose.
+
 ![RunPod GPU Selection](../../assets/images/runpod_deploy.png)
 
-- Select the "FastVideo" or "fastvideo-dev" Pod Template.
+**4. Pick the FastVideo template**
+
+Select the "FastVideo" or "fastvideo-dev" Pod Template. This pulls the pre-built image that includes all dependencies, Flash Attention, and a ready-to-use `uv` environment.
+
 ![RunPod Pod Template Selection](../../assets/images/runpod_create.png)
 
-- Set the Pod name to "`<name>-<FastVideo>-<date>`".
+**5. Name your pod**
+
+Use a memorable name like `yourname-fastvideo-2026-03-28`. This helps if you have multiple pods running.
+
+**6. Add a persistent volume (recommended)**
+
+Attach a network volume to `/root/.cache` or `/models` for storing downloaded model weights. Models can be 10-50 GB each, and re-downloading them every session wastes time and bandwidth.
+
+**7. Deploy**
+
+Click Deploy. The pod takes a few minutes to start while the image pulls. You'll see it transition to "Running" in your dashboard.
+
+## Step 2: Connect via SSH
+
+Once the pod is running, find the "SSH over exposed TCP" connection string in the pod dashboard.
 
-- Finally, once the pod is deployed (will take a few minutes as the image is being pulled), you can SSH into it using "SSH exposed over TCP". You'll need to use the matching private ssh key you provided.
 ![RunPod SSH](../../assets/images/runpod_ssh.png)
 
-## Working with the pod
-
-After SSH'ing into your pod, you'll find the correct `uv` environment already activated and you should be in /FastVideo/ directory. Make sure to use /FastVideo/ for all your work.
-
-To pull in the latest changes from the GitHub repo:
+Connect with:
 
 ```bash
-cd /FastVideo
-git pull
+ssh root@<pod-ip> -p <port> -i ~/.ssh/id_ed25519
 ```
 
-Run your development workflows as usual:
+RunPod also supports VS Code Remote SSH if you prefer an IDE.
 
-```bash
-# Run linters
-pre-commit run --all-files
+### Custom template (advanced)
 
-# Run tests
-pytest tests/
-```
-
-Make sure to push your changes back to the GitHub repo as nothing will be saved to the pod when it is terminated.
-
-After you are done with your work, you can terminate the pod by clicking the "Terminate" and "Delete" buttons. Remember if the pod is not completely deleted, Runpod will keep charging you for it.
-
-## Extra Information:
-If you need to customize the pod template this section has some useful information. For the most part you can leave the defaults of the FastVideo Pod Template.
-
-When creating your pod template, use this image:
+If you're setting up a pod from scratch instead of the FastVideo template, use this image:
 
 ```
 ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-latest
 ```
 
-Paste Container Start Command to support SSH ([RunPod Docs](https://docs.runpod.io/pods/configuration/use-ssh)):
+And paste this as the Container Start Command to enable SSH ([RunPod docs](https://docs.runpod.io/pods/configuration/use-ssh)):
 
 ```bash
 bash -c "apt update;DEBIAN_FRONTEND=noninteractive apt-get install openssh-server -y;mkdir -p ~/.ssh;cd $_;chmod 700 ~/.ssh;echo \"$PUBLIC_KEY\" >> authorized_keys;chmod 700 authorized_keys;service ssh start;sleep infinity"
 ```
 
 ![RunPod template configuration](../../assets/images/runpod_template.png)
+
+## Step 3: Set Up FastVideo
+
+After SSH'ing in, the `uv` virtual environment at `/opt/venv` is already activated (configured in `.bashrc` and `.profile`). You land in the `/FastVideo` directory.
+
+**Clone or pull the repo**
+
+If the pod already has the FastVideo repo:
+
+```bash
+cd /FastVideo
+git pull
+```
+
+If starting from a blank pod:
+
+```bash
+git clone https://github.com/hao-ai-lab/FastVideo.git /FastVideo
+cd /FastVideo
+```
+
+**Install the package**
+
+```bash
+uv pip install -e .[dev]
+```
+
+The Docker image already includes Flash Attention and most heavy dependencies, so this is fast.
+
+**Build the custom kernels (optional)**
+
+VSA and STA attention kernels aren't in the Docker image by default. Build them if you're working on attention backends or need maximum inference performance:
+
+```bash
+cd /FastVideo/fastvideo-kernel
+./build.sh
+```
+
+The build script detects your GPU architecture automatically. An A100 or H100 takes about 5-10 minutes.
+
+**Verify the setup**
+
+```bash
+cd /FastVideo
+python -c "import fastvideo; print('OK')"
+pytest tests/ -q --no-header
+```
+
+## Development Workflow
+
+### Editing code on RunPod
+
+Two common approaches:
+
+**Option A: Edit on RunPod directly**
+
+Use VS Code Remote SSH or `vim`/`nano` on the pod. Commit and push when you're ready:
+
+```bash
+cd /FastVideo
+git add .
+git commit -m "your change"
+git push
+```
+
+**Option B: Edit locally, sync to RunPod**
+
+Work in your local repo, then pull on the pod:
+
+```bash
+# On RunPod:
+cd /FastVideo
+git pull
+```
+
+This keeps your local tools (editor, linters) intact while running GPU workloads on the pod.
+
+### Running linters and tests
+
+```bash
+# Lint
+pre-commit run --all-files
+
+# Full test suite
+pytest tests/
+
+# Just package tests
+pytest fastvideo/tests/ -v
+```
+
+### Storing models
+
+If you attached a persistent volume, point your model downloads there:
+
+```bash
+export HF_HOME=/models/huggingface
+export TRANSFORMERS_CACHE=/models/huggingface
+```
+
+Add these to `/root/.bashrc` so they persist across SSH sessions. The volume survives pod termination, so you only download models once.
+
+### Terminating the pod
+
+When you're done, push any commits you want to keep. RunPod does not save pod storage after termination.
+
+Go to your RunPod dashboard, click "Terminate", then "Delete". A pod that's stopped but not deleted still charges you for storage. Fully delete it to stop all charges.
+
+## GPU Recommendations
+
+| GPU | VRAM | Good for |
+|-----|------|----------|
+| RTX 4090 | 24 GB | Inference testing, small model fine-tuning, quick iteration |
+| A40 | 48 GB | Mid-size training runs, 480p video generation |
+| A100 (40 GB) | 40 GB | Multi-GPU inference, training with sequence parallelism |
+| A100 (80 GB) | 80 GB | Large model training, 720p+ video generation |
+| H100 | 80 GB | Heavy training, benchmarking, kernel development |
+
+For most development work, a single RTX 4090 or A40 is sufficient and cost-effective. Use A100/H100 when you need to reproduce training results at scale or test multi-GPU features.

docs/contributing/pull_requests.md

@@ -0,0 +1,218 @@
+# Contributing via Pull Requests
+
+This guide walks through the PR workflow: title format, labels, CI pipeline, and getting
+your changes merged.
+
+---
+
+## PR Title Format (Required)
+
+Every PR targeting `main` must start with a type tag in square brackets. This is checked by
+Mergify before any merge is allowed.
+
+**Format:**
+
+```
+[type] Short description of the change
+```
+
+**Valid type tags:**
+
+| Tag | When to use |
+|-----|-------------|
+| `[feat]` or `[feature]` | New feature or capability |
+| `[bugfix]` or `[fix]` | Bug fix |
+| `[refactor]` | Code restructuring with no behavior change |
+| `[perf]` | Performance improvement |
+| `[ci]` | CI/CD or build tooling changes |
+| `[doc]` or `[docs]` | Documentation only |
+| `[misc]` or `[chore]` | Housekeeping, dependency bumps, minor cleanup |
+| `[kernel]` | CUDA kernel changes in `fastvideo-kernel/` |
+| `[new-model]` | Adding a new model or pipeline |
+
+**Examples:**
+
+```
+[feat] Add causal Wan 2.2 I2V pipeline
+[bugfix] Fix VAE temporal tiling corruption on H100
+[refactor] Restructure distributed attention dispatch
+[docs] Add LoRA finetuning guide
+[new-model] Port HunyuanVideo 1.5 to FastVideo
+```
+
+If your title is missing the tag, Mergify will post a comment listing the valid formats.
+Update the title and the check will re-evaluate automatically.
+
+---
+
+## Labels
+
+Labels are applied automatically based on your PR title and the files you changed. You don't
+need to set them manually.
+
+**Type label** — set from the `[tag]` in your title:
+`type: feat`, `type: bugfix`, `type: refactor`, `type: perf`, `type: ci`, `type: docs`,
+`type: misc`, `type: new-model`
+
+**Scope labels** — set from which files you modified (multiple labels can apply):
+`scope: training`, `scope: inference`, `scope: attention`, `scope: kernel`, `scope: data`,
+`scope: infra`, `scope: distributed`, `scope: docs`, `scope: ui`, `scope: model`
+
+**Process labels** — set during review and merge:
+
+| Label | Who sets it | Meaning |
+|-------|-------------|---------|
+| `ready` | You (`/merge` comment) or a maintainer | Triggers Full Suite and enables auto-merge |
+| `needs-rebase` | Mergify (automatic) | Your PR has conflicts; rebase against `main` |
+| `do-not-merge` | Maintainer | Blocks merge regardless of CI status |
+
+---
+
+## CI Pipeline
+
+Three tiers run automatically on every PR.
+
+**Tier 1: Pre-commit (~2 min) — runs on every push**
+
+GitHub Actions checks formatting, linting, type correctness, and spelling using pre-commit
+hooks: yapf, ruff, mypy, codespell, pymarkdown, actionlint, and check-filenames.
+
+**Tier 2: Fastcheck (~10-20 min) — runs on every push, path-filtered**
+
+Buildkite runs GPU tests only for the components you changed. If you only modified
+`fastvideo/models/vaes/`, only VAE Tests run. Tests run in parallel.
+
+**Tier 3: Full Suite (~60-90 min) — triggered by the `ready` label**
+
+When you comment `/merge` (or a maintainer adds the `ready` label), Buildkite runs the
+complete test suite on your PR branch: SSIM regression, LoRA inference and training,
+distillation, self-forcing, VSA, VMoBA, performance benchmarks, and API server tests.
+
+---
+
+## Getting Your PR Merged
+
+**Step-by-step:**
+
+1. Open a PR with a title that starts with a valid `[type]` tag.
+2. Push your changes. Pre-commit and Fastcheck run automatically.
+3. Fix any pre-commit failures locally (`pre-commit run --all-files`) and push again.
+4. Wait for at least one approving review.
+5. Once approved and pre-commit is green, comment `/merge` on the PR.
+6. The `ready` label is added, which triggers the Full Suite on your PR branch.
+7. Mergify also auto-rebases your branch against `main` if it is behind and conflict-free.
+8. If all Full Suite tests pass and all merge conditions are met (approval, valid title,
+   pre-commit green, fastcheck green, no draft, no conflicts), Mergify squash-merges to
+   `main` automatically. Your branch is deleted.
+9. If a Full Suite test fails, Mergify removes the `ready` label and posts a comment with a
+   link to the Buildkite build. Fix the issue, push, and comment `/merge` again.
+
+!!! note
+    Only contributors with write permission to the repository can trigger slash commands.
+    If you're an external contributor, ask a maintainer to run `/merge` or add the `ready`
+    label for you.
+
+---
+
+## Running Tests On Demand
+
+Comment on your PR to trigger specific tests independently of the auto-merge flow.
+
+**Trigger the entire Full Suite:**
+
+```
+/test full
+```
+
+**Trigger the Fastcheck suite:**
+
+```
+/test fastcheck
+```
+
+**Trigger individual tests:**
+
+```
+/test encoder          # Encoder component tests
+/test vae              # VAE component tests
+/test transformer      # Transformer / DiT tests
+/test kernel           # CUDA kernel tests
+/test unit             # Unit tests
+
+/test ssim             # SSIM regression tests
+/test training         # Training pipeline tests
+/test lora-inference   # LoRA inference tests
+/test lora-training    # LoRA training tests
+/test distillation     # DMD distillation tests
+/test self-forcing     # Self-Forcing distillation tests
+/test vsa              # VSA training tests
+/test vmoba            # VMoBA inference tests
+/test performance      # Performance benchmarks
+/test api              # API server integration tests
+```
+
+The workflow reacts with a 🚀 emoji to confirm the command was received.
+
+---
+
+## Troubleshooting
+
+### Pre-commit fails
+
+Run locally to reproduce and auto-fix:
+
+```bash
+# Install pre-commit if needed
+uv pip install pre-commit
+pre-commit install
+
+# Run all checks on all files
+pre-commit run --all-files
+```
+
+Common quick fixes:
+
+- **yapf**: `yapf -i <file>` (Python formatting)
+- **ruff**: `ruff check --fix <file>` (linting)
+- **codespell**: `codespell --write-changes <file>` (spelling)
+
+### PR title format check fails
+
+Update your title to start with a valid type tag. The Mergify merge protection check
+re-evaluates automatically after you save the title.
+
+Valid tags: `feat`, `feature`, `bugfix`, `fix`, `refactor`, `perf`, `ci`, `doc`, `docs`,
+`misc`, `chore`, `kernel`, `new-model`
+
+### My PR has merge conflicts (`needs-rebase` label)
+
+Rebase against `main` and force-push:
+
+```bash
+git fetch origin main
+git rebase origin/main
+# Resolve any conflicts, then:
+git push --force-with-lease
+```
+
+Mergify removes the `needs-rebase` label automatically once conflicts are resolved.
+
+### Full Suite failed after `/merge`
+
+The Full Suite found a regression. Mergify removes the `ready` label and posts a comment
+linking to the Buildkite build. Check the failing step's output for assertion errors or
+tracebacks.
+
+Common causes:
+
+- Test failures caused by your code changes
+- Missing dependency in `pyproject.toml`
+- GPU memory issue (some tests require specific hardware like L40S or H100)
+- Kernel build failure (if you changed `fastvideo-kernel/`)
+
+After fixing, push and comment `/merge` again.
+
+### I'm an external contributor without write permission
+
+You can't use slash commands directly. After your PR is approved, ask a maintainer to
+comment `/merge` or add the `ready` label.

docs/contributing/testing.md

@@ -23,13 +23,18 @@ SSIM tests are located in `fastvideo/tests/ssim`. These tests generate videos us
 
 ```
 fastvideo/tests/ssim/
-├── <GPU>_reference_videos/   # Reference videos organized by GPU type (e.g., L40S_reference_videos)
-│   ├── <Model_Name>/
-│   │   ├── <Backend>/        # e.g., FLASH_ATTN, TORCH_SDPA
-│   │   │   └── <Video_File>
+├── reference_videos/
+│   ├── default/
+│   │   └── <GPU>_reference_videos/
+│   │       ├── <Model_Name>/
+│   │       │   ├── <Backend>/        # e.g., FLASH_ATTN, TORCH_SDPA
+│   │       │   │   └── <Video_File>
+│   └── full_quality/
+│       └── <GPU>_reference_videos/
 ├── test_causal_similarity.py
-├── test_inference_similarity.py
-├── update_reference_videos.sh
+├── test_wan_t2v_similarity.py
+├── test_wan_i2v_similarity.py
+├── reference_videos_cli.py
 └── ...
 ```
 
@@ -37,7 +42,7 @@ fastvideo/tests/ssim/
 
 To add a new SSIM test, follow these steps:
 
-1. **Create or Update a Test File**: You can add a new test function to an existing file (like `test_inference_similarity.py`) or create a new one if testing a distinct category of models.
+1. **Create or Update a Test File**: Prefer model-specific files (for example `test_wan_t2v_similarity.py`) and create a new one when testing a distinct model or pipeline.
 
 2. **Define Model Parameters**: Define the configuration for the model you want to test. This includes model path, dimensions, inference steps, and other generation parameters. **Note:** Consider using lower `num_inference_steps` or reduced resolution (e.g., 480p instead of 720p) to keep test execution time reasonable, provided it doesn't compromise the test's ability to detect regression.
 
@@ -81,10 +86,15 @@ To add a new SSIM test, follow these steps:
    ```
 
 4. **Reference Videos**:
-   * When running the test for the first time (or when updating the reference), the test will fail because the reference video is missing. The generated video will be saved in `fastvideo/tests/ssim/generated_videos`.
+   * When running the test for the first time (or when updating the reference), the test will fail because the reference video is missing. The generated video will be saved in `fastvideo/tests/ssim/generated_videos/<quality-tier>/<GPU>_reference_videos`.
    * Inspect the generated video to ensure it meets quality expectations.
-   * Move the generated video to the appropriate reference folder: `fastvideo/tests/ssim/<GPU>_reference_videos/<Model>/<Backend>/`.
-   * You can use the helper script `update_reference_videos.sh` to automate copying videos from `generated_videos` to `L40S_reference_videos`. Note: Check the script to ensure paths match your environment (it defaults to `L40S_reference_videos`).
+   * Move the generated video to the appropriate quality/GPU reference folder:
+     `fastvideo/tests/ssim/reference_videos/<quality-tier>/<GPU>_reference_videos/<Model>/<Backend>/`.
+   * You can use the helper CLI to copy generated videos into a reference folder:
+     `python fastvideo/tests/ssim/reference_videos_cli.py copy-local --quality-tier default --reference-dir fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos`
+   * Upload/download can target both quality tiers and specific GPU folders:
+     `python fastvideo/tests/ssim/reference_videos_cli.py upload --quality-tier all`
+     `python fastvideo/tests/ssim/reference_videos_cli.py download --quality-tier full_quality --device-folder H200_reference_videos`
 
 ### Running Tests Locally
 
@@ -96,36 +106,113 @@ pytest fastvideo/tests/ssim/ -vs
 
 Ensure you have the necessary GPUs available as defined in your test parameters.
 
-## Modal Workflow
+## CI Integration
 
-FastVideo uses [Modal](https://modal.com/) for running tests in a CI environment. The workflow scripts are located in `fastvideo/tests/modal/`.
+FastVideo uses [Modal](https://modal.com/) for running tests in a CI environment. The
+workflow scripts are located in `fastvideo/tests/modal/`.
+
+### Buildkite Pipeline
+
+Tests are orchestrated by Buildkite (`.buildkite/pipeline.yml`) and executed on Modal GPU
+instances. The pipeline runs in two modes:
+
+**Fastcheck** — runs on every PR push, path-filtered. Only tests for the components you
+changed are triggered. Tests run in parallel.
+
+**Full Suite** — runs when a PR enters the Merge Queue (or when triggered manually via
+`/test full`). Covers SSIM regression, training, distillation, inference, and performance.
 
 ### `pr_test.py`
 
-The main entry point for CI tests is `fastvideo/tests/modal/pr_test.py`. This script defines Modal functions that execute the pytest suites on specific hardware (e.g., L40S, H100).
+The main entry point for CI tests is `fastvideo/tests/modal/pr_test.py`. This script defines
+Modal functions that execute the pytest suites on specific hardware (e.g., L40S, H100).
 
 ### Updating Modal Configuration
 
 If you add a new test that requires:
+
 * **Different GPU Hardware**: You may need to change the `@app.function(gpu=...)` decorator.
 * **Longer Execution Time**: Increase the `timeout` parameter.
-* **New Environment Variables/Secrets**: Add them to `secrets=[...]` or the image environment. For example, if your model is gated on Hugging Face, ensure `HF_API_KEY` is passed.
+* **New Environment Variables/Secrets**: Add them to `secrets=[...]` or the image
+  environment. For example, if your model is gated on Hugging Face, ensure `HF_API_KEY`
+  is passed.
 
-For SSIM tests, the `run_ssim_tests` function in `pr_test.py` currently runs:
+For SSIM tests, use `fastvideo/tests/modal/ssim_test.py`:
 
-```python
-@app.function(gpu="L40S:2", image=image, timeout=2700, secrets=[modal.Secret.from_dict({"HF_API_KEY": os.environ.get("HF_API_KEY", "")})])
-def run_ssim_tests():
-    run_test("hf auth login --token $HF_API_KEY && pytest ./fastvideo/tests/ssim -vs")
+```bash
+python -m modal run fastvideo/tests/modal/ssim_test.py::run_ssim_tests
 ```
 
-If your new test file is inside `fastvideo/tests/ssim`, it will automatically be picked up by this command. However, ensure that the `gpu="L40S:2"` configuration is sufficient for your model. If your model requires more GPUs (e.g., 4 or 8), you might need to create a separate Modal function or update the existing one.
+Target specific SSIM files/models:
+
+```bash
+python -m modal run fastvideo/tests/modal/ssim_test.py::run_ssim_tests \
+  --test-files test_wan_t2v_similarity.py \
+  --model-ids Wan2.1-T2V-1.3B-Diffusers
+```
+
+If HF token env vars are not set (`HF_API_KEY` / `HUGGINGFACE_HUB_TOKEN` /
+`HF_TOKEN`), the local entrypoint fails fast. To export raw `generated_videos`
+from Modal to the shared volume:
+
+```bash
+python -m modal run fastvideo/tests/modal/ssim_test.py::run_ssim_tests \
+  --sync-generated-to-volume
+```
+
+The raw export path is quality-tiered:
+
+* default params: `ssim_generated_videos/default/<subdir>/generated_videos`
+* full-quality params: `ssim_generated_videos/full_quality/<subdir>/generated_videos`
+
+The printed `modal volume get` command also downloads into a quality-specific
+local directory under `./generated_videos_modal/<quality-tier>`.
+
+To turn downloaded Modal outputs into local reference videos, use the matching
+quality tier with `copy-local`, for example:
+
+```bash
+python fastvideo/tests/ssim/reference_videos_cli.py copy-local \
+  --quality-tier full_quality \
+  --generated-dir ./generated_videos_modal/full_quality/L40S_reference_videos \
+  --device-folder L40S_reference_videos
+```
 
 ### Workflow Scripts
 
-The shell script that triggers these tests in the CI pipeline is located at `.buildkite/scripts/pr_test.sh`. If you add a new test category (e.g., a new folder outside of `ssim`), you will need to:
+The shell script that triggers tests in CI is `.buildkite/scripts/pr_test.sh`. If you add
+a new test category (e.g., a new folder outside of `ssim`), you will need to:
+
 1. Add a new function in `fastvideo/tests/modal/pr_test.py`.
 2. Add a new case in `.buildkite/scripts/pr_test.sh` to handle the new test type.
 
 !!! note
-    If you are a maintainer, you'll need to finally manually update the workflow script in Buildkite. Otherwise, a maintainer will help you update.
+    If you are a maintainer, update the workflow script in Buildkite after merging. Otherwise,
+    ask a maintainer for help.
+
+## Triggering Tests via Slash Commands
+
+Maintainers and contributors with write access can trigger individual test suites directly
+from a PR comment. The workflow reacts with a 🚀 emoji to confirm the command was received.
+
+```
+/test ssim              # SSIM regression tests
+/test training          # Training pipeline tests
+/test lora-training     # LoRA training tests
+/test lora-inference    # LoRA inference tests
+/test distillation      # DMD distillation tests
+/test self-forcing      # Self-Forcing tests
+/test vsa               # VSA training tests
+/test vmoba             # VMoBA inference tests
+/test performance       # Performance benchmarks
+/test api               # API server integration tests
+/test encoder           # Encoder component tests (Fastcheck)
+/test vae               # VAE component tests (Fastcheck)
+/test transformer       # Transformer / DiT tests (Fastcheck)
+/test kernel            # CUDA kernel tests (Fastcheck)
+/test unit              # Unit tests (Fastcheck)
+/test full              # Entire Full Suite
+/test fastcheck         # Entire Fastcheck suite
+```
+
+See [CI Architecture](ci_architecture.md) for the complete reference.

examples/train/configs/overfit_hunyuan_t2v.yaml

@@ -0,0 +1,78 @@
+# HunyuanVideo T2V overfitting test config.
+#
+# Overfits on 3 short videos (480x832, 77 frames) to verify the
+# training pipeline works end-to-end.
+#
+# Preprocess data first:
+#   CUDA_VISIBLE_DEVICES=0 python scripts/preprocess_hunyuan_overfit.py
+#
+# Run:
+#   bash examples/train/run.sh examples/train/configs/overfit_hunyuan_t2v.yaml
+
+models:
+  student:
+    _target_: fastvideo.train.models.hunyuan.HunyuanModel
+    init_from: hunyuanvideo-community/HunyuanVideo
+    trainable: true
+    enable_gradient_checkpointing_type: full
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/hunyuan_overfit_preprocessed
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 20
+    num_height: 480
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 300
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/hunyuan_overfit
+    training_state_checkpointing_steps: 50
+    checkpoints_total_limit: 2
+
+  tracker:
+    project_name: fastvideo_hunyuan
+    run_name: hunyuan_overfit
+
+  model:
+    precondition_outputs: true
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.hunyuan.hunyuan_pipeline.HunyuanVideoPipeline
+    dataset_file: data/hunyuan_overfit_preprocessed/validation_prompts.json
+    every_steps: 50
+    sampling_steps: [50]
+    guidance_scale: 1.0
+
+pipeline:
+  flow_shift: 7

fastvideo-kernel/pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
 ]
 dependencies = [
     "torch>=2.5.0",
-    "triton>=2.0.0; sys_platform == 'linux'",
+    "triton>=2.0.0",
 ]
 
 [project.urls]

fastvideo/entrypoints/video_generator.py

@@ -52,27 +52,16 @@ class VideoGenerator:
     customization options, similar to popular frameworks like HF Diffusers.
     """
 
-    def __init__(
-        self,
-        fastvideo_args: FastVideoArgs,
-        executor_class: type[Executor],
-        log_stats: bool,
-        *,
-        log_queue=None,
-    ):
+    def __init__(self, fastvideo_args: FastVideoArgs, executor_class: type[Executor], log_stats: bool):
         """
         Initialize the video generator.
-
+        
         Args:
             fastvideo_args: The inference arguments
             executor_class: The executor class to use for inference
-            log_stats: Whether to log statistics
-            log_queue: Optional multiprocessing.Queue to forward worker logs to
         """
         self.fastvideo_args = fastvideo_args
-        self.executor = executor_class(
-            fastvideo_args, log_queue=log_queue
-        )
+        self.executor = executor_class(fastvideo_args)
 
     @classmethod
     def from_pretrained(cls, model_path: str, **kwargs) -> "VideoGenerator":
@@ -91,25 +80,18 @@ class VideoGenerator:
         """
         # If users also provide some kwargs, it will override the FastVideoArgs and PipelineConfig.
         kwargs['model_path'] = model_path
-        log_queue = kwargs.pop("log_queue", None)
         fastvideo_args = FastVideoArgs.from_kwargs(**kwargs)
 
-        return cls.from_fastvideo_args(fastvideo_args, log_queue=log_queue)
+        return cls.from_fastvideo_args(fastvideo_args)
 
     @classmethod
-    def from_fastvideo_args(
-        cls,
-        fastvideo_args: FastVideoArgs,
-        *,
-        log_queue=None,
-    ) -> "VideoGenerator":
+    def from_fastvideo_args(cls, fastvideo_args: FastVideoArgs) -> "VideoGenerator":
         """
         Create a video generator with the specified arguments.
-
+        
         Args:
             fastvideo_args: The inference arguments
-            log_queue: Optional multiprocessing.Queue to forward worker logs to
-
+                
         Returns:
             The created video generator
         """
@@ -121,7 +103,6 @@ class VideoGenerator:
             fastvideo_args=fastvideo_args,
             executor_class=executor_class,
             log_stats=False,  # TODO: implement
-            log_queue=log_queue,
         )
 
     def generate_video(
@@ -159,34 +140,6 @@ class VideoGenerator:
             A metadata dictionary for single-prompt generation, or a list of
             metadata dictionaries for prompt-file batch generation.
         """
-        log_queue = kwargs.pop("log_queue", None)
-        if log_queue:
-            self.executor.set_log_queue(log_queue)
-            
-        try:
-            return self._generate_video_impl(
-                prompt=prompt,
-                sampling_param=sampling_param,
-                mouse_cond=mouse_cond,
-                keyboard_cond=keyboard_cond,
-                grid_sizes=grid_sizes,
-                **kwargs,
-            )
-        finally:
-            if log_queue:
-                self.executor.clear_log_queue()
-
-    def _generate_video_impl(
-        self,
-        prompt: str | None = None,
-        sampling_param: SamplingParam | None = None,
-        mouse_cond: torch.Tensor | None = None,
-        keyboard_cond: torch.Tensor | None = None,
-        grid_sizes: tuple[int, int, int] | list[int] | torch.Tensor
-        | None = None,
-        **kwargs,
-    ) -> dict[str, Any] | list[np.ndarray] | list[dict[str, Any]]:
-        """Internal implementation of generate_video."""
         # Handle batch processing from text file
         if sampling_param is None:
             sampling_param = SamplingParam.from_pretrained(self.fastvideo_args.model_path)
@@ -203,9 +156,8 @@ class VideoGenerator:
 
         if self.fastvideo_args.prompt_txt is not None or sampling_param.prompt_path is not None:
             prompt_txt_path = sampling_param.prompt_path or self.fastvideo_args.prompt_txt
-            if not prompt_txt_path or not os.path.exists(prompt_txt_path):
-                raise FileNotFoundError(
-                    f"Prompt text file not found: {prompt_txt_path}")
+            if not os.path.exists(prompt_txt_path):
+                raise FileNotFoundError(f"Prompt text file not found: {prompt_txt_path}")
 
             # Read prompts from file
             with open(prompt_txt_path, encoding='utf-8') as f:

fastvideo/models/vaes/common.py

@@ -364,10 +364,15 @@ class ParallelTiledVAE(ABC):
             if i > 0:
                 tile = tile[:, :, 1:, :, :]
             row.append(tile)
+        blend_latent_num_frames = max(
+            1,
+            self.blend_num_frames // self.temporal_compression_ratio,
+        )
         result_row = []
         for i, tile in enumerate(row):
             if i > 0:
-                tile = self.blend_t(row[i - 1], tile, self.blend_num_frames)
+                tile = self.blend_t(row[i - 1], tile,
+                                    blend_latent_num_frames)
                 result_row.append(
                     tile[:, :, :tile_latent_stride_num_frames, :, :])
             else:

fastvideo/pipelines/preprocess/preprocess_hunyuan_overfit.py

@@ -0,0 +1,206 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Preprocess HunyuanVideo overfit data into parquet format.
+
+Encodes videos with HunyuanVideo VAE and captions with dual text
+encoders (LLaMA + CLIP) into the t2v parquet schema expected by
+the training framework.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=0 python scripts/preprocess_hunyuan_overfit.py
+"""
+
+import glob
+import json
+import os
+
+import cv2
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import torch
+from safetensors.torch import load_file as safetensors_load_file
+
+from fastvideo.configs.models.vaes import HunyuanVAEConfig
+from fastvideo.configs.pipelines.hunyuan import (
+    clip_preprocess_text,
+    clip_postprocess_text,
+    llama_preprocess_text,
+    llama_postprocess_text,
+)
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
+from fastvideo.models.vaes.hunyuanvae import AutoencoderKLHunyuanVideo
+from fastvideo.utils import maybe_download_model
+
+# --- Config ---
+NUM_FRAMES = 77  # 4k+1 for temporal compression ratio 4
+MAX_HEIGHT = 480
+MAX_WIDTH = 832
+TRAIN_FPS = 16.0
+
+DATA_DIR = "data/hunyuan_overfit"
+OUTPUT_DIR = "data/hunyuan_overfit_preprocessed"
+
+
+def load_video(path: str, num_frames: int) -> torch.Tensor:
+    """Load video as [1, C, T, H, W] in [-1, 1]."""
+    cap = cv2.VideoCapture(path)
+    frames: list[np.ndarray] = []
+    while len(frames) < num_frames:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frames.append(frame)
+    cap.release()
+
+    if len(frames) < num_frames:
+        # Repeat last frame to fill
+        while len(frames) < num_frames:
+            frames.append(frames[-1])
+
+    frames = frames[:num_frames]
+    video = np.stack(frames, axis=0)
+    video = torch.from_numpy(video).float()
+    video = video / 127.5 - 1.0  # [0,255] -> [-1,1]
+    video = video.permute(3, 0, 1, 2).unsqueeze(0)  # [1,C,T,H,W]
+    return video
+
+
+def main() -> None:
+    device = torch.device("cuda:0")
+    model_path = maybe_download_model("hunyuanvideo-community/HunyuanVideo")
+    vae_path = os.path.join(model_path, "vae")
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    # Load captions
+    with open(os.path.join(DATA_DIR, "videos2caption.json")) as f:
+        caption_data = json.load(f)
+
+    # --- Load VAE ---
+    print("Loading HunyuanVideo VAE...")
+    vae_config = HunyuanVAEConfig()
+    vae_config.load_encoder = True
+    vae_config.load_decoder = False
+    vae = AutoencoderKLHunyuanVideo(vae_config)
+    sf_files = glob.glob(os.path.join(vae_path, "*.safetensors"))
+    loaded = {}
+    for sf_file in sf_files:
+        loaded.update(safetensors_load_file(sf_file))
+    vae.load_state_dict(loaded, strict=False)
+    vae = vae.to(device=device, dtype=torch.float16).eval()
+    vae.use_parallel_tiling = False
+    print(f"VAE loaded ({sum(p.numel() for p in vae.parameters())/1e6:.0f}M)")
+
+    # --- Load text encoders ---
+    print("Loading LLaMA text encoder...")
+    from transformers import AutoTokenizer, CLIPTextModel, LlamaModel
+    from fastvideo.configs.pipelines.hunyuan import (LlamaConfig, CLIPTextConfig)
+
+    llama_cfg = LlamaConfig()
+    llama_tok = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"))
+    llama_enc = LlamaModel.from_pretrained(
+        os.path.join(model_path, "text_encoder"),
+        torch_dtype=torch.float16,
+    ).to(device).eval()
+    llama_tok_kwargs = dict(llama_cfg.tokenizer_kwargs)
+
+    print("Loading CLIP text encoder...")
+    clip_cfg = CLIPTextConfig()
+    clip_tok = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer_2"))
+    clip_enc = CLIPTextModel.from_pretrained(
+        os.path.join(model_path, "text_encoder_2"),
+        torch_dtype=torch.float16,
+    ).to(device).eval()
+    clip_tok_kwargs = dict(clip_cfg.tokenizer_kwargs)
+
+    # --- Process each video ---
+    records = []
+    for item in caption_data:
+        video_name = item["path"]
+        caption = item["cap"][0]
+        video_path = os.path.join(DATA_DIR, "videos", video_name)
+
+        print(f"\nProcessing: {video_name}")
+        print(f"  Caption: {caption[:80]}...")
+
+        # Encode video
+        video = load_video(video_path, NUM_FRAMES).to(device=device, dtype=torch.float16)
+        print(f"  Video shape: {video.shape}")
+
+        with torch.no_grad():
+            latent_dist = vae.encode(video)
+            # Cast to fp32 — dataloader hardcodes np.float32
+            latent = latent_dist.mean.squeeze(0).float().cpu()
+        print(f"  Latent shape: {latent.shape}")
+
+        # Encode text with LLaMA
+        llama_text = llama_preprocess_text(caption)
+        with torch.no_grad():
+            llama_inputs = llama_tok(llama_text, **llama_tok_kwargs).to(device)
+            llama_out = llama_enc(**llama_inputs, output_hidden_states=True)
+            llama_embeds = llama_postprocess_text(llama_out).squeeze(0)  # [seq, dim]
+
+        # Encode text with CLIP
+        clip_text = clip_preprocess_text(caption)
+        with torch.no_grad():
+            clip_inputs = clip_tok(clip_text, **clip_tok_kwargs).to(device)
+            clip_out = clip_enc(**clip_inputs)
+            clip_pooled = clip_postprocess_text(clip_out).squeeze(0)  # [dim]
+
+        # Combine: [pooled_clip_row, llama_embeds]
+        llama_dim = llama_embeds.shape[-1]
+        pooled_row = torch.zeros(llama_dim, device=device, dtype=torch.float16)
+        pooled_row[:clip_pooled.shape[-1]] = clip_pooled
+        text_embedding = torch.cat(
+            [pooled_row.unsqueeze(0), llama_embeds],
+            dim=0,
+        ).float().cpu()  # [seq+1, dim]
+        print(f"  Text embedding shape: {text_embedding.shape}")
+
+        record = {
+            "id": video_name,
+            "vae_latent_bytes": latent.numpy().tobytes(),
+            "vae_latent_shape": list(latent.shape),
+            "vae_latent_dtype": str(latent.dtype).replace("torch.", ""),
+            "text_embedding_bytes": text_embedding.numpy().tobytes(),
+            "text_embedding_shape": list(text_embedding.shape),
+            "text_embedding_dtype": str(text_embedding.dtype).replace("torch.", ""),
+            "file_name": video_name,
+            "caption": caption,
+            "media_type": "video",
+            "width": MAX_WIDTH,
+            "height": MAX_HEIGHT,
+            "num_frames": NUM_FRAMES,
+            "duration_sec": NUM_FRAMES / TRAIN_FPS,
+            "fps": TRAIN_FPS,
+        }
+        records.append(record)
+
+    # Clean up encoders
+    del llama_enc, llama_tok, clip_enc, clip_tok, vae
+
+    # Write parquet
+    table = pa.table(
+        {k: [r[k] for r in records]
+         for k in records[0]},
+        schema=pyarrow_schema_t2v,
+    )
+    output_path = os.path.join(OUTPUT_DIR, "data_00000.parquet")
+    pq.write_table(table, output_path)
+    print(f"\nWrote {len(records)} records to {output_path}")
+
+    # Write validation prompts for callback
+    # Wrap in "data" key — ValidationDataset expects field="data"
+    # Use "caption" field — ValidationDataset aliases it to "prompt"
+    val_prompts = {"data": [{"caption": item["cap"][0]} for item in caption_data]}
+    val_path = os.path.join(OUTPUT_DIR, "validation_prompts.json")
+    with open(val_path, "w") as f:
+        json.dump(val_prompts, f, indent=2)
+    print(f"Wrote validation prompts to {val_path}")
+
+    print("\nDone! Use data_path: " + OUTPUT_DIR + " in training config.")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/registry.py

@@ -83,7 +83,7 @@ from fastvideo.configs.sample.wan import (
 )
 from fastvideo.configs.sample.sd35 import SD35SamplingParam
 
-from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
+from fastvideo.fastvideo_args import WorkloadType
 from fastvideo.logger import init_logger
 from fastvideo.utils import (maybe_download_model_index, verify_model_config_and_directory)
 
@@ -135,7 +135,6 @@ class ConfigInfo:
 
     sampling_param_cls: type[SamplingParam] | None
     pipeline_config_cls: type[PipelineConfig]
-    workload_types: tuple[WorkloadType, ...]
 
 
 # The central registry mapping a model name to its configuration information
@@ -151,21 +150,15 @@ _MODEL_NAME_DETECTORS: list[tuple[str, Callable[[str], bool]]] = []
 def register_configs(
     sampling_param_cls: type[SamplingParam] | None,
     pipeline_config_cls: type[PipelineConfig],
-    workload_types: tuple[WorkloadType, ...],
     hf_model_paths: list[str] | None = None,
     model_detectors: list[Callable[[str], bool]] | None = None,
 ) -> None:
-    """Register config classes for a model family.
-
-    workload_types declares which UI workload options this config supports.
-    Use () for configs not exposed as workload options.
-    """
+    """Register config classes for a model family."""
     model_id = str(len(_CONFIG_REGISTRY))
 
     _CONFIG_REGISTRY[model_id] = ConfigInfo(
         sampling_param_cls=sampling_param_cls,
         pipeline_config_cls=pipeline_config_cls,
-        workload_types=workload_types,
     )
 
     if hf_model_paths:
@@ -242,7 +235,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=LTX2BaseSamplingParam,
         pipeline_config_cls=LTX2T2VConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "Lightricks/LTX-2",
             "FastVideo/LTX2-base",
@@ -256,7 +248,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=LTX2DistilledSamplingParam,
         pipeline_config_cls=LTX2T2VConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "FastVideo/LTX2-Distilled-Diffusers",
         ],
@@ -269,7 +260,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Hunyuan15_480P_SamplingParam,
         pipeline_config_cls=Hunyuan15T2V480PConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
         ],
@@ -285,7 +275,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Hunyuan15_480P_StepDistilled_I2V_SamplingParam,
         pipeline_config_cls=Hunyuan15I2V480PStepDistilledConfig,
-        workload_types=(WorkloadType.I2V,),
         hf_model_paths=[
             "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v_step_distilled",
         ],
@@ -293,7 +282,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Hunyuan15_720P_SamplingParam,
         pipeline_config_cls=Hunyuan15T2V720PConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v",
         ],
@@ -301,7 +289,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Hunyuan15_720P_Distilled_I2V_SamplingParam,
         pipeline_config_cls=Hunyuan15I2V720PConfig,
-        workload_types=(WorkloadType.I2V,),
         hf_model_paths=[
             "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v_distilled",
         ],
@@ -309,18 +296,13 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Hunyuan15_SR_1080P_SamplingParam,
         pipeline_config_cls=Hunyuan15SR1080PConfig,
-        workload_types=(),
-        hf_model_paths=[
-            "weizhou03/HunyuanVideo-1.5-Diffusers-1080p",
-            "weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR"
-        ],
+        hf_model_paths=["weizhou03/HunyuanVideo-1.5-Diffusers-1080p", "weizhou03/HunyuanVideo-1.5-Diffusers-1080p-2SR"],
     )
 
     # Hunyuan (excludes gamecraft, hyworld, and versioned models)
     register_configs(
         sampling_param_cls=HunyuanSamplingParam,
         pipeline_config_cls=HunyuanConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "hunyuanvideo-community/HunyuanVideo",
         ],
@@ -332,7 +314,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=FastHunyuanSamplingParam,
         pipeline_config_cls=FastHunyuanConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "FastVideo/FastHunyuan-diffusers",
         ],
@@ -342,7 +323,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=HYWorld_SamplingParam,
         pipeline_config_cls=HYWorldConfig,
-        workload_types=(),
         hf_model_paths=[
             "FastVideo/HY-WorldPlay-Bidirectional-Diffusers",
         ],
@@ -353,7 +333,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=HunyuanGameCraftSamplingParam,
         pipeline_config_cls=HunyuanGameCraftPipelineConfig,
-        workload_types=(WorkloadType.I2V,),
         hf_model_paths=[
             "FastVideo/HunyuanGameCraft-Diffusers",
         ],
@@ -363,7 +342,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=LingBotWorld_SamplingParam,
         pipeline_config_cls=LingBotWorldI2V480PConfig,
-        workload_types=(WorkloadType.I2V,),
         hf_model_paths=[
             "FastVideo/LingBot-World-Base-Cam-Diffusers",
         ],
@@ -374,7 +352,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=PipelineConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "kandinskylab/Kandinsky-5.0-T2V-Lite-sft-5s-Diffusers",
         ],
@@ -383,38 +360,19 @@ def _register_configs() -> None:
         ],
     )
 
-    # LongCat (T2V, I2V, VC use same config; workload varies by path)
+    # LongCat
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=LongCatT2V480PConfig,
-        workload_types=(WorkloadType.T2V,),
-        hf_model_paths=["FastVideo/LongCat-Video-T2V-Diffusers"],
-        model_detectors=[
-            lambda path: "longcat" in path.lower()
-            and "i2v" not in path.lower()
-            and "imagetovideo" not in path.lower()
-            and "vc" not in path.lower()
-            and "videocontinuation" not in path.lower(),
+        hf_model_paths=[
+            "FastVideo/LongCat-Video-T2V-Diffusers",
+            "FastVideo/LongCat-Video-I2V-Diffusers",
+            "FastVideo/LongCat-Video-VC-Diffusers",
         ],
-    )
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=LongCatT2V480PConfig,
-        workload_types=(WorkloadType.I2V,),
-        hf_model_paths=["FastVideo/LongCat-Video-I2V-Diffusers"],
         model_detectors=[
-            lambda path: "longcatimagetovideo" in path.lower()
-            or ("longcat" in path.lower() and "i2v" in path.lower()),
-        ],
-    )
-    register_configs(
-        sampling_param_cls=None,
-        pipeline_config_cls=LongCatT2V480PConfig,
-        workload_types=(),
-        hf_model_paths=["FastVideo/LongCat-Video-VC-Diffusers"],
-        model_detectors=[
-            lambda path: "longcatvideocontinuation" in path.lower()
-            or ("longcat" in path.lower() and "vc" in path.lower()),
+            lambda path: "longcatimagetovideo" in path.lower(),
+            lambda path: "longcatvideocontinuation" in path.lower(),
+            lambda path: "longcat" in path.lower(),
         ],
     )
 
@@ -422,7 +380,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=MatrixGame2_SamplingParam,
         pipeline_config_cls=MatrixGameI2V480PConfig,
-        workload_types=(WorkloadType.I2V,),
         hf_model_paths=[
             "FastVideo/Matrix-Game-2.0-Base-Diffusers",
             "FastVideo/Matrix-Game-2.0-GTA-Diffusers",
@@ -437,7 +394,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Cosmos25SamplingParamBase,
         pipeline_config_cls=Cosmos25Config,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "KyleShao/Cosmos-Predict2.5-2B-Diffusers",
         ],
@@ -454,7 +410,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Cosmos_Predict2_2B_Video2World_SamplingParam,
         pipeline_config_cls=CosmosConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "nvidia/Cosmos-Predict2-2B-Video2World",
         ],
@@ -468,7 +423,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=TurboDiffusionT2V_1_3B_SamplingParam,
         pipeline_config_cls=TurboDiffusionT2V_1_3B_Config,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers",
         ],
@@ -477,7 +431,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=TurboDiffusionT2V_14B_SamplingParam,
         pipeline_config_cls=TurboDiffusionT2V_14B_Config,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "loayrashid/TurboWan2.1-T2V-14B-Diffusers",
         ],
@@ -485,7 +438,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=TurboDiffusionI2V_A14B_SamplingParam,
         pipeline_config_cls=TurboDiffusionI2V_A14B_Config,
-        workload_types=(WorkloadType.I2V,),
         hf_model_paths=[
             "loayrashid/TurboWan2.2-I2V-A14B-Diffusers",
         ],
@@ -495,7 +447,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=WanT2V_1_3B_SamplingParam,
         pipeline_config_cls=WanT2V480PConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
         ],
@@ -504,7 +455,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=WanT2V_14B_SamplingParam,
         pipeline_config_cls=WanT2V720PConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "Wan-AI/Wan2.1-T2V-14B-Diffusers",
             "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers",
@@ -513,7 +463,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=WanI2V_14B_480P_SamplingParam,
         pipeline_config_cls=WanI2V480PConfig,
-        workload_types=(WorkloadType.I2V,),
         hf_model_paths=[
             "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
         ],
@@ -522,7 +471,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=WanI2V_14B_720P_SamplingParam,
         pipeline_config_cls=WanI2V720PConfig,
-        workload_types=(WorkloadType.I2V,),
         hf_model_paths=[
             "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
         ],
@@ -530,7 +478,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Wan2_1_Fun_1_3B_InP_SamplingParam,
         pipeline_config_cls=WanI2V480PConfig,
-        workload_types=(WorkloadType.I2V,),
         hf_model_paths=[
             "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers",
         ],
@@ -538,7 +485,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Wan2_1_Fun_1_3B_Control_SamplingParam,
         pipeline_config_cls=WANV2VConfig,
-        workload_types=(),
         hf_model_paths=[
             "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers",
         ],
@@ -546,7 +492,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=FastWanT2V480P_SamplingParam,
         pipeline_config_cls=FastWan2_1_T2V_480P_Config,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
             "FastVideo/FastWan2.1-T2V-14B-480P-Diffusers",
@@ -556,7 +501,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam,
         pipeline_config_cls=Wan2_2_TI2V_5B_Config,
-        workload_types=(WorkloadType.T2V, WorkloadType.I2V),
         hf_model_paths=[
             "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
         ],
@@ -564,7 +508,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Wan2_2_TI2V_5B_SamplingParam,
         pipeline_config_cls=FastWan2_2_TI2V_5B_Config,
-        workload_types=(WorkloadType.T2V, WorkloadType.I2V),
         hf_model_paths=[
             "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers",
             "FastVideo/FastWan2.2-TI2V-5B-Diffusers",
@@ -573,7 +516,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Wan2_2_T2V_A14B_SamplingParam,
         pipeline_config_cls=Wan2_2_T2V_A14B_Config,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
         ],
@@ -581,7 +523,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=Wan2_2_I2V_A14B_SamplingParam,
         pipeline_config_cls=Wan2_2_I2V_A14B_Config,
-        workload_types=(WorkloadType.I2V,),
         hf_model_paths=[
             "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
         ],
@@ -589,31 +530,17 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam,
         pipeline_config_cls=SelfForcingWanT2V480PConfig,
-        workload_types=(WorkloadType.T2V,),
         hf_model_paths=[
             "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",
         ],
         model_detectors=[lambda path: "wancausaldmdpipeline" in path.lower()],
     )
-    # SFWan2.2: T2V and I2V variants by path
     register_configs(
         sampling_param_cls=SelfForcingWan2_2_T2V_A14B_480P_SamplingParam,
         pipeline_config_cls=SelfForcingWan2_2_T2V480PConfig,
-        workload_types=(WorkloadType.T2V,),
-        hf_model_paths=["rand0nmr/SFWan2.2-T2V-A14B-Diffusers"],
-        model_detectors=[
-            lambda path: ("sfwan2.2" in path.lower() or "sfwan2_2" in path.lower())
-            and "i2v" not in path.lower(),
-        ],
-    )
-    register_configs(
-        sampling_param_cls=SelfForcingWan2_2_T2V_A14B_480P_SamplingParam,
-        pipeline_config_cls=SelfForcingWan2_2_T2V480PConfig,
-        workload_types=(WorkloadType.I2V,),
-        hf_model_paths=["FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers"],
-        model_detectors=[
-            lambda path: ("sfwan2.2" in path.lower() or "sfwan2_2" in path.lower())
-            and "i2v" in path.lower(),
+        hf_model_paths=[
+            "rand0nmr/SFWan2.2-T2V-A14B-Diffusers",
+            "FastVideo/SFWan2.2-I2V-A14B-Preview-Diffusers",
         ],
     )
 
@@ -621,7 +548,6 @@ def _register_configs() -> None:
     register_configs(
         sampling_param_cls=SD35SamplingParam,
         pipeline_config_cls=SD35Config,
-        workload_types=(WorkloadType.T2I,),
         hf_model_paths=[
             "stabilityai/stable-diffusion-3.5-medium",
         ],
@@ -709,52 +635,11 @@ def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None:
 
 _register_configs()
 
-
-def get_registered_model_paths() -> list[str]:
-    """Return all registered HuggingFace model paths.
-
-    Useful for UIs and tooling that need to enumerate supported models.
-    """
-    return sorted(_MODEL_HF_PATH_TO_NAME.keys())
-
-
-def get_registered_models_with_workloads(
-    workload_type: str | None = None,
-) -> list[dict[str, Any]]:
-    """Return models with workload metadata, optionally filtered by workload.
-
-    Args:
-        workload_type: If set (e.g. "t2v", "i2v", "t2i"), only return models
-            that support this workload. If None, return all with workload_types.
-
-    Returns:
-        List of dicts with keys: id, label, workload_types.
-    """
-    result: list[dict[str, Any]] = []
-    for path in sorted(_MODEL_HF_PATH_TO_NAME.keys()):
-        model_id = _MODEL_HF_PATH_TO_NAME[path]
-        config_info = _CONFIG_REGISTRY.get(model_id)
-        if config_info is None:
-            continue
-        workload_values = [w.value for w in config_info.workload_types]
-        if workload_type is not None and workload_type.lower() not in workload_values:
-            continue
-        label = path.split("/")[-1].replace("-", " ").replace("_", " ")
-        result.append({
-            "id": path,
-            "label": label,
-            "workload_types": workload_values,
-        })
-    return result
-
-
 __all__ = [
     "ConfigInfo",
     "ModelInfo",
     "get_model_info",
     "get_pipeline_config_cls_from_name",
-    "get_registered_model_paths",
-    "get_registered_models_with_workloads",
     "get_sampling_param_cls_for_name",
     "get_pipeline_config_classes",
 ]

fastvideo/tests/encoders/test_clip_encoder.py

@@ -21,9 +21,7 @@ os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder_2")
 TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer_2")
 
@@ -34,13 +32,12 @@ def test_clip_encoder():
     Tests compatibility between two different implementations for loading text encoders:
     1. load_text_encoder from fastvideo.models.hunyuan.text_encoder
     2. TextEncoderLoader from fastvideo.models.loader
-    
+
     The test verifies that both implementations:
     - Load models with the same weights and parameters
     - Produce nearly identical outputs for the same input prompts
     """
-    args = FastVideoArgs(model_path="openai/clip-vit-large-patch14",
-                         pipeline_config=HunyuanConfig())
+    args = FastVideoArgs(model_path="openai/clip-vit-large-patch14", pipeline_config=HunyuanConfig())
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
     logger.info("Loading models from %s", args.model_path)
@@ -55,6 +52,7 @@ def test_clip_encoder():
     model1 = CLIPTextModel.from_pretrained(TEXT_ENCODER_PATH).to(torch.float16).to(device).eval()
 
     from fastvideo.models.loader.component_loader import TextEncoderLoader
+
     loader = TextEncoderLoader()
     model2 = loader.load(TEXT_ENCODER_PATH, args)
 
@@ -90,10 +88,7 @@ def test_clip_encoder():
     tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
 
     # Test with some sample prompts
-    prompts = [
-        "a photo of a cat", "a beautiful landscape with mountains",
-        "an astronaut riding a horse on the moon"
-    ]
+    prompts = ["a photo of a cat", "a beautiful landscape with mountains", "an astronaut riding a horse on the moon"]
 
     logger.info("Testing CLIP text encoder with sample prompts")
 
@@ -102,14 +97,11 @@ def test_clip_encoder():
             logger.info("Testing prompt: '%s'", prompt)
 
             # Tokenize the prompt
-            tokens = tokenizer(prompt,
-                               padding="max_length",
-                               max_length=77,
-                               truncation=True,
-                               return_tensors="pt").to(device)
+            tokens = tokenizer(prompt, padding="max_length", max_length=77, truncation=True, return_tensors="pt").to(
+                device
+            )
             # Get embeddings from our implementation
-            outputs1 = model1(input_ids=tokens.input_ids,
-                              output_hidden_states=True)
+            outputs1 = model1(input_ids=tokens.input_ids, output_hidden_states=True)
 
             logger.info("Testing model2")
             print("--------------------------------")
@@ -118,24 +110,25 @@ def test_clip_encoder():
                 outputs2 = model2(
                     input_ids=tokens.input_ids,
                     # attention_mask=tokens.attention_mask,
-                    output_hidden_states=True)
+                    output_hidden_states=True,
+                )
 
             # Compare last hidden states
-            last_hidden_state1 = outputs1.last_hidden_state[
-                tokens.attention_mask == 1]
-            last_hidden_state2 = outputs2.last_hidden_state[
-                tokens.attention_mask == 1]
+            last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask == 1]
+            last_hidden_state2 = outputs2.last_hidden_state[tokens.attention_mask == 1]
             # print("last_hidden_state1", last_hidden_state1)
             # print("last_hidden_state2", last_hidden_state2)
 
-            assert last_hidden_state1.shape == last_hidden_state2.shape, \
+            assert last_hidden_state1.shape == last_hidden_state2.shape, (
                 f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}"
+            )
             # Compare pooler outputs
             pooler_output1 = outputs1.pooler_output
             pooler_output2 = outputs2.pooler_output
 
-            assert pooler_output1.shape == pooler_output2.shape, \
+            assert pooler_output1.shape == pooler_output2.shape, (
                 f"Pooler output shapes don't match: {pooler_output1.shape} vs {pooler_output2.shape}"
+            )
 
             assert_close(pooler_output1, pooler_output2, atol=1e-2, rtol=1e-3)
             assert_close(last_hidden_state1, last_hidden_state2, atol=1e-2, rtol=1e-3)

fastvideo/tests/encoders/test_llama_encoder.py

@@ -15,15 +15,14 @@ from fastvideo.utils import maybe_download_model
 from fastvideo.configs.models.encoders import LlamaConfig
 from torch.distributed.tensor import DTensor
 from torch.testing import assert_close
+
 logger = init_logger(__name__)
 
 os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder")
 TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer")
 
@@ -34,13 +33,12 @@ def test_llama_encoder():
     Tests compatibility between two different implementations for loading text encoders:
     1. load_text_encoder from fastvideo.models.hunyuan.text_encoder
     2. TextEncoderLoader from fastvideo.models.loader
-    
+
     The test verifies that both implementations:
     - Load models with the same weights and parameters
     - Produce nearly identical outputs for the same input prompts
     """
-    args = FastVideoArgs(model_path="meta-llama/Llama-2-7b-hf",
-                         pipeline_config=HunyuanConfig())
+    args = FastVideoArgs(model_path="meta-llama/Llama-2-7b-hf", pipeline_config=HunyuanConfig())
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -68,15 +66,15 @@ def test_llama_encoder():
     logger.info("Model1 has %d parameters", len(params1))
     logger.info("Model2 has %d parameters", len(params2))
 
-
     # check if embed_tokens are the same
     device = model1.embed_tokens.weight.device
-    assert torch.allclose(model1.embed_tokens.weight,
-                          model2.embed_tokens.weight.to_local().to(device) if isinstance(model2.embed_tokens.weight, DTensor) else model2.embed_tokens.weight.to(device))
-    weights = [
-        "layers.{}.input_layernorm.weight",
-        "layers.{}.post_attention_layernorm.weight"
-    ]
+    assert torch.allclose(
+        model1.embed_tokens.weight,
+        model2.embed_tokens.weight.to_local().to(device)
+        if isinstance(model2.embed_tokens.weight, DTensor)
+        else model2.embed_tokens.weight.to(device),
+    )
+    weights = ["layers.{}.input_layernorm.weight", "layers.{}.post_attention_layernorm.weight"]
     for layer_idx in range(hf_config.num_hidden_layers):
         for w in weights:
             name1 = w.format(layer_idx)
@@ -121,33 +119,29 @@ def test_llama_encoder():
             logger.info("Testing prompt: '%s'", prompt)
 
             # Tokenize the prompt
-            tokens = tokenizer(prompt,
-                               padding="max_length",
-                               max_length=512,
-                               truncation=True,
-                               return_tensors="pt").to(device)
+            tokens = tokenizer(prompt, padding="max_length", max_length=512, truncation=True, return_tensors="pt").to(
+                device
+            )
 
             # Get outputs from our implementation
             # filter out padding input_ids
             # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1]
             # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1]
-            outputs1 = model1(input_ids=tokens.input_ids,
-                              output_hidden_states=True)
+            outputs1 = model1(input_ids=tokens.input_ids, output_hidden_states=True)
             print("--------------------------------")
             logger.info("Testing model2")
 
             # Get outputs from HuggingFace implementation
             with set_forward_context(current_timestep=0, attn_metadata=None):
-                outputs2 = model2(input_ids=tokens.input_ids,
-                                  attention_mask=tokens.attention_mask,
-                                  output_hidden_states=True)
+                outputs2 = model2(
+                    input_ids=tokens.input_ids, attention_mask=tokens.attention_mask, output_hidden_states=True
+                )
 
             # Compare last hidden states
-            last_hidden_state1 = outputs1.last_hidden_state[
-                tokens.attention_mask == 1]
-            last_hidden_state2 = outputs2.last_hidden_state[
-                tokens.attention_mask == 1]
+            last_hidden_state1 = outputs1.last_hidden_state[tokens.attention_mask == 1]
+            last_hidden_state2 = outputs2.last_hidden_state[tokens.attention_mask == 1]
 
-            assert last_hidden_state1.shape == last_hidden_state2.shape, \
+            assert last_hidden_state1.shape == last_hidden_state2.shape, (
                 f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}"
+            )
             assert_close(last_hidden_state1, last_hidden_state2, atol=1e-1, rtol=1e-4)

fastvideo/tests/modal/ssim_test.py

@@ -1,7 +1,11 @@
 import ast
+import datetime
 import glob
 import os
 import re
+import shutil
+import subprocess
+import sys
 from dataclasses import dataclass
 from typing import Any
 
@@ -10,7 +14,7 @@ import modal
 app = modal.App()
 
 model_vol = modal.Volume.from_name("hf-model-weights")
-image_version = os.getenv("IMAGE_VERSION")
+image_version = os.getenv("IMAGE_VERSION", "latest")
 image_tag = f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{image_version}"
 print(f"Using image: {image_tag}")
 
@@ -24,30 +28,29 @@ image = (
         "libssl-dev",
         "ffmpeg",
     )
-    .run_commands(
-        "curl --proto '=https' --tlsv1.2 -sSf "
-        "https://sh.rustup.rs | sh -s -- -y --default-toolchain stable"
-    )
+    .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable")
     .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
-    .env({
-        "PATH": "/root/.cargo/bin:$PATH",
-        "BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""),
-        "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
-        "BUILDKITE_PULL_REQUEST": os.environ.get(
-            "BUILDKITE_PULL_REQUEST", ""
-        ),
-        "IMAGE_VERSION": os.environ.get("IMAGE_VERSION", ""),
-    })
+    .env(
+        {
+            "PATH": "/root/.cargo/bin:$PATH",
+            "BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""),
+            "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""),
+            "BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
+            "IMAGE_VERSION": os.environ.get("IMAGE_VERSION", ""),
+        }
+    )
 )
 
 SSIM_NUM_GPUS = 4
 SSIM_TERMINATE_TIMEOUT_S = 30
+HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN")
+RAW_GENERATED_VOLUME_ROOT = "ssim_generated_videos"
+DEFAULT_OUTPUT_QUALITY_TIER = "default"
+FULL_OUTPUT_QUALITY_TIER = "full_quality"
+MODAL_DEVICE_REFERENCE_FOLDER = "L40S_reference_videos"
 SSIM_COMMON_KWARGS = dict(
     image=image,
     timeout=5400,
-    secrets=[modal.Secret.from_dict(
-        {"HF_API_KEY": os.environ.get("HF_API_KEY", "")}
-    )],
     volumes={"/root/data": model_vol},
 )
 
@@ -105,6 +108,194 @@ class _PartitionResult:
     exit_code: int
 
 
+def _split_csv_values(csv_values: str) -> set[str]:
+    return {value.strip() for value in csv_values.split(",") if value.strip()}
+
+
+def _run_git_command(args: list[str]) -> str:
+    try:
+        result = subprocess.run(
+            ["git", *args],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except subprocess.CalledProcessError as error:
+        stderr = (error.stderr or "").strip()
+        raise RuntimeError(f"Failed to run git {' '.join(args)}: {stderr or error}") from error
+    return result.stdout.strip()
+
+
+def _normalize_git_repo_url(git_repo: str) -> str:
+    if git_repo.startswith("git@github.com:"):
+        return "https://github.com/" + git_repo[len("git@github.com:") :]
+    if git_repo.startswith("ssh://git@github.com/"):
+        return "https://github.com/" + git_repo[len("ssh://git@github.com/") :]
+    return git_repo
+
+
+def _resolve_git_repo(git_repo: str) -> str:
+    if git_repo.strip():
+        return _normalize_git_repo_url(git_repo.strip())
+
+    env_repo = os.environ.get("BUILDKITE_REPO", "").strip()
+    if env_repo:
+        return _normalize_git_repo_url(env_repo)
+
+    discovered_repo = _run_git_command(["config", "--get", "remote.origin.url"])
+    if discovered_repo:
+        return _normalize_git_repo_url(discovered_repo)
+
+    raise RuntimeError("Could not resolve git repo URL. Pass --git-repo or set BUILDKITE_REPO.")
+
+
+def _resolve_git_commit(git_commit: str) -> str:
+    if git_commit.strip():
+        return git_commit.strip()
+
+    env_commit = os.environ.get("BUILDKITE_COMMIT", "").strip()
+    if env_commit:
+        return env_commit
+
+    discovered_commit = _run_git_command(["rev-parse", "HEAD"])
+    if discovered_commit:
+        return discovered_commit
+
+    raise RuntimeError("Could not resolve git commit. Pass --git-commit or set BUILDKITE_COMMIT.")
+
+
+def _resolve_pull_request(pr_number: str) -> str:
+    if pr_number.strip():
+        return pr_number.strip()
+    env_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "").strip()
+    if env_pr:
+        return env_pr
+    return "false"
+
+
+def _resolve_hf_api_key(
+    hf_api_key: str,
+) -> str:
+    if hf_api_key.strip():
+        return hf_api_key.strip()
+
+    for key in HF_TOKEN_ENV_KEYS:
+        value = os.environ.get(key, "").strip()
+        if value:
+            return value
+
+    raise RuntimeError(
+        "Hugging Face token is required. Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN; or pass --hf-api-key."
+    )
+
+
+def _sanitize_path_fragment(value: str) -> str:
+    cleaned = re.sub(r"[^A-Za-z0-9_.-]+", "_", value.strip())
+    return cleaned.strip("._-")
+
+
+def _resolve_generated_volume_subdir(
+    requested_subdir: str,
+    git_commit: str,
+) -> str:
+    if requested_subdir.strip():
+        normalized = requested_subdir.strip().strip("/")
+        if not normalized:
+            raise RuntimeError("generated_volume_subdir must not be empty.")
+        return normalized
+
+    short_commit = _sanitize_path_fragment(git_commit[:12]) or "unknown"
+    timestamp = datetime.datetime.utcnow().strftime("%Y%m%d_%H%M%S")
+    return f"{timestamp}_{short_commit}"
+
+
+def _resolve_output_quality_tier(ssim_full_quality: bool) -> str:
+    if ssim_full_quality:
+        return FULL_OUTPUT_QUALITY_TIER
+    return DEFAULT_OUTPUT_QUALITY_TIER
+
+
+def _build_generated_volume_relative_path(
+    *,
+    generated_volume_subdir: str,
+    quality_tier: str,
+) -> str:
+    return os.path.join(
+        RAW_GENERATED_VOLUME_ROOT,
+        quality_tier,
+        generated_volume_subdir,
+        "generated_videos",
+    )
+
+
+def _build_local_generated_download_dir(quality_tier: str) -> str:
+    return os.path.join(".", "generated_videos_modal", quality_tier)
+
+
+def _print_local_reference_copy_command(quality_tier: str) -> None:
+    generated_dir = os.path.join(
+        _build_local_generated_download_dir(quality_tier),
+        MODAL_DEVICE_REFERENCE_FOLDER,
+    )
+    print(
+        "To update local references from the downloaded Modal outputs, run:\n"
+        "  python fastvideo/tests/ssim/reference_videos_cli.py copy-local "
+        f"--quality-tier {quality_tier} "
+        f"--generated-dir {generated_dir} "
+        f"--device-folder {MODAL_DEVICE_REFERENCE_FOLDER}"
+    )
+
+
+def _count_video_files(root: str) -> int:
+    count = 0
+    for current_root, _, filenames in os.walk(root):
+        for filename in filenames:
+            if filename.lower().endswith((".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv")):
+                count += 1
+    return count
+
+
+def _sync_generated_videos_to_volume(
+    repo_root: str,
+    generated_volume_subdir: str,
+    quality_tier: str,
+) -> str | None:
+    generated_root = os.path.join(
+        repo_root,
+        "fastvideo",
+        "tests",
+        "ssim",
+        "generated_videos",
+        quality_tier,
+    )
+    if not os.path.isdir(generated_root):
+        print(
+            f"No generated_videos directory found for quality tier {quality_tier}; skipping raw generated video export."
+        )
+        return None
+
+    relative_dst = _build_generated_volume_relative_path(
+        generated_volume_subdir=generated_volume_subdir,
+        quality_tier=quality_tier,
+    )
+    absolute_dst = os.path.join("/root/data", relative_dst)
+    if os.path.exists(absolute_dst):
+        shutil.rmtree(absolute_dst)
+    os.makedirs(os.path.dirname(absolute_dst), exist_ok=True)
+    shutil.copytree(generated_root, absolute_dst)
+    model_vol.commit()
+
+    num_videos = _count_video_files(absolute_dst)
+    print(f"Raw generated videos exported to Modal volume path: {relative_dst} ({num_videos} video files).")
+    print(
+        "Download command:\n"
+        f"  modal volume get hf-model-weights {relative_dst} "
+        f"{_build_local_generated_download_dir(quality_tier)}"
+    )
+    _print_local_reference_copy_command(quality_tier)
+    return relative_dst
+
+
 def _extract_required_gpus(filepath: str) -> int:
     """Read REQUIRED_GPUS from a test file. Defaults to 1."""
     with open(filepath, encoding="utf-8") as file:
@@ -126,10 +317,7 @@ def _extract_model_ids(filepath: str) -> list[str]:
         value_node = None
 
         if isinstance(node, ast.Assign):
-            target_names = [
-                target.id for target in node.targets
-                if isinstance(target, ast.Name)
-            ]
+            target_names = [target.id for target in node.targets if isinstance(target, ast.Name)]
             value_node = node.value
         elif isinstance(node, ast.AnnAssign):
             if isinstance(node.target, ast.Name):
@@ -138,8 +326,7 @@ def _extract_model_ids(filepath: str) -> list[str]:
 
         if not target_names or value_node is None:
             continue
-        if not any(name.endswith("MODEL_TO_PARAMS")
-                   for name in target_names):
+        if not any(name.endswith("MODEL_TO_PARAMS") for name in target_names):
             continue
         if not isinstance(value_node, ast.Dict):
             continue
@@ -159,35 +346,70 @@ def _extract_model_ids(filepath: str) -> list[str]:
     return unique_model_ids
 
 
-def _discover_ssim_tasks(ssim_dir: str) -> list[SSIMTask]:
+def _discover_ssim_tasks(
+    ssim_dir: str,
+    *,
+    selected_test_files: set[str] | None = None,
+    selected_model_ids: set[str] | None = None,
+) -> list[SSIMTask]:
     tasks = []
     task_id = 0
+    selected_test_files = selected_test_files or set()
+    selected_test_file_names = {os.path.basename(test_file) for test_file in selected_test_files}
+    selected_model_ids = selected_model_ids or set()
+
     test_files = sorted(glob.glob(os.path.join(ssim_dir, "test_*.py")))
+    matched_test_files = set()
+    matched_model_ids = set()
     for filepath in test_files:
+        file_name = os.path.basename(filepath)
+        rel_path = f"./fastvideo/tests/ssim/{file_name}"
+        if selected_test_files and (
+            filepath not in selected_test_files
+            and rel_path not in selected_test_files
+            and file_name not in selected_test_file_names
+        ):
+            continue
+
+        matched_test_files.add(file_name)
         required_gpus = _extract_required_gpus(filepath)
         if required_gpus < 1 or required_gpus > SSIM_NUM_GPUS:
-            raise ValueError(
-                f"{filepath} requires {required_gpus} GPUs, "
-                f"but scheduler supports up to {SSIM_NUM_GPUS}."
-            )
-        rel_path = f"./fastvideo/tests/ssim/{os.path.basename(filepath)}"
+            raise ValueError(f"{filepath} requires {required_gpus} GPUs, but scheduler supports up to {SSIM_NUM_GPUS}.")
         model_ids = _extract_model_ids(filepath)
         if model_ids:
             for model_id in model_ids:
-                tasks.append(SSIMTask(
+                if selected_model_ids and model_id not in selected_model_ids:
+                    continue
+                matched_model_ids.add(model_id)
+                tasks.append(
+                    SSIMTask(
+                        task_id=task_id,
+                        test_file=rel_path,
+                        required_gpus=required_gpus,
+                        model_id=model_id,
+                    )
+                )
+                task_id += 1
+        else:
+            if selected_model_ids:
+                continue
+            tasks.append(
+                SSIMTask(
                     task_id=task_id,
                     test_file=rel_path,
                     required_gpus=required_gpus,
-                    model_id=model_id,
-                ))
-                task_id += 1
-        else:
-            tasks.append(SSIMTask(
-                task_id=task_id,
-                test_file=rel_path,
-                required_gpus=required_gpus,
-            ))
+                )
+            )
             task_id += 1
+
+    if selected_test_files:
+        unmatched_files = sorted(selected_test_file_names - matched_test_files)
+        if unmatched_files:
+            raise RuntimeError("Requested SSIM test file(s) not found: " + ", ".join(unmatched_files))
+    if selected_model_ids:
+        unmatched_model_ids = sorted(selected_model_ids - matched_model_ids)
+        if unmatched_model_ids:
+            raise RuntimeError("Requested SSIM model_id(s) not found: " + ", ".join(unmatched_model_ids))
     return sorted(tasks, key=lambda task: task.sort_key)
 
 
@@ -207,26 +429,24 @@ def _build_checkout_command(git_commit: str, pr_number: str | None) -> str:
         try:
             pr_id = int(pr_number)
         except ValueError as error:
-            raise RuntimeError(
-                f"Invalid BUILDKITE_PULL_REQUEST value: {pr_number}"
-            ) from error
-        return (
-            "git fetch --prune origin "
-            f"refs/pull/{pr_id}/head && "
-            "git checkout FETCH_HEAD"
-        )
+            raise RuntimeError(f"Invalid BUILDKITE_PULL_REQUEST value: {pr_number}") from error
+        return f"git fetch --prune origin refs/pull/{pr_id}/head && git checkout FETCH_HEAD"
     return f"git checkout {shlex.quote(git_commit)}"
 
 
-def _prepare_ssim_workspace() -> tuple[str, list[SSIMTask]]:
+def _prepare_ssim_workspace(
+    *,
+    git_repo: str,
+    git_commit: str,
+    pr_number: str,
+    hf_api_key: str,
+    selected_test_files: set[str] | None = None,
+    selected_model_ids: set[str] | None = None,
+) -> tuple[str, list[SSIMTask]]:
     import shlex
-    import subprocess
 
-    git_repo = os.environ.get("BUILDKITE_REPO")
-    git_commit = os.environ.get("BUILDKITE_COMMIT")
-    pr_number = os.environ.get("BUILDKITE_PULL_REQUEST")
-    if not git_repo or not git_commit:
-        raise RuntimeError("BUILDKITE_REPO and BUILDKITE_COMMIT must be set.")
+    if not hf_api_key.strip():
+        raise RuntimeError("HF API key is required to prepare SSIM workspace.")
 
     checkout_command = _build_checkout_command(git_commit, pr_number)
     repo_root = "/FastVideo"
@@ -235,7 +455,14 @@ def _prepare_ssim_workspace() -> tuple[str, list[SSIMTask]]:
     set -euo pipefail
     source $HOME/.local/bin/env
     source /opt/venv/bin/activate
-    cd {shlex.quote(repo_root)}
+    if [ -d {shlex.quote(repo_root)}/.git ]; then
+      cd {shlex.quote(repo_root)}
+      git remote set-url origin {shlex.quote(git_repo)} || true
+      git fetch --prune origin
+    else
+      git clone {shlex.quote(git_repo)} {shlex.quote(repo_root)}
+      cd {shlex.quote(repo_root)}
+    fi
     {checkout_command}
     git submodule update --init --recursive
     cd fastvideo-kernel
@@ -249,17 +476,22 @@ def _prepare_ssim_workspace() -> tuple[str, list[SSIMTask]]:
         ["/bin/bash", "-lc", command],
         capture_output=True,
         text=True,
+        env={
+            **os.environ,
+            "HF_API_KEY": hf_api_key,
+        },
     )
     if result.returncode != 0:
         print(result.stdout)
         print(result.stderr)
-        raise RuntimeError(
-            "Workspace setup failed with "
-            f"exit code {result.returncode}"
-        )
+        raise RuntimeError(f"Workspace setup failed with exit code {result.returncode}")
 
     ssim_dir = os.path.join(repo_root, "fastvideo", "tests", "ssim")
-    tasks = _discover_ssim_tasks(ssim_dir)
+    tasks = _discover_ssim_tasks(
+        ssim_dir,
+        selected_test_files=selected_test_files,
+        selected_model_ids=selected_model_ids,
+    )
     if not tasks:
         raise RuntimeError("No SSIM test files found.")
     return repo_root, tasks
@@ -271,22 +503,14 @@ def _spawn_ssim_task(
     assigned_gpu_ids: list[str],
     log_dir: str,
     task_index: int,
+    pytest_extra_args: list[str],
 ) -> _RunningTask:
     import shlex
-    import subprocess
 
-    safe_test_name = re.sub(
-        r"[^A-Za-z0-9_.-]+", "_", task.test_name
-    )
-    log_path = os.path.join(
-        log_dir, f"{task_index:03d}_{safe_test_name}.log"
-    )
-    command = (
-        "set -euo pipefail && "
-        "source $HOME/.local/bin/env && "
-        "source /opt/venv/bin/activate && "
-        f"pytest {shlex.quote(task.test_file)} -vs"
-    )
+    safe_test_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", task.test_name)
+    log_path = os.path.join(log_dir, f"{task_index:03d}_{safe_test_name}.log")
+    pytest_command = shlex.join(["pytest", task.test_file, "-vs", *pytest_extra_args])
+    command = f"set -euo pipefail && source $HOME/.local/bin/env && source /opt/venv/bin/activate && {pytest_command}"
     env = os.environ.copy()
     env["HF_HOME"] = "/root/data/.cache"
     # MultiprocExecutor returns CUDA tensors through mp pipes (CUDA IPC).
@@ -327,9 +551,7 @@ def _finalize_running_task(
 ) -> None:
     running_task.log_handle.close()
     available_gpu_ids.extend(running_task.gpu_ids)
-    available_gpu_ids.sort(
-        key=lambda gpu_id: gpu_order.get(gpu_id, len(gpu_order))
-    )
+    available_gpu_ids.sort(key=lambda gpu_id: gpu_order.get(gpu_id, len(gpu_order)))
     results[running_task.task.task_id] = _TaskResult(
         task=running_task.task,
         status=status,
@@ -402,9 +624,30 @@ def _get_visible_gpu_ids() -> list[str]:
     return gpu_ids
 
 
+def _build_pytest_extra_args(
+    *,
+    ssim_full_quality: bool,
+    ssim_reference_repo: str,
+    skip_ssim_reference_download: bool,
+    pytest_k: str,
+) -> list[str]:
+    args = []
+    if ssim_full_quality:
+        args.append("--ssim-full-quality")
+    if ssim_reference_repo.strip():
+        args.extend(["--ssim-reference-repo", ssim_reference_repo.strip()])
+    if skip_ssim_reference_download:
+        args.append("--skip-ssim-reference-download")
+    if pytest_k.strip():
+        args.extend(["-k", pytest_k.strip()])
+    return args
+
+
 def _schedule_ssim_tasks(
     repo_root: str,
     tasks: list[SSIMTask],
+    pytest_extra_args: list[str],
+    fail_fast: bool = True,
 ) -> dict[int, _TaskResult]:
     import tempfile
     import time
@@ -412,16 +655,10 @@ def _schedule_ssim_tasks(
     pending_tasks = list(tasks)
     running_tasks = []
     available_gpu_ids = _get_visible_gpu_ids()
-    gpu_order = {
-        gpu_id: index for index, gpu_id in enumerate(available_gpu_ids)
-    }
-    max_required_gpus = max(
-        (task.required_gpus for task in pending_tasks), default=0
-    )
+    gpu_order = {gpu_id: index for index, gpu_id in enumerate(available_gpu_ids)}
+    max_required_gpus = max((task.required_gpus for task in pending_tasks), default=0)
     if max_required_gpus > len(available_gpu_ids):
-        cuda_visible_devices = os.environ.get(
-            "CUDA_VISIBLE_DEVICES", "<unset>"
-        )
+        cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "<unset>")
         raise RuntimeError(
             "SSIM task requires "
             f"{max_required_gpus} GPUs but only "
@@ -443,19 +680,17 @@ def _schedule_ssim_tasks(
                 break
 
             task = pending_tasks.pop(next_task_index)
-            assigned_gpu_ids = available_gpu_ids[:task.required_gpus]
-            del available_gpu_ids[:task.required_gpus]
+            assigned_gpu_ids = available_gpu_ids[: task.required_gpus]
+            del available_gpu_ids[: task.required_gpus]
             running_task = _spawn_ssim_task(
                 task=task,
                 repo_root=repo_root,
                 assigned_gpu_ids=assigned_gpu_ids,
                 log_dir=log_dir,
                 task_index=task.task_id,
+                pytest_extra_args=pytest_extra_args,
             )
-            print(
-                f"Started {task.test_name} on GPUs "
-                f"{','.join(assigned_gpu_ids)}"
-            )
+            print(f"Started {task.test_name} on GPUs {','.join(assigned_gpu_ids)}")
             running_tasks.append(running_task)
 
         completed_tasks = []
@@ -474,11 +709,8 @@ def _schedule_ssim_tasks(
                 gpu_order=gpu_order,
             )
             running_tasks.remove(running_task)
-            print(
-                f"Finished {running_task.task.test_name} with "
-                f"exit code {returncode}"
-            )
-            if returncode != 0 and not fail_fast_triggered:
+            print(f"Finished {running_task.task.test_name} with exit code {returncode}")
+            if returncode != 0 and fail_fast and not fail_fast_triggered:
                 fail_fast_triggered = True
 
         if fail_fast_triggered and running_tasks:
@@ -517,22 +749,18 @@ def _collect_task_summaries(
     for task in tasks:
         result = results[task.task_id]
         log_content = None
-        if (
-            result.status == "failed"
-            and result.log_path
-            and os.path.exists(result.log_path)
-        ):
-            with open(
-                result.log_path, encoding="utf-8"
-            ) as f:
+        if result.status == "failed" and result.log_path and os.path.exists(result.log_path):
+            with open(result.log_path, encoding="utf-8") as f:
                 log_content = f.read()
-        summaries.append(_TaskSummary(
-            test_name=task.test_name,
-            required_gpus=task.required_gpus,
-            status=result.status,
-            returncode=result.returncode,
-            log_content=log_content,
-        ))
+        summaries.append(
+            _TaskSummary(
+                test_name=task.test_name,
+                required_gpus=task.required_gpus,
+                status=result.status,
+                returncode=result.returncode,
+                log_content=log_content,
+            )
+        )
     return summaries
 
 
@@ -550,18 +778,13 @@ def _print_combined_results(
         if result is None:
             continue
         for s in result.task_summaries:
-            label = (
-                f"[P{result.partition_index}] "
-                f"{s.test_name}"
-            )
+            label = f"[P{result.partition_index}] {s.test_name}"
             if s.status == "passed":
                 passed.append(label)
             elif s.status == "failed":
                 failed.append(label)
                 if first_failure is None:
-                    first_failure = (
-                        result.partition_index, s
-                    )
+                    first_failure = (result.partition_index, s)
             elif s.status == "terminated":
                 terminated.append(label)
             elif s.status == "skipped":
@@ -585,14 +808,8 @@ def _print_combined_results(
             status = "cancelled"
         else:
             count = len(result.task_summaries)
-            status = (
-                "passed" if result.exit_code == 0
-                else "failed"
-            )
-        print(
-            f"  Partition {i}: {status} "
-            f"({count} tasks)"
-        )
+            status = "passed" if result.exit_code == 0 else "failed"
+        print(f"  Partition {i}: {status} ({count} tasks)")
     print(f"  passed: {len(passed)}")
     print(f"  failed: {len(failed)}")
     print(f"  terminated: {len(terminated)}")
@@ -605,40 +822,72 @@ def _print_combined_results(
     if skipped:
         print(f"Skipped: {', '.join(skipped)}")
 
-    has_failures = bool(
-        failed or terminated or skipped
-        or any(r is None for r in partition_results)
-    )
+    has_failures = bool(failed or terminated or skipped or any(r is None for r in partition_results))
     return 1 if has_failures else 0
 
 
 @app.function(gpu=f"L40S:{SSIM_NUM_GPUS}", **SSIM_COMMON_KWARGS)
 def run_ssim_partition(
     partition_index: int,
+    num_partitions: int,
+    git_repo: str,
+    git_commit: str,
+    pr_number: str = "false",
+    hf_api_key: str = "",
+    test_files_csv: str = "",
+    model_ids_csv: str = "",
+    ssim_full_quality: bool = False,
+    ssim_reference_repo: str = "",
+    skip_ssim_reference_download: bool = False,
+    pytest_k: str = "",
+    sync_generated_to_volume: bool = False,
+    generated_volume_subdir: str = "",
+    fail_fast: bool = True,
 ) -> _PartitionResult:
-    repo_root, tasks = _prepare_ssim_workspace()
-    partition = _partition_tasks(tasks, partition_index)
+    selected_test_files = _split_csv_values(test_files_csv)
+    selected_model_ids = _split_csv_values(model_ids_csv)
+    repo_root, tasks = _prepare_ssim_workspace(
+        git_repo=git_repo,
+        git_commit=git_commit,
+        pr_number=pr_number,
+        hf_api_key=hf_api_key,
+        selected_test_files=selected_test_files,
+        selected_model_ids=selected_model_ids,
+    )
+    partition = _partition_tasks(tasks, partition_index, num_partitions)
     if not partition:
-        print(
-            f"Partition {partition_index}: "
-            "no tasks assigned."
-        )
+        print(f"Partition {partition_index}: no tasks assigned.")
         return _PartitionResult(
             partition_index=partition_index,
             task_summaries=[],
             exit_code=0,
         )
-    print(
-        f"Partition {partition_index}: running "
-        f"{len(partition)}/{len(tasks)} tasks"
+    print(f"Partition {partition_index}: running {len(partition)}/{len(tasks)} tasks")
+    pytest_extra_args = _build_pytest_extra_args(
+        ssim_full_quality=ssim_full_quality,
+        ssim_reference_repo=ssim_reference_repo,
+        skip_ssim_reference_download=skip_ssim_reference_download,
+        pytest_k=pytest_k,
     )
-    results = _schedule_ssim_tasks(repo_root, partition)
-    summaries = _collect_task_summaries(
-        partition, results
-    )
-    has_failures = any(
-        s.status != "passed" for s in summaries
+    results = _schedule_ssim_tasks(
+        repo_root,
+        partition,
+        pytest_extra_args=pytest_extra_args,
+        fail_fast=fail_fast,
     )
+    summaries = _collect_task_summaries(partition, results)
+    has_failures = any(s.status != "passed" for s in summaries)
+    if sync_generated_to_volume:
+        quality_tier = _resolve_output_quality_tier(ssim_full_quality)
+        resolved_subdir = _resolve_generated_volume_subdir(
+            generated_volume_subdir,
+            git_commit,
+        )
+        _sync_generated_videos_to_volume(
+            repo_root,
+            resolved_subdir,
+            quality_tier,
+        )
     return _PartitionResult(
         partition_index=partition_index,
         task_summaries=summaries,
@@ -646,18 +895,91 @@ def run_ssim_partition(
     )
 
 
-@app.local_entrypoint()
-def run_ssim_tests():
-    import sys
+NUM_PARTITIONS = 2
 
-    future_0 = run_ssim_partition.spawn(partition_index=0)
-    future_1 = run_ssim_partition.spawn(partition_index=1)
-    results: list[_PartitionResult | None] = [
-        future_0.get(),
-        future_1.get(),
+
+@app.local_entrypoint()
+def run_ssim_tests(
+    git_repo: str = "",
+    git_commit: str = "",
+    pr_number: str = "",
+    hf_api_key: str = "",
+    test_files: str = "",
+    model_ids: str = "",
+    full_quality: bool = False,
+    reference_repo: str = "",
+    skip_reference_download: bool = False,
+    pytest_k: str = "",
+    sync_generated_to_volume: bool = False,
+    generated_volume_subdir: str = "",
+    no_fail_fast: bool = False,
+):
+    resolved_git_repo = _resolve_git_repo(git_repo)
+    resolved_git_commit = _resolve_git_commit(git_commit)
+    resolved_pr_number = _resolve_pull_request(pr_number)
+    resolved_hf_api_key = _resolve_hf_api_key(hf_api_key)
+
+    print(f"Running SSIM on repo: {resolved_git_repo}")
+    print(f"Using commit: {resolved_git_commit}")
+    if resolved_pr_number and resolved_pr_number != "false":
+        print(f"Using PR ref: {resolved_pr_number}")
+    if test_files.strip():
+        print(f"Selected test files: {test_files}")
+    if model_ids.strip():
+        print(f"Selected model ids: {model_ids}")
+    if pytest_k.strip():
+        print(f"Using pytest -k filter: {pytest_k}")
+    quality_tier = _resolve_output_quality_tier(full_quality)
+    if sync_generated_to_volume:
+        resolved_subdir = _resolve_generated_volume_subdir(
+            generated_volume_subdir,
+            resolved_git_commit,
+        )
+        print(
+            "Raw generated videos will be saved to Modal volume path: "
+            f"{_build_generated_volume_relative_path(generated_volume_subdir=resolved_subdir, quality_tier=quality_tier)}"
+        )
+    else:
+        resolved_subdir = ""
+
+    common_kwargs = dict(
+        git_repo=resolved_git_repo,
+        git_commit=resolved_git_commit,
+        pr_number=resolved_pr_number,
+        hf_api_key=resolved_hf_api_key,
+        test_files_csv=test_files,
+        model_ids_csv=model_ids,
+        ssim_full_quality=full_quality,
+        ssim_reference_repo=reference_repo,
+        skip_ssim_reference_download=skip_reference_download,
+        pytest_k=pytest_k,
+        sync_generated_to_volume=sync_generated_to_volume,
+        generated_volume_subdir=resolved_subdir,
+        fail_fast=not no_fail_fast,
+    )
+    futures = [
+        run_ssim_partition.spawn(
+            partition_index=i,
+            num_partitions=NUM_PARTITIONS,
+            **common_kwargs,
+        )
+        for i in range(NUM_PARTITIONS)
     ]
+    results: list[_PartitionResult | None] = [f.get() for f in futures]
 
     exit_code = _print_combined_results(results)
+    if sync_generated_to_volume:
+        download_src = _build_generated_volume_relative_path(
+            generated_volume_subdir=resolved_subdir,
+            quality_tier=quality_tier,
+        )
+        local_download_dir = _build_local_generated_download_dir(quality_tier)
+        print(
+            "To download raw generated videos locally, run:\n"
+            f"  modal volume get hf-model-weights {download_src} "
+            f"{local_download_dir}"
+        )
+        _print_local_reference_copy_command(quality_tier)
     if exit_code != 0:
         sys.exit(exit_code)
     print("All SSIM tasks passed.")

fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py

@@ -3,8 +3,8 @@ from pathlib import Path
 from huggingface_hub import snapshot_download
 import subprocess
 import sys
-from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
 import shutil
+from fastvideo.tests.utils import compute_video_ssim_torchvision
 
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
@@ -181,4 +181,4 @@ def test_e2e_overfit_single_sample():
 
 
 if __name__ == "__main__":
-    test_e2e_overfit_single_sample()
+    test_e2e_overfit_single_sample()

fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py

@@ -4,7 +4,7 @@ from huggingface_hub import snapshot_download
 import shutil
 import subprocess
 import sys
-from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
+from fastvideo.tests.utils import compute_video_ssim_torchvision
 
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
@@ -170,4 +170,4 @@ def test_e2e_overfit_single_sample():
 
 
 if __name__ == "__main__":
-    test_e2e_overfit_single_sample()
+    test_e2e_overfit_single_sample()

fastvideo/tests/nightly/test_e2e_overfit_single_sample.py

@@ -4,7 +4,7 @@ from huggingface_hub import snapshot_download
 import shutil
 import subprocess
 import sys
-from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision
+from fastvideo.tests.utils import compute_video_ssim_torchvision
 
 # Import the training pipeline
 sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
@@ -175,4 +175,4 @@ def test_e2e_overfit_single_sample():
 
 
 if __name__ == "__main__":
-    test_e2e_overfit_single_sample()
+    test_e2e_overfit_single_sample()

fastvideo/tests/ssim/AGENTS.md

@@ -0,0 +1,55 @@
+# SSIM Directory Guidelines
+
+## Scope
+These instructions apply to everything under `fastvideo/tests/ssim/`.
+
+## Purpose
+SSIM tests are GPU-backed end-to-end quality regression checks. They generate
+videos and compare them against device-specific reference videos.
+
+## Key Files
+- `test_*.py`: SSIM regression tests by model/task.
+- `conftest.py`: optional filtering via `FASTVIDEO_SSIM_MODEL_ID`.
+- `reference_videos/<quality-tier>/<GPU>_reference_videos/`: local cache of
+  references split by quality tier and GPU type.
+- `generated_videos/<quality-tier>/<GPU>_reference_videos/`: local outputs and
+  `*_ssim.json` artifacts (git-ignored).
+- `reference_videos_cli.py`: copy/download/upload/ensure reference videos
+  (including HF sync).
+
+## Run Commands
+- Full SSIM suite: `pytest fastvideo/tests/ssim/ -vs`
+- Full SSIM with full-quality params:
+  `pytest fastvideo/tests/ssim/ -vs --ssim-full-quality`
+- Single test file:
+  `pytest fastvideo/tests/ssim/test_wan_t2v_similarity.py -vs`
+- Single model split:
+  `FASTVIDEO_SSIM_MODEL_ID=<model_id> pytest fastvideo/tests/ssim/test_wan_t2v_similarity.py -vs`
+- Modal orchestrator (CI-style scheduling):
+  `modal run fastvideo/tests/modal/ssim_test.py`
+
+## Authoring Rules
+- Name files `test_<feature>_similarity.py`.
+- Set `REQUIRED_GPUS = <int>` near the top of each test module.
+- For multi-model suites, keep configs in `*_MODEL_TO_PARAMS` dictionaries.
+  The Modal scheduler auto-discovers these keys and runs one subprocess per
+  model id.
+- Keep runs deterministic when possible (fixed prompts/seeds/frames/backend).
+- Persist metrics with `write_ssim_results(...)`.
+- Use `pytest.skip(...)` for unsupported hardware, missing assets, or
+  insufficient GPU count.
+
+## Updating Reference Videos
+1. Run the target SSIM test and inspect generated outputs + scores.
+2. Update references only for intentional behavior changes.
+3. Copy local generated outputs into the target reference folder:
+   `python fastvideo/tests/ssim/reference_videos_cli.py copy-local ...`
+4. Upload updated reference folders to HF:
+   `python fastvideo/tests/ssim/reference_videos_cli.py upload ...`
+5. In the PR, explain why references changed and which GPU type generated
+   them.
+
+## PR Expectations
+- Include exact test commands and pass/fail evidence.
+- If thresholds change, include before/after SSIM numbers and rationale.
+- If references change, call out the source commit/model/backend.

fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/steps4_An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali_ssim.json

@@ -1,11 +0,0 @@
-{
-  "mean_ssim": 0.9805062346988254,
-  "min_ssim": 0.9636241793632507,
-  "max_ssim": 0.9951534271240234,
-  "reference_video": "/mnt/fast-disks/hao_lab/loay/FastVideo/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4",
-  "generated_video": "/mnt/fast-disks/hao_lab/loay/FastVideo/fastvideo/tests/ssim/generated_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4",
-  "parameters": {
-    "num_inference_steps": 4,
-    "prompt": "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot."
-  }
-}

fastvideo/tests/ssim/README.md

@@ -1,8 +1,52 @@
-The reference videos in the `*_reference_videos` directory are used as part of an e2e test to ensure consistency in video generation quality across code changes. `test_inference_similarity.py` compares newly generated videos against these references using Structural Similarity Index (SSIM) metrics to detect any regressions in visual quality across code changes.
+The reference videos are used as part of e2e SSIM regression tests.
+Wan inference coverage now lives in `test_wan_t2v_similarity.py` and
+`test_wan_i2v_similarity.py` alongside the other model-specific SSIM files.
+These tests compare newly generated videos against references to detect quality
+regressions.
 
-`A40_reference_videos` are generated on A40s and so on.
+reference layout:
+- `reference_videos/default/<GPU>_reference_videos/<model>/<backend>/...`
+- `reference_videos/full_quality/<GPU>_reference_videos/<model>/<backend>/...`
 
-run `bash update_reference_videos.sh` from inside the `fastvideo/tests/ssim/` directory after running `test_inference_similarity.py` to update reference videos. Note: make sure to update the path to the corresponding device.
+`A40_reference_videos` are generated on A40s and so on. (legacy default layout
+`<ssim_dir>/<GPU>_reference_videos/...` is still read as fallback.)
+
+Before SSIM tests run, missing reference videos are auto-downloaded from a
+public HF repo (configured by `FASTVIDEO_SSIM_REFERENCE_HF_REPO`, default:
+`FastVideo/ssim-reference-videos`).
+
+Use the CLI:
+- `python fastvideo/tests/ssim/reference_videos_cli.py --help`
+- `python fastvideo/tests/ssim/reference_videos_cli.py copy-local --help`
+- `python fastvideo/tests/ssim/reference_videos_cli.py copy-local --quality-tier default --device-folder H200_reference_videos`
+- `python fastvideo/tests/ssim/reference_videos_cli.py copy-local --quality-tier full_quality --device-folder H200_reference_videos`
+- `python fastvideo/tests/ssim/reference_videos_cli.py download --help`
+- `python fastvideo/tests/ssim/reference_videos_cli.py upload --help`
+- `python fastvideo/tests/ssim/reference_videos_cli.py download --quality-tier all`
+- `python fastvideo/tests/ssim/reference_videos_cli.py upload --quality-tier all`
+- `python fastvideo/tests/ssim/reference_videos_cli.py download --quality-tier full_quality --device-folder H200_reference_videos`
+- `python fastvideo/tests/ssim/reference_videos_cli.py upload --quality-tier full_quality --device-folder H200_reference_videos`
+
+For `upload`, the tool reads HF token from `HF_API_KEY` /
+`HUGGINGFACE_HUB_TOKEN` / `HF_TOKEN` and fails fast if none are set.
+
+run `pytest fastvideo/tests/ssim/ -vs --ssim-full-quality` to use the
+`*_FULL_QUALITY_PARAMS` configs (default run keeps the original shortened test
+configs).
+
+to override the HF reference repo at test time:
+`pytest fastvideo/tests/ssim/ -vs --ssim-reference-repo <org/repo>`
+
+to skip auto-download of missing refs:
+`pytest fastvideo/tests/ssim/ -vs --skip-ssim-reference-download`
+
+generated videos are written under:
+- default params: `generated_videos/default/<GPU>_reference_videos/<model_id>/<ATTENTION_BACKEND>/`
+- full-quality params: `generated_videos/full_quality/<GPU>_reference_videos/<model_id>/<ATTENTION_BACKEND>/`
+
+HF repo layout mirrors quality + GPU split:
+- `reference_videos/default/<GPU>_reference_videos/...`
+- `reference_videos/full_quality/<GPU>_reference_videos/...`
 
 reference videos were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb`
 causal videos were generated on commit b318063c0a4618f1d5d99ea82ca67a06aad0d19d

fastvideo/tests/ssim/conftest.py

@@ -1,5 +1,60 @@
 import os
 
+from fastvideo.tests.ssim.reference_utils import (
+    FULL_QUALITY_ENV_VAR,
+    get_output_quality_tier,
+)
+from fastvideo.tests.ssim.reference_videos_cli import (
+    HF_REPO_ENV_KEY,
+    ensure_reference_videos_available,
+)
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--ssim-full-quality",
+        action="store_true",
+        default=False,
+        help=(
+            "Use *_FULL_QUALITY_PARAMS for SSIM tests. "
+            "Default keeps the original CI-friendly params."
+        ),
+    )
+    parser.addoption(
+        "--ssim-reference-repo",
+        default=os.environ.get(HF_REPO_ENV_KEY, ""),
+        help=(
+            "HF repo id for SSIM reference videos "
+            f"(overrides {HF_REPO_ENV_KEY})."
+        ),
+    )
+    parser.addoption(
+        "--skip-ssim-reference-download",
+        action="store_true",
+        default=False,
+        help="Skip auto-download of missing SSIM reference videos from HF.",
+    )
+
+
+def pytest_configure(config):
+    if config.getoption("--ssim-full-quality"):
+        os.environ[FULL_QUALITY_ENV_VAR] = "1"
+
+    repo_id = config.getoption("--ssim-reference-repo")
+    if repo_id:
+        os.environ[HF_REPO_ENV_KEY] = repo_id
+
+    skip_download = config.getoption("--skip-ssim-reference-download")
+    skip_download = skip_download or os.environ.get(
+        "FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD", ""
+    ).strip().lower() in {"1", "true", "yes", "on"}
+
+    if not skip_download:
+        ensure_reference_videos_available(
+            repo_id=repo_id or None,
+            quality_tier=get_output_quality_tier(),
+        )
+
 
 def pytest_collection_modifyitems(config, items):
     """Optionally keep only tests with a matching model_id parameter."""

fastvideo/tests/ssim/inference_similarity_utils.py

@@ -0,0 +1,340 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+from contextlib import contextmanager
+from logging import Logger
+from typing import Iterator
+
+from fastvideo import VideoGenerator
+from fastvideo.tests.ssim.reference_utils import (
+    build_generated_output_dir,
+    build_reference_folder_path,
+    get_cuda_device_name,
+    resolve_device_reference_folder,
+    select_ssim_params,
+)
+from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
+from fastvideo.worker.multiproc_executor import MultiprocExecutor
+
+DEVICE_MAPPINGS = (
+    ("A40", "A40"),
+    ("L40S", "L40S"),
+    ("H100", "H100"),
+    ("H200", "H200"),
+)
+
+
+@contextmanager
+def attention_backend(backend: str) -> Iterator[None]:
+    previous = os.environ.get("FASTVIDEO_ATTENTION_BACKEND")
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = backend
+    try:
+        yield
+    finally:
+        if previous is None:
+            os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None)
+        else:
+            os.environ["FASTVIDEO_ATTENTION_BACKEND"] = previous
+
+
+def shutdown_executor(generator: VideoGenerator | None) -> None:
+    if generator is None:
+        return
+    if isinstance(generator.executor, MultiprocExecutor):
+        generator.executor.shutdown()
+
+
+def resolve_inference_device_reference_folder(logger: Logger) -> str:
+    device_name = get_cuda_device_name()
+    device_reference_folder = resolve_device_reference_folder(
+        DEVICE_MAPPINGS,
+        device_name=device_name,
+        logger=logger,
+    )
+    if device_reference_folder is None:
+        raise ValueError(f"Unsupported device for ssim tests: {device_name}")
+    return device_reference_folder
+
+
+def _find_reference_video(reference_folder: str, prompt: str) -> str:
+    prompt_prefix = prompt[:100].strip()
+    for filename in os.listdir(reference_folder):
+        if filename.endswith(".mp4") and prompt_prefix in filename:
+            return os.path.join(reference_folder, filename)
+    raise FileNotFoundError("Reference video missing")
+
+
+def _assert_similarity(
+    *,
+    logger: Logger,
+    output_dir: str,
+    output_video_name: str,
+    reference_folder: str,
+    prompt: str,
+    num_inference_steps: int,
+    min_acceptable_ssim: float,
+    model_id: str,
+    attention_backend_name: str,
+) -> None:
+    if not os.path.exists(reference_folder):
+        logger.error("Reference folder missing: %s", reference_folder)
+        error_msg = (
+            f"Reference video folder does not exist: {reference_folder}\n"
+            f"To download reference videos, run:\n"
+            f"  python fastvideo/tests/ssim/reference_videos_cli.py download"
+        )
+        raise FileNotFoundError(error_msg)
+
+    try:
+        reference_video_path = _find_reference_video(reference_folder, prompt)
+    except FileNotFoundError as error:
+        logger.error(
+            "Reference video not found for prompt: %s with backend: %s",
+            prompt,
+            attention_backend_name,
+        )
+        raise error
+
+    generated_video_path = os.path.join(output_dir, output_video_name)
+
+    logger.info(
+        "Computing SSIM between %s and %s",
+        reference_video_path,
+        generated_video_path,
+    )
+    ssim_values = compute_video_ssim_torchvision(
+        reference_video_path,
+        generated_video_path,
+        use_ms_ssim=True,
+    )
+
+    mean_ssim = ssim_values[0]
+    logger.info("SSIM mean value: %s", mean_ssim)
+    logger.info("Writing SSIM results to directory: %s", output_dir)
+
+    success = write_ssim_results(
+        output_dir,
+        ssim_values,
+        reference_video_path,
+        generated_video_path,
+        num_inference_steps,
+        prompt,
+    )
+    if not success:
+        logger.error("Failed to write SSIM results to file")
+
+    assert mean_ssim >= min_acceptable_ssim, (
+        f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} "
+        f"for {model_id} with backend {attention_backend_name}"
+    )
+
+
+def _build_init_kwargs(
+    base_params: dict[str, object],
+) -> dict[str, object]:
+    init_kwargs: dict[str, object] = {
+        "num_gpus": base_params["num_gpus"],
+        "sp_size": base_params.get("sp_size", 1),
+        "tp_size": base_params.get("tp_size", 1),
+        "use_fsdp_inference": True,
+        "dit_cpu_offload": False,
+        "dit_layerwise_offload": False,
+    }
+    if "flow_shift" in base_params:
+        init_kwargs["flow_shift"] = base_params["flow_shift"]
+    if base_params.get("vae_sp"):
+        init_kwargs["vae_sp"] = True
+        init_kwargs["vae_tiling"] = True
+    if "text-encoder-precision" in base_params:
+        init_kwargs["text_encoder_precisions"] = base_params["text-encoder-precision"]
+    if base_params.get("ltx2_vae_tiling"):
+        init_kwargs["ltx2_vae_tiling"] = True
+        init_kwargs["ltx2_vae_spatial_tile_size_in_pixels"] = base_params.get(
+            "ltx2_vae_spatial_tile_size_in_pixels", 512
+        )
+        init_kwargs["ltx2_vae_spatial_tile_overlap_in_pixels"] = base_params.get(
+            "ltx2_vae_spatial_tile_overlap_in_pixels", 64
+        )
+        init_kwargs["ltx2_vae_temporal_tile_size_in_frames"] = base_params.get(
+            "ltx2_vae_temporal_tile_size_in_frames", 64
+        )
+        init_kwargs["ltx2_vae_temporal_tile_overlap_in_frames"] = base_params.get(
+            "ltx2_vae_temporal_tile_overlap_in_frames", 24
+        )
+    return init_kwargs
+
+
+def _build_generation_kwargs(
+    base_params: dict[str, object],
+    num_inference_steps: int,
+    output_dir: str,
+) -> dict[str, object]:
+    generation_kwargs: dict[str, object] = {
+        "num_inference_steps": num_inference_steps,
+        "output_path": output_dir,
+        "height": base_params["height"],
+        "width": base_params["width"],
+        "num_frames": base_params["num_frames"],
+        "seed": base_params["seed"],
+    }
+    if "guidance_scale" in base_params:
+        generation_kwargs["guidance_scale"] = base_params["guidance_scale"]
+    if "embedded_cfg_scale" in base_params:
+        generation_kwargs["embedded_cfg_scale"] = base_params["embedded_cfg_scale"]
+    if "fps" in base_params:
+        generation_kwargs["fps"] = base_params["fps"]
+    if "neg_prompt" in base_params:
+        generation_kwargs["neg_prompt"] = base_params["neg_prompt"]
+    return generation_kwargs
+
+
+def run_text_to_video_similarity_test(
+    *,
+    logger: Logger,
+    script_dir: str,
+    device_reference_folder: str,
+    prompt: str,
+    attention_backend_name: str,
+    model_id: str,
+    default_params_map: dict[str, dict[str, object]],
+    full_quality_params_map: dict[str, dict[str, object]],
+    min_acceptable_ssim: float,
+    init_kwargs_override: dict[str, object] | None = None,
+    generation_kwargs_override: dict[str, object] | None = None,
+) -> None:
+    with attention_backend(attention_backend_name):
+        output_dir = build_generated_output_dir(
+            script_dir,
+            device_reference_folder,
+            model_id,
+            attention_backend_name,
+        )
+        output_video_name = f"{prompt[:100].strip()}.mp4"
+        os.makedirs(output_dir, exist_ok=True)
+
+        params_map = select_ssim_params(
+            default_params_map,
+            full_quality_params_map,
+        )
+        base_params = params_map[model_id]
+        num_inference_steps = int(base_params["num_inference_steps"])
+
+        init_kwargs = _build_init_kwargs(base_params)
+        if init_kwargs_override:
+            init_kwargs.update(init_kwargs_override)
+
+        generation_kwargs = _build_generation_kwargs(
+            base_params,
+            num_inference_steps,
+            output_dir,
+        )
+        if generation_kwargs_override:
+            generation_kwargs.update(generation_kwargs_override)
+
+        generator: VideoGenerator | None = None
+        try:
+            generator = VideoGenerator.from_pretrained(
+                model_path=base_params["model_path"],
+                **init_kwargs,
+            )
+            generator.generate_video(prompt, **generation_kwargs)
+        finally:
+            shutdown_executor(generator)
+
+    assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}"
+
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        attention_backend_name,
+    )
+    _assert_similarity(
+        logger=logger,
+        output_dir=output_dir,
+        output_video_name=output_video_name,
+        reference_folder=reference_folder,
+        prompt=prompt,
+        num_inference_steps=num_inference_steps,
+        min_acceptable_ssim=min_acceptable_ssim,
+        model_id=model_id,
+        attention_backend_name=attention_backend_name,
+    )
+
+
+def run_image_to_video_similarity_test(
+    *,
+    logger: Logger,
+    script_dir: str,
+    device_reference_folder: str,
+    prompt: str,
+    image_path: str,
+    attention_backend_name: str,
+    model_id: str,
+    default_params_map: dict[str, dict[str, object]],
+    full_quality_params_map: dict[str, dict[str, object]],
+    min_acceptable_ssim: float,
+    init_kwargs_override: dict[str, object] | None = None,
+    generation_kwargs_override: dict[str, object] | None = None,
+) -> None:
+    with attention_backend(attention_backend_name):
+        output_dir = build_generated_output_dir(
+            script_dir,
+            device_reference_folder,
+            model_id,
+            attention_backend_name,
+        )
+        output_video_name = f"{prompt[:100].strip()}.mp4"
+        os.makedirs(output_dir, exist_ok=True)
+
+        params_map = select_ssim_params(
+            default_params_map,
+            full_quality_params_map,
+        )
+        base_params = params_map[model_id]
+        num_inference_steps = int(base_params["num_inference_steps"])
+
+        init_kwargs = _build_init_kwargs(base_params)
+        if init_kwargs_override:
+            init_kwargs.update(init_kwargs_override)
+
+        generation_kwargs = _build_generation_kwargs(
+            base_params,
+            num_inference_steps,
+            output_dir,
+        )
+        generation_kwargs["image_path"] = image_path
+        if generation_kwargs_override:
+            generation_kwargs.update(generation_kwargs_override)
+
+        generator: VideoGenerator | None = None
+        try:
+            generator = VideoGenerator.from_pretrained(
+                model_path=base_params["model_path"],
+                **init_kwargs,
+            )
+            generator.generate_video(prompt, **generation_kwargs)
+        finally:
+            shutdown_executor(generator)
+
+    assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}"
+
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        attention_backend_name,
+    )
+    _assert_similarity(
+        logger=logger,
+        output_dir=output_dir,
+        output_video_name=output_video_name,
+        reference_folder=reference_folder,
+        prompt=prompt,
+        num_inference_steps=num_inference_steps,
+        min_acceptable_ssim=min_acceptable_ssim,
+        model_id=model_id,
+        attention_backend_name=attention_backend_name,
+    )

fastvideo/tests/ssim/reference_utils.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+from collections.abc import Sequence
+from logging import Logger
+from typing import TypeVar
+
+import torch
+
+DEVICE_REFERENCE_FOLDER_SUFFIX = "_reference_videos"
+FULL_QUALITY_ENV_VAR = "FASTVIDEO_SSIM_FULL_QUALITY"
+DEFAULT_OUTPUT_QUALITY_TIER = "default"
+FULL_OUTPUT_QUALITY_TIER = "full_quality"
+REFERENCE_VIDEOS_DIRNAME = "reference_videos"
+DeviceMapping = tuple[str, str]
+T = TypeVar("T")
+
+
+def build_reference_folder_name(
+    device_prefix: str,
+    suffix: str = DEVICE_REFERENCE_FOLDER_SUFFIX,
+) -> str:
+    return f"{device_prefix}{suffix}"
+
+
+def get_cuda_device_name(index: int = 0) -> str:
+    if not torch.cuda.is_available():
+        return "CPU"
+    return torch.cuda.get_device_name(index)
+
+
+def resolve_device_reference_folder(
+    device_mappings: Sequence[DeviceMapping],
+    *,
+    device_name: str | None = None,
+    fallback_device_prefix: str | None = None,
+    unknown_device_prefix: str | None = None,
+    suffix: str = DEVICE_REFERENCE_FOLDER_SUFFIX,
+    logger: Logger | None = None,
+) -> str | None:
+    resolved_device_name = device_name or get_cuda_device_name()
+    for device_pattern, device_prefix in device_mappings:
+        if device_pattern in resolved_device_name:
+            return build_reference_folder_name(device_prefix, suffix)
+
+    if logger is not None:
+        if fallback_device_prefix is None:
+            logger.warning(
+                "Unsupported device for ssim tests: %s",
+                resolved_device_name,
+            )
+        else:
+            logger.warning(
+                "Unsupported device for ssim tests: %s; using %s references",
+                resolved_device_name,
+                fallback_device_prefix,
+            )
+
+    if unknown_device_prefix is not None:
+        return build_reference_folder_name(unknown_device_prefix, suffix)
+    if fallback_device_prefix is not None:
+        return build_reference_folder_name(fallback_device_prefix, suffix)
+    return None
+
+
+def build_reference_folder_path(
+    script_dir: str,
+    device_reference_folder: str,
+    model_id: str,
+    attention_backend: str,
+) -> str:
+    quality_tier = get_output_quality_tier()
+    tiered_reference_folder = os.path.join(
+        script_dir,
+        REFERENCE_VIDEOS_DIRNAME,
+        quality_tier,
+        device_reference_folder,
+        model_id,
+        attention_backend,
+    )
+    if os.path.exists(tiered_reference_folder):
+        return tiered_reference_folder
+
+    legacy_reference_folder = os.path.join(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        attention_backend,
+    )
+    if quality_tier == DEFAULT_OUTPUT_QUALITY_TIER and os.path.exists(
+        legacy_reference_folder
+    ):
+        return legacy_reference_folder
+
+    return tiered_reference_folder
+
+
+def use_full_quality_configs() -> bool:
+    value = os.environ.get(FULL_QUALITY_ENV_VAR, "").strip().lower()
+    return value in {"1", "true", "yes", "on"}
+
+
+def get_output_quality_tier() -> str:
+    if use_full_quality_configs():
+        return FULL_OUTPUT_QUALITY_TIER
+    return DEFAULT_OUTPUT_QUALITY_TIER
+
+
+def build_generated_output_dir(
+    script_dir: str,
+    device_reference_folder: str,
+    model_id: str,
+    attention_backend: str,
+) -> str:
+    return os.path.join(
+        script_dir,
+        "generated_videos",
+        get_output_quality_tier(),
+        device_reference_folder,
+        model_id,
+        attention_backend,
+    )
+
+
+def select_ssim_params(default_params: T, full_quality_params: T) -> T:
+    if use_full_quality_configs():
+        return full_quality_params
+    return default_params

fastvideo/tests/ssim/reference_videos_cli.py

@@ -0,0 +1,637 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import os
+import shutil
+import sys
+import tempfile
+from collections.abc import Iterable, Sequence
+from pathlib import Path
+
+
+VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv")
+HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN")
+HF_REPO_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO"
+HF_REPO_TYPE_ENV_KEY = "FASTVIDEO_SSIM_REFERENCE_HF_REPO_TYPE"
+
+DEFAULT_REPO_ID = "FastVideo/ssim-reference-videos"
+DEFAULT_REPO_TYPE = "dataset"
+DEFAULT_OUTPUT_QUALITY_TIER = "default"
+FULL_OUTPUT_QUALITY_TIER = "full_quality"
+ALL_OUTPUT_QUALITY_TIERS = "all"
+REFERENCE_VIDEOS_DIRNAME = "reference_videos"
+DEFAULT_DEVICE_REFERENCE_FOLDER = "L40S_reference_videos"
+QUALITY_TIERS = (
+    DEFAULT_OUTPUT_QUALITY_TIER,
+    FULL_OUTPUT_QUALITY_TIER,
+)
+
+
+def _ssim_dir() -> Path:
+    return Path(__file__).resolve().parent
+
+
+def _default_repo_id() -> str:
+    return os.environ.get(HF_REPO_ENV_KEY, DEFAULT_REPO_ID)
+
+
+def _default_repo_type() -> str:
+    return os.environ.get(HF_REPO_TYPE_ENV_KEY, DEFAULT_REPO_TYPE)
+
+
+def _iter_video_files(root: Path) -> Iterable[Path]:
+    for path in root.rglob("*"):
+        if path.is_file() and path.suffix.lower() in VIDEO_EXTENSIONS:
+            yield path
+
+
+def _discover_reference_dirs(base_dir: Path) -> list[Path]:
+    if not base_dir.exists():
+        return []
+    dirs = [p for p in base_dir.iterdir() if p.is_dir() and p.name.endswith("_reference_videos")]
+    return sorted(dirs)
+
+
+@contextlib.contextmanager
+def _exclusive_download_lock(lock_path: Path):
+    """Best-effort cross-process lock for reference video download/copy."""
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+    with lock_path.open("a", encoding="utf-8") as lock_file:
+        if os.name == "posix":
+            import fcntl
+
+            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
+        try:
+            yield
+        finally:
+            if os.name == "posix":
+                import fcntl
+
+                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
+
+
+def _merge_downloaded_reference_content(
+    staging_dir: Path,
+    local_dir: Path,
+) -> int:
+    copied_roots = 0
+
+    tier_root = staging_dir / REFERENCE_VIDEOS_DIRNAME
+    if tier_root.exists():
+        dst_tier_root = local_dir / REFERENCE_VIDEOS_DIRNAME
+        shutil.copytree(tier_root, dst_tier_root, dirs_exist_ok=True)
+        copied_roots += 1
+
+    # Legacy default layout support: <staging>/<GPU>_reference_videos/...
+    legacy_reference_dirs = _discover_reference_dirs(staging_dir)
+    for legacy_dir in legacy_reference_dirs:
+        dst_legacy_dir = local_dir / legacy_dir.name
+        shutil.copytree(legacy_dir, dst_legacy_dir, dirs_exist_ok=True)
+        copied_roots += 1
+
+    return copied_roots
+
+
+def _resolve_device_reference_folder(
+    *,
+    explicit_device_folder: str | None,
+    reference_dir: Path | None,
+) -> str:
+    if explicit_device_folder and explicit_device_folder.strip():
+        return explicit_device_folder.strip()
+    if reference_dir is not None and reference_dir.name.endswith("_reference_videos"):
+        return reference_dir.name
+    return DEFAULT_DEVICE_REFERENCE_FOLDER
+
+
+def _reference_tier_root(base_dir: Path, quality_tier: str) -> Path:
+    return base_dir / REFERENCE_VIDEOS_DIRNAME / quality_tier
+
+
+def _resolve_quality_tiers(quality_tier: str) -> list[str]:
+    if quality_tier == ALL_OUTPUT_QUALITY_TIERS:
+        return list(QUALITY_TIERS)
+    if quality_tier not in QUALITY_TIERS:
+        raise ValueError(f"Unsupported quality tier: {quality_tier}")
+    return [quality_tier]
+
+
+def _discover_reference_dirs_for_tier(
+    base_dir: Path,
+    quality_tier: str,
+) -> list[Path]:
+    tier_root = _reference_tier_root(base_dir, quality_tier)
+    tier_dirs = _discover_reference_dirs(tier_root)
+    if tier_dirs:
+        return tier_dirs
+    if quality_tier == DEFAULT_OUTPUT_QUALITY_TIER:
+        # Keep supporting legacy default refs at <ssim_dir>/<GPU>_reference_videos.
+        return _discover_reference_dirs(base_dir)
+    return []
+
+
+def _has_local_reference_videos(base_dir: Path, quality_tier: str) -> bool:
+    # Check for completion marker file first
+    marker_path = base_dir / REFERENCE_VIDEOS_DIRNAME / quality_tier / f".download_complete_{quality_tier}"
+    if not marker_path.exists():
+        return False
+    # Also verify at least one .mp4 exists
+    for ref_dir in _discover_reference_dirs_for_tier(base_dir, quality_tier):
+        for _ in _iter_video_files(ref_dir):
+            return True
+    return False
+
+
+def _get_hf_token() -> str | None:
+    for key in HF_TOKEN_ENV_KEYS:
+        value = os.environ.get(key, "").strip()
+        if value:
+            return value
+
+    return None
+
+
+def _load_hf_sdk():
+    try:
+        from huggingface_hub import HfApi, snapshot_download
+    except ImportError as exc:
+        raise RuntimeError(
+            "huggingface_hub is required for download/upload.\nInstall with: pip install huggingface_hub"
+        ) from exc
+    return HfApi, snapshot_download
+
+
+def copy_generated_to_reference(
+    generated_dir: Path,
+    reference_dir: Path,
+    *,
+    dry_run: bool = False,
+) -> int:
+    if not generated_dir.exists():
+        raise FileNotFoundError(f"Generated directory not found: {generated_dir}")
+
+    copied = 0
+    for video_file in _iter_video_files(generated_dir):
+        rel = video_file.relative_to(generated_dir)
+        dst = reference_dir / rel
+        if dry_run:
+            print(f"[dry-run] {video_file} -> {dst}")
+        else:
+            dst.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(video_file, dst)
+            print(f"Copied: {rel}")
+        copied += 1
+    return copied
+
+
+def download_reference_videos(
+    *,
+    repo_id: str,
+    repo_type: str,
+    local_dir: Path,
+    quality_tiers: Sequence[str] | None = None,
+    device_folders: Sequence[str] | None = None,
+    token: str | None = None,
+) -> None:
+    _, snapshot_download = _load_hf_sdk()
+
+    local_dir.mkdir(parents=True, exist_ok=True)
+    tiers = tuple(quality_tiers or QUALITY_TIERS)
+    for quality_tier in tiers:
+        if quality_tier not in QUALITY_TIERS:
+            raise ValueError(f"Unsupported quality tier: {quality_tier}")
+    allow_patterns: list[str] = []
+    for quality_tier in tiers:
+        tier_prefix = f"{REFERENCE_VIDEOS_DIRNAME}/{quality_tier}"
+        if device_folders:
+            allow_patterns.extend(f"{tier_prefix}/{folder}/**" for folder in device_folders)
+        else:
+            allow_patterns.append(f"{tier_prefix}/*_reference_videos/**")
+
+        # Keep compatibility for older HF layout where default refs lived at
+        # repo root as <GPU>_reference_videos.
+        if quality_tier == DEFAULT_OUTPUT_QUALITY_TIER:
+            if device_folders:
+                allow_patterns.extend(f"{folder}/**" for folder in device_folders)
+            else:
+                allow_patterns.append("*_reference_videos/**")
+
+    # Work around long-path issues in some huggingface_hub versions:
+    # avoid `local_dir` mode and download into a short cache dir instead,
+    # then merge files into local_dir.
+    with tempfile.TemporaryDirectory(prefix="fv2-ssim-cache-") as tmp_cache_dir:
+        snapshot_path = snapshot_download(
+            repo_id=repo_id,
+            repo_type=repo_type,
+            cache_dir=tmp_cache_dir,
+            token=token,
+            allow_patterns=allow_patterns,
+        )
+        staging_dir = Path(snapshot_path)
+        copied_roots = _merge_downloaded_reference_content(
+            staging_dir=staging_dir,
+            local_dir=local_dir,
+        )
+        if copied_roots == 0:
+            raise RuntimeError(
+                "HF download completed but no reference_videos content was found in downloaded artifacts."
+            )
+        for quality_tier in tiers:
+            marker_path = local_dir / REFERENCE_VIDEOS_DIRNAME / quality_tier / f".download_complete_{quality_tier}"
+            marker_path.parent.mkdir(parents=True, exist_ok=True)
+            marker_path.touch()
+
+
+def upload_reference_videos(
+    *,
+    repo_id: str,
+    repo_type: str,
+    reference_dirs_by_tier: Sequence[tuple[str, Path]],
+    token: str,
+    private: bool,
+) -> None:
+    HfApi, _ = _load_hf_sdk()
+    api = HfApi(token=token)
+    api.create_repo(
+        repo_id=repo_id,
+        repo_type=repo_type,
+        private=private,
+        exist_ok=True,
+    )
+
+    for quality_tier, reference_dir in reference_dirs_by_tier:
+        if not reference_dir.exists():
+            raise FileNotFoundError(f"Reference directory not found: {reference_dir}")
+        path_in_repo = f"{REFERENCE_VIDEOS_DIRNAME}/{quality_tier}/{reference_dir.name}"
+        print(f"Uploading {reference_dir.name} ({quality_tier}) to {repo_id} ...")
+        api.upload_folder(
+            repo_id=repo_id,
+            repo_type=repo_type,
+            folder_path=str(reference_dir),
+            path_in_repo=path_in_repo,
+            token=token,
+        )
+
+
+def _resolve_upload_reference_dirs(
+    *,
+    base_dir: Path,
+    quality_tier: str,
+    explicit_reference_dirs: Sequence[str],
+    device_folders: Sequence[str] | None = None,
+) -> list[tuple[str, Path]]:
+    tiers = _resolve_quality_tiers(quality_tier)
+    selected_device_folders = set(device_folders or [])
+
+    if explicit_reference_dirs:
+        if len(tiers) != 1:
+            raise RuntimeError("--reference-dir requires a single --quality-tier (default or full_quality).")
+        resolved_tier = tiers[0]
+        return [(resolved_tier, Path(p).resolve()) for p in explicit_reference_dirs]
+
+    resolved: list[tuple[str, Path]] = []
+    for tier in tiers:
+        for reference_dir in _discover_reference_dirs_for_tier(base_dir, tier):
+            if selected_device_folders and reference_dir.name not in selected_device_folders:
+                continue
+            resolved.append((tier, reference_dir.resolve()))
+    return resolved
+
+
+def ensure_reference_videos_available(
+    *,
+    local_dir: Path | None = None,
+    repo_id: str | None = None,
+    repo_type: str | None = None,
+    quality_tier: str = DEFAULT_OUTPUT_QUALITY_TIER,
+) -> bool:
+    """Return True if downloaded from HF, False if already present locally."""
+    if quality_tier not in QUALITY_TIERS:
+        raise ValueError(f"Unsupported quality tier: {quality_tier}")
+    target_dir = local_dir or _ssim_dir()
+    lock_path = target_dir / ".reference_videos_download.lock"
+    with _exclusive_download_lock(lock_path):
+        if _has_local_reference_videos(target_dir, quality_tier):
+            print(f"Reference videos ({quality_tier}) already available at {target_dir}")
+            return False
+
+        resolved_repo_id = repo_id or _default_repo_id()
+        resolved_repo_type = repo_type or _default_repo_type()
+        if not resolved_repo_id:
+            raise RuntimeError(
+                f"No local reference videos found and no HF repo configured.\nSet {HF_REPO_ENV_KEY} or pass --repo-id."
+            )
+
+        print(f"Repo ID: {resolved_repo_id}")
+        print(f"Quality tier: {quality_tier}")
+        print(f"No local {quality_tier} reference videos found under {target_dir}. Starting download...")
+        try:
+            download_reference_videos(
+                repo_id=resolved_repo_id,
+                repo_type=resolved_repo_type,
+                local_dir=target_dir,
+                quality_tiers=[quality_tier],
+            )
+            print(f"Download completed for {quality_tier} reference videos.")
+        except Exception as exc:
+            print(f"ERROR: Failed to download {quality_tier} reference videos from {resolved_repo_id}.")
+            print(
+                f"Suggested command to retry: "
+                f"python fastvideo/tests/ssim/reference_videos_cli.py download "
+                f"--quality-tier {quality_tier}"
+            )
+            raise
+
+        if not _has_local_reference_videos(target_dir, quality_tier):
+            raise RuntimeError(f"HF download completed but no {quality_tier} *_reference_videos content found.")
+        return True
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="reference_videos_cli",
+        description=(
+            "Manage SSIM reference videos locally and on Hugging Face.\n"
+            "This tool can copy generated videos into reference folders,\n"
+            "download references from a public HF repo, and upload updates."
+        ),
+        epilog=(
+            "Examples:\n"
+            "  # Copy default generated videos into local L40S refs\n"
+            "  python fastvideo/tests/ssim/reference_videos_cli.py copy-local \\\n"
+            "    --quality-tier default \\\n"
+            "    --device-folder L40S_reference_videos \\\n"
+            "    --reference-dir fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos\n\n"
+            "  # Download both default and full-quality references from HF\n"
+            "  python fastvideo/tests/ssim/reference_videos_cli.py download \\\n"
+            "    --repo-id FastVideo/ssim-reference-videos\n\n"
+            "  # Upload both tiers and all GPU folders (fails if token is unset)\n"
+            "  python fastvideo/tests/ssim/reference_videos_cli.py upload \\\n"
+            "    --repo-id FastVideo/ssim-reference-videos"
+        ),
+        formatter_class=argparse.RawTextHelpFormatter,
+    )
+
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    copy_parser = subparsers.add_parser(
+        "copy-local",
+        help="Copy generated videos into a local reference folder.",
+    )
+    copy_parser.add_argument(
+        "--quality-tier",
+        choices=(DEFAULT_OUTPUT_QUALITY_TIER, FULL_OUTPUT_QUALITY_TIER),
+        default=DEFAULT_OUTPUT_QUALITY_TIER,
+        help="Generated quality tier to copy from.",
+    )
+    copy_parser.add_argument(
+        "--generated-dir",
+        type=Path,
+        default=None,
+        help=("Source generated directory. Default: <ssim_dir>/generated_videos/<quality-tier>/<device-folder>"),
+    )
+    copy_parser.add_argument(
+        "--device-folder",
+        default=None,
+        help=("GPU folder name (e.g., H200_reference_videos) used to build default generated/reference paths."),
+    )
+    copy_parser.add_argument(
+        "--reference-dir",
+        type=Path,
+        default=None,
+        help="Destination reference directory.",
+    )
+    copy_parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print copy operations without writing files.",
+    )
+
+    download_parser = subparsers.add_parser(
+        "download",
+        help="Download reference videos from Hugging Face.",
+    )
+    download_parser.add_argument(
+        "--repo-id",
+        default=_default_repo_id(),
+        help=f"HF repo id. Default: env {HF_REPO_ENV_KEY} or {DEFAULT_REPO_ID}",
+    )
+    download_parser.add_argument(
+        "--repo-type",
+        default=_default_repo_type(),
+        help=f"HF repo type. Default: env {HF_REPO_TYPE_ENV_KEY} or {DEFAULT_REPO_TYPE}",
+    )
+    download_parser.add_argument(
+        "--local-dir",
+        type=Path,
+        default=_ssim_dir(),
+        help=(
+            "Local SSIM directory where references are stored under "
+            "reference_videos/<quality-tier>/<GPU>_reference_videos."
+        ),
+    )
+    download_parser.add_argument(
+        "--quality-tier",
+        choices=(
+            DEFAULT_OUTPUT_QUALITY_TIER,
+            FULL_OUTPUT_QUALITY_TIER,
+            ALL_OUTPUT_QUALITY_TIERS,
+        ),
+        default=ALL_OUTPUT_QUALITY_TIERS,
+        help="Quality tier to download (default: all).",
+    )
+    download_parser.add_argument(
+        "--device-folder",
+        action="append",
+        default=[],
+        help="Specific reference folder to fetch (repeatable), e.g. L40S_reference_videos",
+    )
+
+    upload_parser = subparsers.add_parser(
+        "upload",
+        help="Upload local reference folders to Hugging Face.",
+    )
+    upload_parser.add_argument(
+        "--repo-id",
+        default=_default_repo_id(),
+        help=f"HF repo id. Default: env {HF_REPO_ENV_KEY} or {DEFAULT_REPO_ID}",
+    )
+    upload_parser.add_argument(
+        "--repo-type",
+        default=_default_repo_type(),
+        help=f"HF repo type. Default: env {HF_REPO_TYPE_ENV_KEY} or {DEFAULT_REPO_TYPE}",
+    )
+    upload_parser.add_argument(
+        "--reference-dir",
+        action="append",
+        default=[],
+        help=(
+            "Reference directory to upload (repeatable). "
+            "Requires single --quality-tier. "
+            "Default: discover folders under --base-dir for chosen tier(s)."
+        ),
+    )
+    upload_parser.add_argument(
+        "--quality-tier",
+        choices=(
+            DEFAULT_OUTPUT_QUALITY_TIER,
+            FULL_OUTPUT_QUALITY_TIER,
+            ALL_OUTPUT_QUALITY_TIERS,
+        ),
+        default=ALL_OUTPUT_QUALITY_TIERS,
+        help="Quality tier to upload (default: all).",
+    )
+    upload_parser.add_argument(
+        "--device-folder",
+        action="append",
+        default=[],
+        help=("Specific GPU reference folder to upload (repeatable), e.g. H200_reference_videos."),
+    )
+    upload_parser.add_argument(
+        "--base-dir",
+        type=Path,
+        default=_ssim_dir(),
+        help=("Base SSIM directory that contains reference_videos/<quality-tier>/<GPU>_reference_videos."),
+    )
+    upload_parser.add_argument(
+        "--private",
+        action="store_true",
+        help="Create/use a private repo instead of public.",
+    )
+
+    ensure_parser = subparsers.add_parser(
+        "ensure",
+        help="Download refs only when the selected local quality tier is missing.",
+    )
+    ensure_parser.add_argument(
+        "--repo-id",
+        default=_default_repo_id(),
+        help=f"HF repo id. Default: env {HF_REPO_ENV_KEY} or {DEFAULT_REPO_ID}",
+    )
+    ensure_parser.add_argument(
+        "--repo-type",
+        default=_default_repo_type(),
+        help=f"HF repo type. Default: env {HF_REPO_TYPE_ENV_KEY} or {DEFAULT_REPO_TYPE}",
+    )
+    ensure_parser.add_argument(
+        "--local-dir",
+        type=Path,
+        default=_ssim_dir(),
+        help=("Local SSIM directory that should contain reference_videos/<quality-tier>/<GPU>_reference_videos."),
+    )
+    ensure_parser.add_argument(
+        "--quality-tier",
+        choices=(
+            DEFAULT_OUTPUT_QUALITY_TIER,
+            FULL_OUTPUT_QUALITY_TIER,
+        ),
+        default=DEFAULT_OUTPUT_QUALITY_TIER,
+        help="Quality tier to ensure locally.",
+    )
+
+    return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "copy-local":
+        device_folder = _resolve_device_reference_folder(
+            explicit_device_folder=args.device_folder,
+            reference_dir=args.reference_dir,
+        )
+        generated_dir = args.generated_dir
+        if generated_dir is None:
+            tier_root = _ssim_dir() / "generated_videos" / args.quality_tier
+            tier_device_dir = tier_root / device_folder
+            if tier_device_dir.exists():
+                generated_dir = tier_device_dir
+            else:
+                # Backward compatibility for pre-device-folder layout.
+                generated_dir = tier_root
+        reference_dir = (
+            args.reference_dir
+            if args.reference_dir is not None
+            else (_ssim_dir() / REFERENCE_VIDEOS_DIRNAME / args.quality_tier / device_folder)
+        )
+        copied = copy_generated_to_reference(
+            generated_dir=generated_dir,
+            reference_dir=reference_dir,
+            dry_run=args.dry_run,
+        )
+        print(f"Done. {'Would copy' if args.dry_run else 'Copied'} {copied} video files.")
+        if not args.dry_run and copied > 0:
+            marker = (
+                _ssim_dir() / REFERENCE_VIDEOS_DIRNAME / args.quality_tier / f".download_complete_{args.quality_tier}"
+            )
+            marker.parent.mkdir(parents=True, exist_ok=True)
+            marker.touch()
+        return 0
+
+    if args.command == "download":
+        quality_tiers = _resolve_quality_tiers(args.quality_tier)
+        download_reference_videos(
+            repo_id=args.repo_id,
+            repo_type=args.repo_type,
+            local_dir=args.local_dir,
+            quality_tiers=quality_tiers,
+            device_folders=args.device_folder or None,
+        )
+        print("Download complete for quality tier(s): " + ", ".join(quality_tiers) + ".")
+        return 0
+
+    if args.command == "upload":
+        token = _get_hf_token()
+        if token is None:
+            raise RuntimeError(
+                "Hugging Face API key is required for upload. Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN."
+            )
+
+        reference_dirs_by_tier = _resolve_upload_reference_dirs(
+            base_dir=args.base_dir.resolve(),
+            quality_tier=args.quality_tier,
+            explicit_reference_dirs=args.reference_dir,
+            device_folders=args.device_folder or None,
+        )
+        if not reference_dirs_by_tier:
+            raise RuntimeError(
+                f"No *_reference_videos directories found for the selected quality tier(s) under {args.base_dir}"
+            )
+
+        upload_reference_videos(
+            repo_id=args.repo_id,
+            repo_type=args.repo_type,
+            reference_dirs_by_tier=reference_dirs_by_tier,
+            token=token,
+            private=args.private,
+        )
+        print("Upload complete.")
+        return 0
+
+    if args.command == "ensure":
+        downloaded = ensure_reference_videos_available(
+            local_dir=args.local_dir,
+            repo_id=args.repo_id,
+            repo_type=args.repo_type,
+            quality_tier=args.quality_tier,
+        )
+        if downloaded:
+            print(f"{args.quality_tier} reference videos were missing and have been downloaded.")
+        else:
+            print(f"{args.quality_tier} reference videos already exist locally.")
+        return 0
+
+    parser.error(f"Unknown command: {args.command}")
+    return 2
+
+
+if __name__ == "__main__":
+    try:
+        raise SystemExit(main())
+    except KeyboardInterrupt:
+        print("Interrupted.", file=sys.stderr)
+        raise SystemExit(130)

fastvideo/tests/ssim/test_causal_similarity.py

@@ -1,32 +1,22 @@
 # SPDX-License-Identifier: Apache-2.0
-import json
 import os
 
-import torch
 import pytest
 
-from fastvideo import VideoGenerator
+from fastvideo.configs.sample.wan import (
+    SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam,
+)
 from fastvideo.logger import init_logger
-from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
-from fastvideo.worker.multiproc_executor import MultiprocExecutor
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    resolve_inference_device_reference_folder,
+    run_text_to_video_similarity_test,
+)
 
 logger = init_logger(__name__)
 
 REQUIRED_GPUS = 1
 
-device_name = torch.cuda.get_device_name()
-device_reference_folder_suffix = '_reference_videos'
-
-if "A40" in device_name:
-    device_reference_folder = "A40" + device_reference_folder_suffix
-elif "L40S" in device_name:
-    device_reference_folder = "L40S" + device_reference_folder_suffix
-else:
-    # device_reference_folder = "L40S" + device_reference_folder_suffix
-    logger.warning(f"Unsupported device for ssim tests: {device_name}")
-    # raise ValueError(f"Unsupported device for ssim tests: {device_name}")
-
-# Base parameters from the shell script
+device_reference_folder = resolve_inference_device_reference_folder(logger)
 
 SF_WAN_T2V_PARAMS = {
     "num_gpus": 1,
@@ -40,119 +30,53 @@ SF_WAN_T2V_PARAMS = {
     "tp_size": 1,
 }
 
+_SF_WAN_T2V_FULL_QUALITY_DEFAULTS = SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam()
+SF_WAN_T2V_FULL_QUALITY_PARAMS = {
+    "num_gpus": SF_WAN_T2V_PARAMS["num_gpus"],
+    "model_path": SF_WAN_T2V_PARAMS["model_path"],
+    "height": _SF_WAN_T2V_FULL_QUALITY_DEFAULTS.height,
+    "width": _SF_WAN_T2V_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": SF_WAN_T2V_PARAMS["num_frames"],
+    "num_inference_steps": (_SF_WAN_T2V_FULL_QUALITY_DEFAULTS.num_inference_steps),
+    "guidance_scale": (_SF_WAN_T2V_FULL_QUALITY_DEFAULTS.guidance_scale),
+    "seed": _SF_WAN_T2V_FULL_QUALITY_DEFAULTS.seed,
+    "sp_size": SF_WAN_T2V_PARAMS["sp_size"],
+    "tp_size": SF_WAN_T2V_PARAMS["tp_size"],
+    "neg_prompt": (_SF_WAN_T2V_FULL_QUALITY_DEFAULTS.negative_prompt),
+}
 
 MODEL_TO_PARAMS = {
     "SFWan2.1-T2V-1.3B-Diffusers": SF_WAN_T2V_PARAMS,
 }
-
-I2V_MODEL_TO_PARAMS = {
+FULL_QUALITY_MODEL_TO_PARAMS = {
+    "SFWan2.1-T2V-1.3B-Diffusers": SF_WAN_T2V_FULL_QUALITY_PARAMS,
 }
 
 TEST_PROMPTS = [
-    "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
-    # "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature."
+    "Will Smith casually eats noodles, his relaxed demeanor contrasting "
+    "with the energetic background of a bustling street food market. "
+    "The scene captures a mix of humor and authenticity. Mid-shot "
+    "framing, vibrant lighting.",
 ]
 
-I2V_TEST_PROMPTS = [
-    "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.",
-]
-
-I2V_IMAGE_PATHS = [
-    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg",
-]
 
 @pytest.mark.parametrize("prompt", TEST_PROMPTS)
-@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
+@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"])
 @pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys()))
-def test_causal_similarity(prompt, ATTENTION_BACKEND, model_id):
-    """
-    Test that runs inference with different parameters and compares the output
-    to reference videos using SSIM.
-    """
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
-
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-
-    base_output_dir = os.path.join(script_dir, 'generated_videos', model_id)
-    output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
-    output_video_name = f"{prompt[:100].strip()}.mp4"
-
-    os.makedirs(output_dir, exist_ok=True)
-
-    BASE_PARAMS = MODEL_TO_PARAMS[model_id]
-    num_inference_steps = BASE_PARAMS["num_inference_steps"]
-
-    init_kwargs = {
-        "num_gpus": BASE_PARAMS["num_gpus"],
-        "sp_size": BASE_PARAMS["sp_size"],
-        "tp_size": BASE_PARAMS["tp_size"],
-        "dit_cpu_offload": True,
-    }
-    if BASE_PARAMS.get("vae_sp"):
-        init_kwargs["vae_sp"] = True
-        init_kwargs["vae_tiling"] = True
-    #if "text-encoder-precision" in BASE_PARAMS:
-    #    init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"]
-
-    generation_kwargs = {
-        "num_inference_steps": num_inference_steps,
-        "output_path": output_dir,
-        "height": BASE_PARAMS["height"],
-        "width": BASE_PARAMS["width"],
-        "num_frames": BASE_PARAMS["num_frames"],
-        "seed": BASE_PARAMS["seed"],
-    }
-    if "neg_prompt" in BASE_PARAMS:
-        generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"]
-
-    generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs)
-    generator.generate_video(prompt, **generation_kwargs)
-
-    if isinstance(generator.executor, MultiprocExecutor):
-        generator.executor.shutdown()
-
-    assert os.path.exists(
-        output_dir), f"Output video was not generated at {output_dir}"
-
-    reference_folder = os.path.join(script_dir, device_reference_folder, model_id, ATTENTION_BACKEND)
-
-    if not os.path.exists(reference_folder):
-        logger.error("Reference folder missing")
-        raise FileNotFoundError(
-            f"Reference video folder does not exist: {reference_folder}")
-
-    # Find the matching reference video based on the prompt
-    reference_video_name = None
-
-    for filename in os.listdir(reference_folder):
-        if filename.endswith('.mp4') and prompt[:100].strip() in filename:
-            reference_video_name = filename
-            break
-
-    if not reference_video_name:
-        logger.error(f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}")
-        raise FileNotFoundError(f"Reference video missing")
-
-    reference_video_path = os.path.join(reference_folder, reference_video_name)
-    generated_video_path = os.path.join(output_dir, output_video_name)
-
-    logger.info(
-        f"Computing SSIM between {reference_video_path} and {generated_video_path}"
+def test_causal_similarity(
+    prompt: str,
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    run_text_to_video_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        attention_backend_name=attention_backend_name,
+        model_id=model_id,
+        default_params_map=MODEL_TO_PARAMS,
+        full_quality_params_map=FULL_QUALITY_MODEL_TO_PARAMS,
+        min_acceptable_ssim=0.98,
+        init_kwargs_override={"dit_cpu_offload": True},
     )
-    ssim_values = compute_video_ssim_torchvision(reference_video_path,
-                                                 generated_video_path,
-                                                 use_ms_ssim=True)
-
-    mean_ssim = ssim_values[0]
-    logger.info(f"SSIM mean value: {mean_ssim}")
-    logger.info(f"Writing SSIM results to directory: {output_dir}")
-
-    success = write_ssim_results(output_dir, ssim_values, reference_video_path,
-                                 generated_video_path, num_inference_steps,
-                                 prompt)
-
-    if not success:
-        logger.error("Failed to write SSIM results to file")
-
-    min_acceptable_ssim = 0.98
-    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"

fastvideo/tests/ssim/test_gamecraft_similarity.py

@@ -6,7 +6,8 @@ Generates a video with deterministic seed and camera trajectory,
 then compares against a device-specific reference video via MS-SSIM.
 
 Reference videos must be pre-generated and stored under:
-    <device>_reference_videos/HunyuanGameCraft/<ATTENTION_BACKEND>/
+    reference_videos/<quality-tier>/<device>_reference_videos/HunyuanGameCraft/
+    <ATTENTION_BACKEND>/
 
 To create initial reference videos, run this test once and copy the
 generated videos into the appropriate reference folder.
@@ -17,7 +18,17 @@ import torch
 import pytest
 
 from fastvideo import VideoGenerator
+from fastvideo.configs.sample.hunyuangamecraft import (
+    HunyuanGameCraftSamplingParam,
+)
 from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.reference_utils import (
+    build_generated_output_dir,
+    build_reference_folder_path,
+    get_cuda_device_name,
+    resolve_device_reference_folder,
+    select_ssim_params,
+)
 from fastvideo.tests.utils import (
     compute_video_ssim_torchvision,
     write_ssim_results,
@@ -31,20 +42,18 @@ REQUIRED_GPUS = 1
 # ---------------------------------------------------------------------------
 # Device-dependent reference folder
 # ---------------------------------------------------------------------------
-device_name = torch.cuda.get_device_name()
-device_reference_folder_suffix = "_reference_videos"
-
-if "A40" in device_name:
-    device_reference_folder = "A40" + device_reference_folder_suffix
-elif "L40S" in device_name:
-    device_reference_folder = "L40S" + device_reference_folder_suffix
-elif "A100" in device_name:
-    device_reference_folder = "A100" + device_reference_folder_suffix
-elif "H100" in device_name:
-    device_reference_folder = "H100" + device_reference_folder_suffix
-else:
-    device_reference_folder = "Unknown" + device_reference_folder_suffix
-    logger.warning(f"Unsupported device for ssim tests: {device_name}")
+device_reference_folder = resolve_device_reference_folder(
+    (
+        ("A40", "A40"),
+        ("L40S", "L40S"),
+        ("A100", "A100"),
+        ("H100", "H100"),
+        ("H200", "H200"),
+    ),
+    device_name=get_cuda_device_name(),
+    unknown_device_prefix="Unknown",
+    logger=logger,
+)
 
 # ---------------------------------------------------------------------------
 # Helpers – camera trajectory (self-contained, no official repo dependency)
@@ -82,6 +91,21 @@ GAMECRAFT_T2V_PARAMS = {
     "negative_prompt": "",
 }
 
+_GAMECRAFT_FULL_QUALITY_DEFAULTS = HunyuanGameCraftSamplingParam()
+GAMECRAFT_T2V_FULL_QUALITY_PARAMS = {
+    "num_gpus": GAMECRAFT_T2V_PARAMS["num_gpus"],
+    "model_path": GAMECRAFT_T2V_PARAMS["model_path"],
+    "height": _GAMECRAFT_FULL_QUALITY_DEFAULTS.height,
+    "width": _GAMECRAFT_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": GAMECRAFT_T2V_PARAMS["num_frames"],  # default num_frames: 33
+    "num_inference_steps": _GAMECRAFT_FULL_QUALITY_DEFAULTS.num_inference_steps,
+    "guidance_scale": _GAMECRAFT_FULL_QUALITY_DEFAULTS.guidance_scale,
+    "seed": _GAMECRAFT_FULL_QUALITY_DEFAULTS.seed,
+    "action": GAMECRAFT_T2V_PARAMS["action"],
+    "action_speed": GAMECRAFT_T2V_PARAMS["action_speed"],
+    "negative_prompt": _GAMECRAFT_FULL_QUALITY_DEFAULTS.negative_prompt,
+}
+
 GAMECRAFT_I2V_PARAMS = {
     **GAMECRAFT_T2V_PARAMS,
     "image_path": (
@@ -89,14 +113,24 @@ GAMECRAFT_I2V_PARAMS = {
         "resolve/main/diffusers/astronaut.jpg"
     ),
 }
+GAMECRAFT_I2V_FULL_QUALITY_PARAMS = {
+    **GAMECRAFT_T2V_FULL_QUALITY_PARAMS,
+    "image_path": GAMECRAFT_I2V_PARAMS["image_path"],
+}
 
 MODEL_TO_PARAMS = {
     "HunyuanGameCraft-T2V": GAMECRAFT_T2V_PARAMS,
 }
+FULL_QUALITY_MODEL_TO_PARAMS = {
+    "HunyuanGameCraft-T2V": GAMECRAFT_T2V_FULL_QUALITY_PARAMS,
+}
 
 I2V_MODEL_TO_PARAMS = {
     "HunyuanGameCraft-I2V": GAMECRAFT_I2V_PARAMS,
 }
+FULL_QUALITY_I2V_MODEL_TO_PARAMS = {
+    "HunyuanGameCraft-I2V": GAMECRAFT_I2V_FULL_QUALITY_PARAMS,
+}
 
 TEST_PROMPTS = [
     "A majestic ancient temple stands under a clear blue sky, its grandeur highlighted by towering Doric columns and intricate architectural details.",
@@ -119,13 +153,21 @@ def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id):
     os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
 
     script_dir = os.path.dirname(os.path.abspath(__file__))
-    base_output_dir = os.path.join(script_dir, "generated_videos", model_id)
-    output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
+    output_dir = build_generated_output_dir(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
+    )
     output_video_name = f"{prompt[:100].strip()}.mp4"
 
     os.makedirs(output_dir, exist_ok=True)
 
-    BASE_PARAMS = MODEL_TO_PARAMS[model_id]
+    params_map = select_ssim_params(
+        MODEL_TO_PARAMS,
+        FULL_QUALITY_MODEL_TO_PARAMS,
+    )
+    BASE_PARAMS = params_map[model_id]
     num_inference_steps = BASE_PARAMS["num_inference_steps"]
 
     # Build camera trajectory
@@ -175,8 +217,11 @@ def test_gamecraft_t2v_similarity(prompt, ATTENTION_BACKEND, model_id):
     )
 
     # Compare to reference
-    reference_folder = os.path.join(
-        script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
     )
 
     if not os.path.exists(reference_folder):
@@ -244,13 +289,21 @@ def test_gamecraft_i2v_similarity(prompt, ATTENTION_BACKEND, model_id):
     os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
 
     script_dir = os.path.dirname(os.path.abspath(__file__))
-    base_output_dir = os.path.join(script_dir, "generated_videos", model_id)
-    output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
+    output_dir = build_generated_output_dir(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
+    )
     output_video_name = f"{prompt[:100].strip()}.mp4"
 
     os.makedirs(output_dir, exist_ok=True)
 
-    BASE_PARAMS = I2V_MODEL_TO_PARAMS[model_id]
+    params_map = select_ssim_params(
+        I2V_MODEL_TO_PARAMS,
+        FULL_QUALITY_I2V_MODEL_TO_PARAMS,
+    )
+    BASE_PARAMS = params_map[model_id]
     num_inference_steps = BASE_PARAMS["num_inference_steps"]
 
     # Build camera trajectory
@@ -301,8 +354,11 @@ def test_gamecraft_i2v_similarity(prompt, ATTENTION_BACKEND, model_id):
     )
 
     # Compare to reference
-    reference_folder = os.path.join(
-        script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
     )
 
     if not os.path.exists(reference_folder):

fastvideo/tests/ssim/test_inference_similarity.py

@@ -1,389 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-import json
-import os
-from contextlib import contextmanager
-from typing import Iterator
-
-import torch
-import pytest
-
-from fastvideo import VideoGenerator
-from fastvideo.logger import init_logger
-from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
-from fastvideo.worker.multiproc_executor import MultiprocExecutor
-
-logger = init_logger(__name__)
-
-REQUIRED_GPUS = 2
-
-
-@contextmanager
-def _attention_backend(backend: str) -> Iterator[None]:
-    previous = os.environ.get("FASTVIDEO_ATTENTION_BACKEND")
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = backend
-    try:
-        yield
-    finally:
-        if previous is None:
-            os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None)
-        else:
-            os.environ["FASTVIDEO_ATTENTION_BACKEND"] = previous
-
-
-def _shutdown_executor(generator: VideoGenerator | None) -> None:
-    if generator is None:
-        return
-    if isinstance(generator.executor, MultiprocExecutor):
-        generator.executor.shutdown()
-
-device_name = torch.cuda.get_device_name()
-device_reference_folder_suffix = '_reference_videos'
-
-if "A40" in device_name:
-    device_reference_folder = "A40" + device_reference_folder_suffix
-elif "L40S" in device_name:
-    device_reference_folder = "L40S" + device_reference_folder_suffix
-else:
-    # device_reference_folder = "L40S" + device_reference_folder_suffix
-    logger.warning(f"Unsupported device for ssim tests: {device_name}")
-    # raise ValueError(f"Unsupported device for ssim tests: {device_name}")
-
-# Base parameters from the shell script
-HUNYUAN_PARAMS = {
-    "num_gpus": 4,
-    "model_path": "FastVideo/FastHunyuan-diffusers",
-    "height": 720,
-    "width": 1280,
-    "num_frames": 45,
-    "num_inference_steps": 2,
-    "guidance_scale": 1,
-    "embedded_cfg_scale": 6,
-    "flow_shift": 17,
-    "seed": 1024,
-    "sp_size": 4,
-    "tp_size": 1,
-    "vae_sp": True,
-    "fps": 24,
-}
-
-WAN_T2V_PARAMS = {
-    "num_gpus": 2,
-    "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
-    "height": 480,
-    "width": 832,
-    "num_frames": 45,
-    "num_inference_steps": 4,
-    "guidance_scale": 3,
-    "embedded_cfg_scale": 6,
-    "flow_shift": 7.0,
-    "seed": 1024,
-    "sp_size": 2,
-    "tp_size": 1,
-    "vae_sp": True,
-    "fps": 24,
-    "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
-    "text-encoder-precision": ("fp32",)
-}
-
-WAN_I2V_PARAMS = {
-    "num_gpus": 2,
-    "model_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
-    "height": 480,
-    "width": 832,
-    "num_frames": 45,
-    "num_inference_steps": 2,
-    "guidance_scale": 5.0,
-    "embedded_cfg_scale": 6,
-    "flow_shift": 7.0,
-    "seed": 1024,
-    "sp_size": 2,
-    "tp_size": 1,
-    "vae_sp": True,
-    "fps": 24,
-    "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
-    "text-encoder-precision": ("fp32",)
-}
-
-# LTX-2 distilled one-stage params (no refine/upscale)
-# Official defaults: height=512, width=768, num_frames=121, fps=24, seed=10
-# Using num_frames=41 for faster CI (still valid: 41 = 8×5 + 1)
-LTX2_T2V_PARAMS = {
-    "num_gpus": 2,
-    "model_path": "FastVideo/LTX2-Distilled-Diffusers",
-    "height": 512,
-    "width": 768,
-    "num_frames": 41,  # Shorter for CI; official default is 121
-    "num_inference_steps": 8,  # Distilled uses 8 steps
-    "guidance_scale": 1.0,  # No CFG for distilled
-    "embedded_cfg_scale": 6,
-    "seed": 1024,
-    "sp_size": 2,
-    "tp_size": 1,
-    "fps": 24,
-    "neg_prompt": (
-        "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, "
-        "excessive noise, grainy texture, poor lighting, flickering, motion blur, distorted "
-        "proportions, unnatural skin tones, deformed facial features, asymmetrical face, "
-        "missing facial features, extra limbs, disfigured hands, wrong hand count, artifacts "
-        "around text, inconsistent perspective, camera shake, incorrect depth of field, "
-        "background too sharp, background clutter, distracting reflections, harsh shadows, "
-        "inconsistent lighting direction, color banding, cartoonish rendering, 3D CGI look, "
-        "unrealistic materials, uncanny valley effect, incorrect ethnicity, wrong gender, "
-        "exaggerated expressions, wrong gaze direction, mismatched lip sync, silent or muted "
-        "audio, distorted voice, robotic voice, echo, background noise, off-sync audio, "
-        "incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward "
-        "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, "
-        "flat lighting, inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts."
-    ),
-    "ltx2_vae_tiling": True,
-    "ltx2_vae_spatial_tile_size_in_pixels": 512,
-    "ltx2_vae_spatial_tile_overlap_in_pixels": 64,
-    "ltx2_vae_temporal_tile_size_in_frames": 64,
-    "ltx2_vae_temporal_tile_overlap_in_frames": 24,
-}
-
-MODEL_TO_PARAMS = {
-    # "FastHunyuan-diffusers": HUNYUAN_PARAMS,
-    "Wan2.1-T2V-1.3B-Diffusers": WAN_T2V_PARAMS,
-    # "ltx2_diffusers": LTX2_T2V_PARAMS,
-}
-
-I2V_MODEL_TO_PARAMS = {
-    "Wan2.1-I2V-14B-480P-Diffusers": WAN_I2V_PARAMS,
-}
-
-TEST_PROMPTS = [
-    "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
-    # "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature."
-]
-
-I2V_TEST_PROMPTS = [
-    "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.",
-]
-
-I2V_IMAGE_PATHS = [
-    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg",
-]
-
-
-
-@pytest.mark.parametrize("prompt", I2V_TEST_PROMPTS)
-@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
-@pytest.mark.parametrize("model_id", list(I2V_MODEL_TO_PARAMS.keys()))
-def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
-    """
-    Test that runs inference with different parameters and compares the output
-    to reference videos using SSIM.
-    """
-    assert len(I2V_TEST_PROMPTS) == len(I2V_IMAGE_PATHS), "Expect number of prompts equal to number of images"
-    with _attention_backend(ATTENTION_BACKEND):
-        script_dir = os.path.dirname(os.path.abspath(__file__))
-
-        base_output_dir = os.path.join(script_dir, 'generated_videos', model_id)
-        output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
-        output_video_name = f"{prompt[:100].strip()}.mp4"
-
-        os.makedirs(output_dir, exist_ok=True)
-
-        BASE_PARAMS = I2V_MODEL_TO_PARAMS[model_id]
-        num_inference_steps = BASE_PARAMS["num_inference_steps"]
-        image_path = I2V_IMAGE_PATHS[I2V_TEST_PROMPTS.index(prompt)]
-
-        init_kwargs = {
-            "num_gpus": BASE_PARAMS["num_gpus"],
-            "flow_shift": BASE_PARAMS["flow_shift"],
-            "sp_size": BASE_PARAMS["sp_size"],
-            "tp_size": BASE_PARAMS["tp_size"],
-        }
-        if BASE_PARAMS.get("vae_sp"):
-            init_kwargs["vae_sp"] = True
-            init_kwargs["vae_tiling"] = True
-        if "text-encoder-precision" in BASE_PARAMS:
-            init_kwargs["text_encoder_precisions"] = BASE_PARAMS[
-                "text-encoder-precision"]
-
-        generation_kwargs = {
-            "num_inference_steps": num_inference_steps,
-            "output_path": output_dir,
-            "image_path": image_path,
-            "height": BASE_PARAMS["height"],
-            "width": BASE_PARAMS["width"],
-            "num_frames": BASE_PARAMS["num_frames"],
-            "guidance_scale": BASE_PARAMS["guidance_scale"],
-            "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"],
-            "seed": BASE_PARAMS["seed"],
-            "fps": BASE_PARAMS["fps"],
-        }
-        if "neg_prompt" in BASE_PARAMS:
-            generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"]
-
-        generator: VideoGenerator | None = None
-        try:
-            generator = VideoGenerator.from_pretrained(
-                model_path=BASE_PARAMS["model_path"], **init_kwargs)
-            generator.generate_video(prompt, **generation_kwargs)
-        finally:
-            _shutdown_executor(generator)
-
-    assert os.path.exists(
-        output_dir), f"Output video was not generated at {output_dir}"
-
-    reference_folder = os.path.join(script_dir, device_reference_folder, model_id, ATTENTION_BACKEND)
-
-    if not os.path.exists(reference_folder):
-        logger.error("Reference folder missing")
-        raise FileNotFoundError(
-            f"Reference video folder does not exist: {reference_folder}")
-
-    # Find the matching reference video based on the prompt
-    reference_video_name = None
-
-    for filename in os.listdir(reference_folder):
-        if filename.endswith('.mp4') and prompt[:100].strip() in filename:
-            reference_video_name = filename
-            break
-
-    if not reference_video_name:
-        logger.error(f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}")
-        raise FileNotFoundError(f"Reference video missing")
-
-    reference_video_path = os.path.join(reference_folder, reference_video_name)
-    generated_video_path = os.path.join(output_dir, output_video_name)
-
-    logger.info(
-        f"Computing SSIM between {reference_video_path} and {generated_video_path}"
-    )
-    ssim_values = compute_video_ssim_torchvision(reference_video_path,
-                                                 generated_video_path,
-                                                 use_ms_ssim=True)
-
-    mean_ssim = ssim_values[0]
-    logger.info(f"SSIM mean value: {mean_ssim}")
-    logger.info(f"Writing SSIM results to directory: {output_dir}")
-
-    success = write_ssim_results(output_dir, ssim_values, reference_video_path,
-                                 generated_video_path, num_inference_steps,
-                                 prompt)
-
-    if not success:
-        logger.error("Failed to write SSIM results to file")
-
-    min_acceptable_ssim = 0.97
-    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"
-
-@pytest.mark.parametrize("prompt", TEST_PROMPTS)
-@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN", "TORCH_SDPA"])
-@pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys()))
-def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
-    """
-    Test that runs inference with different parameters and compares the output
-    to reference videos using SSIM.
-    """
-    with _attention_backend(ATTENTION_BACKEND):
-        script_dir = os.path.dirname(os.path.abspath(__file__))
-
-        base_output_dir = os.path.join(script_dir, 'generated_videos', model_id)
-        output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
-        output_video_name = f"{prompt[:100].strip()}.mp4"
-
-        os.makedirs(output_dir, exist_ok=True)
-
-        BASE_PARAMS = MODEL_TO_PARAMS[model_id]
-        num_inference_steps = BASE_PARAMS["num_inference_steps"]
-
-        init_kwargs = {
-            "num_gpus": BASE_PARAMS["num_gpus"],
-            "sp_size": BASE_PARAMS["sp_size"],
-            "tp_size": BASE_PARAMS["tp_size"],
-            "use_fsdp_inference": True,
-            "dit_cpu_offload": False,
-            "dit_layerwise_offload": False,
-        }
-        if "flow_shift" in BASE_PARAMS:
-            init_kwargs["flow_shift"] = BASE_PARAMS["flow_shift"]
-        if BASE_PARAMS.get("vae_sp"):
-            init_kwargs["vae_sp"] = True
-            init_kwargs["vae_tiling"] = True
-        if "text-encoder-precision" in BASE_PARAMS:
-            init_kwargs["text_encoder_precisions"] = BASE_PARAMS[
-                "text-encoder-precision"]
-        # LTX2-specific VAE tiling parameters
-        if BASE_PARAMS.get("ltx2_vae_tiling"):
-            init_kwargs["ltx2_vae_tiling"] = True
-            init_kwargs["ltx2_vae_spatial_tile_size_in_pixels"] = BASE_PARAMS.get(
-                "ltx2_vae_spatial_tile_size_in_pixels", 512)
-            init_kwargs["ltx2_vae_spatial_tile_overlap_in_pixels"] = BASE_PARAMS.get(
-                "ltx2_vae_spatial_tile_overlap_in_pixels", 64)
-            init_kwargs["ltx2_vae_temporal_tile_size_in_frames"] = BASE_PARAMS.get(
-                "ltx2_vae_temporal_tile_size_in_frames", 64)
-            init_kwargs[
-                "ltx2_vae_temporal_tile_overlap_in_frames"] = BASE_PARAMS.get(
-                    "ltx2_vae_temporal_tile_overlap_in_frames", 24)
-
-        generation_kwargs = {
-            "num_inference_steps": num_inference_steps,
-            "output_path": output_dir,
-            "height": BASE_PARAMS["height"],
-            "width": BASE_PARAMS["width"],
-            "num_frames": BASE_PARAMS["num_frames"],
-            "guidance_scale": BASE_PARAMS["guidance_scale"],
-            "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"],
-            "seed": BASE_PARAMS["seed"],
-            "fps": BASE_PARAMS["fps"],
-        }
-        if "neg_prompt" in BASE_PARAMS:
-            generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"]
-
-        generator: VideoGenerator | None = None
-        try:
-            generator = VideoGenerator.from_pretrained(
-                model_path=BASE_PARAMS["model_path"], **init_kwargs)
-            generator.generate_video(prompt, **generation_kwargs)
-        finally:
-            _shutdown_executor(generator)
-
-    assert os.path.exists(
-        output_dir), f"Output video was not generated at {output_dir}"
-
-    reference_folder = os.path.join(script_dir, device_reference_folder, model_id, ATTENTION_BACKEND)
-
-    if not os.path.exists(reference_folder):
-        logger.error("Reference folder missing")
-        raise FileNotFoundError(
-            f"Reference video folder does not exist: {reference_folder}")
-
-    # Find the matching reference video based on the prompt
-    reference_video_name = None
-
-    for filename in os.listdir(reference_folder):
-        if filename.endswith('.mp4') and prompt[:100].strip() in filename:
-            reference_video_name = filename
-            break
-
-    if not reference_video_name:
-        logger.error(f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}")
-        raise FileNotFoundError(f"Reference video missing")
-
-    reference_video_path = os.path.join(reference_folder, reference_video_name)
-    generated_video_path = os.path.join(output_dir, output_video_name)
-
-    logger.info(
-        f"Computing SSIM between {reference_video_path} and {generated_video_path}"
-    )
-    ssim_values = compute_video_ssim_torchvision(reference_video_path,
-                                                 generated_video_path,
-                                                 use_ms_ssim=True)
-
-    mean_ssim = ssim_values[0]
-    logger.info(f"SSIM mean value: {mean_ssim}")
-    logger.info(f"Writing SSIM results to directory: {output_dir}")
-
-    success = write_ssim_results(output_dir, ssim_values, reference_video_path,
-                                 generated_video_path, num_inference_steps,
-                                 prompt)
-
-    if not success:
-        logger.error("Failed to write SSIM results to file")
-
-    min_acceptable_ssim = 0.93
-    assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}"

fastvideo/tests/ssim/test_lingbot_similarity.py

@@ -10,13 +10,21 @@ Note: num_inference_steps is reduced to 4 for faster CI.
 
 import os
 
-import pytest
-import torch
-
-from fastvideo import VideoGenerator
-from fastvideo.logger import init_logger
-from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embedding
-from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
+import pytest
+import torch
+
+from fastvideo import VideoGenerator
+from fastvideo.configs.sample.lingbotworld import LingBotWorld_SamplingParam
+from fastvideo.logger import init_logger
+from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embedding
+from fastvideo.tests.ssim.reference_utils import (
+    build_generated_output_dir,
+    build_reference_folder_path,
+    get_cuda_device_name,
+    resolve_device_reference_folder,
+    select_ssim_params,
+)
+from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
 
 logger = init_logger(__name__)
 
@@ -40,21 +48,20 @@ def _find_lingbotworld_examples_root() -> str | None:
     return None
 
 
-device_name = torch.cuda.get_device_name()
-device_reference_folder_suffix = "_reference_videos"
-
-if "A40" in device_name:
-    device_reference_folder = "A40" + device_reference_folder_suffix
-elif "L40S" in device_name:
-    device_reference_folder = "L40S" + device_reference_folder_suffix
-elif "H100" in device_name:
-    device_reference_folder = "H100" + device_reference_folder_suffix
-else:
-    device_reference_folder = None
-    logger.warning("Unsupported device for ssim tests: %s", device_name)
+device_name = get_cuda_device_name()
+device_reference_folder = resolve_device_reference_folder(
+    (
+        ("A40", "A40"),
+        ("L40S", "L40S"),
+        ("H100", "H100"),
+        ("H200", "H200"),
+    ),
+    device_name=device_name,
+    logger=logger,
+)
 
 
-LINGBOT_PARAMS = {
+LINGBOT_PARAMS = {
     "model_path": "FastVideo/LingBot-World-Base-Cam-Diffusers",
     "num_gpus": 2,
     "height": 256,
@@ -81,8 +88,28 @@ LINGBOT_PARAMS = {
         "镜头晃动，画面闪烁，模糊，噪点，水印，签名，文字，变形，扭曲，液化，不合逻辑的结构，卡顿，"
         "PPT幻灯片感，过暗，欠曝，低对比度，霓虹灯光感，过度锐化，3D渲染感，人物，行人，游客，身体，"
         "皮肤，肢体，面部特征，汽车，电线"
-    ),
-}
+    ),
+}
+_LINGBOT_FULL_QUALITY_DEFAULTS = LingBotWorld_SamplingParam()
+LINGBOT_FULL_QUALITY_PARAMS = {
+    "model_path": LINGBOT_PARAMS["model_path"],
+    "num_gpus": LINGBOT_PARAMS["num_gpus"],
+    "height": _LINGBOT_FULL_QUALITY_DEFAULTS.height,
+    "width": _LINGBOT_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": LINGBOT_PARAMS["num_frames"],  # default num_frames: 125
+    "num_inference_steps": _LINGBOT_FULL_QUALITY_DEFAULTS.num_inference_steps,
+    "guidance_scale": _LINGBOT_FULL_QUALITY_DEFAULTS.guidance_scale,
+    "guidance_scale_2": _LINGBOT_FULL_QUALITY_DEFAULTS.guidance_scale_2,
+    "embedded_cfg_scale": LINGBOT_PARAMS["embedded_cfg_scale"],
+    "flow_shift": LINGBOT_PARAMS["flow_shift"],
+    "boundary_ratio": _LINGBOT_FULL_QUALITY_DEFAULTS.boundary_ratio,
+    "seed": _LINGBOT_FULL_QUALITY_DEFAULTS.seed,
+    "fps": _LINGBOT_FULL_QUALITY_DEFAULTS.fps,
+    "spatial_scale": LINGBOT_PARAMS["spatial_scale"],
+    "example_case": LINGBOT_PARAMS["example_case"],
+    "image_path": LINGBOT_PARAMS["image_path"],
+    "negative_prompt": _LINGBOT_FULL_QUALITY_DEFAULTS.negative_prompt,
+}
 
 TEST_PROMPTS = [
     "The video presents a soaring journey through a fantasy jungle. The wind "
@@ -95,74 +122,80 @@ TEST_PROMPTS = [
 
 @pytest.mark.parametrize("prompt", TEST_PROMPTS)
 @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"])
-def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str):
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
-
-    if device_reference_folder is None:
-        pytest.skip(f"Unsupported device for LingBot SSIM test: {device_name}")
-    if torch.cuda.device_count() < LINGBOT_PARAMS["num_gpus"]:
-        pytest.skip(
-            f"LingBot SSIM test requires {LINGBOT_PARAMS['num_gpus']} GPUs, "
-            f"but only {torch.cuda.device_count()} detected."
-        )
+def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str):
+    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
+
+    params = select_ssim_params(LINGBOT_PARAMS, LINGBOT_FULL_QUALITY_PARAMS)
+
+    if device_reference_folder is None:
+        pytest.skip(f"Unsupported device for LingBot SSIM test: {device_name}")
+    if torch.cuda.device_count() < params["num_gpus"]:
+        pytest.skip(
+            f"LingBot SSIM test requires {params['num_gpus']} GPUs, "
+            f"but only {torch.cuda.device_count()} detected."
+        )
 
     examples_root = _find_lingbotworld_examples_root()
     if examples_root is None:
         pytest.skip(
             "lingbotworld_examples not found under examples/inference/basic.")
 
-    action_path = os.path.join(examples_root, LINGBOT_PARAMS["example_case"])
+    action_path = os.path.join(examples_root, params["example_case"])
     if not (os.path.exists(os.path.join(action_path, "poses.npy"))
             and os.path.exists(os.path.join(action_path, "intrinsics.npy"))):
         pytest.skip(f"Missing camera npy files under {action_path}")
 
-    c2ws_plucker_emb, aligned_num_frames = prepare_camera_embedding(
-        action_path=action_path,
-        num_frames=LINGBOT_PARAMS["num_frames"],
-        height=LINGBOT_PARAMS["height"],
-        width=LINGBOT_PARAMS["width"],
-        spatial_scale=LINGBOT_PARAMS["spatial_scale"],
-    )
+    c2ws_plucker_emb, aligned_num_frames = prepare_camera_embedding(
+        action_path=action_path,
+        num_frames=params["num_frames"],
+        height=params["height"],
+        width=params["width"],
+        spatial_scale=params["spatial_scale"],
+    )
 
     script_dir = os.path.dirname(os.path.abspath(__file__))
     model_id = "LingBot-World-Base-Cam-Diffusers"
-    output_dir = os.path.join(script_dir, "generated_videos", model_id,
-                              ATTENTION_BACKEND)
+    output_dir = build_generated_output_dir(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
+    )
     output_video_name = f"{prompt[:100].strip()}.mp4"
     os.makedirs(output_dir, exist_ok=True)
 
-    init_kwargs = {
-        "num_gpus": LINGBOT_PARAMS["num_gpus"],
-        "flow_shift": LINGBOT_PARAMS["flow_shift"],
-        "boundary_ratio": LINGBOT_PARAMS["boundary_ratio"],
-        "use_fsdp_inference": True,
-        "dit_cpu_offload": True,
+    init_kwargs = {
+        "num_gpus": params["num_gpus"],
+        "flow_shift": params["flow_shift"],
+        "boundary_ratio": params["boundary_ratio"],
+        "use_fsdp_inference": True,
+        "dit_cpu_offload": True,
         "dit_layerwise_offload": False,
         "text_encoder_cpu_offload": True,
         "vae_cpu_offload": False,
         "pin_cpu_memory": True,
     }
-    generation_kwargs = {
-        "output_path": output_dir,
-        "image_path": LINGBOT_PARAMS["image_path"],
-        "height": LINGBOT_PARAMS["height"],
-        "width": LINGBOT_PARAMS["width"],
-        "num_frames": aligned_num_frames,
-        "num_inference_steps": LINGBOT_PARAMS["num_inference_steps"],
-        "guidance_scale": LINGBOT_PARAMS["guidance_scale"],
-        "guidance_scale_2": LINGBOT_PARAMS["guidance_scale_2"],
-        "embedded_cfg_scale": LINGBOT_PARAMS["embedded_cfg_scale"],
-        "seed": LINGBOT_PARAMS["seed"],
-        "fps": LINGBOT_PARAMS["fps"],
-        "negative_prompt": LINGBOT_PARAMS["negative_prompt"],
-        "c2ws_plucker_emb": c2ws_plucker_emb,
-    }
+    generation_kwargs = {
+        "output_path": output_dir,
+        "image_path": params["image_path"],
+        "height": params["height"],
+        "width": params["width"],
+        "num_frames": aligned_num_frames,
+        "num_inference_steps": params["num_inference_steps"],
+        "guidance_scale": params["guidance_scale"],
+        "guidance_scale_2": params["guidance_scale_2"],
+        "embedded_cfg_scale": params["embedded_cfg_scale"],
+        "seed": params["seed"],
+        "fps": params["fps"],
+        "negative_prompt": params["negative_prompt"],
+        "c2ws_plucker_emb": c2ws_plucker_emb,
+    }
 
     generator: VideoGenerator | None = None
     try:
-        generator = VideoGenerator.from_pretrained(
-            model_path=LINGBOT_PARAMS["model_path"], **init_kwargs)
-        generator.generate_video(prompt, **generation_kwargs)
+        generator = VideoGenerator.from_pretrained(
+            model_path=params["model_path"], **init_kwargs)
+        generator.generate_video(prompt, **generation_kwargs)
     finally:
         if generator is not None:
             generator.shutdown()
@@ -171,8 +204,12 @@ def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str):
     assert os.path.exists(generated_video_path), (
         f"Output video was not generated at {generated_video_path}")
 
-    reference_folder = os.path.join(script_dir, device_reference_folder, model_id,
-                                    ATTENTION_BACKEND)
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
+    )
     if not os.path.exists(reference_folder):
         raise FileNotFoundError(
             f"Reference video folder does not exist: {reference_folder}")
@@ -196,9 +233,9 @@ def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str):
     mean_ssim = ssim_values[0]
     logger.info("SSIM mean value: %s", mean_ssim)
 
-    write_ssim_results(output_dir, ssim_values, reference_video_path,
-                       generated_video_path,
-                       LINGBOT_PARAMS["num_inference_steps"], prompt)
+    write_ssim_results(output_dir, ssim_values, reference_video_path,
+                       generated_video_path,
+                       params["num_inference_steps"], prompt)
 
     min_acceptable_ssim = 0.90
     assert mean_ssim >= min_acceptable_ssim, (

fastvideo/tests/ssim/test_longcat_similarity.py

@@ -20,7 +20,15 @@ import pytest
 import torch
 
 from fastvideo import VideoGenerator
+from fastvideo.configs.sample.base import SamplingParam
 from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.reference_utils import (
+    build_generated_output_dir,
+    build_reference_folder_path,
+    get_cuda_device_name,
+    resolve_device_reference_folder,
+    select_ssim_params,
+)
 from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
 
 logger = init_logger(__name__)
@@ -28,17 +36,16 @@ logger = init_logger(__name__)
 REQUIRED_GPUS = 1
 
 # Device-specific reference folder
-device_name = torch.cuda.get_device_name()
-device_reference_folder_suffix = "_reference_videos"
-
-if "A40" in device_name:
-    device_reference_folder = "A40" + device_reference_folder_suffix
-elif "L40S" in device_name:
-    device_reference_folder = "L40S" + device_reference_folder_suffix
-elif "H100" in device_name:
-    device_reference_folder = "H100" + device_reference_folder_suffix
-else:
-    logger.warning(f"Unsupported device for ssim tests: {device_name}")
+device_reference_folder = resolve_device_reference_folder(
+    (
+        ("A40", "A40"),
+        ("L40S", "L40S"),
+        ("H100", "H100"),
+        ("H200", "H200"),
+    ),
+    device_name=get_cuda_device_name(),
+    logger=logger,
+)
 
 # Common negative prompt from example scripts
 NEGATIVE_PROMPT = (
@@ -64,6 +71,21 @@ LONGCAT_T2V_PARAMS = {
     "seed": 42,
     "negative_prompt": NEGATIVE_PROMPT,
 }
+_LONGCAT_T2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(
+    LONGCAT_T2V_PARAMS["model_path"]
+)
+LONGCAT_T2V_FULL_QUALITY_PARAMS = {
+    "num_gpus": LONGCAT_T2V_PARAMS["num_gpus"],
+    "model_path": LONGCAT_T2V_PARAMS["model_path"],
+    "height": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.height,
+    "width": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": LONGCAT_T2V_PARAMS["num_frames"],  # default num_frames: 125
+    "num_inference_steps": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.num_inference_steps,
+    "guidance_scale": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.guidance_scale,
+    "fps": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.fps,
+    "seed": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.seed,
+    "negative_prompt": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.negative_prompt,
+}
 
 # =============================================================================
 # LongCat I2V Parameters (from basic_longcat_i2v.py)
@@ -80,6 +102,21 @@ LONGCAT_I2V_PARAMS = {
     "seed": 42,
     "negative_prompt": NEGATIVE_PROMPT,
 }
+_LONGCAT_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(
+    LONGCAT_I2V_PARAMS["model_path"]
+)
+LONGCAT_I2V_FULL_QUALITY_PARAMS = {
+    "num_gpus": LONGCAT_I2V_PARAMS["num_gpus"],
+    "model_path": LONGCAT_I2V_PARAMS["model_path"],
+    "height": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.height,
+    "width": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": LONGCAT_I2V_PARAMS["num_frames"],  # default num_frames: 125
+    "num_inference_steps": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.num_inference_steps,
+    "guidance_scale": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.guidance_scale,
+    "fps": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.fps,
+    "seed": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.seed,
+    "negative_prompt": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.negative_prompt,
+}
 
 # =============================================================================
 # LongCat VC Parameters (from basic_longcat_vc.py)
@@ -97,6 +134,22 @@ LONGCAT_VC_PARAMS = {
     "num_cond_frames": 13,
     "negative_prompt": NEGATIVE_PROMPT,
 }
+_LONGCAT_VC_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained(
+    LONGCAT_VC_PARAMS["model_path"]
+)
+LONGCAT_VC_FULL_QUALITY_PARAMS = {
+    "num_gpus": LONGCAT_VC_PARAMS["num_gpus"],
+    "model_path": LONGCAT_VC_PARAMS["model_path"],
+    "height": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.height,
+    "width": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": LONGCAT_VC_PARAMS["num_frames"],  # default num_frames: 125
+    "num_inference_steps": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.num_inference_steps,
+    "guidance_scale": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.guidance_scale,
+    "fps": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.fps,
+    "seed": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.seed,
+    "num_cond_frames": LONGCAT_VC_PARAMS["num_cond_frames"],
+    "negative_prompt": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.negative_prompt,
+}
 
 # Test prompts
 T2V_TEST_PROMPTS = [
@@ -153,15 +206,22 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str):
     """
     os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
 
+    params = select_ssim_params(LONGCAT_T2V_PARAMS, LONGCAT_T2V_FULL_QUALITY_PARAMS)
+
     script_dir = os.path.dirname(os.path.abspath(__file__))
     model_id = "LongCat-Video-T2V"
     
-    output_dir = os.path.join(script_dir, "generated_videos", model_id, ATTENTION_BACKEND)
+    output_dir = build_generated_output_dir(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
+    )
     output_video_name = f"{prompt[:100].strip()}.mp4"
     os.makedirs(output_dir, exist_ok=True)
 
     init_kwargs = {
-        "num_gpus": LONGCAT_T2V_PARAMS["num_gpus"],
+        "num_gpus": params["num_gpus"],
         "use_fsdp_inference": True,
         "dit_cpu_offload": True,
         "vae_cpu_offload": True,
@@ -171,18 +231,18 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str):
 
     generation_kwargs = {
         "output_path": output_dir,
-        "height": LONGCAT_T2V_PARAMS["height"],
-        "width": LONGCAT_T2V_PARAMS["width"],
-        "num_frames": LONGCAT_T2V_PARAMS["num_frames"],
-        "num_inference_steps": LONGCAT_T2V_PARAMS["num_inference_steps"],
-        "guidance_scale": LONGCAT_T2V_PARAMS["guidance_scale"],
-        "fps": LONGCAT_T2V_PARAMS["fps"],
-        "seed": LONGCAT_T2V_PARAMS["seed"],
-        "negative_prompt": LONGCAT_T2V_PARAMS["negative_prompt"],
+        "height": params["height"],
+        "width": params["width"],
+        "num_frames": params["num_frames"],
+        "num_inference_steps": params["num_inference_steps"],
+        "guidance_scale": params["guidance_scale"],
+        "fps": params["fps"],
+        "seed": params["seed"],
+        "negative_prompt": params["negative_prompt"],
     }
 
     generator = VideoGenerator.from_pretrained(
-        model_path=LONGCAT_T2V_PARAMS["model_path"], **init_kwargs
+        model_path=params["model_path"], **init_kwargs
     )
     generator.generate_video(prompt, **generation_kwargs)
     generator.shutdown()
@@ -193,8 +253,11 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str):
     )
 
     # Find reference video
-    reference_folder = os.path.join(
-        script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
     )
     if not os.path.exists(reference_folder):
         raise FileNotFoundError(
@@ -224,7 +287,7 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str):
 
     write_ssim_results(
         output_dir, ssim_values, reference_video_path, generated_video_path,
-        LONGCAT_T2V_PARAMS["num_inference_steps"], prompt
+        params["num_inference_steps"], prompt
     )
 
     min_acceptable_ssim = 0.90
@@ -244,10 +307,17 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str):
     """
     os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
 
+    params = select_ssim_params(LONGCAT_I2V_PARAMS, LONGCAT_I2V_FULL_QUALITY_PARAMS)
+
     script_dir = os.path.dirname(os.path.abspath(__file__))
     model_id = "LongCat-Video-I2V"
     
-    output_dir = os.path.join(script_dir, "generated_videos", model_id, ATTENTION_BACKEND)
+    output_dir = build_generated_output_dir(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
+    )
     output_video_name = f"{prompt[:100].strip()}.mp4"
     os.makedirs(output_dir, exist_ok=True)
 
@@ -256,7 +326,7 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str):
     image_path = _resolve_asset_path(I2V_IMAGE_PATHS[prompt_idx])
 
     init_kwargs = {
-        "num_gpus": LONGCAT_I2V_PARAMS["num_gpus"],
+        "num_gpus": params["num_gpus"],
         "use_fsdp_inference": True,
         "dit_cpu_offload": True,
         "vae_cpu_offload": True,
@@ -267,18 +337,18 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str):
     generation_kwargs = {
         "output_path": output_dir,
         "image_path": image_path,
-        "height": LONGCAT_I2V_PARAMS["height"],
-        "width": LONGCAT_I2V_PARAMS["width"],
-        "num_frames": LONGCAT_I2V_PARAMS["num_frames"],
-        "num_inference_steps": LONGCAT_I2V_PARAMS["num_inference_steps"],
-        "guidance_scale": LONGCAT_I2V_PARAMS["guidance_scale"],
-        "fps": LONGCAT_I2V_PARAMS["fps"],
-        "seed": LONGCAT_I2V_PARAMS["seed"],
-        "negative_prompt": LONGCAT_I2V_PARAMS["negative_prompt"],
+        "height": params["height"],
+        "width": params["width"],
+        "num_frames": params["num_frames"],
+        "num_inference_steps": params["num_inference_steps"],
+        "guidance_scale": params["guidance_scale"],
+        "fps": params["fps"],
+        "seed": params["seed"],
+        "negative_prompt": params["negative_prompt"],
     }
 
     generator = VideoGenerator.from_pretrained(
-        model_path=LONGCAT_I2V_PARAMS["model_path"], **init_kwargs
+        model_path=params["model_path"], **init_kwargs
     )
     generator.generate_video(prompt, **generation_kwargs)
     generator.shutdown()
@@ -289,8 +359,11 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str):
     )
 
     # Find reference video
-    reference_folder = os.path.join(
-        script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
     )
     if not os.path.exists(reference_folder):
         raise FileNotFoundError(
@@ -320,7 +393,7 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str):
 
     write_ssim_results(
         output_dir, ssim_values, reference_video_path, generated_video_path,
-        LONGCAT_I2V_PARAMS["num_inference_steps"], prompt
+        params["num_inference_steps"], prompt
     )
 
     min_acceptable_ssim = 0.90
@@ -340,10 +413,17 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str):
     """
     os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
 
+    params = select_ssim_params(LONGCAT_VC_PARAMS, LONGCAT_VC_FULL_QUALITY_PARAMS)
+
     script_dir = os.path.dirname(os.path.abspath(__file__))
     model_id = "LongCat-Video-VC"
     
-    output_dir = os.path.join(script_dir, "generated_videos", model_id, ATTENTION_BACKEND)
+    output_dir = build_generated_output_dir(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
+    )
     output_video_name = f"{prompt[:100].strip()}.mp4"
     os.makedirs(output_dir, exist_ok=True)
 
@@ -355,7 +435,7 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str):
         pytest.skip(f"Input video not found at {video_path}")
 
     init_kwargs = {
-        "num_gpus": LONGCAT_VC_PARAMS["num_gpus"],
+        "num_gpus": params["num_gpus"],
         "use_fsdp_inference": False,
         "dit_cpu_offload": False,
         "vae_cpu_offload": True,
@@ -367,19 +447,19 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str):
     generation_kwargs = {
         "output_path": output_dir,
         "video_path": video_path,
-        "num_cond_frames": LONGCAT_VC_PARAMS["num_cond_frames"],
-        "height": LONGCAT_VC_PARAMS["height"],
-        "width": LONGCAT_VC_PARAMS["width"],
-        "num_frames": LONGCAT_VC_PARAMS["num_frames"],
-        "num_inference_steps": LONGCAT_VC_PARAMS["num_inference_steps"],
-        "guidance_scale": LONGCAT_VC_PARAMS["guidance_scale"],
-        "fps": LONGCAT_VC_PARAMS["fps"],
-        "seed": LONGCAT_VC_PARAMS["seed"],
-        "negative_prompt": LONGCAT_VC_PARAMS["negative_prompt"],
+        "num_cond_frames": params["num_cond_frames"],
+        "height": params["height"],
+        "width": params["width"],
+        "num_frames": params["num_frames"],
+        "num_inference_steps": params["num_inference_steps"],
+        "guidance_scale": params["guidance_scale"],
+        "fps": params["fps"],
+        "seed": params["seed"],
+        "negative_prompt": params["negative_prompt"],
     }
 
     generator = VideoGenerator.from_pretrained(
-        model_path=LONGCAT_VC_PARAMS["model_path"], **init_kwargs
+        model_path=params["model_path"], **init_kwargs
     )
     generator.generate_video(prompt, **generation_kwargs)
     generator.shutdown()
@@ -390,8 +470,11 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str):
     )
 
     # Find reference video
-    reference_folder = os.path.join(
-        script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
     )
     if not os.path.exists(reference_folder):
         raise FileNotFoundError(
@@ -421,7 +504,7 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str):
 
     write_ssim_results(
         output_dir, ssim_values, reference_video_path, generated_video_path,
-        LONGCAT_VC_PARAMS["num_inference_steps"], prompt
+        params["num_inference_steps"], prompt
     )
 
     min_acceptable_ssim = 0.90

fastvideo/tests/ssim/test_matrixgame_similarity.py

@@ -5,8 +5,16 @@ import torch
 import pytest
 
 from fastvideo import VideoGenerator
-from fastvideo.models.dits.matrixgame.utils import create_action_presets
+from fastvideo.configs.sample.wan import MatrixGame2_SamplingParam
 from fastvideo.logger import init_logger
+from fastvideo.models.dits.matrixgame.utils import create_action_presets
+from fastvideo.tests.ssim.reference_utils import (
+    build_generated_output_dir,
+    build_reference_folder_path,
+    get_cuda_device_name,
+    resolve_device_reference_folder,
+    select_ssim_params,
+)
 from fastvideo.tests.utils import (
     compute_video_ssim_torchvision,
     write_ssim_results,
@@ -17,17 +25,15 @@ logger = init_logger(__name__)
 
 REQUIRED_GPUS = 1
 
-device_name = torch.cuda.get_device_name()
-device_reference_folder_suffix = "_reference_videos"
-
-if "A40" in device_name:
-    device_reference_folder = "A40" + device_reference_folder_suffix
-elif "L40S" in device_name:
-    device_reference_folder = "L40S" + device_reference_folder_suffix
-else:
-    # device_reference_folder = "L40S" + device_reference_folder_suffix
-    logger.warning(f"Unsupported device for ssim tests: {device_name}")
-    # raise ValueError(f"Unsupported device for ssim tests: {device_name}")
+device_reference_folder = resolve_device_reference_folder(
+    (
+        ("A40", "A40"),
+        ("L40S", "L40S"),
+        ("H200", "H200"),
+    ),
+    device_name=get_cuda_device_name(),
+    logger=logger,
+)
 
 # Base parameters from the shell script
 
@@ -41,11 +47,26 @@ MATRIXGAME_PARAMS = {
     "seed": 1024,
     "keyboard_dim": 4,
 }
+_MATRIXGAME_FULL_QUALITY_DEFAULTS = MatrixGame2_SamplingParam()
+MATRIXGAME_FULL_QUALITY_PARAMS = {
+    "num_gpus": MATRIXGAME_PARAMS["num_gpus"],
+    "model_path": MATRIXGAME_PARAMS["model_path"],
+    "height": _MATRIXGAME_FULL_QUALITY_DEFAULTS.height,
+    "width": _MATRIXGAME_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": MATRIXGAME_PARAMS["num_frames"],  # default num_frames: 57
+    "num_inference_steps": _MATRIXGAME_FULL_QUALITY_DEFAULTS.num_inference_steps,
+    "guidance_scale": _MATRIXGAME_FULL_QUALITY_DEFAULTS.guidance_scale,
+    "seed": _MATRIXGAME_FULL_QUALITY_DEFAULTS.seed,
+    "keyboard_dim": MATRIXGAME_PARAMS["keyboard_dim"],
+}
 
 
 MODEL_TO_PARAMS = {
     "Matrix-Game-2.0-Diffusers-Base": MATRIXGAME_PARAMS,
 }
+FULL_QUALITY_MODEL_TO_PARAMS = {
+    "Matrix-Game-2.0-Diffusers-Base": MATRIXGAME_FULL_QUALITY_PARAMS,
+}
 
 I2V_MODEL_TO_PARAMS = {}
 
@@ -70,19 +91,26 @@ def test_matrixgame_similarity(prompt, ATTENTION_BACKEND, model_id):
 
     script_dir = os.path.dirname(os.path.abspath(__file__))
 
-    base_output_dir = os.path.join(script_dir, "generated_videos", model_id)
-    output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
-    output_video_name = f"output.mp4"
+    output_dir = build_generated_output_dir(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
+    )
+    output_video_name = "output.mp4"
 
     os.makedirs(output_dir, exist_ok=True)
 
-    BASE_PARAMS = MODEL_TO_PARAMS[model_id]
+    params_map = select_ssim_params(
+        MODEL_TO_PARAMS,
+        FULL_QUALITY_MODEL_TO_PARAMS,
+    )
+    BASE_PARAMS = params_map[model_id]
     num_inference_steps = BASE_PARAMS["num_inference_steps"]
 
     # Create action conditions for MatrixGame
     actions = create_action_presets(
-        BASE_PARAMS["num_frames"], keyboard_dim=BASE_PARAMS["keyboard_dim"],
-        seed=BASE_PARAMS["seed"]
+        BASE_PARAMS["num_frames"], keyboard_dim=BASE_PARAMS["keyboard_dim"], seed=BASE_PARAMS["seed"]
     )
     latent_frames = (BASE_PARAMS["num_frames"] - 1) // 4 + 1
     grid_sizes = torch.tensor([latent_frames, 44, 80])
@@ -117,19 +145,18 @@ def test_matrixgame_similarity(prompt, ATTENTION_BACKEND, model_id):
     if isinstance(generator.executor, MultiprocExecutor):
         generator.executor.shutdown()
 
-    assert os.path.exists(output_dir), (
-        f"Output video was not generated at {output_dir}"
-    )
+    assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}"
 
-    reference_folder = os.path.join(
-        script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
     )
 
     if not os.path.exists(reference_folder):
         logger.error("Reference folder missing")
-        raise FileNotFoundError(
-            f"Reference video folder does not exist: {reference_folder}"
-        )
+        raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}")
 
     # Find the matching reference video based on the prompt
     reference_video_name = None
@@ -140,20 +167,14 @@ def test_matrixgame_similarity(prompt, ATTENTION_BACKEND, model_id):
             break
 
     if not reference_video_name:
-        logger.error(
-            f"Reference video not found for model: {model_id} with backend: {ATTENTION_BACKEND}"
-        )
+        logger.error(f"Reference video not found for model: {model_id} with backend: {ATTENTION_BACKEND}")
         raise FileNotFoundError("Reference video missing")
 
     reference_video_path = os.path.join(reference_folder, reference_video_name)
     generated_video_path = os.path.join(output_dir, output_video_name)
 
-    logger.info(
-        f"Computing SSIM between {reference_video_path} and {generated_video_path}"
-    )
-    ssim_values = compute_video_ssim_torchvision(
-        reference_video_path, generated_video_path, use_ms_ssim=True
-    )
+    logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}")
+    ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True)
 
     mean_ssim = ssim_values[0]
     logger.info(f"SSIM mean value: {mean_ssim}")

fastvideo/tests/ssim/test_sd35_similarity.py

@@ -1,41 +1,42 @@
 # SPDX-License-Identifier: Apache-2.0
-
 from __future__ import annotations
 
 import os
-import shlex
-import logging
 
 import pytest
 import torch
 
-logger = logging.getLogger(__name__)
+from fastvideo.configs.sample.sd35 import SD35SamplingParam
+from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    run_text_to_video_similarity_test,
+)
+from fastvideo.tests.ssim.reference_utils import (
+    get_cuda_device_name,
+    resolve_device_reference_folder,
+)
+
+logger = init_logger(__name__)
 
 REQUIRED_GPUS = 1
 
-def _device_reference_folder() -> str:
-    """Pick a reference folder name based on the current CUDA device."""
-    suffix = "_reference_videos"
-    device_name = torch.cuda.get_device_name(0)
-
-    if "A40" in device_name:
-        return "A40" + suffix
-    if "L40S" in device_name:
-        return "L40S" + suffix
-    if "H100" in device_name:
-        return "H100" + suffix
-    if "RTX 4090" in device_name or "4090" in device_name:
-        return "RTX4090" + suffix
-
-    logger.warning(
-        "Unsupported device for ssim tests: %s; using L40S references", device_name
-    )
-    return "L40S" + suffix
-
-
 SD35_MODEL_PATH = os.getenv(
     "SD35_MODEL_DIR",
-    "/FastVideo/official_weights/stabilityai__stable-diffusion-3.5-medium",
+    "stabilityai/stable-diffusion-3.5-medium",
+)
+
+device_reference_folder = resolve_device_reference_folder(
+    (
+        ("A40", "A40"),
+        ("L40S", "L40S"),
+        ("H100", "H100"),
+        ("H200", "H200"),
+        ("RTX 4090", "RTX4090"),
+        ("4090", "RTX4090"),
+    ),
+    device_name=get_cuda_device_name(),
+    fallback_device_prefix="L40S",
+    logger=logger,
 )
 
 MODEL_ID = "stabilityai__stable-diffusion-3.5-medium"
@@ -44,163 +45,80 @@ TEST_PROMPTS = [
     "a photo of a cat",
 ]
 
-pytestmark = pytest.mark.filterwarnings(
-    "ignore:.*torch.jit.script_method.*:DeprecationWarning",
+SD35_PARAMS = {
+    "num_gpus": 1,
+    "model_path": SD35_MODEL_PATH,
+    "sp_size": 1,
+    "tp_size": 1,
+    "height": 256,
+    "width": 256,
+    "num_frames": 1,
+    "fps": 1,
+    "num_inference_steps": 8,
+    "guidance_scale": 6.0,
+    "seed": 0,
+    "neg_prompt": "lowres, blurry",
+}
+
+_SD35_FULL_QUALITY_DEFAULTS = SD35SamplingParam()
+SD35_FULL_QUALITY_PARAMS = {
+    "num_gpus": 1,
+    "model_path": SD35_MODEL_PATH,
+    "sp_size": 1,
+    "tp_size": 1,
+    "height": _SD35_FULL_QUALITY_DEFAULTS.height,
+    "width": _SD35_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": SD35_PARAMS["num_frames"],
+    "fps": _SD35_FULL_QUALITY_DEFAULTS.fps,
+    "num_inference_steps": (_SD35_FULL_QUALITY_DEFAULTS.num_inference_steps),
+    "guidance_scale": _SD35_FULL_QUALITY_DEFAULTS.guidance_scale,
+    "seed": _SD35_FULL_QUALITY_DEFAULTS.seed,
+    "neg_prompt": _SD35_FULL_QUALITY_DEFAULTS.negative_prompt,
+}
+
+SD35_MODEL_TO_PARAMS = {
+    MODEL_ID: SD35_PARAMS,
+}
+SD35_FULL_QUALITY_MODEL_TO_PARAMS = {
+    MODEL_ID: SD35_FULL_QUALITY_PARAMS,
+}
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="SD3.5 SSIM test requires CUDA",
 )
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="SD3.5 SSIM test requires CUDA")
-@pytest.mark.parametrize("ATTENTION_BACKEND", ["TORCH_SDPA"])
 @pytest.mark.parametrize("prompt", TEST_PROMPTS)
-def test_sd35_similarity(prompt: str, ATTENTION_BACKEND: str) -> None:
-    from fastvideo import VideoGenerator
-    from fastvideo.tests.utils import (
-        compute_video_ssim_torchvision,
-        write_ssim_results,
-    )
+@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
+@pytest.mark.parametrize("model_id", list(SD35_MODEL_TO_PARAMS.keys()))
+def test_sd35_similarity(
+    prompt: str,
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    is_hf_repo = "/" in SD35_MODEL_PATH and not SD35_MODEL_PATH.startswith("/")
+    if not is_hf_repo and not os.path.isdir(SD35_MODEL_PATH):
+        pytest.skip(f"SD3.5 weights not found at {SD35_MODEL_PATH} (set SD35_MODEL_DIR to override)")
 
-    if not os.path.isdir(SD35_MODEL_PATH):
-        pytest.skip(
-            f"SD3.5 weights not found at {SD35_MODEL_PATH} (set SD35_MODEL_DIR to override)"
-        )
-
-    old_backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND")
-    old_transformers_verbosity = os.environ.get("TRANSFORMERS_VERBOSITY")
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
-    os.environ["TRANSFORMERS_VERBOSITY"] = "error"
-    try:
-        script_dir = os.path.dirname(os.path.abspath(__file__))
-        output_dir = os.path.join(
-            script_dir, "generated_videos", MODEL_ID, ATTENTION_BACKEND
-        )
-        os.makedirs(output_dir, exist_ok=True)
-
-        prompt_prefix = prompt[:100].strip()
-        output_video_name = f"{prompt_prefix}.mp4"
-        expected_video_path = os.path.join(output_dir, output_video_name)
-
-        for filename in os.listdir(output_dir):
-            if filename.endswith(".mp4") and filename.startswith(prompt_prefix):
-                try:
-                    os.remove(os.path.join(output_dir, filename))
-                except FileNotFoundError:
-                    pass
-
-        init_kwargs = {
-            "num_gpus": 1,
-            # Save a single-frame MP4 to reuse the existing video-SSIM harness.
+    run_text_to_video_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        attention_backend_name=attention_backend_name,
+        model_id=model_id,
+        default_params_map=SD35_MODEL_TO_PARAMS,
+        full_quality_params_map=SD35_FULL_QUALITY_MODEL_TO_PARAMS,
+        min_acceptable_ssim=0.98,
+        init_kwargs_override={
             "workload_type": "t2v",
-            "sp_size": 1,
-            "tp_size": 1,
-            "dit_cpu_offload": False,
-            "dit_layerwise_offload": False,
+            "use_fsdp_inference": False,
             "text_encoder_cpu_offload": False,
             "vae_cpu_offload": False,
             "image_encoder_cpu_offload": False,
             "pin_cpu_memory": False,
-            "use_fsdp_inference": False,
-        }
-
-        num_inference_steps = 8
-        generation_kwargs = {
-            "output_path": output_dir,
-            "height": 256,
-            "width": 256,
-            "num_frames": 1,
-            "fps": 1,
-            "num_inference_steps": num_inference_steps,
-            "guidance_scale": 6.0,
-            "seed": 0,
-            "negative_prompt": "lowres, blurry",
+        },
+        generation_kwargs_override={
             "save_video": True,
-        }
-
-        generator = VideoGenerator.from_pretrained(model_path=SD35_MODEL_PATH, **init_kwargs)
-        try:
-            generator.generate_video(prompt, **generation_kwargs)
-        finally:
-            generator.shutdown()
-
-        generated_video_path = None
-        if os.path.exists(expected_video_path):
-            generated_video_path = expected_video_path
-        else:
-            candidates = [
-                os.path.join(output_dir, f)
-                for f in os.listdir(output_dir)
-                if f.endswith(".mp4") and f.startswith(prompt_prefix)
-            ]
-            if candidates:
-                generated_video_path = max(candidates, key=os.path.getmtime)
-
-        assert generated_video_path is not None and os.path.exists(generated_video_path), (
-            f"Output video was not generated under {output_dir} for prompt '{prompt}'"
-        )
-
-        device_reference_folder = _device_reference_folder()
-        reference_folder = os.path.join(script_dir, device_reference_folder, MODEL_ID, ATTENTION_BACKEND)
-
-        if not os.path.exists(reference_folder):
-            bless_cmd = (
-                f"mkdir -p {shlex.quote(reference_folder)} && "
-                f"cp {shlex.quote(generated_video_path)} {shlex.quote(reference_folder)}/"
-            )
-            pytest.fail(
-                f"Reference folder does not exist: {reference_folder}\n"
-                f"Generated video saved at: {generated_video_path}\n"
-                "To bless references, copy the generated mp4 into the reference folder with a matching name.\n"
-                f"Example:\n  {bless_cmd}"
-            )
-
-        reference_video_path = os.path.join(reference_folder, output_video_name)
-        if not os.path.exists(reference_video_path):
-            reference_video_name = None
-            for filename in os.listdir(reference_folder):
-                if filename.endswith(".mp4") and filename.startswith(prompt_prefix):
-                    reference_video_name = filename
-                    break
-            if not reference_video_name:
-                bless_cmd = (
-                    f"cp {shlex.quote(generated_video_path)} {shlex.quote(reference_folder)}/"
-                )
-                pytest.fail(
-                    f"Reference video not found for prompt '{prompt}' under: {reference_folder}\n"
-                    f"Expected name: {output_video_name}\n"
-                    f"Generated video saved at: {generated_video_path}\n"
-                    "To bless references, copy the generated mp4 into the reference folder.\n"
-                    f"Example:\n  {bless_cmd}"
-                )
-            reference_video_path = os.path.join(reference_folder, reference_video_name)
-
-        logger.info("Computing SSIM between %s and %s", reference_video_path, generated_video_path)
-        ssim_values = compute_video_ssim_torchvision(
-            reference_video_path, generated_video_path, use_ms_ssim=True
-        )
-
-        mean_ssim = ssim_values[0]
-        logger.info("SSIM mean value: %s", mean_ssim)
-
-        write_ssim_results(
-            output_dir,
-            ssim_values,
-            reference_video_path,
-            generated_video_path,
-            num_inference_steps,
-            prompt,
-        )
-
-        min_acceptable_ssim = 0.98
-        assert mean_ssim >= min_acceptable_ssim, (
-            f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} "
-            f"for {MODEL_ID} with backend {ATTENTION_BACKEND}"
-        )
-
-    finally:
-        if old_backend is None:
-            os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None)
-        else:
-            os.environ["FASTVIDEO_ATTENTION_BACKEND"] = old_backend
-        if old_transformers_verbosity is None:
-            os.environ.pop("TRANSFORMERS_VERBOSITY", None)
-        else:
-            os.environ["TRANSFORMERS_VERBOSITY"] = old_transformers_verbosity
+        },
+    )

fastvideo/tests/ssim/test_turbodiffusion_similarity.py

@@ -1,184 +1,144 @@
 # SPDX-License-Identifier: Apache-2.0
-"""
-SSIM-based similarity test for TurboDiffusion inference.
-
-TurboDiffusion uses the SLA (Sparse-Linear Attention) backend and RCM scheduler
-for 1-4 step video generation.
-"""
 import os
 
-import torch
 import pytest
 
-from fastvideo import VideoGenerator
+from fastvideo.configs.sample.turbodiffusion import (
+    TurboDiffusionI2V_A14B_SamplingParam,
+    TurboDiffusionT2V_1_3B_SamplingParam,
+)
 from fastvideo.logger import init_logger
-from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
-from fastvideo.worker.multiproc_executor import MultiprocExecutor
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    DEVICE_MAPPINGS,
+    resolve_inference_device_reference_folder,
+    run_image_to_video_similarity_test,
+    run_text_to_video_similarity_test,
+)
+from fastvideo.tests.ssim.reference_utils import (
+    get_cuda_device_name,
+    resolve_device_reference_folder,
+)
 
 logger = init_logger(__name__)
 
 REQUIRED_GPUS = 4
 
-device_name = torch.cuda.get_device_name()
-device_reference_folder_suffix = '_reference_videos'
-
-if "A40" in device_name:
-    device_reference_folder = "A40" + device_reference_folder_suffix
-elif "L40S" in device_name:
-    device_reference_folder = "L40S" + device_reference_folder_suffix
-else:
-    # device_reference_folder = "L40S" + device_reference_folder_suffix
-    logger.warning(f"Unsupported device for ssim tests: {device_name}, using L40S references")
+device_name = get_cuda_device_name()
+device_reference_folder = resolve_device_reference_folder(
+    DEVICE_MAPPINGS,
+    device_name=device_name,
+)
+if device_reference_folder is None:
     raise ValueError(f"Unsupported device for ssim tests: {device_name}")
 
-
-# TurboDiffusion parameters (1-4 step generation with RCM scheduler + SLA attention)
 TURBODIFFUSION_PARAMS = {
     "num_gpus": 4,
     "model_path": "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers",
     "height": 480,
     "width": 832,
     "num_frames": 81,
-    "num_inference_steps": 4,  # TurboDiffusion uses 1-4 steps
-    "guidance_scale": 1.0,  # No CFG for TurboDiffusion
+    "num_inference_steps": 4,
+    "guidance_scale": 1.0,
     "seed": 42,
     "sp_size": 4,
     "tp_size": 1,
     "fps": 24,
 }
+_TURBODIFFUSION_FULL_QUALITY_DEFAULTS = TurboDiffusionT2V_1_3B_SamplingParam()
+TURBODIFFUSION_FULL_QUALITY_PARAMS = {
+    "num_gpus": TURBODIFFUSION_PARAMS["num_gpus"],
+    "model_path": TURBODIFFUSION_PARAMS["model_path"],
+    "height": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.height,
+    "width": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": TURBODIFFUSION_PARAMS["num_frames"],
+    "num_inference_steps": (_TURBODIFFUSION_FULL_QUALITY_DEFAULTS.num_inference_steps),
+    "guidance_scale": (_TURBODIFFUSION_FULL_QUALITY_DEFAULTS.guidance_scale),
+    "seed": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.seed,
+    "sp_size": TURBODIFFUSION_PARAMS["sp_size"],
+    "tp_size": TURBODIFFUSION_PARAMS["tp_size"],
+    "fps": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.fps,
+}
 
 TURBODIFFUSION_MODEL_TO_PARAMS = {
     "TurboWan2.1-T2V-1.3B-Diffusers": TURBODIFFUSION_PARAMS,
 }
+FULL_QUALITY_TURBODIFFUSION_MODEL_TO_PARAMS = {
+    "TurboWan2.1-T2V-1.3B-Diffusers": (TURBODIFFUSION_FULL_QUALITY_PARAMS),
+}
 
 TURBODIFFUSION_TEST_PROMPTS = [
-    "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
+    "Will Smith casually eats noodles, his relaxed demeanor contrasting "
+    "with the energetic background of a bustling street food market. "
+    "The scene captures a mix of humor and authenticity. Mid-shot "
+    "framing, vibrant lighting.",
 ]
 
 
 @pytest.mark.parametrize("prompt", TURBODIFFUSION_TEST_PROMPTS)
-@pytest.mark.parametrize("model_id", list(TURBODIFFUSION_MODEL_TO_PARAMS.keys()))
-def test_turbodiffusion_inference_similarity(prompt, model_id):
-    """
-    Test that runs TurboDiffusion inference with SLA attention and RCM scheduler,
-    then compares the output to reference videos using SSIM.
-    """
-    # TurboDiffusion requires SLA attention backend
-    ATTENTION_BACKEND = "SLA_ATTN"
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
-
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-
-    base_output_dir = os.path.join(script_dir, 'generated_videos', model_id)
-    output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
-    output_video_name = f"{prompt[:100].strip()}.mp4"
-
-    os.makedirs(output_dir, exist_ok=True)
-
-    BASE_PARAMS = TURBODIFFUSION_MODEL_TO_PARAMS[model_id]
-    num_inference_steps = BASE_PARAMS["num_inference_steps"]
-
-    init_kwargs = {
-        "num_gpus": BASE_PARAMS["num_gpus"],
-        "sp_size": BASE_PARAMS["sp_size"],
-        "tp_size": BASE_PARAMS["tp_size"],
-        "override_pipeline_cls_name": "TurboDiffusionPipeline",
-    }
-
-    generation_kwargs = {
-        "num_inference_steps": num_inference_steps,
-        "output_path": output_dir,
-        "height": BASE_PARAMS["height"],
-        "width": BASE_PARAMS["width"],
-        "num_frames": BASE_PARAMS["num_frames"],
-        "guidance_scale": BASE_PARAMS["guidance_scale"],
-        "seed": BASE_PARAMS["seed"],
-        "fps": BASE_PARAMS["fps"],
-    }
-
-    generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs)
-    generator.generate_video(prompt, **generation_kwargs)
-
-    if isinstance(generator.executor, MultiprocExecutor):
-        generator.executor.shutdown()
-
-    assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}"
-
-    reference_folder = os.path.join(
-        script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
-    )
-
-    if not os.path.exists(reference_folder):
-        logger.error("Reference folder missing")
-        raise FileNotFoundError(
-            f"Reference video folder does not exist: {reference_folder}"
-        )
-
-    # Find the matching reference video based on the prompt
-    reference_video_name = None
-
-    for filename in os.listdir(reference_folder):
-        if filename.endswith('.mp4') and prompt[:100].strip() in filename:
-            reference_video_name = filename
-            break
-
-    if not reference_video_name:
-        logger.error(
-            f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}"
-        )
-        raise FileNotFoundError(f"Reference video missing")
-
-    reference_video_path = os.path.join(reference_folder, reference_video_name)
-    generated_video_path = os.path.join(output_dir, output_video_name)
-
-    logger.info(
-        f"Computing SSIM between {reference_video_path} and {generated_video_path}"
-    )
-    ssim_values = compute_video_ssim_torchvision(
-        reference_video_path, generated_video_path, use_ms_ssim=True
-    )
-
-    mean_ssim = ssim_values[0]
-    logger.info(f"SSIM mean value: {mean_ssim}")
-    logger.info(f"Writing SSIM results to directory: {output_dir}")
-
-    success = write_ssim_results(
-        output_dir, ssim_values, reference_video_path,
-        generated_video_path, num_inference_steps, prompt
-    )
-
-    if not success:
-        logger.error("Failed to write SSIM results to file")
-
-    # TurboDiffusion uses fewer steps, may have slightly lower SSIM
-    min_acceptable_ssim = 0.95
-    assert mean_ssim >= min_acceptable_ssim, (
-        f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} "
-        f"for {model_id} with backend {ATTENTION_BACKEND}"
+@pytest.mark.parametrize(
+    "model_id",
+    list(TURBODIFFUSION_MODEL_TO_PARAMS.keys()),
+)
+def test_turbodiffusion_inference_similarity(
+    prompt: str,
+    model_id: str,
+) -> None:
+    run_text_to_video_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        attention_backend_name="SLA_ATTN",
+        model_id=model_id,
+        default_params_map=TURBODIFFUSION_MODEL_TO_PARAMS,
+        full_quality_params_map=(FULL_QUALITY_TURBODIFFUSION_MODEL_TO_PARAMS),
+        min_acceptable_ssim=0.95,
+        init_kwargs_override={
+            "override_pipeline_cls_name": "TurboDiffusionPipeline",
+        },
     )
 
 
-# TurboDiffusion I2V parameters (dual-model with RCM scheduler + SLA attention)
 TURBODIFFUSION_I2V_PARAMS = {
     "num_gpus": 4,
     "model_path": "loayrashid/TurboWan2.2-I2V-A14B-Diffusers",
     "height": 480,
     "width": 832,
     "num_frames": 45,
-    "num_inference_steps": 4,  # TurboDiffusion uses 1-4 steps
-    "guidance_scale": 1.0,  # No CFG for TurboDiffusion
+    "num_inference_steps": 4,
+    "guidance_scale": 1.0,
     "seed": 42,
     "sp_size": 4,
     "tp_size": 1,
     "fps": 24,
 }
+_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS = TurboDiffusionI2V_A14B_SamplingParam()
+TURBODIFFUSION_I2V_FULL_QUALITY_PARAMS = {
+    "num_gpus": TURBODIFFUSION_I2V_PARAMS["num_gpus"],
+    "model_path": TURBODIFFUSION_I2V_PARAMS["model_path"],
+    "height": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.height,
+    "width": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": TURBODIFFUSION_I2V_PARAMS["num_frames"],
+    "num_inference_steps": (_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.num_inference_steps),
+    "guidance_scale": (_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.guidance_scale),
+    "seed": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.seed,
+    "sp_size": TURBODIFFUSION_I2V_PARAMS["sp_size"],
+    "tp_size": TURBODIFFUSION_I2V_PARAMS["tp_size"],
+    "fps": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.fps,
+}
 
 TURBODIFFUSION_I2V_MODEL_TO_PARAMS = {
     "TurboWan2.2-I2V-A14B-Diffusers": TURBODIFFUSION_I2V_PARAMS,
 }
+FULL_QUALITY_TURBODIFFUSION_I2V_MODEL_TO_PARAMS = {
+    "TurboWan2.2-I2V-A14B-Diffusers": (TURBODIFFUSION_I2V_FULL_QUALITY_PARAMS),
+}
 
 TURBODIFFUSION_I2V_TEST_PROMPTS = [
-    "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.",
+    "An astronaut hatching from an egg, on the surface of the moon, "
+    "the darkness and depth of space realised in the background. "
+    "High quality, ultrarealistic detail and breath-taking "
+    "movie-like camera shot.",
 ]
 
 TURBODIFFUSION_I2V_IMAGE_PATHS = [
@@ -187,114 +147,27 @@ TURBODIFFUSION_I2V_IMAGE_PATHS = [
 
 
 @pytest.mark.parametrize("prompt", TURBODIFFUSION_I2V_TEST_PROMPTS)
-@pytest.mark.parametrize("model_id", list(TURBODIFFUSION_I2V_MODEL_TO_PARAMS.keys()))
-def test_turbodiffusion_i2v_inference_similarity(prompt, model_id):
-    """
-    Test that runs TurboDiffusion I2V inference with dual-model switching,
-    then compares the output to reference videos using SSIM.
-    """
-    # TurboDiffusion requires SLA attention backend
-    ATTENTION_BACKEND = "SLA_ATTN"
-    os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
-
-    assert len(TURBODIFFUSION_I2V_TEST_PROMPTS) == len(TURBODIFFUSION_I2V_IMAGE_PATHS), \
-        "Expect number of prompts equal to number of images"
-
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-
-    base_output_dir = os.path.join(script_dir, 'generated_videos', model_id)
-    output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
-    output_video_name = f"{prompt[:100].strip()}.mp4"
-
-    os.makedirs(output_dir, exist_ok=True)
-
-    BASE_PARAMS = TURBODIFFUSION_I2V_MODEL_TO_PARAMS[model_id]
-    num_inference_steps = BASE_PARAMS["num_inference_steps"]
+@pytest.mark.parametrize(
+    "model_id",
+    list(TURBODIFFUSION_I2V_MODEL_TO_PARAMS.keys()),
+)
+def test_turbodiffusion_i2v_inference_similarity(
+    prompt: str,
+    model_id: str,
+) -> None:
     image_path = TURBODIFFUSION_I2V_IMAGE_PATHS[TURBODIFFUSION_I2V_TEST_PROMPTS.index(prompt)]
-
-    init_kwargs = {
-        "num_gpus": BASE_PARAMS["num_gpus"],
-        "sp_size": BASE_PARAMS["sp_size"],
-        "tp_size": BASE_PARAMS["tp_size"],
-        "override_pipeline_cls_name": "TurboDiffusionI2VPipeline",
-        # Keep both transformers in VRAM - avoids CPU RAM bottleneck
-        "dit_cpu_offload": False,
-        "use_fsdp_inference": True,
-        "dit_layerwise_offload": False,
-    }
-
-    generation_kwargs = {
-        "num_inference_steps": num_inference_steps,
-        "output_path": output_dir,
-        "image_path": image_path,
-        "height": BASE_PARAMS["height"],
-        "width": BASE_PARAMS["width"],
-        "num_frames": BASE_PARAMS["num_frames"],
-        "guidance_scale": BASE_PARAMS["guidance_scale"],
-        "seed": BASE_PARAMS["seed"],
-        "fps": BASE_PARAMS["fps"],
-    }
-
-    generator = VideoGenerator.from_pretrained(
-        model_path=BASE_PARAMS["model_path"],
-        **init_kwargs
-    )
-    generator.generate_video(prompt, **generation_kwargs)
-
-    if isinstance(generator.executor, MultiprocExecutor):
-        generator.executor.shutdown()
-
-    assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}"
-
-    reference_folder = os.path.join(
-        script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
-    )
-
-    if not os.path.exists(reference_folder):
-        logger.error("Reference folder missing")
-        raise FileNotFoundError(
-            f"Reference video folder does not exist: {reference_folder}"
-        )
-
-    # Find the matching reference video based on the prompt
-    reference_video_name = None
-
-    for filename in os.listdir(reference_folder):
-        if filename.endswith('.mp4') and prompt[:100].strip() in filename:
-            reference_video_name = filename
-            break
-
-    if not reference_video_name:
-        logger.error(
-            f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}"
-        )
-        raise FileNotFoundError(f"Reference video missing")
-
-    reference_video_path = os.path.join(reference_folder, reference_video_name)
-    generated_video_path = os.path.join(output_dir, output_video_name)
-
-    logger.info(
-        f"Computing SSIM between {reference_video_path} and {generated_video_path}"
-    )
-    ssim_values = compute_video_ssim_torchvision(
-        reference_video_path, generated_video_path, use_ms_ssim=True
-    )
-
-    mean_ssim = ssim_values[0]
-    logger.info(f"SSIM mean value: {mean_ssim}")
-    logger.info(f"Writing SSIM results to directory: {output_dir}")
-
-    success = write_ssim_results(
-        output_dir, ssim_values, reference_video_path,
-        generated_video_path, num_inference_steps, prompt
-    )
-
-    if not success:
-        logger.error("Failed to write SSIM results to file")
-
-    # TurboDiffusion I2V uses fewer steps, may have slightly lower SSIM
-    min_acceptable_ssim = 0.95
-    assert mean_ssim >= min_acceptable_ssim, (
-        f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} "
-        f"for {model_id} with backend {ATTENTION_BACKEND}"
+    run_image_to_video_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        image_path=image_path,
+        attention_backend_name="SLA_ATTN",
+        model_id=model_id,
+        default_params_map=TURBODIFFUSION_I2V_MODEL_TO_PARAMS,
+        full_quality_params_map=(FULL_QUALITY_TURBODIFFUSION_I2V_MODEL_TO_PARAMS),
+        min_acceptable_ssim=0.95,
+        init_kwargs_override={
+            "override_pipeline_cls_name": ("TurboDiffusionI2VPipeline"),
+        },
     )

fastvideo/tests/ssim/test_wan_i2v_similarity.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+
+from fastvideo.configs.sample.wan import WanI2V_14B_480P_SamplingParam
+from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    resolve_inference_device_reference_folder,
+    run_image_to_video_similarity_test,
+)
+
+logger = init_logger(__name__)
+
+REQUIRED_GPUS = 2
+
+device_reference_folder = resolve_inference_device_reference_folder(logger)
+
+WAN_I2V_PARAMS = {
+    "num_gpus": 2,
+    "model_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
+    "height": 480,
+    "width": 832,
+    "num_frames": 45,
+    "num_inference_steps": 2,
+    "guidance_scale": 5.0,
+    "embedded_cfg_scale": 6,
+    "flow_shift": 7.0,
+    "seed": 1024,
+    "sp_size": 2,
+    "tp_size": 1,
+    "vae_sp": True,
+    "fps": 24,
+    "neg_prompt": (
+        "Bright tones, overexposed, static, blurred details, subtitles, style, "
+        "works, paintings, images, static, overall gray, worst quality, low "
+        "quality, JPEG compression residue, ugly, incomplete, extra fingers, "
+        "poorly drawn hands, poorly drawn faces, deformed, disfigured, "
+        "misshapen limbs, fused fingers, still picture, messy background, "
+        "three legs, many people in the background, walking backwards"
+    ),
+    "text-encoder-precision": ("fp32",),
+}
+_WAN_I2V_FULL_QUALITY_DEFAULTS = WanI2V_14B_480P_SamplingParam()
+WAN_I2V_FULL_QUALITY_PARAMS = {
+    "num_gpus": WAN_I2V_PARAMS["num_gpus"],
+    "model_path": WAN_I2V_PARAMS["model_path"],
+    "height": _WAN_I2V_FULL_QUALITY_DEFAULTS.height,
+    "width": _WAN_I2V_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": WAN_I2V_PARAMS["num_frames"],  # default num_frames: 81
+    "num_inference_steps": _WAN_I2V_FULL_QUALITY_DEFAULTS.num_inference_steps,
+    "guidance_scale": _WAN_I2V_FULL_QUALITY_DEFAULTS.guidance_scale,
+    "embedded_cfg_scale": WAN_I2V_PARAMS["embedded_cfg_scale"],
+    "flow_shift": WAN_I2V_PARAMS["flow_shift"],
+    "seed": _WAN_I2V_FULL_QUALITY_DEFAULTS.seed,
+    "sp_size": WAN_I2V_PARAMS["sp_size"],
+    "tp_size": WAN_I2V_PARAMS["tp_size"],
+    "vae_sp": WAN_I2V_PARAMS["vae_sp"],
+    "fps": _WAN_I2V_FULL_QUALITY_DEFAULTS.fps,
+    "neg_prompt": _WAN_I2V_FULL_QUALITY_DEFAULTS.negative_prompt,
+    "text-encoder-precision": WAN_I2V_PARAMS["text-encoder-precision"],
+}
+
+WAN_I2V_MODEL_TO_PARAMS = {
+    "Wan2.1-I2V-14B-480P-Diffusers": WAN_I2V_PARAMS,
+}
+FULL_QUALITY_WAN_I2V_MODEL_TO_PARAMS = {
+    "Wan2.1-I2V-14B-480P-Diffusers": WAN_I2V_FULL_QUALITY_PARAMS,
+}
+
+WAN_I2V_TEST_CASES = [
+    (
+        "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.",
+        "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg",
+    ),
+]
+
+
+@pytest.mark.parametrize(("prompt", "image_path"), WAN_I2V_TEST_CASES)
+@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"])
+@pytest.mark.parametrize("model_id", list(WAN_I2V_MODEL_TO_PARAMS.keys()))
+def test_wan_i2v_inference_similarity(
+    prompt: str,
+    image_path: str,
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    run_image_to_video_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        image_path=image_path,
+        attention_backend_name=attention_backend_name,
+        model_id=model_id,
+        default_params_map=WAN_I2V_MODEL_TO_PARAMS,
+        full_quality_params_map=FULL_QUALITY_WAN_I2V_MODEL_TO_PARAMS,
+        min_acceptable_ssim=0.97,
+    )

fastvideo/tests/ssim/test_wan_t2v_similarity.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+
+from fastvideo.configs.sample.wan import WanT2V_1_3B_SamplingParam
+from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    resolve_inference_device_reference_folder,
+    run_text_to_video_similarity_test,
+)
+
+logger = init_logger(__name__)
+
+REQUIRED_GPUS = 2
+
+device_reference_folder = resolve_inference_device_reference_folder(logger)
+
+WAN_T2V_PARAMS = {
+    "num_gpus": 2,
+    "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    "height": 480,
+    "width": 832,
+    "num_frames": 45,
+    "num_inference_steps": 4,
+    "guidance_scale": 3,
+    "embedded_cfg_scale": 6,
+    "flow_shift": 7.0,
+    "seed": 1024,
+    "sp_size": 2,
+    "tp_size": 1,
+    "vae_sp": True,
+    "fps": 24,
+    "neg_prompt": (
+        "Bright tones, overexposed, static, blurred details, subtitles, style, "
+        "works, paintings, images, static, overall gray, worst quality, low "
+        "quality, JPEG compression residue, ugly, incomplete, extra fingers, "
+        "poorly drawn hands, poorly drawn faces, deformed, disfigured, "
+        "misshapen limbs, fused fingers, still picture, messy background, "
+        "three legs, many people in the background, walking backwards"
+    ),
+    "text-encoder-precision": ("fp32",),
+}
+_WAN_T2V_FULL_QUALITY_DEFAULTS = WanT2V_1_3B_SamplingParam()
+WAN_T2V_FULL_QUALITY_PARAMS = {
+    "num_gpus": WAN_T2V_PARAMS["num_gpus"],
+    "model_path": WAN_T2V_PARAMS["model_path"],
+    "height": _WAN_T2V_FULL_QUALITY_DEFAULTS.height,
+    "width": _WAN_T2V_FULL_QUALITY_DEFAULTS.width,
+    "num_frames": WAN_T2V_PARAMS["num_frames"],  # default num_frames: 81
+    "num_inference_steps": _WAN_T2V_FULL_QUALITY_DEFAULTS.num_inference_steps,
+    "guidance_scale": _WAN_T2V_FULL_QUALITY_DEFAULTS.guidance_scale,
+    "embedded_cfg_scale": WAN_T2V_PARAMS["embedded_cfg_scale"],
+    "flow_shift": WAN_T2V_PARAMS["flow_shift"],
+    "seed": _WAN_T2V_FULL_QUALITY_DEFAULTS.seed,
+    "sp_size": WAN_T2V_PARAMS["sp_size"],
+    "tp_size": WAN_T2V_PARAMS["tp_size"],
+    "vae_sp": WAN_T2V_PARAMS["vae_sp"],
+    "fps": _WAN_T2V_FULL_QUALITY_DEFAULTS.fps,
+    "neg_prompt": _WAN_T2V_FULL_QUALITY_DEFAULTS.negative_prompt,
+    "text-encoder-precision": WAN_T2V_PARAMS["text-encoder-precision"],
+}
+
+WAN_T2V_MODEL_TO_PARAMS = {
+    "Wan2.1-T2V-1.3B-Diffusers": WAN_T2V_PARAMS,
+}
+FULL_QUALITY_WAN_T2V_MODEL_TO_PARAMS = {
+    "Wan2.1-T2V-1.3B-Diffusers": WAN_T2V_FULL_QUALITY_PARAMS,
+}
+
+WAN_T2V_TEST_PROMPTS = [
+    "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.",
+]
+
+
+@pytest.mark.parametrize("prompt", WAN_T2V_TEST_PROMPTS)
+@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN", "TORCH_SDPA"])
+@pytest.mark.parametrize("model_id", list(WAN_T2V_MODEL_TO_PARAMS.keys()))
+def test_wan_t2v_inference_similarity(
+    prompt: str,
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    run_text_to_video_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=prompt,
+        attention_backend_name=attention_backend_name,
+        model_id=model_id,
+        default_params_map=WAN_T2V_MODEL_TO_PARAMS,
+        full_quality_params_map=FULL_QUALITY_WAN_T2V_MODEL_TO_PARAMS,
+        min_acceptable_ssim=0.93,
+    )

fastvideo/tests/ssim/update_reference_videos.sh

@@ -1,63 +0,0 @@
-#!/bin/bash
-
-# Script to update reference videos using videos from generated_videos directory
-# Both directories should exist in the same directory as this script
-
-set -e  # Exit on any error
-
-# Define directory paths
-GENERATED_DIR="generated_videos"
-REFERENCE_DIR="L40S_reference_videos"
-
-# Colors for output
-RED='\033[0;31m'
-GREEN='\033[0;32m'
-YELLOW='\033[1;33m'
-NC='\033[0m' # No Color
-
-echo -e "${YELLOW}Starting reference video update...${NC}"
-
-# Check if generated_videos directory exists
-if [ ! -d "$GENERATED_DIR" ]; then
-    echo -e "${RED}Error: $GENERATED_DIR directory not found!${NC}"
-    exit 1
-fi
-
-# Check if reference_videos directory exists
-if [ ! -d "$REFERENCE_DIR" ]; then
-    echo -e "${RED}Error: $REFERENCE_DIR directory not found!${NC}"
-    exit 1
-fi
-
-# Function to copy videos recursively
-copy_videos() {
-    local src_dir="$1"
-    local dst_dir="$2"
-    
-    # Find all video files in the source directory
-    find "$src_dir" -type f \( -name "*.mp4" -o -name "*.avi" -o -name "*.mov" -o -name "*.mkv" -o -name "*.webm" -o -name "*.flv" \) | while read -r video_file; do
-        # Get relative path from source directory
-        relative_path="${video_file#$src_dir/}"
-        
-        # Construct destination path
-        dst_file="$dst_dir/$relative_path"
-        
-        # Create destination directory if it doesn't exist
-        dst_file_dir=$(dirname "$dst_file")
-        mkdir -p "$dst_file_dir"
-        
-        # Copy the video file
-        echo -e "${GREEN}Copying: $relative_path${NC}"
-        cp "$video_file" "$dst_file"
-    done
-}
-
-# Perform the copy operation
-echo -e "${YELLOW}Copying videos from $GENERATED_DIR to $REFERENCE_DIR...${NC}"
-copy_videos "$GENERATED_DIR" "$REFERENCE_DIR"
-
-echo -e "${GREEN}Reference videos updated successfully!${NC}"
-
-# Show summary
-video_count=$(find "$GENERATED_DIR" -type f \( -name "*.mp4" -o -name "*.avi" -o -name "*.mov" -o -name "*.mkv" -o -name "*.webm" -o -name "*.flv" \) | wc -l)
-echo -e "${YELLOW}Total videos processed: $video_count${NC}"

fastvideo/tests/transformers/test_hunyuanvideo.py

@@ -7,9 +7,7 @@ import pytest
 import torch
 
 from fastvideo.configs.pipelines.base import PipelineConfig
-from fastvideo.distributed.parallel_state import (
-    get_sp_parallel_rank,
-    get_sp_world_size)
+from fastvideo.distributed.parallel_state import get_sp_parallel_rank, get_sp_world_size
 from fastvideo.logger import init_logger
 from fastvideo.models.loader.component_loader import TransformerLoader
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -24,9 +22,7 @@ os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
 CONFIG_PATH = os.path.join(TRANSFORMER_PATH, "config.json")
 
@@ -41,9 +37,7 @@ REFERENCE_LATENT = 93.23263549804688
 @pytest.mark.usefixtures("distributed_setup")
 @pytest.mark.skip(reason="Skipping test_hunyuanvideo_distributed")
 def test_hunyuanvideo_distributed():
-    logger.info(
-        f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}"
-    )
+    logger.info(f"Initializing process: rank={RANK}, local_rank={LOCAL_RANK}, world_size={WORLD_SIZE}")
 
     torch.cuda.set_device(f"cuda:{LOCAL_RANK}")
 
@@ -51,9 +45,7 @@ def test_hunyuanvideo_distributed():
     sp_rank = get_sp_parallel_rank()
     sp_world_size = get_sp_world_size()
 
-    logger.info(
-        f"Process rank {RANK} initialized with SP rank {sp_rank} in SP world size {sp_world_size}"
-    )
+    logger.info(f"Process rank {RANK} initialized with SP rank {sp_rank} in SP world size {sp_world_size}")
 
     config = json.load(open(CONFIG_PATH))
     # remove   "_class_name": "HunyuanVideoTransformer3DModel",   "_diffusers_version": "0.32.0.dev0",
@@ -62,9 +54,11 @@ def test_hunyuanvideo_distributed():
     config.pop("_diffusers_version")
 
     precision_str = "bf16"
-    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         dit_cpu_offload=True,
-                         pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str))
+    args = FastVideoArgs(
+        model_path=TRANSFORMER_PATH,
+        dit_cpu_offload=True,
+        pipeline_config=PipelineConfig(dit_config=HunyuanVideoConfig(), dit_precision=precision_str),
+    )
     args.device = torch.device(f"cuda:{LOCAL_RANK}")
 
     loader = TransformerLoader()
@@ -78,23 +72,12 @@ def test_hunyuanvideo_distributed():
     device = torch.device(f"cuda:{LOCAL_RANK}")
 
     # Video latents [B, C, T, H, W]
-    hidden_states = torch.randn(batch_size,
-                                16,
-                                8,
-                                16,
-                                16,
-                                device=device,
-                                dtype=torch.bfloat16)
+    hidden_states = torch.randn(batch_size, 16, 8, 16, 16, device=device, dtype=torch.bfloat16)
     chunk_per_rank = hidden_states.shape[2] // sp_world_size
-    hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) *
-                                  chunk_per_rank]
+    hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank : (sp_rank + 1) * chunk_per_rank]
 
     # Text embeddings [B, L, D] (including global token)
-    encoder_hidden_states = torch.randn(batch_size,
-                                        seq_len + 1,
-                                        4096,
-                                        device=device,
-                                        dtype=torch.bfloat16)
+    encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 4096, device=device, dtype=torch.bfloat16)
 
     # Timestep
     timestep = torch.tensor([500], device=device, dtype=torch.bfloat16)
@@ -117,7 +100,5 @@ def test_hunyuanvideo_distributed():
 
     # Check if latents are similar
     diff_output_latents = abs(REFERENCE_LATENT - latent)
-    logger.info(
-        f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}"
-    )
+    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
     assert diff_output_latents < 1e-4, f"Output latents differ significantly: max diff = {diff_output_latents}"

fastvideo/tests/transformers/test_wanvideo.py

@@ -23,9 +23,7 @@ os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
 
 
@@ -34,22 +32,25 @@ def test_wan_transformer():
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     precision = torch.bfloat16
     precision_str = "bf16"
-    args = FastVideoArgs(model_path=TRANSFORMER_PATH,
-                         dit_cpu_offload=True,
-                         pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str))
+    args = FastVideoArgs(
+        model_path=TRANSFORMER_PATH,
+        dit_cpu_offload=True,
+        pipeline_config=PipelineConfig(dit_config=WanVideoConfig(), dit_precision=precision_str),
+    )
     args.device = device
 
     loader = TransformerLoader()
     model2 = loader.load(TRANSFORMER_PATH, args).to(dtype=precision)
 
-    model1 = WanTransformer3DModel.from_pretrained(
-        TRANSFORMER_PATH, device=device,
-        torch_dtype=precision).to(device, dtype=precision).requires_grad_(False)
+    model1 = (
+        WanTransformer3DModel.from_pretrained(TRANSFORMER_PATH, device=device, torch_dtype=precision)
+        .to(device, dtype=precision)
+        .requires_grad_(False)
+    )
 
     total_params = sum(p.numel() for p in model1.parameters())
     # Calculate weight sum for model1 (converting to float64 to avoid overflow)
-    weight_sum_model1 = sum(
-        p.to(torch.float64).sum().item() for p in model1.parameters())
+    weight_sum_model1 = sum(p.to(torch.float64).sum().item() for p in model1.parameters())
     # Also calculate mean for more stable comparison
     weight_mean_model1 = weight_sum_model1 / total_params
     logger.info("Model 1 weight sum: %s", weight_sum_model1)
@@ -57,8 +58,7 @@ def test_wan_transformer():
 
     # Calculate weight sum for model2 (converting to float64 to avoid overflow)
     total_params_model2 = sum(p.numel() for p in model2.parameters())
-    weight_sum_model2 = sum(
-        p.to(torch.float64).sum().item() for p in model2.parameters())
+    weight_sum_model2 = sum(p.to(torch.float64).sum().item() for p in model2.parameters())
     # Also calculate mean for more stable comparison
     weight_mean_model2 = weight_sum_model2 / total_params_model2
     logger.info("Model 2 weight sum: %s", weight_sum_model2)
@@ -78,20 +78,10 @@ def test_wan_transformer():
     seq_len = 30
 
     # Video latents [B, C, T, H, W]
-    hidden_states = torch.randn(batch_size,
-                                16,
-                                21,
-                                160,
-                                90,
-                                device=device,
-                                dtype=precision)
+    hidden_states = torch.randn(batch_size, 16, 21, 160, 90, device=device, dtype=precision)
 
     # Text embeddings [B, L, D] (including global token)
-    encoder_hidden_states = torch.randn(batch_size,
-                                        seq_len + 1,
-                                        4096,
-                                        device=device,
-                                        dtype=precision)
+    encoder_hidden_states = torch.randn(batch_size, seq_len + 1, 4096, device=device, dtype=precision)
 
     # Timestep
     timestep = torch.tensor([500], device=device, dtype=precision)
@@ -100,7 +90,7 @@ def test_wan_transformer():
         data_type="dummy",
     )
 
-    with torch.amp.autocast('cuda', dtype=precision):
+    with torch.amp.autocast("cuda", dtype=precision):
         output1 = model1(
             hidden_states=hidden_states,
             encoder_hidden_states=encoder_hidden_states,
@@ -108,13 +98,13 @@ def test_wan_transformer():
             return_dict=False,
         )[0]
         with set_forward_context(
-                current_timestep=0,
-                attn_metadata=None,
-                forward_batch=forward_batch,
+            current_timestep=0,
+            attn_metadata=None,
+            forward_batch=forward_batch,
         ):
-            output2 = model2(hidden_states=hidden_states,
-                             encoder_hidden_states=encoder_hidden_states,
-                             timestep=timestep)
+            output2 = model2(
+                hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, timestep=timestep
+            )
 
     # Check if outputs have the same shape
     assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}"

fastvideo/tests/utils.py

@@ -7,11 +7,42 @@ from fastvideo.logger import init_logger
 import numpy as np
 import torch
 from pytorch_msssim import ms_ssim, ssim
-from torchvision.io import read_video
 
 logger = init_logger(__name__)
 
 
+def _read_video_frames(path: str) -> torch.Tensor:
+    """Read video frames as a ``(T, C, H, W)`` uint8 tensor.
+
+    Tries ``torchvision.io.read_video`` first (available in older
+    torchvision) and falls back to PyAV when the function has been
+    removed (torchvision >= 0.26).
+    """
+    try:
+        from torchvision.io import read_video
+
+        frames, _, _ = read_video(
+            path,
+            pts_unit="sec",
+            output_format="TCHW",
+        )
+        return frames
+    except (ImportError, AttributeError):
+        pass
+
+    import av
+
+    container = av.open(path)
+    frames = []
+    for frame in container.decode(video=0):
+        arr = frame.to_ndarray(format="rgb24")
+        frames.append(torch.from_numpy(arr).permute(2, 0, 1))
+    container.close()
+    if not frames:
+        raise RuntimeError(f"No video frames decoded from {path}")
+    return torch.stack(frames)
+
+
 def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
     """
     Compute SSIM between two videos.
@@ -27,12 +58,8 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
     if not os.path.exists(video2_path):
         raise FileNotFoundError(f"Video2 not found: {video2_path}")
 
-    frames1, _, _ = read_video(video1_path,
-                               pts_unit='sec',
-                               output_format="TCHW")
-    frames2, _, _ = read_video(video2_path,
-                               pts_unit='sec',
-                               output_format="TCHW")
+    frames1 = _read_video_frames(video1_path)
+    frames2 = _read_video_frames(video2_path)
 
     # Ensure same number of frames
     min_frames = min(frames1.shape[0], frames2.shape[0])
@@ -50,8 +77,8 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
 
     # Process each frame individually
     for i in range(min_frames):
-        img1 = frames1[i:i + 1]
-        img2 = frames2[i:i + 1]
+        img1 = frames1[i : i + 1]
+        img2 = frames2[i : i + 1]
 
         with torch.no_grad():
             value = ms_ssim(img1, img2, data_range=1.0) if use_ms_ssim else ssim(img1, img2, data_range=1.0)
@@ -71,14 +98,14 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True):
 
         return mean_ssim, min_ssim, max_ssim
     else:
-        print('No SSIM values calculated')
+        print("No SSIM values calculated")
         return 0, 0, 0
 
 
 def compare_folders(reference_folder, generated_folder, use_ms_ssim=True):
     """
     Compare videos with the same filename between reference_folder and generated_folder
-    
+
     Example usage:
         results = compare_folders(reference_folder, generated_folder,
                               args.use_ms_ssim)
@@ -97,10 +124,8 @@ def compare_folders(reference_folder, generated_folder, use_ms_ssim=True):
         else:
             print("\nNo valid SSIM values to average")
     """
-    
-    reference_videos = [
-        f for f in os.listdir(reference_folder) if f.endswith('.mp4')
-    ]
+
+    reference_videos = [f for f in os.listdir(reference_folder) if f.endswith(".mp4")]
 
     results = {}
 
@@ -111,27 +136,23 @@ def compare_folders(reference_folder, generated_folder, use_ms_ssim=True):
         if os.path.exists(gen_path):
             print(f"\nComparing {video_name}...")
             try:
-                ssim_value = compute_video_ssim_torchvision(
-                    ref_path, gen_path, use_ms_ssim)
+                ssim_value = compute_video_ssim_torchvision(ref_path, gen_path, use_ms_ssim)
                 results[video_name] = ssim_value
             except Exception as e:
                 print(f"Error comparing {video_name}: {e}")
                 results[video_name] = None
         else:
-            print(
-                f"\nSkipping {video_name} - no matching file in generated folder"
-            )
+            print(f"\nSkipping {video_name} - no matching file in generated folder")
 
     return results
 
-def write_ssim_results(output_dir, ssim_values, reference_path, generated_path,
-                       num_inference_steps, prompt):
+
+def write_ssim_results(output_dir, ssim_values, reference_path, generated_path, num_inference_steps, prompt):
     """
     Write SSIM results to a JSON file in the same directory as the generated videos.
     """
     try:
-        logger.info(
-            f"Attempting to write SSIM results to directory: {output_dir}")
+        logger.info(f"Attempting to write SSIM results to directory: {output_dir}")
 
         if not os.path.exists(output_dir):
             os.makedirs(output_dir, exist_ok=True)
@@ -144,20 +165,17 @@ def write_ssim_results(output_dir, ssim_values, reference_path, generated_path,
             "max_ssim": max_ssim,
             "reference_video": reference_path,
             "generated_video": generated_path,
-            "parameters": {
-                "num_inference_steps": num_inference_steps,
-                "prompt": prompt
-            }
+            "parameters": {"num_inference_steps": num_inference_steps, "prompt": prompt},
         }
 
         test_name = f"steps{num_inference_steps}_{prompt[:100]}"
         result_file = os.path.join(output_dir, f"{test_name}_ssim.json")
         logger.info(f"Writing JSON results to: {result_file}")
-        with open(result_file, 'w') as f:
+        with open(result_file, "w") as f:
             json.dump(result, f, indent=2)
 
         logger.info(f"SSIM results written to {result_file}")
         return True
     except Exception as e:
         logger.error(f"ERROR writing SSIM results: {str(e)}")
-        return False
+        return False

fastvideo/tests/vaes/test_hunyuan_vae.py

@@ -9,6 +9,7 @@ from safetensors.torch import load_file
 
 from fastvideo.configs.pipelines import PipelineConfig
 from fastvideo.logger import init_logger
+
 # from fastvideo.models.vaes.hunyuanvae import (
 #     AutoencoderKLHunyuanVideo as MyHunyuanVAE)
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -22,9 +23,7 @@ os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 VAE_PATH = os.path.join(MODEL_PATH, "vae")
 CONFIG_PATH = os.path.join(VAE_PATH, "config.json")
 
@@ -42,30 +41,28 @@ def test_hunyuan_vae():
     device = torch.device("cuda:0")
     precision = torch.bfloat16
     precision_str = "bf16"
-    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=HunyuanVAEConfig(), vae_precision=precision_str))
+    args = FastVideoArgs(
+        model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=HunyuanVAEConfig(), vae_precision=precision_str)
+    )
     args.device = device
     args.vae_cpu_offload = False
 
     loader = VAELoader()
     model = loader.load(VAE_PATH, args)
 
-    model.enable_tiling(tile_sample_min_height=32,
-                         tile_sample_min_width=32,
-                         tile_sample_min_num_frames=8,
-                         tile_sample_stride_height=16,
-                         tile_sample_stride_width=16,
-                         tile_sample_stride_num_frames=4)
+    model.enable_tiling(
+        tile_sample_min_height=32,
+        tile_sample_min_width=32,
+        tile_sample_min_num_frames=8,
+        tile_sample_stride_height=16,
+        tile_sample_stride_width=16,
+        tile_sample_stride_num_frames=4,
+    )
 
     batch_size = 1
 
     # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size,
-                               3,
-                               21,
-                               64,
-                               64,
-                               device=device,
-                               dtype=torch.bfloat16)
+    input_tensor = torch.randn(batch_size, 3, 21, 64, 64, device=device, dtype=torch.bfloat16)
 
     # Disable gradients for inference
     with torch.no_grad():
@@ -81,8 +78,5 @@ def test_hunyuan_vae():
         raise ValueError(f"Unknown device: {device_name}")
 
     diff_encoded_latents = abs(REFERENCE_LATENT - latent)
-    logger.info(
-        f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}"
-    )
+    logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}")
     assert diff_encoded_latents < 1e-4, f"Encoded latents differ significantly: max diff = {diff_encoded_latents}"
-

fastvideo/tests/vaes/test_wan_vae.py

@@ -20,9 +20,7 @@ os.environ["MASTER_ADDR"] = "localhost"
 os.environ["MASTER_PORT"] = "29503"
 
 BASE_MODEL_PATH = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
-                                  local_dir=os.path.join("data", BASE_MODEL_PATH) # store in the large /workspace disk on Runpod
-                                  )
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH))
 VAE_PATH = os.path.join(MODEL_PATH, "vae")
 
 
@@ -31,28 +29,23 @@ def test_wan_vae():
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     precision = torch.float32
     precision_str = "fp32"
-    args = FastVideoArgs(model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str))
+    args = FastVideoArgs(
+        model_path=VAE_PATH, pipeline_config=PipelineConfig(vae_config=WanVAEConfig(), vae_precision=precision_str)
+    )
     args.device = device
     args.vae_cpu_offload = False
 
     loader = VAELoader()
     model2 = loader.load(VAE_PATH, args)
-    assert model2.use_feature_cache # Default to use the original WanVAE algorithm
+    assert model2.use_feature_cache  # Default to use the original WanVAE algorithm
 
-    model1 = AutoencoderKLWan.from_pretrained(
-        VAE_PATH, torch_dtype=precision).to(device).eval()
+    model1 = AutoencoderKLWan.from_pretrained(VAE_PATH, torch_dtype=precision).to(device).eval()
 
     # Create identical inputs for both models
     batch_size = 1
 
     # Video input [B, C, T, H, W]
-    input_tensor = torch.randn(batch_size,
-                               3,
-                               81,
-                               32,
-                               32,
-                               device=device,
-                               dtype=precision)
+    input_tensor = torch.randn(batch_size, 3, 81, 32, 32, device=device, dtype=precision)
     # latent_tensor = torch.randn(batch_size,
     #                             16,
     #                             21,
@@ -69,18 +62,22 @@ def test_wan_vae():
         print("--------------------------------")
         latent2 = model2.encode(input_tensor)
         # Check if latents have the same shape
-        assert latent1.mean.shape == latent2.mean.shape, f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
+        assert latent1.mean.shape == latent2.mean.shape, (
+            f"Latent shapes don't match: {latent1.mean.shape} vs {latent2.mean.shape}"
+        )
         # Check if latents are similar
         assert_close(latent1.mean, latent2.mean, atol=1e-4, rtol=1e-4)
         # Test decoding
         logger.info("Testing decoding...")
         latent1_tensor = latent1.mode()
-        mean1 = (torch.tensor(model1.config.latents_mean).view(
-            1, model1.config.z_dim, 1, 1, 1).to(input_tensor.device,
-                                                input_tensor.dtype))
-        std1 = (1.0 / torch.tensor(model1.config.latents_std).view(
-            1, model1.config.z_dim, 1, 1, 1)).to(input_tensor.device,
-                                                input_tensor.dtype)
+        mean1 = (
+            torch.tensor(model1.config.latents_mean)
+            .view(1, model1.config.z_dim, 1, 1, 1)
+            .to(input_tensor.device, input_tensor.dtype)
+        )
+        std1 = (1.0 / torch.tensor(model1.config.latents_std).view(1, model1.config.z_dim, 1, 1, 1)).to(
+            input_tensor.device, input_tensor.dtype
+        )
         latent1_tensor = latent1_tensor / std1 + mean1
         output1 = model1.decode(latent1_tensor).sample
 

fastvideo/worker/executor.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 from abc import ABC, abstractmethod
 from collections.abc import Callable
-from queue import Queue
 from typing import Any, TypeVar, cast
 
 from fastvideo.fastvideo_args import FastVideoArgs
@@ -15,14 +14,8 @@ _R = TypeVar("_R")
 
 class Executor(ABC):
 
-    def __init__(
-        self,
-        fastvideo_args: FastVideoArgs,
-        *,
-        log_queue=None,
-    ):
+    def __init__(self, fastvideo_args: FastVideoArgs):
         self.fastvideo_args = fastvideo_args
-        self._log_queue = log_queue
 
         self._init_executor()
 
@@ -104,16 +97,6 @@ class Executor(ABC):
         """
         raise NotImplementedError
 
-    @abstractmethod
-    def set_log_queue(self, log_queue: Queue | None) -> None:
-        """Forward worker logs to the given queue. Call before generate_video."""
-        self.collective_rpc("set_log_queue", kwargs={"log_queue": log_queue})
-
-    @abstractmethod
-    def clear_log_queue(self) -> None:
-        """Stop forwarding worker logs to the queue. Call after generate_video."""
-        self.collective_rpc("clear_log_queue")
-
     @abstractmethod
     def shutdown(self) -> None:
         """

fastvideo/worker/multiproc_executor.py

@@ -7,8 +7,6 @@ import contextlib
 from dataclasses import dataclass
 from enum import Enum
 import faulthandler
-import logging
-import logging.handlers
 import multiprocessing as mp
 from multiprocessing.connection import Connection
 from multiprocessing.queues import Queue
@@ -36,11 +34,6 @@ from fastvideo.worker.worker_base import WorkerWrapperBase
 logger = init_logger(__name__)
 
 
-def _make_queue_log_handler(log_queue: Queue) -> logging.Handler:
-    """Create a QueueHandler that forwards fastvideo logs to a multiprocessing queue."""
-    return logging.handlers.QueueHandler(log_queue)
-
-
 class StreamingTaskType(str, Enum):
     """
     Enumeration for different streaming task types.
@@ -104,7 +97,6 @@ class MultiprocExecutor(Executor):
                         distributed_init_method=distributed_init_method,
                         streaming_input_queue=self._streaming_input_queue,
                         streaming_output_queue=self._streaming_output_queue,
-                        log_queue=self._log_queue,
                     ))
 
             # Workers must be created before wait_for_ready to avoid
@@ -254,14 +246,6 @@ class MultiprocExecutor(Executor):
             if response["status"] != "lora_adapter_merged":
                 raise RuntimeError(f"Worker {i} failed to merge LoRA weights")
 
-    def set_log_queue(self, log_queue: Queue | None) -> None:
-        """Forward worker logs to the given queue. Call before generate_video."""
-        self.collective_rpc("set_log_queue", kwargs={"log_queue": log_queue})
-
-    def clear_log_queue(self) -> None:
-        """Stop forwarding worker logs to the queue. Call after generate_video."""
-        self.collective_rpc("clear_log_queue")
-
     def collective_rpc(self,
                        method: str | Callable,
                        timeout: float | None = None,
@@ -314,12 +298,6 @@ class MultiprocExecutor(Executor):
             return  # Prevent multiple shutdown calls
 
         logger.info("Shutting down MultiprocExecutor...")
-
-        # Check if workers were initialized (they might not be if initialization failed)
-        if not hasattr(self, 'workers') or not self.workers:
-            logger.info("No workers to shut down.")
-            return
-
         self.shutting_down = True
 
         # First try gentle termination
@@ -449,14 +427,11 @@ class WorkerMultiprocProc:
         pipe: Connection,
         streaming_input_queue: Queue | None = None,
         streaming_output_queue: Queue | None = None,
-        _initial_log_handler: logging.Handler | None = None,
-        **kwargs: Any,
     ):
         self.rank = rank
         self.pipe = pipe
         self.streaming_input_queue = streaming_input_queue
         self.streaming_output_queue = streaming_output_queue
-        self._initial_log_handler = _initial_log_handler
         wrapper = WorkerWrapperBase(fastvideo_args=fastvideo_args, rpc_rank=rank)
 
         all_kwargs: list[dict] = [{} for _ in range(fastvideo_args.num_gpus)]
@@ -483,7 +458,6 @@ class WorkerMultiprocProc:
         distributed_init_method: str,
         streaming_input_queue: Queue | None = None,
         streaming_output_queue: Queue | None = None,
-        log_queue: Queue | None = None,
     ) -> UnreadyWorkerProcHandle:
         context = get_mp_context()
         executor_pipe, worker_pipe = context.Pipe(duplex=True)
@@ -498,7 +472,6 @@ class WorkerMultiprocProc:
             "ready_pipe": writer,
             "streaming_input_queue": streaming_input_queue,
             "streaming_output_queue": streaming_output_queue,
-            "log_queue": log_queue,
         }
         # Run EngineCore busy loop in background process.
         proc = context.Process(target=WorkerMultiprocProc.worker_main,
@@ -515,13 +488,6 @@ class WorkerMultiprocProc:
         """ Worker initialization and execution loops.
         This runs a background process """
 
-        log_queue = kwargs.pop("log_queue", None)
-        # Add log handler before model loading so we capture fsdp_load, cuda, etc.
-        if log_queue is not None:
-            _handler = _make_queue_log_handler(log_queue)
-            logging.getLogger("fastvideo").addHandler(_handler)
-            kwargs["_initial_log_handler"] = _handler
-
         # Signal handler used for graceful termination.
         # SystemExit exception is only raised once to allow this and worker
         # processes to terminate without error
@@ -557,21 +523,9 @@ class WorkerMultiprocProc:
 
             worker.worker_busy_loop()
 
-        except Exception as exc:
+        except Exception:
             if ready_pipe is not None:
                 logger.exception("WorkerMultiprocProc failed to start.")
-                # Send error status to parent before closing pipe
-                try:
-                    traceback_str = get_exception_traceback()
-                    ready_pipe.send({
-                        "status": "ERROR",
-                        "error": str(exc),
-                        "traceback": traceback_str,
-                        "rank": rank,
-                    })
-                except Exception:
-                    # If sending fails, at least log it
-                    pass
             else:
                 logger.exception("WorkerMultiprocProc failed.")
 
@@ -581,8 +535,7 @@ class WorkerMultiprocProc:
             shutdown_requested = True
             traceback = get_exception_traceback()
             logger.error("Worker %d hit an exception: %s", rank, traceback)
-            if parent_process:
-                parent_process.send_signal(signal.SIGQUIT)
+            parent_process.send_signal(signal.SIGQUIT)
 
         finally:
             if ready_pipe is not None:
@@ -599,8 +552,7 @@ class WorkerMultiprocProc:
                       "See stack trace for root cause.")
 
         pipes = {handle.ready_pipe: handle for handle in unready_proc_handles}
-        ready_proc_handles: list[WorkerProcHandle | None] = [None] * len(unready_proc_handles)
-        worker_errors: list[str] = []
+        ready_proc_handles: list[WorkerProcHandle | None] = ([None] * len(unready_proc_handles))
         while pipes:
             ready = mp.connection.wait(pipes.keys())
             for pipe in ready:
@@ -609,40 +561,19 @@ class WorkerMultiprocProc:
                     # Wait until the WorkerProc is ready.
                     unready_proc_handle = pipes.pop(pipe)
                     response: dict[str, Any] = pipe.recv()
-                    if response["status"] == "ERROR":
-                        # Worker sent error details
-                        error_msg = response.get("error", "Unknown error")
-                        traceback_str = response.get("traceback", "")
-                        rank = response.get("rank", "unknown")
-                        error_info = f"Worker {rank} error: {error_msg}"
-                        if traceback_str:
-                            error_info += f"\n{traceback_str}"
-                        worker_errors.append(error_info)
-                        # Log a concise error message (full traceback will be in the exception)
-                        logger.error("Worker %s initialization failed: %s", rank, error_msg)
-                        # Continue to check other workers, but we'll fail at the end
-                    elif response["status"] != "READY":
-                        worker_errors.append(f"Worker returned unexpected status: {response.get('status', 'unknown')}")
+                    if response["status"] != "READY":
+                        raise e
 
-                    if response["status"] == "READY":
-                        ready_proc_handles[unready_proc_handle.rank] = (
-                            WorkerProcHandle.from_unready_handle(unready_proc_handle))
+                    ready_proc_handles[unready_proc_handle.rank] = (
+                        WorkerProcHandle.from_unready_handle(unready_proc_handle))
 
                 except EOFError:
-                    # Pipe closed without sending status - worker crashed
-                    worker_errors.append(f"Worker process crashed (pipe closed unexpectedly)")
                     e.__suppress_context__ = True
                     raise e from None
 
                 finally:
                     # Close connection.
                     pipe.close()
-        
-        # If any workers failed, raise exception with details
-        if worker_errors:
-            error_msg = "WorkerMultiprocProc initialization failed due to exceptions in background processes:\n"
-            error_msg += "\n".join(f"  - {err}" for err in worker_errors)
-            raise Exception(error_msg) from None
 
         logger.info("%d workers ready", len(ready_proc_handles))
         return cast(list[WorkerProcHandle], ready_proc_handles)
@@ -666,14 +597,6 @@ class WorkerMultiprocProc:
                         with contextlib.suppress(Exception):
                             self.pipe.send(response)
                         break
-                    if method == "set_log_queue":
-                        self._set_log_queue(kwargs.get("log_queue"))
-                        self.pipe.send({"status": "ok"})
-                        continue
-                    if method == "clear_log_queue":
-                        self._clear_log_queue()
-                        self.pipe.send({"status": "ok"})
-                        continue
                     if method == "start_streaming_queue_loop":
                         self.pipe.send({"status": "streaming_queue_loop_started"})
                         self.streaming_queue_loop()
@@ -742,30 +665,6 @@ class WorkerMultiprocProc:
                 logger.error("Worker %d queue loop error: %s", self.rank, e)
                 self.streaming_output_queue.put(StreamingResult(task_type=StreamingTaskType.STEP, error=e))
 
-    _log_queue_handler: logging.Handler | None = None
-
-    def _set_log_queue(self, log_queue: Queue | None) -> None:
-        """Add a handler that forwards fastvideo logs to the given queue."""
-        self._clear_log_queue()
-        if log_queue is None:
-            return
-        # Remove initial handler if present (from worker_main) to avoid duplicates
-        if self._initial_log_handler is not None:
-            logging.getLogger("fastvideo").removeHandler(
-                self._initial_log_handler
-            )
-            self._initial_log_handler = None
-        self._log_queue_handler = _make_queue_log_handler(log_queue)
-        logging.getLogger("fastvideo").addHandler(self._log_queue_handler)
-
-    def _clear_log_queue(self) -> None:
-        """Remove the log queue handler."""
-        if self._log_queue_handler is not None:
-            logging.getLogger("fastvideo").removeHandler(
-                self._log_queue_handler
-            )
-            self._log_queue_handler = None
-
     @staticmethod
     def setup_proc_title_and_log_prefix() -> None:
         dp_size = get_dp_group().world_size

mkdocs.yml

@@ -167,6 +167,8 @@ nav:
       - Index: contributing/developer_env/index.md
       - Docker: contributing/developer_env/docker.md
       - RunPod: contributing/developer_env/runpod.md
+    - Pull Requests: contributing/pull_requests.md
+    - CI Architecture: contributing/ci_architecture.md
     - Testing: contributing/testing.md
     - Profiling: contributing/profiling.md
     - Coding Agents: contributing/coding_agents.md

pyproject.toml

@@ -80,10 +80,6 @@ dependencies = [
     "torchcodec",
     "ray>=2.49.1",
     "ftfy==6.3.1",
-
-    # Job Runner
-    "fastapi==0.129.0",
-    "uvicorn==0.41.0"
 ]
 
 [tool.uv]
@@ -150,7 +146,7 @@ check_untyped_defs = true
 follow_imports = "silent"
 
 [tool.codespell]
-skip = "./data,./wandb,./ui/package-lock.json"
+skip ="./data,./wandb"
 
 [tool.ruff]
 # Allow lines to be as long as 120.

ui/.env.example

@@ -1 +0,0 @@
-PUBLIC_API_BASE_URL=http://localhost:8189/api

ui/.gitignore

@@ -1,42 +0,0 @@
-# See https://help.github.com/articles/ignoring-files/ for more about ignoring files.
-
-# dependencies
-/node_modules
-/.pnp
-.pnp.*
-.yarn/*
-!.yarn/patches
-!.yarn/plugins
-!.yarn/releases
-!.yarn/versions
-
-# testing
-/coverage
-
-# next.js
-/.next/
-/out/
-
-# production
-/build
-
-# misc
-.DS_Store
-*.pem
-
-# debug
-npm-debug.log*
-yarn-debug.log*
-yarn-error.log*
-.pnpm-debug.log*
-
-# env files (can opt-in for committing if needed)
-.env.local
-
-# vercel
-.vercel
-
-# typescript
-*.tsbuildinfo
-next-env.d.ts
-.svelte-kit