Update app.py

app.py

@@ -1,16 +1,15 @@
-from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
-import streamlit as st
-import base64
-import io
 from huggingface_hub import InferenceClient
 from audiorecorder import audiorecorder
 import speech_recognition as sr
+from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
 from datasets import load_dataset
+import streamlit as st
+import base64
+import io
 import torch
+import os
 
-processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
-model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
-vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
+model = Wav2Vec2ForCTC.from_pretrained("facebook/mms-tts-spa")
 pre_prompt_text = "You are a behavioral AI, your answers should be brief, stoic and humanistic."
 
 if "history" not in st.session_state:
@@ -18,7 +17,6 @@ if "history" not in st.session_state:
 
 if "pre_prompt_sent" not in st.session_state:
     st.session_state.pre_prompt_sent = False
-
 def recognize_speech(audio_data, show_messages=True):
     recognizer = sr.Recognizer()
     audio_recording = sr.AudioFile(audio_data)
@@ -78,17 +76,14 @@ def generate(audio_text, history, temperature=None, max_new_tokens=512, top_p=0.
         response += response_token.token.text
     
     response = ' '.join(response.split()).replace('</s>', '')
-    audio_file = text_to_speech(response)
+    audio_file = text_to_speech(response, speed=1.3)
     return response, audio_file
 
 def text_to_speech(text):
-    inputs = processor(text=text, return_tensors="pt", target_language="es")
+    with torch.no_grad():
+        logits = model.generate(text)
 
-    embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
-    speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
-    speech = model.generate_speech(inputs.input_ids, speaker_embeddings, vocoder=vocoder)
-    with sf.SoundFile("speech.wav", "w", 16000, 1) as file:
-        file.write(speech.cpu().numpy().astype('float32'))
+    sf.write("speech.wav", logits.numpy(), samplerate=16000)
 
     with open("speech.wav", "rb") as audio_file:
         encoded_audio = base64.b64encode(audio_file.read()).decode()