import gradio as gr
from audiocraft.models import MusicGen, AudioGen
from audiocraft.data.audio import audio_write
import os
import torch

# Verificar dispositivo
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Usando dispositivo: {device}")

# Carregar modelos corretamente
try:
    # Não usamos .to(device) aqui - o AudioCraft cuida disso internamente
    musicgen = MusicGen.get_pretrained('facebook/musicgen-small')
    audiogen = AudioGen.get_pretrained('facebook/audiogen-medium')
    
    # Configurar parâmetros de geração
    musicgen.set_generation_params(duration=5)
    audiogen.set_generation_params(duration=5)
    
except Exception as e:
    print(f"Erro ao carregar modelos: {e}")
    raise

def generate_audio(prompt, model, model_type):
    try:

        # Caso 1: Se vier da API (JSON com lista), pega o primeiro item
        if isinstance(prompt, list):
            prompt = prompt[0]
        
        print(f"Gerando {model_type} para:", prompt)
        wav = model.generate([prompt])[0]  # Gera o áudio
        
        # Diretório para salvar os arquivos
        temp_dir = "generated_audio"
        os.makedirs(temp_dir, exist_ok=True)
        
        # Caminho do arquivo
        output_path = os.path.join(temp_dir, f"{model_type}_{hash(prompt)}.wav")
        
        # Salvar o áudio
        audio_write(
            output_path[:-4],  # Remove a extensão .wav
            wav.cpu(),        # Garante que está na CPU para salvar
            model.sample_rate,
            strategy="loudness",
            loudness_compressor=True
        )
        
        print(f"Áudio gerado com sucesso em: {output_path}")
        return output_path
        
    except Exception as e:
        print(f"Erro ao gerar áudio: {e}")
        raise gr.Error(f"Falha na geração: {str(e)}")
        
def generate_music(prompt):
    return generate_audio(prompt, musicgen, "música")

def generate_sound_effect(prompt):
    return generate_audio(prompt, audiogen, "efeito sonoro")

# Configuração da interface
with gr.Blocks() as demo:
    gr.Markdown("# 🎵 Gerador de Áudio com MusicGen e AudioGen")
    with gr.Tabs():
        with gr.TabItem("MusicGen 🎵"):
            gr.Markdown("## Gerar Música")
            with gr.Row():
                music_input = gr.Text(label="Descreva a música que deseja")
                music_output = gr.Audio(label="Música Gerada", type="filepath")
            music_button = gr.Button("Gerar Música")
            music_button.click(fn=generate_music, inputs=music_input, outputs=music_output)
            
        with gr.TabItem("AudioGen 🔊"):
            gr.Markdown("## Gerar Efeitos Sonoros")
            with gr.Row():
                sound_input = gr.Text(label="Descreva o efeito sonoro que deseja")
                sound_output = gr.Audio(label="Som Gerado", type="filepath")
            sound_button = gr.Button("Gerar Som")
            sound_button.click(fn=generate_sound_effect, inputs=sound_input, outputs=sound_output)

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