Create README.md

README.md

@@ -0,0 +1,40 @@
+```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(...)
+```