app.py

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)