Update app.py

app.py

@@ -17,19 +17,23 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 asr_pipe = pipeline("automatic-speech-recognition", model="asapp/sew-d-mid-400k-ft-ls100h", device=device)
 
 # Translate en to rus text
-translation_en_to_rus = pipeline("translation", model="t5-small")
+translation_en_to_rus = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ru")
 
 # Create speech from rus text
-vits_model = VitsModel.from_pretrained("facebook/mms-tts-rus")
-vits_tokenizer = VitsTokenizer.from_pretrained("facebook/mms-tts-rus")
+#vits_model = VitsModel.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
+#vits_tokenizer = VitsTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
+
+model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
+tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ru-en")
 
 def transform_audio_to_speech_en(audio):
     outputs = asr_pipe(audio, max_new_tokens=256, generate_kwargs={"task": "translate"})
     return outputs["text"]
 
 def translator(text):
-  translated_text = translation_en_to_rus(text)
-  return translated_text[0]['translation_text']
+    translated = translator_model.generate(**translator_tokenizer(text, return_tensors="pt", padding=True))
+    translated_text = translator_tokenizer.decode(translated[0], skip_special_tokens=True)
+    return translated_text
 
 def synthesise(translated_text):
     translated_text = translator(translated_text)