Add 2nd order heun scheduler #1336

patrickvonplaten · 2022-11-17T19:20:37Z

This PR does two things:

1. It adds the first second order scheduler to diffusers without changing the logic of the scheduler / model API:
- a. The loop is still defined by iterating over self.scheduler.timesteps
- b. The scheduler step function still returns only the computed model values
1. It fixes the progress bar of PNDM and also makes it compatible with 2nd order schedulers. Now PNDM shows 50 update steps instead of 51.

Note that this design might evolve in the future as discussed in #1308

patrickvonplaten · 2022-11-17T19:26:55Z

src/diffusers/schedulers/scheduling_heun.py

+        self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
+
+        timesteps = torch.from_numpy(timesteps)
+        timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2), timesteps[-1:]])


@patil-suraj @pcuenca @anton-l - this just repeats timesteps for the second order

That's cool! But why do we need to append the final sigma again? We end up with three trailing zeros here, not two.

Also, I think repeat_interleave was not compatible / efficient with mps and maybe onnx; let's just take a note to deal with that later.

Same question as Pedro about the final sigma.

True I should remove it

patrickvonplaten · 2022-11-17T19:27:05Z

src/diffusers/schedulers/scheduling_heun.py

+        sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
+        sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32)
+        sigmas = torch.from_numpy(sigmas).to(device=device)
+        self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])


@patil-suraj @pcuenca @anton-l - this just repeats sigmas for the second order

patrickvonplaten · 2022-11-17T19:27:20Z

src/diffusers/schedulers/scheduling_heun.py

+
+    @property
+    def state_in_first_order(self):
+        return self.dt is None


simple property defining the mode of the scheduler

HuggingFaceDocBuilderDev · 2022-11-17T19:27:32Z

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

keturn · 2022-11-17T21:32:20Z

src/diffusers/schedulers/scheduling_heun.py

