#!/usr/bin/env python3
#!/usr/bin/env python3
from diffusers import StableDiffusionXLPipeline, AutoencoderKL
import torch

path = "stabilityai/stable-diffusion-xl-base-1.0"
vae_path = "madebyollin/sdxl-vae-fp16-fix"

# vae = AutoencoderKL.from_pretrained(vae_path, torch_dtype=torch.float16)
# pipe = StableDiffusionXLPipeline.from_pretrained(path, torch_dtype=torch.float16, vae=vae, variant="fp16", use_safetensors=True, local_files_only=True, add_watermarker=False)
pipe = StableDiffusionXLPipeline.from_pretrained(path, torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
pipe.to("cuda")

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

for _ in range(5):
    image = pipe(prompt=prompt, num_inference_steps=steps).images[0]