#!/usr/bin/env python3
import torch

from diffusers import AutoencoderKL, StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, KDPM2AncestralDiscreteScheduler

vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix",
    torch_dtype=torch.float16
)

base = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
    vae=vae,
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
)

scheduler = KDPM2AncestralDiscreteScheduler.from_config(base.scheduler.config, use_karras_sigmas=True)
base.scheduler = scheduler

base.to("cuda")

def print_step(s, t, latents):
    print(s)

generator=torch.manual_seed(1111)

images = base(
    prompt="LOVE",
    num_inference_steps=10,
    generator=generator,
    callback=print_step
).images