integrate the new samplers PR

5 files changed, 36 insertions, 87 deletions

diff --git a/modules/processing.py b/modules/processing.py
index e01c8b3f..e567956c 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -477,7 +477,7 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
             self.firstphase_height_truncated = int(scale * self.height)

 

     def sample(self, conditioning, unconditional_conditioning, seeds, subseeds, subseed_strength):

-        self.sampler = sd_samplers.samplers[self.sampler_index].constructor(self.sd_model)

+        self.sampler = sd_samplers.create_sampler_with_index(sd_samplers.samplers, self.sampler_index, self.sd_model)

 

         if not self.enable_hr:

             x = create_random_tensors([opt_C, self.height // opt_f, self.width // opt_f], seeds=seeds, subseeds=subseeds, subseed_strength=self.subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

@@ -520,7 +520,8 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 

         shared.state.nextjob()

 

-        self.sampler = sd_samplers.samplers[self.sampler_index].constructor(self.sd_model)

+        self.sampler = sd_samplers.create_sampler_with_index(sd_samplers.samplers, self.sampler_index, self.sd_model)

+

         noise = create_random_tensors(samples.shape[1:], seeds=seeds, subseeds=subseeds, subseed_strength=subseed_strength, seed_resize_from_h=self.seed_resize_from_h, seed_resize_from_w=self.seed_resize_from_w, p=self)

 

         # GC now before running the next img2img to prevent running out of memory

@@ -555,7 +556,7 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
         self.nmask = None

 

     def init(self, all_prompts, all_seeds, all_subseeds):

-        self.sampler = sd_samplers.samplers_for_img2img[self.sampler_index].constructor(self.sd_model)

+        self.sampler = sd_samplers.create_sampler_with_index(sd_samplers.samplers_for_img2img, self.sampler_index, self.sd_model)

         crop_region = None

 

         if self.image_mask is not None:

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 8d6eb762..497df943 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -13,46 +13,46 @@ from modules.shared import opts, cmd_opts, state
 import modules.shared as shared

 

 

-SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases'])

+SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options'])

 

 samplers_k_diffusion = [

-    ('Euler a', 'sample_euler_ancestral', ['k_euler_a']),

-    ('Euler', 'sample_euler', ['k_euler']),

-    ('LMS', 'sample_lms', ['k_lms']),

-    ('Heun', 'sample_heun', ['k_heun']),

-    ('DPM2', 'sample_dpm_2', ['k_dpm_2']),

-    ('DPM2 a', 'sample_dpm_2_ancestral', ['k_dpm_2_a']),

-    ('DPM fast', 'sample_dpm_fast', ['k_dpm_fast']),

-    ('DPM adaptive', 'sample_dpm_adaptive', ['k_dpm_ad']),

+    ('Euler a', 'sample_euler_ancestral', ['k_euler_a'], {}),

+    ('Euler', 'sample_euler', ['k_euler'], {}),

+    ('LMS', 'sample_lms', ['k_lms'], {}),

+    ('Heun', 'sample_heun', ['k_heun'], {}),

+    ('DPM2', 'sample_dpm_2', ['k_dpm_2'], {}),

+    ('DPM2 a', 'sample_dpm_2_ancestral', ['k_dpm_2_a'], {}),

+    ('DPM fast', 'sample_dpm_fast', ['k_dpm_fast'], {}),

+    ('DPM adaptive', 'sample_dpm_adaptive', ['k_dpm_ad'], {}),

+    ('LMS Karras', 'sample_lms', ['k_lms_ka'], {'scheduler': 'karras'}),

+    ('DPM2 Karras', 'sample_dpm_2', ['k_dpm_2_ka'], {'scheduler': 'karras'}),

+    ('DPM2 a Karras', 'sample_dpm_2_ancestral', ['k_dpm_2_a_ka'], {'scheduler': 'karras'}),

 ]

 

-if opts.show_karras_scheduler_variants:

-    k_diffusion.sampling.sample_dpm_2_ka = k_diffusion.sampling.sample_dpm_2

-    k_diffusion.sampling.sample_dpm_2_ancestral_ka = k_diffusion.sampling.sample_dpm_2_ancestral

-    k_diffusion.sampling.sample_lms_ka = k_diffusion.sampling.sample_lms

-    samplers_k_diffusion_ka = [

-        ('LMS K Scheduling', 'sample_lms_ka', ['k_lms_ka']),

-        ('DPM2 K Scheduling', 'sample_dpm_2_ka', ['k_dpm_2_ka']),

-        ('DPM2 a K Scheduling', 'sample_dpm_2_ancestral_ka', ['k_dpm_2_a_ka']),

-    ]

-    samplers_k_diffusion.extend(samplers_k_diffusion_ka)

-

 samplers_data_k_diffusion = [

-    SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases)

-    for label, funcname, aliases in samplers_k_diffusion

+    SamplerData(label, lambda model, funcname=funcname: KDiffusionSampler(funcname, model), aliases, options)

+    for label, funcname, aliases, options in samplers_k_diffusion

     if hasattr(k_diffusion.sampling, funcname)

 ]

 

 all_samplers = [

     *samplers_data_k_diffusion,

-    SamplerData('DDIM', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.ddim.DDIMSampler, model), []),

-    SamplerData('PLMS', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.plms.PLMSSampler, model), []),

