Need images, audio, video, or 3D structures from a description? A diffusion model turns noise into an output through a multistep denoising process. In JAX, the diffusers library makes pretrained pipelines practical without requiring you to implement the model from scratch. section 11.3.3
A diffusers workflow combines:
- A textual encoder and optionally an image encoder to represent conditioning input. section 11.3.3
- A multistep diffusion process controlled by interchangeable noise schedulers. section 11.3.3
- A pretrained model that can serve as a building block for an end-to-end system. section 11.3.3
Generate images with Stable Diffusion
Load FlaxStableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5") with JAX bfloat16 parameters, then use the pipeline for image generation. The example checks the JAX backend and prepares inputs for accelerator execution. 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
Use the installation command and complete image-generation listing as a practical starting point for a Flax-based diffusion pipeline. section 11.3.3
