import os
import gc
import gradio as gr
import numpy as np
import torch
import json
import spaces
import config
import utils
import logging
from PIL import Image, PngImagePlugin
from datetime import datetime
from diffusers.models import AutoencoderKL
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

DESCRIPTION = "Animagine XL 3.1"
if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU. </p>"
IS_COLAB = utils.is_google_colab() or os.getenv("IS_COLAB") == "1"
HF_TOKEN = os.getenv("HF_TOKEN")
CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
MIN_IMAGE_SIZE = int(os.getenv("MIN_IMAGE_SIZE", "512"))
MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./outputs")

MODEL = os.getenv(
    "MODEL",
    "https://huggingface.co/cagliostrolab/animagine-xl-3.1/blob/main/animagine-xl-3.1.safetensors",
)

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

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

def load_pipeline(model_name):
    vae = AutoencoderKL.from_pretrained(
        "madebyollin/sdxl-vae-fp16-fix",
        torch_dtype=torch.float16,
    )
    pipeline = (
        StableDiffusionXLPipeline.from_single_file
        if MODEL.endswith(".safetensors")
        else StableDiffusionXLPipeline.from_pretrained
    )
    
    img_pipeline = (
        StableDiffusionXLImg2ImgPipeline.from_single_file
        if MODEL.endswith(".safetensors")
        else StableDiffusionXLImg2ImgPipeline.from_pretrained
    )

    pipe = pipeline(
        model_name,
        vae=vae,
        torch_dtype=torch.float16,
        custom_pipeline="lpw_stable_diffusion_xl",
        use_safetensors=True,
        add_watermarker=False,
        use_auth_token=HF_TOKEN,
    )

    img_pipe = img_pipeline(
        model_name,
        vae=vae,
        torch_dtype=torch.float16,
        custom_pipeline="lpw_stable_diffusion_xl",
        use_safetensors=True,
        add_watermarker=False,
        use_auth_token=HF_TOKEN,
    )

    pipe.to(device)
    img_pipe.to(device)
    return pipe, img_pipe

def load_img(resize_width,img: str):
    img = Image.open(img)
    width, height = img.size
    scale = resize_width / width
    resize_height = int(height * scale)
    
    img = img.resize((resize_width, resize_height), Image.Resampling.LANCZOS)
    return img, resize_width, resize_height

@spaces.GPU
def example_generate(
    prompt: str,
    negative_prompt: str = "",
    seed: int = 0,
    custom_width: int = 1024,
    custom_height: int = 1024,
    guidance_scale: float = 7.0,
    num_inference_steps: int = 28,
    sampler: str = "Euler a",
    aspect_ratio_selector: str = "896 x 1152",
    style_selector: str = "(None)",
    quality_selector: str = "Standard v3.1",
    use_upscaler: bool = False,
    upscaler_strength: float = 0.55,
    upscale_by: float = 1.5,
    add_quality_tags: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    generator = utils.seed_everything(seed)

    width, height = utils.aspect_ratio_handler(
        aspect_ratio_selector,
        custom_width,
        custom_height,
    )

    prompt = utils.add_wildcard(prompt, wildcard_files)

    prompt, negative_prompt = utils.preprocess_prompt(
        quality_prompt, quality_selector, prompt, negative_prompt, add_quality_tags
    )
    prompt, negative_prompt = utils.preprocess_prompt(
        styles, style_selector, prompt, negative_prompt
    )

    width, height = utils.preprocess_image_dimensions(width, height)

    backup_scheduler = pipe.scheduler
    pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)

    if use_upscaler:
        upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
    metadata = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "resolution": f"{width} x {height}",
        "guidance_scale": guidance_scale,
        "num_inference_steps": num_inference_steps,
        "seed": seed,
        "sampler": sampler,
        "sdxl_style": style_selector,
        "add_quality_tags": add_quality_tags,
        "quality_tags": quality_selector,
    }

    if use_upscaler:
        new_width = int(width * upscale_by)
        new_height = int(height * upscale_by)
        metadata["use_upscaler"] = {
            "upscale_method": "nearest-exact",
            "upscaler_strength": upscaler_strength,
            "upscale_by": upscale_by,
            "new_resolution": f"{new_width} x {new_height}",
        }
    else:
        metadata["use_upscaler"] = None
        metadata["Model"] = {
            "Model": DESCRIPTION,
            "Model hash": "e3c47aedb0",
        }
    
    logger.info(json.dumps(metadata, indent=4))

    try:
        if use_upscaler:
            latents = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                generator=generator,
                output_type="latent",
            ).images
            upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
            images = upscaler_pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                image=upscaled_latents,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                strength=upscaler_strength,
                generator=generator,
                output_type="pil",
            ).images
        else:
            images = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                guidance_scale=guidance_scale,
                num_inference_steps=num_inference_steps,
                generator=generator,
                output_type="pil",
            ).images

        if images:
            image_paths = [
                utils.save_image(image, metadata, OUTPUT_DIR, IS_COLAB)
                for image in images
            ]

            for image_path in image_paths:
                logger.info(f"Image saved as {image_path} with metadata")

        return image_paths, metadata
    except Exception as e:
        logger.exception(f"An error occurred: {e}")
        raise
    finally:
        if use_upscaler:
            del upscaler_pipe
        pipe.scheduler = backup_scheduler
        utils.free_memory()

