File size: 2,110 Bytes
7496177 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | from __future__ import annotations
import argparse
import logging
import os
import shutil
from pathlib import Path
logger = logging.getLogger(__name__)
def _resolve_path(name_or_path: str) -> str:
if os.path.isdir(name_or_path):
return name_or_path
from huggingface_hub import snapshot_download
return snapshot_download(name_or_path)
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
prog="omnivoice-prefetch",
description="Cache auxiliary OmniVoice assets for offline/container use.",
)
parser.add_argument(
"--model-dir",
default="/app/model",
help="Directory containing the OmniVoice model files.",
)
parser.add_argument(
"--audio-tokenizer",
default="eustlb/higgs-audio-v2-tokenizer",
help="Audio tokenizer repo id or local path.",
)
parser.add_argument(
"--asr-model",
default="openai/whisper-large-v3-turbo",
help="ASR model repo id or local path.",
)
parser.add_argument(
"--copy-audio-tokenizer",
action="store_true",
help="Copy the tokenizer into <model-dir>/audio_tokenizer.",
)
parser.add_argument(
"--prefetch-asr",
action="store_true",
help="Download the ASR model into the Hugging Face cache.",
)
return parser
def main() -> None:
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
args = build_parser().parse_args()
tokenizer_path = _resolve_path(args.audio_tokenizer)
if args.copy_audio_tokenizer:
model_dir = Path(args.model_dir)
target_dir = model_dir / "audio_tokenizer"
target_dir.mkdir(parents=True, exist_ok=True)
shutil.copytree(tokenizer_path, target_dir, dirs_exist_ok=True)
logger.info("Copied audio tokenizer to %s", target_dir)
else:
logger.info("Cached audio tokenizer at %s", tokenizer_path)
if args.prefetch_asr:
asr_path = _resolve_path(args.asr_model)
logger.info("Cached ASR model at %s", asr_path)
|