Update app.py

app.py

@@ -3,56 +3,56 @@ import os
 import whisper
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from gtts import gTTS
-import numpy as np
+import IPython.display as ipd
 
-# Load models
-model_stt = whisper.load_model("base")
-model_translation = AutoModelForSeq2SeqLM.from_pretrained("alirezamsh/small100")
-tokenizer_translation = AutoTokenizer.from_pretrained("alirezamsh/small100")
+# Load Whisper STT model
+whisper_model = whisper.load_model("base")
 
-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)
+# Load translation models
+tokenizer = AutoTokenizer.from_pretrained("alirezamsh/small100")
+model = AutoModelForSeq2SeqLM.from_pretrained("alirezamsh/small100")
 
-    # Speech-to-Text (STT)
-    audio = whisper.load_audio(input_file)
+def translate_speech(audio_file, target_lang):
+    # Load audio
+    audio = whisper.load_audio(audio_file)
     audio = whisper.pad_or_trim(audio)
-    mel = whisper.log_mel_spectrogram(audio).to(model_stt.device)
-    _, probs = model_stt.detect_language(mel)
+    mel = whisper.log_mel_spectrogram(audio).to(whisper_model.device)
+
+    # Detect language
+    _, probs = whisper_model.detect_language(mel)
+    lang = max(probs, key=probs.get)
+
+    # Decode audio into text
     options = whisper.DecodingOptions()
-    result = whisper.decode(model_stt, mel, options)
+    result = whisper.decode(whisper_model, mel, options)
     text = result.text
-    lang = max(probs, key=probs.get)
 
-    # Translate
-    tokenizer_translation.src_lang = lang
-    tokenizer_translation.tgt_lang = to_lang
-    encoded_bg = tokenizer_translation(text, return_tensors="pt")
-    generated_tokens = model_translation.generate(**encoded_bg)
-    translated_text = tokenizer_translation.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-
-    # Text-to-Speech (TTS)
-    tts = gTTS(text=translated_text, lang=to_lang)
-    output_file = "output_audio.mp3"
-    tts.save(output_file)
-
-    # Load output audio as numpy array
-    audio_np = np.array(output_file)
-
-    return audio_np
-
-languages = ["ru", "fr", "es", "de"]  # Example languages: Russian, French, Spanish, German
-file_input = gr.inputs.File(label="Upload Audio", accept="audio/*")
-dropdown = gr.inputs.Dropdown(languages, label="Translation Language")
-audio_output = gr.outputs.Audio(label="Translated Voice", type="numpy")
-
-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()
+    # Translate text
+    tokenizer.src_lang = lang
+    encoded_text = tokenizer(text, return_tensors="pt")
+    generated_tokens = model.generate(**encoded_text)
+    translated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+
+    # Text-to-speech (TTS)
+    tts = gTTS(text=translated_text, lang=target_lang)
+    audio_path = "translated_audio.mp3"
+    tts.save(audio_path)
+
+    return audio_path
+
+def translate_speech_interface(audio, target_lang):
+    audio_path = "recorded_audio.wav"
+    with open(audio_path, "wb") as f:
+        f.write(audio.read())
+
+    translated_audio = translate_speech(audio_path, target_lang)
+    translated_audio = open(translated_audio, "rb")
+
+    return translated_audio
+
+# Define the Gradio interface
+audio_recording = gr.inputs.Audio(source="microphone", type="wav", label="Record your speech")
+target_language = gr.inputs.Dropdown(["en", "ru", "fr"], label="Target Language")
+output_audio = gr.outputs.Audio(type="audio/mpeg", label="Translated Audio")
 
+gr.Interface(fn=translate_speech_interface, inputs=[audio_recording, target_language], outputs=output_audio, title="Speech Translator").launch()