import gradio as gr
from transformers import pipeline
from datasets import load_dataset
import torchaudio
import torch


# Cargar el modelo de traducción
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-es-en")
synthesiser = pipeline("text-to-speech", "microsoft/speecht5_tts")

# Cargar voces preentrenadas
embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")


# Función principal
def translate_and_speak(text):
    # Traducir de español a inglés
    translated = translator(text, src_lang="es", tgt_lang="en")[0]["translation_text"]

    # Usar una voz preentrenada
    speaker_embeddings = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)

    # Generar audio
    speech = synthesiser(translated, forward_params={"speaker_embeddings": speaker_embeddings})

    # Extraer waveform y sampling_rate del resultado
    waveform = torch.tensor(speech["audio"])  # Convertir a tensor de PyTorch
    sampling_rate = speech["sampling_rate"]

    # Guardar el audio temporalmente
    output_path = "output.wav"
    torchaudio.save(output_path, waveform.unsqueeze(0), sample_rate=sampling_rate)  # Unsqueeze para añadir batch dimension

    return translated, output_path


# Crear la interfaz
interface = gr.Interface(
    fn=translate_and_speak,
    inputs=gr.Textbox(label="Texto en Español"),
    outputs=[
        gr.Textbox(label="Texto Traducido (Inglés)"),
        gr.Audio(label="Audio Generado (Inglés)")
    ],
    title="Traductor y Generador de Voz",
    description="Escribe un texto en español y este será traducido al inglés y leído en voz alta usando IA."
)

# Lanzar la aplicación
if __name__ == "__main__":
    interface.launch()