import gradio as gr
from transformers import BarkModel, AutoProcessor
import torch
from scipy.io.wavfile import write as write_wav
import os


### if you run on GPU use the following code: ###
# device = "cuda" if torch.cuda.is_available() else "cpu"
# model = BarkModel.from_pretrained("suno/bark-small", torch_dtype=torch.float16).to(device)
# model.enable_cpu_offload()

### if you run on CPU use the following code: ###
device = "cpu"
# load in fp16
model = BarkModel.from_pretrained("suno/bark-small").to(device)
processor = AutoProcessor.from_pretrained("suno/bark")

voice_preset = "v2/en_speaker_3"

def generate_audio(text, preset, output_file_name="bark_generation"):
    file_name = output_file_name + ".wav"
    inputs = processor(text, voice_preset=preset)
    audio_array = model.generate(**inputs)
    audio_array = audio_array.cpu().numpy().squeeze()   
    sample_rate = model.generation_config.sample_rate
    write_wav(file_name, sample_rate, audio_array)
    return file_name


#Bark Presets List
presets = ["v2/en_speaker_0","v2/en_speaker_1", "v2/en_speaker_2", "v2/en_speaker_3",  "v2/en_speaker_4",  "v2/en_speaker_5",  "v2/en_speaker_6"]

#Gradio Interface
iface = gr.Interface(fn=generate_audio, inputs=["text", gr.components.Dropdown(choices=presets), "text"], outputs="audio")
iface.launch()