#!/usr/bin/env python

from __future__ import annotations

import requests
import os
import random

import gradio as gr
import numpy as np
import spaces
import torch
import cv2
from PIL import Image
from io import BytesIO
from diffusers.utils import load_image
from diffusers import StableDiffusionXLControlNetPipeline, StableDiffusionXLControlNetInpaintPipeline, ControlNetModel, AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, AutoPipelineForInpainting, UNet2DConditionModel
from controlnet_aux import HEDdetector

DESCRIPTION = "# Run any LoRA or SD Model"
if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>⚠️ This space is running on the CPU. This demo doesn't work on CPU 😞! Run on a GPU by duplicating this space or test our website for free and unlimited by <a href='https://squaadai.com'>clicking here</a>, which provides these and more options.</p>"

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1824"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
ENABLE_USE_LORA = os.getenv("ENABLE_USE_LORA", "1") == "1"
ENABLE_USE_LORA2 = os.getenv("ENABLE_USE_LORA2", "1") == "1"
ENABLE_USE_VAE = os.getenv("ENABLE_USE_VAE", "1") == "1"
ENABLE_USE_IMG2IMG = os.getenv("ENABLE_USE_IMG2IMG", "1") == "1"
ENABLE_USE_CONTROLNET = os.getenv("ENABLE_USE_CONTROLNET", "1") == "1"
ENABLE_USE_CONTROLNETINPAINT = os.getenv("ENABLE_USE_CONTROLNETINPAINT", "1") == "1"

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

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

@spaces.GPU
def generate(
    prompt: str = "",
    negative_prompt: str = "",
    prompt_2: str = "",
    negative_prompt_2: str = "",
    use_negative_prompt: bool = False,
    use_prompt_2: bool = False,
    use_negative_prompt_2: bool = False,
    seed: int = 0,
    width: int = 1024,
    height: int = 1024,
    guidance_scale_base: float = 5.0,
    num_inference_steps_base: int = 25,
    controlnet_conditioning_scale: float = 1,
    control_guidance_start: float = 0,
    control_guidance_end: float = 1,
    strength_img2img: float = 0.7,
    use_vae: bool = False,
    use_lora: bool = False,
    use_lora2: bool = False,
    model = 'stabilityai/stable-diffusion-xl-base-1.0',
    vaecall = 'madebyollin/sdxl-vae-fp16-fix',
    lora = '',
    lora2 = '',
    controlnet_model = 'diffusers/controlnet-canny-sdxl-1.0',
    lora_scale: float = 0.7,
    lora_scale2: float = 0.7,
    use_img2img: bool = False,
    use_controlnet: bool = False,
    use_controlnetinpaint: bool = False,
    url = '',
    controlnet_img = '',
    controlnet_inpaint = '',
):
    if torch.cuda.is_available():

        if not use_img2img:
            pipe = DiffusionPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                        
            if use_vae:  
                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                
        if use_img2img:
            pipe = AutoPipelineForImage2Image.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)

            init_image = load_image(url)

            if use_vae:  
                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)

        if use_controlnet:
            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
            pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)

            image = load_image(controlnet_img)
            
            image = np.array(image)
            image = cv2.Canny(image, 250, 255)
            image = image[:, :, None]
            image = np.concatenate([image, image, image], axis=2)
            image = Image.fromarray(image)

            if use_vae:
                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, vae=vae, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                
        if use_controlnetinpaint:
            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
            pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)

            image_start = load_image(controlnet_img)
            image = load_image(controlnet_img)
            image_mask = load_image(controlnet_img2img)
            
            image = np.array(image)
            image = cv2.Canny(image, 100, 200)
            image = image[:, :, None]
            image = np.concatenate([image, image, image], axis=2)
            image = Image.fromarray(image)
            
            if use_vae:
                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, vae=vae, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                
        if use_lora:
            pipe.load_lora_weights(lora, adapter_name="1")
            pipe.set_adapters("1", adapter_weights=[lora_scale])
            
        if use_lora2:
            pipe.load_lora_weights(lora, adapter_name="1")
            pipe.load_lora_weights(lora2, adapter_name="2")
            pipe.set_adapters(["1", "2"], adapter_weights=[lora_scale, lora_scale2])

    reset_device_map()
    pipe.enable_model_cpu_offload()     
    generator = torch.Generator().manual_seed(seed)
    
    if not use_negative_prompt:
        negative_prompt = None  # type: ignore
    if not use_prompt_2:
        prompt_2 = None  # type: ignore
    if not use_negative_prompt_2:
        negative_prompt_2 = None  # type: ignore

    if use_controlnetinpaint:
        image = pipe(
            prompt=prompt,
            strength=strength_img2img,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            eta=0.0,
            mask_image=image_mask,
            image=image_start,
            control_image=image,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            guidance_scale=guidance_scale_base,
            num_inference_steps=num_inference_steps_base,
            generator=generator,
        ).images[0]
        return image
    if use_controlnet:
        image = pipe(
            prompt=prompt,
            controlnet_conditioning_scale=controlnet_conditioning_scale,
            control_guidance_start=control_guidance_start,
            control_guidance_end=control_guidance_end,
            image=image,
            negative_prompt=negative_prompt,
            prompt_2=prompt_2,
            width=width,
            height=height,
            negative_prompt_2=negative_prompt_2,
            guidance_scale=guidance_scale_base,
            num_inference_steps=num_inference_steps_base,
            generator=generator,
        ).images[0]
        return image
    elif use_img2img:
        images = pipe(
            prompt=prompt,
            image=init_image,
            strength=strength_img2img,
            negative_prompt=negative_prompt,
            prompt_2=prompt_2,
            negative_prompt_2=negative_prompt_2,
            width=width,
            height=height,
            guidance_scale=guidance_scale_base,
            num_inference_steps=num_inference_steps_base,
            generator=generator,
        ).images[0]
        return images
    else:
        return pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            prompt_2=prompt_2,
            negative_prompt_2=negative_prompt_2,
            width=width,
            height=height,
            guidance_scale=guidance_scale_base,
            num_inference_steps=num_inference_steps_base,
            generator=generator,
        ).images[0]

