app.py

import gradio as gr
import torch
import numpy as np
import librosa
from transformers import pipeline

# --------------------------------------------------
# ASR Pipeline (for English transcription)
# --------------------------------------------------
asr = pipeline(
    "automatic-speech-recognition",
    model="facebook/wav2vec2-base-960h"
)

# --------------------------------------------------
# Mapping for Target Languages (Spanish, Chinese, Japanese)
# --------------------------------------------------
translation_models = {
    "Spanish": "Helsinki-NLP/opus-mt-en-es",
    "Chinese": "Helsinki-NLP/opus-mt-en-zh",
    "Japanese": "Helsinki-NLP/opus-mt-en-ja"
}

translation_tasks = {
    "Spanish": "translation_en_to_es",
    "Chinese": "translation_en_to_zh",
    "Japanese": "translation_en_to_ja"
}

tts_models = {
    "Spanish": "facebook/mms-tts-spa",
    "Chinese": "facebook/mms-tts-che",     
    "Japanese": "esnya/japanese_speecht5_tts"
}

# --------------------------------------------------
# Caches for translator and TTS pipelines
# --------------------------------------------------
translator_cache = {}
tts_cache = {}

def get_translator(target_language):
    if target_language in translator_cache:
        return translator_cache[target_language]
    model_name = translation_models[target_language]
    task_name = translation_tasks[target_language]
    translator = pipeline(task_name, model=model_name)
    translator_cache[target_language] = translator
    return translator

def get_tts(target_language):
    if target_language in tts_cache:
        return tts_cache[target_language]
    model_name = tts_models.get(target_language)
    if model_name is None:
        raise ValueError(f"No TTS model available for {target_language}.")
    try:
        tts_pipeline = pipeline("text-to-speech", model=model_name)
    except Exception as e:
        raise ValueError(f"Failed to load TTS model for {target_language} with model '{model_name}'.\nError: {e}")
    tts_cache[target_language] = tts_pipeline
    return tts_pipeline

# --------------------------------------------------
# Prediction Function
# --------------------------------------------------
def predict(audio, text, target_language):
    # Step 1: Obtain English text from text input if provided, otherwise use ASR.
    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)
        input_audio = {"array": audio_data, "sampling_rate": 16000}
        asr_result = asr(input_audio)
        english_text = asr_result["text"]
    else:
        return "No input provided.", "", None

    # Step 2: Translate the English text to the target language.
    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: Synthesize speech using the TTS pipeline.
    try:
        tts_pipeline = get_tts(target_language)
        tts_result = tts_pipeline(translated_text)
        synthesized_audio = (tts_result["sample_rate"], tts_result["wav"])
    except Exception as e:
        return english_text, translated_text, f"TTS error: {e}"

    return english_text, translated_text, synthesized_audio

# --------------------------------------------------
# Gradio Interface Setup
# --------------------------------------------------
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=list(translation_models.keys()), value="Spanish", label="Target Language")
    ],
    outputs=[
        gr.Textbox(label="English Transcription"),
        gr.Textbox(label="Translation (Target Language)"),
        gr.Audio(label="Synthesized Speech in Target Language")
    ],
    title="Multimodal Language Learning Aid",
    description=(
        "This app provides three outputs:\n"
        "1. English transcription (from ASR or text input),\n"
        "2. Translation to Spanish, Chinese, or Japanese (using Helsinki-NLP models), and\n"
        "3. Synthetic speech in the target language (using Facebook MMS TTS or equivalent).\n\n"
        "Either record/upload an English audio sample or enter English text directly."
    ),
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch()