from __future__ import annotations

import math
import random

import gradio as gr
import torch
from PIL import Image, ImageOps
from diffusers import StableDiffusionSAGPipeline


help_text = """
Self-Attention Guidance (SAG) - Condition-Agnostic Diffusion Guidance Using the Internal Self-Attention
"""


examples = [
    [
        ' ',
        50,
        "Fix Seed",
        123,
        7.5,
        1.0,
    ],
    [
        '.',
        50,
        "Fix Seed",
        456,
        7.5,
        1.0,
    ],
    [
        'A cute Scottish Fold playing with a ball',
        50,
        "Fix Seed",
        789,
        5.0,
        1.0,
    ],
    [
        'A person with a happy dog',
        50,
        "Fix Seed",
        901,
        5.0,
        1.0,
    ],
]


model_id = "runwayml/stable-diffusion-v1-5"

def main():
    pipe = StableDiffusionSAGPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to('cuda')

    def generate(
        prompt: str,
        steps: int,
        randomize_seed: bool,
        seed: int,
        cfg_scale: float,
        sag_scale: float,
    ):
        seed = random.randint(0, 100000) if randomize_seed else seed

        generator = torch.manual_seed(seed)
        ori_image = pipe(prompt, generator=generator, guidance_scale=cfg_scale, sag_scale=0.0).images[0]
        generator = torch.manual_seed(seed)
        sag_image = pipe(prompt, generator=generator, guidance_scale=cfg_scale, sag_scale=sag_scale).images[0]
        return [ori_image, sag_image]

    def reset():
        return [0, "Randomize Seed", 1371, 5.0, 0.75, None, None]

    with gr.Blocks() as demo:
        gr.HTML("""<h1 style="font-weight: 900; margin-bottom: 7px;">
           Self-Attention Guidance (SAG) Demo
        """)
        with gr.Row():
            with gr.Column(scale=5):
                prompt = gr.Textbox(lines=1, label="Enter your prompt", interactive=True)
            with gr.Column(scale=1, min_width=60):
                generate_button = gr.Button("Generate")
            with gr.Column(scale=1, min_width=60):
                reset_button = gr.Button("Reset")

        with gr.Row():
            steps = gr.Number(value=50, precision=0, label="Steps", interactive=True)
            randomize_seed = gr.Radio(
                ["Fix Seed", "Randomize Seed"],
                label="Seed Type",
                value="Fix Seed",
                type="index",
                show_label=False,
                interactive=True,
            )
            seed = gr.Number(value=8978, precision=0, label="Seed", interactive=True)
            
        with gr.Row():
            cfg_scale = gr.Slider(
                label="Guidance Scale", minimum=0, maximum=10, value=5.0, step=0.1
            )
            sag_scale = gr.Slider(
                label="Self-Attention Guidance Scale", minimum=0, maximum=1.0, value=0.75, step=0.05
            )

        with gr.Row():
            ori_image = gr.Image(label="CFG", type="pil", interactive=False)
            sag_image = gr.Image(label="SAG + CFG", type="pil", interactive=False)
            ori_image.style(height=512, width=512)
            sag_image.style(height=512, width=512)

            
        ex = gr.Examples(
            examples=examples,
            fn=generate,
            inputs=[
                prompt,
                steps,
                randomize_seed,
                seed,
                cfg_scale,
                sag_scale,
            ],
            outputs=[ori_image, sag_image],
            cache_examples=False,
        )

        gr.Markdown(help_text)
        
        generate_button.click(
            fn=generate,
            inputs=[
                prompt,
                steps,
                randomize_seed,
                seed,
                cfg_scale,
                sag_scale,
            ],
            outputs=[ori_image, sag_image],
        )
        reset_button.click(
            fn=reset,
            inputs=[],
            outputs=[steps, randomize_seed, seed, cfg_scale, sag_scale, ori_image, sag_image],
        )

    demo.queue(concurrency_count=1)
    demo.launch(share=False)


if __name__ == "__main__":
    main()