+            [`~schedulers.scheduling_utils.SchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is the sample tensor.
+        """
+        step_index = (self.timesteps == timestep).nonzero().item()


Does this need to be adjusted to account for duplicate entries in timesteps? because == will match more than one thing and .item() complains about that.

Then the result is ambiguous, but I guess we can take the first and then add one if not state_in_first_order?

I agree with @keturn, see index_for_timestep in a comment above.

Suggested change

step_index = (self.timesteps == timestep).nonzero().item()

step_index = self.index_for_timestep(timestep)

The hack to get the timestep index seems to become even harder to read and debug here. Should we maybe return the timestep index arguments at some point? (planting a thought for 1.0.0)

src/diffusers/schedulers/scheduling_heun.py

pcuenca

This is a great effort, thanks a lot! I find it much easier to discuss based on some code than on the ether, so it's very helpful.

I think it's a bit finicky to get it right with all the indexing and stuff, but the end result looks really understandable. I could get what was going in a first pass even though I didn't realize there were some details a bit off until I looked more carefully. I think this is acceptable for the "write-once, read-many" approach we want to achieve here.

As an exercise, I'll try to find some time to build the alternative example with the scheduler returning a tuple that I proposed in the other thread, and see how it compares to this.

pcuenca · 2022-11-18T08:45:22Z

src/diffusers/schedulers/scheduling_heun.py

+        self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
+
+        timesteps = torch.from_numpy(timesteps)
+        timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2), timesteps[-1:]])


That's cool! But why do we need to append the final sigma again? We end up with three trailing zeros here, not two.

src/diffusers/schedulers/scheduling_heun.py

pcuenca · 2022-11-18T08:48:41Z

src/diffusers/schedulers/scheduling_heun.py

+            sample (`torch.FloatTensor`): input sample timestep (`int`, optional): current timestep
+        Returns:
+            `torch.FloatTensor`: scaled input sample
+        """


Compute the index here instead of receiving as an argument?

Suggested change

"""

"""

step_index = self.index_for_timestep(timestep)

where index_for_timestep would be something like:

def index_for_timestep(self, timestep): pos = -1 if self.state_in_first_order else 0 return (self.timesteps == timestep).nonzero()[pos].item()

It's better to use a function because it's a bit non-trivial and step requires the index too.

An alternative would be to call scale_model_input from the scheduler, but that's an API breaking change. (Why didn't we do it like that?)

Agree with Pedro, the scale_model_input in other schedulers like LMSDiscreteScheduler don't take step_index as input, so would be nice to follow the same API.

Also @pcuenca

An alternative would be to call scale_model_input from the scheduler, but that's an API breaking change. (Why didn't we do it like that?)

scale_model_input is not called from scheduler because the scaled input needs to be passed to the model, and in the first iteration the model call happens before the scheduler step.

Sorry yeah this was a bug

src/diffusers/schedulers/scheduling_heun.py

pcuenca · 2022-11-18T08:53:23Z

src/diffusers/schedulers/scheduling_heun.py

+        self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
+
+        timesteps = torch.from_numpy(timesteps)
+        timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2), timesteps[-1:]])


Also, I think repeat_interleave was not compatible / efficient with mps and maybe onnx; let's just take a note to deal with that later.

pcuenca · 2022-11-18T08:54:53Z

src/diffusers/schedulers/scheduling_heun.py

+            [`~schedulers.scheduling_utils.SchedulerOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is the sample tensor.
+        """
+        step_index = (self.timesteps == timestep).nonzero().item()


I agree with @keturn, see index_for_timestep in a comment above.

Suggested change

step_index = (self.timesteps == timestep).nonzero().item()

step_index = self.index_for_timestep(timestep)

pcuenca · 2022-11-18T08:55:59Z

src/diffusers/schedulers/scheduling_heun.py

+        # currently only gamma=0 is supported. This usually works best anyways.
+        # We can support gamma in the future but then need to scale the timestep before
+        # passing it to the model which requires a change in API
+        gamma = 0


Undecided if it's better to just remove the gamma var (but leave the comment) or keep it for clarity.

src/diffusers/schedulers/scheduling_heun.py

patil-suraj

This is really nice! Pretty much the same comments as @pcuenca about the step_index.
Also let's try to compare this with alternate approach that @pcuenca proposed and then we could decide the final API. Will try this more and see if I have any other comments.

Also, my main comment here is that the formula for prev_sample seems wrong here.

Instead of prev_sample = model_output + derivative * self.dt

it should be prev_sample = sample + derivative * self.dt as per the paper and k-diffusion.

patil-suraj · 2022-11-18T10:29:18Z

src/diffusers/schedulers/scheduling_heun.py

+            sample (`torch.FloatTensor`): input sample timestep (`int`, optional): current timestep
+        Returns:
+            `torch.FloatTensor`: scaled input sample
+        """


Agree with Pedro, the scale_model_input in other schedulers like LMSDiscreteScheduler don't take step_index as input, so would be nice to follow the same API.

Also @pcuenca

An alternative would be to call scale_model_input from the scheduler, but that's an API breaking change. (Why didn't we do it like that?)

scale_model_input is not called from scheduler because the scaled input needs to be passed to the model, and in the first iteration the model call happens before the scheduler step.

patil-suraj · 2022-11-18T10:35:33Z

src/diffusers/schedulers/scheduling_heun.py

+        self.sigmas = torch.cat([sigmas[:1], sigmas[1:-1].repeat_interleave(2), sigmas[-1:]])
+
+        timesteps = torch.from_numpy(timesteps)
+        timesteps = torch.cat([timesteps[:1], timesteps[1:].repeat_interleave(2), timesteps[-1:]])


Same question as Pedro about the final sigma.

src/diffusers/schedulers/scheduling_heun.py

…add_heun

pcuenca

Submitted changes to make it work (I think).

src/diffusers/schedulers/scheduling_heun.py

patrickvonplaten · 2022-11-21T17:54:17Z

Thanks a lot for the corrections @pcuenca @patil-suraj @keturn.

The PR as is now is functional and gives 1-to-1 the same results as k-diffusion:

#!/usr/bin/env python3
from diffusers import DiffusionPipeline, StableDiffusionPipeline, HeunDiscreteScheduler
import torch

