import gradio as gr
import requests
# from IPython.display import Audio
import tempfile
import os

def generate_music(prompt, duration=10):
    API_URL = "https://api-inference.huggingface.co/models/facebook/musicgen-small"
    headers = {"Authorization": f"Bearer {os.getenv('HF_API_KEY')}"}
    
    payload = {
        "inputs": prompt
    }
    
    try:
        response = requests.post(API_URL, headers=headers, json=payload, timeout=300)
        if response.status_code != 200:
            return None, f"Error: API returned status code {response.status_code}"
            
        # Save the audio bytes to a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as tmp_file:
            tmp_file.write(response.content)
            tmp_file_path = tmp_file.name
            
        return tmp_file_path, "Music generated successfully!"
    except Exception as e:
        return None, f"Error: {str(e)}"

# Create the Gradio interface
def gradio_interface(prompt, duration=10):
    audio_path, message = generate_music(prompt, duration)
    if audio_path:
        return audio_path, message
    else:
        raise gr.Error(message)

# Create the Gradio app
with gr.Blocks() as demo:
    gr.Markdown("""
    # MusicGen Audio Generator
    Generate music from text descriptions using Meta's MusicGen model.
    """)
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Enter your music description",
                placeholder="e.g., liquid drum and bass, atmospheric synths, airy sounds",
                lines=3
            )
            duration_slider = gr.Slider(
                minimum=5,
                maximum=30,
                value=10,
                step=5,
                label="Duration (seconds)"
            )
            generate_button = gr.Button("Generate Music")
        
        with gr.Column():
            audio_output = gr.Audio(label="Generated Music")
            message_output = gr.Textbox(label="Status")
    
    generate_button.click(
        fn=gradio_interface,
        inputs=[prompt_input, duration_slider],
        outputs=[audio_output, message_output]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()