import os
import gradio as gr
import torch
from fairseq.checkpoint_utils import load_model_ensemble_and_task_from_hf_hub
from huggingface_hub import login

# Obtener el token desde las variables de entorno
hf_token = os.getenv("HF_TOKEN")
if hf_token is None:
    raise ValueError("Debe proporcionar un token de Hugging Face en las variables de entorno.")

# Ingresar el token
login(hf_token)

# Cargar el modelo
models, cfg, task = load_model_ensemble_and_task_from_hf_hub(
    "gitgato/speecht5_tts_mabama_es"
)
model = models[0]

# Función para generar la salida de texto a voz
def text_to_speech(text):
    # Preprocesamiento del texto
    tokens = task.source_dictionary.encode_line(text, add_if_not_exist=False)
    
    # Generar salida de audio
    with torch.no_grad():
        sample = {"net_input": {"src_tokens": tokens.unsqueeze(0).long()}}
        generator = task.build_generator([model], cfg.generation)
        audio = task.inference_step(generator, [model], sample)
    
    return audio[0][0].numpy()

# Crear interfaz de Gradio
iface = gr.Interface(
    fn=text_to_speech,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Ingrese el texto aquí..."),
    outputs=gr.outputs.Audio(type="numpy", label="Output Audio")
)

if __name__ == "__main__":
    iface.launch()