import gradio as gr
import spaces
import torchaudio
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write

model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')
model.set_generation_params(duration=8)  # generate 8 seconds.

@spaces.GPU(duration=120)  # Specify duration if the function is expected to take more than 60s
def generate_music(description, audio_file):
    if audio_file is None:
        wav = model.generate([description])  # generates 1 sample based on the provided description
    else:
        melody, sr = torchaudio.load(audio_file)
        wav = model.generate_with_chroma([description], melody[None], sr)  # generates using the melody from the given audio and the provided description
    
    audio_write('output', wav[0].cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)
    return 'output.wav'

iface = gr.Interface(
    fn=generate_music,
    inputs=[
        gr.Text(label="Description"),
        gr.Audio(type="filepath", label="Audio File (optional)")
    ],
    outputs=gr.Audio(type="filepath"),
    title="MusicGen",
    description="Generate music using the MusicGen model. Provide a description and optionally an audio file for melody.",
)

iface.launch()