readme: ```python
  #!/usr/bin/env python3
  from diffusers import FlaxStableDiffusionPipeline
  from jax import pmap
  import numpy as np
  import jax
  from flax.jax_utils import replicate
  from flax.training.common_utils import shard
  
  
  prng_seed = jax.random.PRNGKey(0)
  num_inference_steps = 50
  
  pipeline, params = FlaxStableDiffusionPipeline.from_pretrained("fusing/stable-diffusion-flax-new", use_auth_token=True)
  del params["safety_checker"]
  
  # pmap
  p_sample = pmap(pipeline.__call__, static_broadcasted_argnums=(3,))
  
  # prep prompts
  prompt = "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of field, close up, split lighting, cinematic"
  num_samples = jax.device_count()
  prompt = num_samples * [prompt]
  prompt_ids = pipeline.prepare_inputs(prompt)
  
  # replicate
  params = replicate(params)
  prng_seed = jax.random.split(prng_seed, 8)
  prompt_ids = shard(prompt_ids)
  
  # run
  images = p_sample(prompt_ids, params, prng_seed, num_inference_steps).images
  
  # get pil images
  images_pil = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))
  
  import ipdb; ipdb.set_trace()
  print("Images should be good")
  # images_pil[0].save(...)
  ```