diffusers/tools

1#!/usr/bin/env python32from diffusers import FlaxStableDiffusionXLPipeline3import numpy as np4import jax.numpy as jnp5import jax6 7path = "hf-internal-testing/tiny-stable-diffusion-xl-pipe"8 9pipe, params = FlaxStableDiffusionXLPipeline.from_pretrained(path)10 11prompt = "An astronaut riding a green horse on Mars"12steps = 313 14batch_size, height, width, ch = 1, 32, 32, 415num_elems = batch_size * height * width * ch16rng = jax.random.PRNGKey(0)17latents = (jnp.arange(num_elems) / num_elems)[:, None, None, None].reshape(batch_size, ch, width, height)18 19print("latents", np.abs(np.asarray(latents)).sum())20 21prompt_embeds = pipe.prepare_inputs(prompt)22 23image = pipe(prompt_embeds, params, rng, latents=latents, num_inference_steps=3, output_type="np").images[0]24 25print(np.abs(np.asarray(image)).sum())26