from diffusers import StableDiffusionPipeline
import gradio as gr
import torch
import os
from PIL import Image

from share_btn import community_icon_html, loading_icon_html, share_js

if torch.cuda.is_available():
  torchfloat = torch.float16
else:
  torchfloat = torch.float32


models = [
  "DGSpitzer/Cyberpunk-Anime-Diffusion"
]

prompt_prefixes = {
  models[0]: "dgs illustration style "
}

current_model = models[0]

#auth_token = os.environ.get("test") or True
#pipe = StableDiffusionPipeline.from_pretrained(current_model, use_auth_token=auth_token, torch_dtype=torchfloat, revision="fp16")
pipe = StableDiffusionPipeline.from_pretrained(current_model, torch_dtype=torchfloat)

if torch.cuda.is_available():
  pipe = pipe.to("cuda")
else:
  pipe = pipe.to("cpu")

device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"

def on_model_change(model):

    global current_model
    global pipe
    if model != current_model:
        current_model = model
        pipe = StableDiffusionPipeline.from_pretrained(current_model, torch_dtype=torchfloat)
        if torch.cuda.is_available():
            pipe = pipe.to("cuda")

def inference(prompt, negPrompt, guidance, steps, width_input, height_input):
    print("Generated image with prompt: " + prompt)
    if negPrompt != "":
        print("Negative prompt: " + negPrompt)
    promptPrev = prompt
    prompt = prompt_prefixes[current_model] + prompt
    results = pipe(prompt, negative_prompt = negPrompt, num_inference_steps=int(steps), guidance_scale=guidance, width=width_input, height=height_input)
    image = results.images[0] if not results.nsfw_content_detected[0] else Image.open("nsfw_placeholder.jpg")
    return image, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)

def inference_example(prompt, negPrompt, guidance, steps, width_input, height_input):
    
    prompt = prompt_prefixes[current_model] + prompt
    results = pipe(prompt, negative_prompt = negPrompt, num_inference_steps=int(steps), guidance_scale=guidance, width=width_input, height=height_input)
    image = results.images[0] if not results.nsfw_content_detected[0] else Image.open("nsfw_placeholder.jpg")
    return image

css = """
#col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
a {text-decoration-line: underline; font-weight: 600;}
.animate-spin {
    animation: spin 1s linear infinite;
}
@keyframes spin {
    from {
        transform: rotate(0deg);
    }
    to {
        transform: rotate(360deg);
    }
}
#share-btn-container {
    display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
}
#share-btn {
    all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.25rem !important; padding-bottom: 0.25rem !important;right:0;
}
#share-btn * {
    all: unset;
}
#share-btn-container div:nth-child(-n+2){
    width: auto !important;
    min-height: 0px !important;
}
#share-btn-container .wrap {
    display: none !important;
}
"""

