Add custom k_diffusion scheduler in settings

modules/generation_parameters_copypaste.py

@@ -306,6 +306,18 @@ def parse_generation_parameters(x: str):
     if "RNG" not in res:
         res["RNG"] = "GPU"
 
+    if "KDiff Schedule Type" not in res:
+        res["KDiff Schedule Type"] = "Automatic"
+
+    if "KDiff Schedule max sigma" not in res:
+        res["KDiff Schedule max sigma"] = 14.6
+
+    if "KDiff Schedule min sigma" not in res:
+        res["KDiff Schedule min sigma"] = 0.3
+
+    if "KDiff Schedule rho" not in res:
+        res["KDiff Schedule rho"] = 7.0
+
     return res
 
 
@@ -318,6 +330,10 @@ def parse_generation_parameters(x: str):
     ('Conditional mask weight', 'inpainting_mask_weight'),
     ('Model hash', 'sd_model_checkpoint'),
     ('ENSD', 'eta_noise_seed_delta'),
+    ('KDiff Schedule Type', 'k_sched_type'),
+    ('KDiff Schedule max sigma', 'sigma_max'),
+    ('KDiff Schedule min sigma', 'sigma_min'),
+    ('KDiff Schedule rho', 'rho'),
     ('Noise multiplier', 'initial_noise_multiplier'),
     ('Eta', 'eta_ancestral'),
     ('Eta DDIM', 'eta_ddim'),

modules/sd_samplers_kdiffusion.py

@@ -44,6 +44,14 @@
     'sample_dpm_2': ['s_churn', 's_tmin', 's_tmax', 's_noise'],
 }
 
+k_diffusion_samplers_map = {x.name: x for x in samplers_data_k_diffusion}
+k_diffusion_scheduler = {
+    'Automatic': None,
+    'karras': k_diffusion.sampling.get_sigmas_karras,
+    'exponential': k_diffusion.sampling.get_sigmas_exponential,
+    'polyexponential': k_diffusion.sampling.get_sigmas_polyexponential
+}
+
 
 class CFGDenoiser(torch.nn.Module):
     """
@@ -265,6 +273,13 @@ def launch_sampling(self, steps, func):
 
         try:
             return func()
+        except RecursionError:
+            print(
+                'Encountered RecursionError during sampling, returning last latent. '
+                'rho >5 with a polyexponential scheduler may cause this error. '
+                'You should try to use a smaller rho value instead.'
+            )
+            return self.last_latent
         except sd_samplers_common.InterruptedException:
             return self.last_latent
 
@@ -304,6 +319,29 @@ def get_sigmas(self, p, steps):
 
         if p.sampler_noise_scheduler_override:
             sigmas = p.sampler_noise_scheduler_override(steps)
+        elif opts.k_sched_type != "Automatic":
+            m_sigma_min, m_sigma_max = (self.model_wrap.sigmas[0].item(), self.model_wrap.sigmas[-1].item())
+            sigma_min, sigma_max = (0.1, 10)
+            sigmas_kwargs = {
+                'sigma_min': sigma_min if opts.use_old_karras_scheduler_sigmas else m_sigma_min,
+                'sigma_max': sigma_max if opts.use_old_karras_scheduler_sigmas else m_sigma_max
+            }
+
+            sigmas_func = k_diffusion_scheduler[opts.k_sched_type]
+            p.extra_generation_params["KDiff Schedule Type"] = opts.k_sched_type
+
+            if opts.sigma_min != 0.3:
+                # take 0.0 as model default
+                sigmas_kwargs['sigma_min'] = opts.sigma_min or m_sigma_min
+                p.extra_generation_params["KDiff Schedule min sigma"] = opts.sigma_min
+            if opts.sigma_max != 14.6:
+                sigmas_kwargs['sigma_max'] = opts.sigma_max or m_sigma_max
+                p.extra_generation_params["KDiff Schedule max sigma"] = opts.sigma_max
+            if opts.k_sched_type != 'exponential':
+                sigmas_kwargs['rho'] = opts.rho
+                p.extra_generation_params["KDiff Schedule rho"] = opts.rho
+
+            sigmas = sigmas_func(n=steps, **sigmas_kwargs, device=shared.device)
         elif self.config is not None and self.config.options.get('scheduler', None) == 'karras':
             sigma_min, sigma_max = (0.1, 10) if opts.use_old_karras_scheduler_sigmas else (self.model_wrap.sigmas[0].item(), self.model_wrap.sigmas[-1].item())
 

modules/shared.py

@@ -517,6 +517,10 @@ def list_samplers():
     's_churn': OptionInfo(0.0, "sigma churn", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
     's_tmin':  OptionInfo(0.0, "sigma tmin",  gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
     's_noise': OptionInfo(1.0, "sigma noise", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+    'k_sched_type':  OptionInfo("Automatic", "scheduler type", gr.Dropdown, {"choices": ["Automatic", "karras", "exponential", "polyexponential"]}),
+    'sigma_max': OptionInfo(14.6, "sigma max", gr.Number).info("the maximum noise strength for the scheduler. Set to 0 to use the same value which 'xxx karras' samplers use."),
+    'sigma_min': OptionInfo(0.3, "sigma min", gr.Number).info("the minimum noise strength for the scheduler. Set to 0 to use the same value which 'xxx karras' samplers use."),
+    'rho':  OptionInfo(7.0, "rho", gr.Number).info("higher will make a more steep noise scheduler (decrease faster). default for karras is 7.0, for polyexponential is 1.0"),
     'eta_noise_seed_delta': OptionInfo(0, "Eta noise seed delta", gr.Number, {"precision": 0}).info("ENSD; does not improve anything, just produces different results for ancestral samplers - only useful for reproducing images"),
     'always_discard_next_to_last_sigma': OptionInfo(False, "Always discard next-to-last sigma").link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/6044"),
     'uni_pc_variant': OptionInfo("bh1", "UniPC variant", gr.Radio, {"choices": ["bh1", "bh2", "vary_coeff"]}),

scripts/xyz_grid.py

@@ -10,7 +10,7 @@
 import modules.scripts as scripts
 import gradio as gr
 
-from modules import images, sd_samplers, processing, sd_models, sd_vae
+from modules import images, sd_samplers, processing, sd_models, sd_vae, sd_samplers_kdiffusion
 from modules.processing import process_images, Processed, StableDiffusionProcessingTxt2Img
 from modules.shared import opts, state
 import modules.shared as shared
@@ -220,6 +220,10 @@ def __init__(self, *args, **kwargs):
     AxisOption("Sigma min", float, apply_field("s_tmin")),
     AxisOption("Sigma max", float, apply_field("s_tmax")),
     AxisOption("Sigma noise", float, apply_field("s_noise")),
+    AxisOption("KDiff Schedule Type", str, apply_override("k_sched_type"), choices=lambda: list(sd_samplers_kdiffusion.k_diffusion_scheduler)),
+    AxisOption("KDiff Schedule min sigma", float, apply_override("sigma_min")),
+    AxisOption("KDiff Schedule max sigma", float, apply_override("sigma_max")),
+    AxisOption("KDiff Schedule rho", float, apply_override("rho")),
     AxisOption("Eta", float, apply_field("eta")),
     AxisOption("Clip skip", int, apply_clip_skip),
     AxisOption("Denoising", float, apply_field("denoising_strength")),