From 125319988984987801dc4b4ab1e5ed36e9b211c5 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Fri, 10 Feb 2023 03:30:20 -0800
Subject: Working UniPC (for batch size 1)

---
 modules/models/diffusion/uni_pc/sampler.py | 85 ++++++++++++++++++++++++++++++
 1 file changed, 85 insertions(+)
 create mode 100644 modules/models/diffusion/uni_pc/sampler.py

(limited to 'modules/models/diffusion/uni_pc/sampler.py')

diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py
new file mode 100644
index 00000000..7cccd8a2
--- /dev/null
+++ b/modules/models/diffusion/uni_pc/sampler.py
@@ -0,0 +1,85 @@
+"""SAMPLING ONLY."""
+
+import torch
+
+from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC
+
+class UniPCSampler(object):
+    def __init__(self, model, **kwargs):
+        super().__init__()
+        self.model = model
+        to_torch = lambda x: x.clone().detach().to(torch.float32).to(model.device)
+        self.before_sample = None
+        self.after_sample = None
+        self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod))
+
+    def register_buffer(self, name, attr):
+        if type(attr) == torch.Tensor:
+            if attr.device != torch.device("cuda"):
+                attr = attr.to(torch.device("cuda"))
+        setattr(self, name, attr)
+
+    def set_hooks(self, before, after):
+        self.before_sample = before
+        self.after_sample = after
+
+    @torch.no_grad()
+    def sample(self,
+               S,
+               batch_size,
+               shape,
+               conditioning=None,
+               callback=None,
+               normals_sequence=None,
+               img_callback=None,
+               quantize_x0=False,
+               eta=0.,
+               mask=None,
+               x0=None,
+               temperature=1.,
+               noise_dropout=0.,
+               score_corrector=None,
+               corrector_kwargs=None,
+               verbose=True,
+               x_T=None,
+               log_every_t=100,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,
+               # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+               **kwargs
+               ):
+        if conditioning is not None:
+            if isinstance(conditioning, dict):
+                cbs = conditioning[list(conditioning.keys())[0]].shape[0]
+                if cbs != batch_size:
+                    print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+            else:
+                if conditioning.shape[0] != batch_size:
+                    print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
+
+        # sampling
+        C, H, W = shape
+        size = (batch_size, C, H, W)
+
+        device = self.model.betas.device
+        if x_T is None:
+            img = torch.randn(size, device=device)
+        else:
+            img = x_T
+
+        ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)
+
+        model_fn = model_wrapper(
+            lambda x, t, c: self.model.apply_model(x, t, c),
+            ns,
+            model_type="noise",
+            guidance_type="classifier-free",
+            #condition=conditioning,
+            #unconditional_condition=unconditional_conditioning,
+            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)
+        x = uni_pc.sample(img, steps=S, skip_type="time_uniform", method="multistep", order=3, lower_order_final=True)
+
+        return x.to(device), None
-- 
cgit v1.2.3


From 21880eb9e57b884635a07d2360831b4186afddf4 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Fri, 10 Feb 2023 04:47:08 -0800
Subject: Fix logspam and live previews

---
 modules/models/diffusion/uni_pc/sampler.py | 20 ++++++++++++++-----
 modules/models/diffusion/uni_pc/uni_pc.py  | 32 ++++++++++++++----------------
 modules/sd_samplers_compvis.py             | 20 ++++++++++---------
 3 files changed, 41 insertions(+), 31 deletions(-)

(limited to 'modules/models/diffusion/uni_pc/sampler.py')

diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py
index 7cccd8a2..219e9862 100644
--- a/modules/models/diffusion/uni_pc/sampler.py
+++ b/modules/models/diffusion/uni_pc/sampler.py
@@ -19,9 +19,10 @@ class UniPCSampler(object):
                 attr = attr.to(torch.device("cuda"))
         setattr(self, name, attr)
 
-    def set_hooks(self, before, after):
-        self.before_sample = before
-        self.after_sample = after
+    def set_hooks(self, before_sample, after_sample, after_update):
+        self.before_sample = before_sample
+        self.after_sample = after_sample
+        self.after_update = after_update
 
     @torch.no_grad()
     def sample(self,
@@ -50,9 +51,17 @@ class UniPCSampler(object):
                ):
         if conditioning is not None:
             if isinstance(conditioning, dict):
-                cbs = conditioning[list(conditioning.keys())[0]].shape[0]
+                ctmp = conditioning[list(conditioning.keys())[0]]
+                while isinstance(ctmp, list): ctmp = ctmp[0]
+                cbs = ctmp.shape[0]
                 if cbs != batch_size:
                     print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+
