import gradio as gr
from transformers import pipeline
import torch
import numpy as np

# Load Tom's brain (MusicGen-Small)
# Using cpu for Hugging Face Free Tier compatibility
synthesiser = pipeline("text-to-audio", "facebook/musicgen-small", device="cpu")

# Your requested Clean UI CSS
custom_css = """
footer {display: none !important;}
.show-api {display: none !important;}
button.reference {display: none !important;}
"""

def tom_make_music(prompt, duration):
    if not prompt:
        return None
    
    # 50 tokens = ~1 second of audio
    max_tokens = int(duration * 50)
    
    output = synthesiser(
        prompt, 
        forward_params={
            "do_sample": True, 
            "max_new_tokens": max_tokens
        }
    )
    
    return (output["sampling_rate"], output["audio"].squeeze())

# Build the Interface
# NOTE: css removed from here for Gradio 6.0
with gr.Blocks(title="Tom The Music AI") as demo:
    gr.Markdown("# 🎹 Tom-The-Ai-Music-AI")
    gr.Markdown("I'm Tom. Describe a song and I'll generate it for you. (IMPORTANT: The is Beta Version, Which means it might not work as you want)")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="What should I play?", 
                placeholder="e.g., Chill lo-fi beats for studying",
                lines=3
            )
            duration_slider = gr.Slider(
                minimum=2, maximum=15, value=8, step=1, label="Seconds"
            )
            btn = gr.Button("Generate with Tom", variant="primary")
            
        with gr.Column():
            output_audio = gr.Audio(label="Tom's Masterpiece")

    # api_visibility="undocumented" hides this specific function from the API docs
    btn.click(
        fn=tom_make_music, 
        inputs=[input_text, duration_slider], 
        outputs=output_audio,
        api_visibility="undocumented"
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0", 
        server_port=7860,
        css=custom_css,      # CSS is now passed here
        footer_links=[]      # Replaces show_api=False to hide the entire footer
    )