Update app.py

app.py

@@ -1,225 +1,13 @@
-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)
+torch
+transformers>=4.33.0
+gradio
+librosa
+numpy
+scipy
+accelerate
+sentencepiece
+soundfile
+datasets
+TTS
+git+https://github.com/myshell-ai/MeloTTS-Chinese.git
+git+https://github.com/myshell-ai/MeloTTS-Japanese.git