Spaces:
Running
Running
| from __future__ import annotations | |
| import threading | |
| from functools import lru_cache | |
| from pathlib import Path | |
| import soundfile as sf | |
| import torch | |
| from pocket_tts import TTSModel | |
| _INFERENCE_LOCK = threading.Lock() | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| def _load_model() -> TTSModel: | |
| model = TTSModel.load_model(language="english") | |
| model.to(DEVICE) | |
| model.eval() | |
| return model | |
| def _load_voice_state(voice: str): | |
| return _load_model().get_state_for_audio_prompt(voice) | |
| def load_model() -> None: | |
| """Load the model now, so that the first request does not wait on it.""" | |
| _load_model() | |
| def synthesize(text: str, voice: str, output_path: Path) -> Path: | |
| model = _load_model() | |
| voice_state = _load_voice_state(voice) | |
| # generate_audio is not thread-safe. | |
| with _INFERENCE_LOCK: | |
| audio = model.generate_audio(voice_state, text, copy_state=True) | |
| waveform = audio.detach().cpu().float().squeeze().numpy() | |
| sf.write(str(output_path), waveform, model.sample_rate, subtype="PCM_16") | |
| return output_path | |