diff --git a/modules/sd_schedulers.py b/modules/sd_schedulers.py index 4ddb778501a..0c09af8d0b5 100644 --- a/modules/sd_schedulers.py +++ b/modules/sd_schedulers.py @@ -61,6 +61,14 @@ def loglinear_interp(t_steps, num_steps): return torch.FloatTensor(sigmas).to(device) +def kl_optimal(n, sigma_min, sigma_max, device): + alpha_min = torch.arctan(torch.tensor(sigma_min, device=device)) + alpha_max = torch.arctan(torch.tensor(sigma_max, device=device)) + step_indices = torch.arange(n + 1, device=device) + sigmas = torch.tan(step_indices / n * alpha_min + (1.0 - step_indices / n) * alpha_max) + return sigmas + + schedulers = [ Scheduler('automatic', 'Automatic', None), Scheduler('uniform', 'Uniform', uniform, need_inner_model=True), @@ -68,6 +76,7 @@ def loglinear_interp(t_steps, num_steps): Scheduler('exponential', 'Exponential', k_diffusion.sampling.get_sigmas_exponential), Scheduler('polyexponential', 'Polyexponential', k_diffusion.sampling.get_sigmas_polyexponential, default_rho=1.0), Scheduler('sgm_uniform', 'SGM Uniform', sgm_uniform, need_inner_model=True, aliases=["SGMUniform"]), + Scheduler('kl_optimal', 'KL Optimal', kl_optimal), Scheduler('align_your_steps', 'Align Your Steps', get_align_your_steps_sigmas), ]