Update app.py

app.py

@@ -2,81 +2,54 @@ import gradio as gr
 from transformers import pipeline
 from gtts import gTTS
 import io
-import os
-import soundfile as sf
+from pydub import AudioSegment
+import numpy as np
 
-# Load Amharic to English model
+# Load translation models
 am_en_translator = pipeline("text2text-generation", model="Atnafu/Amharic-English-MT")
-
-# Load English to Amharic model
 en_am_translator = pipeline("text2text-generation", model="Atnafu/English-Amharic-MT")
-
-# Load Speech-to-Text (STT) models
-# For Amharic STT
-am_stt_pipeline = pipeline("automatic-speech-recognition", model="speechbrain/asr-wav2vec2-commonvoice-amharic")
-# For English STT
+# English STT model
 en_stt_pipeline = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h")
 
-def text_to_audio_en(text):
-    if not text:
-        return None
-    tts = gTTS(text=text, lang='en')
-    audio_fp = io.BytesIO()
-    tts.write_to_fp(audio_fp)
-    audio_fp.seek(0)
-    # Gradio expects (sample_rate, audio_data) for audio output
-    # We'll save to a temp file and load it to get sample rate and data
-    # Or, more directly, use soundfile to read the BytesIO into the correct format if possible
-    # For simplicity, let's return the BytesIO directly and let Gradio handle it.
-    # In some Gradio versions, returning BytesIO of mp3 works directly.
-    # If not, a temp .wav file conversion might be needed.
-    
-    # Let's try to convert to wav using pydub for better Gradio compatibility if needed, 
-    # but for now, return BytesIO (which Gradio usually handles)
-    return audio_fp.getvalue() # Returning bytes directly
-
-def text_to_audio_am(text):
+def text_to_audio(text, lang):
     if not text:
         return None
-    tts = gTTS(text=text, lang='am')
-    audio_fp = io.BytesIO()
-    tts.write_to_fp(audio_fp)
-    audio_fp.seek(0)
-    return audio_fp.getvalue() # Returning bytes directly
+    tts = gTTS(text=text, lang=lang)
+    mp3_bytes = io.BytesIO()
+    tts.write_to_fp(mp3_bytes)
+    mp3_bytes.seek(0)
+    # Convert to WAV for Gradio compatibility
+    audio = AudioSegment.from_file(mp3_bytes, format="mp3")
+    wav_bytes = io.BytesIO()
+    audio.export(wav_bytes, format="wav")
+    wav_bytes.seek(0)
+    # Gradio expects sample_rate, np.array
+    sample_rate = audio.frame_rate  # e.g., 22050
+    audio_np = np.frombuffer(wav_bytes.read(), np.int16)
+    return (sample_rate, audio_np)
 
 def am_to_en_full(amharic_text_input, amharic_audio_input):
-    if amharic_audio_input is not None:
-        # Transcribe audio to text
-        amharic_text = am_stt_pipeline(amharic_audio_input)['text']
-    elif amharic_text_input:
+    # Only support text input for Amharic
+    if amharic_text_input:
         amharic_text = amharic_text_input
+    elif amharic_audio_input:
+        return "Sorry, Amharic voice feature not supported yet.", None
     else:
-        return "", None # No input, return empty text and no audio
-
-    # Translate Amharic to English
+        return "", None
     translated_en_text = am_en_translator(amharic_text)[0]['generated_text']
-    
-    # Convert translated English text to audio
-    translated_en_audio = text_to_audio_en(translated_en_text)
-    
-    return translated_en_text, (44100, translated_en_audio) # Assuming 44100 Hz sample rate for gTTS, adjust if actual is different
+    translated_en_audio = text_to_audio(translated_en_text, 'en')
+    return translated_en_text, translated_en_audio
 
 def en_to_am_full(english_text_input, english_audio_input):
     if english_audio_input is not None:
-        # Transcribe audio to text
         english_text = en_stt_pipeline(english_audio_input)['text']
     elif english_text_input:
         english_text = english_text_input
     else:
-        return "", None # No input, return empty text and no audio
-
-    # Translate English to Amharic
+        return "", None
     translated_am_text = en_am_translator(english_text)[0]['generated_text']
-    
-    # Convert translated Amharic text to audio
-    translated_am_audio = text_to_audio_am(translated_am_text)
-
-    return translated_am_text, (44100, translated_am_audio) # Assuming 44100 Hz sample rate for gTTS, adjust if actual is different
+    translated_am_audio = text_to_audio(translated_am_text, 'am')
+    return translated_am_text, translated_am_audio
 
 with gr.Blocks(title="Amharic-English Two-Way Translator with Voice") as demo:
     gr.Markdown("# Amharic-English Two-Way Translator with Voice")
@@ -98,7 +71,6 @@ with gr.Blocks(title="Amharic-English Two-Way Translator with Voice") as demo:
             am_text_out = gr.Textbox(lines=3, label="Amharic Text Output", interactive=False)
             am_audio_out = gr.Audio(label="Amharic Audio Output")
 
-    # Event listeners
     am_en_btn.click(
         am_to_en_full,
         inputs=[am_text_in, am_audio_in],
@@ -110,4 +82,4 @@ with gr.Blocks(title="Amharic-English Two-Way Translator with Voice") as demo:
         outputs=[am_text_out, am_audio_out]
     )
 
-demo.launch(share=True)
+demo.launch()