 shard(processed_image)

>>> output = pipe(
...     prompt_ids=prompt_ids,
...     image=processed_image,
...     params=p_params,
...     prng_seed=rng,
...     num_inference_steps=50,
...     neg_prompt_ids=negative_prompt_ids,
...     jit=True,
... ).images

>>> output_images = pipe.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:])))
>>> output_images = make_image_grid(output_images, num_samples // 4, 4)
>>> output_i