Repository: huggingface/diffusers
Pull request: Flax: avoid recompilation when params change
Fixes #1015 

Problem statement:
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.

Hints from issue 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) 🎉
> 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'd put it on a flax branch :-) 
Hey @camenduru, 

Thanks for opening the issue on the Hub. I answered there: https://huggingface.co/nitrosocke/mo-di-diffusion/discussions/8#6362a3672c0bb59ecc191a07 . Please ping me if you need my help going forward :-) 
If ok for the author of the repo, I'd open a branch directly on main actually. Flax and PyTorch weights can live in the same repo branch side-by-side.
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/diffusers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Commit message: Merge remote-tracking branch 'origin/main' into flax-no-recompile

Patch to review:
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
index 23148dcfe2e3..335d46ef09ea 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
@@ -187,7 +187,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,13 +259,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:
-            # run with python for loop
-            for i in range(num_inference_steps):
-                latents, scheduler_state = loop_body(i, (latents, scheduler_state))
-        else:
-            latents, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state))
+        latents, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state))
 
         # scale and decode the image latents with vae
         latents = 1 / 0.18215 * latents
@@ -283,11 +276,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 +326,12 @@ 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):
+            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 +343,6 @@ def __call__(
                 width,
                 guidance_scale,
                 latents,
-                debug,
                 neg_prompt_ids,
             )
         else:
@@ -358,7 +355,6 @@ def __call__(
                 width,
                 guidance_scale,
                 latents,
-                debug,
                 neg_prompt_ids,
             )
 
@@ -389,7 +385,11 @@ def __call__(
 
 
 # TODO: maybe use a config dict instead of so many static argnums
-@partial(jax.pmap, static_broadcasted_argnums=(0, 4, 5, 6, 7, 9))
+@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 +400,6 @@ def _p_generate(
     width,
     guidance_scale,
     latents,
-    debug,
     neg_prompt_ids,
 ):
     return pipe._generate(
@@ -412,7 +411,6 @@ def _p_generate(
         width,
         guidance_scale,
         latents,
-        debug,
         neg_prompt_ids,
     )
 


Code review comments:
diff_hunk: @@ -260,13 +259,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:
-            # run with python for loop
-            for i in range(num_inference_steps):
-                latents, scheduler_state = loop_body(i, (latents, scheduler_state))
path: src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py
text: @user1:
Think it's nice to have this option for easy debugging.

@author:
The problem is that `debug` has to be passed from everywhere and it adds noise to the very busy `partial` here: https://github.com/huggingface/diffusers/blob/7db319bdae4023bbd6c7a7ca3d16abd873dcdbd0/src/diffusers/pipelines/stable_diffusion/pipeline_flax_stable_diffusion.py#LL388-L391

How do you feel about using a `DEBUG` constant instead, that we only use when we are really debugging?

@user1:
I see, yeah setting DEBUG as a constant sounds good too!