  return grid


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


>>> rng = create_key(0)

>>> # get canny image
>>> canny_image = load_image(
...     "https://huggingface.co/datasets/YiYiXu/test-doc-assets/resolve/main/blog_post_cell_10_output_0.jpeg"
... )

>>> prompts = "best quality, extremely detailed"
>>> negative_prompts = "monochrome, lowres, bad anatomy, worst quality, low quality"

>>> # load control net and stable