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 ) # For Coqui TTS 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 (3 languages) # ------------------------------------------------------ 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" } # ------------------------------------------------------ # 3. TTS Config: # - Spanish: MMS TTS (facebook/mms-tts-spa) # - Chinese, Japanese: Coqui XTTS-v2 (tts_models/multilingual/multi-dataset/xtts_v2) # ------------------------------------------------------ SPANISH = "Spanish" CHINESE = "Chinese" JAPANESE = "Japanese" # For Spanish (MMS) mms_spanish_config = { "model_id": "facebook/mms-tts-spa", "architecture": "vits" } # We'll map Chinese/Japanese to Coqui language codes coqui_lang_map = { CHINESE: "zh", JAPANESE: "ja" } # ------------------------------------------------------ # 4. Global Caches # ------------------------------------------------------ translator_cache = {} spanish_vits_cache = None coqui_tts_cache = None def get_translator(lang): """ Return a cached MarianMT translator for the specified language. """ 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 # ------------------------------------------------------ # 5. Spanish TTS: MMS (VITS) # ------------------------------------------------------ def load_spanish_vits(): """ Load and cache the Spanish MMS TTS model (VITS). """ global spanish_vits_cache if spanish_vits_cache is not None: return spanish_vits_cache try: model = VitsModel.from_pretrained(mms_spanish_config["model_id"]) tokenizer = AutoTokenizer.from_pretrained(mms_spanish_config["model_id"]) spanish_vits_cache = (model, tokenizer) except Exception as e: raise RuntimeError(f"Failed to load Spanish TTS model {mms_spanish_config['model_id']}: {e}") return spanish_vits_cache def run_spanish_tts(text): """ Run MMS TTS (VITS) for Spanish text. Returns (sample_rate, waveform). """ model, tokenizer = load_spanish_vits() inputs = tokenizer(text, return_tensors="pt") with torch.no_grad(): output = model(**inputs) if not hasattr(output, "waveform"): raise RuntimeError("Spanish TTS model output does not contain 'waveform'.") waveform = output.waveform.squeeze().cpu().numpy() sample_rate = 16000 return sample_rate, waveform # ------------------------------------------------------ # 6. Chinese/Japanese TTS: Coqui XTTS-v2 # ------------------------------------------------------ def load_coqui_tts(): """ Load and cache the Coqui XTTS-v2 model (multilingual). """ global coqui_tts_cache if coqui_tts_cache is not None: return coqui_tts_cache try: # If you have a GPU on HF Spaces, you can set gpu=True. # If not, set gpu=False to run on CPU (slower). coqui_tts_cache = CoquiTTS("tts_models/multilingual/multi-dataset/xtts_v2", gpu=False) except Exception as e: raise RuntimeError("Failed to load Coqui XTTS-v2 TTS: %s" % e) return coqui_tts_cache def run_coqui_tts(text, lang): """ Run Coqui TTS for Chinese or Japanese text. We specify the language code from coqui_lang_map. Returns (sample_rate, waveform). """ coqui_tts = load_coqui_tts() lang_code = coqui_lang_map[lang] # "zh" or "ja" # We must output to a file, then read it back. # Use a temporary file to store the wave. 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 # no speaker_wav, default voice ) data, sr = sf.read(tmp_name) finally: # Cleanup the temporary file if os.path.exists(tmp_name): os.remove(tmp_name) return sr, data # ------------------------------------------------------ # 7. Main Prediction Function # ------------------------------------------------------ def predict(audio, text, target_language): """ 1. Get English text (ASR if audio provided, else text). 2. Translate to target_language. 3. TTS with the chosen approach: - Spanish -> MMS TTS (VITS) - Chinese/Japanese -> Coqui XTTS-v2 """ # Step 1: English text if text.strip(): english_text = text.strip() elif audio is not None: sample_rate, audio_data = audio # Convert to float32 if needed if audio_data.dtype not in [np.float32, np.float64]: audio_data = audio_data.astype(np.float32) # Stereo -> mono if len(audio_data.shape) > 1 and audio_data.shape[1] > 1: audio_data = np.mean(audio_data, axis=1) # Resample to 16k if needed 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: if target_language == SPANISH: sr, waveform = run_spanish_tts(translated_text) else: # Chinese or Japanese sr, waveform = run_coqui_tts(translated_text, target_language) except Exception as e: return english_text, translated_text, f"TTS error: {e}" return english_text, translated_text, (sr, waveform) # ------------------------------------------------------ # 8. Gradio Interface # ------------------------------------------------------ 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=[SPANISH, CHINESE, JAPANESE], 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=( "1. Transcribes English speech using Wav2Vec2 (or takes English text).\n" "2. Translates to Spanish, Chinese, or Japanese (via Helsinki-NLP).\n" "3. Synthesizes speech:\n" " - Spanish -> facebook/mms-tts-spa (VITS)\n" " - Chinese & Japanese -> Coqui XTTS-v2 (multilingual TTS)\n\n" "Note: The Coqui model is 'tts_models/multilingual/multi-dataset/xtts_v2' and expects language codes.\n" "If you need voice cloning, set `speaker_wav` in `tts_to_file()`. By default, it uses a single generic voice." ), allow_flagging="never" ) if __name__ == "__main__": iface.launch(server_name="0.0.0.0", server_port=7860)