diff --git a/.gitignore b/.gitignore index 5010b2cd3..5ed80a76b 100644 --- a/.gitignore +++ b/.gitignore @@ -85,6 +85,8 @@ docs/distillation/examples/ dmd_t2v_output/ preprocess_output_text/ + .claude/ .codex/ openspec/ +fastvideo/tests/ssim/reference_videos/** diff --git a/docs/contributing/testing.md b/docs/contributing/testing.md index b3debb7b6..5239b53c0 100644 --- a/docs/contributing/testing.md +++ b/docs/contributing/testing.md @@ -23,13 +23,18 @@ SSIM tests are located in `fastvideo/tests/ssim`. These tests generate videos us ``` fastvideo/tests/ssim/ -├── _reference_videos/ # Reference videos organized by GPU type (e.g., L40S_reference_videos) -│ ├── / -│ │ ├── / # e.g., FLASH_ATTN, TORCH_SDPA -│ │ │ └── +├── reference_videos/ +│ ├── default/ +│ │ └── _reference_videos/ +│ │ ├── / +│ │ │ ├── / # e.g., FLASH_ATTN, TORCH_SDPA +│ │ │ │ └── +│ └── full_quality/ +│ └── _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//_reference_videos`. * Inspect the generated video to ensure it meets quality expectations. - * Move the generated video to the appropriate reference folder: `fastvideo/tests/ssim/_reference_videos///`. - * 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//_reference_videos///`. + * 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//generated_videos` +* full-quality params: `ssim_generated_videos/full_quality//generated_videos` + +The printed `modal volume get` command also downloads into a quality-specific +local directory under `./generated_videos_modal/`. + +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 diff --git a/fastvideo/tests/modal/ssim_test.py b/fastvideo/tests/modal/ssim_test.py index 23e016873..9eca9c273 100644 --- a/fastvideo/tests/modal/ssim_test.py +++ b/fastvideo/tests/modal/ssim_test.py @@ -1,7 +1,11 @@ import ast +import datetime import glob import os import re +import shutil +import subprocess +import sys from dataclasses import dataclass from typing import Any @@ -10,7 +14,7 @@ import modal app = modal.App() model_vol = modal.Volume.from_name("hf-model-weights") -image_version = os.getenv("IMAGE_VERSION") +image_version = os.getenv("IMAGE_VERSION", "latest") image_tag = f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{image_version}" print(f"Using image: {image_tag}") @@ -24,30 +28,29 @@ image = ( "libssl-dev", "ffmpeg", ) - .run_commands( - "curl --proto '=https' --tlsv1.2 -sSf " - "https://sh.rustup.rs | sh -s -- -y --default-toolchain stable" - ) + .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable") .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc") - .env({ - "PATH": "/root/.cargo/bin:$PATH", - "BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""), - "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""), - "BUILDKITE_PULL_REQUEST": os.environ.get( - "BUILDKITE_PULL_REQUEST", "" - ), - "IMAGE_VERSION": os.environ.get("IMAGE_VERSION", ""), - }) + .env( + { + "PATH": "/root/.cargo/bin:$PATH", + "BUILDKITE_REPO": os.environ.get("BUILDKITE_REPO", ""), + "BUILDKITE_COMMIT": os.environ.get("BUILDKITE_COMMIT", ""), + "BUILDKITE_PULL_REQUEST": os.environ.get("BUILDKITE_PULL_REQUEST", ""), + "IMAGE_VERSION": os.environ.get("IMAGE_VERSION", ""), + } + ) ) SSIM_NUM_GPUS = 4 SSIM_TERMINATE_TIMEOUT_S = 30 +HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN") +RAW_GENERATED_VOLUME_ROOT = "ssim_generated_videos" +DEFAULT_OUTPUT_QUALITY_TIER = "default" +FULL_OUTPUT_QUALITY_TIER = "full_quality" +MODAL_DEVICE_REFERENCE_FOLDER = "L40S_reference_videos" SSIM_COMMON_KWARGS = dict( image=image, timeout=5400, - secrets=[modal.Secret.from_dict( - {"HF_API_KEY": os.environ.get("HF_API_KEY", "")} - )], volumes={"/root/data": model_vol}, ) @@ -105,6 +108,194 @@ class _PartitionResult: exit_code: int +def _split_csv_values(csv_values: str) -> set[str]: + return {value.strip() for value in csv_values.split(",") if value.strip()} + + +def _run_git_command(args: list[str]) -> str: + try: + result = subprocess.run( + ["git", *args], + check=True, + capture_output=True, + text=True, + ) + except subprocess.CalledProcessError as error: + stderr = (error.stderr or "").strip() + raise RuntimeError(f"Failed to run git {' '.join(args)}: {stderr or error}") from error + return result.stdout.strip() + + +def _normalize_git_repo_url(git_repo: str) -> str: + if git_repo.startswith("git@github.com:"): + return "https://github.com/" + git_repo[len("git@github.com:") :] + if git_repo.startswith("ssh://git@github.com/"): + return "https://github.com/" + git_repo[len("ssh://git@github.com/") :] + return git_repo + + +def _resolve_git_repo(git_repo: str) -> str: + if git_repo.strip(): + return _normalize_git_repo_url(git_repo.strip()) + + env_repo = os.environ.get("BUILDKITE_REPO", "").strip() + if env_repo: + return _normalize_git_repo_url(env_repo) + + discovered_repo = _run_git_command(["config", "--get", "remote.origin.url"]) + if discovered_repo: + return _normalize_git_repo_url(discovered_repo) + + raise RuntimeError("Could not resolve git repo URL. Pass --git-repo or set BUILDKITE_REPO.") + + +def _resolve_git_commit(git_commit: str) -> str: + if git_commit.strip(): + return git_commit.strip() + + env_commit = os.environ.get("BUILDKITE_COMMIT", "").strip() + if env_commit: + return env_commit + + discovered_commit = _run_git_command(["rev-parse", "HEAD"]) + if discovered_commit: + return discovered_commit + + raise RuntimeError("Could not resolve git commit. Pass --git-commit or set BUILDKITE_COMMIT.") + + +def _resolve_pull_request(pr_number: str) -> str: + if pr_number.strip(): + return pr_number.strip() + env_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "").strip() + if env_pr: + return env_pr + return "false" + + +def _resolve_hf_api_key( + hf_api_key: str, +) -> str: + if hf_api_key.strip(): + return hf_api_key.strip() + + for key in HF_TOKEN_ENV_KEYS: + value = os.environ.get(key, "").strip() + if value: + return value + + raise RuntimeError( + "Hugging Face token is required. Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN; or pass --hf-api-key." + ) + + +def _sanitize_path_fragment(value: str) -> str: + cleaned = re.sub(r"[^A-Za-z0-9_.-]+", "_", value.strip()) + return cleaned.strip("._-") + + +def _resolve_generated_volume_subdir( + requested_subdir: str, + git_commit: str, +) -> str: + if requested_subdir.strip(): + normalized = requested_subdir.strip().strip("/") + if not normalized: + raise RuntimeError("generated_volume_subdir must not be empty.") + return normalized + + short_commit = _sanitize_path_fragment(git_commit[:12]) or "unknown" + timestamp = datetime.datetime.utcnow().strftime("%Y%m%d_%H%M%S") + return f"{timestamp}_{short_commit}" + + +def _resolve_output_quality_tier(ssim_full_quality: bool) -> str: + if ssim_full_quality: + return FULL_OUTPUT_QUALITY_TIER + return DEFAULT_OUTPUT_QUALITY_TIER + + +def _build_generated_volume_relative_path( + *, + generated_volume_subdir: str, + quality_tier: str, +) -> str: + return os.path.join( + RAW_GENERATED_VOLUME_ROOT, + quality_tier, + generated_volume_subdir, + "generated_videos", + ) + + +def _build_local_generated_download_dir(quality_tier: str) -> str: + return os.path.join(".", "generated_videos_modal", quality_tier) + + +def _print_local_reference_copy_command(quality_tier: str) -> None: + generated_dir = os.path.join( + _build_local_generated_download_dir(quality_tier), + MODAL_DEVICE_REFERENCE_FOLDER, + ) + print( + "To update local references from the downloaded Modal outputs, run:\n" + " python fastvideo/tests/ssim/reference_videos_cli.py copy-local " + f"--quality-tier {quality_tier} " + f"--generated-dir {generated_dir} " + f"--device-folder {MODAL_DEVICE_REFERENCE_FOLDER}" + ) + + +def _count_video_files(root: str) -> int: + count = 0 + for current_root, _, filenames in os.walk(root): + for filename in filenames: + if filename.lower().endswith((".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv")): + count += 1 + return count + + +def _sync_generated_videos_to_volume( + repo_root: str, + generated_volume_subdir: str, + quality_tier: str, +) -> str | None: + generated_root = os.path.join( + repo_root, + "fastvideo", + "tests", + "ssim", + "generated_videos", + quality_tier, + ) + if not os.path.isdir(generated_root): + print( + f"No generated_videos directory found for quality tier {quality_tier}; skipping raw generated video export." + ) + return None + + relative_dst = _build_generated_volume_relative_path( + generated_volume_subdir=generated_volume_subdir, + quality_tier=quality_tier, + ) + absolute_dst = os.path.join("/root/data", relative_dst) + if os.path.exists(absolute_dst): + shutil.rmtree(absolute_dst) + os.makedirs(os.path.dirname(absolute_dst), exist_ok=True) + shutil.copytree(generated_root, absolute_dst) + model_vol.commit() + + num_videos = _count_video_files(absolute_dst) + print(f"Raw generated videos exported to Modal volume path: {relative_dst} ({num_videos} video files).") + print( + "Download command:\n" + f" modal volume get hf-model-weights {relative_dst} " + f"{_build_local_generated_download_dir(quality_tier)}" + ) + _print_local_reference_copy_command(quality_tier) + return relative_dst + + def _extract_required_gpus(filepath: str) -> int: """Read REQUIRED_GPUS from a test file. Defaults to 1.""" with open(filepath, encoding="utf-8") as file: @@ -126,10 +317,7 @@ def _extract_model_ids(filepath: str) -> list[str]: value_node = None if isinstance(node, ast.Assign): - target_names = [ - target.id for target in node.targets - if isinstance(target, ast.Name) - ] + target_names = [target.id for target in node.targets if isinstance(target, ast.Name)] value_node = node.value elif isinstance(node, ast.AnnAssign): if isinstance(node.target, ast.Name): @@ -138,8 +326,7 @@ def _extract_model_ids(filepath: str) -> list[str]: if not target_names or value_node is None: continue - if not any(name.endswith("MODEL_TO_PARAMS") - for name in target_names): + if not any(name.endswith("MODEL_TO_PARAMS") for name in target_names): continue if not isinstance(value_node, ast.Dict): continue @@ -159,35 +346,70 @@ def _extract_model_ids(filepath: str) -> list[str]: return unique_model_ids -def _discover_ssim_tasks(ssim_dir: str) -> list[SSIMTask]: +def _discover_ssim_tasks( + ssim_dir: str, + *, + selected_test_files: set[str] | None = None, + selected_model_ids: set[str] | None = None, +) -> list[SSIMTask]: tasks = [] task_id = 0 + selected_test_files = selected_test_files or set() + selected_test_file_names = {os.path.basename(test_file) for test_file in selected_test_files} + selected_model_ids = selected_model_ids or set() + test_files = sorted(glob.glob(os.path.join(ssim_dir, "test_*.py"))) + matched_test_files = set() + matched_model_ids = set() for filepath in test_files: + file_name = os.path.basename(filepath) + rel_path = f"./fastvideo/tests/ssim/{file_name}" + if selected_test_files and ( + filepath not in selected_test_files + and rel_path not in selected_test_files + and file_name not in selected_test_file_names + ): + continue + + matched_test_files.add(file_name) required_gpus = _extract_required_gpus(filepath) if required_gpus < 1 or required_gpus > SSIM_NUM_GPUS: - raise ValueError( - f"{filepath} requires {required_gpus} GPUs, " - f"but scheduler supports up to {SSIM_NUM_GPUS}." - ) - rel_path = f"./fastvideo/tests/ssim/{os.path.basename(filepath)}" + raise ValueError(f"{filepath} requires {required_gpus} GPUs, but scheduler supports up to {SSIM_NUM_GPUS}.") model_ids = _extract_model_ids(filepath) if model_ids: for model_id in model_ids: - tasks.append(SSIMTask( + if selected_model_ids and model_id not in selected_model_ids: + continue + matched_model_ids.add(model_id) + tasks.append( + SSIMTask( + task_id=task_id, + test_file=rel_path, + required_gpus=required_gpus, + model_id=model_id, + ) + ) + task_id += 1 + else: + if selected_model_ids: + continue + tasks.append( + SSIMTask( task_id=task_id, test_file=rel_path, required_gpus=required_gpus, - model_id=model_id, - )) - task_id += 1 - else: - tasks.append(SSIMTask( - task_id=task_id, - test_file=rel_path, - required_gpus=required_gpus, - )) + ) + ) task_id += 1 + + if selected_test_files: + unmatched_files = sorted(selected_test_file_names - matched_test_files) + if unmatched_files: + raise RuntimeError("Requested SSIM test file(s) not found: " + ", ".join(unmatched_files)) + if selected_model_ids: + unmatched_model_ids = sorted(selected_model_ids - matched_model_ids) + if unmatched_model_ids: + raise RuntimeError("Requested SSIM model_id(s) not found: " + ", ".join(unmatched_model_ids)) return sorted(tasks, key=lambda task: task.sort_key) @@ -207,26 +429,24 @@ def _build_checkout_command(git_commit: str, pr_number: str | None) -> str: try: pr_id = int(pr_number) except ValueError as error: - raise RuntimeError( - f"Invalid BUILDKITE_PULL_REQUEST value: {pr_number}" - ) from error - return ( - "git fetch --prune origin " - f"refs/pull/{pr_id}/head && " - "git checkout FETCH_HEAD" - ) + raise RuntimeError(f"Invalid BUILDKITE_PULL_REQUEST value: {pr_number}") from error + return f"git fetch --prune origin refs/pull/{pr_id}/head && git checkout FETCH_HEAD" return f"git checkout {shlex.quote(git_commit)}" -def _prepare_ssim_workspace() -> tuple[str, list[SSIMTask]]: +def _prepare_ssim_workspace( + *, + git_repo: str, + git_commit: str, + pr_number: str, + hf_api_key: str, + selected_test_files: set[str] | None = None, + selected_model_ids: set[str] | None = None, +) -> tuple[str, list[SSIMTask]]: import shlex - import subprocess - git_repo = os.environ.get("BUILDKITE_REPO") - git_commit = os.environ.get("BUILDKITE_COMMIT") - pr_number = os.environ.get("BUILDKITE_PULL_REQUEST") - if not git_repo or not git_commit: - raise RuntimeError("BUILDKITE_REPO and BUILDKITE_COMMIT must be set.") + if not hf_api_key.strip(): + raise RuntimeError("HF API key is required to prepare SSIM workspace.") checkout_command = _build_checkout_command(git_commit, pr_number) repo_root = "/FastVideo" @@ -235,7 +455,14 @@ def _prepare_ssim_workspace() -> tuple[str, list[SSIMTask]]: set -euo pipefail source $HOME/.local/bin/env source /opt/venv/bin/activate - cd {shlex.quote(repo_root)} + if [ -d {shlex.quote(repo_root)}/.git ]; then + cd {shlex.quote(repo_root)} + git remote set-url origin {shlex.quote(git_repo)} || true + git fetch --prune origin + else + git clone {shlex.quote(git_repo)} {shlex.quote(repo_root)} + cd {shlex.quote(repo_root)} + fi {checkout_command} git submodule update --init --recursive cd fastvideo-kernel @@ -249,17 +476,22 @@ def _prepare_ssim_workspace() -> tuple[str, list[SSIMTask]]: ["/bin/bash", "-lc", command], capture_output=True, text=True, + env={ + **os.environ, + "HF_API_KEY": hf_api_key, + }, ) if result.returncode != 0: print(result.stdout) print(result.stderr) - raise RuntimeError( - "Workspace setup failed with " - f"exit code {result.returncode}" - ) + raise RuntimeError(f"Workspace setup failed with exit code {result.returncode}") ssim_dir = os.path.join(repo_root, "fastvideo", "tests", "ssim") - tasks = _discover_ssim_tasks(ssim_dir) + tasks = _discover_ssim_tasks( + ssim_dir, + selected_test_files=selected_test_files, + selected_model_ids=selected_model_ids, + ) if not tasks: raise RuntimeError("No SSIM test files found.") return repo_root, tasks @@ -271,22 +503,14 @@ def _spawn_ssim_task( assigned_gpu_ids: list[str], log_dir: str, task_index: int, + pytest_extra_args: list[str], ) -> _RunningTask: import shlex - import subprocess - safe_test_name = re.sub( - r"[^A-Za-z0-9_.-]+", "_", task.test_name - ) - log_path = os.path.join( - log_dir, f"{task_index:03d}_{safe_test_name}.log" - ) - command = ( - "set -euo pipefail && " - "source $HOME/.local/bin/env && " - "source /opt/venv/bin/activate && " - f"pytest {shlex.quote(task.test_file)} -vs" - ) + safe_test_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", task.test_name) + log_path = os.path.join(log_dir, f"{task_index:03d}_{safe_test_name}.log") + pytest_command = shlex.join(["pytest", task.test_file, "-vs", *pytest_extra_args]) + command = f"set -euo pipefail && source $HOME/.local/bin/env && source /opt/venv/bin/activate && {pytest_command}" env = os.environ.copy() env["HF_HOME"] = "/root/data/.cache" # MultiprocExecutor returns CUDA tensors through mp pipes (CUDA IPC). @@ -327,9 +551,7 @@ def _finalize_running_task( ) -> None: running_task.log_handle.close() available_gpu_ids.extend(running_task.gpu_ids) - available_gpu_ids.sort( - key=lambda gpu_id: gpu_order.get(gpu_id, len(gpu_order)) - ) + available_gpu_ids.sort(key=lambda gpu_id: gpu_order.get(gpu_id, len(gpu_order))) results[running_task.task.task_id] = _TaskResult( task=running_task.task, status=status, @@ -402,9 +624,30 @@ def _get_visible_gpu_ids() -> list[str]: return gpu_ids +def _build_pytest_extra_args( + *, + ssim_full_quality: bool, + ssim_reference_repo: str, + skip_ssim_reference_download: bool, + pytest_k: str, +) -> list[str]: + args = [] + if ssim_full_quality: + args.append("--ssim-full-quality") + if ssim_reference_repo.strip(): + args.extend(["--ssim-reference-repo", ssim_reference_repo.strip()]) + if skip_ssim_reference_download: + args.append("--skip-ssim-reference-download") + if pytest_k.strip(): + args.extend(["-k", pytest_k.strip()]) + return args + + def _schedule_ssim_tasks( repo_root: str, tasks: list[SSIMTask], + pytest_extra_args: list[str], + fail_fast: bool = True, ) -> dict[int, _TaskResult]: import tempfile import time @@ -412,16 +655,10 @@ def _schedule_ssim_tasks( pending_tasks = list(tasks) running_tasks = [] available_gpu_ids = _get_visible_gpu_ids() - gpu_order = { - gpu_id: index for index, gpu_id in enumerate(available_gpu_ids) - } - max_required_gpus = max( - (task.required_gpus for task in pending_tasks), default=0 - ) + gpu_order = {gpu_id: index for index, gpu_id in enumerate(available_gpu_ids)} + max_required_gpus = max((task.required_gpus for task in pending_tasks), default=0) if max_required_gpus > len(available_gpu_ids): - cuda_visible_devices = os.environ.get( - "CUDA_VISIBLE_DEVICES", "" - ) + cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "") 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.") diff --git a/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py b/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py index 64c6cd59c..1adf817f3 100644 --- a/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py +++ b/fastvideo/tests/nightly/test_e2e_dmd_t2v_crush_smol.py @@ -3,8 +3,8 @@ from pathlib import Path from huggingface_hub import snapshot_download import subprocess import sys -from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision import shutil +from fastvideo.tests.utils import compute_video_ssim_torchvision # Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) @@ -181,4 +181,4 @@ def test_e2e_overfit_single_sample(): if __name__ == "__main__": - test_e2e_overfit_single_sample() \ No newline at end of file + test_e2e_overfit_single_sample() diff --git a/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py index 3caa67e5d..88c409f40 100644 --- a/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py +++ b/fastvideo/tests/nightly/test_e2e_i2v_overfit_single_sample.py @@ -4,7 +4,7 @@ from huggingface_hub import snapshot_download import shutil import subprocess import sys -from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision +from fastvideo.tests.utils import compute_video_ssim_torchvision # Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) @@ -170,4 +170,4 @@ def test_e2e_overfit_single_sample(): if __name__ == "__main__": - test_e2e_overfit_single_sample() \ No newline at end of file + test_e2e_overfit_single_sample() diff --git a/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py b/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py index 2b646ac30..93f44eeb3 100644 --- a/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py +++ b/fastvideo/tests/nightly/test_e2e_overfit_single_sample.py @@ -4,7 +4,7 @@ from huggingface_hub import snapshot_download import shutil import subprocess import sys -from fastvideo.tests.ssim.test_inference_similarity import compute_video_ssim_torchvision +from fastvideo.tests.utils import compute_video_ssim_torchvision # Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) @@ -175,4 +175,4 @@ def test_e2e_overfit_single_sample(): if __name__ == "__main__": - test_e2e_overfit_single_sample() \ No newline at end of file + test_e2e_overfit_single_sample() diff --git a/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 b/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 deleted file mode 100644 index 10cecbdb3..000000000 Binary files a/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 deleted file mode 100644 index df7cda500..000000000 Binary files a/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 b/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 deleted file mode 100644 index db8b1f61e..000000000 Binary files a/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 deleted file mode 100644 index cb3c33f88..000000000 Binary files a/fastvideo/tests/ssim/A40_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 b/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 deleted file mode 100644 index c4ea3c07e..000000000 Binary files a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/TORCH_SDPA/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 b/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/TORCH_SDPA/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 deleted file mode 100644 index 36ae18a79..000000000 Binary files a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/TORCH_SDPA/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 b/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 deleted file mode 100644 index c5f0bb925..000000000 Binary files a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 deleted file mode 100644 index b6decf1be..000000000 Binary files a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 b/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 deleted file mode 100644 index cd9be1dd2..000000000 Binary files a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscap.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 b/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 deleted file mode 100644 index 107ade650..000000000 Binary files a/fastvideo/tests/ssim/A40_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of .mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/AGENTS.md b/fastvideo/tests/ssim/AGENTS.md new file mode 100644 index 000000000..71817090f --- /dev/null +++ b/fastvideo/tests/ssim/AGENTS.md @@ -0,0 +1,55 @@ +# SSIM Directory Guidelines + +## Scope +These instructions apply to everything under `fastvideo/tests/ssim/`. + +## Purpose +SSIM tests are GPU-backed end-to-end quality regression checks. They generate +videos and compare them against device-specific reference videos. + +## Key Files +- `test_*.py`: SSIM regression tests by model/task. +- `conftest.py`: optional filtering via `FASTVIDEO_SSIM_MODEL_ID`. +- `reference_videos//_reference_videos/`: local cache of + references split by quality tier and GPU type. +- `generated_videos//_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= 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__similarity.py`. +- Set `REQUIRED_GPUS = ` 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. diff --git a/fastvideo/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 deleted file mode 100644 index cb83ba11a..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 deleted file mode 100644 index 63c9d13e6..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/FastHunyuan-diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/LingBot-World-Base-Cam-Diffusers/FLASH_ATTN/The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/LingBot-World-Base-Cam-Diffusers/FLASH_ATTN/The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue.mp4 deleted file mode 100644 index 9c0f925b8..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/LingBot-World-Base-Cam-Diffusers/FLASH_ATTN/The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-I2V/FLASH_ATTN/A woman sits at a wooden table by the window in a cozy café. She reaches out with her right hand, pi.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-I2V/FLASH_ATTN/A woman sits at a wooden table by the window in a cozy café. She reaches out with her right hand, pi.mp4 deleted file mode 100644 index f17abf832..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-I2V/FLASH_ATTN/A woman sits at a wooden table by the window in a cozy café. She reaches out with her right hand, pi.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-I2V/TORCH_SDPA/A woman sits at a wooden table by the window in a cozy café. She reaches out with her right hand, pi.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-I2V/TORCH_SDPA/A woman sits at a wooden table by the window in a cozy café. She reaches out with her right hand, pi.mp4 deleted file mode 100644 index 74bb31dfb..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-I2V/TORCH_SDPA/A woman sits at a wooden table by the window in a cozy café. She reaches out with her right hand, pi.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-T2V/FLASH_ATTN/In a realistic photography style, a white boy around seven or eight years old sits on a park bench,.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-T2V/FLASH_ATTN/In a realistic photography style, a white boy around seven or eight years old sits on a park bench,.mp4 deleted file mode 100644 index e6a2ec0df..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-T2V/FLASH_ATTN/In a realistic photography style, a white boy around seven or eight years old sits on a park bench,.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-T2V/TORCH_SDPA/In a realistic photography style, a white boy around seven or eight years old sits on a park bench,.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-T2V/TORCH_SDPA/In a realistic photography style, a white boy around seven or eight years old sits on a park bench,.mp4 deleted file mode 100644 index d4403e860..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-T2V/TORCH_SDPA/In a realistic photography style, a white boy around seven or eight years old sits on a park bench,.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-VC/FLASH_ATTN/A person rides a motorcycle along a long, straight road that stretches between a body of water and a.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-VC/FLASH_ATTN/A person rides a motorcycle along a long, straight road that stretches between a body of water and a.mp4 deleted file mode 100644 index eac144587..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-VC/FLASH_ATTN/A person rides a motorcycle along a long, straight road that stretches between a body of water and a.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-VC/TORCH_SDPA/A person rides a motorcycle along a long, straight road that stretches between a body of water and a.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-VC/TORCH_SDPA/A person rides a motorcycle along a long, straight road that stretches between a body of water and a.mp4 deleted file mode 100644 index 9386e7050..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/LongCat-Video-VC/TORCH_SDPA/A person rides a motorcycle along a long, straight road that stretches between a body of water and a.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/Matrix-Game-2.0-Diffusers-Base/FLASH_ATTN/video.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/Matrix-Game-2.0-Diffusers-Base/FLASH_ATTN/video.mp4 deleted file mode 100644 index deca1faad..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/Matrix-Game-2.0-Diffusers-Base/FLASH_ATTN/video.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/SFWan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/SFWan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 deleted file mode 100644 index b9cb854e5..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/SFWan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 deleted file mode 100644 index 83eede69b..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of_1.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of_1.mp4 deleted file mode 100644 index 83eede69b..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of_1.mp4 and /dev/null differ diff --git a/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 b/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 deleted file mode 100644 index b12da4d45..000000000 Binary files a/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 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/steps4_An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali_ssim.json b/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/steps4_An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali_ssim.json deleted file mode 100644 index db690f38d..000000000 --- a/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.2-I2V-A14B-Diffusers/SLA_ATTN/steps4_An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali_ssim.json +++ /dev/null @@ -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." - } -} \ No newline at end of file diff --git a/fastvideo/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 deleted file mode 100644 index 9a22a7088..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/Wan2.1-I2V-14B-480P-Diffusers/FLASH_ATTN/An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space reali.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 deleted file mode 100644 index 2ad844c66..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 deleted file mode 100644 index 282754528..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/Wan2.1-T2V-1.3B-Diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 deleted file mode 100644 index b9856eb9b..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/FLASH_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 deleted file mode 100644 index adf743b3f..000000000 Binary files a/fastvideo/tests/ssim/L40S_reference_videos/ltx2_diffusers/TORCH_SDPA/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 and /dev/null differ diff --git a/fastvideo/tests/ssim/README.md b/fastvideo/tests/ssim/README.md index 5246ed1d9..198f5bfe8 100644 --- a/fastvideo/tests/ssim/README.md +++ b/fastvideo/tests/ssim/README.md @@ -1,8 +1,52 @@ -The reference videos in the `*_reference_videos` directory are used as part of an e2e test to ensure consistency in video generation quality across code changes. `test_inference_similarity.py` compares newly generated videos against these references using Structural Similarity Index (SSIM) metrics to detect any regressions in visual quality across code changes. +The reference videos are used as part of e2e SSIM regression tests. +Wan inference coverage now lives in `test_wan_t2v_similarity.py` and +`test_wan_i2v_similarity.py` alongside the other model-specific SSIM files. +These tests compare newly generated videos against references to detect quality +regressions. -`A40_reference_videos` are generated on A40s and so on. +reference layout: +- `reference_videos/default/_reference_videos///...` +- `reference_videos/full_quality/_reference_videos///...` -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 +`/_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 ` + +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/_reference_videos///` +- full-quality params: `generated_videos/full_quality/_reference_videos///` + +HF repo layout mirrors quality + GPU split: +- `reference_videos/default/_reference_videos/...` +- `reference_videos/full_quality/_reference_videos/...` reference videos were generated on commit `4aeabbc629e0edf91477e80e795e7bb1823c71cb` causal videos were generated on commit b318063c0a4618f1d5d99ea82ca67a06aad0d19d diff --git a/fastvideo/tests/ssim/conftest.py b/fastvideo/tests/ssim/conftest.py index 318ad66e8..01b658a23 100644 --- a/fastvideo/tests/ssim/conftest.py +++ b/fastvideo/tests/ssim/conftest.py @@ -1,5 +1,60 @@ import os +from fastvideo.tests.ssim.reference_utils import ( + FULL_QUALITY_ENV_VAR, + get_output_quality_tier, +) +from fastvideo.tests.ssim.reference_videos_cli import ( + HF_REPO_ENV_KEY, + ensure_reference_videos_available, +) + + +def pytest_addoption(parser): + parser.addoption( + "--ssim-full-quality", + action="store_true", + default=False, + help=( + "Use *_FULL_QUALITY_PARAMS for SSIM tests. " + "Default keeps the original CI-friendly params." + ), + ) + parser.addoption( + "--ssim-reference-repo", + default=os.environ.get(HF_REPO_ENV_KEY, ""), + help=( + "HF repo id for SSIM reference videos " + f"(overrides {HF_REPO_ENV_KEY})." + ), + ) + parser.addoption( + "--skip-ssim-reference-download", + action="store_true", + default=False, + help="Skip auto-download of missing SSIM reference videos from HF.", + ) + + +def pytest_configure(config): + if config.getoption("--ssim-full-quality"): + os.environ[FULL_QUALITY_ENV_VAR] = "1" + + repo_id = config.getoption("--ssim-reference-repo") + if repo_id: + os.environ[HF_REPO_ENV_KEY] = repo_id + + skip_download = config.getoption("--skip-ssim-reference-download") + skip_download = skip_download or os.environ.get( + "FASTVIDEO_SSIM_SKIP_REFERENCE_DOWNLOAD", "" + ).strip().lower() in {"1", "true", "yes", "on"} + + if not skip_download: + ensure_reference_videos_available( + repo_id=repo_id or None, + quality_tier=get_output_quality_tier(), + ) + def pytest_collection_modifyitems(config, items): """Optionally keep only tests with a matching model_id parameter.""" diff --git a/fastvideo/tests/ssim/inference_similarity_utils.py b/fastvideo/tests/ssim/inference_similarity_utils.py new file mode 100644 index 000000000..c6059ac9b --- /dev/null +++ b/fastvideo/tests/ssim/inference_similarity_utils.py @@ -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, + ) diff --git a/fastvideo/tests/ssim/reference_utils.py b/fastvideo/tests/ssim/reference_utils.py new file mode 100644 index 000000000..78ee60b2f --- /dev/null +++ b/fastvideo/tests/ssim/reference_utils.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +from collections.abc import Sequence +from logging import Logger +from typing import TypeVar + +import torch + +DEVICE_REFERENCE_FOLDER_SUFFIX = "_reference_videos" +FULL_QUALITY_ENV_VAR = "FASTVIDEO_SSIM_FULL_QUALITY" +DEFAULT_OUTPUT_QUALITY_TIER = "default" +FULL_OUTPUT_QUALITY_TIER = "full_quality" +REFERENCE_VIDEOS_DIRNAME = "reference_videos" +DeviceMapping = tuple[str, str] +T = TypeVar("T") + + +def build_reference_folder_name( + device_prefix: str, + suffix: str = DEVICE_REFERENCE_FOLDER_SUFFIX, +) -> str: + return f"{device_prefix}{suffix}" + + +def get_cuda_device_name(index: int = 0) -> str: + if not torch.cuda.is_available(): + return "CPU" + return torch.cuda.get_device_name(index) + + +def resolve_device_reference_folder( + device_mappings: Sequence[DeviceMapping], + *, + device_name: str | None = None, + fallback_device_prefix: str | None = None, + unknown_device_prefix: str | None = None, + suffix: str = DEVICE_REFERENCE_FOLDER_SUFFIX, + logger: Logger | None = None, +) -> str | None: + resolved_device_name = device_name or get_cuda_device_name() + for device_pattern, device_prefix in device_mappings: + if device_pattern in resolved_device_name: + return build_reference_folder_name(device_prefix, suffix) + + if logger is not None: + if fallback_device_prefix is None: + logger.warning( + "Unsupported device for ssim tests: %s", + resolved_device_name, + ) + else: + logger.warning( + "Unsupported device for ssim tests: %s; using %s references", + resolved_device_name, + fallback_device_prefix, + ) + + if unknown_device_prefix is not None: + return build_reference_folder_name(unknown_device_prefix, suffix) + if fallback_device_prefix is not None: + return build_reference_folder_name(fallback_device_prefix, suffix) + return None + + +def build_reference_folder_path( + script_dir: str, + device_reference_folder: str, + model_id: str, + attention_backend: str, +) -> str: + quality_tier = get_output_quality_tier() + tiered_reference_folder = os.path.join( + script_dir, + REFERENCE_VIDEOS_DIRNAME, + quality_tier, + device_reference_folder, + model_id, + attention_backend, + ) + if os.path.exists(tiered_reference_folder): + return tiered_reference_folder + + legacy_reference_folder = os.path.join( + script_dir, + device_reference_folder, + model_id, + attention_backend, + ) + if quality_tier == DEFAULT_OUTPUT_QUALITY_TIER and os.path.exists( + legacy_reference_folder + ): + return legacy_reference_folder + + return tiered_reference_folder + + +def use_full_quality_configs() -> bool: + value = os.environ.get(FULL_QUALITY_ENV_VAR, "").strip().lower() + return value in {"1", "true", "yes", "on"} + + +def get_output_quality_tier() -> str: + if use_full_quality_configs(): + return FULL_OUTPUT_QUALITY_TIER + return DEFAULT_OUTPUT_QUALITY_TIER + + +def build_generated_output_dir( + script_dir: str, + device_reference_folder: str, + model_id: str, + attention_backend: str, +) -> str: + return os.path.join( + script_dir, + "generated_videos", + get_output_quality_tier(), + device_reference_folder, + model_id, + attention_backend, + ) + + +def select_ssim_params(default_params: T, full_quality_params: T) -> T: + if use_full_quality_configs(): + return full_quality_params + return default_params diff --git a/fastvideo/tests/ssim/reference_videos_cli.py b/fastvideo/tests/ssim/reference_videos_cli.py new file mode 100644 index 000000000..57577bf67 --- /dev/null +++ b/fastvideo/tests/ssim/reference_videos_cli.py @@ -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: /_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 /_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 _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: /generated_videos//" + ), + ) + 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//_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//_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//_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) diff --git a/fastvideo/tests/ssim/test_causal_similarity.py b/fastvideo/tests/ssim/test_causal_similarity.py index 6a7b6e809..18b104854 100644 --- a/fastvideo/tests/ssim/test_causal_similarity.py +++ b/fastvideo/tests/ssim/test_causal_similarity.py @@ -1,32 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 -import json import os -import torch import pytest -from fastvideo import VideoGenerator +from fastvideo.configs.sample.wan import ( + SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam, +) from fastvideo.logger import init_logger -from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results -from fastvideo.worker.multiproc_executor import MultiprocExecutor +from fastvideo.tests.ssim.inference_similarity_utils import ( + resolve_inference_device_reference_folder, + run_text_to_video_similarity_test, +) logger = init_logger(__name__) REQUIRED_GPUS = 1 -device_name = torch.cuda.get_device_name() -device_reference_folder_suffix = '_reference_videos' - -if "A40" in device_name: - device_reference_folder = "A40" + device_reference_folder_suffix -elif "L40S" in device_name: - device_reference_folder = "L40S" + device_reference_folder_suffix -else: - # device_reference_folder = "L40S" + device_reference_folder_suffix - logger.warning(f"Unsupported device for ssim tests: {device_name}") - # raise ValueError(f"Unsupported device for ssim tests: {device_name}") - -# Base parameters from the shell script +device_reference_folder = resolve_inference_device_reference_folder(logger) SF_WAN_T2V_PARAMS = { "num_gpus": 1, @@ -40,119 +30,53 @@ SF_WAN_T2V_PARAMS = { "tp_size": 1, } +_SF_WAN_T2V_FULL_QUALITY_DEFAULTS = SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam() +SF_WAN_T2V_FULL_QUALITY_PARAMS = { + "num_gpus": SF_WAN_T2V_PARAMS["num_gpus"], + "model_path": SF_WAN_T2V_PARAMS["model_path"], + "height": _SF_WAN_T2V_FULL_QUALITY_DEFAULTS.height, + "width": _SF_WAN_T2V_FULL_QUALITY_DEFAULTS.width, + "num_frames": SF_WAN_T2V_PARAMS["num_frames"], + "num_inference_steps": (_SF_WAN_T2V_FULL_QUALITY_DEFAULTS.num_inference_steps), + "guidance_scale": (_SF_WAN_T2V_FULL_QUALITY_DEFAULTS.guidance_scale), + "seed": _SF_WAN_T2V_FULL_QUALITY_DEFAULTS.seed, + "sp_size": SF_WAN_T2V_PARAMS["sp_size"], + "tp_size": SF_WAN_T2V_PARAMS["tp_size"], + "neg_prompt": (_SF_WAN_T2V_FULL_QUALITY_DEFAULTS.negative_prompt), +} MODEL_TO_PARAMS = { "SFWan2.1-T2V-1.3B-Diffusers": SF_WAN_T2V_PARAMS, } - -I2V_MODEL_TO_PARAMS = { +FULL_QUALITY_MODEL_TO_PARAMS = { + "SFWan2.1-T2V-1.3B-Diffusers": SF_WAN_T2V_FULL_QUALITY_PARAMS, } TEST_PROMPTS = [ - "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.", - # "A lone hiker stands atop a towering cliff, silhouetted against the vast horizon. The rugged landscape stretches endlessly beneath, its earthy tones blending into the soft blues of the sky. The scene captures the spirit of exploration and human resilience. High angle, dynamic framing, with soft natural lighting emphasizing the grandeur of nature." + "Will Smith casually eats noodles, his relaxed demeanor contrasting " + "with the energetic background of a bustling street food market. " + "The scene captures a mix of humor and authenticity. Mid-shot " + "framing, vibrant lighting.", ] -I2V_TEST_PROMPTS = [ - "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.", -] - -I2V_IMAGE_PATHS = [ - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg", -] @pytest.mark.parametrize("prompt", TEST_PROMPTS) -@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) +@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"]) @pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) -def test_causal_similarity(prompt, ATTENTION_BACKEND, model_id): - """ - Test that runs inference with different parameters and compares the output - to reference videos using SSIM. - """ - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND - - script_dir = os.path.dirname(os.path.abspath(__file__)) - - base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) - output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) - output_video_name = f"{prompt[:100].strip()}.mp4" - - os.makedirs(output_dir, exist_ok=True) - - BASE_PARAMS = MODEL_TO_PARAMS[model_id] - num_inference_steps = BASE_PARAMS["num_inference_steps"] - - init_kwargs = { - "num_gpus": BASE_PARAMS["num_gpus"], - "sp_size": BASE_PARAMS["sp_size"], - "tp_size": BASE_PARAMS["tp_size"], - "dit_cpu_offload": True, - } - if BASE_PARAMS.get("vae_sp"): - init_kwargs["vae_sp"] = True - init_kwargs["vae_tiling"] = True - #if "text-encoder-precision" in BASE_PARAMS: - # init_kwargs["text_encoder_precisions"] = BASE_PARAMS["text-encoder-precision"] - - generation_kwargs = { - "num_inference_steps": num_inference_steps, - "output_path": output_dir, - "height": BASE_PARAMS["height"], - "width": BASE_PARAMS["width"], - "num_frames": BASE_PARAMS["num_frames"], - "seed": BASE_PARAMS["seed"], - } - if "neg_prompt" in BASE_PARAMS: - generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"] - - generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs) - generator.generate_video(prompt, **generation_kwargs) - - if isinstance(generator.executor, MultiprocExecutor): - generator.executor.shutdown() - - assert os.path.exists( - output_dir), f"Output video was not generated at {output_dir}" - - reference_folder = os.path.join(script_dir, device_reference_folder, model_id, ATTENTION_BACKEND) - - if not os.path.exists(reference_folder): - logger.error("Reference folder missing") - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}") - - # Find the matching reference video based on the prompt - reference_video_name = None - - for filename in os.listdir(reference_folder): - if filename.endswith('.mp4') and prompt[:100].strip() in filename: - reference_video_name = filename - break - - if not reference_video_name: - logger.error(f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}") - raise FileNotFoundError(f"Reference video missing") - - reference_video_path = os.path.join(reference_folder, reference_video_name) - generated_video_path = os.path.join(output_dir, output_video_name) - - logger.info( - f"Computing SSIM between {reference_video_path} and {generated_video_path}" +def test_causal_similarity( + prompt: str, + attention_backend_name: str, + model_id: str, +) -> None: + run_text_to_video_similarity_test( + logger=logger, + script_dir=os.path.dirname(os.path.abspath(__file__)), + device_reference_folder=device_reference_folder, + prompt=prompt, + attention_backend_name=attention_backend_name, + model_id=model_id, + default_params_map=MODEL_TO_PARAMS, + full_quality_params_map=FULL_QUALITY_MODEL_TO_PARAMS, + min_acceptable_ssim=0.98, + init_kwargs_override={"dit_cpu_offload": True}, ) - ssim_values = compute_video_ssim_torchvision(reference_video_path, - generated_video_path, - use_ms_ssim=True) - - mean_ssim = ssim_values[0] - logger.info(f"SSIM mean value: {mean_ssim}") - logger.info(f"Writing SSIM results to directory: {output_dir}") - - success = write_ssim_results(output_dir, ssim_values, reference_video_path, - generated_video_path, num_inference_steps, - prompt) - - if not success: - logger.error("Failed to write SSIM results to file") - - min_acceptable_ssim = 0.98 - assert mean_ssim >= min_acceptable_ssim, f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}" diff --git a/fastvideo/tests/ssim/test_gamecraft_similarity.py b/fastvideo/tests/ssim/test_gamecraft_similarity.py index 8ed59beb2..4eabed58d 100644 --- a/fastvideo/tests/ssim/test_gamecraft_similarity.py +++ b/fastvideo/tests/ssim/test_gamecraft_similarity.py @@ -6,7 +6,8 @@ Generates a video with deterministic seed and camera trajectory, then compares against a device-specific reference video via MS-SSIM. Reference videos must be pre-generated and stored under: - _reference_videos/HunyuanGameCraft// + reference_videos//_reference_videos/HunyuanGameCraft/ + / 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): diff --git a/fastvideo/tests/ssim/test_inference_similarity.py b/fastvideo/tests/ssim/test_inference_similarity.py deleted file mode 100644 index d90a0bc88..000000000 --- a/fastvideo/tests/ssim/test_inference_similarity.py +++ /dev/null @@ -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}" diff --git a/fastvideo/tests/ssim/test_lingbot_similarity.py b/fastvideo/tests/ssim/test_lingbot_similarity.py index 457f2bcb4..20fc6140e 100644 --- a/fastvideo/tests/ssim/test_lingbot_similarity.py +++ b/fastvideo/tests/ssim/test_lingbot_similarity.py @@ -10,13 +10,21 @@ Note: num_inference_steps is reduced to 4 for faster CI. import os -import pytest -import torch - -from fastvideo import VideoGenerator -from fastvideo.logger import init_logger -from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embedding -from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results +import pytest +import torch + +from fastvideo import VideoGenerator +from fastvideo.configs.sample.lingbotworld import LingBotWorld_SamplingParam +from fastvideo.logger import init_logger +from fastvideo.models.dits.lingbotworld.cam_utils import prepare_camera_embedding +from fastvideo.tests.ssim.reference_utils import ( + build_generated_output_dir, + build_reference_folder_path, + get_cuda_device_name, + resolve_device_reference_folder, + select_ssim_params, +) +from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results logger = init_logger(__name__) @@ -40,21 +48,20 @@ def _find_lingbotworld_examples_root() -> str | None: return None -device_name = torch.cuda.get_device_name() -device_reference_folder_suffix = "_reference_videos" - -if "A40" in device_name: - device_reference_folder = "A40" + device_reference_folder_suffix -elif "L40S" in device_name: - device_reference_folder = "L40S" + device_reference_folder_suffix -elif "H100" in device_name: - device_reference_folder = "H100" + device_reference_folder_suffix -else: - device_reference_folder = None - logger.warning("Unsupported device for ssim tests: %s", device_name) +device_name = get_cuda_device_name() +device_reference_folder = resolve_device_reference_folder( + ( + ("A40", "A40"), + ("L40S", "L40S"), + ("H100", "H100"), + ("H200", "H200"), + ), + device_name=device_name, + logger=logger, +) -LINGBOT_PARAMS = { +LINGBOT_PARAMS = { "model_path": "FastVideo/LingBot-World-Base-Cam-Diffusers", "num_gpus": 2, "height": 256, @@ -81,8 +88,28 @@ LINGBOT_PARAMS = { "镜头晃动，画面闪烁，模糊，噪点，水印，签名，文字，变形，扭曲，液化，不合逻辑的结构，卡顿，" "PPT幻灯片感，过暗，欠曝，低对比度，霓虹灯光感，过度锐化，3D渲染感，人物，行人，游客，身体，" "皮肤，肢体，面部特征，汽车，电线" - ), -} + ), +} +_LINGBOT_FULL_QUALITY_DEFAULTS = LingBotWorld_SamplingParam() +LINGBOT_FULL_QUALITY_PARAMS = { + "model_path": LINGBOT_PARAMS["model_path"], + "num_gpus": LINGBOT_PARAMS["num_gpus"], + "height": _LINGBOT_FULL_QUALITY_DEFAULTS.height, + "width": _LINGBOT_FULL_QUALITY_DEFAULTS.width, + "num_frames": LINGBOT_PARAMS["num_frames"], # default num_frames: 125 + "num_inference_steps": _LINGBOT_FULL_QUALITY_DEFAULTS.num_inference_steps, + "guidance_scale": _LINGBOT_FULL_QUALITY_DEFAULTS.guidance_scale, + "guidance_scale_2": _LINGBOT_FULL_QUALITY_DEFAULTS.guidance_scale_2, + "embedded_cfg_scale": LINGBOT_PARAMS["embedded_cfg_scale"], + "flow_shift": LINGBOT_PARAMS["flow_shift"], + "boundary_ratio": _LINGBOT_FULL_QUALITY_DEFAULTS.boundary_ratio, + "seed": _LINGBOT_FULL_QUALITY_DEFAULTS.seed, + "fps": _LINGBOT_FULL_QUALITY_DEFAULTS.fps, + "spatial_scale": LINGBOT_PARAMS["spatial_scale"], + "example_case": LINGBOT_PARAMS["example_case"], + "image_path": LINGBOT_PARAMS["image_path"], + "negative_prompt": _LINGBOT_FULL_QUALITY_DEFAULTS.negative_prompt, +} TEST_PROMPTS = [ "The video presents a soaring journey through a fantasy jungle. The wind " @@ -95,74 +122,80 @@ TEST_PROMPTS = [ @pytest.mark.parametrize("prompt", TEST_PROMPTS) @pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) -def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND - - if device_reference_folder is None: - pytest.skip(f"Unsupported device for LingBot SSIM test: {device_name}") - if torch.cuda.device_count() < LINGBOT_PARAMS["num_gpus"]: - pytest.skip( - f"LingBot SSIM test requires {LINGBOT_PARAMS['num_gpus']} GPUs, " - f"but only {torch.cuda.device_count()} detected." - ) +def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + + params = select_ssim_params(LINGBOT_PARAMS, LINGBOT_FULL_QUALITY_PARAMS) + + if device_reference_folder is None: + pytest.skip(f"Unsupported device for LingBot SSIM test: {device_name}") + if torch.cuda.device_count() < params["num_gpus"]: + pytest.skip( + f"LingBot SSIM test requires {params['num_gpus']} GPUs, " + f"but only {torch.cuda.device_count()} detected." + ) examples_root = _find_lingbotworld_examples_root() if examples_root is None: pytest.skip( "lingbotworld_examples not found under examples/inference/basic.") - action_path = os.path.join(examples_root, LINGBOT_PARAMS["example_case"]) + action_path = os.path.join(examples_root, params["example_case"]) if not (os.path.exists(os.path.join(action_path, "poses.npy")) and os.path.exists(os.path.join(action_path, "intrinsics.npy"))): pytest.skip(f"Missing camera npy files under {action_path}") - c2ws_plucker_emb, aligned_num_frames = prepare_camera_embedding( - action_path=action_path, - num_frames=LINGBOT_PARAMS["num_frames"], - height=LINGBOT_PARAMS["height"], - width=LINGBOT_PARAMS["width"], - spatial_scale=LINGBOT_PARAMS["spatial_scale"], - ) + c2ws_plucker_emb, aligned_num_frames = prepare_camera_embedding( + action_path=action_path, + num_frames=params["num_frames"], + height=params["height"], + width=params["width"], + spatial_scale=params["spatial_scale"], + ) script_dir = os.path.dirname(os.path.abspath(__file__)) model_id = "LingBot-World-Base-Cam-Diffusers" - output_dir = os.path.join(script_dir, "generated_videos", model_id, - ATTENTION_BACKEND) + output_dir = build_generated_output_dir( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, + ) output_video_name = f"{prompt[:100].strip()}.mp4" os.makedirs(output_dir, exist_ok=True) - init_kwargs = { - "num_gpus": LINGBOT_PARAMS["num_gpus"], - "flow_shift": LINGBOT_PARAMS["flow_shift"], - "boundary_ratio": LINGBOT_PARAMS["boundary_ratio"], - "use_fsdp_inference": True, - "dit_cpu_offload": True, + init_kwargs = { + "num_gpus": params["num_gpus"], + "flow_shift": params["flow_shift"], + "boundary_ratio": params["boundary_ratio"], + "use_fsdp_inference": True, + "dit_cpu_offload": True, "dit_layerwise_offload": False, "text_encoder_cpu_offload": True, "vae_cpu_offload": False, "pin_cpu_memory": True, } - generation_kwargs = { - "output_path": output_dir, - "image_path": LINGBOT_PARAMS["image_path"], - "height": LINGBOT_PARAMS["height"], - "width": LINGBOT_PARAMS["width"], - "num_frames": aligned_num_frames, - "num_inference_steps": LINGBOT_PARAMS["num_inference_steps"], - "guidance_scale": LINGBOT_PARAMS["guidance_scale"], - "guidance_scale_2": LINGBOT_PARAMS["guidance_scale_2"], - "embedded_cfg_scale": LINGBOT_PARAMS["embedded_cfg_scale"], - "seed": LINGBOT_PARAMS["seed"], - "fps": LINGBOT_PARAMS["fps"], - "negative_prompt": LINGBOT_PARAMS["negative_prompt"], - "c2ws_plucker_emb": c2ws_plucker_emb, - } + generation_kwargs = { + "output_path": output_dir, + "image_path": params["image_path"], + "height": params["height"], + "width": params["width"], + "num_frames": aligned_num_frames, + "num_inference_steps": params["num_inference_steps"], + "guidance_scale": params["guidance_scale"], + "guidance_scale_2": params["guidance_scale_2"], + "embedded_cfg_scale": params["embedded_cfg_scale"], + "seed": params["seed"], + "fps": params["fps"], + "negative_prompt": params["negative_prompt"], + "c2ws_plucker_emb": c2ws_plucker_emb, + } generator: VideoGenerator | None = None try: - generator = VideoGenerator.from_pretrained( - model_path=LINGBOT_PARAMS["model_path"], **init_kwargs) - generator.generate_video(prompt, **generation_kwargs) + generator = VideoGenerator.from_pretrained( + model_path=params["model_path"], **init_kwargs) + generator.generate_video(prompt, **generation_kwargs) finally: if generator is not None: generator.shutdown() @@ -171,8 +204,12 @@ def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): assert os.path.exists(generated_video_path), ( f"Output video was not generated at {generated_video_path}") - reference_folder = os.path.join(script_dir, device_reference_folder, model_id, - ATTENTION_BACKEND) + reference_folder = build_reference_folder_path( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, + ) if not os.path.exists(reference_folder): raise FileNotFoundError( f"Reference video folder does not exist: {reference_folder}") @@ -196,9 +233,9 @@ def test_lingbot_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): mean_ssim = ssim_values[0] logger.info("SSIM mean value: %s", mean_ssim) - write_ssim_results(output_dir, ssim_values, reference_video_path, - generated_video_path, - LINGBOT_PARAMS["num_inference_steps"], prompt) + write_ssim_results(output_dir, ssim_values, reference_video_path, + generated_video_path, + params["num_inference_steps"], prompt) min_acceptable_ssim = 0.90 assert mean_ssim >= min_acceptable_ssim, ( diff --git a/fastvideo/tests/ssim/test_longcat_similarity.py b/fastvideo/tests/ssim/test_longcat_similarity.py index bf52ec5bf..7266f0482 100644 --- a/fastvideo/tests/ssim/test_longcat_similarity.py +++ b/fastvideo/tests/ssim/test_longcat_similarity.py @@ -20,7 +20,15 @@ import pytest import torch from fastvideo import VideoGenerator +from fastvideo.configs.sample.base import SamplingParam from fastvideo.logger import init_logger +from fastvideo.tests.ssim.reference_utils import ( + build_generated_output_dir, + build_reference_folder_path, + get_cuda_device_name, + resolve_device_reference_folder, + select_ssim_params, +) from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results logger = init_logger(__name__) @@ -28,17 +36,16 @@ logger = init_logger(__name__) REQUIRED_GPUS = 1 # Device-specific reference folder -device_name = torch.cuda.get_device_name() -device_reference_folder_suffix = "_reference_videos" - -if "A40" in device_name: - device_reference_folder = "A40" + device_reference_folder_suffix -elif "L40S" in device_name: - device_reference_folder = "L40S" + device_reference_folder_suffix -elif "H100" in device_name: - device_reference_folder = "H100" + device_reference_folder_suffix -else: - logger.warning(f"Unsupported device for ssim tests: {device_name}") +device_reference_folder = resolve_device_reference_folder( + ( + ("A40", "A40"), + ("L40S", "L40S"), + ("H100", "H100"), + ("H200", "H200"), + ), + device_name=get_cuda_device_name(), + logger=logger, +) # Common negative prompt from example scripts NEGATIVE_PROMPT = ( @@ -64,6 +71,21 @@ LONGCAT_T2V_PARAMS = { "seed": 42, "negative_prompt": NEGATIVE_PROMPT, } +_LONGCAT_T2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( + LONGCAT_T2V_PARAMS["model_path"] +) +LONGCAT_T2V_FULL_QUALITY_PARAMS = { + "num_gpus": LONGCAT_T2V_PARAMS["num_gpus"], + "model_path": LONGCAT_T2V_PARAMS["model_path"], + "height": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.height, + "width": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.width, + "num_frames": LONGCAT_T2V_PARAMS["num_frames"], # default num_frames: 125 + "num_inference_steps": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.num_inference_steps, + "guidance_scale": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.guidance_scale, + "fps": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.fps, + "seed": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.seed, + "negative_prompt": _LONGCAT_T2V_FULL_QUALITY_DEFAULTS.negative_prompt, +} # ============================================================================= # LongCat I2V Parameters (from basic_longcat_i2v.py) @@ -80,6 +102,21 @@ LONGCAT_I2V_PARAMS = { "seed": 42, "negative_prompt": NEGATIVE_PROMPT, } +_LONGCAT_I2V_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( + LONGCAT_I2V_PARAMS["model_path"] +) +LONGCAT_I2V_FULL_QUALITY_PARAMS = { + "num_gpus": LONGCAT_I2V_PARAMS["num_gpus"], + "model_path": LONGCAT_I2V_PARAMS["model_path"], + "height": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.height, + "width": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.width, + "num_frames": LONGCAT_I2V_PARAMS["num_frames"], # default num_frames: 125 + "num_inference_steps": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.num_inference_steps, + "guidance_scale": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.guidance_scale, + "fps": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.fps, + "seed": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.seed, + "negative_prompt": _LONGCAT_I2V_FULL_QUALITY_DEFAULTS.negative_prompt, +} # ============================================================================= # LongCat VC Parameters (from basic_longcat_vc.py) @@ -97,6 +134,22 @@ LONGCAT_VC_PARAMS = { "num_cond_frames": 13, "negative_prompt": NEGATIVE_PROMPT, } +_LONGCAT_VC_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( + LONGCAT_VC_PARAMS["model_path"] +) +LONGCAT_VC_FULL_QUALITY_PARAMS = { + "num_gpus": LONGCAT_VC_PARAMS["num_gpus"], + "model_path": LONGCAT_VC_PARAMS["model_path"], + "height": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.height, + "width": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.width, + "num_frames": LONGCAT_VC_PARAMS["num_frames"], # default num_frames: 125 + "num_inference_steps": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.num_inference_steps, + "guidance_scale": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.guidance_scale, + "fps": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.fps, + "seed": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.seed, + "num_cond_frames": LONGCAT_VC_PARAMS["num_cond_frames"], + "negative_prompt": _LONGCAT_VC_FULL_QUALITY_DEFAULTS.negative_prompt, +} # Test prompts T2V_TEST_PROMPTS = [ @@ -153,15 +206,22 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str): """ os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + params = select_ssim_params(LONGCAT_T2V_PARAMS, LONGCAT_T2V_FULL_QUALITY_PARAMS) + script_dir = os.path.dirname(os.path.abspath(__file__)) model_id = "LongCat-Video-T2V" - output_dir = os.path.join(script_dir, "generated_videos", model_id, ATTENTION_BACKEND) + output_dir = build_generated_output_dir( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, + ) output_video_name = f"{prompt[:100].strip()}.mp4" os.makedirs(output_dir, exist_ok=True) init_kwargs = { - "num_gpus": LONGCAT_T2V_PARAMS["num_gpus"], + "num_gpus": params["num_gpus"], "use_fsdp_inference": True, "dit_cpu_offload": True, "vae_cpu_offload": True, @@ -171,18 +231,18 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str): generation_kwargs = { "output_path": output_dir, - "height": LONGCAT_T2V_PARAMS["height"], - "width": LONGCAT_T2V_PARAMS["width"], - "num_frames": LONGCAT_T2V_PARAMS["num_frames"], - "num_inference_steps": LONGCAT_T2V_PARAMS["num_inference_steps"], - "guidance_scale": LONGCAT_T2V_PARAMS["guidance_scale"], - "fps": LONGCAT_T2V_PARAMS["fps"], - "seed": LONGCAT_T2V_PARAMS["seed"], - "negative_prompt": LONGCAT_T2V_PARAMS["negative_prompt"], + "height": params["height"], + "width": params["width"], + "num_frames": params["num_frames"], + "num_inference_steps": params["num_inference_steps"], + "guidance_scale": params["guidance_scale"], + "fps": params["fps"], + "seed": params["seed"], + "negative_prompt": params["negative_prompt"], } generator = VideoGenerator.from_pretrained( - model_path=LONGCAT_T2V_PARAMS["model_path"], **init_kwargs + model_path=params["model_path"], **init_kwargs ) generator.generate_video(prompt, **generation_kwargs) generator.shutdown() @@ -193,8 +253,11 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str): ) # Find reference video - reference_folder = os.path.join( - script_dir, device_reference_folder, model_id, ATTENTION_BACKEND + reference_folder = build_reference_folder_path( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, ) if not os.path.exists(reference_folder): raise FileNotFoundError( @@ -224,7 +287,7 @@ def test_longcat_t2v_similarity(prompt: str, ATTENTION_BACKEND: str): write_ssim_results( output_dir, ssim_values, reference_video_path, generated_video_path, - LONGCAT_T2V_PARAMS["num_inference_steps"], prompt + params["num_inference_steps"], prompt ) min_acceptable_ssim = 0.90 @@ -244,10 +307,17 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): """ os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + params = select_ssim_params(LONGCAT_I2V_PARAMS, LONGCAT_I2V_FULL_QUALITY_PARAMS) + script_dir = os.path.dirname(os.path.abspath(__file__)) model_id = "LongCat-Video-I2V" - output_dir = os.path.join(script_dir, "generated_videos", model_id, ATTENTION_BACKEND) + output_dir = build_generated_output_dir( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, + ) output_video_name = f"{prompt[:100].strip()}.mp4" os.makedirs(output_dir, exist_ok=True) @@ -256,7 +326,7 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): image_path = _resolve_asset_path(I2V_IMAGE_PATHS[prompt_idx]) init_kwargs = { - "num_gpus": LONGCAT_I2V_PARAMS["num_gpus"], + "num_gpus": params["num_gpus"], "use_fsdp_inference": True, "dit_cpu_offload": True, "vae_cpu_offload": True, @@ -267,18 +337,18 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): generation_kwargs = { "output_path": output_dir, "image_path": image_path, - "height": LONGCAT_I2V_PARAMS["height"], - "width": LONGCAT_I2V_PARAMS["width"], - "num_frames": LONGCAT_I2V_PARAMS["num_frames"], - "num_inference_steps": LONGCAT_I2V_PARAMS["num_inference_steps"], - "guidance_scale": LONGCAT_I2V_PARAMS["guidance_scale"], - "fps": LONGCAT_I2V_PARAMS["fps"], - "seed": LONGCAT_I2V_PARAMS["seed"], - "negative_prompt": LONGCAT_I2V_PARAMS["negative_prompt"], + "height": params["height"], + "width": params["width"], + "num_frames": params["num_frames"], + "num_inference_steps": params["num_inference_steps"], + "guidance_scale": params["guidance_scale"], + "fps": params["fps"], + "seed": params["seed"], + "negative_prompt": params["negative_prompt"], } generator = VideoGenerator.from_pretrained( - model_path=LONGCAT_I2V_PARAMS["model_path"], **init_kwargs + model_path=params["model_path"], **init_kwargs ) generator.generate_video(prompt, **generation_kwargs) generator.shutdown() @@ -289,8 +359,11 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): ) # Find reference video - reference_folder = os.path.join( - script_dir, device_reference_folder, model_id, ATTENTION_BACKEND + reference_folder = build_reference_folder_path( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, ) if not os.path.exists(reference_folder): raise FileNotFoundError( @@ -320,7 +393,7 @@ def test_longcat_i2v_similarity(prompt: str, ATTENTION_BACKEND: str): write_ssim_results( output_dir, ssim_values, reference_video_path, generated_video_path, - LONGCAT_I2V_PARAMS["num_inference_steps"], prompt + params["num_inference_steps"], prompt ) min_acceptable_ssim = 0.90 @@ -340,10 +413,17 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): """ os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + params = select_ssim_params(LONGCAT_VC_PARAMS, LONGCAT_VC_FULL_QUALITY_PARAMS) + script_dir = os.path.dirname(os.path.abspath(__file__)) model_id = "LongCat-Video-VC" - output_dir = os.path.join(script_dir, "generated_videos", model_id, ATTENTION_BACKEND) + output_dir = build_generated_output_dir( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, + ) output_video_name = f"{prompt[:100].strip()}.mp4" os.makedirs(output_dir, exist_ok=True) @@ -355,7 +435,7 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): pytest.skip(f"Input video not found at {video_path}") init_kwargs = { - "num_gpus": LONGCAT_VC_PARAMS["num_gpus"], + "num_gpus": params["num_gpus"], "use_fsdp_inference": False, "dit_cpu_offload": False, "vae_cpu_offload": True, @@ -367,19 +447,19 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): generation_kwargs = { "output_path": output_dir, "video_path": video_path, - "num_cond_frames": LONGCAT_VC_PARAMS["num_cond_frames"], - "height": LONGCAT_VC_PARAMS["height"], - "width": LONGCAT_VC_PARAMS["width"], - "num_frames": LONGCAT_VC_PARAMS["num_frames"], - "num_inference_steps": LONGCAT_VC_PARAMS["num_inference_steps"], - "guidance_scale": LONGCAT_VC_PARAMS["guidance_scale"], - "fps": LONGCAT_VC_PARAMS["fps"], - "seed": LONGCAT_VC_PARAMS["seed"], - "negative_prompt": LONGCAT_VC_PARAMS["negative_prompt"], + "num_cond_frames": params["num_cond_frames"], + "height": params["height"], + "width": params["width"], + "num_frames": params["num_frames"], + "num_inference_steps": params["num_inference_steps"], + "guidance_scale": params["guidance_scale"], + "fps": params["fps"], + "seed": params["seed"], + "negative_prompt": params["negative_prompt"], } generator = VideoGenerator.from_pretrained( - model_path=LONGCAT_VC_PARAMS["model_path"], **init_kwargs + model_path=params["model_path"], **init_kwargs ) generator.generate_video(prompt, **generation_kwargs) generator.shutdown() @@ -390,8 +470,11 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): ) # Find reference video - reference_folder = os.path.join( - script_dir, device_reference_folder, model_id, ATTENTION_BACKEND + reference_folder = build_reference_folder_path( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, ) if not os.path.exists(reference_folder): raise FileNotFoundError( @@ -421,7 +504,7 @@ def test_longcat_vc_similarity(prompt: str, ATTENTION_BACKEND: str): write_ssim_results( output_dir, ssim_values, reference_video_path, generated_video_path, - LONGCAT_VC_PARAMS["num_inference_steps"], prompt + params["num_inference_steps"], prompt ) min_acceptable_ssim = 0.90 diff --git a/fastvideo/tests/ssim/test_matrixgame_similarity.py b/fastvideo/tests/ssim/test_matrixgame_similarity.py index 3554a8f9a..533e80295 100644 --- a/fastvideo/tests/ssim/test_matrixgame_similarity.py +++ b/fastvideo/tests/ssim/test_matrixgame_similarity.py @@ -5,8 +5,16 @@ import torch import pytest from fastvideo import VideoGenerator -from fastvideo.models.dits.matrixgame.utils import create_action_presets +from fastvideo.configs.sample.wan import MatrixGame2_SamplingParam from fastvideo.logger import init_logger +from fastvideo.models.dits.matrixgame.utils import create_action_presets +from fastvideo.tests.ssim.reference_utils import ( + build_generated_output_dir, + build_reference_folder_path, + get_cuda_device_name, + resolve_device_reference_folder, + select_ssim_params, +) from fastvideo.tests.utils import ( compute_video_ssim_torchvision, write_ssim_results, @@ -17,17 +25,15 @@ logger = init_logger(__name__) REQUIRED_GPUS = 1 -device_name = torch.cuda.get_device_name() -device_reference_folder_suffix = "_reference_videos" - -if "A40" in device_name: - device_reference_folder = "A40" + device_reference_folder_suffix -elif "L40S" in device_name: - device_reference_folder = "L40S" + device_reference_folder_suffix -else: - # device_reference_folder = "L40S" + device_reference_folder_suffix - logger.warning(f"Unsupported device for ssim tests: {device_name}") - # raise ValueError(f"Unsupported device for ssim tests: {device_name}") +device_reference_folder = resolve_device_reference_folder( + ( + ("A40", "A40"), + ("L40S", "L40S"), + ("H200", "H200"), + ), + device_name=get_cuda_device_name(), + logger=logger, +) # Base parameters from the shell script @@ -41,11 +47,26 @@ MATRIXGAME_PARAMS = { "seed": 1024, "keyboard_dim": 4, } +_MATRIXGAME_FULL_QUALITY_DEFAULTS = MatrixGame2_SamplingParam() +MATRIXGAME_FULL_QUALITY_PARAMS = { + "num_gpus": MATRIXGAME_PARAMS["num_gpus"], + "model_path": MATRIXGAME_PARAMS["model_path"], + "height": _MATRIXGAME_FULL_QUALITY_DEFAULTS.height, + "width": _MATRIXGAME_FULL_QUALITY_DEFAULTS.width, + "num_frames": MATRIXGAME_PARAMS["num_frames"], # default num_frames: 57 + "num_inference_steps": _MATRIXGAME_FULL_QUALITY_DEFAULTS.num_inference_steps, + "guidance_scale": _MATRIXGAME_FULL_QUALITY_DEFAULTS.guidance_scale, + "seed": _MATRIXGAME_FULL_QUALITY_DEFAULTS.seed, + "keyboard_dim": MATRIXGAME_PARAMS["keyboard_dim"], +} MODEL_TO_PARAMS = { "Matrix-Game-2.0-Diffusers-Base": MATRIXGAME_PARAMS, } +FULL_QUALITY_MODEL_TO_PARAMS = { + "Matrix-Game-2.0-Diffusers-Base": MATRIXGAME_FULL_QUALITY_PARAMS, +} I2V_MODEL_TO_PARAMS = {} @@ -70,19 +91,26 @@ def test_matrixgame_similarity(prompt, ATTENTION_BACKEND, model_id): script_dir = os.path.dirname(os.path.abspath(__file__)) - base_output_dir = os.path.join(script_dir, "generated_videos", model_id) - output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) - output_video_name = f"output.mp4" + output_dir = build_generated_output_dir( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, + ) + output_video_name = "output.mp4" os.makedirs(output_dir, exist_ok=True) - BASE_PARAMS = MODEL_TO_PARAMS[model_id] + params_map = select_ssim_params( + MODEL_TO_PARAMS, + FULL_QUALITY_MODEL_TO_PARAMS, + ) + BASE_PARAMS = params_map[model_id] num_inference_steps = BASE_PARAMS["num_inference_steps"] # Create action conditions for MatrixGame actions = create_action_presets( - BASE_PARAMS["num_frames"], keyboard_dim=BASE_PARAMS["keyboard_dim"], - seed=BASE_PARAMS["seed"] + BASE_PARAMS["num_frames"], keyboard_dim=BASE_PARAMS["keyboard_dim"], seed=BASE_PARAMS["seed"] ) latent_frames = (BASE_PARAMS["num_frames"] - 1) // 4 + 1 grid_sizes = torch.tensor([latent_frames, 44, 80]) @@ -117,19 +145,18 @@ def test_matrixgame_similarity(prompt, ATTENTION_BACKEND, model_id): if isinstance(generator.executor, MultiprocExecutor): generator.executor.shutdown() - assert os.path.exists(output_dir), ( - f"Output video was not generated at {output_dir}" - ) + assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}" - reference_folder = os.path.join( - script_dir, device_reference_folder, model_id, ATTENTION_BACKEND + reference_folder = build_reference_folder_path( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, ) if not os.path.exists(reference_folder): logger.error("Reference folder missing") - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}" - ) + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") # Find the matching reference video based on the prompt reference_video_name = None @@ -140,20 +167,14 @@ def test_matrixgame_similarity(prompt, ATTENTION_BACKEND, model_id): break if not reference_video_name: - logger.error( - f"Reference video not found for model: {model_id} with backend: {ATTENTION_BACKEND}" - ) + logger.error(f"Reference video not found for model: {model_id} with backend: {ATTENTION_BACKEND}") raise FileNotFoundError("Reference video missing") reference_video_path = os.path.join(reference_folder, reference_video_name) generated_video_path = os.path.join(output_dir, output_video_name) - logger.info( - f"Computing SSIM between {reference_video_path} and {generated_video_path}" - ) - ssim_values = compute_video_ssim_torchvision( - reference_video_path, generated_video_path, use_ms_ssim=True - ) + logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") + ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) mean_ssim = ssim_values[0] logger.info(f"SSIM mean value: {mean_ssim}") diff --git a/fastvideo/tests/ssim/test_sd35_similarity.py b/fastvideo/tests/ssim/test_sd35_similarity.py index f0cc5b3d1..98ce90e59 100644 --- a/fastvideo/tests/ssim/test_sd35_similarity.py +++ b/fastvideo/tests/ssim/test_sd35_similarity.py @@ -1,41 +1,42 @@ # SPDX-License-Identifier: Apache-2.0 - from __future__ import annotations import os -import shlex -import logging import pytest import torch -logger = logging.getLogger(__name__) +from fastvideo.configs.sample.sd35 import SD35SamplingParam +from fastvideo.logger import init_logger +from fastvideo.tests.ssim.inference_similarity_utils import ( + run_text_to_video_similarity_test, +) +from fastvideo.tests.ssim.reference_utils import ( + get_cuda_device_name, + resolve_device_reference_folder, +) + +logger = init_logger(__name__) REQUIRED_GPUS = 1 -def _device_reference_folder() -> str: - """Pick a reference folder name based on the current CUDA device.""" - suffix = "_reference_videos" - device_name = torch.cuda.get_device_name(0) - - if "A40" in device_name: - return "A40" + suffix - if "L40S" in device_name: - return "L40S" + suffix - if "H100" in device_name: - return "H100" + suffix - if "RTX 4090" in device_name or "4090" in device_name: - return "RTX4090" + suffix - - logger.warning( - "Unsupported device for ssim tests: %s; using L40S references", device_name - ) - return "L40S" + suffix - - SD35_MODEL_PATH = os.getenv( "SD35_MODEL_DIR", - "/FastVideo/official_weights/stabilityai__stable-diffusion-3.5-medium", + "stabilityai/stable-diffusion-3.5-medium", +) + +device_reference_folder = resolve_device_reference_folder( + ( + ("A40", "A40"), + ("L40S", "L40S"), + ("H100", "H100"), + ("H200", "H200"), + ("RTX 4090", "RTX4090"), + ("4090", "RTX4090"), + ), + device_name=get_cuda_device_name(), + fallback_device_prefix="L40S", + logger=logger, ) MODEL_ID = "stabilityai__stable-diffusion-3.5-medium" @@ -44,163 +45,80 @@ TEST_PROMPTS = [ "a photo of a cat", ] -pytestmark = pytest.mark.filterwarnings( - "ignore:.*torch.jit.script_method.*:DeprecationWarning", +SD35_PARAMS = { + "num_gpus": 1, + "model_path": SD35_MODEL_PATH, + "sp_size": 1, + "tp_size": 1, + "height": 256, + "width": 256, + "num_frames": 1, + "fps": 1, + "num_inference_steps": 8, + "guidance_scale": 6.0, + "seed": 0, + "neg_prompt": "lowres, blurry", +} + +_SD35_FULL_QUALITY_DEFAULTS = SD35SamplingParam() +SD35_FULL_QUALITY_PARAMS = { + "num_gpus": 1, + "model_path": SD35_MODEL_PATH, + "sp_size": 1, + "tp_size": 1, + "height": _SD35_FULL_QUALITY_DEFAULTS.height, + "width": _SD35_FULL_QUALITY_DEFAULTS.width, + "num_frames": SD35_PARAMS["num_frames"], + "fps": _SD35_FULL_QUALITY_DEFAULTS.fps, + "num_inference_steps": (_SD35_FULL_QUALITY_DEFAULTS.num_inference_steps), + "guidance_scale": _SD35_FULL_QUALITY_DEFAULTS.guidance_scale, + "seed": _SD35_FULL_QUALITY_DEFAULTS.seed, + "neg_prompt": _SD35_FULL_QUALITY_DEFAULTS.negative_prompt, +} + +SD35_MODEL_TO_PARAMS = { + MODEL_ID: SD35_PARAMS, +} +SD35_FULL_QUALITY_MODEL_TO_PARAMS = { + MODEL_ID: SD35_FULL_QUALITY_PARAMS, +} + + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="SD3.5 SSIM test requires CUDA", ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="SD3.5 SSIM test requires CUDA") -@pytest.mark.parametrize("ATTENTION_BACKEND", ["TORCH_SDPA"]) @pytest.mark.parametrize("prompt", TEST_PROMPTS) -def test_sd35_similarity(prompt: str, ATTENTION_BACKEND: str) -> None: - from fastvideo import VideoGenerator - from fastvideo.tests.utils import ( - compute_video_ssim_torchvision, - write_ssim_results, - ) +@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"]) +@pytest.mark.parametrize("model_id", list(SD35_MODEL_TO_PARAMS.keys())) +def test_sd35_similarity( + prompt: str, + attention_backend_name: str, + model_id: str, +) -> None: + is_hf_repo = "/" in SD35_MODEL_PATH and not SD35_MODEL_PATH.startswith("/") + if not is_hf_repo and not os.path.isdir(SD35_MODEL_PATH): + pytest.skip(f"SD3.5 weights not found at {SD35_MODEL_PATH} (set SD35_MODEL_DIR to override)") - if not os.path.isdir(SD35_MODEL_PATH): - pytest.skip( - f"SD3.5 weights not found at {SD35_MODEL_PATH} (set SD35_MODEL_DIR to override)" - ) - - old_backend = os.environ.get("FASTVIDEO_ATTENTION_BACKEND") - old_transformers_verbosity = os.environ.get("TRANSFORMERS_VERBOSITY") - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND - os.environ["TRANSFORMERS_VERBOSITY"] = "error" - try: - script_dir = os.path.dirname(os.path.abspath(__file__)) - output_dir = os.path.join( - script_dir, "generated_videos", MODEL_ID, ATTENTION_BACKEND - ) - os.makedirs(output_dir, exist_ok=True) - - prompt_prefix = prompt[:100].strip() - output_video_name = f"{prompt_prefix}.mp4" - expected_video_path = os.path.join(output_dir, output_video_name) - - for filename in os.listdir(output_dir): - if filename.endswith(".mp4") and filename.startswith(prompt_prefix): - try: - os.remove(os.path.join(output_dir, filename)) - except FileNotFoundError: - pass - - init_kwargs = { - "num_gpus": 1, - # Save a single-frame MP4 to reuse the existing video-SSIM harness. + run_text_to_video_similarity_test( + logger=logger, + script_dir=os.path.dirname(os.path.abspath(__file__)), + device_reference_folder=device_reference_folder, + prompt=prompt, + attention_backend_name=attention_backend_name, + model_id=model_id, + default_params_map=SD35_MODEL_TO_PARAMS, + full_quality_params_map=SD35_FULL_QUALITY_MODEL_TO_PARAMS, + min_acceptable_ssim=0.98, + init_kwargs_override={ "workload_type": "t2v", - "sp_size": 1, - "tp_size": 1, - "dit_cpu_offload": False, - "dit_layerwise_offload": False, + "use_fsdp_inference": False, "text_encoder_cpu_offload": False, "vae_cpu_offload": False, "image_encoder_cpu_offload": False, "pin_cpu_memory": False, - "use_fsdp_inference": False, - } - - num_inference_steps = 8 - generation_kwargs = { - "output_path": output_dir, - "height": 256, - "width": 256, - "num_frames": 1, - "fps": 1, - "num_inference_steps": num_inference_steps, - "guidance_scale": 6.0, - "seed": 0, - "negative_prompt": "lowres, blurry", + }, + generation_kwargs_override={ "save_video": True, - } - - generator = VideoGenerator.from_pretrained(model_path=SD35_MODEL_PATH, **init_kwargs) - try: - generator.generate_video(prompt, **generation_kwargs) - finally: - generator.shutdown() - - generated_video_path = None - if os.path.exists(expected_video_path): - generated_video_path = expected_video_path - else: - candidates = [ - os.path.join(output_dir, f) - for f in os.listdir(output_dir) - if f.endswith(".mp4") and f.startswith(prompt_prefix) - ] - if candidates: - generated_video_path = max(candidates, key=os.path.getmtime) - - assert generated_video_path is not None and os.path.exists(generated_video_path), ( - f"Output video was not generated under {output_dir} for prompt '{prompt}'" - ) - - device_reference_folder = _device_reference_folder() - reference_folder = os.path.join(script_dir, device_reference_folder, MODEL_ID, ATTENTION_BACKEND) - - if not os.path.exists(reference_folder): - bless_cmd = ( - f"mkdir -p {shlex.quote(reference_folder)} && " - f"cp {shlex.quote(generated_video_path)} {shlex.quote(reference_folder)}/" - ) - pytest.fail( - f"Reference folder does not exist: {reference_folder}\n" - f"Generated video saved at: {generated_video_path}\n" - "To bless references, copy the generated mp4 into the reference folder with a matching name.\n" - f"Example:\n {bless_cmd}" - ) - - reference_video_path = os.path.join(reference_folder, output_video_name) - if not os.path.exists(reference_video_path): - reference_video_name = None - for filename in os.listdir(reference_folder): - if filename.endswith(".mp4") and filename.startswith(prompt_prefix): - reference_video_name = filename - break - if not reference_video_name: - bless_cmd = ( - f"cp {shlex.quote(generated_video_path)} {shlex.quote(reference_folder)}/" - ) - pytest.fail( - f"Reference video not found for prompt '{prompt}' under: {reference_folder}\n" - f"Expected name: {output_video_name}\n" - f"Generated video saved at: {generated_video_path}\n" - "To bless references, copy the generated mp4 into the reference folder.\n" - f"Example:\n {bless_cmd}" - ) - reference_video_path = os.path.join(reference_folder, reference_video_name) - - logger.info("Computing SSIM between %s and %s", reference_video_path, generated_video_path) - ssim_values = compute_video_ssim_torchvision( - reference_video_path, generated_video_path, use_ms_ssim=True - ) - - mean_ssim = ssim_values[0] - logger.info("SSIM mean value: %s", mean_ssim) - - write_ssim_results( - output_dir, - ssim_values, - reference_video_path, - generated_video_path, - num_inference_steps, - prompt, - ) - - min_acceptable_ssim = 0.98 - assert mean_ssim >= min_acceptable_ssim, ( - f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {MODEL_ID} with backend {ATTENTION_BACKEND}" - ) - - finally: - if old_backend is None: - os.environ.pop("FASTVIDEO_ATTENTION_BACKEND", None) - else: - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = old_backend - if old_transformers_verbosity is None: - os.environ.pop("TRANSFORMERS_VERBOSITY", None) - else: - os.environ["TRANSFORMERS_VERBOSITY"] = old_transformers_verbosity + }, + ) diff --git a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py index db6030a53..08f6c603c 100644 --- a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py +++ b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py @@ -1,184 +1,144 @@ # SPDX-License-Identifier: Apache-2.0 -""" -SSIM-based similarity test for TurboDiffusion inference. - -TurboDiffusion uses the SLA (Sparse-Linear Attention) backend and RCM scheduler -for 1-4 step video generation. -""" import os -import torch import pytest -from fastvideo import VideoGenerator +from fastvideo.configs.sample.turbodiffusion import ( + TurboDiffusionI2V_A14B_SamplingParam, + TurboDiffusionT2V_1_3B_SamplingParam, +) from fastvideo.logger import init_logger -from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results -from fastvideo.worker.multiproc_executor import MultiprocExecutor +from fastvideo.tests.ssim.inference_similarity_utils import ( + DEVICE_MAPPINGS, + resolve_inference_device_reference_folder, + run_image_to_video_similarity_test, + run_text_to_video_similarity_test, +) +from fastvideo.tests.ssim.reference_utils import ( + get_cuda_device_name, + resolve_device_reference_folder, +) logger = init_logger(__name__) REQUIRED_GPUS = 4 -device_name = torch.cuda.get_device_name() -device_reference_folder_suffix = '_reference_videos' - -if "A40" in device_name: - device_reference_folder = "A40" + device_reference_folder_suffix -elif "L40S" in device_name: - device_reference_folder = "L40S" + device_reference_folder_suffix -else: - # device_reference_folder = "L40S" + device_reference_folder_suffix - logger.warning(f"Unsupported device for ssim tests: {device_name}, using L40S references") +device_name = get_cuda_device_name() +device_reference_folder = resolve_device_reference_folder( + DEVICE_MAPPINGS, + device_name=device_name, +) +if device_reference_folder is None: raise ValueError(f"Unsupported device for ssim tests: {device_name}") - -# TurboDiffusion parameters (1-4 step generation with RCM scheduler + SLA attention) TURBODIFFUSION_PARAMS = { "num_gpus": 4, "model_path": "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers", "height": 480, "width": 832, "num_frames": 81, - "num_inference_steps": 4, # TurboDiffusion uses 1-4 steps - "guidance_scale": 1.0, # No CFG for TurboDiffusion + "num_inference_steps": 4, + "guidance_scale": 1.0, "seed": 42, "sp_size": 4, "tp_size": 1, "fps": 24, } +_TURBODIFFUSION_FULL_QUALITY_DEFAULTS = TurboDiffusionT2V_1_3B_SamplingParam() +TURBODIFFUSION_FULL_QUALITY_PARAMS = { + "num_gpus": TURBODIFFUSION_PARAMS["num_gpus"], + "model_path": TURBODIFFUSION_PARAMS["model_path"], + "height": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.height, + "width": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.width, + "num_frames": TURBODIFFUSION_PARAMS["num_frames"], + "num_inference_steps": (_TURBODIFFUSION_FULL_QUALITY_DEFAULTS.num_inference_steps), + "guidance_scale": (_TURBODIFFUSION_FULL_QUALITY_DEFAULTS.guidance_scale), + "seed": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.seed, + "sp_size": TURBODIFFUSION_PARAMS["sp_size"], + "tp_size": TURBODIFFUSION_PARAMS["tp_size"], + "fps": _TURBODIFFUSION_FULL_QUALITY_DEFAULTS.fps, +} TURBODIFFUSION_MODEL_TO_PARAMS = { "TurboWan2.1-T2V-1.3B-Diffusers": TURBODIFFUSION_PARAMS, } +FULL_QUALITY_TURBODIFFUSION_MODEL_TO_PARAMS = { + "TurboWan2.1-T2V-1.3B-Diffusers": (TURBODIFFUSION_FULL_QUALITY_PARAMS), +} TURBODIFFUSION_TEST_PROMPTS = [ - "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.", + "Will Smith casually eats noodles, his relaxed demeanor contrasting " + "with the energetic background of a bustling street food market. " + "The scene captures a mix of humor and authenticity. Mid-shot " + "framing, vibrant lighting.", ] @pytest.mark.parametrize("prompt", TURBODIFFUSION_TEST_PROMPTS) -@pytest.mark.parametrize("model_id", list(TURBODIFFUSION_MODEL_TO_PARAMS.keys())) -def test_turbodiffusion_inference_similarity(prompt, model_id): - """ - Test that runs TurboDiffusion inference with SLA attention and RCM scheduler, - then compares the output to reference videos using SSIM. - """ - # TurboDiffusion requires SLA attention backend - ATTENTION_BACKEND = "SLA_ATTN" - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND - - script_dir = os.path.dirname(os.path.abspath(__file__)) - - base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) - output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) - output_video_name = f"{prompt[:100].strip()}.mp4" - - os.makedirs(output_dir, exist_ok=True) - - BASE_PARAMS = TURBODIFFUSION_MODEL_TO_PARAMS[model_id] - num_inference_steps = BASE_PARAMS["num_inference_steps"] - - init_kwargs = { - "num_gpus": BASE_PARAMS["num_gpus"], - "sp_size": BASE_PARAMS["sp_size"], - "tp_size": BASE_PARAMS["tp_size"], - "override_pipeline_cls_name": "TurboDiffusionPipeline", - } - - generation_kwargs = { - "num_inference_steps": num_inference_steps, - "output_path": output_dir, - "height": BASE_PARAMS["height"], - "width": BASE_PARAMS["width"], - "num_frames": BASE_PARAMS["num_frames"], - "guidance_scale": BASE_PARAMS["guidance_scale"], - "seed": BASE_PARAMS["seed"], - "fps": BASE_PARAMS["fps"], - } - - generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs) - generator.generate_video(prompt, **generation_kwargs) - - if isinstance(generator.executor, MultiprocExecutor): - generator.executor.shutdown() - - assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}" - - reference_folder = os.path.join( - script_dir, device_reference_folder, model_id, ATTENTION_BACKEND - ) - - if not os.path.exists(reference_folder): - logger.error("Reference folder missing") - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}" - ) - - # Find the matching reference video based on the prompt - reference_video_name = None - - for filename in os.listdir(reference_folder): - if filename.endswith('.mp4') and prompt[:100].strip() in filename: - reference_video_name = filename - break - - if not reference_video_name: - logger.error( - f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}" - ) - raise FileNotFoundError(f"Reference video missing") - - reference_video_path = os.path.join(reference_folder, reference_video_name) - generated_video_path = os.path.join(output_dir, output_video_name) - - logger.info( - f"Computing SSIM between {reference_video_path} and {generated_video_path}" - ) - ssim_values = compute_video_ssim_torchvision( - reference_video_path, generated_video_path, use_ms_ssim=True - ) - - mean_ssim = ssim_values[0] - logger.info(f"SSIM mean value: {mean_ssim}") - logger.info(f"Writing SSIM results to directory: {output_dir}") - - success = write_ssim_results( - output_dir, ssim_values, reference_video_path, - generated_video_path, num_inference_steps, prompt - ) - - if not success: - logger.error("Failed to write SSIM results to file") - - # TurboDiffusion uses fewer steps, may have slightly lower SSIM - min_acceptable_ssim = 0.95 - assert mean_ssim >= min_acceptable_ssim, ( - f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}" +@pytest.mark.parametrize( + "model_id", + list(TURBODIFFUSION_MODEL_TO_PARAMS.keys()), +) +def test_turbodiffusion_inference_similarity( + prompt: str, + model_id: str, +) -> None: + run_text_to_video_similarity_test( + logger=logger, + script_dir=os.path.dirname(os.path.abspath(__file__)), + device_reference_folder=device_reference_folder, + prompt=prompt, + attention_backend_name="SLA_ATTN", + model_id=model_id, + default_params_map=TURBODIFFUSION_MODEL_TO_PARAMS, + full_quality_params_map=(FULL_QUALITY_TURBODIFFUSION_MODEL_TO_PARAMS), + min_acceptable_ssim=0.95, + init_kwargs_override={ + "override_pipeline_cls_name": "TurboDiffusionPipeline", + }, ) -# TurboDiffusion I2V parameters (dual-model with RCM scheduler + SLA attention) TURBODIFFUSION_I2V_PARAMS = { "num_gpus": 4, "model_path": "loayrashid/TurboWan2.2-I2V-A14B-Diffusers", "height": 480, "width": 832, "num_frames": 45, - "num_inference_steps": 4, # TurboDiffusion uses 1-4 steps - "guidance_scale": 1.0, # No CFG for TurboDiffusion + "num_inference_steps": 4, + "guidance_scale": 1.0, "seed": 42, "sp_size": 4, "tp_size": 1, "fps": 24, } +_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS = TurboDiffusionI2V_A14B_SamplingParam() +TURBODIFFUSION_I2V_FULL_QUALITY_PARAMS = { + "num_gpus": TURBODIFFUSION_I2V_PARAMS["num_gpus"], + "model_path": TURBODIFFUSION_I2V_PARAMS["model_path"], + "height": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.height, + "width": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.width, + "num_frames": TURBODIFFUSION_I2V_PARAMS["num_frames"], + "num_inference_steps": (_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.num_inference_steps), + "guidance_scale": (_TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.guidance_scale), + "seed": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.seed, + "sp_size": TURBODIFFUSION_I2V_PARAMS["sp_size"], + "tp_size": TURBODIFFUSION_I2V_PARAMS["tp_size"], + "fps": _TURBODIFFUSION_I2V_FULL_QUALITY_DEFAULTS.fps, +} TURBODIFFUSION_I2V_MODEL_TO_PARAMS = { "TurboWan2.2-I2V-A14B-Diffusers": TURBODIFFUSION_I2V_PARAMS, } +FULL_QUALITY_TURBODIFFUSION_I2V_MODEL_TO_PARAMS = { + "TurboWan2.2-I2V-A14B-Diffusers": (TURBODIFFUSION_I2V_FULL_QUALITY_PARAMS), +} TURBODIFFUSION_I2V_TEST_PROMPTS = [ - "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.", + "An astronaut hatching from an egg, on the surface of the moon, " + "the darkness and depth of space realised in the background. " + "High quality, ultrarealistic detail and breath-taking " + "movie-like camera shot.", ] TURBODIFFUSION_I2V_IMAGE_PATHS = [ @@ -187,114 +147,27 @@ TURBODIFFUSION_I2V_IMAGE_PATHS = [ @pytest.mark.parametrize("prompt", TURBODIFFUSION_I2V_TEST_PROMPTS) -@pytest.mark.parametrize("model_id", list(TURBODIFFUSION_I2V_MODEL_TO_PARAMS.keys())) -def test_turbodiffusion_i2v_inference_similarity(prompt, model_id): - """ - Test that runs TurboDiffusion I2V inference with dual-model switching, - then compares the output to reference videos using SSIM. - """ - # TurboDiffusion requires SLA attention backend - ATTENTION_BACKEND = "SLA_ATTN" - os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND - - assert len(TURBODIFFUSION_I2V_TEST_PROMPTS) == len(TURBODIFFUSION_I2V_IMAGE_PATHS), \ - "Expect number of prompts equal to number of images" - - script_dir = os.path.dirname(os.path.abspath(__file__)) - - base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) - output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) - output_video_name = f"{prompt[:100].strip()}.mp4" - - os.makedirs(output_dir, exist_ok=True) - - BASE_PARAMS = TURBODIFFUSION_I2V_MODEL_TO_PARAMS[model_id] - num_inference_steps = BASE_PARAMS["num_inference_steps"] +@pytest.mark.parametrize( + "model_id", + list(TURBODIFFUSION_I2V_MODEL_TO_PARAMS.keys()), +) +def test_turbodiffusion_i2v_inference_similarity( + prompt: str, + model_id: str, +) -> None: image_path = TURBODIFFUSION_I2V_IMAGE_PATHS[TURBODIFFUSION_I2V_TEST_PROMPTS.index(prompt)] - - init_kwargs = { - "num_gpus": BASE_PARAMS["num_gpus"], - "sp_size": BASE_PARAMS["sp_size"], - "tp_size": BASE_PARAMS["tp_size"], - "override_pipeline_cls_name": "TurboDiffusionI2VPipeline", - # Keep both transformers in VRAM - avoids CPU RAM bottleneck - "dit_cpu_offload": False, - "use_fsdp_inference": True, - "dit_layerwise_offload": False, - } - - generation_kwargs = { - "num_inference_steps": num_inference_steps, - "output_path": output_dir, - "image_path": image_path, - "height": BASE_PARAMS["height"], - "width": BASE_PARAMS["width"], - "num_frames": BASE_PARAMS["num_frames"], - "guidance_scale": BASE_PARAMS["guidance_scale"], - "seed": BASE_PARAMS["seed"], - "fps": BASE_PARAMS["fps"], - } - - generator = VideoGenerator.from_pretrained( - model_path=BASE_PARAMS["model_path"], - **init_kwargs - ) - generator.generate_video(prompt, **generation_kwargs) - - if isinstance(generator.executor, MultiprocExecutor): - generator.executor.shutdown() - - assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}" - - reference_folder = os.path.join( - script_dir, device_reference_folder, model_id, ATTENTION_BACKEND - ) - - if not os.path.exists(reference_folder): - logger.error("Reference folder missing") - raise FileNotFoundError( - f"Reference video folder does not exist: {reference_folder}" - ) - - # Find the matching reference video based on the prompt - reference_video_name = None - - for filename in os.listdir(reference_folder): - if filename.endswith('.mp4') and prompt[:100].strip() in filename: - reference_video_name = filename - break - - if not reference_video_name: - logger.error( - f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}" - ) - raise FileNotFoundError(f"Reference video missing") - - reference_video_path = os.path.join(reference_folder, reference_video_name) - generated_video_path = os.path.join(output_dir, output_video_name) - - logger.info( - f"Computing SSIM between {reference_video_path} and {generated_video_path}" - ) - ssim_values = compute_video_ssim_torchvision( - reference_video_path, generated_video_path, use_ms_ssim=True - ) - - mean_ssim = ssim_values[0] - logger.info(f"SSIM mean value: {mean_ssim}") - logger.info(f"Writing SSIM results to directory: {output_dir}") - - success = write_ssim_results( - output_dir, ssim_values, reference_video_path, - generated_video_path, num_inference_steps, prompt - ) - - if not success: - logger.error("Failed to write SSIM results to file") - - # TurboDiffusion I2V uses fewer steps, may have slightly lower SSIM - min_acceptable_ssim = 0.95 - assert mean_ssim >= min_acceptable_ssim, ( - f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " - f"for {model_id} with backend {ATTENTION_BACKEND}" + run_image_to_video_similarity_test( + logger=logger, + script_dir=os.path.dirname(os.path.abspath(__file__)), + device_reference_folder=device_reference_folder, + prompt=prompt, + image_path=image_path, + attention_backend_name="SLA_ATTN", + model_id=model_id, + default_params_map=TURBODIFFUSION_I2V_MODEL_TO_PARAMS, + full_quality_params_map=(FULL_QUALITY_TURBODIFFUSION_I2V_MODEL_TO_PARAMS), + min_acceptable_ssim=0.95, + init_kwargs_override={ + "override_pipeline_cls_name": ("TurboDiffusionI2VPipeline"), + }, ) diff --git a/fastvideo/tests/ssim/test_wan_i2v_similarity.py b/fastvideo/tests/ssim/test_wan_i2v_similarity.py new file mode 100644 index 000000000..b6eb97b0d --- /dev/null +++ b/fastvideo/tests/ssim/test_wan_i2v_similarity.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +import os + +import pytest + +from fastvideo.configs.sample.wan import WanI2V_14B_480P_SamplingParam +from fastvideo.logger import init_logger +from fastvideo.tests.ssim.inference_similarity_utils import ( + resolve_inference_device_reference_folder, + run_image_to_video_similarity_test, +) + +logger = init_logger(__name__) + +REQUIRED_GPUS = 2 + +device_reference_folder = resolve_inference_device_reference_folder(logger) + +WAN_I2V_PARAMS = { + "num_gpus": 2, + "model_path": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers", + "height": 480, + "width": 832, + "num_frames": 45, + "num_inference_steps": 2, + "guidance_scale": 5.0, + "embedded_cfg_scale": 6, + "flow_shift": 7.0, + "seed": 1024, + "sp_size": 2, + "tp_size": 1, + "vae_sp": True, + "fps": 24, + "neg_prompt": ( + "Bright tones, overexposed, static, blurred details, subtitles, style, " + "works, paintings, images, static, overall gray, worst quality, low " + "quality, JPEG compression residue, ugly, incomplete, extra fingers, " + "poorly drawn hands, poorly drawn faces, deformed, disfigured, " + "misshapen limbs, fused fingers, still picture, messy background, " + "three legs, many people in the background, walking backwards" + ), + "text-encoder-precision": ("fp32",), +} +_WAN_I2V_FULL_QUALITY_DEFAULTS = WanI2V_14B_480P_SamplingParam() +WAN_I2V_FULL_QUALITY_PARAMS = { + "num_gpus": WAN_I2V_PARAMS["num_gpus"], + "model_path": WAN_I2V_PARAMS["model_path"], + "height": _WAN_I2V_FULL_QUALITY_DEFAULTS.height, + "width": _WAN_I2V_FULL_QUALITY_DEFAULTS.width, + "num_frames": WAN_I2V_PARAMS["num_frames"], # default num_frames: 81 + "num_inference_steps": _WAN_I2V_FULL_QUALITY_DEFAULTS.num_inference_steps, + "guidance_scale": _WAN_I2V_FULL_QUALITY_DEFAULTS.guidance_scale, + "embedded_cfg_scale": WAN_I2V_PARAMS["embedded_cfg_scale"], + "flow_shift": WAN_I2V_PARAMS["flow_shift"], + "seed": _WAN_I2V_FULL_QUALITY_DEFAULTS.seed, + "sp_size": WAN_I2V_PARAMS["sp_size"], + "tp_size": WAN_I2V_PARAMS["tp_size"], + "vae_sp": WAN_I2V_PARAMS["vae_sp"], + "fps": _WAN_I2V_FULL_QUALITY_DEFAULTS.fps, + "neg_prompt": _WAN_I2V_FULL_QUALITY_DEFAULTS.negative_prompt, + "text-encoder-precision": WAN_I2V_PARAMS["text-encoder-precision"], +} + +WAN_I2V_MODEL_TO_PARAMS = { + "Wan2.1-I2V-14B-480P-Diffusers": WAN_I2V_PARAMS, +} +FULL_QUALITY_WAN_I2V_MODEL_TO_PARAMS = { + "Wan2.1-I2V-14B-480P-Diffusers": WAN_I2V_FULL_QUALITY_PARAMS, +} + +WAN_I2V_TEST_CASES = [ + ( + "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background. High quality, ultrarealistic detail and breath-taking movie-like camera shot.", + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg", + ), +] + + +@pytest.mark.parametrize(("prompt", "image_path"), WAN_I2V_TEST_CASES) +@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN"]) +@pytest.mark.parametrize("model_id", list(WAN_I2V_MODEL_TO_PARAMS.keys())) +def test_wan_i2v_inference_similarity( + prompt: str, + image_path: str, + attention_backend_name: str, + model_id: str, +) -> None: + run_image_to_video_similarity_test( + logger=logger, + script_dir=os.path.dirname(os.path.abspath(__file__)), + device_reference_folder=device_reference_folder, + prompt=prompt, + image_path=image_path, + attention_backend_name=attention_backend_name, + model_id=model_id, + default_params_map=WAN_I2V_MODEL_TO_PARAMS, + full_quality_params_map=FULL_QUALITY_WAN_I2V_MODEL_TO_PARAMS, + min_acceptable_ssim=0.97, + ) diff --git a/fastvideo/tests/ssim/test_wan_t2v_similarity.py b/fastvideo/tests/ssim/test_wan_t2v_similarity.py new file mode 100644 index 000000000..2f4f0bd57 --- /dev/null +++ b/fastvideo/tests/ssim/test_wan_t2v_similarity.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +import os + +import pytest + +from fastvideo.configs.sample.wan import WanT2V_1_3B_SamplingParam +from fastvideo.logger import init_logger +from fastvideo.tests.ssim.inference_similarity_utils import ( + resolve_inference_device_reference_folder, + run_text_to_video_similarity_test, +) + +logger = init_logger(__name__) + +REQUIRED_GPUS = 2 + +device_reference_folder = resolve_inference_device_reference_folder(logger) + +WAN_T2V_PARAMS = { + "num_gpus": 2, + "model_path": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + "height": 480, + "width": 832, + "num_frames": 45, + "num_inference_steps": 4, + "guidance_scale": 3, + "embedded_cfg_scale": 6, + "flow_shift": 7.0, + "seed": 1024, + "sp_size": 2, + "tp_size": 1, + "vae_sp": True, + "fps": 24, + "neg_prompt": ( + "Bright tones, overexposed, static, blurred details, subtitles, style, " + "works, paintings, images, static, overall gray, worst quality, low " + "quality, JPEG compression residue, ugly, incomplete, extra fingers, " + "poorly drawn hands, poorly drawn faces, deformed, disfigured, " + "misshapen limbs, fused fingers, still picture, messy background, " + "three legs, many people in the background, walking backwards" + ), + "text-encoder-precision": ("fp32",), +} +_WAN_T2V_FULL_QUALITY_DEFAULTS = WanT2V_1_3B_SamplingParam() +WAN_T2V_FULL_QUALITY_PARAMS = { + "num_gpus": WAN_T2V_PARAMS["num_gpus"], + "model_path": WAN_T2V_PARAMS["model_path"], + "height": _WAN_T2V_FULL_QUALITY_DEFAULTS.height, + "width": _WAN_T2V_FULL_QUALITY_DEFAULTS.width, + "num_frames": WAN_T2V_PARAMS["num_frames"], # default num_frames: 81 + "num_inference_steps": _WAN_T2V_FULL_QUALITY_DEFAULTS.num_inference_steps, + "guidance_scale": _WAN_T2V_FULL_QUALITY_DEFAULTS.guidance_scale, + "embedded_cfg_scale": WAN_T2V_PARAMS["embedded_cfg_scale"], + "flow_shift": WAN_T2V_PARAMS["flow_shift"], + "seed": _WAN_T2V_FULL_QUALITY_DEFAULTS.seed, + "sp_size": WAN_T2V_PARAMS["sp_size"], + "tp_size": WAN_T2V_PARAMS["tp_size"], + "vae_sp": WAN_T2V_PARAMS["vae_sp"], + "fps": _WAN_T2V_FULL_QUALITY_DEFAULTS.fps, + "neg_prompt": _WAN_T2V_FULL_QUALITY_DEFAULTS.negative_prompt, + "text-encoder-precision": WAN_T2V_PARAMS["text-encoder-precision"], +} + +WAN_T2V_MODEL_TO_PARAMS = { + "Wan2.1-T2V-1.3B-Diffusers": WAN_T2V_PARAMS, +} +FULL_QUALITY_WAN_T2V_MODEL_TO_PARAMS = { + "Wan2.1-T2V-1.3B-Diffusers": WAN_T2V_FULL_QUALITY_PARAMS, +} + +WAN_T2V_TEST_PROMPTS = [ + "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.", +] + + +@pytest.mark.parametrize("prompt", WAN_T2V_TEST_PROMPTS) +@pytest.mark.parametrize("attention_backend_name", ["FLASH_ATTN", "TORCH_SDPA"]) +@pytest.mark.parametrize("model_id", list(WAN_T2V_MODEL_TO_PARAMS.keys())) +def test_wan_t2v_inference_similarity( + prompt: str, + attention_backend_name: str, + model_id: str, +) -> None: + run_text_to_video_similarity_test( + logger=logger, + script_dir=os.path.dirname(os.path.abspath(__file__)), + device_reference_folder=device_reference_folder, + prompt=prompt, + attention_backend_name=attention_backend_name, + model_id=model_id, + default_params_map=WAN_T2V_MODEL_TO_PARAMS, + full_quality_params_map=FULL_QUALITY_WAN_T2V_MODEL_TO_PARAMS, + min_acceptable_ssim=0.93, + ) diff --git a/fastvideo/tests/ssim/update_reference_videos.sh b/fastvideo/tests/ssim/update_reference_videos.sh deleted file mode 100644 index e137c1b44..000000000 --- a/fastvideo/tests/ssim/update_reference_videos.sh +++ /dev/null @@ -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}" diff --git a/fastvideo/tests/utils.py b/fastvideo/tests/utils.py index 930a4720a..77408231c 100644 --- a/fastvideo/tests/utils.py +++ b/fastvideo/tests/utils.py @@ -7,11 +7,42 @@ from fastvideo.logger import init_logger import numpy as np import torch from pytorch_msssim import ms_ssim, ssim -from torchvision.io import read_video logger = init_logger(__name__) +def _read_video_frames(path: str) -> torch.Tensor: + """Read video frames as a ``(T, C, H, W)`` uint8 tensor. + + Tries ``torchvision.io.read_video`` first (available in older + torchvision) and falls back to PyAV when the function has been + removed (torchvision >= 0.26). + """ + try: + from torchvision.io import read_video + + frames, _, _ = read_video( + path, + pts_unit="sec", + output_format="TCHW", + ) + return frames + except (ImportError, AttributeError): + pass + + import av + + container = av.open(path) + frames = [] + for frame in container.decode(video=0): + arr = frame.to_ndarray(format="rgb24") + frames.append(torch.from_numpy(arr).permute(2, 0, 1)) + container.close() + if not frames: + raise RuntimeError(f"No video frames decoded from {path}") + return torch.stack(frames) + + def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True): """ Compute SSIM between two videos. @@ -27,12 +58,8 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True): if not os.path.exists(video2_path): raise FileNotFoundError(f"Video2 not found: {video2_path}") - frames1, _, _ = read_video(video1_path, - pts_unit='sec', - output_format="TCHW") - frames2, _, _ = read_video(video2_path, - pts_unit='sec', - output_format="TCHW") + frames1 = _read_video_frames(video1_path) + frames2 = _read_video_frames(video2_path) # Ensure same number of frames min_frames = min(frames1.shape[0], frames2.shape[0]) @@ -50,8 +77,8 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True): # Process each frame individually for i in range(min_frames): - img1 = frames1[i:i + 1] - img2 = frames2[i:i + 1] + img1 = frames1[i : i + 1] + img2 = frames2[i : i + 1] with torch.no_grad(): value = ms_ssim(img1, img2, data_range=1.0) if use_ms_ssim else ssim(img1, img2, data_range=1.0) @@ -71,14 +98,14 @@ def compute_video_ssim_torchvision(video1_path, video2_path, use_ms_ssim=True): return mean_ssim, min_ssim, max_ssim else: - print('No SSIM values calculated') + print("No SSIM values calculated") return 0, 0, 0 def compare_folders(reference_folder, generated_folder, use_ms_ssim=True): """ Compare videos with the same filename between reference_folder and generated_folder - + Example usage: results = compare_folders(reference_folder, generated_folder, args.use_ms_ssim) @@ -97,10 +124,8 @@ def compare_folders(reference_folder, generated_folder, use_ms_ssim=True): else: print("\nNo valid SSIM values to average") """ - - reference_videos = [ - f for f in os.listdir(reference_folder) if f.endswith('.mp4') - ] + + reference_videos = [f for f in os.listdir(reference_folder) if f.endswith(".mp4")] results = {} @@ -111,27 +136,23 @@ def compare_folders(reference_folder, generated_folder, use_ms_ssim=True): if os.path.exists(gen_path): print(f"\nComparing {video_name}...") try: - ssim_value = compute_video_ssim_torchvision( - ref_path, gen_path, use_ms_ssim) + ssim_value = compute_video_ssim_torchvision(ref_path, gen_path, use_ms_ssim) results[video_name] = ssim_value except Exception as e: print(f"Error comparing {video_name}: {e}") results[video_name] = None else: - print( - f"\nSkipping {video_name} - no matching file in generated folder" - ) + print(f"\nSkipping {video_name} - no matching file in generated folder") return results -def write_ssim_results(output_dir, ssim_values, reference_path, generated_path, - num_inference_steps, prompt): + +def write_ssim_results(output_dir, ssim_values, reference_path, generated_path, num_inference_steps, prompt): """ Write SSIM results to a JSON file in the same directory as the generated videos. """ try: - logger.info( - f"Attempting to write SSIM results to directory: {output_dir}") + logger.info(f"Attempting to write SSIM results to directory: {output_dir}") if not os.path.exists(output_dir): os.makedirs(output_dir, exist_ok=True) @@ -144,20 +165,17 @@ def write_ssim_results(output_dir, ssim_values, reference_path, generated_path, "max_ssim": max_ssim, "reference_video": reference_path, "generated_video": generated_path, - "parameters": { - "num_inference_steps": num_inference_steps, - "prompt": prompt - } + "parameters": {"num_inference_steps": num_inference_steps, "prompt": prompt}, } test_name = f"steps{num_inference_steps}_{prompt[:100]}" result_file = os.path.join(output_dir, f"{test_name}_ssim.json") logger.info(f"Writing JSON results to: {result_file}") - with open(result_file, 'w') as f: + with open(result_file, "w") as f: json.dump(result, f, indent=2) logger.info(f"SSIM results written to {result_file}") return True except Exception as e: logger.error(f"ERROR writing SSIM results: {str(e)}") - return False \ No newline at end of file + return False