From 79ffb9453f8eddbdd4e316b9d9c75812b0eea4e1 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Fri, 10 Feb 2023 05:27:05 -0800
Subject: Add UniPC sampler settings

---
 modules/models/diffusion/uni_pc/sampler.py | 5 +++--
 modules/models/diffusion/uni_pc/uni_pc.py  | 2 +-
 modules/shared.py                          | 5 +++++
 3 files changed, 9 insertions(+), 3 deletions(-)

(limited to 'modules')

diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py
index 219e9862..e66a21e3 100644
--- a/modules/models/diffusion/uni_pc/sampler.py
+++ b/modules/models/diffusion/uni_pc/sampler.py
@@ -3,6 +3,7 @@
 import torch
 
 from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC
+from modules import shared
 
 class UniPCSampler(object):
     def __init__(self, model, **kwargs):
@@ -89,7 +90,7 @@ class UniPCSampler(object):
             guidance_scale=unconditional_guidance_scale,
         )
 
-        uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample, after_update=self.after_update)
-        x = uni_pc.sample(img, steps=S, skip_type="time_uniform", method="multistep", order=3, lower_order_final=True)
+        uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=shared.opts.uni_pc_thresholding, variant=shared.opts.uni_pc_variant, condition=conditioning, unconditional_condition=unconditional_conditioning, before_sample=self.before_sample, after_sample=self.after_sample, after_update=self.after_update)
+        x = uni_pc.sample(img, steps=S, skip_type=shared.opts.uni_pc_skip_type, method="multistep", order=shared.opts.uni_pc_order, lower_order_final=shared.opts.uni_pc_lower_order_final)
 
         return x.to(device), None
diff --git a/modules/models/diffusion/uni_pc/uni_pc.py b/modules/models/diffusion/uni_pc/uni_pc.py
index 31ee81a6..df63d1bc 100644
--- a/modules/models/diffusion/uni_pc/uni_pc.py
+++ b/modules/models/diffusion/uni_pc/uni_pc.py
@@ -750,7 +750,7 @@ class UniPC:
         if method == 'multistep':
             assert steps >= order, "UniPC order must be < sampling steps"
             timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
-            print(f"Running UniPC Sampling with {timesteps.shape[0]} timesteps")
+            print(f"Running UniPC Sampling with {timesteps.shape[0]} timesteps, order {order}")
             assert timesteps.shape[0] - 1 == steps
             with torch.no_grad():
                 vec_t = timesteps[0].expand((x.shape[0]))
diff --git a/modules/shared.py b/modules/shared.py
index 79fbf724..34242073 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -480,6 +480,11 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters"
     's_noise': OptionInfo(1.0, "sigma noise", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
     'eta_noise_seed_delta': OptionInfo(0, "Eta noise seed delta", gr.Number, {"precision": 0}),
     'always_discard_next_to_last_sigma': OptionInfo(False, "Always discard next-to-last sigma"),
+    'uni_pc_variant': OptionInfo("bh1", "UniPC variant", gr.Radio, {"choices": ["bh1", "vary_coeff"]}),
+    'uni_pc_skip_type': OptionInfo("time_uniform", "UniPC skip type", gr.Radio, {"choices": ["time_uniform", "time_quadratic", "logSNR"]}),
+    'uni_pc_order': OptionInfo(3, "UniPC order (must be < sampling steps)", gr.Slider, {"minimum": 1, "maximum": 150 - 1, "step": 1}),
+    'uni_pc_thresholding': OptionInfo(False, "UniPC thresholding"),
+    'uni_pc_lower_order_final': OptionInfo(True, "UniPC lower order final"),
 }))
 
 options_templates.update(options_section(('postprocessing', "Postprocessing"), {
-- 
cgit v1.2.3