import spaces
import gradio as gr
import torchaudio
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
import logging
import os
import uuid
from torch.cuda.amp import autocast
import torch

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

logging.info("Iniciando y cargando el modelo preentrenado.")
model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')
logging.info("Modelo cargado exitosamente.")

@spaces.GPU(queue=False)
def generar_musica(descripcion, melodia_audio, duracion):
    torch.cuda.empty_cache()
    torch.cuda.synchronize()

    with autocast():
        logging.info("Iniciando la generación de música.")
        model.set_generation_params(duration=duracion)

        if descripcion:
            descripcion = [descripcion]
            if melodia_audio:
                logging.info(f"Cargando la melodía de audio desde: {melodia_audio}")
                melodia, sr = torchaudio.load(melodia_audio)
                logging.info("Generando música con descripción y melodía.")
                wav = model.generate_with_chroma(descripcion, melodia[None], sr)
            else:
                logging.info("Generando música solo con descripción.")
                wav = model.generate(descripcion)
        else:
            logging.info("Generando música de manera incondicional.")
            wav = model.generate_unconditional(1)

        filename = f'{str(uuid.uuid4())}.wav'
        logging.info(f"Guardando la música generada con el nombre: {filename}")
        path = audio_write(filename, wav[0].cpu().to(torch.float32), model.sample_rate, strategy="loudness", loudness_compressor=True)
        logging.info(f"La forma del tensor de audio generado: {wav[0].shape}")
        logging.info("Música generada y guardada con éxito.")
        if not os.path.exists(path):
            raise ValueError(f'No se pudo guardar el audio en {path}')

        return path

descripcion = gr.Textbox(label="Descripción", placeholder="acústico, guitarra, melodía, trap, re menor, 90 bpm")
melodia_audio = gr.Audio(label="Melodía de audio (opcional)", type="filepath")
duracion = gr.Slider(label="Duración (segundos)", minimum=10, maximum=600, step=10, value=30)
output_path = gr.Audio(label="Música generada", type="filepath")

gr.Interface(
    fn=generar_musica,
    inputs=[descripcion, melodia_audio, duracion],
    outputs=output_path,
    title="Generador de Música",
    description="Genera música utilizando el modelo MusicGen.",
).queue(concurrency_count=10).launch(server_port=7860, server_name="0.0.0.0")