#!/usr/bin/env python

from __future__ import annotations

import os
import random

import gradio as gr
import numpy as np
import PIL.Image
import torch
from diffusers.models import AutoencoderKL
from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler

DESCRIPTION = '# Animagine XL'
if not torch.cuda.is_available():
    DESCRIPTION += '\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>'

MAX_SEED = np.iinfo(np.int32).max
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv(
    'CACHE_EXAMPLES') == '1'
USE_TORCH_COMPILE = os.getenv('USE_TORCH_COMPILE') == '1'
ENABLE_CPU_OFFLOAD = os.getenv('ENABLE_CPU_OFFLOAD') == '1'

MODEL = "Linaqruf/animagine-xl"

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
if torch.cuda.is_available():
    pipe = StableDiffusionXLPipeline.from_pretrained(
        MODEL,
        torch_dtype=torch.float16,
        use_safetensors=True,
        variant='fp16')

    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
    
    if ENABLE_CPU_OFFLOAD:
        pipe.enable_model_cpu_offload()
    else:
        pipe.to(device)

    if USE_TORCH_COMPILE:
        pipe.unet = torch.compile(pipe.unet,
                                  mode='reduce-overhead',
                                  fullgraph=True)
else:
    pipe = None


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


def generate(prompt: str,
             negative_prompt: str = '',
             prompt_2: str = '',
             negative_prompt_2: str = '',
             use_prompt_2: bool = False,
             seed: int = 0,
             width: int = 1024,
             height: int = 1024,
             target_width: int = 1024,
             target_height: int = 1024,
             original_width: int = 4096,
             original_height: int = 4096,
             guidance_scale_base: float = 12.0,
             num_inference_steps_base: int = 50) -> PIL.Image.Image:
    
    generator = torch.Generator().manual_seed(seed)

    if negative_prompt == '':
        negative_prompt = None  # type: ignore
    if not use_prompt_2:
        prompt_2 = None  # type: ignore
        negative_prompt_2 = None  # type: ignore
    if negative_prompt_2 == '':
        negative_prompt_2 = None  

    return pipe(prompt=prompt,
                    negative_prompt=negative_prompt,
                    prompt_2=prompt_2,
                    negative_prompt_2=negative_prompt_2,
                    width=width,
                    height=height,
                    target_size=(target_width, target_height),
                    original_size=(original_width, original_height),
                    guidance_scale=guidance_scale_base,
                    num_inference_steps=num_inference_steps_base,
                    generator=generator,
                    output_type='pil').images[0]


examples = [
    'face focus, cute, masterpiece, best quality, 1girl, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck',
    'face focus, bishounen, masterpiece, best quality, 1boy, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck',
]

# choices = [
#     "Vertical (9:16)",
#     "Portrait (4:5)",
#     "Square (1:1)",
#     "Photo (4:3)",
#     "Landscape (3:2)",
#     "Widescreen (16:9)",
#     "Cinematic (21:9)",
# ]

# choice_to_size = {
#     "Vertical (9:16)": (768, 1344),
#     "Portrait (4:5)": (912, 1144),
#     "Square (1:1)": (1024, 1024),
#     "Photo (4:3)": (1184, 888),
#     "Landscape (3:2)": (1256, 832),
#     "Widescreen (16:9)": (1368, 768),
#     "Cinematic (21:9)": (1568, 672),
# }

with gr.Blocks(css='style.css') as demo:
    gr.Markdown(DESCRIPTION)
    gr.DuplicateButton(value='Duplicate Space for private use',
                       elem_id='duplicate-button',
                       visible=os.getenv('SHOW_DUPLICATE_BUTTON') == '1')
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Text(
                label='Prompt',
                max_lines=1,
                placeholder='Enter your prompt',
            )
            negative_prompt = gr.Text(
                label='Negative Prompt',
                max_lines=1,
                placeholder='Enter a negative prompt',
                value='lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry',
            )
            use_prompt_2 = gr.Checkbox(
                label='Use prompt 2', 
                value=False
            )   
            prompt_2 = gr.Text(
                label='Prompt 2',
                max_lines=1,
                placeholder='Enter your prompt',
                visible=False,
            )
            negative_prompt_2 = gr.Text(
                label='Negative prompt 2',
                max_lines=1,
                placeholder='Enter a negative prompt',
                visible=False,
            )

            # with gr.Row():
            #     aspect_ratio = gr.Dropdown(choices=choices, label="Aspect Ratio Preset", value=choices[2])
            with gr.Row():
                width = gr.Slider(
                    label='Width',
                    minimum=256,
                    maximum=4096,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label='Height',
                    minimum=256,
                    maximum=4096,
                    step=32,
                    value=1024,
                )
            with gr.Accordion(label='Advanced Config', open=False):
                with gr.Accordion(label='Conditioning Resolution', open=False):
                    with gr.Row():
                        original_width = gr.Slider(
                            label='Original Width',
                            minimum=1024,
                            maximum=4096,
                            step=32,
                            value=4096,
                        )
                        original_height = gr.Slider(
                            label='Original Height',
                            minimum=1024,
                            maximum=4096,
                            step=32,
                            value=4096,
                        )
                    with gr.Row():
                        target_width = gr.Slider(
                            label='Target Width',
                            minimum=1024,
                            maximum=4096,
                            step=32,
                            value=1024,
                            )
                        target_height = gr.Slider(
                            label='Target Height',
                            minimum=1024,
                            maximum=4096,
                            step=32,
                            value=1024,
                        )
                seed = gr.Slider(label='Seed',
                                minimum=0,
                                maximum=MAX_SEED,
                                step=1,
                                value=0)
                
                randomize_seed = gr.Checkbox(label='Randomize seed', value=True)
                with gr.Row():
                    guidance_scale_base = gr.Slider(
                        label='Guidance scale',
                        minimum=1,
                        maximum=20,
                        step=0.1,
                        value=12.0)
                    num_inference_steps_base = gr.Slider(
                        label='Number of inference steps',
                        minimum=10,
                        maximum=100,
                        step=1,
                        value=50)

        with gr.Column(scale=2):
            with gr.Blocks():
                run_button = gr.Button('Generate')
            result = gr.Image(label='Result', show_label=False)

    gr.Examples(examples=examples,
                inputs=prompt,
                outputs=result,
                fn=generate,
                cache_examples=CACHE_EXAMPLES)
    
    use_prompt_2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_prompt_2,
        outputs=prompt_2,
        queue=False,
        api_name=False,
    )
    use_prompt_2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_prompt_2,
        outputs=negative_prompt_2,
        queue=False,
        api_name=False,
    )

    inputs = [
        prompt,
        negative_prompt,
        prompt_2,
        negative_prompt_2,
        use_prompt_2,
        seed,
        width,
        height,
        target_width,
        target_height,
        original_width,
        original_height,
        guidance_scale_base,
        num_inference_steps_base,
    ]
    prompt.submit(
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=result,
        api_name='run',
    )
    negative_prompt.submit(
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=result,
        api_name=False,
    )
    prompt_2.submit(
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=result,
        api_name=False,
    )
    negative_prompt_2.submit(
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=result,
        api_name=False,
    )
    run_button.click(
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=inputs,
        outputs=result,
        api_name=False,
    )

demo.queue(max_size=20).launch()