import gradio as gr
import torchaudio  
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
import spaces
import logging
import os
import uuid
import tempfile

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

@spaces.GPU(duration=120)  
def generate_music(description, melody_audio):
    logging.info("Iniciando a geração de música.")

    # Carrega o modelo pré-treinado
    logging.info("Carregando o modelo pré-treinado.")
    model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')
    model.set_generation_params(duration=8)

    if description:
        description = [description]
        if melody_audio:
            logging.info(f"Carregando a melodia de áudio de: {melody_audio}")
            melody, sr = torchaudio.load(melody_audio)
            logging.info("Gerando música com descrição e melodia.")
            wav = model.generate_with_chroma(description, melody[None], sr)
        else:
            logging.info("Gerando música apenas com descrição.")
            wav = model.generate(description)
    else:
        logging.info("Gerando música de forma incondicional.")
        wav = model.generate_unconditional(1)
        
    # Gera nome de arquivo único 
    filename = f'{str(uuid.uuid4())}.wav'
    
    # Salva na pasta temporária
    output_path = tempfile.mktemp(suffix='.wav')  
    logging.info(f"Salvando a música gerada em: {output_path}")
    
    audio_write(output_path, wav[0].cpu(), model.sample_rate, 
                strategy="loudness", loudness_compressor=True)
    
    # Verifica se o arquivo existe
    if not os.access(output_path, os.R_OK):
        raise ValueError(f'Failed to save audio to {output_path}')

    logging.info(f"A forma do tensor de áudio gerado: {wav[0].shape}")
    logging.info("Música gerada e salva com sucesso.")

    return output_path
    
# Define a interface Gradio
description = gr.Textbox(label="Description", placeholder="acoustic, guitar, melody, trap, d minor, 90 bpm")
melody_audio = gr.Audio(label="Melody Audio (optional)", type="filepath")
output_path = gr.Audio(label="Generated Music", type="filepath")

gr.Interface(
    fn=generate_music,
    inputs=[description, melody_audio],
    outputs=output_path,
    title="MusicGen Demo",
    description="Generate music using the MusicGen model.",
    examples=[
        ["trap, synthesizer, songstarters, dark, G# minor, 140 bpm", "./assets/kalhonaho.mp3"],
        ["upbeat, electronic, synth, dance, 120 bpm", None]
    ]
).launch()