Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import numpy as np
 import torch
 from transformers import pipeline, VitsModel, AutoTokenizer, AutoTokenizer
+from transformers import SpeechT5ForTextToSpeech, SpeechT5HifiGan, SpeechT5Processor
 
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
@@ -42,8 +43,19 @@ translation_models = {
 
 asr_pipe = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device)
 
-vist_model = VitsModel.from_pretrained("facebook/mms-tts-spa")
-vist_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-spa")
+#vist_model = VitsModel.from_pretrained("facebook/mms-tts-spa")
+#vist_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-spa")
+
+
+model = SpeechT5ForTextToSpeech.from_pretrained(
+    "juangtzi/speecht5_finetuned_voxpopuli_es"
+)
+checkpoint = "microsoft/speecht5_tts"
+processor = SpeechT5Processor.from_pretrained(checkpoint)
+vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
+
+speaker_embeddings2 = np.load('speaker_embeddings.npy')
+speaker_embeddings2 = torch.tensor(speaker_embeddings2)
 
 lang_detector = pipeline("text-classification", model="papluca/xlm-roberta-base-language-detection")
 
@@ -73,16 +85,18 @@ def synthesise(text):
     else:
         text = text
     print(text)
-    inputs = vist_tokenizer(text, return_tensors="pt")
-    with torch.no_grad():
-        output = vist_model(**inputs).waveform[0]
+    inputs = processor(text, return_tensors="pt")
+    output = model.generate_speech(inputs["input_ids"], speaker_embeddings2, vocoder=vocoder)
     return output
 
 def speech_to_speech_translation(audio):
     translated_text = translate(audio)
     synthesised_speech = synthesise(translated_text)
-    synthesised_speech = (synthesised_speech.numpy() * 32767).astype(np.int16)
-    return 16000, synthesised_speech
+    audio_data = synthesised_speech.cpu().numpy()
+    audio_data = np.squeeze(audio_data)
+    audio_data = audio_data / np.max(np.abs(audio_data))
+    sample_rate = 16000
+    return (sample_rate, audio_data)
 
 title = "Cascaded STST"
 description = """