scripts/publish_hf_gguf.py

"""Upload an Objectverse Diary GGUF model file to Hugging Face Hub."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any


def validate_gguf_file(gguf_file: Path) -> dict[str, object]:
    if not gguf_file.exists() or not gguf_file.is_file():
        raise FileNotFoundError(f"GGUF file does not exist: {gguf_file}")
    if gguf_file.suffix.lower() != ".gguf":
        raise ValueError(f"GGUF file must use .gguf suffix: {gguf_file}")
    size_bytes = gguf_file.stat().st_size
    if size_bytes <= 0:
        raise ValueError(f"GGUF file is empty: {gguf_file}")
    return {
        "gguf_file": str(gguf_file),
        "size_bytes": size_bytes,
    }


def upload_gguf(
    *,
    gguf_file: Path,
    repo_id: str,
    path_in_repo: str,
    private: bool,
    commit_message: str,
    dry_run: bool,
) -> dict[str, object]:
    summary = validate_gguf_file(gguf_file)
    summary.update(
        {
            "repo_id": repo_id,
            "path_in_repo": path_in_repo,
            "private": private,
            "commit_message": commit_message,
            "dry_run": dry_run,
        }
    )
    if dry_run:
        summary["uploaded"] = False
        return summary

    from huggingface_hub import HfApi

    api = HfApi()
    api.create_repo(repo_id=repo_id, repo_type="model", private=private, exist_ok=True)
    api.upload_file(
        path_or_fileobj=str(gguf_file),
        path_in_repo=path_in_repo,
        repo_id=repo_id,
        repo_type="model",
        commit_message=commit_message,
    )
    summary["uploaded"] = True
    summary["url"] = f"https://huggingface.co/{repo_id}/blob/main/{path_in_repo}"
    return summary


def _print_json(payload: dict[str, Any]) -> None:
    print(json.dumps(payload, indent=2, sort_keys=True), flush=True)


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gguf-file", type=Path, required=True)
    parser.add_argument("--repo-id", required=True)
    parser.add_argument("--path-in-repo", required=True)
    parser.add_argument("--private", action="store_true")
    parser.add_argument(
        "--commit-message",
        default="Upload Objectverse Diary GGUF model",
    )
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = _parse_args()
    _print_json(
        upload_gguf(
            gguf_file=args.gguf_file,
            repo_id=args.repo_id,
            path_in_repo=args.path_in_repo,
            private=args.private,
            commit_message=args.commit_message,
            dry_run=args.dry_run,
        )
    )


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        raise SystemExit(str(exc)) from exc