[feat]: overhaul SSIM test infrastructure — partition scheduling, helper migration, CI fixes (#1185)

Co-authored-by: Will Lin <wlsaidhi@gmail.com>
2026-03-29 23:59:25 +00:00
parent 2085a4fc4a
commit e294ca011c
54 changed files with 2701 additions and 1389 deletions
@@ -85,6 +85,8 @@ docs/distillation/examples/
 dmd_t2v_output/
 preprocess_output_text/

+
 .claude/
 .codex/
 openspec/
+fastvideo/tests/ssim/reference_videos/**
@@ -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

@@ -111,15 +121,46 @@ If you add a new test that requires:
 * **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.

-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

@@ -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.")
@@ -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()
@@ -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()
@@ -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()
@@ -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.
@@ -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."
-  }
-}
@@ -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
@@ -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."""
@@ -0,0 +1,335 @@
+# 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")
+        raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}")
+
+    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,
+    )
@@ -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
@@ -0,0 +1,668 @@
+# 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:
+    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.\n"
+            "Install 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."
+            )
+
+
+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):
+            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(
+                "No local reference videos found and no HF repo configured.\n"
+                f"Set {HF_REPO_ENV_KEY} or pass --repo-id."
+            )
+
+        print(
+            f"No local {quality_tier} reference videos found under {target_dir}. "
+            f"Downloading from HF repo {resolved_repo_id} ..."
+        )
+        download_reference_videos(
+            repo_id=resolved_repo_id,
+            repo_type=resolved_repo_type,
+            local_dir=target_dir,
+            quality_tiers=[quality_tier],
+        )
+
+        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."
+        )
+        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(
+                "No *_reference_videos directories found for the selected "
+                f"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)
@@ -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}"
@@ -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):
@@ -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}"
@@ -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, (
@@ -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
@@ -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}")
@@ -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
+        },
+    )
@@ -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"),
+        },
    )
@@ -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,
+    )
@@ -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,
+    )
@@ -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}"
@@ -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