Update app.py

app.py

@@ -1,13 +1,9 @@
 import os
 import gradio as gr
-import whisper
-import IPython
+import soundfile as sf
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from gtts import gTTS
 
-# Load the ASR model
-asr_model = whisper.load_model("base")
-
 # Load the translation model
 translation_tokenizer = AutoTokenizer.from_pretrained("alirezamsh/small100")
 translation_model = AutoModelForSeq2SeqLM.from_pretrained("alirezamsh/small100")
@@ -24,25 +20,19 @@ available_languages = {
 def translate_audio(audio_file, target_language):
     to_lang = available_languages[target_language]
 
-    # Auto to text (ASR)
-    audio = whisper.load_audio(audio_file.name)
-    audio = whisper.pad_or_trim(audio)
-    mel = whisper.log_mel_spectrogram(audio).to(asr_model.device)
-    _, probs = asr_model.detect_language(mel)
-    options = whisper.DecodingOptions()
-    result = whisper.decode(asr_model, mel, options)
-    text = result.text
+    # Load audio
+    audio, sample_rate = sf.read(audio_file.name)
 
     # Translate the text
     translation_tokenizer.src_lang = to_lang
-    encoded_bg = translation_tokenizer(text, return_tensors="pt")
+    encoded_bg = translation_tokenizer(audio, return_tensors="pt", padding=True, truncation=True)
     generated_tokens = translation_model.generate(**encoded_bg)
-    translated_text = translation_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    translated_audio = translation_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+
+    # Save translated audio
+    output_file = "translated_audio.wav"
+    sf.write(output_file, translated_audio, sample_rate)
 
-    # Text-to-audio (TTS)
-    tts = gTTS(text=translated_text, lang=to_lang)
-    output_file = "translated_audio.mp3"
-    tts.save(output_file)
     return output_file
 
 # Gradio interface