Update app.py

app.py

@@ -1,15 +1,18 @@
-import streamlit as st
-import base64
-import io
 from huggingface_hub import InferenceClient
 from audiorecorder import audiorecorder
 import speech_recognition as sr
-from pydub import AudioSegment
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import streamlit as st
+import base64
+import io
 import torch
+from transformers import VitsProcessor, VitsModel, AutoTokenizer
+import numpy as np
+from scipy.io.wavfile import write
+from pydub import AudioSegment
 
+processor = VitsProcessor.from_pretrained("facebook/mms-tts-spa")
+model = VitsModel.from_pretrained("facebook/mms-tts-spa")
 tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-spa")
-model = AutoModelForSeq2SeqLM.from_pretrained("facebook/mms-tts-spa")
 
 pre_prompt_text = "Eres una IA conductual, tus respuestas deberán ser breves, estóicas y humanistas."
 
@@ -55,42 +58,13 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def generate(audio_text, history, temperature=None, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
-    client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
-
-    temperature = float(temperature) if temperature is not None else 0.9
-    temperature = max(temperature, 1e-2)
-    top_p = float(top_p)
-
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=42)
-
+def generate(audio_text, history):
     formatted_prompt = format_prompt(audio_text, history)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=True)
-    response = ""
-
-    for response_token in stream:
-        response += response_token.token.text
-    
-    response = ' '.join(response.split()).replace('</s>', '')
-    audio_file = text_to_speech(response, speed=1.3)
-    return response, audio_file
-
-def text_to_speech(text, speed=1.3):
-    inputs = tokenizer(text, return_tensors="pt", padding="longest", truncation=True)
+    inputs = tokenizer(formatted_prompt, return_tensors="pt", padding=True, truncation=True)
     with torch.no_grad():
         output = model.generate(**inputs)
-    audio = output[0]
-    audio_bytes = io.BytesIO()
-    audio_tensor = torch.tensor(audio).float()
-    torch.save(audio_tensor, audio_bytes)
-    audio_bytes.seek(0)
-    return base64.b64encode(audio_bytes.read()).decode()
+    audio = output['audio']
+    return audio
 
 def main():
     audio_data = audiorecorder("Presiona para hablar", "Deteniendo la grabación...")
@@ -101,12 +75,24 @@ def main():
         audio_text = recognize_speech("audio.wav")
 
         if audio_text:
-            response, audio_file = generate(audio_text, history=st.session_state.history)  
+            audio_file = generate(audio_text, history=st.session_state.history)
 
             if audio_file is not None:
-                st.markdown(
-                    f"""<audio autoplay="autoplay" controls="controls" src="data:audio/mp3;base64,{base64.b64encode(audio_file.read()).decode()}" type="audio/mp3" id="audio_player"></audio>""",
-                    unsafe_allow_html=True)
+                # Guardar el archivo WAV
+                write("output.wav", processor.feature_extractor.sampling_rate, audio_file)
+
+                # Convertir el archivo WAV a MP3 utilizando pydub
+                audio = AudioSegment.from_wav("output.wav")
+                audio.export("output.mp3", format="mp3")
+
+                # Leer el archivo MP3 y mostrarlo en Streamlit
+                with open("output.mp3", "rb") as file:
+                    audio_bytes = file.read()
+                st.audio(audio_bytes, format="audio/mp3")
+
+                # Eliminar archivos temporales (opcional)
+                os.remove("output.wav")
+                os.remove("output.mp3")
 
 if __name__ == "__main__":
-    main()
+    main()