Update app.py

app.py

@@ -1,13 +1,20 @@
 import gradio as gr
 import numpy as np
 import torch
+from vosk import KaldiRecognizer, Model
 from datasets import load_dataset
-from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor, pipeline
+from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
-# Load a lightweight automatic speech recognition model (vosk)
-asr_pipe = pipeline("automatic-speech-recognition", model="alphacep/kaldi-ru", device=device)
+# Load Vosk automatic speech recognition model
+vosk_model = Model("elizabetvaganova/speech-to-speech-translation-vaganova") 
+
+def recognize_speech(audio):
+    recognizer = KaldiRecognizer(vosk_model, 16000)
+    recognizer.AcceptWaveform(audio.data)
+    result = recognizer.FinalResult()
+    return result["text"]
 
 # Load a lightweight text-to-speech checkpoint and speaker embeddings
 processor = SpeechT5Processor.from_pretrained("ttskit/ttskit-tts-ljspeech")
@@ -19,8 +26,10 @@ embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validat
 speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
 
 def translate(audio):
-    outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"})
-    return outputs["text"]
+    recognizer = KaldiRecognizer(vosk_model, 16000)
+    recognizer.AcceptWaveform(audio.data)
+    result = recognizer.FinalResult()
+    return result["text"]
 
 def synthesise(text):
     inputs = processor(text=text, return_tensors="pt")
@@ -60,4 +69,4 @@ file_translate = gr.Interface(
 with demo:
     gr.TabbedInterface([mic_translate, file_translate], ["Microphone", "Audio File"])
 
-demo.launch()
+demo.launch()