Update app.py

app.py

@@ -4,40 +4,42 @@ import whisper
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from gtts import gTTS
 
+# Load models
+model_stt = whisper.load_model("base")
+model_translation = AutoModelForSeq2SeqLM.from_pretrained("alirezamsh/small100")
+tokenizer_translation = AutoTokenizer.from_pretrained("alirezamsh/small100")
+
 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(model.device)
-    _, probs = model.detect_language(mel)
+    mel = whisper.log_mel_spectrogram(audio).to(model_stt.device)
+    _, probs = model_stt.detect_language(mel)
     options = whisper.DecodingOptions()
-    result = whisper.decode(model, mel, options)
+    result = whisper.decode(model_stt, mel, options)
     text = result.text
     lang = max(probs, key=probs.get)
-    
+
     # 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]
-    
+    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)
-    
+
     return output_file
 
 languages = ["ru", "fr", "es", "de"]  # Example languages: Russian, French, Spanish, German
-file_input = gr.inputs.File(label="Upload Audio", type="audio")
+file_input = gr.inputs.File(label="Upload Audio", accept="audio/*")
 dropdown = gr.inputs.Dropdown(languages, label="Translation Language")
 audio_output = gr.outputs.Audio(type="file", label="Translated Voice")