#!/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
import xformers
import triton
from PIL import Image
from io import BytesIO
from diffusers.utils import load_image
from diffusers import StableDiffusionXLControlNetPipeline, StableDiffusionXLControlNetInpaintPipeline, ControlNetModel, AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, AutoPipelineForInpainting, EulerDiscreteScheduler, DPMSolverMultistepScheduler

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_CONTROLNETIMG2IMG = os.getenv("ENABLE_USE_CONTROLNET", "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_controlnetimg2img: bool = False,
    url = '',
    controlnet_img = '',
    controlnet_img2img = '',
):
    if torch.cuda.is_available():

        if not use_img2img:
            scheduler = DPMSolverMultistepScheduler.from_pretrained(model, subfolder="scheduler")
            pipe = DiffusionPipeline.from_pretrained(model, scheduler=scheduler, torch_dtype=torch.float16)

            pipe.to(device)
            
            if use_vae:  
                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)

                pipe.to(device)
                
        if use_img2img:
            pipe = AutoPipelineForImage2Image.from_pretrained(model, torch_dtype=torch.float16)

            init_image = load_image(url)

            if use_vae:  
                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
                pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, torch_dtype=torch.float16)

        if use_controlnet:
            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
            pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, controlnet=controlnet, torch_dtype=torch.float16)

            image = load_image(controlnet_img)
            
            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)
                pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, controlnet=controlnet, vae=vae, torch_dtype=torch.float16)

        if use_controlnetimg2img:
            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
            pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, controlnet=controlnet, torch_dtype=torch.float16)

            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)
                pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, controlnet=controlnet, vae=vae, torch_dtype=torch.float16)
                
        if use_lora:
            pipe.load_lora_weights(lora)
            pipe.fuse_lora(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])
            
    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_controlnetimg2img:
        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,
            output_type="pil",
        ).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,
            output_type="pil",
        ).images[0]

theme = gr.themes.Monochrome(
    text_size=gr.themes.Size(lg="18px", md="15px", sm="13px", xl="22px", xs="12px", xxl="24px", xxs="9px"),
    font=[gr.themes.GoogleFont('Source Sans Pro'), 'ui-sans-serif', 'system-ui', 'sans-serif'],
)

with gr.Blocks(theme=theme, css="style.css") as demo:
    gr.HTML()
    gr.Markdown("""# Squaad AI 🧙
### Run Stable Diffusion for free, in seconds
<small>Squaad AI was a free Artificial Intelligence website that brought together the best of the AI world in one place. Unfortunately, it had to cease operations after facing compatibility issues. This space is intended to make its legacy available as open source and allow anyone to run SD in a simple and free way.</small>.

<div style="display: flex; align-items: center;">
            <a target="_blank" href="https://discord.gg/CqXVwGDqFC"><img src="https://img.shields.io/badge/discord-join%20now-blue.svg?logo=discord" alt="discord.gg/CqXVwGDqFC"></a>
        </div>""", elem_id="main_title")
    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)', placeholder='e.g https://example.com/image.png')
        controlnet_img = gr.Text(label='URL (Controlnet)', placeholder='e.g https://example.com/image.png')
        controlnet_img2img = 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_controlnetimg2img = gr.Checkbox(label='Use Controlnet Img2Img', value=False, visible=ENABLE_USE_CONTROLNETIMG2IMG)
            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_controlnetimg2img.change(
        fn=lambda x: gr.update(visible=x),
        inputs=use_controlnetimg2img,
        outputs=controlnet_img2img,
        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_controlnetimg2img,
            url,
            controlnet_img,
            controlnet_img2img,
        ],
        outputs=result,
        api_name="run",
    )

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