[Scheduler design] The pragmatic approach #719

anton-l · 2022-10-04T13:28:24Z

This schedulers API redesign addresses concerns raised in #336 and starts to make the schedulers interchangeable without making scheduler class-dependent customizations to pipelines.

Now every scheduler contains an init_noise_sigma parameter to scale the normal distribution of the initial noise. While it is just 1.0 for DDPM, DDIM and PNDM, it is customized for the VE and K-LMS schedulers.
Example usage:

sample = torch.randn(*shape, generator=generator) * self.scheduler.init_noise_sigma

Every scheduler needs to implement the scale_model_input(sample, timestep) (even if it just returns the sample) that scales the denoising model's input based on the current timestep. The method should be called before every model() call.
Example usage:

sample = self.scheduler.scale_model_input(sample, t)
output = model(sample, t)

Note: the decision to not make it a base class method is intentional, as suggested by @patrickvonplaten

Closes #336

HuggingFaceDocBuilderDev · 2022-10-04T13:31:28Z

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

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py

patrickvonplaten

Awesome, love this API design - think it's slim and should solve 99% of our problems! Actually maybe no need after all to have a SchedulerType.CONTINUOUS then after all in a first step 😍

anton-l · 2022-10-04T17:22:35Z

The slow tests pass ✔️

anton-l · 2022-10-04T17:24:44Z

src/diffusers/schedulers/scheduling_lms_discrete.py

        timesteps = timesteps.to(original_samples.device)
+        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]


Really dislike what we have to do here, but unfortunately there's no good vectorized alternative to search for multiple indices and keep the order the same.

don't think it's that bad honestly

patrickvonplaten · 2022-10-04T22:07:39Z

Really cool - like this new design a lot!

Think we can merge this tomorrow 😍

Some final TODOs / suggestions:

add the design now to all other schedulers and their scripts (or leave as clear TODOs)
think we could add a test that forces every new scheduler to have:
- a init_noise_sigma variable
- a scale_model_input function
- a step function
- => this could ensure that we follow this design in the future. If some schedulers don't follow this design, let's exempt them for this test maybe for now with a big TODO to fix it
Give scale_model_input nice docstrings and make sure it's displayed well in the docs. Also let's maybe add a short comment to DDIM, PNDM and DDPM stating that they don't need model scaling
Add a big ⚠️ ⚠️ to this PR that it's backwards breaking and let's maybe try to make it easy for users to fix their code by:
- If the first timestep that is passed to LMS is an int and is 0 => then it's very likely to be wrong (let's throw a warning here)
- If LMS step function is used before having called scale_model_input it's most likely wrong (let's throw an maybe even an error here)

anton-l · 2022-10-05T12:07:53Z

src/diffusers/schedulers/scheduling_lms_discrete.py

+        if not self.is_scale_input_called:
+            warnings.warn(
+                "The `scale_model_input` function should be called before `step` to ensure correct denoising. "
+                "See `StableDiffusionPipeline` for a usage example."
+            )


This will pop up in existing community pipelines but won't break them like an exception would. The legacy pipelines can continue using the manual scaling code 👍

patrickvonplaten · 2022-10-05T12:21:50Z

tests/test_scheduler.py

@@ -226,6 +226,27 @@ def recursive_check(tuple_object, dict_object):

            recursive_check(outputs_tuple, outputs_dict)

+    def test_scheduler_public_api(self):


patrickvonplaten

Very nice! Happy to merge and help you update existing notebooks / docs / blog post now:

Update PR description that states exactly what people have to change if they have written their own custom loop. E.g. If you have been using the K-LMS scheduler, please make sure to do the following:
If you have been using other schedulers, no need to change anything, but we recomend for generality to always make use of init_sigma and scale_model_input
All blog posts
All notebooks
All training examples

* init * improve add_noise * [debug start] run slow test * [debug end] * quick revert * Add docstrings and warnings + API tests * Make the warning less spammy

init

159e15c

anton-l commented Oct 4, 2022

View reviewed changes

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py Outdated Show resolved Hide resolved

patrickvonplaten reviewed Oct 4, 2022

View reviewed changes

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py Outdated Show resolved Hide resolved

anton-l commented Oct 4, 2022

View reviewed changes

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py Outdated Show resolved Hide resolved

patrickvonplaten reviewed Oct 4, 2022

View reviewed changes

anton-l added 2 commits October 4, 2022 17:30

improve add_noise

74ae717

[debug start] run slow test

fa9667f

[debug end]

c06af2b

anton-l commented Oct 4, 2022

View reviewed changes

quick revert

9325ca4

patrickvonplaten mentioned this pull request Oct 4, 2022

[SchedulerDesign] Alternative scheduler design #711

Closed

anton-l mentioned this pull request Oct 5, 2022

[Schedulers Refactoring] Phase 1: timesteps and scaling #637

Closed

anton-l added 3 commits October 5, 2022 14:00

Add docstrings and warnings + API tests

46ceb10

Merge main

f004100

Make the warning less spammy

7a4abb0

anton-l commented Oct 5, 2022

View reviewed changes

patrickvonplaten reviewed Oct 5, 2022

View reviewed changes

patrickvonplaten approved these changes Oct 5, 2022

View reviewed changes

anton-l merged commit 6b09f37 into main Oct 5, 2022

patrickvonplaten deleted the scheduler-refactor-pragmatic branch October 5, 2022 13:13

AbdullahAlfaraj mentioned this pull request Oct 11, 2022

New Scheduler: add Euler Ancestral Scheduler to StableDiffusionPipeline #636

Closed

3 tasks

yuxu915 mentioned this pull request Oct 17, 2023

resume_from_checkpoint training fail in train_dreambooth_lora_sdxl.py #5412

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Scheduler design] The pragmatic approach #719

[Scheduler design] The pragmatic approach #719

Uh oh!

anton-l commented Oct 4, 2022 •

edited

Loading

Uh oh!

HuggingFaceDocBuilderDev commented Oct 4, 2022 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

patrickvonplaten left a comment

Uh oh!

anton-l commented Oct 4, 2022

Uh oh!

anton-l Oct 4, 2022 •

edited

Loading

Uh oh!

patrickvonplaten Oct 4, 2022

Uh oh!

patrickvonplaten commented Oct 4, 2022

Uh oh!

anton-l Oct 5, 2022

Uh oh!

patrickvonplaten Oct 5, 2022

Uh oh!

patrickvonplaten left a comment

Uh oh!

Uh oh!

		timesteps = timesteps.to(original_samples.device)
		step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]

		@@ -226,6 +226,27 @@ def recursive_check(tuple_object, dict_object):

		recursive_check(outputs_tuple, outputs_dict)

		def test_scheduler_public_api(self):

[Scheduler design] The pragmatic approach #719

[Scheduler design] The pragmatic approach #719

Uh oh!

Conversation

anton-l commented Oct 4, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

anton-l commented Oct 4, 2022

Uh oh!

anton-l Oct 4, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten Oct 4, 2022

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten commented Oct 4, 2022

Uh oh!

anton-l Oct 5, 2022

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten Oct 5, 2022

Choose a reason for hiding this comment

Uh oh!

patrickvonplaten left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

anton-l commented Oct 4, 2022 •

edited

Loading

HuggingFaceDocBuilderDev commented Oct 4, 2022 •

edited

Loading

anton-l Oct 4, 2022 •

edited

Loading