File size: 5,076 Bytes
4fcd019 41edec1 4fcd019 41edec1 ee52b8e 41edec1 ee52b8e 4fcd019 | 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """
Text-to-speech for TalkToDoc.
Converts translated text responses into audible speech in the patient's
selected language.
English, Yoruba, Igbo, and Hausa use YarnGPT, a local model trained
specifically on Nigerian-accented speech.
Nigerian Pidgin uses MMS-TTS (facebook/mms-tts-pcm) instead, since YarnGPT
has no Pidgin support. MMS-TTS is the only free option found with a
dedicated Pidgin checkpoint.
Note on the WavTokenizer checkpoint: YarnGPT downloads this itself on
first import, using a bare requests.get() with no error checking, which
can silently save a corrupted file if the download hiccups. The
Dockerfile pre-downloads the same two files at build time using
huggingface_hub's properly tested download function instead, so by the
time this file is imported, YarnGPT finds them already in place and
skips its own fragile download step. Locally (not in Docker), the first
import still triggers YarnGPT's own download as normal, this only
matters for the deployed container.
Performance note: YarnGPT's own generate_speech() function reloads its
full model from disk on every call, which is too slow for a live app.
This file loads the model once and reuses it, using the same generation
steps YarnGPT's own function uses internally, just without the reload.
Call preload_models() once when the app starts, so the first real request
isn't slow either.
"""
import torch
import torchaudio
import scipy.io.wavfile
from transformers import VitsModel, AutoTokenizer
from yarngpt.core import load_model_and_tokenizer, SPEAKER_MAPPING, AVAILABLE_SPEAKERS
# Default speaker used for each YarnGPT-supported language.
YARNGPT_SPEAKERS = {
"english": "idera",
"yoruba": "abayomi",
"igbo": "chioma",
"hausa": "amina",
}
_yarngpt_model = None
_yarngpt_tokenizer = None
_pidgin_model = None
_pidgin_tokenizer = None
def _get_yarngpt_model():
global _yarngpt_model, _yarngpt_tokenizer
if _yarngpt_model is None:
_yarngpt_model, _yarngpt_tokenizer = load_model_and_tokenizer()
return _yarngpt_model, _yarngpt_tokenizer
def _get_pidgin_model():
global _pidgin_model, _pidgin_tokenizer
if _pidgin_model is None:
_pidgin_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-pcm")
_pidgin_model = VitsModel.from_pretrained("facebook/mms-tts-pcm")
return _pidgin_model, _pidgin_tokenizer
def preload_models():
"""Loads both TTS backends into memory ahead of time. Call this once
when the Flask app starts, so the first real request isn't slow."""
_get_yarngpt_model()
_get_pidgin_model()
def _generate_yarngpt_speech(text, speaker, language, temperature=0.1, repetition_penalty=1.1, max_length=4000):
"""
Same steps as yarngpt's own generate_speech(), but reuses the model
already loaded by _get_yarngpt_model() instead of reloading it.
"""
model_speaker = SPEAKER_MAPPING.get(speaker, speaker)
if model_speaker not in AVAILABLE_SPEAKERS:
raise ValueError(f"Unknown speaker: {speaker}")
model, audio_tokenizer = _get_yarngpt_model()
prompt = audio_tokenizer.create_prompt(text, language, model_speaker)
input_ids = audio_tokenizer.tokenize_prompt(prompt)
attention_mask = torch.ones_like(input_ids)
output = model.generate(
input_ids=input_ids,
attention_mask=attention_mask,
do_sample=True,
temperature=temperature,
repetition_penalty=repetition_penalty,
max_length=max_length,
pad_token_id=model.config.eos_token_id,
eos_token_id=model.config.eos_token_id,
)
codes = audio_tokenizer.get_codes(output)
audio = audio_tokenizer.get_audio(codes)
return audio
def synthesize_speech(text, language, output_path):
"""
text: the text to speak
language: one of "english", "yoruba", "hausa", "igbo", "pidgin"
output_path: where to save the resulting .wav file
Returns output_path.
"""
language = language.lower()
if language in YARNGPT_SPEAKERS:
speaker = YARNGPT_SPEAKERS[language]
audio = _generate_yarngpt_speech(text, speaker=speaker, language=language)
torchaudio.save(output_path, audio, sample_rate=24000)
return output_path
if language == "pidgin":
model, tokenizer = _get_pidgin_model()
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
output = model(**inputs).waveform
scipy.io.wavfile.write(
output_path,
rate=model.config.sampling_rate,
data=output.numpy().squeeze(),
)
return output_path
raise ValueError(f"Unsupported language: {language}")
if __name__ == "__main__":
import sys
if len(sys.argv) < 3:
print('Usage: python tts.py "text" language [output_file]')
else:
input_text = sys.argv[1]
selected_language = sys.argv[2]
output_file = sys.argv[3] if len(sys.argv) > 3 else "output.wav"
synthesize_speech(input_text, selected_language, output_file)
print("Saved to", output_file)
|