import glob
import os
import random

import gradio as gr
import numpy as np
import PIL.Image
import spaces
import torch
from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    StableDiffusionXLPipeline,
)
from torchvision.transforms import ToTensor

# pyright: reportPrivateImportUsage=false


DESCRIPTION = f"""
    # 🎨 Inversion-InstantStyle 🎨  
    This is an interactive demo of [Inversion-InstantStyle](https://gojasper.github.io/style-rank-project/#inversion_instantstyle), which combines DDIM inversion and renoising with the [Instant-Style](https://instantstyle.github.io/) styling method.
    It was proposed in the context of our [Style-Rank](https://gojasper.github.io/style-rank-project) benchmark which evaluates training-free styling methods, by *Eyal Benaroche, Clément Chadebec, Onur Tasar, and Benjamin Aubin* from [Jasper Research](https://www.jasper.ai/) and [Ecole Polytechnique](https://www.polytechnique.edu/).
    """

OPEN_SOURCE_PROMO = f"""
    If you enjoy the space, please also promote *open-source* by giving a ⭐ to our repo [![GitHub Stars](https://img.shields.io/github/stars/gojasper/style-rank?style=social)](https://github.com/gojasper/style-rank)
    """

DISCLAIMER = f"""
This demo is only for research purpose. Jasper cannot be held responsible for the generation of NSFW (Not Safe For Work) content through the use of this demo. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards. Jasper provides the tools, but the responsibility for their use lies with the individual user."""

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
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1024"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"

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

if gr.NO_RELOAD:
    if torch.cuda.is_available():
        vae = AutoencoderKL.from_pretrained(
            "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
        )
        pipe = StableDiffusionXLPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0",
            vae=vae,
            torch_dtype=torch.float16,
            use_safetensors=True,
            variant="fp16",
        )
        pipe.load_ip_adapter(
            "h94/IP-Adapter",
            subfolder="sdxl_models",
            weight_name="ip-adapter_sdxl.safetensors",
        )
        pipe.to(device)

        forward_scheduler = DDIMScheduler.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler"
        )
        invert_scheduler = DDIMInverseScheduler(**forward_scheduler.config)
    else:
        raise ValueError("This demo does not work on CPU.")

css = """
    h1 {
        text-align: center;
        display:block;
    }
    p {
        text-align: justify;
        display:block;
    }
"""


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


def img_to_latents(x: torch.Tensor, vae: AutoencoderKL):
    x = 2.0 * x - 1.0
    posterior = vae.encode(x).latent_dist
    latents = posterior.mean * 0.18215
    return latents


def invert_image(model, image: np.ndarray, n_steps: int, width: int, height: int):

    model.scheduler = invert_scheduler

    image = PIL.Image.fromarray(image).resize((width, height))
    image_tensor = ToTensor()(image).to(model.device, dtype=torch.float16)
    image_tensor = image_tensor.unsqueeze(0)
    latent = img_to_latents(image_tensor, model.vae)

    model.set_ip_adapter_scale(0)

    inv_latents = model(
        prompt="",
        negative_prompt="",
        ip_adapter_image=image,
        guidance_scale=1.0,
        output_type="latent",
        return_dict=False,
        num_inference_steps=n_steps,
        latents=latent,
    )[0]

    return inv_latents


@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,
    style_image_value=None,
    noise_scale: float = 1.5,
) -> PIL.Image.Image:
    torch.manual_seed(seed)

    if style_image_value is None:
        gr.Error("Please provide a style image")

    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

    # Add scaled noise to the latent
    noise = torch.randn(1, 4, width // 8, height // 8).to(device, dtype=torch.float16)

    # Invert the image and get the latent
    if style_image_value is not None:
        latent = invert_image(pipe, style_image_value, 30, width, height)

        latent = latent + noise_scale * noise
        latent = latent / torch.sqrt(
            torch.tensor(1 + noise_scale**2).to(device, dtype=torch.float16)
        )

    else:
        latent = noise

    scale = {
        "up": {"block_0": [0.0, 1.0, 0.0]},
    }
    pipe.set_ip_adapter_scale(scale)

    pipe.scheduler = forward_scheduler
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        ip_adapter_image=style_image_value,
        latents=latent,
        prompt_2=prompt_2,
        negative_prompt_2=negative_prompt_2,
        guidance_scale=guidance_scale_base,
        num_inference_steps=num_inference_steps_base,
        output_type="pil",
    ).images[0]

    return image


examples_prompts = [
    "Astronaut in a jungle, detailed, 8k",
    "A bird",
    "A tiger",
    "A cat",
    "A cactus",
    "A panda",
    "A duck",
    "An elephant",
    "A dragon head",
]

examples_images = glob.glob("./images/*.png")
assert len(examples_images) == len(
    examples_prompts
), "Number of example images and prompts should match"
examples = [[prompt, image] for prompt, image in zip(examples_prompts, examples_images)]

with gr.Blocks(css=css) as demo:
    gr.Markdown(DESCRIPTION)

    gr.Markdown(OPEN_SOURCE_PROMO)

    with gr.Row():

        with gr.Blocks():

            with gr.Column():

                with gr.Row():
                    prompt = gr.Text(
                        label="Target prompt",
                        show_label=False,
                        max_lines=1,
                        placeholder="Enter your prompt to generate new content",
                        container=False,
                    )
                    run_button = gr.Button("Run", scale=0)

                style_image = gr.Image(label="Style reference image")

                noise_scale = gr.Slider(
                    label="Noise Scale",
                    minimum=0,
                    maximum=5,
                    step=0.1,
                    value=1.5,
                )

        with gr.Blocks():

            with gr.Column():
                result = gr.Image(label="Result", show_label=False)

    with gr.Accordion("Advanced options", open=False):
        with gr.Row():
            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(
            label="Negative prompt",
            max_lines=1,
            placeholder="Enter a negative prompt",
            visible=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,
        )

        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(
                label="Guidance scale for base",
                minimum=1,
                maximum=20,
                step=0.1,
                value=5.0,
            )
            num_inference_steps_base = gr.Slider(
                label="Number of inference steps for base",
                minimum=10,
                maximum=100,
                step=1,
                value=25,
            )
        with gr.Row(visible=False) as refiner_params:
            guidance_scale_refiner = gr.Slider(
                label="Guidance scale for refiner",
                minimum=1,
                maximum=20,
                step=0.1,
                value=5.0,
            )
            num_inference_steps_refiner = gr.Slider(
                label="Number of inference steps for refiner",
                minimum=10,
                maximum=100,
                step=1,
                value=25,
            )

    gr.Examples(
        examples=examples,
        inputs=[prompt, style_image],
        outputs=result,
        fn=None,
        run_on_click=False,
    )

    gr.Markdown("## Disclaimer")
    gr.Markdown(DISCLAIMER)

    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,
    )

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit,
            prompt_2.submit,
            negative_prompt_2.submit,
            run_button.click,  # Only trigger on 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,
            style_image,
            noise_scale,
        ],
        outputs=result,
        api_name="run",
    )

if __name__ == "__main__":
    demo.launch()