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

# Load the pre-trained MusicGen model
model = MusicGen.get_pretrained('facebook/musicgen-small')
model.set_generation_params(duration=8)  # 8 seconds of audio

def generate_music(prompt):
    wav = model.generate([prompt])  # Generate waveform from prompt
    output_path = "output.wav"
    audio_write(output_path, wav[0], model.sample_rate, strategy="loudness", format="wav")
    return output_path

# Gradio UI
interface = gr.Interface(
    fn=generate_music,
    inputs=gr.Textbox(lines=2, placeholder="Enter a music prompt...", label="Prompt"),
    outputs=gr.Audio(type="filepath", label="Generated Music"),
    title="🎵 AI Music Composer using MusicGen",
    description="Enter a text prompt and generate music using Meta's MusicGen (Small model)."
)

interface.launch()