From 73d1caf8f28a387f2db5a77a8892edad8ed505a0 Mon Sep 17 00:00:00 2001 From: Logan Date: Fri, 10 May 2024 12:38:10 +1000 Subject: [PATCH 1/3] Add Align Your Steps to available schedulers * Include both SDXL and SD 1.5 variants (https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html) --- modules/sd_schedulers.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/modules/sd_schedulers.py b/modules/sd_schedulers.py index 75eb3ac03..2131eae46 100644 --- a/modules/sd_schedulers.py +++ b/modules/sd_schedulers.py @@ -4,6 +4,7 @@ import torch import k_diffusion +import numpy as np @dataclasses.dataclass class Scheduler: @@ -30,6 +31,35 @@ def sgm_uniform(n, sigma_min, sigma_max, inner_model, device): sigs += [0.0] return torch.FloatTensor(sigs).to(device) +def get_align_your_steps_sigmas(n, device, sigma_id): + # https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html + def loglinear_interp(t_steps, num_steps): + """ + Performs log-linear interpolation of a given array of decreasing numbers. + """ + xs = np.linspace(0, 1, len(t_steps)) + ys = np.log(t_steps[::-1]) + + new_xs = np.linspace(0, 1, num_steps) + new_ys = np.interp(new_xs, xs, ys) + + interped_ys = np.exp(new_ys)[::-1].copy() + return interped_ys + + if sigma_id == "sdxl": + sigmas = [14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.029] + elif sigma_id == "sd15": + sigmas = [14.615, 6.475, 3.861, 2.697, 1.886, 1.396, 0.963, 0.652, 0.399, 0.152, 0.029] + else: + print(f'Align Your Steps sigma identifier "{sigma_id}" not recognized, defaulting to SD 1.5.') + sigmas = [14.615, 6.475, 3.861, 2.697, 1.886, 1.396, 0.963, 0.652, 0.399, 0.152, 0.029] + + if n != len(sigmas): + sigmas = np.append(loglinear_interp(sigmas, n), [0.0]) + else: + sigmas.append(0.0) + + return torch.FloatTensor(sigmas).to(device) schedulers = [ Scheduler('automatic', 'Automatic', None), @@ -38,6 +68,8 @@ schedulers = [ 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('align_your_steps_sdxl', 'Align Your Steps (SDXL)', lambda n, sigma_min, sigma_max, device: get_align_your_steps_sigmas(n, device, "sdxl")), + Scheduler('align_your_steps_sd15', 'Align Your Steps (SD 1.5)', lambda n, sigma_min, sigma_max, device: get_align_your_steps_sigmas(n, device, "sd15")), ] schedulers_map = {**{x.name: x for x in schedulers}, **{x.label: x for x in schedulers}} From d6b4444069d36cf7554eb9932061ecf43e9b1335 Mon Sep 17 00:00:00 2001 From: Logan Date: Fri, 10 May 2024 18:05:45 +1000 Subject: [PATCH 2/3] Use shared.sd_model.is_sdxl to determine base AYS sigmas --- modules/sd_schedulers.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/modules/sd_schedulers.py b/modules/sd_schedulers.py index 2131eae46..0ac1f7a21 100644 --- a/modules/sd_schedulers.py +++ b/modules/sd_schedulers.py @@ -6,6 +6,8 @@ import k_diffusion import numpy as np +from modules import shared + @dataclasses.dataclass class Scheduler: name: str @@ -31,7 +33,7 @@ def sgm_uniform(n, sigma_min, sigma_max, inner_model, device): sigs += [0.0] return torch.FloatTensor(sigs).to(device) -def get_align_your_steps_sigmas(n, device, sigma_id): +def get_align_your_steps_sigmas(n, sigma_min, sigma_max, device): # https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html def loglinear_interp(t_steps, num_steps): """ @@ -46,12 +48,10 @@ def get_align_your_steps_sigmas(n, device, sigma_id): interped_ys = np.exp(new_ys)[::-1].copy() return interped_ys - if sigma_id == "sdxl": + if shared.sd_model.is_sdxl: sigmas = [14.615, 6.315, 3.771, 2.181, 1.342, 0.862, 0.555, 0.380, 0.234, 0.113, 0.029] - elif sigma_id == "sd15": - sigmas = [14.615, 6.475, 3.861, 2.697, 1.886, 1.396, 0.963, 0.652, 0.399, 0.152, 0.029] else: - print(f'Align Your Steps sigma identifier "{sigma_id}" not recognized, defaulting to SD 1.5.') + # Default to SD 1.5 sigmas. sigmas = [14.615, 6.475, 3.861, 2.697, 1.886, 1.396, 0.963, 0.652, 0.399, 0.152, 0.029] if n != len(sigmas): @@ -68,8 +68,7 @@ schedulers = [ 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('align_your_steps_sdxl', 'Align Your Steps (SDXL)', lambda n, sigma_min, sigma_max, device: get_align_your_steps_sigmas(n, device, "sdxl")), - Scheduler('align_your_steps_sd15', 'Align Your Steps (SD 1.5)', lambda n, sigma_min, sigma_max, device: get_align_your_steps_sigmas(n, device, "sd15")), + Scheduler('align_your_steps', 'Align Your Steps', get_align_your_steps_sigmas), ] schedulers_map = {**{x.name: x for x in schedulers}, **{x.label: x for x in schedulers}} From 1d7448281751ea3223c681a82de8219a6fbe1d22 Mon Sep 17 00:00:00 2001 From: Logan Date: Sat, 18 May 2024 09:09:57 +1000 Subject: [PATCH 3/3] Default device for sigma tensor to CPU * Consistent with implementations in k-diffusion. * Makes this compatible with https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/15823 --- modules/sd_schedulers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/sd_schedulers.py b/modules/sd_schedulers.py index 0ac1f7a21..4ddb77850 100644 --- a/modules/sd_schedulers.py +++ b/modules/sd_schedulers.py @@ -33,7 +33,7 @@ def sgm_uniform(n, sigma_min, sigma_max, inner_model, device): sigs += [0.0] return torch.FloatTensor(sigs).to(device) -def get_align_your_steps_sigmas(n, sigma_min, sigma_max, device): +def get_align_your_steps_sigmas(n, sigma_min, sigma_max, device='cpu'): # https://research.nvidia.com/labs/toronto-ai/AlignYourSteps/howto.html def loglinear_interp(t_steps, num_steps): """