from transformers import AutoProcessor, BarkModel # Para poder seleccionar voz de hombre o de mujer y poder seleccionar el idioma
import gradio as gr

model_name="suno/bark-small"

processor = AutoProcessor.from_pretrained(model_name)
model = BarkModel.from_pretrained(model_name)

# Selección de voz - Valores: https://suno-ai.notion.site/8b8e8749ed514b0cbf3f699013548683
voice_preset = "v2/es_speaker_0"

def tts(frase):
    # Obtenemos el procesador
    inputs = processor(frase, voice_preset=voice_preset) # la frase de entrada se procesa usando la voz seleccionada

    # Realiza la inferencia a partir de procesador y el modelo
    audio_array = model.generate(**inputs, pad_token_id=100) # El modelo genera el audio a partir de los inputs procesados. Pad_token es el relleno que sirve para cuando trabajamos con secuencias de texto de diferentes longitudes y transformarlas para que todas tengan la misma longitud y poder procesarlas en lotes. Pad_token_id=100 especifica el identificador del token de relleno
    audio_array = audio_array.cpu().numpy().squeeze() # El array del audio se castea a un array de numpy

    sample_rate = model.generation_config.sample_rate # Obtenemos la tasa de muestreo

    return sample_rate,audio_array

demo = gr.Interface(
    tts,
    inputs=gr.Text(label="Teclea el texto a pronunciar"), # Entrada campo de texto
    outputs=gr.Audio(label="audio generado"), # Salida audio
    title="De texto a voz",
)

demo.launch()