Partitions a pytree of arrays across local devices so parallel code can feed each device its slice of the input. section 11.3.3
It is used with replicated model parameters and tokenized prompts in a Flax Stable Diffusion pipeline before pmap-driven image generation. section 11.3.3
Parallel image generation input prep
A pipeline loads pretrained weights, replicates parameters with replicate(), shreds prompt arrays with shard(), then generates several images in parallel on local devices. section 11.3.3 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
Shows shard() in context beside replicate() and a Flax Stable Diffusion pipeline setup. section 11.3.3
