hao-ai-lab-FastVideo/fastvideo/tests/ssim/reference_videos_cli.py

# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import argparse
import contextlib
import importlib.util
import os
import shutil
import sys
import tempfile
from collections.abc import Iterable, Sequence
from pathlib import Path
from types import ModuleType


def _load_envs_registry() -> ModuleType:
    """Load fastvideo/envs.py as a standalone module, without importing the fastvideo package.

    Importing fastvideo.envs runs fastvideo/__init__.py, which imports torch and the
    generation stack. This CLI runs as a plain script that needs only the standard
    library, so it loads the registry file by path.
    """
    registry_path = Path(__file__).resolve().parents[2] / "envs.py"
    spec = importlib.util.spec_from_file_location("_fastvideo_envs_registry", registry_path)
    assert spec is not None and spec.loader is not None
    registry = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(registry)
    return registry


envs = _load_envs_registry()

VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv")
# Additional artefact types stored under the same reference folders. Latent
# tensors (`.pt`) back the latent-slice regression tests used by flaky
# pixel-space models (e.g. LTX-2 distilled). Keeping them in the same upload
# flow means seeding a new test only requires one HF round-trip.
LATENT_EXTENSIONS = (".pt", )
REFERENCE_EXTENSIONS = VIDEO_EXTENSIONS + LATENT_EXTENSIONS + (".png", )
HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN")
HF_REPO_ENV_KEY = "FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO"
HF_REPO_TYPE_ENV_KEY = "FASTVIDEO_TEST_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"
DRAFTS_DIRNAME = "drafts"
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 envs.FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO.get()


def _default_repo_type() -> str:
    return envs.FASTVIDEO_TEST_SSIM_REFERENCE_HF_REPO_TYPE.get()


def _iter_reference_files(root: Path) -> Iterable[Path]:
    """Yield supported references under `root`.

    Used by copy-local and the "has local references" marker probe so that
    image- and latent-only tests (no mp4) still satisfy readiness checks.
    """
    for path in root.rglob("*"):
        if path.is_file() and path.suffix.lower() in REFERENCE_EXTENSIONS:
            yield path


def _discover_reference_dirs(base_dir: Path) -> list[Path]:
    if not base_dir.exists():
        return []
    dirs = [p for p in base_dir.iterdir() if p.is_dir() and p.name.endswith("_reference_videos")]
    return sorted(dirs)


@contextlib.contextmanager
def _exclusive_download_lock(lock_path: Path):
    """Best-effort cross-process lock for reference video download/copy."""
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    with lock_path.open("a", encoding="utf-8") as lock_file:
        if os.name == "posix":
            import fcntl

            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
        try:
            yield
        finally:
            if os.name == "posix":
                import fcntl

                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)


def _merge_downloaded_reference_content(
    staging_dir: Path,
    local_dir: Path,
) -> int:
    copied_roots = 0

    tier_root = staging_dir / REFERENCE_VIDEOS_DIRNAME
    if tier_root.exists():
        dst_tier_root = local_dir / REFERENCE_VIDEOS_DIRNAME
        shutil.copytree(tier_root, dst_tier_root, dirs_exist_ok=True)
        copied_roots += 1

    # Legacy default layout support: <staging>/<GPU>_reference_videos/...
    legacy_reference_dirs = _discover_reference_dirs(staging_dir)
    for legacy_dir in legacy_reference_dirs:
        dst_legacy_dir = local_dir / legacy_dir.name
        shutil.copytree(legacy_dir, dst_legacy_dir, dirs_exist_ok=True)
        copied_roots += 1

    return copied_roots


def _resolve_device_reference_folder(
    *,
    explicit_device_folder: str | None,
    reference_dir: Path | None,
) -> str:
    if explicit_device_folder and explicit_device_folder.strip():
        return explicit_device_folder.strip()
    if reference_dir is not None and reference_dir.name.endswith("_reference_videos"):
        return reference_dir.name
    return DEFAULT_DEVICE_REFERENCE_FOLDER


def _reference_tier_root(base_dir: Path, quality_tier: str) -> Path:
    return base_dir / REFERENCE_VIDEOS_DIRNAME / quality_tier


def _resolve_quality_tiers(quality_tier: str) -> list[str]:
    if quality_tier == ALL_OUTPUT_QUALITY_TIERS:
        return list(QUALITY_TIERS)
    if quality_tier not in QUALITY_TIERS:
        raise ValueError(f"Unsupported quality tier: {quality_tier}")
    return [quality_tier]


