import gradio as gr
import spaces
import torch
from acestep.pipeline import AceStepPipeline

model_pipeline = None

def load_model():
    global model_pipeline
    if model_pipeline is None:
        model_pipeline = AceStepPipeline.from_pretrained(
            "ace-step/acestep-5Hz-lm-4B",
            torch_dtype=torch.bfloat16
        )
    return model_pipeline

@spaces.GPU(duration=120)
def generate_audio(prompt, duration_sec):
    pipe = load_model()
    result = pipe(
        prompt=prompt,
        duration=duration_sec
    )
    return result

with gr.Blocks(title="ACE-Step 1.5 Generator") as demo:
    gr.Markdown("# ACE-Step 1.5 Music Generator")
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Prompt", 
                placeholder="A pop song with acoustic guitar and lively drums"
            )
            duration_slider = gr.Slider(
                minimum=5, 
                maximum=60, 
                value=15, 
                step=5, 
                label="Duration (seconds)"
            )
            generate_btn = gr.Button("Generate")
        with gr.Column():
            audio_output = gr.Audio(label="Output Audio")

    generate_btn.click(
        fn=generate_audio,
        inputs=[prompt_input, duration_slider],
        outputs=[audio_output]
    )

if __name__ == "__main__":
    demo.launch()