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 afteradd 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 readcontributing wiki page
• I have performed a self-review of my own code
• My code follows thestyle guidelines
• My code passestests