def _discover_reference_dirs_for_tier(
    base_dir: Path,
    quality_tier: str,
) -> list[Path]:
    tier_root = _reference_tier_root(base_dir, quality_tier)
    tier_dirs = _discover_reference_dirs(tier_root)
    if tier_dirs:
        return tier_dirs
    if quality_tier == DEFAULT_OUTPUT_QUALITY_TIER:
        # Keep supporting legacy default refs at <ssim_dir>/<GPU>_reference_videos.
        return _discover_reference_dirs(base_dir)
    return []


def _has_local_reference_videos(base_dir: Path, quality_tier: str) -> bool:
    marker_path = base_dir / REFERENCE_VIDEOS_DIRNAME / quality_tier / f".download_complete_{quality_tier}"
    if not marker_path.exists():
        return False
    tier_root = _reference_tier_root(base_dir, quality_tier)
    for ref_dir in _discover_reference_dirs(tier_root):
        for _ in _iter_reference_files(ref_dir):
            return True
    return False


def _get_hf_token() -> str | None:
    for key in HF_TOKEN_ENV_KEYS:
        value = os.environ.get(key, "").strip()
        if value:
            return value

    return None


def _load_hf_sdk():
    try:
        from huggingface_hub import HfApi, snapshot_download
    except ImportError as exc:
        raise RuntimeError(
            "huggingface_hub is required for download/upload.\nInstall with: uv 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 ref_file in _iter_reference_files(generated_dir):
        rel = ref_file.relative_to(generated_dir)
        dst = reference_dir / rel
        if dry_run:
            print(f"[dry-run] {ref_file} -> {dst}")
        else:
            dst.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(ref_file, dst)
            print(f"Copied: {rel}")
        copied += 1
    return copied


def download_reference_videos(
    *,
    repo_id: str,
    repo_type: str,
    local_dir: Path,
    quality_tiers: Sequence[str] | None = None,
    device_folders: Sequence[str] | None = None,
    token: str | None = None,
) -> None:
    _, snapshot_download = _load_hf_sdk()

    local_dir.mkdir(parents=True, exist_ok=True)
    tiers = tuple(quality_tiers or QUALITY_TIERS)
    for quality_tier in tiers:
        if quality_tier not in QUALITY_TIERS:
            raise ValueError(f"Unsupported quality tier: {quality_tier}")
    allow_patterns: list[str] = []
    for quality_tier in tiers:
        tier_prefix = f"{REFERENCE_VIDEOS_DIRNAME}/{quality_tier}"
        if device_folders:
            allow_patterns.extend(f"{tier_prefix}/{folder}/**" for folder in device_folders)
        else:
            allow_patterns.append(f"{tier_prefix}/*_reference_videos/**")

        # Keep compatibility for older HF layout where default refs lived at
        # repo root as <GPU>_reference_videos.
        if quality_tier == DEFAULT_OUTPUT_QUALITY_TIER:
            if device_folders:
                allow_patterns.extend(f"{folder}/**" for folder in device_folders)
            else:
                allow_patterns.append("*_reference_videos/**")

    # Work around long-path issues in some huggingface_hub versions:
    # avoid `local_dir` mode and download into a short cache dir instead,
    # then merge files into local_dir.
    with tempfile.TemporaryDirectory(prefix="fv2-ssim-cache-") as tmp_cache_dir:
        snapshot_path = snapshot_download(
            repo_id=repo_id,
            repo_type=repo_type,
            cache_dir=tmp_cache_dir,
            token=token,
            allow_patterns=allow_patterns,
        )
        staging_dir = Path(snapshot_path)
        copied_roots = _merge_downloaded_reference_content(
            staging_dir=staging_dir,
            local_dir=local_dir,
        )
        if copied_roots == 0:
            raise RuntimeError(
                "HF download completed but no reference_videos content was found in downloaded artifacts.")
        for quality_tier in tiers:
            marker_path = local_dir / REFERENCE_VIDEOS_DIRNAME / quality_tier / f".download_complete_{quality_tier}"
            marker_path.parent.mkdir(parents=True, exist_ok=True)
            marker_path.touch()


def upload_reference_videos(
    *,
    repo_id: str,
    repo_type: str,
    reference_dirs_by_tier: Sequence[tuple[str, Path]],
    token: str,
    private: bool,
    model_id: str | None = None,
    force: bool = False,
) -> 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,
    )

    try:
        existing_repo_files = set(api.list_repo_files(repo_id=repo_id, repo_type=repo_type))
    except Exception:
        # Fresh repo or list failure — treat as empty so upload can proceed.
        existing_repo_files = set()

    for quality_tier, reference_dir in reference_dirs_by_tier:
        if not reference_dir.exists():
            raise FileNotFoundError(f"Reference directory not found: {reference_dir}")
        base_in_repo = f"{REFERENCE_VIDEOS_DIRNAME}/{quality_tier}/{reference_dir.name}"
        if model_id:
            folder_path = reference_dir / model_id
            if not folder_path.exists():
                raise FileNotFoundError(f"Model subfolder not found for upload: {folder_path}")
            path_in_repo = f"{base_in_repo}/{model_id}"
        else:
            folder_path = reference_dir
            path_in_repo = base_in_repo

        conflicts = sorted(f for f in existing_repo_files if f.startswith(f"{path_in_repo}/") or f == path_in_repo)
        if conflicts and not force:
            preview = "\n".join(f"  - {c}" for c in conflicts[:10])
            more = f"\n  ... and {len(conflicts) - 10} more" if len(conflicts) > 10 else ""
            raise RuntimeError(f"Refusing to overwrite existing HF files under {path_in_repo} "
                               f"({len(conflicts)} file(s) already present):\n{preview}{more}\n"
                               f"Re-run with --force to overwrite.")

        target_desc = f"{reference_dir.name}/{model_id}" if model_id else reference_dir.name
        print(f"Uploading {target_desc} ({quality_tier}) to {repo_id}/{path_in_repo} ...")
        api.upload_folder(
            repo_id=repo_id,
            repo_type=repo_type,
            folder_path=str(folder_path),
            path_in_repo=path_in_repo,
            token=token,
        )