@spaces.GPU
def generate(
    prompt: str,
    negative_prompt: str = "",
    seed: int = 0,
    custom_width: int = 1024,
    custom_height: int = 1024,
    guidance_scale: float = 7.0,
    num_inference_steps: int = 28,
    sampler: str = "Euler a",
    aspect_ratio_selector: str = "896 x 1152",
    style_selector: str = "(None)",
    quality_selector: str = "Standard v3.1",
    use_upscaler: bool = False,
    upscaler_strength: float = 0.55,
    upscale_by: float = 1.5,
    add_quality_tags: bool = True,
    isImg2Img: bool = True,
    img_path: str= "",
    img2img_strength: float=0.65,
    
    progress=gr.Progress(track_tqdm=True),
):
    generator = utils.seed_everything(seed)

    width, height = utils.aspect_ratio_handler(
        aspect_ratio_selector,
        custom_width,
        custom_height,
    )

    prompt = utils.add_wildcard(prompt, wildcard_files)

    prompt, negative_prompt = utils.preprocess_prompt(
        quality_prompt, quality_selector, prompt, negative_prompt, add_quality_tags
    )
    prompt, negative_prompt = utils.preprocess_prompt(
        styles, style_selector, prompt, negative_prompt
    )

    width, height = utils.preprocess_image_dimensions(width, height)

    backup_scheduler = pipe.scheduler
    pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)

    img_backup_scheduler = img_pipe.scheduler
    img_pipe.scheduler = utils.get_scheduler(img_pipe.scheduler.config, sampler)

    if use_upscaler:
        upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
    metadata = {
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "resolution": f"{width} x {height}",
        "guidance_scale": guidance_scale,
        "num_inference_steps": num_inference_steps,
        "seed": seed,
        "sampler": sampler,
        "sdxl_style": style_selector,
        "add_quality_tags": add_quality_tags,
        "quality_tags": quality_selector,
        "isImg2Img": isImg2Img,
        "img_path": img_path,
        "img2img_strength": img2img_strength
    }

    if use_upscaler:
        new_width = int(width * upscale_by)
        new_height = int(height * upscale_by)
        metadata["use_upscaler"] = {
            "upscale_method": "nearest-exact",
            "upscaler_strength": upscaler_strength,
            "upscale_by": upscale_by,
            "new_resolution": f"{new_width} x {new_height}",
        }
    else:
        metadata["use_upscaler"] = None
        metadata["Model"] = {
            "Model": DESCRIPTION,
            "Model hash": "e3c47aedb0",
        }
    
    logger.info(json.dumps(metadata, indent=4))

    try:
        if use_upscaler:
            if isImg2Img:
                print("Img2Img")
                img, img_width, img_height = load_img(1024, img_path)
                latents = img_pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    width=img_width,
                    height=img_height,
                    image=img,
                    strength=img2img_strength,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    generator=generator,
                    output_type="latent",
                ).images
                upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
                images = upscaler_pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    image=upscaled_latents,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    strength=upscaler_strength,
                    generator=generator,
                    output_type="pil",
                ).images
            else:
                latents = pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    width=width,
                    height=height,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    generator=generator,
                    output_type="latent",
                ).images
                upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
                images = upscaler_pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    image=upscaled_latents,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    strength=upscaler_strength,
                    generator=generator,
                    output_type="pil",
                ).images
        else:
            if isImg2Img:
                print("Img2Img")
                img, img_width, img_height = load_img(512, img_path)
                images = img_pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    width=img_width,
                    height=img_height,
                    image=img,
                    strength=img2img_strength,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    generator=generator,
                    output_type="pil",
                ).images
            else:
                images = pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    width=width,
                    height=height,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_inference_steps,
                    generator=generator,
                    output_type="pil",
                ).images

        if images:
            image_paths = [
                utils.save_image(image, metadata, OUTPUT_DIR, IS_COLAB)
                for image in images
            ]

            for image_path in image_paths:
                logger.info(f"Image saved as {image_path} with metadata")

        return image_paths, metadata
    except Exception as e:
        logger.exception(f"An error occurred: {e}")
        raise
    finally:
        if use_upscaler:
            del upscaler_pipe

        if isImg2Img:
            img_pipe.scheduler = img_backup_scheduler
        else:
            pipe.scheduler = backup_scheduler
        
        utils.free_memory()

