Need to generate images without building a diffusion system from scratch? diffusers provides pretrained pipelines for images, audio, video, molecule structures, and related outputs. It also supports training custom diffusion models. section 11.3.3
Its main building blocks are:
- Ready-to-run pipelines for text-to-image, variation, inpainting, upscaling, 3D generation, text-to-audio, and text-to-video tasks. section 11.3.3
- Interchangeable noise schedulers that trade off diffusion speed and output quality. section 11.3.3
- Pretrained models that can be combined with schedulers to construct end-to-end diffusion systems. section 11.3.3
pip install --upgrade diffusers[flax]
Install the library with Flax support using pip before loading a JAX-compatible pipeline. section 11.3.3
Generate images with Stable Diffusion
A JAX and Flax workflow loads FlaxStableDiffusionPipeline.from_pretrained with runwayml/stable-diffusion-v1-5, selects the bf16 revision, and uses jnp.bfloat16 parameters. The example also checks the JAX backend, which reports TPU in the demonstrated run. 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
Browse the complete image-generation listing, including imports, pretrained pipeline loading, revision selection, and parameter configuration. section 11.3.3
