#!/usr/bin/env python3
from diffusers import DiffusionPipeline, KandinskyPriorPipeline, DDPMScheduler, DDIMScheduler
import torch

import os
from huggingface_hub import HfApi
from pathlib import Path

api = HfApi()

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

prompt = "A alien cheeseburger creature eating itself, claymation, cinematic, moody lighting"
negative_prompt = "low quality, bad quality"

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

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

images = t2i_pipe(prompt, num_images_per_prompt=4, image_embeds=image_embeds, negative_image_embeds=negative_image_embeds, negative_prompt=negative_prompt).images

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")