with gr.Blocks(css=css) as demo:
    gr.HTML(
        """
            <div style="text-align: center; max-width: 700px; margin: 0 auto;">
              <div
                style="
                  display: inline-flex;
                  align-items: center;
                  gap: 0.8rem;
                  font-size: 1.75rem;
                "
              >
                <h1 style="font-weight: 900; margin-bottom: 7px;">
                  DGS Diffusion Space
                </h1>
              </div>
              <p style="margin-bottom: 10px; font-size: 94%">
               Demo for Cyberpunk Anime Diffusion. Based of the projects by anzorq <a href="https://twitter.com/hahahahohohe"></a> and fffiloni <a href="https://twitter.com/fffiloni"></a>
              </p>
            </div>
            <div
               style="
                  gap: 1.2rem;
                  font-size: 1rem;
                "
             >
              <p>Another online version without queue: <a href="https://colab.research.google.com/github/HelixNGC7293/cyberpunk-anime-diffusion/blob/main/cyberpunk_anime_diffusion.ipynb"><img data-canonical-src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab" src="https://camo.githubusercontent.com/84f0493939e0c4de4e6dbe113251b4bfb5353e57134ffd9fcab6b8714514d4d1/68747470733a2f2f636f6c61622e72657365617263682e676f6f676c652e636f6d2f6173736574732f636f6c61622d62616467652e737667"></a></p>
            </div>
        """
    )
    gr.Markdown('''
      👇 Buy me a coffee if you like ♥ this project!
      [![Buy me a coffee](https://badgen.net/badge/icon/Buy%20Me%20A%20Coffee?icon=buymeacoffee&label)](https://www.buymeacoffee.com/dgspitzer)
    ''')
    with gr.Row():
        
        with gr.Column():
            model = gr.Dropdown(label="Model", choices=models, value=models[0])
            prompt = gr.Textbox(label="Prompt", placeholder="{} is added automatically".format(prompt_prefixes[current_model]), elem_id="input-prompt")
            negPrompt = gr.Textbox(label="Negative Prompt", placeholder="Enter what you don't want to generate", elem_id="input-negPrompt")
            guidance = gr.Slider(label="Guidance scale", value=7, maximum=8)
            steps = gr.Slider(label="Steps", value=20, maximum=30, minimum=2)
            width_input = gr.Slider(label="Width", value=576, maximum=768, minimum=384, step=64)
            height_input = gr.Slider(label="Height", value=576, maximum=768, minimum=384, step=64)
            run = gr.Button(value="Run")
            gr.Markdown(f"Running on: {device}")
        with gr.Column():
            image_out = gr.Image(height=512, type="filepath", elem_id="output-img")

    with gr.Column(elem_id="col-container"):
        with gr.Group(elem_id="share-btn-container"):
          community_icon = gr.HTML(community_icon_html, visible=False)
          loading_icon = gr.HTML(loading_icon_html, visible=False)
          share_button = gr.Button("Share to community", elem_id="share-btn", visible=False)
          
    model.change(on_model_change, inputs=model, outputs=[])
    run.click(inference, inputs=[prompt, negPrompt, guidance, steps, width_input, height_input], outputs=[image_out, share_button, community_icon, loading_icon])
    
    share_button.click(None, [], [], _js=share_js)
    
    gr.Examples([
        ["perfect face portrait of beautiful smile girl,  clean face, wears hoody, half body, soldier working in a cyberpunk city, cleavage, intricate, 8k, highly detailed, digital painting, intense, sharp focus", "out of focus, scary, creepy, evil, disfigured, missing limbs, ugly, gross, missing fingers", 7, 20, 512, 704],
        ["portrait of a beautiful fancy gorgeous anime girl, intricate details", "out of focus, scary, creepy, evil, disfigured, missing limbs, ugly, gross, missing fingers", 7, 20, 448, 640],
        ["a beautiful perfect face girl, Anime fine details portrait of school girl in front of modern tokyo city landscape on the background deep bokeh, anime masterpiece by studio ghibli, 8k, sharp high quality anime, artstation", "out of focus, scary, creepy, evil, disfigured, missing limbs, ugly, gross, missing fingers", 7, 20, 704, 704],
        ["city landscape with fancy car, racing on the road, gopro, intricate details, 4k, cyberpunk", "out of focus, scary, creepy, evil, disfigured, missing limbs, ugly, gross, missing fingers", 7, 20, 704, 704],
        ["portrait of liu yifei girl, soldier working in a cyberpunk city, cleavage, intricate, 8k, highly detailed, digital painting, intense, sharp focus", "out of focus, scary, creepy, evil, disfigured, missing limbs, ugly, gross, missing fingers", 7, 20, 704, 704],
        ["portrait of a muscular beard male in dgs illustration style, half-body, holding robot arms, strong chest", "out of focus, scary, creepy, evil, disfigured, missing limbs, ugly, gross, missing fingers", 7, 20, 512, 640],
    ], [prompt, negPrompt, guidance, steps, width_input, height_input], image_out, inference_example, cache_examples=torch.cuda.is_available())
    gr.Markdown('''
      Models and Space by [@DGSpitzer](https://www.youtube.com/channel/UCzzsYBF4qwtMwJaPJZ5SuPg)❤️ [@大谷的游戏创作小屋](https://space.bilibili.com/176003)
      [![Twitter Follow](https://img.shields.io/twitter/follow/DGSpitzer?label=%40DGSpitzer&style=social)](https://twitter.com/DGSpitzer)
      ![visitors](https://visitor-badge.glitch.me/badge?page_id=dgspitzer_DGS_Diffusion_Space)
      
      ![Model Views](https://visitor-badge.glitch.me/badge?page_id=Cyberpunk_Anime_Diffusion)
      
    ''')

demo.queue()
demo.launch()