#!/usr/bin/env python3
from diffusers import DiffusionPipeline
import torch
from diffusers.models.attention_processor import AttnAddedKVProcessor2_0, XFormersAttnAddedKVProcessor, AttnAddedKVProcessor

import time
import os
from huggingface_hub import HfApi
from pathlib import Path

api = HfApi()
prev_time = time.time()


prompt = "a picture of elon musk next to a rocket"
negative_prompt = "low quality, ugly"

pipe_prior = DiffusionPipeline.from_pretrained(
    "kandinsky-community/kandinsky-2-1-prior", torch_dtype=torch.float16
)
pipe_prior.to("cuda")
t2i_pipe = DiffusionPipeline.from_pretrained("kandinsky-community/kandinsky-2-1", torch_dtype=torch.float16)
t2i_pipe.to("cuda")

t2i_pipe.unet.set_attn_processor(AttnAddedKVProcessor2_0())
t2i_pipe.unet.to(memory_format=torch.channels_last)
t2i_pipe.unet = torch.compile(t2i_pipe.unet, mode="reduce-overhead", fullgraph=True)

next_time = time.time()
print("Loading", next_time - prev_time)
prev_time = next_time

generator = torch.Generator(device="cuda").manual_seed(12)
image_embeds, negative_image_embeds = pipe_prior(prompt, negative_prompt=negative_prompt, generator=generator).to_tuple()

next_time = time.time()
print("Prior", next_time - prev_time)
prev_time = next_time

for _ in range(3):
    images = t2i_pipe(prompt, image_embeds=image_embeds, negative_image_embeds=negative_image_embeds, negative_prompt=negative_prompt, num_inference_steps=50, generator=generator).images

next_time = time.time()
print("Text-to-image", next_time - prev_time)
prev_time = next_time

for i, image in enumerate(images):
    path = os.path.join(Path.home(), "images", f"aa_{i}.png")
    image.save(path)

    api.upload_file(
        path_or_fileobj=path,
        path_in_repo=path.split("/")[-1],
        repo_id="patrickvonplaten/images",
        repo_type="dataset",
    )
    print(f"https://huggingface.co/datasets/patrickvonplaten/images/blob/main/aa_{i}.png")