# ==============================
# Basic Interface
# ==============================

#import gradio as gr
#demo = gr.load("mswhite/miamivice", src="models")  
#demo.launch()

# ==============================
# Detailed Interface
# ==============================

import os
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")

# CPU config
#pipe = StableDiffusionPipeline.from_pretrained("mswhite/miamivice") 

# 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"

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:

    #girl_image    = gr.Image(label = None, value = os.path.join(os.path.dirname(__file__), "headshot.jpg"), visible=False)
    
#    with gr.Row():
#        logo_image    = gr.Image(label = None, value = os.path.join(os.path.dirname(__file__), "miami.jpg"), interactive=True).style(height=720)
#        #gr.Markdown("<img src=https://huggingface.co/spaces/mswhite/miamivice/resolve/main/miami.jpg width=720px>")
#        output        = gr.Image(label="Output").style(height=720)
#    with gr.Row():
#        art_to_draw   = gr.Textbox(label="Prompt to Draw: e.g. woman with red hair")
#        flashb_btn    = gr.Button("Flashback to the 80s")
#        #back80_btn    = gr.Button("Send me back to the 80s")

#    flashb_btn.click(fn=nagelize,    inputs=[art_to_draw], outputs=[output]) 
#    #back80_btn.click(fn=nagelize,    inputs=[art_to_draw], outputs=[output])    

# Mobile 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").style(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()