| 1 | #!/usr/bin/env python3 |
| 2 | """Download pinned model files once, with resumable transfers and SHA-256 checks.""" |
| 3 | import argparse |
| 4 | import fcntl |
| 5 | import hashlib |
| 6 | import json |
| 7 | import os |
| 8 | from pathlib import Path |
| 9 | import subprocess |
| 10 | |
| 11 | |
| 12 | def checksum(path): |
| 13 | with path.open("rb") as source: |
| 14 | return hashlib.file_digest(source, "sha256").hexdigest() |
| 15 | |
| 16 | |
| 17 | def download(root, runtime_only=False): |
| 18 | manifest = json.loads(Path(__file__).with_name("models.json").read_text()) |
| 19 | directory = root / manifest["directory"] |
| 20 | directory.mkdir(parents=True, exist_ok=True) |
| 21 | with (directory / ".download.lock").open("a") as lock: |
| 22 | fcntl.flock(lock, fcntl.LOCK_EX) |
| 23 | for spec in manifest["files"]: |
| 24 | if runtime_only and not spec["name"].startswith("runtime/"): |
| 25 | continue |
| 26 | target = directory / spec["name"] |
| 27 | target.parent.mkdir(parents=True, exist_ok=True) |
| 28 | if target.exists(): |
| 29 | if target.stat().st_size == spec["bytes"] and checksum(target) == spec["sha256"]: |
| 30 | print(f"Verified {target.name}", flush=True) |
| 31 | continue |
| 32 | raise ValueError(f"Checksum mismatch: {target}. Move this file aside and retry.") |
| 33 | pending = target.with_suffix(target.suffix + ".part") |
| 34 | url = f"https://huggingface.co/{manifest['repository']}/resolve/{manifest['revision']}/{spec['name']}" |
| 35 | print(f"Downloading {target.name}", flush=True) |
| 36 | subprocess.run([ |
| 37 | "curl", "--fail", "--location", "--silent", "--show-error", |
| 38 | "--continue-at", "-", "--retry", "3", "--retry-delay", "3", |
| 39 | "--connect-timeout", "15", "--max-time", "28800", |
| 40 | "--output", str(pending), url, |
| 41 | ], check=True) |
| 42 | if pending.stat().st_size != spec["bytes"] or checksum(pending) != spec["sha256"]: |
| 43 | raise ValueError(f"Checksum mismatch: {pending}. Move this file aside and retry.") |
| 44 | pending.chmod(0o640) |
| 45 | pending.replace(target) |
| 46 | print(f"Verified {target.name}", flush=True) |
| 47 | (directory / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") |
| 48 | |
| 49 | |
| 50 | if __name__ == "__main__": |
| 51 | parser = argparse.ArgumentParser(description="Download local AI models") |
| 52 | parser.add_argument("--root", type=Path, default=Path("/srv/clover/Media/AI/LLM")) |
| 53 | parser.add_argument("--runtime-only", action="store_true") |
| 54 | args = parser.parse_args() |
| 55 | os.umask(0o027) |
| 56 | download(args.root, args.runtime_only) |