[Flax] Stateless schedulers, fixes and refactors #1661

skirsten · 2022-12-11T21:54:17Z

Hi, here some fixes and improvements for the Flax schedulers. Let me know what you think!
Sorry that its a huge PR with a single commit 😅

Refactor schedulers to be completely stateless (all the state is in the params)

No more state in the scheduler class
No more Implicit transfers
Extracted common state (common state can also be reused from other schedulers)
The shape and dtypes of the state returned by set_timestamp is now final and wont be changed by step (this reduces the amount of jit misses if jitting the scheduler separately)

Added dtype param to schedulers

Leave it at fp32 though unless you want to lose all details in the image.

Fix copy paste error in `add_noise` function

The add_noise function and thus img2img were not working in DDIM and DPMSolverMultistep
Extracted common logic so it can't happen again

Removed all jax conditionals to fix performance bottleneck

Using jax.lax.cond and jax.lax.switch causes the CPU to have to wait for the pred (even when jitted) causing the GPU pipeline to stall (not enough kernels scheduled). More info here.
If you notice that PNDM and DPMSolverMultistep were slower than DDIM, this was the reason.
Usually this is most noticeable on fast GPU + slow CPU combo or if running a splitkernel (separately jitted scheduler instead of the megakernel as in this repo).
Evaluating all branches instead has no noticeable performance impact.

Fixed small bugs and improvements

Added v_prediction where it was missing
Made DDPM jitable. Though I'm not sure sure if it works correctly.
Fixed DPMSolverMultistep not being able to start in the middle of a schedule. This caused img2img not to work.
Made LMSDiscrete run. Its not jitable and I always get back a black image though.
Probably some other stuff that I forgot about

Validation (outdated)

I messed up so Pytorch is fp16 and Flax is bf16

name	Pytorch	Flax v0.10.2	Flax this PR
`DDIM`
`DPMSolverMultistep`
`PNDM`

HuggingFaceDocBuilderDev · 2022-12-11T21:58:32Z

The documentation is not available anymore as the PR was closed or merged.

pcuenca · 2022-12-13T06:25:44Z

Hi @skirsten, this looks amazing! I see you are tweaking stuff, let me know when you want a review :)

skirsten · 2022-12-13T21:23:26Z

Hi @pcuenca, It should be ready for review now 😅

pcuenca · 2022-12-13T21:35:43Z

Awesome, will do this week!

patrickvonplaten · 2022-12-19T12:05:24Z

src/diffusers/schedulers/scheduling_common_flax.py

@@ -0,0 +1,106 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.


Can we maybe add this to scheduling_utils_flax.py instead ? :-) We usually don't have _common_ files in src/diffusers

Done. I moved it to scheduling_utils_flax.py

patrickvonplaten

This looks super nice! Thanks a lot for working on this @skirsten :-)

@pcuenca it'd be amazing if you could give this a try on a TPU

pcuenca

This is great! Much clearer code (and more efficient, according to the previous comments).

I tested PNDM, DDPM, DPM Solver and LMS Discrete on TPU v3-8 and these were my results:

DPM Solver produces identical results as the previous version.
There are some minor visual differences in PNDM. I haven't found the reason: tried to force dtype to float32 when computing the betas but it didn't make a difference.
LMS does not work in either version.
DDPM produces noise in the new version. It crashed in the previous one.

I think we should merge this as it's so much better. My approach would be:

Deal with DDPM and LMS in a followup PR, as they didn't work anyway in the previous implementation.
If we can easily find a reason for the minor discrepancies in PNDM, let's try to apply it. I already spent a couple hours and couldn't find it, so I wouldn't spend much more time.

What do you think?

pcuenca · 2022-12-19T19:09:02Z

src/diffusers/schedulers/scheduling_pndm_flax.py

+        model_output = jax.lax.select(
+            (state.counter % 4) != 3,
+            model_output,  # remainder 0, 1, 2
+            state.cur_model_output + 1 / 6 * model_output,  # remainder 3
+        )


These changes are all much cleaner this way. Thanks a lot!

pcuenca · 2022-12-19T19:09:32Z

src/diffusers/schedulers/scheduling_pndm_flax.py

