Want to generate images from text without training a model from scratch? Stable Diffusion is a pretrained text-to-image diffusion model used through pipelines such as Hugging Face Diffusers. A textual encoder represents the prompt, then a multistep diffusion process generates the image. section 11.3.3
A Diffusers workflow combines:
- A pretrained Stable Diffusion pipeline for inference. section 11.3.3
- Interchangeable noise schedulers that trade off diffusion speed and output quality. section 11.3.3
- Model components that can be recombined into custom end-to-end diffusion systems. section 11.3.3
Loading the model with JAX
A Flax pipeline loads the runwayml/stable-diffusion-v1-5 checkpoint with bfloat16 parameters. The surrounding example checks the JAX backend and generates eight images, demonstrating TPU-oriented inference. section 11.3.3
Listing 11.15 Generating images using the diffusers library
from jax.lib import xla_bridge #1
print(xla_bridge.get_backend().platform) #1
>>> tpu
import jax #2
import jax.numpy as jnp #2
from flax.jax_utils import replicate #2
from flax.training.common_utils import shard #2
from PIL import Image #2
from diffusers import FlaxStableDiffusionPipeline #2
pipeline, params = FlaxStableDiffusionPipeline.from_pretrained( #3
"runwayml/stable-diffusion-v1-5", #3
revision="bf16", #4
dtype=jnp.bfloat16, #4
)
prompts = [ #5
"Three leopards holding a scarlet rose and in the centre a green rat.",
"A large glass bottle full of tiny elephants of different colors",
"HAL-9000 in the style of Picasso",
"Pink panther drinking a coffee. ",
"Golden tiger fighting vikings in a modern city, Van Gogh style",
"A blue Axolotl sitting on an emerald throne.",
"A pumpkin-shape spaceship spinning around the Moon",
"Black panther doing his homework",
]
p_params = replicate(params) #6
prompt_ids = pipeline.prepare_inputs(prompts) #7
prompt_ids = shard(prompt_ids) #7
seed = 42 #8
rng = jax.random.PRNGKey(seed) #8
rng = jax.random.split(rng, jax.device_count()) #8
images = pipeline(prompt_ids, p_params, rng, jit=True).images #9
def image_grid(imgs, rows, cols): #10
w,h = imgs[0].size
grid = Image.new('RGB', size=(cols*w, rows*h))
for i, img in enumerate(imgs): grid.paste(img, box=(i%cols*w, i//cols*h))
return grid
images = images.reshape((images.shape[0], ) + images.shape[-3:]) #11
images = pipeline.numpy_to_pil(images) #12
image_grid(images, 2, 4) #12
Inspect the complete JAX/Flax listing for loading the checkpoint, configuring precision, and generating images. section 11.3.3
