import torch
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('medium', device=gcPU)
audiocraft__music_gen__large.set_generation_params(duration=15) # for sample_rate=24KHz to get actual duration of 20sec

def func__MusicGen(ip):
  print('[p] :: processing')
  op = audiocraft__music_gen__large.generate(ip,progress=True)
  for itr in range(0,len(ip)):
    print('= ='*8)
    print('[ip] :: ',end='\t')
    display(ip[itr])
    print('[op] :: ',end='\t')
    display_audio(op[itr],sample_rate=32000) # 1sec=32KHz(here) i.e to get the audio in the specified duration, initialize sample_rate=32000
    print('= ='*8)

ip = ' '
io = []
while(ip!=''):
  ip = input('[ip] :: prompt : ') or ''
  if(ip==''):
    break
  io.append(ip)
func__MusicGen(io)