Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| """Fetch the model checkpoints that are deliberately kept out of git. | |
| A `git clone` of PhonePaint is ~26 MB and contains no weights: `weights/best.pt` | |
| is blocked by the top-level `.gitignore` `*.pt` rule and the MAPS checkpoints by | |
| `tools/maps/.gitignore`. This script puts them back. | |
| Two backends: | |
| --from DIR copy from another PhonePaint checkout or asset tree. Offline, and | |
| useful when a machine already has the files somewhere. | |
| (default) download from the Hugging Face model repo. Needs read access; the | |
| repo is private for now, so export HF_TOKEN or `hf auth login`. | |
| Paths in the HF repo mirror the paths in this checkout exactly, so the mapping | |
| is 1:1 in both directions and nothing has to be renamed on arrival. | |
| Usage: | |
| python tools/fetch_assets.py # download whatever is missing | |
| python tools/fetch_assets.py --check # report only, exit 1 if short | |
| python tools/fetch_assets.py --from ~/PhonePaint # copy from a local checkout | |
| python tools/fetch_assets.py --force # re-fetch even if present | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import shutil | |
| import sys | |
| from pathlib import Path | |
| REPO_ROOT = Path(__file__).resolve().parents[1] | |
| DEFAULT_HF_REPO = "Tom-Brenner/phonepaint-assets" | |
| # Repo-root-relative. Same strings are used as path_in_repo on the Hub. | |
| ASSETS: tuple[str, ...] = ( | |
| "weights/best.pt", | |
| "tools/maps/torch_models/timbuck_eng.pt", | |
| *( | |
| f"tools/maps/torch_models/ensemble_model/timbuck_eng_{i}.pt" | |
| for i in range(1, 11) | |
| ), | |
| ) | |
| def _present(path: Path) -> bool: | |
| """A truncated or zero-byte file counts as missing, not as present.""" | |
| return path.is_file() and path.stat().st_size > 0 | |
| def _missing(dest_root: Path) -> list[str]: | |
| return [rel for rel in ASSETS if not _present(dest_root / rel)] | |
| def _report(dest_root: Path) -> int: | |
| missing = _missing(dest_root) | |
| for rel in ASSETS: | |
| path = dest_root / rel | |
| if _present(path): | |
| print(f" ok {rel} ({path.stat().st_size / 2**20:.1f} MiB)") | |
| else: | |
| print(f" MISSING {rel}") | |
| if missing: | |
| print(f"\n{len(missing)} of {len(ASSETS)} assets missing under {dest_root}") | |
| return 1 | |
| print(f"\nall {len(ASSETS)} assets present under {dest_root}") | |
| return 0 | |
| def _fetch_local(source_root: Path, dest_root: Path, wanted: list[str]) -> None: | |
| for rel in wanted: | |
| src = source_root / rel | |
| if not _present(src): | |
| raise SystemExit(f"source is missing {rel}: {src}") | |
| dst = dest_root / rel | |
| dst.parent.mkdir(parents=True, exist_ok=True) | |
| shutil.copy2(src, dst) | |
| print(f" copied {rel} ({dst.stat().st_size / 2**20:.1f} MiB)") | |
| def _fetch_hf( | |
| repo_id: str, | |
| revision: str | None, | |
| token: str | None, | |
| dest_root: Path, | |
| wanted: list[str], | |
| ) -> None: | |
| try: | |
| from huggingface_hub import hf_hub_download | |
| except ImportError as exc: # pragma: no cover - environment problem, not logic | |
| raise SystemExit( | |
| "huggingface_hub is required for the download backend " | |
| "(it ships with the PhonePaint env). Use --from DIR to copy " | |
| f"from a local checkout instead. Import error: {exc}" | |
| ) from exc | |
| from huggingface_hub.errors import HfHubHTTPError | |
| for rel in wanted: | |
| try: | |
| cached = hf_hub_download( | |
| repo_id=repo_id, | |
| filename=rel, | |
| revision=revision, | |
| token=token, | |
| repo_type="model", | |
| ) | |
| except HfHubHTTPError as exc: | |
| raise SystemExit( | |
| f"failed to download {rel} from {repo_id}: {exc}\n" | |
| "If the repo is private, export HF_TOKEN or run `hf auth login` " | |
| "with a token that has read access." | |
| ) from exc | |
| dst = dest_root / rel | |
| dst.parent.mkdir(parents=True, exist_ok=True) | |
| # Copy out of the shared HF cache: the pipeline expects real files at | |
| # these paths, and a cache entry can be garbage-collected underneath us. | |
| shutil.copy2(cached, dst) | |
| print(f" fetched {rel} ({dst.stat().st_size / 2**20:.1f} MiB)") | |
| def main() -> int: | |
| ap = argparse.ArgumentParser( | |
| description=__doc__, | |
| formatter_class=argparse.RawDescriptionHelpFormatter, | |
| ) | |
| ap.add_argument( | |
| "--check", | |
| action="store_true", | |
| help="Report which assets are present and exit 1 if any are missing", | |
| ) | |
| ap.add_argument( | |
| "--from", | |
| dest="source", | |
| type=Path, | |
| default=None, | |
| metavar="DIR", | |
| help="Copy from this local checkout instead of downloading", | |
| ) | |
| ap.add_argument( | |
| "--repo", | |
| default=DEFAULT_HF_REPO, | |
| help=f"Hugging Face model repo (default: {DEFAULT_HF_REPO})", | |
| ) | |
| ap.add_argument( | |
| "--revision", | |
| default=None, | |
| help="Branch, tag, or commit on the Hub (default: the repo default branch)", | |
| ) | |
| ap.add_argument( | |
| "--token", | |
| default=None, | |
| help="HF token; otherwise HF_TOKEN or a cached `hf auth login` is used", | |
| ) | |
| ap.add_argument( | |
| "--dest", | |
| type=Path, | |
| default=REPO_ROOT, | |
| help=f"Checkout to populate (default: {REPO_ROOT})", | |
| ) | |
| ap.add_argument( | |
| "--force", | |
| action="store_true", | |
| help="Re-fetch assets that are already present", | |
| ) | |
| args = ap.parse_args() | |
| dest_root = args.dest.expanduser().resolve() | |
| if args.check: | |
| return _report(dest_root) | |
| wanted = list(ASSETS) if args.force else _missing(dest_root) | |
| if not wanted: | |
| print(f"all {len(ASSETS)} assets already present under {dest_root}") | |
| return 0 | |
| print(f"{len(wanted)} asset(s) to fetch into {dest_root}") | |
| if args.source is not None: | |
| source_root = args.source.expanduser().resolve() | |
| if source_root == dest_root: | |
| raise SystemExit("--from source and --dest are the same directory") | |
| print(f"source: {source_root} (local copy)\n") | |
| _fetch_local(source_root, dest_root, wanted) | |
| else: | |
| print(f"source: {args.repo} (Hugging Face)\n") | |
| _fetch_hf(args.repo, args.revision, args.token, dest_root, wanted) | |
| still_missing = _missing(dest_root) | |
| if still_missing: | |
| print("\nstill missing after fetch:", ", ".join(still_missing)) | |
| return 1 | |
| print(f"\nall {len(ASSETS)} assets present under {dest_root}") | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |