import gradio as gr
import torch
from transformers import AutoProcessor, MusicgenForConditionalGeneration

# Load Meta's AudioGen model
model_id = "facebook/musicgen-small"
device = "cuda" if torch.cuda.is_available() else "cpu"

processor = AutoProcessor.from_pretrained(model_id)
model = MusicgenForConditionalGeneration.from_pretrained(model_id).to(device)

def generate_audio(text_prompt, duration=5):
    """Generate music based on a text prompt."""
    inputs = processor(
        text=[text_prompt], 
        return_tensors="pt"
    ).to(device)

    with torch.no_grad():
        audio_waveform = model.generate(**inputs, max_new_tokens=duration * 50)

    return audio_waveform[0].cpu().numpy(), 32000  # Sample rate

# Gradio UI
iface = gr.Interface(
    fn=generate_audio,
    inputs=[
        gr.Textbox(label="Music Prompt", placeholder="e.g., A calm piano melody"),
        gr.Slider(1, 30, value=5, step=1, label="Duration (seconds)")
    ],
    outputs=gr.Audio(label="Generated Music"),
    title="Text-to-Music Generator",
    description="Enter a text prompt to generate music using Meta's MusicGen model.",
)

if __name__ == "__main__":
    iface.launch()