# import torch
import gradio as gr

from transformers import pipeline

tts = pipeline(model="suno/bark-small", task= 'text-to-speech')

def text_to_speech(text):
    output = tts(text)
    audio = output["audio"]
    sampling_rate = output["sampling_rate"]

    # VERY IMPORT BUG FIX: gradio expect the transpose of the audio
    # ref: https://github.com/gradio-app/gradio/issues/5836
    return sampling_rate, audio.T

demo = gr.Interface(fn=text_to_speech, 
                    inputs="text", 
                    outputs=gr.Audio(),
                    examples=[["The King comes here tonight."]])
    
if __name__ == "__main__":
    demo.launch(show_api=False)