+            elif isinstance(conditioning, list):
+                for ctmp in conditioning:
+                    if ctmp.shape[0] != batch_size:
+                        print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+
             else:
                 if conditioning.shape[0] != batch_size:
                     print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
@@ -60,6 +69,7 @@ class UniPCSampler(object):
         # sampling
         C, H, W = shape
         size = (batch_size, C, H, W)
+        print(f'Data shape for UniPC sampling is {size}, eta {eta}')
 
         device = self.model.betas.device
         if x_T is None:
@@ -79,7 +89,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)
+        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)
 
         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 ec6b37da..31ee81a6 100644
--- a/modules/models/diffusion/uni_pc/uni_pc.py
+++ b/modules/models/diffusion/uni_pc/uni_pc.py
@@ -378,7 +378,8 @@ class UniPC:
         condition=None,
         unconditional_condition=None,
         before_sample=None,
-        after_sample=None
+        after_sample=None,
+        after_update=None
     ):
         """Construct a UniPC.
 
@@ -394,6 +395,7 @@ class UniPC:
         self.unconditional_condition = unconditional_condition
         self.before_sample = before_sample
         self.after_sample = after_sample
+        self.after_update = after_update
 
     def dynamic_thresholding_fn(self, x0, t=None):
         """
@@ -434,15 +436,6 @@ class UniPC:
         noise = self.noise_prediction_fn(x, t)
         dims = x.dim()
         alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
-        from pprint import pp
-        print("X:")
-        pp(x)
-        print("sigma_t:")
-        pp(sigma_t)
-        print("noise:")
-        pp(noise)
-        print("alpha_t:")
-        pp(alpha_t)
         x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
         if self.thresholding:
             p = 0.995   # A hyperparameter in the paper of "Imagen" [1].
@@ -524,7 +517,7 @@ class UniPC:
             return self.multistep_uni_pc_vary_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
 
     def multistep_uni_pc_vary_update(self, x, model_prev_list, t_prev_list, t, order, use_corrector=True):
-        print(f'using unified predictor-corrector with order {order} (solver type: vary coeff)')
+        #print(f'using unified predictor-corrector with order {order} (solver type: vary coeff)')
         ns = self.noise_schedule
         assert order <= len(model_prev_list)
 
@@ -568,7 +561,7 @@ class UniPC:
             A_p = C_inv_p
 
         if use_corrector:
-            print('using corrector')
+            #print('using corrector')
             C_inv = torch.linalg.inv(C)
             A_c = C_inv
 
@@ -627,7 +620,7 @@ class UniPC:
         return x_t, model_t
 
     def multistep_uni_pc_bh_update(self, x, model_prev_list, t_prev_list, t, order, x_t=None, use_corrector=True):
-        print(f'using unified predictor-corrector with order {order} (solver type: B(h))')
+        #print(f'using unified predictor-corrector with order {order} (solver type: B(h))')
         ns = self.noise_schedule
         assert order <= len(model_prev_list)
         dims = x.dim()
@@ -695,7 +688,7 @@ class UniPC:
             D1s = None
 
         if use_corrector:
-            print('using corrector')
+            #print('using corrector')
             # for order 1, we use a simplified version
             if order == 1:
                 rhos_c = torch.tensor([0.5], device=b.device)
@@ -755,8 +748,9 @@ class UniPC:
         t_T = self.noise_schedule.T if t_start is None else t_start
         device = x.device
         if method == 'multistep':
-            assert steps >= order
+            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")
             assert timesteps.shape[0] - 1 == steps
             with torch.no_grad():
                 vec_t = timesteps[0].expand((x.shape[0]))
@@ -768,6 +762,8 @@ class UniPC:
                     x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, init_order, use_corrector=True)
                     if model_x is None:
                         model_x = self.model_fn(x, vec_t)
+                    if self.after_update is not None:
+                        self.after_update(x, model_x)
                     model_prev_list.append(model_x)
                     t_prev_list.append(vec_t)
                 for step in range(order, steps + 1):
@@ -776,13 +772,15 @@ class UniPC:
                         step_order = min(order, steps + 1 - step)
                     else:
                         step_order = order
-                    print('this step order:', step_order)
+                    #print('this step order:', step_order)
                     if step == steps:
-                        print('do not run corrector at the last step')
+                        #print('do not run corrector at the last step')
                         use_corrector = False
                     else:
                         use_corrector = True
                     x, model_x =  self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, step_order, use_corrector=use_corrector)
