-
Notifications
You must be signed in to change notification settings - Fork 6.5k
Add experimental Heun scheduler #1356
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
Outdated
Show resolved
Hide resolved
| i = 0 | ||
| t = self.scheduler.timesteps[0] | ||
| while t > 0: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We need to keep i and t, which is very ugly. i is only required for the callback.
| if not return_dict: | ||
| return (prev_sample, self.timesteps[step_index]) | ||
|
|
||
| return SchedulerOutput(prev_sample=prev_sample, timestep=self.timesteps[step_index]) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Similar implementation as in #1336 except that we return the next timestep.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. |
|
|
||
| # compute the previous noisy sample x_t -> x_t-1 | ||
| latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample | ||
| latents, t = self.scheduler.step(noise_pred, t, latents, return_dict=False, **extra_step_kwargs) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hmm this makes the step function a bit more black-box - not suuuper happy about it at this point of the library, but maybe at some point we have to let the scheduler adapt t anyways. However would still be happier with exposing the whole timestamp array @patil-suraj @anton-l
Without duplicating the sigmas and timesteps (reference: #1336)
➕
No duplication of sigmas and timesteps.
Does not require ugly function
index_for_timestep(or confusing oneliner repeated twice in the code)➖
Requires modification of all pipelines.
The pipeline loop is uglier.
All in all I don't think we gain any more clarity than in #1336, so I recommend we adopt #1336 unless it's important that the timesteps and sigmas are not duplicated.
This implementation does work. I'll push a fix for #1336 so it works too.