# ==============================
# Program
# ==============================

import spaces
import os
import gradio as gr
import torch
from diffusers import DiffusionPipeline, StableDiffusionPipeline, StableDiffusionImg2ImgPipeline

pipe = StableDiffusionPipeline.from_pretrained("mswhite/graffiti_model", torch_dtype=torch.float16) #, safety_checker=None)
pipe = pipe.to("cuda")
pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))

@spaces.GPU
def legoize(text):
    prompt = "brit_graf " + text
    return pipe(prompt).images[0]

with gr.Blocks() as demo:
    
  gr.Markdown("<img src=https://huggingface.co/spaces/mswhite/Graffiti/resolve/main/0071.jpg width=720px>")    
  output        = gr.Image(label="Output", width=720)
  art_to_draw   = gr.Textbox(label="Prompt to Draw: e.g. ")
  flashb_btn    = gr.Button("Spray Paint the Wall")
  flashb_btn.click(fn=legoize, inputs=[art_to_draw], outputs=[output]) 

demo.launch()