+                    if self.after_update is not None:
+                        self.after_update(x, model_x)
                     for i in range(order - 1):
                         t_prev_list[i] = t_prev_list[i + 1]
                         model_prev_list[i] = model_prev_list[i + 1]
diff --git a/modules/sd_samplers_compvis.py b/modules/sd_samplers_compvis.py
index 86fa1c5b..946079ae 100644
--- a/modules/sd_samplers_compvis.py
+++ b/modules/sd_samplers_compvis.py
@@ -103,16 +103,11 @@ class VanillaStableDiffusionSampler:
 
         return x, ts, cond, unconditional_conditioning
 
-    def after_sample(self, x, ts, cond, uncond, res):
-        if self.is_unipc:
-            # unipc model_fn returns (pred_x0)
-            # p_sample_ddim returns (x_prev, pred_x0)
-            res = (None, res[0])
-
+    def update_step(self, last_latent):
         if self.mask is not None:
-            self.last_latent = self.init_latent * self.mask + self.nmask * res[1]
+            self.last_latent = self.init_latent * self.mask + self.nmask * last_latent
         else:
-            self.last_latent = res[1]
+            self.last_latent = last_latent
 
         sd_samplers_common.store_latent(self.last_latent)
 
@@ -120,8 +115,15 @@ class VanillaStableDiffusionSampler:
         state.sampling_step = self.step
         shared.total_tqdm.update()
 
+    def after_sample(self, x, ts, cond, uncond, res):
+        if not self.is_unipc:
+            self.update_step(res[1])
+
         return x, ts, cond, uncond, res
 
+    def unipc_after_update(self, x, model_x):
+        self.update_step(x)
+
     def initialize(self, p):
         self.eta = p.eta if p.eta is not None else shared.opts.eta_ddim
         if self.eta != 0.0:
@@ -131,7 +133,7 @@ class VanillaStableDiffusionSampler:
             if hasattr(self.sampler, fieldname):
                 setattr(self.sampler, fieldname, self.p_sample_ddim_hook)
         if self.is_unipc:
-            self.sampler.set_hooks(lambda x, t, c, u: self.before_sample(x, t, c, u), lambda x, t, c, u, r: self.after_sample(x, t, c, u, r))
+            self.sampler.set_hooks(lambda x, t, c, u: self.before_sample(x, t, c, u), lambda x, t, c, u, r: self.after_sample(x, t, c, u, r), lambda x, mx: self.unipc_after_update(x, mx))
 
         self.mask = p.mask if hasattr(p, 'mask') else None
         self.nmask = p.nmask if hasattr(p, 'nmask') else None
-- 
cgit v1.2.3


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 +++++
 scripts/xyz_grid.py                        | 7 +++++++
 4 files changed, 16 insertions(+), 3 deletions(-)

(limited to 'modules/models/diffusion/uni_pc/sampler.py')

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"), {
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index 5982cfba..72421e0c 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -126,6 +126,10 @@ def apply_styles(p: StableDiffusionProcessingTxt2Img, x: str, _):
     p.styles.extend(x.split(','))
 
 
+def apply_uni_pc_order(p, x, xs):
+    opts.data["uni_pc_order"] = min(x, p.steps - 1)
+
+
 def format_value_add_label(p, opt, x):
     if type(x) == float:
         x = round(x, 8)
@@ -202,6 +206,7 @@ axis_options = [
     AxisOptionImg2Img("Cond. Image Mask Weight", float, apply_field("inpainting_mask_weight")),
     AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: list(sd_vae.vae_dict)),
     AxisOption("Styles", str, apply_styles, choices=lambda: list(shared.prompt_styles.styles)),
+    AxisOption("UniPC Order", int, apply_uni_pc_order, cost=0.5),
 ]
 
 
@@ -310,9 +315,11 @@ class SharedSettingsStackHelper(object):
     def __enter__(self):
         self.CLIP_stop_at_last_layers = opts.CLIP_stop_at_last_layers
         self.vae = opts.sd_vae
+        self.uni_pc_order = opts.uni_pc_order
   
     def __exit__(self, exc_type, exc_value, tb):
         opts.data["sd_vae"] = self.vae
+        opts.data["uni_pc_order"] = self.uni_pc_order
         modules.sd_models.reload_model_weights()
         modules.sd_vae.reload_vae_weights()
 
-- 
cgit v1.2.3


From fb274229b2c5c1a89dac0b3da28c08c92d71fd95 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Fri, 10 Feb 2023 14:30:35 -0800
Subject: bug fix

---
 modules/models/diffusion/uni_pc/sampler.py | 2 +-
 modules/processing.py                      | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

