Repository: huggingface/diffusers

Issue:
JAX / Flax: avoid recompilation on params change: guidance scale, size, number of steps
This is required to improve the demo(s), we currently have the sliders disabled.

Some users are also requesting it in discord and the forums. Our demo notebook is very fast on TPU hardware but users can't quickly explore different parameters.


Discussion:
Please make a script that can convert original stable diffusion to flax 🚀
and please make an auto-converter if someone publishes it like this  https://huggingface.co/nitrosocke/modern-disney-diffusion it will automatically convert to all different models 🤖🧨
Hey @camenduru it should be as easy as doing: 

```python
from diffusers import FlaxStableDiffusionPipeline

pipeline, params = FlaxStableDiffusionPipeline.from_pretrained("<path-to-stable-diffusion-checkpoint>", from_pt=True)
pipeline.save_pretrained("<path/to/hub/repo/with/flax/branch>", params=params)
```
We should add better docs for this though and an automated converter tool sounds like a great idea indeed! 
@patrickvonplaten ✨ thanks I converted https://huggingface.co/nitrosocke/mo-di-diffusion to flax but how to push to flax branch I opened PR but [PR](https://huggingface.co/nitrosocke/mo-di-diffusion/discussions/8) going to main branch 🆘
and I converted like this 
```
pipeline, params = FlaxStableDiffusionPipeline.from_pretrained("nitrosocke/mo-di-diffusion", from_pt=True)
pipeline.save_pretrained("/home/camenduru/mo-di-flax", params=params)
```
should I put this to `bf16` branch or `flax` branch ?
I also converted waifu-diffusion ✨ haru will add [here](https://huggingface.co/hakurei/waifu-diffusion) 🎉

Resolving patch:
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
@@ -42,6 +42,9 @@
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
+# Set to True to use python for loop instead of jax.fori_loop for easier debugging
+DEBUG = False
+
 
 class FlaxStableDiffusionPipeline(FlaxDiffusionPipeline):
     r"""
@@ -187,7 +190,6 @@ def _generate(
         width: Optional[int] = None,
         guidance_scale: float = 7.5,
         latents: Optional[jnp.array] = None,
-        debug: bool = False,
         neg_prompt_ids: jnp.array = None,
     ):
         # 0. Default height and width to unet
@@ -260,8 +262,7 @@ def loop_body(step, args):
 
         # scale the initial noise by the standard deviation required by the scheduler
         latents = latents * self.scheduler.init_noise_sigma
-
-        if debug:
+        if DEBUG:
             # run with python for loop
             for i in range(num_inference_steps):
                 latents, scheduler_state = loop_body(i, (latents, scheduler_state))
@@ -283,11 +284,10 @@ def __call__(
         num_inference_steps: int = 50,
         height: Optional[int] = None,
         width: Optional[int] = None,
-        guidance_scale: float = 7.5,
+        guidance_scale: Union[float, jnp.array] = 7.5,
         latents: jnp.array = None,
         return_dict: bool = True,
         jit: bool = False,
-        debug: bool = False,
         neg_prompt_ids: jnp.array = None,
     ):
         r"""
@@ -334,6 +334,14 @@ def __call__(
         height = height or self.unet.config.sample_size * self.vae_scale_factor
         width = width or self.unet.config.sample_size * self.vae_scale_factor
 
+        if isinstance(guidance_scale, float):
+            # Convert to a tensor so each device gets a copy. Follow the prompt_ids for
+            # shape information, as they may be sharded (when `jit` is `True`), or not.
+            guidance_scale = jnp.array([guidance_scale] * prompt_ids.shape[0])
+            if len(prompt_ids.shape) > 2:
+                # Assume sharded
+                guidance_scale = guidance_scale.reshape(prompt_ids.shape[:2])
+
         if jit:
             images = _p_generate(
                 self,
@@ -345,7 +353,6 @@ def __call__(
                 width,
                 guidance_scale,
                 latents,
-                debug,
                 neg_prompt_ids,
             )
         else:
@@ -358,7 +365,6 @@ def __call__(
                 width,
                 guidance_scale,
                 latents,
-                debug,
                 neg_prompt_ids,
             )
 
@@ -388,8 +394,13 @@ def __call__(
         return FlaxStableDiffusionPipelineOutput(images=images, nsfw_content_detected=has_nsfw_concept)
 
 
-# TODO: maybe use a config dict instead of so many static argnums
-@partial(jax.pmap, static_broadcasted_argnums=(0, 4, 5, 6, 7, 9))
+# Static argnums are pipe, num_inference_steps, height, width. A change would trigger recompilation.
+# Non-static args are (sharded) input tensors mapped over their first dimension (hence, `0`).
+@partial(
+    jax.pmap,
+    in_axes=(None, 0, 0, 0, None, None, None, 0, 0, 0),
+    static_broadcasted_argnums=(0, 4, 5, 6),
+)
 def _p_generate(
     pipe,
     prompt_ids,
@@ -400,7 +411,6 @@ def _p_generate(
     width,
     guidance_scale,
     latents,
-    debug,
     neg_prompt_ids,
 ):
     return pipe._generate(
@@ -412,7 +422,6 @@ def _p_generate(
         width,
         guidance_scale,
         latents,
-        debug,
         neg_prompt_ids,
     )
 


Test patch: