import os
import gradio as gr
from transformers import BarkModel, AutoProcessor
import torch

# función para el título animado 
js = """
function createGradioAnimation() {
    var container = document.createElement('div');
    container.id = 'gradio-animation';
    container.style.fontSize = '2em';
    container.style.fontWeight = 'bold';
    container.style.textAlign = 'center';
    container.style.marginBottom = '20px';

    var text = 'Audio Generation!';
    for (var i = 0; i < text.length; i++) {
        (function(i){
            setTimeout(function(){
                var letter = document.createElement('span');
                letter.style.opacity = '0';
                letter.style.transition = 'opacity 0.5s';
                letter.innerText = text[i];

                container.appendChild(letter);

                setTimeout(function() {
                    letter.style.opacity = '1';
                }, 50);
            }, i * 250);
        })(i);
    }

    var gradioContainer = document.querySelector('.gradio-container');
    gradioContainer.insertBefore(container, gradioContainer.firstChild);

    return 'Animation created';
}
"""
# Desactivar paralelismo en tokenizers para evitar advertencias
os.environ["TOKENIZERS_PARALLELISM"] = "false"

device = "cuda" if torch.cuda.is_available() else "cpu"

def say_something(text):
    processor = AutoProcessor.from_pretrained("suno/bark-small")

    model = BarkModel.from_pretrained("suno/bark-small").to(device)

    inputs = processor(text)
    audio_array = model.generate(**inputs)
    audio_array = audio_array.cpu().numpy().squeeze()

    # Devuelve el audio como un numpy array junto con el sample rate
    sample_rate = model.generation_config.sample_rate
    return (sample_rate, audio_array)
    

with gr.Blocks(js=js) as demo:
    textBox = gr.Textbox(label="Text", value="Hello! [laugths]. This is a test!")
    button = gr.Button("Generate Audio", variant="primary")
    audio_output = gr.Audio(label="Generated Audio")
    button.click(say_something, inputs=textBox, outputs=audio_output)


demo.launch()