mswhite/Bricksburg

1# ==============================2# Program3# ==============================4 5import spaces6import os7import gradio as gr8import torch9from diffusers import DiffusionPipeline, StableDiffusionPipeline, StableDiffusionImg2ImgPipeline10 11pipe = StableDiffusionPipeline.from_pretrained("mswhite/bricksburg_model", torch_dtype=torch.float16) #, safety_checker=None)12pipe = pipe.to("cuda")13pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))14 15@spaces.GPU16def legoize(text):17    prompt = "mini_mw_lgo" + text18    return pipe(prompt).images[0]19 20with gr.Blocks() as demo:21    22  gr.Markdown("<img src=https://huggingface.co/spaces/mswhite/Bricksburg/resolve/main/Emmet.jpg width=720px>")    23  output        = gr.Image(label="Output", width=720)24  art_to_draw   = gr.Textbox(label="Prompt to Draw: e.g. ")25  flashb_btn    = gr.Button("Everything is Awesome in Bricksburg")26  flashb_btn.click(fn=legoize, inputs=[art_to_draw], outputs=[output]) 27 28demo.launch()29