When parallel execution needs the same model state on every device, replicate() copies Flax parameters or other arrays across devices. In a Stable Diffusion workflow, replicate model parameters and preprocessed prompts before generating images in parallel; use shard() instead when inputs must be divided among local devices. section 11.3.3
Replicating a diffusion pipeline
Load FlaxStableDiffusionPipeline and its external parameters, prepare prompts, then call replicate(params) and replicate(preprocessed_prompts). The resulting device-replicated values support the pipeline's pmap()-based generation functions. 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 listing showing replicate() alongside shard(), pipeline loading, prompt preparation, and parallel image generation. section 11.3.3
