Update app.py

app.py

@@ -1,46 +1,61 @@
 import gradio as gr
 import whisper
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 from TTS.api import TTS
 
-# Load models
-print("Loading Whisper (for Hindi STT)...")
-whisper_model = whisper.load_model("small")
+# Load Whisper model (better accuracy with 'medium')
+asr_model = whisper.load_model("medium")
 
-print("Loading MarianMT (for Hindi to English)...")
-translator_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-hi-en")
-translator_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-hi-en")
+# 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)
 
-print("Loading TTS model (Tacotron2 + HiFi-GAN)...")
+# Load English TTS model
 tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False, gpu=False)
 
-# Pipeline function
-def hindi_speech_to_english_audio(audio):
-    # Step 1: Hindi Speech to Text
-    result = whisper_model.transcribe(audio, language="hi")
-    hindi_text = result["text"]
+# Utility: Get Hindi text from audio
+def speech_to_text(audio_path):
+    result = asr_model.transcribe(audio_path, language="hi")
+    return result["text"]
 
-    # Step 2: Hindi to English Translation
-    inputs = translator_tokenizer(hindi_text, return_tensors="pt", padding=True)
-    translated_tokens = translator_model.generate(**inputs)
-    english_text = translator_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+# 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
 
-    # Step 3: English Text to Speech
-    tts.tts_to_file(text=english_text, file_path="output.wav")
+# 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)
 
-    return english_text, "output.wav"
+    # Step 2: Translate to English
+    english_text = translate_hi_to_en(hindi_text)
 
-# Gradio Interface
+    # Step 3: Generate English speech
+    english_audio_path = "output.wav"
+    tts.tts_to_file(text=english_text, file_path=english_audio_path)
+
+    return hindi_text, english_text, english_audio_path
+
+# Gradio UI
 interface = gr.Interface(
-    fn=hindi_speech_to_english_audio,
-    inputs=gr.Audio(sources=["microphone"], type="filepath", label="🎙️ Speak in Hindi"),
+    fn=translate_audio,
+    inputs=gr.Audio(source="microphone", type="filepath"),
     outputs=[
-        gr.Textbox(label="🔤 Translated English Text"),
-        gr.Audio(type="filepath", label="🗣️ English Audio Output")
+        gr.Textbox(label="Hindi Transcript"),
+        gr.Textbox(label="English Translation"),
+        gr.Audio(label="English Speech")
     ],
-    title="Hindi Speech to English Audio Translator",
-    description="🎧 Speak in Hindi and hear it back in English!",
+    title="Hindi to English Speech Translator",
+    description="🎤 Speak in Hindi → 📄 Translated English Text → 🔊 Spoken English Output"
 )
 
-# Run app
-interface.launch()
+if __name__ == "__main__":
+    interface.launch()