+    SamplerData('DDIM', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.ddim.DDIMSampler, model), [], {}),

+    SamplerData('PLMS', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.plms.PLMSSampler, model), [], {}),

 ]

 

 samplers = []

 samplers_for_img2img = []

 

 

+def create_sampler_with_index(list_of_configs, index, model):

+    config = list_of_configs[index]

+    sampler = config.constructor(model)

+    sampler.config = config

+    

+    return sampler

+

+

 def set_samplers():

     global samplers, samplers_for_img2img

 

@@ -130,6 +130,7 @@ class VanillaStableDiffusionSampler:
         self.step = 0

         self.eta = None

         self.default_eta = 0.0

+        self.config = None

 

     def number_of_needed_noises(self, p):

         return 0

@@ -291,6 +292,7 @@ class KDiffusionSampler:
         self.stop_at = None

         self.eta = None

         self.default_eta = 1.0

+        self.config = None

 

     def callback_state(self, d):

         store_latent(d["denoised"])

@@ -355,11 +357,12 @@ class KDiffusionSampler:
         steps = steps or p.steps

 

         if p.sampler_noise_scheduler_override:

-          sigmas = p.sampler_noise_scheduler_override(steps)

-        elif self.funcname.endswith('ka'):

-          sigmas = k_diffusion.sampling.get_sigmas_karras(n=steps, sigma_min=0.1, sigma_max=10, device=shared.device)

+            sigmas = p.sampler_noise_scheduler_override(steps)

+        elif self.config is not None and self.config.options.get('scheduler', None) == 'karras':

+            sigmas = k_diffusion.sampling.get_sigmas_karras(n=steps, sigma_min=0.1, sigma_max=10, device=shared.device)

         else:

-          sigmas = self.model_wrap.get_sigmas(steps)

+            sigmas = self.model_wrap.get_sigmas(steps)

+

         x = x * sigmas[0]

 

         extra_params_kwargs = self.initialize(p)

