scripts/download_checkpoints.py

#!/usr/bin/env python3
"""CLI: download the SAM 3.1 and TAPNext++ checkpoints this pipeline needs.

Both land in the standard Hugging Face cache (``~/.cache/huggingface``) via
``huggingface_hub`` rather than being copied into the repo -- ``tapnextpp_512
.ckpt`` alone is 2.53 GB, and re-downloading (or duplicating) it per-run or
per-clone would be wasteful when the cache already dedupes by content hash.
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402

logger = get_logger(__name__)

SAM3_REPO_ID = "facebook/sam3.1"
#: The weight file is required; the rest are best-effort (their exact set has
#: varied across SAM 3.x releases).
SAM3_REQUIRED_FILE = "sam3.1_multiplex.pt"
SAM3_OPTIONAL_FILES = ("config.json", "tokenizer.json", "tokenizer_config.json")

TAPNEXTPP_HF_REPO_ID = "google/tapnet"
TAPNEXTPP_FILENAME = "tapnextpp_512.ckpt"
TAPNEXTPP_GCS_URL = "https://storage.googleapis.com/gresearch/tapnextpp/tapnextpp_512.ckpt"

_HF_TOKEN_HELP = (
    "Hugging Face token with access to the gated facebook/sam3.1 repo. Falls "
    "back to the HF_TOKEN environment variable if not passed -- set that "
    "instead of passing a token on the command line (it avoids the token "
    "ending up in shell history / process listings). This project's HF "
    "account has a token with sam3.1 access; export it as HF_TOKEN rather "
    "than hardcoding it anywhere in source."
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--hf-token", default=None, help=_HF_TOKEN_HELP)
    parser.add_argument(
        "--checkpoints-dir",
        type=Path,
        default=Path("checkpoints"),
        help="only used for the GCS fallback download of tapnextpp; HF downloads use the HF cache",
    )
    parser.add_argument("--skip-sam3", action="store_true")
    parser.add_argument("--skip-tapnext", action="store_true")
    parser.add_argument(
        "--tapnext-source",
        choices=["gcs", "hf"],
        default="gcs",
        help="tapnextpp_512.ckpt is available both from its original GCS bucket (no auth) "
        "and from the google/tapnet HF repo",
    )
    parser.add_argument("--log-level", default="INFO")
    return parser.parse_args()


def _resolve_token(cli_token: str | None) -> str | None:
    return cli_token or os.environ.get("HF_TOKEN")


def download_sam3(token: str | None) -> None:
    from huggingface_hub import hf_hub_download
    from huggingface_hub.utils import HfHubHTTPError

    path = hf_hub_download(repo_id=SAM3_REPO_ID, filename=SAM3_REQUIRED_FILE, token=token)
    logger.info("sam3.1 %s -> %s", SAM3_REQUIRED_FILE, path)

    for filename in SAM3_OPTIONAL_FILES:
        try:
            path = hf_hub_download(repo_id=SAM3_REPO_ID, filename=filename, token=token)
            logger.info("sam3.1 %s -> %s", filename, path)
        except HfHubHTTPError as exc:
            logger.warning("sam3.1 %s not found in %s, skipping: %s", filename, SAM3_REPO_ID, exc)


def download_tapnextpp(source: str, checkpoints_dir: Path) -> None:
    if source == "hf":
        from huggingface_hub import hf_hub_download

        path = hf_hub_download(repo_id=TAPNEXTPP_HF_REPO_ID, filename=TAPNEXTPP_FILENAME)
        logger.info("tapnextpp -> %s", path)
        return

    from fpgm.utils.io import download_with_resume, ensure_dir

    dest = ensure_dir(checkpoints_dir) / TAPNEXTPP_FILENAME
    download_with_resume(TAPNEXTPP_GCS_URL, dest)
    logger.info("tapnextpp -> %s", dest)


def main() -> int:
    args = parse_args()
    setup_logging(args.log_level)
    token = _resolve_token(args.hf_token)

    if not args.skip_sam3:
        if token is None:
            logger.warning(
                "no HF token given (pass --hf-token or set HF_TOKEN); "
                "facebook/sam3.1 is gated and this will likely fail"
            )
        download_sam3(token)

    if not args.skip_tapnext:
        download_tapnextpp(args.tapnext_source, args.checkpoints_dir)

    return 0


if __name__ == "__main__":
    raise SystemExit(main())