import gradio as gr
import torch

from diffusers import DiffusionPipeline

MODEL_ID = "PYY2001/BizGen"

pipe = DiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
)
pipe.to("cpu")

def generate(prompt):
    out = pipe(prompt).images[0]
    return out

demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Image(),
    title="BizGen Image Generator"
)

demo.launch()