(limited to 'modules/models/diffusion/uni_pc/sampler.py')

diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py
index e66a21e3..0bef6eed 100644
--- a/modules/models/diffusion/uni_pc/sampler.py
+++ b/modules/models/diffusion/uni_pc/sampler.py
@@ -70,7 +70,7 @@ class UniPCSampler(object):
         # sampling
         C, H, W = shape
         size = (batch_size, C, H, W)
-        print(f'Data shape for UniPC sampling is {size}, eta {eta}')
+        print(f'Data shape for UniPC sampling is {size}')
 
         device = self.model.betas.device
         if x_T is None:
diff --git a/modules/processing.py b/modules/processing.py
index b7cf5357..0ca15491 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -884,7 +884,9 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
 
         shared.state.nextjob()
 
-        img2img_sampler_name = 'DDIM'  # PLMS/UniPC does not support img2img so we just silently switch ot DDIM
+        img2img_sampler_name = self.sampler_name
+        if self.sampler_name in ['PLMS', 'UniPC']:  # PLMS/UniPC do not support img2img so we just silently switch to DDIM
+            img2img_sampler_name = 'DDIM'
         self.sampler = sd_samplers.create_sampler(img2img_sampler_name, self.sd_model)
 
         samples = samples[:, :, self.truncate_y//2:samples.shape[2]-(self.truncate_y+1)//2, self.truncate_x//2:samples.shape[3]-(self.truncate_x+1)//2]
-- 
cgit v1.2.3


From 716a69237cefb385f71105dbbf50e92d664e0f42 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Sat, 11 Feb 2023 06:18:34 -0800
Subject: support SD2.X models

---
 modules/models/diffusion/uni_pc/sampler.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

(limited to 'modules/models/diffusion/uni_pc/sampler.py')

diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py
index 0bef6eed..708a9b2b 100644
--- a/modules/models/diffusion/uni_pc/sampler.py
+++ b/modules/models/diffusion/uni_pc/sampler.py
@@ -80,10 +80,13 @@ class UniPCSampler(object):
 
         ns = NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)
 
+        # SD 1.X is "noise", SD 2.X is "v"
+        model_type = "v" if self.model.parameterization == "v" else "noise"
+
         model_fn = model_wrapper(
             lambda x, t, c: self.model.apply_model(x, t, c),
             ns,
-            model_type="noise",
+            model_type=model_type,
             guidance_type="classifier-free",
             #condition=conditioning,
             #unconditional_condition=unconditional_conditioning,
-- 
cgit v1.2.3


From 5fef67f6ee949a61826a3a043ea8610fd89fc371 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Fri, 10 Mar 2023 19:56:14 -0500
Subject: Requested changes

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

(limited to 'modules/models/diffusion/uni_pc/sampler.py')

diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py
index 708a9b2b..6bb3bb21 100644
--- a/modules/models/diffusion/uni_pc/sampler.py
+++ b/modules/models/diffusion/uni_pc/sampler.py
@@ -93,7 +93,7 @@ class UniPCSampler(object):
             guidance_scale=unconditional_guidance_scale,
         )
 
-        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)
+        uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, 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/sd_samplers_compvis.py b/modules/sd_samplers_compvis.py
index ad39ab2b..7d07c4a5 100644
--- a/modules/sd_samplers_compvis.py
+++ b/modules/sd_samplers_compvis.py
@@ -140,10 +140,12 @@ class VanillaStableDiffusionSampler:
 
     def adjust_steps_if_invalid(self, p, num_steps):
         if ((self.config.name == 'DDIM') and p.ddim_discretize == 'uniform') or (self.config.name == 'PLMS') or (self.config.name == 'UniPC'):
+            if self.config.name == 'UniPC' and num_steps < shared.opts.uni_pc_order:
+                num_steps = shared.opts.uni_pc_order
             valid_step = 999 / (1000 // num_steps)
             if valid_step == math.floor(valid_step):
                 return int(valid_step) + 1
-        
+
         return num_steps
 
     def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
diff --git a/modules/shared.py b/modules/shared.py
index 7c559fa4..29f8dccb 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -485,10 +485,9 @@ 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_variant': OptionInfo("bh1", "UniPC variant", gr.Radio, {"choices": ["bh1", "bh2", "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": 50, "step": 1}),
-    'uni_pc_thresholding': OptionInfo(False, "UniPC thresholding"),
     'uni_pc_lower_order_final': OptionInfo(True, "UniPC lower order final"),
 }))
 
-- 
cgit v1.2.3