#!/usr/bin/env python3
from diffusers import StableDiffusionXLImg2ImgPipeline, EulerDiscreteScheduler, UniPCMultistepScheduler, DEISMultistepScheduler, HeunDiscreteScheduler, DPMSolverMultistepScheduler
import time
import hf_image_uploader as hiu
import numpy as np
from huggingface_hub import HfApi
import torch
import sys

path = sys.argv[1]

api = HfApi()
start_time = time.time()
pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(path, torch_dtype=torch.float16)
pipe.to("cuda")


prompt = "An astronaut riding a green horse on Mars"


init_image = torch.from_numpy(np.load("/home/patrick/images/xl_latents.npy")).to("cuda")





for scheduler_cls, kwargs in [
        (EulerDiscreteScheduler, {}), 
        (EulerDiscreteScheduler, {"use_karras_sigmas": True}), 
        (UniPCMultistepScheduler, {}),
        (UniPCMultistepScheduler, {"use_karras_sigmas": True}), 
        (DEISMultistepScheduler, {}),
        (DEISMultistepScheduler, {"use_karras_sigmas": True}), 
        (HeunDiscreteScheduler, {}),
        (HeunDiscreteScheduler, {"use_karras_sigmas": True}), 
        (DPMSolverMultistepScheduler, {}),
        (DPMSolverMultistepScheduler, {"use_karras_sigmas": True, "algorithm_type": "sd-dpmsolver++"}),
]:
    for num_steps in [11,12]:
        pipe.scheduler = scheduler_cls.from_config(pipe.scheduler.config, **kwargs)
        image = pipe(prompt=prompt, num_inference_steps=num_steps, image=init_image).images[0]
        hiu.upload(image, "patrickvonplaten/images")