import gradio as gr
from transformers import pipeline
import scipy.io.wavfile as wavfile
import numpy as np

# Load the small version of MusicGen (Fits in Free Tier memory)
model_id = "facebook/musicgen-small"
synthesizer = pipeline("text-to-audio", model=model_id)

def generate_music(prompt):
    # The model generates the audio based on your prompt
    output = synthesizer(prompt, forward_params={"do_sample": True})
    
    # Extract the audio data and sampling rate
    sampling_rate = output["sampling_rate"]
    audio_data = output["audio"]
    
    # Hugging Face/Gradio expects a tuple of (rate, data)
    # We ensure the data is in the correct format (float32 or int16)
    return (sampling_rate, audio_data.T)

# Create the beautiful web interface
with gr.Blocks() as demo:
    gr.Markdown("# 🎵 Open Hearts Music Generator")
    gr.Markdown("Enter your cinematic prompt below to generate a unique track.")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Your Music Prompt", 
                placeholder="A cinematic pop-folk crossover with piano and strings...",
                lines=3
            )
            generate_btn = gr.Button("Generate Music", variant="primary")
        
        with gr.Column():
            audio_output = gr.Audio(label="Resulting Audio")

    # Connect the button to the function
    generate_btn.click(fn=generate_music, inputs=input_text, outputs=audio_output)

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