study-buddy / app /benchmarks /tts_asset_setup.py
GitHub Actions
deploy d092bea3608b7a29952f16357fda39b7a29e399b
2e818da
Raw
History Blame Contribute Delete
3.32 kB
"""Pinned, resumable asset setup for ResearchMate's pocket-tts default voice and Piper fallback voice."""
from __future__ import annotations
import argparse
import hashlib
import json
from dataclasses import asdict, dataclass
from pathlib import Path
@dataclass(frozen=True)
class DownloadSpec:
name: str
repo_id: str
revision: str
destination: str
allow_patterns: tuple[str, ...] = ()
def download_specs() -> list[DownloadSpec]:
return [
DownloadSpec(
"pocket-tts",
"kyutai/pocket-tts",
"4c8ad48f8a003909bc4f1122cbe88a4252124621",
"pocket",
),
DownloadSpec(
"piper-en-us",
"rhasspy/piper-voices",
"e21c7de8d4eab79b902f0d61e662b3f21664b8d2",
"piper",
(
"en/en_US/lessac/medium/en_US-lessac-medium.onnx",
"en/en_US/lessac/medium/en_US-lessac-medium.onnx.json",
"en/en_US/lessac/medium/MODEL_CARD",
),
),
]
def setup_assets(root: Path) -> dict:
from huggingface_hub import snapshot_download
root.mkdir(parents=True, exist_ok=True)
rows: list[dict] = []
for spec in download_specs():
destination = root / spec.destination
try:
snapshot_download(
repo_id=spec.repo_id,
revision=spec.revision,
allow_patterns=list(spec.allow_patterns) or None,
local_dir=destination,
)
rows.append({
**asdict(spec),
"status": "installed",
"bytes": _directory_size(destination),
"checksums": _checksums(destination),
})
except Exception as exc:
rows.append({**asdict(spec), "status": "error", "error": str(exc)})
manifest = {"schema_version": 1, "root": str(root), "assets": rows, "total_bytes": _directory_size(root)}
(root / "installation-manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
return manifest
def _directory_size(path: Path) -> int:
return sum(item.stat().st_size for item in path.rglob("*") if item.is_file()) if path.exists() else 0
def _sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def _checksums(path: Path) -> dict[str, str]:
return {
item.relative_to(path).as_posix(): _sha256(item)
for item in path.rglob("*") if item.is_file() and not item.name.startswith(".cache")
}
def main() -> None:
parser = argparse.ArgumentParser(description="Download pinned Pocket TTS and Piper assets.")
parser.add_argument("--root", type=Path, default=Path.home() / ".studybuddy" / "models")
parser.add_argument("--yes", action="store_true", help="Accept the model download without prompting")
args = parser.parse_args()
if not args.yes:
answer = input("Download Pocket TTS and Piper assets? [y/N] ")
if answer.strip().lower() not in {"y", "yes"}:
return
manifest = setup_assets(args.root)
print(json.dumps(manifest, indent=2))
if __name__ == "__main__":
main()