# ==============================
# Main Code
# ==============================

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

# Nvidia GPU config

pipe = StableDiffusionPipeline.from_pretrained("mswhite/miamivice", torch_dtype=torch.float16) #, safety_checker=None)
pipe = pipe.to("cuda")

# Optional Img2Img pipeline

#pipe_img = StableDiffusionImg2ImgPipeline.from_pretrained("mswhite/miamivice", torch_dtype=torch.float16, safety_checker=None)
#pipe_img = pipe_img.to("cuda")

# Sample Prompt - "bright clean clear illustration side profile of a woman wearing pink lipstick p_nagel"

@spaces.GPU
def nagelize(text):
    prompt = "p_nagel detailed illustration of " + text + " bright solid colors"
    return pipe(prompt).images[0]

#def nagelize_me(source_image):
#    prompt = "p_nagel"
#    return pipe_img(prompt, image=source_image, strength=0.6, num_inference_steps=500, guidance_scale=40.0).images[0]

# Web Version of Blocks()

with gr.Blocks() as demo:
    
  gr.Markdown("<img src=https://huggingface.co/spaces/mswhite/miamivice/resolve/main/miami.jpg width=720px>")
  output        = gr.Image(label="Output", width=720)
  art_to_draw   = gr.Textbox(label="Prompt to Draw: e.g. woman with red hair")
  flashb_btn    = gr.Button("Flashback to the 80s")
  flashb_btn.click(fn=nagelize, inputs=[art_to_draw], outputs=[output]) 

demo.launch()