Update app.py

app.py

@@ -1,61 +1,49 @@
 import gradio as gr
-import whisper
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+import whisper
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
 from TTS.api import TTS
 
-# Load Whisper model (better accuracy with 'medium')
-asr_model = whisper.load_model("medium")
+# Load Whisper for Hindi speech recognition
+asr_model = whisper.load_model("base")
 
-# Load NLLB Hindi to English translator
-translation_model_name = "facebook/nllb-200-distilled-600M"
-translator_tokenizer = AutoTokenizer.from_pretrained(translation_model_name)
-translator_model = AutoModelForSeq2SeqLM.from_pretrained(translation_model_name)
+# Load translation model and tokenizer
+translation_model_name = "facebook/m2m100_418M"
+translator_tokenizer = M2M100Tokenizer.from_pretrained(translation_model_name)
+translator_model = M2M100ForConditionalGeneration.from_pretrained(translation_model_name)
 
-# Load English TTS model
-tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False, gpu=False)
+# Load TTS model for English speech
+tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False, gpu=torch.cuda.is_available())
 
-# Utility: Get Hindi text from audio
-def speech_to_text(audio_path):
+def translate_hindi_speech(audio_path):
+    # Transcribe Hindi speech
     result = asr_model.transcribe(audio_path, language="hi")
-    return result["text"]
-
-# Utility: Translate Hindi to English
-def translate_hi_to_en(text_hi):
-    inputs = translator_tokenizer(text_hi, return_tensors="pt")
-    translated_tokens = translator_model.generate(**inputs, forced_bos_token_id=translator_tokenizer.lang_code_to_id["eng_Latn"])
-    translated_text = translator_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
-    return translated_text
-
-# Main app logic
-def translate_audio(audio):
-    if audio is None:
-        return "No audio input", "", None
-    
-    # Step 1: Convert Hindi speech to Hindi text
-    hindi_text = speech_to_text(audio)
+    hindi_text = result["text"]
 
-    # Step 2: Translate to English
-    english_text = translate_hi_to_en(hindi_text)
+    # Translate Hindi to English
+    translator_tokenizer.src_lang = "hi"
+    encoded = translator_tokenizer(hindi_text, return_tensors="pt")
+    generated_tokens = translator_model.generate(**encoded, forced_bos_token_id=translator_tokenizer.get_lang_id("en"))
+    english_text = translator_tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
 
-    # Step 3: Generate English speech
-    english_audio_path = "output.wav"
-    tts.tts_to_file(text=english_text, file_path=english_audio_path)
+    # Convert English text to speech
+    output_audio_path = "output_en.wav"
+    tts.tts_to_file(text=english_text, file_path=output_audio_path)
 
-    return hindi_text, english_text, english_audio_path
+    return english_text, output_audio_path
 
-# Gradio UI
+# Gradio interface
 interface = gr.Interface(
-    fn=translate_audio,
-    inputs=gr.Audio(source="microphone", type="filepath"),
+    fn=translate_hindi_speech,
+    inputs=gr.Audio(type="filepath", label="Speak in Hindi"),
     outputs=[
-        gr.Textbox(label="Hindi Transcript"),
-        gr.Textbox(label="English Translation"),
-        gr.Audio(label="English Speech")
+        gr.Text(label="English Translation"),
+        gr.Audio(type="filepath", label="English Speech Output")
     ],
-    title="Hindi to English Speech Translator",
-    description="🎤 Speak in Hindi → 📄 Translated English Text → 🔊 Spoken English Output"
+    title="🎙️ Hindi to English Speech Translator",
+    description="Speak in Hindi → See English translation + hear English output.",
+    live=False,
 )
 
-if __name__ == "__main__":
-    interface.launch()
+# Launch the app
+interface.launch()