def _reference_folder_repo_relative(reference_folder: Path, base_dir: Path | None = None) -> Path:
    resolved_reference_folder = reference_folder.resolve()
    resolved_base_dir = (base_dir or _ssim_dir()).resolve()
    tiered_root = resolved_base_dir / REFERENCE_VIDEOS_DIRNAME
    try:
        relative = resolved_reference_folder.relative_to(tiered_root)
    except ValueError:
        try:
            legacy_relative = resolved_reference_folder.relative_to(resolved_base_dir)
        except ValueError as exc:
            raise RuntimeError(
                f"Reference folder {reference_folder} is not under SSIM directory {resolved_base_dir}") from exc
        if not legacy_relative.parts or not legacy_relative.parts[0].endswith("_reference_videos"):
            raise RuntimeError(
                f"Reference folder {reference_folder} does not match reference_videos/<tier>/<device>/<model>/<backend>"
            )
        relative = Path(DEFAULT_OUTPUT_QUALITY_TIER, *legacy_relative.parts)

    if len(relative.parts) < 4:
        raise RuntimeError(f"Reference folder {reference_folder} does not include <quality>/<device>/<model>/<backend>")
    return relative


def upload_draft_reference_artifact(
    *,
    repo_id: str,
    repo_type: str,
    generated_artifact_path: Path,
    reference_folder: Path,
    token: str | None = None,
    base_dir: Path | None = None,
    private: bool = False,
) -> str:
    resolved_token = token or _get_hf_token()
    if resolved_token is None:
        raise RuntimeError("Hugging Face API key is required for SSIM bootstrap draft upload. "
                           "Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN.")
    if not generated_artifact_path.exists():
        raise FileNotFoundError(f"Generated artifact not found: {generated_artifact_path}")

    relative_reference_folder = _reference_folder_repo_relative(
        reference_folder=reference_folder,
        base_dir=base_dir,
    )
    path_in_repo = f"{DRAFTS_DIRNAME}/{relative_reference_folder.as_posix()}/{generated_artifact_path.name}"

    HfApi, _ = _load_hf_sdk()
    api = HfApi(token=resolved_token)
    api.create_repo(
        repo_id=repo_id,
        repo_type=repo_type,
        private=private,
        exist_ok=True,
    )
    print(f"Uploading SSIM draft reference to {repo_id}/{path_in_repo} ...")
    api.upload_file(
        repo_id=repo_id,
        repo_type=repo_type,
        path_or_fileobj=str(generated_artifact_path),
        path_in_repo=path_in_repo,
        token=resolved_token,
    )
    print("Draft reference uploaded. Review before promotion.")
    return path_in_repo


