Michael Treadgold
73 lines
2.0 KiB
Python
73 lines
2.0 KiB
Python
"""Download Inflect-Nano-v1 model weights from HuggingFace.
|
|
|
|
Run this once after cloning the repo:
|
|
python download_weights.py
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Use requests or urllib — try both
|
|
try:
|
|
import requests
|
|
_has_requests = True
|
|
except ImportError:
|
|
import urllib.request
|
|
import urllib.error
|
|
_has_requests = False
|
|
|
|
WEIGHTS_DIR = Path(__file__).resolve().parent / "weights"
|
|
HF_BASE = "https://huggingface.co/owensong/Inflect-Nano-v1/resolve/main/weights"
|
|
|
|
FILES = [
|
|
"inflect_nano_v1_acoustic.pt",
|
|
"inflect_nano_v1_vocoder.pt",
|
|
]
|
|
|
|
|
|
def download_file(url: str, dest: Path) -> None:
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
print(f"Downloading {dest.name}...", end=" ", flush=True)
|
|
|
|
if _has_requests:
|
|
resp = requests.get(url, stream=True)
|
|
resp.raise_for_status()
|
|
total = int(resp.headers.get("content-length", 0))
|
|
with dest.open("wb") as f:
|
|
downloaded = 0
|
|
for chunk in resp.iter_content(chunk_size=8192):
|
|
f.write(chunk)
|
|
downloaded += len(chunk)
|
|
if total:
|
|
pct = downloaded * 100 // total
|
|
print(f"\rDownloading {dest.name}... {pct}%", end="", flush=True)
|
|
print(" done.")
|
|
else:
|
|
urllib.request.urlretrieve(url, str(dest))
|
|
print("done.")
|
|
|
|
|
|
def main() -> None:
|
|
missing = [f for f in FILES if not (WEIGHTS_DIR / f).is_file()]
|
|
if not missing:
|
|
print("All weights already present.")
|
|
return
|
|
|
|
print(f"Missing {len(missing)} weight file(s). Downloading from HuggingFace...\n")
|
|
for filename in missing:
|
|
url = f"{HF_BASE}/{filename}"
|
|
dest = WEIGHTS_DIR / filename
|
|
try:
|
|
download_file(url, dest)
|
|
except Exception as e:
|
|
print(f"\n FAILED: {e}")
|
|
print(f" Manual download: {url}")
|
|
sys.exit(1)
|
|
|
|
print(f"\nAll weights downloaded to {WEIGHTS_DIR}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|