import torch
import gradio as gr
from audiocraft.models import musicgen
from audiocraft.utils.notebook import display_audio

gcPU = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

audiocraft__music_gen__large = musicgen.MusicGen.get_pretrained('facebook/musicgen-medium', device=gcPU)
audiocraft__music_gen__large.set_generation_params(duration=15) # for sample_rate=24KHz to get actual duration of 20sec

def func__genA(ip):
    print('[p] :: ... processing ...')
    print(f"[logs] :: processing on '{gcPU}' device")
    try:
      op = audiocraft__music_gen__large.generate(ip,progress=True)
      return 'genA',(24000,op)  # 1sec=32KHz(here-audiocraft) i.e to get the audio in the specified duration, initialize sample_rate=32000
    except Exception as expn:
      return expn,gr.Audio()
        
with gr.Blocks() as io:
  gr.Markdown('<center>[ip] :: text</center>')
  ip = gr.Textbox(placeholder="sep : ||")
  gr.Markdown("<hr>")
  p = gr.Button('genA') # generate Audio[Music]
  gr.Markdown("<hr>")
  gr.Markdown('<center>[op] :: text-audio</center>')
  op0 = gr.Textbox()
  op1 = gr.Audio(type='numpy')
  p.click(func__genA,inputs=[ip], outputs=[op0,op1])

io.launch(share=False,debug=False)