with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
    gr.HTML(
        "<p><center>📙 For any additional support, join our <a href='https://discord.gg/JprjXpjt9K'>Discord</a></center></p>"
    )
    gr.Markdown(DESCRIPTION, elem_id="description")
    with gr.Group():
        model = gr.Text(label='Model', placeholder='e.g. stabilityai/stable-diffusion-xl-base-1.0')
        vaecall = gr.Text(label='VAE', placeholder='e.g. madebyollin/sdxl-vae-fp16-fix')
        lora = gr.Text(label='LoRA 1', placeholder='e.g. nerijs/pixel-art-xl')
        lora2 = gr.Text(label='LoRA 2', placeholder='e.g. nerijs/pixel-art-xl')
        controlnet_model = gr.Text(label='Controlnet', placeholder='e.g diffusers/controlnet-canny-sdxl-1.0')
        lora_scale = gr.Slider(
                info="The closer to 1, the more it will resemble LoRA, but errors may be visible.",
                label="Lora Scale 1",
                minimum=0.01,
                maximum=1,
                step=0.01,
                value=0.7,
            )
        lora_scale2 = gr.Slider(
                info="The closer to 1, the more it will resemble LoRA, but errors may be visible.",
                label="Lora Scale 2",
                minimum=0.01,
                maximum=1,
                step=0.01,
                value=0.7,
            )
        url = gr.Text(label='URL (Img2Img)')
        controlnet_img = gr.Text(label='URL (Controlnet)', placeholder='e.g https://example.com/image.png')
        controlnet_inpaint = gr.Text(label='URL (Controlnet - IMG2IMG)', placeholder='e.g https://example.com/image.png')
        with gr.Row():
            prompt = gr.Text(
                placeholder="Input prompt",
                label="Prompt",
                show_label=False,
                max_lines=1,
                container=False,
            )
            run_button = gr.Button("Run", scale=0)
        result = gr.Image(label="Result", show_label=False)
    with gr.Accordion("Advanced options", open=False):
        with gr.Row():
            use_controlnet = gr.Checkbox(label='Use Controlnet', value=False, visible=ENABLE_USE_CONTROLNET)
            use_controlnetinpaint = gr.Checkbox(label='Use Controlnet Img2Img', value=False, visible=ENABLE_USE_CONTROLNETINPAINT)
            use_img2img = gr.Checkbox(label='Use Img2Img', value=False, visible=ENABLE_USE_IMG2IMG)
            use_vae = gr.Checkbox(label='Use VAE', value=False, visible=ENABLE_USE_VAE)
            use_lora = gr.Checkbox(label='Use Lora 1', value=False, visible=ENABLE_USE_LORA)
            use_lora2 = gr.Checkbox(label='Use Lora 2', value=False, visible=ENABLE_USE_LORA2)
            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
            use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False)
            use_negative_prompt_2 = gr.Checkbox(label="Use negative prompt 2", value=False)
        negative_prompt = gr.Text(
            placeholder="Input Negative Prompt",
            label="Negative prompt",
            max_lines=1,
            visible=False,
        )
        prompt_2 = gr.Text(
            placeholder="Input Prompt 2",
            label="Prompt 2",
            max_lines=1,
            visible=False,
        )
        negative_prompt_2 = gr.Text(
            placeholder="Input Negative Prompt 2",
            label="Negative prompt 2",
            max_lines=1,
            visible=False,
        )

        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():
            width = gr.Slider(
                label="Width",
                minimum=256,
                maximum=MAX_IMAGE_SIZE,
                step=32,
                value=1024,
            )
            height = gr.Slider(
                label="Height",
                minimum=256,
                maximum=MAX_IMAGE_SIZE,
                step=32,
                value=1024,
            )
            
        with gr.Row():
            guidance_scale_base = gr.Slider(
                info="Scale for classifier-free guidance",
                label="Guidance scale",
                minimum=1,
                maximum=20,
                step=0.1,
                value=5.0,
            )
        with gr.Row():
            num_inference_steps_base = gr.Slider(
                info="Number of denoising steps",
                label="Number of inference steps",
                minimum=10,
                maximum=100,
                step=1,
                value=25,
            )
        with gr.Row():
            controlnet_conditioning_scale = gr.Slider(
                info="controlnet_conditioning_scale",
                label="controlnet_conditioning_scale",
                minimum=0.01,
                maximum=2,
                step=0.01,
                value=1,
            )
        with gr.Row():
            control_guidance_start = gr.Slider(
                info="control_guidance_start",
                label="control_guidance_start",
                minimum=0.01,
                maximum=1,
                step=0.01,
                value=0,
            )
        with gr.Row():
            control_guidance_end = gr.Slider(
                info="control_guidance_end",
                label="control_guidance_end",
                minimum=0.01,
                maximum=1,
                step=0.01,
                value=1,
            )
        with gr.Row():
            strength_img2img = gr.Slider(
                info="Strength for Img2Img",
                label="Strength",
                minimum=0,
                maximum=1,
                step=0.01,
                value=0.7,
            )

    use_negative_prompt.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt,
        outputs=negative_prompt,
        queue=False,
        api_name=False,
    )
    use_prompt_2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_prompt_2,
        outputs=prompt_2,
        queue=False,
        api_name=False,
    )
    use_negative_prompt_2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_negative_prompt_2,
        outputs=negative_prompt_2,
        queue=False,
        api_name=False,
    )
    use_vae.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_vae,
        outputs=vaecall,
        queue=False,
        api_name=False,
    )
    use_lora.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_lora,
        outputs=lora,
        queue=False,
        api_name=False,
    )
    use_lora2.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_lora2,
        outputs=lora2,
        queue=False,
        api_name=False,
    )
    use_img2img.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_img2img,
        outputs=url,
        queue=False,
        api_name=False,
    )
    use_controlnet.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_controlnet,
        outputs=controlnet_img,
        queue=False,
        api_name=False,
    )
    use_controlnetinpaint.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_controlnetinpaint,
        outputs=controlnet_inpaint,
        queue=False,
        api_name=False,
    )

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit,
            prompt_2.submit,
            negative_prompt_2.submit,
            run_button.click,
        ],
        fn=randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            prompt_2,
            negative_prompt_2,
            use_negative_prompt,
            use_prompt_2,
            use_negative_prompt_2,
            seed,
            width,
            height,
            guidance_scale_base,
            num_inference_steps_base,
            controlnet_conditioning_scale,
            control_guidance_start,
            control_guidance_end,
            strength_img2img,
            use_vae,
            use_lora,
            use_lora2,
            model,
            vaecall,
            lora,
            lora2,
            controlnet_model,
            lora_scale,
            lora_scale2,
            use_img2img,
            use_controlnet,
            use_controlnetinpaint,
            url,
            controlnet_img,
            controlnet_inpaint,
        ],
        outputs=result,
        api_name="run",
    )

if __name__ == "__main__":
    demo.queue(max_size=20, default_concurrency_limit=2).launch()