mkbk
120 lines
4.0 KiB
Python
120 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tarfile
|
|
import tempfile
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
SRC_DIR = PROJECT_ROOT / "src"
|
|
if str(SRC_DIR) not in sys.path:
|
|
sys.path.insert(0, str(SRC_DIR))
|
|
|
|
from owner_voice_pet.speech_models import (
|
|
DEFAULT_STT_DIR,
|
|
DEFAULT_STT_URL,
|
|
DEFAULT_VAD_URL,
|
|
check_speech_models,
|
|
write_default_manifest,
|
|
)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description="Download local speech models for owner_voice_pet.")
|
|
parser.add_argument("--dir", default="models", help="Target model directory. Default: models")
|
|
parser.add_argument("--force", action="store_true", help="Re-download and replace existing model files")
|
|
args = parser.parse_args(argv)
|
|
|
|
target = Path(args.dir)
|
|
target.mkdir(parents=True, exist_ok=True)
|
|
download_vad(target, force=args.force)
|
|
download_stt(target, force=args.force)
|
|
manifest = write_default_manifest(target)
|
|
status = check_speech_models(target, require_sherpa=False)
|
|
print(f"models_dir={target}")
|
|
print(f"manifest={manifest}")
|
|
print(f"missing_files={list(status.missing_files)}")
|
|
return 0 if not status.missing_files else 1
|
|
|
|
|
|
def download_vad(models_dir: Path, *, force: bool = False) -> Path:
|
|
output = models_dir / "vad" / "silero_vad.onnx"
|
|
if output.exists() and not force:
|
|
print(f"skip existing {output}")
|
|
return output
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
download_file(DEFAULT_VAD_URL, output)
|
|
return output
|
|
|
|
|
|
def download_stt(models_dir: Path, *, force: bool = False) -> Path:
|
|
output_dir = models_dir / "stt" / DEFAULT_STT_DIR
|
|
if output_dir.exists() and any(output_dir.iterdir()) and not force:
|
|
print(f"skip existing {output_dir}")
|
|
return output_dir
|
|
output_dir.parent.mkdir(parents=True, exist_ok=True)
|
|
last_error: Exception | None = None
|
|
for attempt in range(1, 4):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
archive = Path(tmp) / f"{DEFAULT_STT_DIR}.tar.bz2"
|
|
try:
|
|
download_file(DEFAULT_STT_URL, archive)
|
|
extract_tar_bz2(archive, output_dir.parent)
|
|
return output_dir
|
|
except Exception as exc:
|
|
last_error = exc
|
|
print(f"download/extract attempt {attempt} failed: {exc}")
|
|
if output_dir.exists():
|
|
shutil.rmtree(output_dir)
|
|
if last_error is not None:
|
|
raise last_error
|
|
return output_dir
|
|
|
|
|
|
def download_file(url: str, output: Path) -> None:
|
|
tmp_output = output.with_suffix(output.suffix + ".part")
|
|
print(f"download {url}")
|
|
if shutil.which("curl"):
|
|
subprocess.run(
|
|
[
|
|
"curl",
|
|
"-L",
|
|
"--fail",
|
|
"--retry",
|
|
"3",
|
|
"--connect-timeout",
|
|
"20",
|
|
"--max-time",
|
|
"600",
|
|
"-o",
|
|
str(tmp_output),
|
|
url,
|
|
],
|
|
check=True,
|
|
)
|
|
else:
|
|
with urllib.request.urlopen(url, timeout=600) as response, tmp_output.open("wb") as handle:
|
|
shutil.copyfileobj(response, handle)
|
|
tmp_output.replace(output)
|
|
print(f"wrote {output} ({output.stat().st_size} bytes)")
|
|
|
|
|
|
def extract_tar_bz2(archive: Path, target_dir: Path) -> None:
|
|
with tarfile.open(archive, "r:bz2") as tar:
|
|
members = tar.getmembers()
|
|
resolved_target = target_dir.resolve()
|
|
for member in members:
|
|
destination = (target_dir / member.name).resolve()
|
|
if resolved_target != destination and resolved_target not in destination.parents:
|
|
raise RuntimeError(f"refusing unsafe archive member: {member.name}")
|
|
tar.extractall(target_dir, members=members)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|