import gradio as gr
from audiocraft.models import MusicGen
import torch
import numpy as np

# Carga el modelo
model = MusicGen.get_pretrained('small')
model.set_generation_params(duration=10)

def generate_music(prompt):
    audio = model.generate([prompt], progress=True)
    audio_cpu = audio[0].cpu().numpy()
    return (32000, audio_cpu)

interface = gr.Interface(
    fn=generate_music,
    inputs=gr.Textbox(label="Describe tu música (ej: rock épico)"),
    outputs=gr.Audio(label="Tu música generada"),
    title="Generador de Música con MusicGen",
    description="Escribe un prompt y genera música con IA."
)
interface.launch()