from gradio import gr
from datetime import datetime
from pathlib import Path

from audiocraft.models import musicgen
from audiocraft.data.audio import audio_write
#from audiocraft.utils.notebook import display_audio

model = musicgen.MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.1', device='cpu')

# path to save our samples.
out_dir = Path("./samples")
out_dir.mkdir(exist_ok=True, parents=True)

text = "hip hop, soul, piano, chords, jazz, neo jazz, G# minor, 140 bpm"
duration=15
use_sampling=True
temperature=1.0
top_k=250
cfg_coef=3.0

def music(text,duration=duration,use_sampling=use_sampling,temperature=temperature,top_k=top_k,cfg_coef=cfg_coef):
    text = "hip hop, soul, piano, chords, jazz, neo jazz, G# minor, 140 bpm"


    model.set_generation_params(
        duration=duration,
        use_sampling=use_sampling,
        temperature=temperature,
        top_k=top_k,
        cfg_coef=cfg_coef,
    )
    N = 4
    out = model.generate(
        [text] * N,
        progress=True,
    )
    
    # Write to files
    dt_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    for i in range(N):
        audio_write(
            out_dir / f"{dt_str}_{i:02d}",
            out[i].cpu(),
            model.sample_rate,
            strategy="loudness",
        )
        print(out_dir / f"{dt_str}_{i:02d}")
    
    # Or, if in a notebook, display audio widgets
    # display_audio(out, model.sample_rate)
with gr.Blocks() as app:
    inp=gr.Textbox()
    b=gr.Button()
    b.click(music,inp,None)
app.launch()