>>> from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel


>>> def image_grid(imgs, rows, cols):
...     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


>>> def create_key(seed=0):
...     return jax.random.PRNGKey(seed)


>>> rng = create_key(0)

>>> # get canny i