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

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

# Limitar el uso de memoria de la GPU
torch.cuda.set_per_process_memory_fraction(0.8)

# Cargar el modelo preentrenado
logging.info("Cargando el modelo preentrenado.")
model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')

@spaces.GPU(duration=120)
def generar_musica(descripcion, melodia_audio, duracion):
    # Limpiar la caché de la GPU antes de generar
    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)
        print("Música guardada en", path, ".")
        # Verifica la forma del tensor de audio generado y si fue guardado correctamente
        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
    
# Definir la interfaz Gradio
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)  # Máximo 10 minutos (600 segundos)
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.",
    examples=[
        ["trap, sintetizador, songstarters, oscuro, G# menor, 140 bpm", "./assets/kalhonaho.mp3", 30],
        ["alegre, electrónica, sintetizador, dance, 120 bpm", None, 60]
    ]
).launch()