Feature: Beta scheduler

[v0xie]

Description

• This PR implements the scheduler described in "Beta Sampling is All You Need: Efficient Image Generation Strategy for Diffusion Models using Stepwise Spectral Analysis" (2024, Lee et. al). The basic conclusion of the paper is spending more time at the beginning and end of denoising improves image quality.

• Output images are usually worse than other schedulers at low step count (<= 10), but improved at higher step counts ( >= 20 ).

• Beta scheduler with alpha = beta = 1.0 is supposed to be equivalent to Uniform scheduler, but in practice the Uniform scheduler produces different sigmas / results.

Changes

• Add new scheduler "Beta"
• Adds two options in Sampler Parameters: "Beta scheduler - alpha" and "Beta scheduler - beta"
• Add "Beta schedule alpha" and "Beta schedule beta" to XYZ plot options
• Writes alpha/beta to infotext when scheduler is "Beta"

Additional links:

The authors' paper describing the method: https://arxiv.org/abs/2407.12173

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

Screenshots/videos:

SD XL

[ibrainventures]

If i load a image by PNG Info Tab, there are correctly the Beta schedule sub-params included:

Beta schedule alpha: 0.45, Beta schedule beta: 0.8, Version: v1.10.1

But those params are not respected by System after pressing "Send to txt2img".
There are still the "system saved params" (which are as example 0.5 / 0.7 saved) active.
Which gives a different output..

Is it possible to include / resepect those image-specific params?

Thanks for this scheduler, which is really a great scheduler type.

[heftig]

Beta scheduler with alpha = beta = 1.0 is supposed to be equivalent to Uniform scheduler, but in practice the Uniform scheduler produces different sigmas / results.

Are you sure your implementation is right? It seems to be producing beta-distributed sigmas instead of beta-distributed timesteps as explained in the paper.

I think we want something more like this (adapted from the 'normal' scheduler):

def beta_scheduler(n, sigma_min, sigma_max, inner_model, device, sgm=False):
    start = inner_model.sigma_to_t(torch.tensor(sigma_max))
    end = inner_model.sigma_to_t(torch.tensor(sigma_min))

    # From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)
    alpha = shared.opts.beta_dist_alpha
    beta = shared.opts.beta_dist_beta
    if sgm:
        timesteps = np.linspace(1, 0, n + 1)[:-1]
    else:
        timesteps = np.linspace(1, 0, n)
    timesteps = [stats.beta.ppf(x, alpha, beta) for x in timesteps]

    sigs = []
    for x in range(len(timesteps)):
        ts = end + (timesteps[x] * (start - end))
        sigs.append(inner_model.t_to_sigma(ts))
    sigs += [0.0]
    return torch.FloatTensor(sigs).to(device)

With this, using alpha = beta = 1.0 makes it almost identical to the uniform scheduler.