>>> init_image = download_image(img_url).resize((512, 512))
>>> mask_image = download_image(mask_url).resize((512, 512))

>>> pipeline, params = FlaxStableDiffusionInpaintPipeline.from_pretrained(
...     "xvjiarui/stable-diffusion-2-inpainting"
... )

>>> prompt = "Face of a yellow cat, high resolution, sitting on a park bench"
>>> prng_seed = jax.random.PRNGKey(0)
>>> num_inference_steps = 50

>>> num_samples = jax.device_count()
>>> prompt =