Create alternate_sampler_noise_schedules.py

scripts/alternate_sampler_noise_schedules.py

+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)