import gradio as gr
import spaces
import torchaudio
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write

# Carrega o modelo pré-treinado
model = MusicGen.get_pretrained('nateraw/musicgen-songstarter-v0.2')
model.set_generation_params(duration=8)  # Configura a duração da geração para 8 segundos

@spaces.GPU(duration=120)  # Habilita o uso de GPU
def generate_audio(prompt, audio_input=None):
    descriptions = [prompt]  # Usa o prompt como descrição
    if audio_input is None:
        # Geração incondicional com descrições
        wav = model.generate(descriptions)
    else:
        # Carrega o áudio de entrada e gera com base nele e nas descrições
        melody, sr = torchaudio.load(audio_input)
        wav = model.generate_with_chroma(descriptions, melody[None], sr)

    # Salva o arquivo de áudio gerado
    output_path = 'generated_audio.wav'
    audio_write(output_path, wav[0].cpu(), model.sample_rate, strategy="loudness", loudness_compressor=True)

    return output_path

# Cria a interface de usuário com Gradio
iface = gr.Interface(
    fn=generate_audio,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter a description..."),
        gr.File(label="Audio Input (optional)", type="filepath", optional=True)],
    outputs=gr.File(label="Generated Audio")
)

iface.launch()