#!/usr/bin/env python3
import torch
import numpy as np
from diffusers import StableDiffusionXLPipeline

path = "hf-internal-testing/tiny-stable-diffusion-xl-pipe"

pipe = StableDiffusionXLPipeline.from_pretrained(path)
pipe.unet.set_default_attn_processor()

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

batch_size, height, width, ch = 1, 32, 32, 4
num_elems = batch_size * height * width * ch
latents = (torch.arange(num_elems) / num_elems)[:, None, None, None].reshape(batch_size, ch, width, height)
print("latents", latents.abs().sum())

image = pipe(prompt, latents=latents, num_inference_steps=3, output_type="np", guidance_scale=7.5).images[0]

print(np.abs(image).sum())