From a5f66b5003527508c2e2c49c79360ab20071fec2 Mon Sep 17 00:00:00 2001 From: v0xie <28695009+v0xie@users.noreply.github.com> Date: Thu, 18 Jul 2024 15:53:54 -0700 Subject: [PATCH 1/8] feature: beta scheduler --- modules/sd_schedulers.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/modules/sd_schedulers.py b/modules/sd_schedulers.py index 84b0abb6a..19d294c1d 100644 --- a/modules/sd_schedulers.py +++ b/modules/sd_schedulers.py @@ -2,6 +2,7 @@ import dataclasses import torch import k_diffusion import numpy as np +from scipy import stats from modules import shared @@ -115,6 +116,17 @@ def ddim_scheduler(n, sigma_min, sigma_max, inner_model, device): return torch.FloatTensor(sigs).to(device) +def beta_scheduler(n, sigma_min, sigma_max, inner_model, device): + # From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024) """ + alpha = 0.6 + beta = 0.6 + timesteps = 1 - np.linspace(0, 1, n) + timesteps = [stats.beta.ppf(x, alpha, beta) for x in timesteps] + sigmas = [sigma_min + ((x)*(sigma_max-sigma_min)) for x in timesteps] + [0.0] + sigmas = torch.FloatTensor(sigmas).to(device) + return sigmas + + schedulers = [ Scheduler('automatic', 'Automatic', None), Scheduler('uniform', 'Uniform', uniform, need_inner_model=True), @@ -127,6 +139,7 @@ schedulers = [ Scheduler('simple', 'Simple', simple_scheduler, need_inner_model=True), Scheduler('normal', 'Normal', normal_scheduler, need_inner_model=True), Scheduler('ddim', 'DDIM', ddim_scheduler, need_inner_model=True), + Scheduler('beta', 'Beta', beta_scheduler, need_inner_model=True), ] schedulers_map = {**{x.name: x for x in schedulers}, **{x.label: x for x in schedulers}} From 7e1bd3e3c30e1d7e95f719a925f11d9225251f7c Mon Sep 17 00:00:00 2001 From: v0xie <28695009+v0xie@users.noreply.github.com> Date: Fri, 19 Jul 2024 13:44:22 -0700 Subject: [PATCH 2/8] refactor: syntax and add 0.0 on new line --- modules/sd_schedulers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modules/sd_schedulers.py b/modules/sd_schedulers.py index 19d294c1d..6f2dc6308 100644 --- a/modules/sd_schedulers.py +++ b/modules/sd_schedulers.py @@ -122,9 +122,9 @@ def beta_scheduler(n, sigma_min, sigma_max, inner_model, device): beta = 0.6 timesteps = 1 - np.linspace(0, 1, n) timesteps = [stats.beta.ppf(x, alpha, beta) for x in timesteps] - sigmas = [sigma_min + ((x)*(sigma_max-sigma_min)) for x in timesteps] + [0.0] - sigmas = torch.FloatTensor(sigmas).to(device) - return sigmas + sigmas = [sigma_min + (x * (sigma_max-sigma_min)) for x in timesteps] + sigmas += [0.0] + return torch.FloatTensor(sigmas).to(device) schedulers = [ From 3a5a66775c4c9dd03ffbf3c1696ae0db64a71793 Mon Sep 17 00:00:00 2001 From: v0xie <28695009+v0xie@users.noreply.github.com> Date: Fri, 19 Jul 2024 14:08:08 -0700 Subject: [PATCH 3/8] add new options 'beta_dist_alpha', 'beta_dist_beta' --- modules/shared_options.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modules/shared_options.py b/modules/shared_options.py index 096366e0a..203f7a32b 100644 --- a/modules/shared_options.py +++ b/modules/shared_options.py @@ -404,6 +404,8 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters" 'uni_pc_lower_order_final': OptionInfo(True, "UniPC lower order final", infotext='UniPC lower order final'), 'sd_noise_schedule': OptionInfo("Default", "Noise schedule for sampling", gr.Radio, {"choices": ["Default", "Zero Terminal SNR"]}, infotext="Noise Schedule").info("for use with zero terminal SNR trained models"), 'skip_early_cond': OptionInfo(0.0, "Ignore negative prompt during early sampling", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}, infotext="Skip Early CFG").info("disables CFG on a proportion of steps at the beginning of generation; 0=skip none; 1=skip all; can both improve sample diversity/quality and speed up sampling"), + 'beta_dist_alpha': OptionInfo(0.6, "Beta scheduler - alpha", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}, infotext='Beta scheduler alpha').info('Default = 0.6; the alpha parameter of the beta distribution used in Beta sampling'), + 'beta_dist_beta': OptionInfo(0.6, "Beta scheduler - beta", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}, infotext='Beta scheduler beta').info('Default = 0.6; the beta parameter of the beta distribution used in Beta sampling'), })) options_templates.update(options_section(('postprocessing', "Postprocessing", "postprocessing"), { From f6f055a93df4fb1a59fa8f28e3b8f092fd7ba511 Mon Sep 17 00:00:00 2001 From: v0xie <28695009+v0xie@users.noreply.github.com> Date: Fri, 19 Jul 2024 14:08:44 -0700 Subject: [PATCH 4/8] use configured alpha/beta values in Beta scheduling --- modules/sd_schedulers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/sd_schedulers.py b/modules/sd_schedulers.py index 6f2dc6308..f4d16e309 100644 --- a/modules/sd_schedulers.py +++ b/modules/sd_schedulers.py @@ -118,8 +118,8 @@ def ddim_scheduler(n, sigma_min, sigma_max, inner_model, device): def beta_scheduler(n, sigma_min, sigma_max, inner_model, device): # From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024) """ - alpha = 0.6 - beta = 0.6 + alpha = shared.opts.beta_dist_alpha + beta = shared.opts.beta_dist_beta timesteps = 1 - np.linspace(0, 1, n) timesteps = [stats.beta.ppf(x, alpha, beta) for x in timesteps] sigmas = [sigma_min + (x * (sigma_max-sigma_min)) for x in timesteps] From e285af6e4817a34c86212d1b640aad7225f0c022 Mon Sep 17 00:00:00 2001 From: v0xie <28695009+v0xie@users.noreply.github.com> Date: Fri, 19 Jul 2024 14:15:10 -0700 Subject: [PATCH 5/8] add beta schedule opts to xyz options --- scripts/xyz_grid.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py index 566493266..6a42a04d9 100644 --- a/scripts/xyz_grid.py +++ b/scripts/xyz_grid.py @@ -259,6 +259,8 @@ axis_options = [ AxisOption("Schedule min sigma", float, apply_override("sigma_min")), AxisOption("Schedule max sigma", float, apply_override("sigma_max")), AxisOption("Schedule rho", float, apply_override("rho")), + AxisOption("Beta schedule alpha", float, apply_override("beta_dist_alpha")), + AxisOption("Beta schedule beta", float, apply_override("beta_dist_beta")), AxisOption("Eta", float, apply_field("eta")), AxisOption("Clip skip", int, apply_override('CLIP_stop_at_last_layers')), AxisOption("Denoising", float, apply_field("denoising_strength")), From 94275b115c2a6c3c273baa92caf0b5f4ff3cc43f Mon Sep 17 00:00:00 2001 From: v0xie <28695009+v0xie@users.noreply.github.com> Date: Fri, 19 Jul 2024 14:15:55 -0700 Subject: [PATCH 6/8] enforce beta_dist_alpha / beta_dist_beta > 0 to avoid nan --- modules/shared_options.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/shared_options.py b/modules/shared_options.py index 203f7a32b..4ff7f51c0 100644 --- a/modules/shared_options.py +++ b/modules/shared_options.py @@ -404,8 +404,8 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters" 'uni_pc_lower_order_final': OptionInfo(True, "UniPC lower order final", infotext='UniPC lower order final'), 'sd_noise_schedule': OptionInfo("Default", "Noise schedule for sampling", gr.Radio, {"choices": ["Default", "Zero Terminal SNR"]}, infotext="Noise Schedule").info("for use with zero terminal SNR trained models"), 'skip_early_cond': OptionInfo(0.0, "Ignore negative prompt during early sampling", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}, infotext="Skip Early CFG").info("disables CFG on a proportion of steps at the beginning of generation; 0=skip none; 1=skip all; can both improve sample diversity/quality and speed up sampling"), - 'beta_dist_alpha': OptionInfo(0.6, "Beta scheduler - alpha", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}, infotext='Beta scheduler alpha').info('Default = 0.6; the alpha parameter of the beta distribution used in Beta sampling'), - 'beta_dist_beta': OptionInfo(0.6, "Beta scheduler - beta", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}, infotext='Beta scheduler beta').info('Default = 0.6; the beta parameter of the beta distribution used in Beta sampling'), + 'beta_dist_alpha': OptionInfo(0.6, "Beta scheduler - alpha", gr.Slider, {"minimum": 0.01, "maximum": 1.0, "step": 0.01}, infotext='Beta scheduler alpha').info('Default = 0.6; the alpha parameter of the beta distribution used in Beta sampling'), + 'beta_dist_beta': OptionInfo(0.6, "Beta scheduler - beta", gr.Slider, {"minimum": 0.01, "maximum": 1.0, "step": 0.01}, infotext='Beta scheduler beta').info('Default = 0.6; the beta parameter of the beta distribution used in Beta sampling'), })) options_templates.update(options_section(('postprocessing', "Postprocessing", "postprocessing"), { From 9de7084884f4266a19e4ed3f503687927c845c4d Mon Sep 17 00:00:00 2001 From: v0xie <28695009+v0xie@users.noreply.github.com> Date: Fri, 19 Jul 2024 14:54:24 -0700 Subject: [PATCH 7/8] always add alpha/beta to extra_generation_params when schedule is Beta --- modules/sd_samplers_kdiffusion.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py index 95a354dac..4e5310bc6 100644 --- a/modules/sd_samplers_kdiffusion.py +++ b/modules/sd_samplers_kdiffusion.py @@ -119,6 +119,10 @@ class KDiffusionSampler(sd_samplers_common.Sampler): if scheduler.need_inner_model: sigmas_kwargs['inner_model'] = self.model_wrap + + if scheduler.label == 'Beta': + p.extra_generation_params["Beta schedule alpha"] = opts.beta_dist_alpha + p.extra_generation_params["Beta schedule beta"] = opts.beta_dist_beta sigmas = scheduler.function(n=steps, **sigmas_kwargs, device=devices.cpu) From 874954060297d847bf30cc5d220effe80ac18968 Mon Sep 17 00:00:00 2001 From: v0xie <28695009+v0xie@users.noreply.github.com> Date: Fri, 19 Jul 2024 15:33:07 -0700 Subject: [PATCH 8/8] fix lint --- modules/sd_samplers_kdiffusion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py index 4e5310bc6..0c94d100d 100644 --- a/modules/sd_samplers_kdiffusion.py +++ b/modules/sd_samplers_kdiffusion.py @@ -119,7 +119,7 @@ class KDiffusionSampler(sd_samplers_common.Sampler): if scheduler.need_inner_model: sigmas_kwargs['inner_model'] = self.model_wrap - + if scheduler.label == 'Beta': p.extra_generation_params["Beta schedule alpha"] = opts.beta_dist_alpha p.extra_generation_params["Beta schedule beta"] = opts.beta_dist_beta