File size: 8,919 Bytes
824ebbf |
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 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 |
import gradio as gr
import torch
import numpy as np
import librosa
import soundfile as sf
import tempfile
import os
from transformers import pipeline, VitsModel, AutoTokenizer
from datasets import load_dataset
# For MeloTTS (Chinese and Japanese)
try:
from melo.api import TTS as MeloTTS
except ImportError:
raise ImportError("Please install the MeloTTS package (e.g., pip install myshell-ai/MeloTTS-Chinese)")
# ------------------------------------------------------
# 1. ASR Pipeline (English) using Wav2Vec2
# ------------------------------------------------------
asr = pipeline(
"automatic-speech-recognition",
model="facebook/wav2vec2-base-960h"
)
# ------------------------------------------------------
# 2. Translation Models (8 languages)
# ------------------------------------------------------
translation_models = {
"Spanish": "Helsinki-NLP/opus-mt-en-es",
"Vietnamese": "Helsinki-NLP/opus-mt-en-vi",
"Indonesian": "Helsinki-NLP/opus-mt-en-id",
"Turkish": "Helsinki-NLP/opus-mt-en-trk",
"Portuguese": "Helsinki-NLP/opus-mt-tc-big-en-pt",
"Korean": "Helsinki-NLP/opus-mt-tc-big-en-ko",
"Chinese": "Helsinki-NLP/opus-mt-en-zh",
"Japanese": "Helsinki-NLP/opus-mt-en-jap"
}
translation_tasks = {
"Spanish": "translation_en_to_es",
"Vietnamese": "translation_en_to_vi",
"Indonesian": "translation_en_to_id",
"Turkish": "translation_en_to_tr",
"Portuguese": "translation_en_to_pt",
"Korean": "translation_en_to-ko",
"Chinese": "translation_en_to_zh",
"Japanese": "translation_en_to_ja"
}
# ------------------------------------------------------
# 3. TTS Configuration
# - MMS TTS (VITS) for: Spanish, Vietnamese, Indonesian, Turkish, Portuguese, Korean
# - MeloTTS for: Chinese and Japanese
# ------------------------------------------------------
tts_config = {
"Spanish": {"model_id": "facebook/mms-tts-spa", "architecture": "vits", "type": "mms"},
"Vietnamese": {"model_id": "facebook/mms-tts-vie", "architecture": "vits", "type": "mms"},
"Indonesian": {"model_id": "facebook/mms-tts-ind", "architecture": "vits", "type": "mms"},
"Turkish": {"model_id": "facebook/mms-tts-tur", "architecture": "vits", "type": "mms"},
"Portuguese": {"model_id": "facebook/mms-tts-por", "architecture": "vits", "type": "mms"},
"Korean": {"model_id": "facebook/mms-tts-kor", "architecture": "vits", "type": "mms"},
"Chinese": {"type": "melo"},
"Japanese": {"type": "melo"}
}
# ------------------------------------------------------
# 4. Global Caches for Translators and TTS Models
# ------------------------------------------------------
translator_cache = {}
mms_tts_cache = {} # For MMS (VITS-based) TTS models
melo_tts_cache = {} # For MeloTTS models (Chinese/Japanese)
# ------------------------------------------------------
# 5. Translator Helper
# ------------------------------------------------------
def get_translator(lang):
if lang in translator_cache:
return translator_cache[lang]
model_name = translation_models[lang]
task_name = translation_tasks[lang]
translator = pipeline(task_name, model=model_name)
translator_cache[lang] = translator
return translator
# ------------------------------------------------------
# 6. MMS TTS (VITS) Helper for languages using MMS TTS
# ------------------------------------------------------
def load_mms_tts(lang):
if lang in mms_tts_cache:
return mms_tts_cache[lang]
config = tts_config[lang]
try:
model = VitsModel.from_pretrained(config["model_id"])
tokenizer = AutoTokenizer.from_pretrained(config["model_id"])
mms_tts_cache[lang] = (model, tokenizer)
except Exception as e:
raise RuntimeError(f"Failed to load MMS TTS model for {lang} ({config['model_id']}): {e}")
return mms_tts_cache[lang]
def run_mms_tts(text, lang):
model, tokenizer = load_mms_tts(lang)
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
output = model(**inputs)
if not hasattr(output, "waveform"):
raise RuntimeError(f"MMS TTS model output for {lang} does not contain 'waveform'.")
waveform = output.waveform.squeeze().cpu().numpy()
sample_rate = 16000
return sample_rate, waveform
# ------------------------------------------------------
# 7. MeloTTS Helper for Chinese and Japanese
# ------------------------------------------------------
def run_melo_tts(text, lang):
"""
Uses the myshell-ai MeloTTS model.
For Chinese, use language parameter 'ZH'; for Japanese, use 'JP'.
"""
device = 'cuda' if torch.cuda.is_available() else 'cpu'
lang_param = 'ZH' if lang == "Chinese" else 'JP'
if lang not in melo_tts_cache:
try:
model = MeloTTS(language=lang_param, device=device)
melo_tts_cache[lang] = model
except Exception as e:
raise RuntimeError(f"Failed to load MeloTTS model for {lang}: {e}")
else:
model = melo_tts_cache[lang]
speaker_ids = model.hps.data.spk2id
# Assume the speaker key is the same as lang_param
speaker_key = lang_param
speed = 1.0
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
tmp_name = tmp.name
try:
model.tts_to_file(text, speaker_ids[speaker_key], tmp_name, speed=speed)
data, sr = sf.read(tmp_name)
finally:
if os.path.exists(tmp_name):
os.remove(tmp_name)
return sr, data
# ------------------------------------------------------
# 8. Main Prediction Function
# ------------------------------------------------------
def predict(audio, text, target_language):
"""
1. Obtain English text (via ASR if audio provided, else text).
2. Translate the English text to target_language.
3. Generate TTS audio using either MMS TTS (VITS) or MeloTTS.
"""
# Step 1: Get English text.
if text.strip():
english_text = text.strip()
elif audio is not None:
sample_rate, audio_data = audio
if audio_data.dtype not in [np.float32, np.float64]:
audio_data = audio_data.astype(np.float32)
if len(audio_data.shape) > 1 and audio_data.shape[1] > 1:
audio_data = np.mean(audio_data, axis=1)
if sample_rate != 16000:
audio_data = librosa.resample(audio_data, orig_sr=sample_rate, target_sr=16000)
asr_input = {"array": audio_data, "sampling_rate": 16000}
asr_result = asr(asr_input)
english_text = asr_result["text"]
else:
return "No input provided.", "", None
# Step 2: Translate.
translator = get_translator(target_language)
try:
translation_result = translator(english_text)
translated_text = translation_result[0]["translation_text"]
except Exception as e:
return english_text, f"Translation error: {e}", None
# Step 3: TTS.
try:
tts_type = tts_config[target_language]["type"]
if tts_type == "mms":
sr, waveform = run_mms_tts(translated_text, target_language)
elif tts_type == "melo":
sr, waveform = run_melo_tts(translated_text, target_language)
else:
raise RuntimeError("Unknown TTS type for target language.")
except Exception as e:
return english_text, translated_text, f"TTS error: {e}"
return english_text, translated_text, (sr, waveform)
# ------------------------------------------------------
# 9. Gradio Interface
# ------------------------------------------------------
language_choices = [
"Spanish", "Vietnamese", "Indonesian", "Turkish", "Portuguese", "Korean", "Chinese", "Japanese"
]
iface = gr.Interface(
fn=predict,
inputs=[
gr.Audio(type="numpy", label="Record/Upload English Audio (optional)"),
gr.Textbox(lines=4, placeholder="Or enter English text here", label="English Text Input (optional)"),
gr.Dropdown(choices=language_choices, value="Spanish", label="Target Language")
],
outputs=[
gr.Textbox(label="English Transcription"),
gr.Textbox(label="Translation (Target Language)"),
gr.Audio(label="Synthesized Speech")
],
title="Multimodal Language Learning Aid",
description=(
"This app performs the following steps:\n"
"1. Transcribes English speech using Wav2Vec2 (or accepts text input).\n"
"2. Translates the English text to the target language using Helsinki-NLP MarianMT models.\n"
"3. Synthesizes speech:\n"
" - For Spanish, Vietnamese, Indonesian, Turkish, Portuguese, and Korean: uses Facebook MMS TTS (VITS-based).\n"
" - For Chinese and Japanese: uses myshell-ai MeloTTS models.\n"
"\nSelect your target language from the dropdown."
),
allow_flagging="never"
)
if __name__ == "__main__":
iface.launch(server_name="0.0.0.0", server_port=7860)
|