Update app.py

app.py

@@ -1,64 +1,45 @@
-import os
-import tempfile
 import gradio as gr
+import os
 import whisper
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from gtts import gTTS
-import base64
-
-# Load the models and tokenizer
-whisper_model = whisper.load_model("base")
-tokenizer = AutoTokenizer.from_pretrained("alirezamsh/small100")
-model = AutoModelForSeq2SeqLM.from_pretrained("alirezamsh/small100")
 
-def translate_audio(input_file, to_lang):
-    # Load the audio file
+def speech_to_speech(input_audio, to_lang):
+    # Save the uploaded audio file
+    input_file = "input_audio" + os.path.splitext(input_audio.name)[1]
+    input_audio.save(input_file)
+    
+    # Speech-to-Text (STT)
+    model = whisper.load_model("base")
     audio = whisper.load_audio(input_file)
     audio = whisper.pad_or_trim(audio)
-    mel = whisper.log_mel_spectrogram(audio).to(whisper_model.device)
-    
-    # Detect language using Whisper
-    _, probs = whisper_model.detect_language(mel)
-    lang = max(probs, key=probs.get)
-    
-    # Convert audio to text
+    mel = whisper.log_mel_spectrogram(audio).to(model.device)
+    _, probs = model.detect_language(mel)
     options = whisper.DecodingOptions()
-    result = whisper.decode(whisper_model, mel, options)
+    result = whisper.decode(model, mel, options)
     text = result.text
+    lang = max(probs, key=probs.get)
     
-    # Translate the text
+    # Translate
+    tokenizer = AutoTokenizer.from_pretrained("alirezamsh/small100")
+    model = AutoModelForSeq2SeqLM.from_pretrained("alirezamsh/small100")
     tokenizer.src_lang = lang
+    tokenizer.tgt_lang = to_lang
     encoded_bg = tokenizer(text, return_tensors="pt")
     generated_tokens = model.generate(**encoded_bg)
     translated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
     
-    # Convert translated text to audio
+    # Text-to-Speech (TTS)
     tts = gTTS(text=translated_text, lang=to_lang)
-    temp_output_file = tempfile.NamedTemporaryFile(suffix=".mp3").name
-    tts.save(temp_output_file)
-    
-    # Load audio data from file
-    audio_data = open(temp_output_file, "rb").read()
+    output_file = "output_audio.mp3"
+    tts.save(output_file)
     
-    # Convert audio data to Base64 string
-    audio_base64 = base64.b64encode(audio_data).decode("utf-8")
-    
-    return audio_base64
-
-def translate_audio_interface(input_file, to_lang):
-    return translate_audio(input_file, to_lang)
+    return output_file
 
-iface = gr.Interface(
-    fn=translate_audio_interface,
-    inputs=["file", "text"],
-    outputs="text",
-    title="Audio Translation",
-    description="Uploadd an MP3 file and select the target language for translation.",
-    examples=[
-        ["audio_example.mp3", "en"],
-        ["speech_sample.mp3", "fr"],
-    ]
-)
+languages = ["ru", "fr", "es", "de"]  # Example languages: Russian, French, Spanish, German
+file_input = gr.inputs.File(label="Upload Audio", type="audio")
+dropdown = gr.inputs.Dropdown(languages, label="Translation Language")
+audio_output = gr.outputs.Audio(type="file", label="Translated Voice")
 
-iface.launch()
+gr.Interface(fn=speech_to_speech, inputs=[file_input, dropdown], outputs=audio_output, title="Speech-to-Speech Translator", description="Upload an audio file (MP3, WAV, or FLAC) and choose the target language for translation.", theme="default").launch()