MathMai / tts.py
Heterogeneity2025's picture
Upload 2 files
5789671 verified
Raw
History Blame Contribute Delete
2.79 kB
"""Text-to-speech using a Hugging Face model (facebook/mms-tts-eng).
The model is loaded lazily and cached so the first call pays the download/load
cost and later calls are fast. Every public function is wrapped so that if the
model (or torch) is unavailable the game keeps working silently instead of
crashing -- ``speak()`` simply returns ``None`` and the UI shows no audio.
``speak()`` returns a path to a freshly written ``.wav`` file. Browsers play a
real audio file far more reliably than an in-memory array, which matters for
gr.Audio autoplay.
"""
import os
import tempfile
MODEL_ID = "facebook/mms-tts-eng"
# Reuse one temp directory for the generated clips so they don't litter the disk.
_AUDIO_DIR = os.path.join(tempfile.gettempdir(), "math_adventure_audio")
os.makedirs(_AUDIO_DIR, exist_ok=True)
_clip_counter = 0
_model = None
_tokenizer = None
_load_failed = False
def _load():
"""Load and cache the TTS model + tokenizer. Returns True on success."""
global _model, _tokenizer, _load_failed
if _model is not None:
return True
if _load_failed:
return False
try:
from transformers import VitsModel, AutoTokenizer
_model = VitsModel.from_pretrained(MODEL_ID)
_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
_model.eval()
return True
except Exception as exc: # pragma: no cover - environment dependent
print(f"[tts] could not load {MODEL_ID}: {exc}. Audio disabled.")
_load_failed = True
return False
def warm_up():
"""Pre-load the model at app startup (optional; speeds up the first round)."""
_load()
def speak(text):
"""Synthesize ``text`` and return the path to a ``.wav`` file for gr.Audio.
Returns ``None`` if synthesis is unavailable, which gr.Audio renders as
"no audio" rather than erroring.
"""
global _clip_counter
if not text or not _load():
return None
try:
import numpy as np
import torch
from scipy.io import wavfile
inputs = _tokenizer(text, return_tensors="pt")
with torch.no_grad():
waveform = _model(**inputs).waveform
audio = waveform.squeeze().cpu().numpy()
sr = _model.config.sampling_rate
pcm = np.int16(np.clip(audio, -1.0, 1.0) * 32767)
# Rotate over a small set of filenames to bound disk use while still
# giving each clip a fresh URL so the browser doesn't serve a stale one.
_clip_counter = (_clip_counter + 1) % 8
path = os.path.join(_AUDIO_DIR, f"clip_{_clip_counter}.wav")
wavfile.write(path, sr, pcm)
return path
except Exception as exc: # pragma: no cover - environment dependent
print(f"[tts] synthesis failed: {exc}")
return None