#!/usr/bin/env python3
#from diffusers import DiffusionPipeline, StableDiffusionPipeline, KDPM2DiscreteScheduler, KDPM2AncestralDiscreteScheduler, HeunDiscreteScheduler
from diffusers import DiffusionPipeline, StableDiffusionPipeline, HeunDiscreteScheduler
import torch
import os

seed = 33
inference_steps = 25

old_pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", custom_pipeline="sd_text2img_k_diffusion")
old_pipe = old_pipe.to("cuda")

for prompt in ["astronaut riding horse", "whale falling from sky", "magical forest", "highly photorealistic picture of johnny depp"]:
    # for sampler in ["sample_dpm_2_ancestral", "sample_dpm_2", "sample_heun"]:
#    for sampler in ["heun", "sample_dpm_2_ancestral", "sample_dpm_2"]:
    for sampler in ["sample_heun"]:
        old_pipe.set_sampler(sampler)
        torch.manual_seed(0)
        image = old_pipe(prompt, num_inference_steps=inference_steps).images[0]
        folder = f"a_{'_'.join(prompt.split())}"
        os.makedirs(f"/home/patrick_huggingface_co/images/{folder}", exist_ok=True)
        image.save(f"/home/patrick_huggingface_co/images/{folder}/{sampler}.png")

        pipe = StableDiffusionPipeline(**old_pipe.components)
        pipe = pipe.to("cuda")

#        if sampler == "sample_dpm_2":
#            pipe.scheduler = KDPM2DiscreteScheduler.from_config(pipe.scheduler.config)
#        elif sampler == "sample_dpm_2_ancestral":
#            pipe.scheduler = KDPM2AncestralDiscreteScheduler.from_config(pipe.scheduler.config)
        if sampler == "sample_heun":
            pipe.scheduler = HeunDiscreteScheduler.from_config(pipe.scheduler.config)

        torch.manual_seed(0)
        image = pipe(prompt, num_inference_steps=inference_steps).images[0]

        image.save(f"/home/patrick_huggingface_co/images/{folder}/hf_{sampler}.png")
    break