import torch
import spaces
import gradio as gr

from transformers import MusicgenForConditionalGeneration
model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
sampling_rate = model.config.audio_encoder.sampling_rate

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

from transformers import AutoProcessor
processor = AutoProcessor.from_pretrained("facebook/musicgen-small")

@spaces.GPU
def generate_music(desc):
    inputs = processor(text=[desc], padding=True, return_tensors="pt")
    audio_values = model.generate(**inputs.to(device), do_sample=True, guidance_scale=3, max_new_tokens=256)
    return sampling_rate, audio_values[0][0].cpu().numpy()
    
# gr.Interface(fn=generate_music, inputs="text", outputs="audio").launch()

with gr.Blocks() as app:
    with gr.Row():
        music_desc = gr.TextArea(label="Music Description")
        music_player = gr.Audio(label="Play My Tune")

    gen_btn = gr.Button("Get Some Tune!!")
    gen_btn.click(fn=generate_music, inputs=[music_desc], outputs=[music_player])
    
if __name__ == '__main__':
    app.launch()