import gradio as gr
import torch
import torchaudio
from diffusers import AudioLDMPipeline

# Cargar el modelo de AudioLDM
pipe = AudioLDMPipeline.from_pretrained("cvssp/audioldm-s-full-v2")
pipe.to("cpu")  # Ejecutar en CPU

# Función para generar música
def generar_musica(descripcion):
    try:
        audio = pipe(descripcion, num_inference_steps=25)  # Reducir pasos para mejor rendimiento
        audio_path = "musica_generada.wav"
        torchaudio.save(audio_path, audio["audio"].unsqueeze(0).cpu(), 16000)
        return audio_path
    except Exception as e:
        return f"Error: {e}"

# Interfaz Gradio
with gr.Blocks() as interfaz:
    gr.Markdown("# Generador de Música AI")
    gr.Markdown("Introduce una descripción y genera una pista de música.")

    descripcion = gr.Textbox(label="Descripción de la música")
    salida = gr.Audio(label="Música Generada")

    boton = gr.Button("Generar Música")

    boton.click(fn=generar_musica, inputs=descripcion, outputs=salida)

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