+            cur_model_output=jax.lax.select_n(
+                state.counter % 4,
+                state.cur_model_output + 1 / 6 * model_output,  # remainder 0
+                state.cur_model_output + 1 / 3 * model_output,  # remainder 1
+                state.cur_model_output + 1 / 3 * model_output,  # remainder 2
+                jnp.zeros_like(state.cur_model_output),  # remainder 3
+            ),


src/diffusers/schedulers/scheduling_pndm_flax.py

pcuenca · 2022-12-19T19:21:58Z

Maybe @patil-suraj wants to take a quick look too.

Co-authored-by: Pedro Cuenca <[email protected]>

patrickvonplaten · 2022-12-20T00:42:37Z

Cool, let's merge as this is a clear improvement to what we had previously. More than happy to fix scheduler one-by-one in the future.

* [Flax] Stateless schedulers, fixes and refactors * Remove scheduling_common_flax and some renames * Update src/diffusers/schedulers/scheduling_pndm_flax.py Co-authored-by: Pedro Cuenca <[email protected]> Co-authored-by: Pedro Cuenca <[email protected]>

skirsten force-pushed the flax/stateless-schedulers-and-improvements branch 2 times, most recently from 3a64fbf to 0108066 Compare December 12, 2022 00:37

skirsten force-pushed the flax/stateless-schedulers-and-improvements branch from 3f48e8e to c6357ed Compare December 13, 2022 21:13

patrickvonplaten reviewed Dec 19, 2022

View reviewed changes

skirsten force-pushed the flax/stateless-schedulers-and-improvements branch from c6357ed to 9315088 Compare December 19, 2022 12:57

pcuenca approved these changes Dec 19, 2022

View reviewed changes

skirsten added 2 commits December 19, 2022 20:35

[Flax] Stateless schedulers, fixes and refactors

4ee3fce

Remove scheduling_common_flax and some renames

752fb76

skirsten force-pushed the flax/stateless-schedulers-and-improvements branch from 9315088 to 752fb76 Compare December 19, 2022 19:37

Update src/diffusers/schedulers/scheduling_pndm_flax.py

685551f

Co-authored-by: Pedro Cuenca <[email protected]>

patrickvonplaten merged commit f106ab4 into huggingface:main Dec 20, 2022

skirsten mentioned this pull request Dec 21, 2022

Add Flax stable diffusion img2img pipeline #1355

Merged

2 tasks

patil-suraj mentioned this pull request Dec 23, 2022

Add v-prediction to Flax training examples #1777

Closed

This was referenced Jan 31, 2023

latest release breaks Flax pipes for Stable Diffusion when using bfloat16 and/or non-default scheduler #2155

Closed

Make FlaxLMSDiscreteScheduler jittable #2180

Open

		@@ -0,0 +1,106 @@
		# Copyright 2022 The HuggingFace Team. All rights reserved.

[Flax] Stateless schedulers, fixes and refactors #1661

[Flax] Stateless schedulers, fixes and refactors #1661

Uh oh!

Conversation

skirsten commented Dec 11, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Refactor schedulers to be completely stateless (all the state is in the params)

Added dtype param to schedulers

Fix copy paste error in add_noise function

Removed all jax conditionals to fix performance bottleneck

Fixed small bugs and improvements

Validation (outdated)

Uh oh!

HuggingFaceDocBuilderDev commented Dec 11, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pcuenca commented Dec 13, 2022

Uh oh!

skirsten commented Dec 13, 2022

Uh oh!

pcuenca commented Dec 13, 2022

Uh oh!

patrickvonplaten Dec 19, 2022

Choose a reason for hiding this comment

Uh oh!

skirsten Dec 19, 2022

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

pcuenca left a comment

Choose a reason for hiding this comment

Uh oh!

pcuenca Dec 19, 2022

Choose a reason for hiding this comment

Uh oh!

pcuenca Dec 19, 2022

Choose a reason for hiding this comment

Uh oh!

Uh oh!

pcuenca commented Dec 19, 2022

Uh oh!

patrickvonplaten commented Dec 20, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

skirsten commented Dec 11, 2022 •

edited

Loading

Fix copy paste error in `add_noise` function

HuggingFaceDocBuilderDev commented Dec 11, 2022 •

edited

Loading