def fake_generate(*args,use_upscaler=False,**kwargs):
    args = ",".join(args)

    #result, metadata = generate(args, use_upscaler=use_upscaler)
    
    return None, None

if torch.cuda.is_available():
    pipe, img_pipe = load_pipeline(MODEL)
    logger.info("Loaded on Device!")
else:
    pipe, img_pipe = None, None

styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in config.style_list}
quality_prompt = {
    k["name"]: (k["prompt"], k["negative_prompt"]) for k in config.quality_prompt_list
}

wildcard_files = utils.load_wildcard_files("wildcard")

with gr.Blocks(css="style.css", theme="NoCrypt/miku@1.2.1") as demo:
    title = gr.HTML(
        f"""<h1><span>{DESCRIPTION}</span></h1>""",
        elem_id="title",
    )
    gr.Markdown(
        f"""Gradio demo for [cagliostrolab/animagine-xl-3.1](https://huggingface.co/cagliostrolab/animagine-xl-3.1)""",
        elem_id="subtitle",
    )
    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=2):
            with gr.Tab("Input"):
                with gr.Group():
                    prompt = gr.Text(
                        label="Prompt",
                        max_lines=5,
                        placeholder="Enter your prompt",
                    )
                    negative_prompt = gr.Text(
                        label="Negative Prompt",
                        max_lines=5,
                        placeholder="Enter a negative prompt",
                    )
                    with gr.Accordion(label="Img2Img", open=False):
                        isImg2Img = gr.Checkbox(
                            label="Enable Img2Img", value=False
                        )
                        image = gr.Image(
                            sources=["upload", "webcam", "clipboard"],
                            type="filepath",
                            #visible=False,
                            label="Choose Image"
                        )
                        img2img_strength = gr.Slider(
                            minimum=0.05,
                            maximum=1,
                            step=0.05,
                            value=0.65,
                            label="Strength",
                            #visible=False
                        )
                    with gr.Accordion(label="Quality Tags", open=True):
                        add_quality_tags = gr.Checkbox(
                            label="Add Quality Tags", value=True
                        )
                        quality_selector = gr.Dropdown(
                            label="Quality Tags Presets",
                            interactive=True,
                            choices=list(quality_prompt.keys()),
                            value="Standard v3.1",
                        )
            with gr.Tab("ControlNet"):
                with gr.Group():
                    gr.Label(label="ControlNet is not available now")
                    use_controlnet = gr.Checkbox(
                        label="Use ControlNet",
                        value=False
                    )
                    controlnet_type = gr.Radio(
                        label="ControlNet",
                        choices=["Canny", "Depth", "OpenPose"]                        
                    )
                    
            with gr.Tab("Advanced Settings"):
                with gr.Group():
                    style_selector = gr.Radio(
                        label="Style Preset",
                        container=True,
                        interactive=True,
                        choices=list(styles.keys()),
                        value="(None)",
                    )
                with gr.Group():
                    aspect_ratio_selector = gr.Radio(
                        label="Aspect Ratio",
                        choices=config.aspect_ratios,
                        value="896 x 1152",
                        container=True,
                    )
                with gr.Group(visible=False) as custom_resolution:
                    with gr.Row():
                        custom_width = gr.Slider(
                            label="Width",
                            minimum=MIN_IMAGE_SIZE,
                            maximum=MAX_IMAGE_SIZE,
                            step=8,
                            value=1024,
                        )
                        custom_height = gr.Slider(
                            label="Height",
                            minimum=MIN_IMAGE_SIZE,
                            maximum=MAX_IMAGE_SIZE,
                            step=8,
                            value=1024,
                        )
                with gr.Group():
                    use_upscaler = gr.Checkbox(label="Use Upscaler", value=False)
                    with gr.Row() as upscaler_row:
                        upscaler_strength = gr.Slider(
                            label="Strength",
                            minimum=0,
                            maximum=1,
                            step=0.05,
                            value=0.55,
                            visible=False,
                        )
                        upscale_by = gr.Slider(
                            label="Upscale by",
                            minimum=1,
                            maximum=1.5,
                            step=0.1,
                            value=1.5,
                            visible=False,
                        )
                with gr.Group():
                    sampler = gr.Dropdown(
                        label="Sampler",
                        choices=config.sampler_list,
                        interactive=True,
                        value="Euler a",
                    )
                with gr.Group():
                    seed = gr.Slider(
                        label="Seed", minimum=0, maximum=utils.MAX_SEED, step=1, value=0
                    )
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                with gr.Group():
                    with gr.Row():
                        guidance_scale = gr.Slider(
                            label="Guidance scale",
                            minimum=1,
                            maximum=12,
                            step=0.1,
                            value=7.0,
                        )
                        num_inference_steps = gr.Slider(
                            label="Number of inference steps",
                            minimum=1,
                            maximum=50,
                            step=1,
                            value=28,
                        )
        with gr.Column(scale=3):
            with gr.Blocks():
                run_button = gr.Button("Generate", variant="primary")
            result = gr.Gallery(
                label="Result", 
                columns=1, 
                height='100%', 
                preview=True, 
                show_label=False
            )
            with gr.Accordion(label="Generation Parameters", open=False):
                gr_metadata = gr.JSON(label="metadata", show_label=False)
            gr.Examples(
                examples=config.examples,
                inputs=prompt,
                outputs=[result, gr_metadata],
                fn=lambda *args, **kwargs: example_generate(*args, use_upscaler=True, **kwargs),
                cache_examples=CACHE_EXAMPLES,
            )
    use_upscaler.change(
        fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
        inputs=use_upscaler,
        outputs=[upscaler_strength, upscale_by],
        queue=False,
        api_name=False,
    )
    #isImg2Img.change(
    #    fn=lambda x: [gr.update(visible=x), gr.update(visible=x)],
    #    inputs=isImg2Img,
    #    outputs=[image, img2img_strength],
    #    queue=False,
    #    api_name=False,
    #)
    aspect_ratio_selector.change(
        fn=lambda x: gr.update(visible=x == "Custom"),
        inputs=aspect_ratio_selector,
        outputs=custom_resolution,
        queue=False,
        api_name=False,
    )

    gr.on(
        triggers=[
            prompt.submit,
            negative_prompt.submit,
            run_button.click,
        ],
        fn=utils.randomize_seed_fn,
        inputs=[seed, randomize_seed],
        outputs=seed,
        queue=False,
        api_name=False,
    ).then(
        fn=generate,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            custom_width,
            custom_height,
            guidance_scale,
            num_inference_steps,
            sampler,
            aspect_ratio_selector,
            style_selector,
            quality_selector,
            use_upscaler,
            upscaler_strength,
            upscale_by,
            add_quality_tags,
            isImg2Img,
            image,
            img2img_strength
        ],
        outputs=[result, gr_metadata],
        api_name="run",
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(debug=IS_COLAB, share=IS_COLAB)