bizgen-api / app.py
Saosri Ghosal
added changes
ca917a8
raw
history blame contribute delete
433 Bytes
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()