diff --git a/modules/shared.py b/modules/shared.py
index 9e4860a2..ca2e4c74 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -236,7 +236,6 @@ options_templates.update(options_section(('ui', "User interface"), {
     "font": OptionInfo("", "Font for image grids that have text"),

     "js_modal_lightbox": OptionInfo(True, "Enable full page image viewer"),

     "js_modal_lightbox_initialy_zoomed": OptionInfo(True, "Show images zoomed in by default in full page image viewer"),

-    "show_karras_scheduler_variants": OptionInfo(True, "Show Karras scheduling variants for select samplers. Try these variants if your K sampled images suffer from excessive noise."),

 }))

 

 options_templates.update(options_section(('sampler-params', "Sampler parameters"), {

diff --git a/scripts/alternate_sampler_noise_schedules.py b/scripts/alternate_sampler_noise_schedules.py
deleted file mode 100644
index 4f3ed8fb..00000000
--- a/scripts/alternate_sampler_noise_schedules.py
+++ /dev/null
@@ -1,53 +0,0 @@
-import inspect
-from modules.processing import Processed, process_images
-import gradio as gr
-import modules.scripts as scripts
-import k_diffusion.sampling
-import torch
-
-
-class Script(scripts.Script):
-
-    def title(self):
-        return "Alternate Sampler Noise Schedules"
-
-    def ui(self, is_img2img):
-      noise_scheduler = gr.Dropdown(label="Noise Scheduler", choices=['Default','Karras','Exponential', 'Variance Preserving'], value='Default', type="index")
-      sched_smin     = gr.Slider(value=0.1,   label="Sigma min",                  minimum=0.0,   maximum=100.0, step=0.5,)
-      sched_smax     = gr.Slider(value=10.0,  label="Sigma max",                  minimum=0.0,   maximum=100.0, step=0.5)
-      sched_rho      = gr.Slider(value=7.0,   label="Sigma rho (Karras only)",    minimum=7.0,   maximum=100.0, step=0.5)
-      sched_beta_d   = gr.Slider(value=19.9,  label="Beta distribution (VP only)",minimum=0.0,   maximum=40.0, step=0.5)
-      sched_beta_min = gr.Slider(value=0.1,   label="Beta min (VP only)",         minimum=0.0,   maximum=40.0, step=0.1)
-      sched_eps_s    = gr.Slider(value=0.001, label="Epsilon (VP only)",          minimum=0.001, maximum=1.0,   step=0.001)
-
-      return [noise_scheduler, sched_smin, sched_smax, sched_rho, sched_beta_d, sched_beta_min, sched_eps_s]
-
-    def run(self, p, noise_scheduler, sched_smin, sched_smax, sched_rho, sched_beta_d, sched_beta_min, sched_eps_s):
-
-      noise_scheduler_func_name = ['-','get_sigmas_karras','get_sigmas_exponential','get_sigmas_vp'][noise_scheduler]
-
-      base_params = {
-        "sigma_min":sched_smin, 
-        "sigma_max":sched_smax, 
-        "rho":sched_rho, 
-        "beta_d":sched_beta_d, 
-        "beta_min":sched_beta_min, 
-        "eps_s":sched_eps_s,
-        "device":"cuda" if torch.cuda.is_available() else "cpu"
-      }
-
-      if hasattr(k_diffusion.sampling,noise_scheduler_func_name):
-
-        sigma_func  = getattr(k_diffusion.sampling,noise_scheduler_func_name)
-        sigma_func_kwargs = {}
-
-        for k,v in base_params.items():
-          if k in inspect.signature(sigma_func).parameters:
-            sigma_func_kwargs[k] = v
-
-        def substitute_noise_scheduler(n):
-          return sigma_func(n,**sigma_func_kwargs)
-
-        p.sampler_noise_scheduler_override = substitute_noise_scheduler
-
-      return process_images(p)
diff --git a/scripts/img2imgalt.py b/scripts/img2imgalt.py
index 0ef137f7..f9894cb0 100644
--- a/scripts/img2imgalt.py
+++ b/scripts/img2imgalt.py
@@ -8,7 +8,6 @@ import gradio as gr
 

 from modules import processing, shared, sd_samplers, prompt_parser

 from modules.processing import Processed

-from modules.sd_samplers import samplers

 from modules.shared import opts, cmd_opts, state

 

 import torch

@@ -159,7 +158,7 @@ class Script(scripts.Script):
             

             combined_noise = ((1 - randomness) * rec_noise + randomness * rand_noise) / ((randomness**2 + (1-randomness)**2) ** 0.5)

             

-            sampler = samplers[p.sampler_index].constructor(p.sd_model)

+            sampler = sd_samplers.create_sampler_with_index(sd_samplers.samplers, p.sampler_index, p.sd_model)

 

             sigmas = sampler.model_wrap.get_sigmas(p.steps)