import gradio as gr
from transformers import AutoTokenizer, AutoModel

# Load model and tokenizer for CPU
ckpt = "Mar2Ding/songcomposer_pretrain"
tokenizer = AutoTokenizer.from_pretrained(ckpt, trust_remote_code=True)
model = AutoModel.from_pretrained(ckpt, trust_remote_code=True)

def generate_song(prompt):
    try:
        output = model.inference_pretrain(prompt, tokenizer)
        return output
    except Exception as e:
        return f"Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# 🎶 EchoSongGen")
    gr.Markdown("Generate melody and lyrics from a text prompt using SongComposer (CPU version).")

    prompt = gr.Textbox(label="Song Prompt", placeholder="e.g. A sad love song in the rain")
    generate_btn = gr.Button("Generate")
    output = gr.Textbox(label="Generated Song")

    generate_btn.click(fn=generate_song, inputs=prompt, outputs=output)

demo.launch()