"""Local inference script (same model settings as the API)."""
from pathlib import Path

from omnivoice import OmniVoice
import soundfile as sf
import torch

MODEL_DIR = Path(__file__).resolve().parent / "KhanhTTS-OmniVoice"

model = OmniVoice.from_pretrained(
    str(MODEL_DIR),
    device_map="cpu",
    dtype=torch.float32,
    local_files_only=True,
)
REF_AUDIO = MODEL_DIR.parent / "refs" / "reference.wav"
REF_TEXT = (MODEL_DIR.parent / "refs" / "reference.txt").read_text(encoding="utf-8").strip() if (MODEL_DIR.parent / "refs" / "reference.txt").is_file() else None

if REF_AUDIO.is_file():
    voice = model.create_voice_clone_prompt(ref_audio=str(REF_AUDIO), ref_text=REF_TEXT)
    audio = model.generate(text="Xin chào các bạn.", voice_clone_prompt=voice)
else:
    raise SystemExit(f"Đặt file giọng mẫu tại: {REF_AUDIO}")

sf.write("out.wav", audio[0], 24000)