import tempfile
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
import gradio as gr
import torch
import uuid
import os
from scipy.io.wavfile import write

model = MusicGen.get_pretrained("facebook/musicgen-small")
model.set_generation_params(duration=5)

def generate_music(description):
    # This line was not indented properly, fixed by adding indentation
    wav = model.generate([description])  
    audio_array = wav.cpu().numpy().squeeze()  
    sample_rate = model.sample_rate

    # Generate a unique file path
    file_id = uuid.uuid1()

    file_path = os.path.join(
        tempfile.gettempdir(),
        f'{file_id}.wav'
    )

    print(f"Temporary directory: {tempfile.gettempdir()}")
    print(f"File path: {file_path}")

    # Write the audio file to the temporary path
    write(file_path, rate=sample_rate, data=audio_array)

    return file_path

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_music, 
    inputs="text", 
    outputs=gr.components.Audio(type="filepath", label="Audio"), 
    title="Text to Audio Generation", 
    description="Generate audio based on text descriptions.", 
    live=False
)

# Launch the Gradio interface
iface.launch(debug=True)