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 Coqui TTS (XTTS-v2) used for Chinese and Japanese
try:
from TTS.api import TTS as CoquiTTS
except ImportError:
raise ImportError("Please install Coqui TTS via pip install TTS.")
# ------------------------------------------------------
# 1. ASR Pipeline (English) using Wav2Vec2
# ------------------------------------------------------
asr = pipeline(
"automatic-speech-recognition",
model="facebook/wav2vec2-base-960h"
)
# ------------------------------------------------------
# 2. Translation Models (9 languages)
# ------------------------------------------------------
translation_models = {
"French": "Helsinki-NLP/opus-mt-en-fr",
"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 = {
"French": "translation_en_to_fr",
"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: French, Spanish, Vietnamese, Indonesian, Turkish, Portuguese, Korean
# - Coqui XTTS-v2 for: Chinese and Japanese
# ------------------------------------------------------
tts_config = {
"French": {"model_id": "facebook/mms-tts-fra", "architecture": "vits", "type": "mms"},
"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": "coqui"},
"Japanese": {"type": "coqui"}
}
# For Coqui, map languages to expected language codes.
coqui_lang_map = {
"Chinese": "zh",
"Japanese": "ja"
}
# ------------------------------------------------------
# 4. Global Caches for Translators and TTS Models
# ------------------------------------------------------
translator_cache = {}
mms_tts_cache = {}
coqui_tts_cache = None
# ------------------------------------------------------
# 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. Coqui TTS Helper for Chinese and Japanese
# ------------------------------------------------------
def load_coqui_tts():
global coqui_tts_cache
if coqui_tts_cache is not None:
return coqui_tts_cache
try:
coqui_tts_cache = CoquiTTS("tts_models/multilingual/multi-dataset/xtts_v2", gpu=False)
except Exception as e:
raise RuntimeError(f"Failed to load Coqui XTTS-v2 TTS: {e}")
return coqui_tts_cache
def run_coqui_tts(text, lang):
coqui_tts = load_coqui_tts()
lang_code = coqui_lang_map[lang]
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
tmp_name = tmp.name
try:
coqui_tts.tts_to_file(
text=text,
file_path=tmp_name,
language=lang_code
)
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 English text to target_language.
3. Generate TTS audio using either MMS TTS (VITS) or Coqui XTTS-v2.
"""
# 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"].lower()
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 == "coqui":
sr, waveform = run_coqui_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 = [
"French", "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="French", 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 tasks:\n"
"1. Transcribes English speech using Wav2Vec2 (accepts text input as well).\n"
"2. Translates the English text to the target language using Helsinki-NLP models.\n"
"3. Provides speech:\n"
" - For French, Spanish, Vietnamese, Indonesian, Turkish, Portuguese, and Korean: uses Facebook MMS TTS (VITS-based).\n"
" - For Chinese and Japanese: uses myshell-ai MeloTTS models (work-in-progress).\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)