import spaces

import torch
import gradio as gr
from PIL import Image

from transformers import (
    CLIPTokenizer,
    CLIPTextModel
)

from diffusers import (
    UNet2DConditionModel,
    DDPMScheduler
)

from safetensors.torch import load_file


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


tokenizer = CLIPTokenizer.from_pretrained(
    "openai/clip-vit-base-patch32"
)

text_encoder = CLIPTextModel.from_pretrained(
    "openai/clip-vit-base-patch32"
)

text_encoder.eval()


unet = UNet2DConditionModel(
    sample_size=128,

    in_channels=3,
    out_channels=3,

    layers_per_block=2,

    block_out_channels=(
        128,
        256,
        512,
        512
    ),

    down_block_types=(
        "DownBlock2D",
        "DownBlock2D",
        "CrossAttnDownBlock2D",
        "DownBlock2D"
    ),

    up_block_types=(
        "UpBlock2D",
        "CrossAttnUpBlock2D",
        "UpBlock2D",
        "UpBlock2D"
    ),

    cross_attention_dim=512
)


state = load_file(
    "model.safetensors"
)

unet.load_state_dict(
    state
)

unet.eval()


scheduler = DDPMScheduler(
    num_train_timesteps=1000,
    beta_schedule="scaled_linear",
    prediction_type="epsilon"
)



@spaces.GPU
def generate(prompt, steps, seed, progress=gr.Progress()):

    if seed == 0:
        seed = torch.randint(
            0,
            2**32 - 1,
            (1,)
        ).item()

    generator = torch.Generator(
        device="cuda"
    ).manual_seed(
        int(seed)
    )
    unet.to("cuda")
    text_encoder.to("cuda")


    tokens = tokenizer(
        prompt,
        padding="max_length",
        max_length=77,
        truncation=True,
        return_tensors="pt"
    )

    tokens = {
        k:v.to("cuda")
        for k,v in tokens.items()
    }


    with torch.no_grad():

        text = text_encoder(
            **tokens
        ).last_hidden_state


    image = torch.randn(
        (1,3,128,128),
        generator=generator,
        device="cuda"
    )


    scheduler.set_timesteps(steps)


    with torch.no_grad():

        for i, t in enumerate(scheduler.timesteps):

            progress(
                i / len(scheduler.timesteps),
                desc=f"Diffusion step {i+1}/{steps}"
            )

            noise_pred = unet(
                image,
                t,
                encoder_hidden_states=text
            ).sample


            image = scheduler.step(
                noise_pred,
                t,
                image
            ).prev_sample



    image = (
        image
        .clamp(-1,1)
        .add(1)
        .div(2)
    )


    image = (
        image[0]
        .permute(1,2,0)
        .cpu()
        .numpy()
    )


    image = (
        image * 255
    ).astype("uint8")


    result = Image.fromarray(image)


    unet.to("cpu")
    text_encoder.to("cpu")
    torch.cuda.empty_cache()


    return result, seed




# ==========================
# UI
# ==========================

examples = [
    ["Eevee, a cute brown fox-like pokemon with fluffy fur"],
    ["Pikachu, a small yellow electric pokemon with red cheeks"],
    ["Charizard, a large orange dragon pokemon with blue wings"],
    ["a legendary ice pokemon covered in crystals and snow"],
    ["a robotic electric pokemon made of metal and technology"],
    ["a tiny cute bug pokemon sitting on a leaf"],
]


with gr.Blocks(
    theme=gr.themes.Soft()
) as demo:

    history = gr.State([])

    gr.Markdown(
        """
        # ⚡ Pikadiffusion
        
        A 128×128 Pokémon text-to-image diffusion model.

        Generate Pokémon-style creatures from text prompts.
        """
    )


    with gr.Row():

        with gr.Column():

            prompt = gr.Textbox(
                label="Pokemon prompt",
                value="Eevee, a cute brown fox-like pokemon with fluffy fur",
                lines=3
            )


            steps = gr.Slider(
                minimum=50,
                maximum=1000,
                value=250,
                step=50,
                label="Diffusion steps"
            )


            seed = gr.Number(
                value=0,
                label="Seed (0 = random)"
            )


            generate_btn = gr.Button(
                "⚡ Generate",
                variant="primary"
            )


        with gr.Column():

            output = gr.Image(
                label="Generated Pokemon",
                type="pil"
            )

            seed_out = gr.Number(
                label="Used seed"
            )


    gr.Examples(
        examples=examples,
        inputs=prompt
    )


    generate_btn.click(
        fn=generate,
        inputs=[
            prompt,
            steps,
            seed
        ],
        outputs=[
            output,
            seed_out,
        ]
    )
demo.launch()