def promote_draft_references(
    *,
    repo_id: str,
    repo_type: str,
    quality_tier: str,
    device_folder: str,
    model_id: str,
    token: str,
    attention_backend: str | None = None,
    private: bool = False,
    force: bool = False,
) -> None:
    HfApi, snapshot_download = _load_hf_sdk()
    api = HfApi(token=token)
    api.create_repo(
        repo_id=repo_id,
        repo_type=repo_type,
        private=private,
        exist_ok=True,
    )
    draft_prefix_parts = [
        DRAFTS_DIRNAME,
        quality_tier,
        device_folder,
        model_id,
    ]
    reference_prefix_parts = [
        REFERENCE_VIDEOS_DIRNAME,
        quality_tier,
        device_folder,
        model_id,
    ]
    if attention_backend:
        draft_prefix_parts.append(attention_backend)
        reference_prefix_parts.append(attention_backend)
    draft_prefix = "/".join(draft_prefix_parts)
    reference_prefix = "/".join(reference_prefix_parts)

    try:
        existing_repo_files = set(api.list_repo_files(repo_id=repo_id, repo_type=repo_type))
    except Exception:
        # Fresh repo or list failure — treat as empty so promotion can report
        # a clear "no drafts found" error instead of surfacing HF internals.
        existing_repo_files = set()
    draft_files = sorted(f for f in existing_repo_files if f.startswith(f"{draft_prefix}/"))
    if not draft_files:
        raise RuntimeError(f"No draft references found under {repo_id}/{draft_prefix}")

    conflicts = sorted(f for f in existing_repo_files if f.startswith(f"{reference_prefix}/") or f == reference_prefix)
    if conflicts and not force:
        preview = "\n".join(f"  - {c}" for c in conflicts[:10])
        more = f"\n  ... and {len(conflicts) - 10} more" if len(conflicts) > 10 else ""
        raise RuntimeError(f"Refusing to overwrite existing HF files under {reference_prefix} "
                           f"({len(conflicts)} file(s) already present):\n{preview}{more}\n"
                           f"Re-run with --force to overwrite.")

    with tempfile.TemporaryDirectory(prefix="fv2-ssim-draft-") as tmp_cache_dir:
        snapshot_path = Path(
            snapshot_download(
                repo_id=repo_id,
                repo_type=repo_type,
                cache_dir=tmp_cache_dir,
                token=token,
                allow_patterns=[f"{draft_prefix}/**"],
            ))
        draft_root = snapshot_path / draft_prefix
        if not draft_root.exists():
            raise RuntimeError(f"Downloaded draft path is missing: {draft_root}")
        print(f"Promoting {len(draft_files)} draft reference file(s) to {repo_id}/{reference_prefix} ...")
        api.upload_folder(
            repo_id=repo_id,
            repo_type=repo_type,
            folder_path=str(draft_root),
            path_in_repo=reference_prefix,
            token=token,
        )


def _resolve_upload_reference_dirs(
    *,
    base_dir: Path,
    quality_tier: str,
    explicit_reference_dirs: Sequence[str],
    device_folders: Sequence[str] | None = None,
) -> list[tuple[str, Path]]:
    tiers = _resolve_quality_tiers(quality_tier)
    selected_device_folders = set(device_folders or [])

    if explicit_reference_dirs:
        if len(tiers) != 1:
            raise RuntimeError("--reference-dir requires a single --quality-tier (default or full_quality).")
        resolved_tier = tiers[0]
        return [(resolved_tier, Path(p).resolve()) for p in explicit_reference_dirs]

    resolved: list[tuple[str, Path]] = []
    for tier in tiers:
        for reference_dir in _discover_reference_dirs_for_tier(base_dir, tier):
            if selected_device_folders and reference_dir.name not in selected_device_folders:
                continue
            resolved.append((tier, reference_dir.resolve()))
    return resolved


