Blame for source/scripts/deploy_huggingface_space.py

#!/usr/bin/env python3
"""Build, check, and optionally publish the browser demo as a static HF Space."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import tempfile


REPO = Path(__file__).resolve().parents[1]
GRAPH_NAMES = ("encoder", "cross_cache", "decoder_step")
# ORT chooses JSEP or Asyncify from browser support. Keep both module/binary
# pairs; Vite's hashed binary asset alone does not cover wasmPaths=/ort/.
ORT_FILES = tuple(
    f"ort-wasm-simd-threaded.{variant}.{suffix}"
    for variant in ("jsep", "asyncify") for suffix in ("mjs", "wasm")
)


def sha256(path: Path) -> str:
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def copy_file(source: Path, target: Path) -> None:
    target.parent.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(source, target)


def package_space(web_dir: Path, graph_dir: Path, ort_dir: Path, output: Path) -> dict:
    """Copy a small allowlist; never upload the working checkout or audio cache."""
    manifest = json.loads((graph_dir / "manifest.json").read_text())
    if manifest["precision"] != "fp32":
        raise ValueError("The public Space requires the checked FP32 export")
    for name in GRAPH_NAMES:
        graph = graph_dir / f"{name}.onnx"
        record = manifest["graphs"][name]
        if graph.stat().st_size != record["bytes"] or sha256(graph) != record["sha256"]:
            raise ValueError(f"Graph does not match the export manifest: {graph}")

    if output.exists() and any(output.iterdir()):
        raise ValueError(f"Output must be empty: {output}")
    output.mkdir(parents=True, exist_ok=True)
    copy_file(web_dir / "local.html", output / "index.html")
    copy_file(web_dir / "local.html", output / "local.html")
    copy_file(web_dir / "favicon.png", output / "favicon.png")

    # The Vite manifest follows the local entry's imports and runtime assets.
    # This excludes compare/server pages and unreferenced public audio.
    build_manifest = json.loads((web_dir / ".vite/manifest.json").read_text())
    visited: set[str] = set()

    def copy_entry(key: str) -> None:
        if key in visited:
            return
        visited.add(key)
        entry = build_manifest[key]
        for field in ("imports", "dynamicImports"):
            for dependency in entry.get(field, []):
                copy_entry(dependency)
        files = [entry["file"], *entry.get("css", []), *entry.get("assets", [])]
        for relative in files:
            source = (web_dir / relative).resolve()
            if not source.is_relative_to(web_dir.resolve()):
                raise ValueError(f"Invalid Vite asset path: {relative}")
            copy_file(source, output / relative)

    copy_entry("local.html")
    manifest["checkpoint"]["path"] = Path(manifest["checkpoint"]["path"]).name
    (output / "models/fp32").mkdir(parents=True)
    (output / "models/fp32/manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    for name in GRAPH_NAMES:
        copy_file(graph_dir / f"{name}.onnx", output / f"models/fp32/{name}.onnx")
    for filename in ORT_FILES:
        copy_file(ort_dir / filename, output / "ort" / filename)

    # Only publish the licensed catalog's exact audio/notes, with graph and
    # file hashes checked. Never copy the private local examples directory.
    catalog = json.loads((web_dir / "examples/catalog.json").read_text())
    public_files = ["examples/catalog.json", "examples/CREDITS.md", "soundfonts/GeneralUser_GS.sf2", "soundfonts/LICENSE.txt"]
    for example in catalog["examples"]:
        for field, digest_field in (("audio_url", "audio_sha256"), ("notes_url", "notes_sha256")):
            relative = example[field].lstrip("/")
            source = (web_dir / relative).resolve()
            if not source.is_relative_to((web_dir / "examples").resolve()):
                raise ValueError(f"Invalid example asset: {relative}")
            if sha256(source) != example[digest_field]:
                raise ValueError(f"Example asset hash mismatch: {relative}")
            public_files.append(relative)
        prediction = json.loads((web_dir / example["notes_url"].lstrip("/")).read_text())
        if prediction["checkpoint_sha256"] != manifest["checkpoint"]["sha256"]:
            raise ValueError(f"Example prediction uses another model: {example['id']}")
    for relative in public_files:
        copy_file(web_dir / relative, output / relative)

    copy_file(REPO / "deploy/huggingface/README.md", output / "README.md")
    # Keep the exact app source next to the build for review and reproduction.
    shutil.copytree(REPO / "demo/web/src", output / "source/demo/web/src")
    for filename in ("package.json", "pnpm-lock.yaml", "vite.config.ts", "tsconfig.json", "local.html", "index.html", "compare.html"):
        copy_file(REPO / "demo/web" / filename, output / "source/demo/web" / filename)
    copy_file(REPO / "demo/web/public/favicon.png", output / "source/demo/web/public/favicon.png")
    copy_file(Path(__file__), output / "source/scripts/deploy_huggingface_space.py")
    copy_file(REPO / "scripts/check_huggingface_space.mjs", output / "source/scripts/check_huggingface_space.mjs")
    copy_file(REPO / "scripts/prepare_public_demo.py", output / "source/scripts/prepare_public_demo.py")
    copy_file(REPO / "scripts/decode_demo_tokens.mjs", output / "source/scripts/decode_demo_tokens.mjs")
    copy_file(REPO / "configs/demo/public_examples.json", output / "source/configs/demo/public_examples.json")
    copy_file(REPO / "deploy/huggingface/GeneralUser-LICENSE.txt", output / "source/deploy/huggingface/GeneralUser-LICENSE.txt")
    copy_file(REPO / "deploy/huggingface/README.md", output / "source/deploy/huggingface/README.md")
    for relative in public_files:
        copy_file(web_dir / relative, output / "source/demo/web/public" / relative)

    def git(*args: str) -> str:
        return subprocess.check_output(["git", *args], cwd=REPO, text=True).strip()

    receipt = {
        "schema_version": 1,
        "source_commit": git("rev-parse", "HEAD"),
        "source_branch": git("branch", "--show-current"),
        "source_dirty": bool(git("status", "--porcelain")),
        "build_command": "VITE_HENDRIX_SPACE=1 npx --yes pnpm@10.20.0 --dir demo/web build --manifest",
        "checkpoint_sha256": manifest["checkpoint"]["sha256"],
        "example_catalog": {"count": len(catalog["examples"]), "sha256": sha256(web_dir / "examples/catalog.json")},
        "files": {
            path.relative_to(output).as_posix(): {"bytes": path.stat().st_size, "sha256": sha256(path)}
            for path in sorted(output.rglob("*")) if path.is_file()
        },
    }
    (output / "deployment.json").write_text(json.dumps(receipt, indent=2) + "\n")
    return receipt


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-id", default="jackboyla/hendrix")
    parser.add_argument("--graph-dir", type=Path, default=REPO / "progress/evaluations/yptf_web_fp32_d26")
    parser.add_argument("--output-dir", type=Path, help="An empty staging directory; default: a new directory under progress/evaluations")
    parser.add_argument("--upload", action="store_true", help="Create/update a public static Space with the current HF login")
    args = parser.parse_args()

    api = None
    if args.upload:
        from huggingface_hub import HfApi, get_token

        if not get_token():
            raise SystemExit("Log in first with: hf auth login")
        api = HfApi()
        api.whoami()  # Fail before building if credentials are invalid.

    subprocess.run(
        ["npx", "--yes", "pnpm@10.20.0", "--dir", "demo/web", "build", "--manifest"],
        cwd=REPO, env={**os.environ, "VITE_HENDRIX_SPACE": "1"}, check=True,
    )
    evaluations = REPO / "progress/evaluations"
    evaluations.mkdir(parents=True, exist_ok=True)
    output = args.output_dir.resolve() if args.output_dir else Path(tempfile.mkdtemp(prefix="hf_space_", dir=evaluations))
    receipt = package_space(
        REPO / "demo/web/dist", args.graph_dir.resolve(),
        REPO / "demo/web/node_modules/onnxruntime-web/dist", output,
    )
    total = sum(record["bytes"] for record in receipt["files"].values())
    print(f"Checked package: {output} ({total / 1_000_000:.1f} MB)", flush=True)

    if api is not None:
        from huggingface_hub.errors import RepositoryNotFoundError

        try:
            existing = api.space_info(args.repo_id)
        except RepositoryNotFoundError:
            api.create_repo(args.repo_id, repo_type="space", space_sdk="static", private=False)
        else:
            if existing.private or existing.sdk != "static":
                raise SystemExit("Existing Space must already be public and static; no settings were changed")
        # No deletion patterns: preserve unrelated files in an existing Space.
        commit = api.upload_folder(
            repo_id=args.repo_id, repo_type="space", folder_path=output,
            commit_message="Update Hendrix music transcription demo",
        )
        print(f"Space: https://huggingface.co/spaces/{args.repo_id}")
        print(f"Commit: {commit.commit_url}")


if __name__ == "__main__":
    main()