Skip to content

Add experimental Heun scheduler - #1356

Closed
pcuenca wants to merge 2 commits into
mainfrom
heun-no-dups
Closed

pcuenca wants to merge 2 commits into
mainfrom
heun-no-dups

Conversation

@pcuenca

@pcuenca pcuenca commented Nov 21, 2022 •

Copy link
Copy Markdown
Member

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.

Comment thread src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py Outdated
Comment on lines +498 to +500
i = 0
t = self.scheduler.timesteps[0]
while t > 0:

Copy link
Copy Markdown
Member Author

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.

Comment on lines +196 to +199
if not return_dict:
return (prev_sample, self.timesteps[step_index])

return SchedulerOutput(prev_sample=prev_sample, timestep=self.timesteps[step_index])

Copy link
Copy Markdown
Member Author

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.

@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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)

Copy link
Copy Markdown
Contributor

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants