from einops import rearrange
import torch

from audiocraft.data.audio import audio_write
from audiocraft.models import MusicGen


#!pip install audiocraft
#pip install git+https://github.com/facebookresearch/audiocraft.git
import torch
import torchaudio
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
from audiocraft.utils.notebook import display_audio

#version='facebook/musicgen-melody'
version='facebook/musicgen-small'

print("Loading model", version)

#get the pre-trained model
model = MusicGen.get_pretrained(version, device='cuda')

#def predict_full(model, model_path, decoder, text, melody, duration, topk, topp, temperature, cfg_coef, progress=gr.Progress()):


#The pre trained models are:
#facebook/musicgen-small: 300M model, text to music only - 🤗 Hub
#facebook/musicgen-medium: 1.5B model, text to music only - 🤗 Hub
#facebook/musicgen-melody: 1.5B model, text to music and text+melody to music - 🤗 Hub
#facebook/musicgen-large: 3.3B model, text to music only - 🤗 Hub
    
model.set_generation_params(duration=90)  # generate N seconds.
wav = model.generate_unconditional(4)    # generates 4 unconditional audio samples for each description

#Give the descriptions Text to music 
descriptions = ['chillstep, calm synth, droplets', 
                'new wave 80s chillstep, strong emotion', 
                'darkest depths of space, feeling uneasy', 
               'uplifting, orchestral, drums, classic, instrumental',]
   
# Epic orchestral music, large orchestras and grand sweeping melody, atmosphere
# ambient theme, a swelling violin until the beat rushes into urgency
# a slow orchestral buildup where there is an epic choir near the end
# orchestral, classic, instrumental
   
#Generate the music
wav = model.generate(descriptions,progress=True)  

#display the audio
#display_audio(wav, 32000)

#write audio(s)
for idx, one_wav in enumerate(wav):
    # Will save under {idx}.wav, with loudness normalization at -14 db LUFS.
    audio_write(f'{idx}', one_wav.cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)