Add KL Optimal scheduler

[drhead]

Description

• Adds a new scheduler called "KL Optimal", based on what is noted as an optimal scheduler in Sabour, et. al, "Align Your Steps", theorem 3.1.
• (technobabble) "The optimal schedule $t^*$ minimizing the KL-divergence between $p(\mathrm{x}, t_{\text{min}})$ and the distribution of $\bar{\mathrm{x}}_{t_{\text{min}}}$"
• Seems to substantially improve sample quality over the uniform and EDM/Karras scheduler, and also converges much faster, especially on ZSNR models.
• Might need CFG skipped on the first sampling step on ZSNR models on some prompts (and doing this often improves sample quality regardless). Recommend waiting until after add code for skipping CFG on early steps #15607 is merged to merge this.

Screenshots/videos:

Note: These grids have CFG skipped on the first step unless otherwise specified.
Test grid for 50 steps: Comparing different noise schedulers, and also including sigma max of 4500 (approximately the default for ZSNR) and 160 in the comparison (an alternative that I find works well for other schedulers, spending less time at extremely high sigmas):

All samplers perform adequately on this one. Most notably, KL Optimal functions better at the higher sigma max value, where Karras is somewhat washed out/smoother.

Test grid for 25 steps:

At 25 steps, Karras is an absolute mess at the higher default sigma max value. Others all still look acceptable.

Test grid for 10 steps:

Complete failure case for Uniform and Karras. KL Optimal doesn't look fantastic, but frankly it's better looking than a 10-step sample from a non-distilled model has any right to be.

Failure case when not skipping CFG (50 steps). Probably ZSNR specific:

Not all prompts will do this, but this happens to be one of the ones that does. The black "doors" surrounding the images in the KL Optimal column will render as a black border that encroaches on the whole image at lower step counts. Higher step counts might stabilize this depending on prompt. Skipping CFG on the first timestep is enough to reliably fix issues such as this.

Checklist:

• I have read contributing wiki page
• I have performed a self-review of my own code
• My code follows the style guidelines
• My code passes tests

[zcatharisis]

Thanks for implementing this to A1111. Could I ask for your preferred settings for SD 1.5 non-ZSNR models, especially for Schedule Sigma Max and the Negative Guidance minimum sigma? Those are the only variables I haven't been able to dial in in my testing.

[drhead]

Thanks for implementing this to A1111. Could I ask for your preferred settings for SD 1.5 non-ZSNR models, especially for Schedule Sigma Max and the Negative Guidance minimum sigma? Those are the only variables I haven't been able to dial in in my testing.

For regular models, leave schedule sigma max at default (0/automatic), the only reason I mess with it on ZSNR models is because it should be infinity there but that doesn't/can't work properly. For negative guidance minimum sigma, I usually find that 0.75 works fine (set to skip all steps instead of every other). Turn it down if images get too smoothed over.

[AUTOMATIC1111]

Is that it? Isn't just just a theoretical example unrelated to SD?

[drhead]

Is that it? Isn't just just a theoretical example unrelated to SD?

It's an optimal schedule derived for DDIM so it is as related to SD as anything that is generally applicable to diffusion models is. I and several other people have tested the schedule fairly thoroughly since this was filed, and we've found it to work very well in practice, and it has similar effects to the advertised effects of Align Your Steps. The "theoretical" part, as I understand it, is more in that this schedule as it is would be assuming that the diffusion model behaves somewhat perfectly (which it does not) and that training a schedule as the paper recommends is aligning the schedule to the imperfect diffusion model (which would not be unique to this schedule).