openpi/scripts/upload_to_hf.py

#!/usr/bin/env python3
"""Upload a local folder to a private Hugging Face repo."""

from __future__ import annotations

import argparse
import os
import pathlib
import sys

from huggingface_hub import HfApi

DEFAULT_IGNORE_PATTERNS = (
    ".cache/**",
    ".git/**",
    "__pycache__/**",
    "*.pyc",
    "openpi/.git/**",
    "openpi/.venv/**",
    "openpi/checkpoints/**",
    "openpi/.pytest_cache/**",
    "openpi/.ruff_cache/**",
)


def _resolve_token() -> str | None:
    for key in ("HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HUGGING_FACE_HUB_TOKEN"):
        token = os.getenv(key)
        if token:
            return token
    return None


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--folder",
        default="/mnt/c/users/luiss/pi05tests",
        help="Local folder to upload.",
    )
    parser.add_argument(
        "--repo-id",
        default=None,
        help="Target HF repo id (namespace/name). If omitted, uses <username>/pi05tests-openpi-multiarm.",
    )
    parser.add_argument(
        "--commit-message",
        default="Upload pi05 multiarm workspace",
        help="Commit message for upload.",
    )
    parser.add_argument(
        "--include-checkpoints",
        action="store_true",
        help="If set, do not ignore openpi/checkpoints/**.",
    )
    args = parser.parse_args()

    folder = pathlib.Path(args.folder).resolve()
    if not folder.exists():
        raise FileNotFoundError(f"Folder does not exist: {folder}")

    token = _resolve_token()
    if not token:
        print(
            "No HF token found in env. Set one of: HF_TOKEN, HUGGINGFACE_HUB_TOKEN, HUGGING_FACE_HUB_TOKEN.",
            file=sys.stderr,
        )
        return 2

    api = HfApi(token=token)
    who = api.whoami(token=token)
    username = who["name"]
    repo_id = args.repo_id or f"{username}/pi05tests-openpi-multiarm"

    ignore_patterns = list(DEFAULT_IGNORE_PATTERNS)
    if args.include_checkpoints:
        ignore_patterns = [p for p in ignore_patterns if p != "openpi/checkpoints/**"]

    print(f"Uploading folder: {folder}")
    print(f"Repo: {repo_id} (private)")
    print(f"Ignore patterns: {ignore_patterns}")

    api.create_repo(
        repo_id=repo_id,
        private=True,
        exist_ok=True,
        repo_type="model",
        token=token,
    )
    api.upload_folder(
        folder_path=str(folder),
        repo_id=repo_id,
        repo_type="model",
        token=token,
        commit_message=args.commit_message,
        ignore_patterns=ignore_patterns,
    )
    print(f"Upload completed: https://huggingface.co/{repo_id}")
    return 0


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