import random
import os

import spaces
import numpy as np
import torch
from PIL import Image
import huggingface_hub
import gradio as gr

from src.pipeline_flux_kontext_nag import NAGFluxKontextPipeline
from src.transformer_flux import NAGFluxTransformer2DModel


MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048
DEFAULT_GUIDANCE_SCALE = 2.5
DEFAULT_NEGATIVE_PROMPT = "Low resolution, blurry, lack of details"


transformer = NAGFluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-Kontext-dev",
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
)
pipe = NAGFluxKontextPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Kontext-dev",
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)

device = "cuda"
pipe = pipe.to(device)

examples = [
    ["./assets/monster.png", "Transform to 1960s pop art poster style.", "Use a bright pink, green and blue color palette.", 5],
    ["./assets/rabbit.jpg", "Using this elegant style, create a portrait of a cute Godzilla wearing a pearl tiara and lace collar, maintaining the same refined quality and soft color tones.", DEFAULT_NEGATIVE_PROMPT, 5],
]


def get_duration(
        input_image,
        prompt,
        negative_prompt, guidance_scale,
        nag_negative_prompt, nag_scale,
        width, height,
        num_inference_steps,
        seed, randomize_seed,
        compare,
):
    duration = int(num_inference_steps) * 1.5 + 5
    if compare:
        duration *= 1.7
    return duration


@spaces.GPU(duration=get_duration)
def sample(
        input_image,
        prompt,
        negative_prompt=None, guidance_scale=DEFAULT_GUIDANCE_SCALE,
        nag_negative_prompt=None, nag_scale=5.0,
        width=1024, height=1024,
        num_inference_steps=25,
        seed=2025, randomize_seed=False,
        compare=True,
):
    prompt = prompt.strip()
    negative_prompt = negative_prompt.strip() if negative_prompt and negative_prompt.strip() else None
    guidance_scale = float(guidance_scale)
    width, height = int(width), int(height)
    num_inference_steps = int(num_inference_steps)
    
    if (randomize_seed):
        seed = random.randint(0, MAX_SEED)
    else:
        seed = int(seed)

    if input_image is not None:
        input_image = input_image.convert("RGB")

    generator = torch.Generator(device="cuda").manual_seed(seed)
    if input_image is not None:
        image_nag = pipe(
            prompt=prompt,
            image=input_image,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            nag_negative_prompt=nag_negative_prompt,
            nag_scale=nag_scale,
            generator=generator,
            width=input_image.size[0],
            height=input_image.size[1],
            num_inference_steps=num_inference_steps,
        ).images[0]
    else:
        image_nag = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            nag_negative_prompt=nag_negative_prompt,
            nag_scale=nag_scale,
            generator=generator,
            width=width,
            height=height,
            num_inference_steps=num_inference_steps,
        ).images[0]

    if compare:
        generator = torch.Generator(device="cuda").manual_seed(seed)
        if input_image is not None:
            image_normal = pipe(
                prompt=prompt,
                image=input_image,
                negative_prompt=negative_prompt,
                guidance_scale=guidance_scale,
                generator=generator,
                width=input_image.size[0],
                height=input_image.size[1],
                num_inference_steps=num_inference_steps,
            ).images[0]
        else:
            image_normal = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                guidance_scale=guidance_scale,
                generator=generator,
                width=width,
                height=height,
                num_inference_steps=num_inference_steps,
            ).images[0]
    else:
        image_normal = Image.new("RGB", image_nag.size, color=(0, 0, 0))

    return (image_normal, image_nag), seed


def sample_example(
        input_image,
        prompt,
        nag_negative_prompt,
        nag_scale,
):
    outputs, seed = sample(
        input_image=input_image,
        prompt=prompt,
        negative_prompt=None, guidance_scale=DEFAULT_GUIDANCE_SCALE,
        nag_negative_prompt=nag_negative_prompt, nag_scale=nag_scale,
        width=1024, height=1024,
        num_inference_steps=25,
        seed=2025, randomize_seed=False,
        compare=True,
    )
    return outputs, DEFAULT_GUIDANCE_SCALE, 1024, 1024, 25, seed, True


css="""
#col-container {
    margin: 0 auto;
    max-width: 960;
}
"""


with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown('''# Normalized Attention Guidance (NAG) Flux-Kontext-Dev
        NAG demos: [LTX Video Fast](https://huggingface.co/spaces/ChenDY/NAG_ltx-video-distilled), [Wan2.1-T2V-14B](https://huggingface.co/spaces/ChenDY/NAG_wan2-1-fast), [FLUX.1-dev](https://huggingface.co/spaces/ChenDY/NAG_FLUX.1-dev)
        
        Implementation of [Normalized Attention Guidance](https://chendaryen.github.io/NAG.github.io/)
        
        [Paper](https://arxiv.org/abs/2505.21179), [GitHub](https://github.com/ChenDarYen/Normalized-Attention-Guidance), [ComfyUI](https://github.com/ChenDarYen/ComfyUI-NAG)
        ''')
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Upload the image for editing", type="pil")
                prompt = gr.Textbox(
                    label="Prompt",
                    max_lines=3,
                    placeholder="Enter your prompt",
                )
                nag_negative_prompt = gr.Textbox(
                    label="Negative Prompt for NAG",
                    value=DEFAULT_NEGATIVE_PROMPT,
                    max_lines=3,
                )
                nag_scale = gr.Slider(label="NAG Scale", minimum=1., maximum=20., step=0.25, value=5.)
                compare = gr.Checkbox(label="Compare with baseline", info="If unchecked, only sample with NAG will be generated.", value=True)
                button = gr.Button("Generate", min_width=120)
                with gr.Accordion("Advanced Settings", open=False):
                    negative_prompt = gr.Textbox(label="Negative Prompt", value=None, visible=False)
                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=1., maximum=15., step=0.1, value=DEFAULT_GUIDANCE_SCALE)
                    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,
                        )
                    num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=25)
                    seed = gr.Slider(label="Seed", minimum=1, maximum=MAX_SEED, step=1, randomize=True)
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

            with gr.Column():
                output = gr.ImageSlider(label="Left: Baseline, Right: With NAG", interactive=False)

        gr.Examples(
            examples=examples,
            fn=sample_example,
            inputs=[
                input_image,
                prompt,
                nag_negative_prompt,
                nag_scale,
            ],
            outputs=[output, guidance_scale, width, height, num_inference_steps, seed, compare],
            cache_examples="lazy",
        )

    gr.on(
        triggers=[
            button.click,
            prompt.submit
        ],
        fn=sample,
        inputs=[
            input_image,
            prompt,
            negative_prompt, guidance_scale,
            nag_negative_prompt, nag_scale,
            width, height,
            num_inference_steps,
            seed, randomize_seed,
            compare,
        ],
        outputs=[output, seed],
    )


if __name__ == "__main__":
    huggingface_hub.login(os.getenv('HF_TOKEN'))
    demo.launch(share=True)