seed = 33

pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", custom_pipeline="sd_text2img_k_diffusion")
pipe = pipe.to("cuda")

prompt = "an astronaut riding a horse on mars"
pipe.set_sampler("sample_heun")
generator = torch.Generator(device="cuda").manual_seed(seed)
image = pipe(prompt, generator=generator, num_inference_steps=20).images[0]

image.save("./astronaut_heun_k_diffusion_comp.png")

pipe = StableDiffusionPipeline(**pipe.components)
pipe = pipe.to("cuda")
pipe.scheduler = HeunDiscreteScheduler.from_config(pipe.scheduler.config)
generator = torch.Generator(device="cuda").manual_seed(seed)
image = pipe(prompt, generator=generator, num_inference_steps=20).images[0]

image.save("./astronaut_heun_comp.png")

K Diffusion

This PR (Diffusers)

…nto add_heun

…add_heun

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

…sion.py

…nto add_heun

* Add heun * Finish first version of heun * remove bogus * finish * finish * improve * up * up * fix more * change progress bar * Update src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py * finish * up * up * up

patrickvonplaten added 5 commits November 17, 2022 18:20

Add heun

9a4d722

Finish first version of heun

7564993

remove bogus

5bb32f8

finish

c53a983

finish

22129a3

patrickvonplaten commented Nov 17, 2022

View reviewed changes

patrickvonplaten requested review from anton-l, patil-suraj and pcuenca November 17, 2022 19:27

patrickvonplaten mentioned this pull request Nov 17, 2022

[Design discussion] Diffusers schedulers/samplers & 2nd Order Schedulers #1308

Closed

keturn reviewed Nov 17, 2022

View reviewed changes

src/diffusers/schedulers/scheduling_heun.py Show resolved Hide resolved

pcuenca reviewed Nov 18, 2022

View reviewed changes

patil-suraj reviewed Nov 18, 2022

View reviewed changes

Merge branch 'main' of https://github.com/huggingface/diffusers into …

2128772

…add_heun

pcuenca mentioned this pull request Nov 21, 2022

Add experimental Heun scheduler #1356

Closed

pcuenca reviewed Nov 21, 2022

View reviewed changes

src/diffusers/schedulers/scheduling_heun.py Show resolved Hide resolved

src/diffusers/schedulers/scheduling_heun.py Outdated Show resolved Hide resolved

patrickvonplaten added 4 commits November 21, 2022 17:28

improve

e5e62ff

up

d9d8e84

up

2a24c62

Merge branch 'main' into add_heun

91acb29

patrickvonplaten added 4 commits November 21, 2022 17:54

fix more

719b749

Merge branch 'add_heun' of https://github.com/huggingface/diffusers i…

f50ee61

…nto add_heun

change progress bar

ed2c233

Merge branch 'main' of https://github.com/huggingface/diffusers into …

2f8a2d8

…add_heun

patrickvonplaten commented Nov 28, 2022

View reviewed changes

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

patrickvonplaten added 6 commits November 28, 2022 21:24

Update src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffu…

69eea0c

…sion.py

finish

f00cbc9

Merge branch 'add_heun' of https://github.com/huggingface/diffusers i…

ae39bdb

…nto add_heun

up

b30c326

up

e684d59

up

fa3ed43

patrickvonplaten merged commit 4c54519 into main Nov 28, 2022

patrickvonplaten deleted the add_heun branch November 28, 2022 21:56

	step_index = (self.timesteps == timestep).nonzero().item()
	step_index = self.index_for_timestep(timestep)

Add 2nd order heun scheduler #1336

Add 2nd order heun scheduler #1336

Uh oh!

Conversation

patrickvonplaten commented Nov 17, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

pcuenca left a comment • 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!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

patil-suraj left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pcuenca left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

patrickvonplaten commented Nov 21, 2022

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

patrickvonplaten commented Nov 17, 2022 •

edited

Loading

HuggingFaceDocBuilderDev commented Nov 17, 2022 •

edited

Loading

pcuenca left a comment •

edited

Loading