#!/usr/bin/env python
# Patch3.09
import os
import random
import uuid
import gdown

import gradio as gr
import numpy as np
from PIL import Image
import spaces
import torch
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler

DESCRIPTION = """ """

def save_image(img):
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)
    return unique_name

def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    return seed

MAX_SEED = np.iinfo(np.int32).max

if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU, This Space may not work on CPU.</p>"

USE_TORCH_COMPILE = 0
ENABLE_CPU_OFFLOAD = 0

if torch.cuda.is_available():
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        torch_dtype=torch.float16,
        use_safetensors=True,
    )
    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)

    # Download the LoRA weights from Google Drive
    drive_folder_url = "https://drive.google.com/drive/folders/1ExL5VNChyYWXho1QbgNbOkTK3xc8mhHW"
    weight_file_id = "18n6gF7Jda92MpqK7cYs0Gv2IqLAVnltZ"
    weight_file_name = "pytorch_lora_weights.safetensors"
    # Use gdown to download the file
    gdown.download(f"https://drive.google.com/uc?id={weight_file_id}", weight_file_name, quiet=False)
    
    pipe.load_lora_weights(weight_file_name, adapter_name="icon")
    pipe.set_adapters("icon")

    pipe.to("cuda")

@spaces.GPU(enable_queue=True)
def generate(
    prompt: str,
    negative_prompt: str = "",
    use_negative_prompt: bool = False,
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 3,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    seed = int(randomize_seed_fn(seed, randomize_seed))

    if not use_negative_prompt:
        negative_prompt = ""  # type: ignore

    images = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=25,
        num_images_per_prompt=1,
        cross_attention_kwargs={"scale": 0.65},
        output_type="pil",
    ).images
    image_paths = [save_image(img) for img in images]
    print(image_paths)
    return image_paths, seed

examples = [
    "1boy, male focus, sky, star (sky), night, pointing up, night sky, hood down, starry sky, hood, blue theme, outdoors, long sleeves, shooting star, hoodie, short hair, jacket, scenery, cloud, from behind, blue eyes, best quality, amazing quality, best aesthetic, absurdres",
    "1boy, male focus, bishounen, holding sword, holding weapon, katana, sword, japanese clothes, haori, east asian architecture, solo, looking at viewer, expressionless, blue hair, purple eyes, long hair, best quality, amazing quality, best aesthetic, absurdres",
    "1boy, male focus, holding drink, holding, drink, toned male, toned, pectorals, jacket, open jacket, open clothes, tank top, chain necklace, necklace, stud earrings, earrings, jewelry, cafe, plant, indoors, lens flare, solo, looking at viewer, open mouth, fang, white hair, yellow eyes, short hair, best quality, amazing quality, best aesthetic, absurdres, year 2023",
    "1boy, male focus, dark-skinned male, dark skin, squatting, heart hands, bara, wooden floor, floor, indoors, gym uniform, sneakers, shoes, solo, looking at viewer, frown, sweatdrop, very short hair, best quality, amazing quality, best aesthetic, absurdres, year 2023",
    "1boy, male focus, short hair, blue hair, blue eyes, graphic t-shirt, punk t-shirt, digital illustration, cyan and black, looking at viewer,  busy city street, belt, black pants, atmospheric lighting, midriff peek, night, blurry, best quality, amazing quality, best aesthetic, absurdres",
    "Ultra realistic close up portrait ((beautiful pale cyberpunk female with heavy black eyeliner)), blue eyes, shaved side haircut, hyper detail, cinematic lighting, magic neon, dark red city, Canon EOS R3, nikon, f/1.4, ISO 200, 1/160s, 8K, RAW, unedited, symmetrical balance, in-frame, 8K"
]

css = '''
.gradio-container{max-width: 600px !important}
h1{text-align:center}
footer {
    visibility: hidden
}
'''
with gr.Blocks(css=css, theme="ParityError/Anime") as demo:
    gr.Markdown(DESCRIPTION)
    gr.DuplicateButton(
        value="Duplicate Space for private use",
        elem_id="duplicate-button",
        visible=False,
    )

    with gr.Group():
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0)
        result = gr.Gallery(label="Result", columns=1, preview=True, show_label=False)
    with gr.Accordion("Advanced options", open=False):
        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
        negative_prompt = gr.Text(
            label="Negative prompt",
            lines=4,
            max_lines=6,
            value="""(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation""",
            placeholder="Enter a negative prompt",
            visible=True,
        )
        seed = gr.Slider(
            label="Seed",
            minimum=0,
            maximum=MAX_SEED,
            step=1,
            value=0,
            visible=True
        )
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        with gr.Row(visible=True):
            width = gr.Slider(
                label="Width",
                minimum=512,
                maximum=2048,
                step=8,
                value=1024,
            )
            height = gr.Slider(
                label="Height",
                minimum=512,
                maximum=2048,
                step=8,
                value=1024,
            )
        with gr.Row():
            guidance_scale = gr.Slider(
                label="Guidance Scale",
                minimum=0.1,
                maximum=20.0,
                step=0.1,
                value=6,
            )

    gr.Examples(
        examples=examples,
        inputs=prompt,
        outputs=[result, seed],
        fn=generate,
        cache_examples=False,
    )

    use_negative_prompt.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt,
        outputs=negative_prompt,
        api_name=False,
    )

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit,
            run_button.click,
        ],
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            use_negative_prompt,
            seed,
            width,
            height,
            guidance_scale,
            randomize_seed,
        ],
        outputs=[result, seed],
        api_name="run",
    )
    
if __name__ == "__main__":
    demo.queue(max_size=20).launch(show_api=False, debug=False)