def ensure_reference_videos_available(
    *,
    local_dir: Path | None = None,
    repo_id: str | None = None,
    repo_type: str | None = None,
    quality_tier: str = DEFAULT_OUTPUT_QUALITY_TIER,
) -> bool:
    """Return True if downloaded from HF, False if already present locally."""
    if quality_tier not in QUALITY_TIERS:
        raise ValueError(f"Unsupported quality tier: {quality_tier}")
    target_dir = local_dir or _ssim_dir()
    lock_path = target_dir / ".reference_videos_download.lock"
    with _exclusive_download_lock(lock_path):
        if _has_local_reference_videos(target_dir, quality_tier):
            print(f"Reference videos ({quality_tier}) already available at {target_dir}")
            return False

        resolved_repo_id = repo_id or _default_repo_id()
        resolved_repo_type = repo_type or _default_repo_type()
        if not resolved_repo_id:
            raise RuntimeError(
                f"No local reference videos found and no HF repo configured.\nSet {HF_REPO_ENV_KEY} or pass --repo-id.")

        print(f"Repo ID: {resolved_repo_id}")
        print(f"Quality tier: {quality_tier}")
        print(f"No local {quality_tier} reference videos found under {target_dir}. Starting download...")
        try:
            download_reference_videos(
                repo_id=resolved_repo_id,
                repo_type=resolved_repo_type,
                local_dir=target_dir,
                quality_tiers=[quality_tier],
            )
            print(f"Download completed for {quality_tier} reference videos.")
        except Exception as exc:
            print(f"ERROR: Failed to download {quality_tier} reference videos from {resolved_repo_id}.")
            print(f"Suggested command to retry: "
                  f"python fastvideo/tests/ssim/reference_videos_cli.py download "
                  f"--quality-tier {quality_tier}")
            raise

        if not _has_local_reference_videos(target_dir, quality_tier):
            raise RuntimeError(f"HF download completed but no {quality_tier} *_reference_videos content found.")
        return True


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="reference_videos_cli",
        description=("Manage SSIM reference videos locally and on Hugging Face.\n"
                     "This tool can copy generated videos into reference folders,\n"
                     "download references from a public HF repo, and upload updates."),
        epilog=("Examples:\n"
                "  # Copy default generated videos into local L40S refs\n"
                "  python fastvideo/tests/ssim/reference_videos_cli.py copy-local \\\n"
                "    --quality-tier default \\\n"
                "    --device-folder L40S_reference_videos \\\n"
                "    --reference-dir fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos\n\n"
                "  # Download both default and full-quality references from HF\n"
                "  python fastvideo/tests/ssim/reference_videos_cli.py download \\\n"
                "    --repo-id FastVideo/ssim-reference-videos\n\n"
                "  # Upload both tiers and all GPU folders (fails if token is unset)\n"
                "  python fastvideo/tests/ssim/reference_videos_cli.py upload \\\n"
                "    --repo-id FastVideo/ssim-reference-videos"),
        formatter_class=argparse.RawTextHelpFormatter,
    )

    subparsers = parser.add_subparsers(dest="command", required=True)

    copy_parser = subparsers.add_parser(
        "copy-local",
        help="Copy generated videos into a local reference folder.",
    )
    copy_parser.add_argument(
        "--quality-tier",
        choices=(DEFAULT_OUTPUT_QUALITY_TIER, FULL_OUTPUT_QUALITY_TIER),
        default=DEFAULT_OUTPUT_QUALITY_TIER,
        help="Generated quality tier to copy from.",
    )
    copy_parser.add_argument(
        "--generated-dir",
        type=Path,
        default=None,
        help=("Source generated directory. Default: <ssim_dir>/generated_videos/<quality-tier>/<device-folder>"),
    )
    copy_parser.add_argument(
        "--device-folder",
        default=None,
        help=("GPU folder name (e.g., H200_reference_videos) used to build default generated/reference paths."),
    )
    copy_parser.add_argument(
        "--reference-dir",
        type=Path,
        default=None,
        help="Destination reference directory.",
    )
    copy_parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print copy operations without writing files.",
    )

    download_parser = subparsers.add_parser(
        "download",
        help="Download reference videos from Hugging Face.",
    )
    download_parser.add_argument(
        "--repo-id",
        default=_default_repo_id(),
        help=f"HF repo id. Default: env {HF_REPO_ENV_KEY} or {DEFAULT_REPO_ID}",
    )
    download_parser.add_argument(
        "--repo-type",
        default=_default_repo_type(),
        help=f"HF repo type. Default: env {HF_REPO_TYPE_ENV_KEY} or {DEFAULT_REPO_TYPE}",
    )
    download_parser.add_argument(
        "--local-dir",
        type=Path,
        default=_ssim_dir(),
        help=("Local SSIM directory where references are stored under "
              "reference_videos/<quality-tier>/<GPU>_reference_videos."),
    )
    download_parser.add_argument(
        "--quality-tier",
        choices=(
            DEFAULT_OUTPUT_QUALITY_TIER,
            FULL_OUTPUT_QUALITY_TIER,
            ALL_OUTPUT_QUALITY_TIERS,
        ),
        default=ALL_OUTPUT_QUALITY_TIERS,
        help="Quality tier to download (default: all).",
    )
    download_parser.add_argument(
        "--device-folder",
        action="append",
        default=[],
        help="Specific reference folder to fetch (repeatable), e.g. L40S_reference_videos",
    )

    upload_parser = subparsers.add_parser(
        "upload",
        help="Upload local reference folders to Hugging Face.",
    )
    upload_parser.add_argument(
        "--repo-id",
        default=_default_repo_id(),
        help=f"HF repo id. Default: env {HF_REPO_ENV_KEY} or {DEFAULT_REPO_ID}",
    )
    upload_parser.add_argument(
        "--repo-type",
        default=_default_repo_type(),
        help=f"HF repo type. Default: env {HF_REPO_TYPE_ENV_KEY} or {DEFAULT_REPO_TYPE}",
    )
    upload_parser.add_argument(
        "--reference-dir",
        action="append",
        default=[],
        help=("Reference directory to upload (repeatable). "
              "Requires single --quality-tier. "
              "Default: discover folders under --base-dir for chosen tier(s)."),
    )
    upload_parser.add_argument(
        "--quality-tier",
        choices=(
            DEFAULT_OUTPUT_QUALITY_TIER,
            FULL_OUTPUT_QUALITY_TIER,
            ALL_OUTPUT_QUALITY_TIERS,
        ),
        default=ALL_OUTPUT_QUALITY_TIERS,
        help="Quality tier to upload (default: all).",
    )
    upload_parser.add_argument(
        "--device-folder",
        action="append",
        default=[],
        help=("Specific GPU reference folder to upload (repeatable), e.g. H200_reference_videos."),
    )
    upload_parser.add_argument(
        "--base-dir",
        type=Path,
        default=_ssim_dir(),
        help=("Base SSIM directory that contains reference_videos/<quality-tier>/<GPU>_reference_videos."),
    )
    upload_parser.add_argument(
        "--private",
        action="store_true",
        help="Create/use a private repo instead of public.",
    )
    upload_parser.add_argument(
        "--model-id",
        default=None,
        help=("Restrict upload to a single model subfolder "
              "(reference_videos/<tier>/<device>/<model_id>). "
              "Use when seeding references for a single new test."),
    )
    upload_parser.add_argument(
        "--force",
        action="store_true",
        help=("Allow overwriting files that already exist at the target path on "
              "Hugging Face. Off by default so seeding a new test cannot "
              "clobber existing references."),
    )

    promote_parser = subparsers.add_parser(
        "promote-draft",
        help="Promote reviewed draft references into the canonical HF reference layout.",
    )
    promote_parser.add_argument(
        "--repo-id",
        default=_default_repo_id(),
        help=f"HF repo id. Default: env {HF_REPO_ENV_KEY} or {DEFAULT_REPO_ID}",
    )
    promote_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}",
    )
    promote_parser.add_argument(
        "--quality-tier",
        choices=(
            DEFAULT_OUTPUT_QUALITY_TIER,
            FULL_OUTPUT_QUALITY_TIER,
        ),
        default=DEFAULT_OUTPUT_QUALITY_TIER,
        help="Quality tier to promote.",
    )
    promote_parser.add_argument(
        "--device-folder",
        default=DEFAULT_DEVICE_REFERENCE_FOLDER,
        help="GPU reference folder to promote, e.g. L40S_reference_videos.",
    )
    promote_parser.add_argument(
        "--model-id",
        required=True,
        help="Model subfolder to promote from drafts.",
    )
    promote_parser.add_argument(
        "--attention-backend",
        default=None,
        help="Optional backend subfolder to promote. Default promotes all backends under the model.",
    )
    promote_parser.add_argument(
        "--private",
        action="store_true",
        help="Create/use a private repo instead of public.",
    )
    promote_parser.add_argument(
        "--force",
        action="store_true",
        help="Allow overwriting existing canonical references. Off by default.",
    )

    ensure_parser = subparsers.add_parser(
        "ensure",
        help="Download refs only when the selected local quality tier is missing.",
    )
    ensure_parser.add_argument(
        "--repo-id",
        default=_default_repo_id(),
        help=f"HF repo id. Default: env {HF_REPO_ENV_KEY} or {DEFAULT_REPO_ID}",
    )
    ensure_parser.add_argument(
        "--repo-type",
        default=_default_repo_type(),
        help=f"HF repo type. Default: env {HF_REPO_TYPE_ENV_KEY} or {DEFAULT_REPO_TYPE}",
    )
    ensure_parser.add_argument(
        "--local-dir",
        type=Path,
        default=_ssim_dir(),
        help=("Local SSIM directory that should contain reference_videos/<quality-tier>/<GPU>_reference_videos."),
    )
    ensure_parser.add_argument(
        "--quality-tier",
        choices=(
            DEFAULT_OUTPUT_QUALITY_TIER,
            FULL_OUTPUT_QUALITY_TIER,
        ),
        default=DEFAULT_OUTPUT_QUALITY_TIER,
        help="Quality tier to ensure locally.",
    )

    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)

    if args.command == "copy-local":
        device_folder = _resolve_device_reference_folder(
            explicit_device_folder=args.device_folder,
            reference_dir=args.reference_dir,
        )
        generated_dir = args.generated_dir
        if generated_dir is None:
            tier_root = _ssim_dir() / "generated_videos" / args.quality_tier
            tier_device_dir = tier_root / device_folder
            if tier_device_dir.exists():
                generated_dir = tier_device_dir
            else:
                # Backward compatibility for pre-device-folder layout.
                generated_dir = tier_root
        reference_dir = (args.reference_dir if args.reference_dir is not None else
                         (_ssim_dir() / REFERENCE_VIDEOS_DIRNAME / args.quality_tier / device_folder))
        copied = copy_generated_to_reference(
            generated_dir=generated_dir,
            reference_dir=reference_dir,
            dry_run=args.dry_run,
        )
        print(f"Done. {'Would copy' if args.dry_run else 'Copied'} {copied} video files.")
        if not args.dry_run and copied > 0:
            marker = (_ssim_dir() / REFERENCE_VIDEOS_DIRNAME / args.quality_tier /
                      f".download_complete_{args.quality_tier}")
            marker.parent.mkdir(parents=True, exist_ok=True)
            marker.touch()
        return 0

    if args.command == "download":
        quality_tiers = _resolve_quality_tiers(args.quality_tier)
        download_reference_videos(
            repo_id=args.repo_id,
            repo_type=args.repo_type,
            local_dir=args.local_dir,
            quality_tiers=quality_tiers,
            device_folders=args.device_folder or None,
        )
        print("Download complete for quality tier(s): " + ", ".join(quality_tiers) + ".")
        return 0

    if args.command == "upload":
        token = _get_hf_token()
        if token is None:
            raise RuntimeError(
                "Hugging Face API key is required for upload. Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN.")

        reference_dirs_by_tier = _resolve_upload_reference_dirs(
            base_dir=args.base_dir.resolve(),
            quality_tier=args.quality_tier,
            explicit_reference_dirs=args.reference_dir,
            device_folders=args.device_folder or None,
        )
        if not reference_dirs_by_tier:
            raise RuntimeError(
                f"No *_reference_videos directories found for the selected quality tier(s) under {args.base_dir}")

        upload_reference_videos(
            repo_id=args.repo_id,
            repo_type=args.repo_type,
            reference_dirs_by_tier=reference_dirs_by_tier,
            token=token,
            private=args.private,
            model_id=args.model_id,
            force=args.force,
        )
        print("Upload complete.")
        return 0

    if args.command == "promote-draft":
        token = _get_hf_token()
        if token is None:
            raise RuntimeError("Hugging Face API key is required for promote-draft. "
                               "Set HF_API_KEY, HUGGINGFACE_HUB_TOKEN, or HF_TOKEN.")
        promote_draft_references(
            repo_id=args.repo_id,
            repo_type=args.repo_type,
            quality_tier=args.quality_tier,
            device_folder=args.device_folder,
            model_id=args.model_id,
            attention_backend=args.attention_backend,
            token=token,
            private=args.private,
            force=args.force,
        )
        print("Promotion 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)