jank hacky experimental UniPC impl

This is definitely not the ideal way to do this. But. Uh. It exists now.

dynthres_core.py

@@ -48,7 +48,9 @@ class DynThresh:
 
         ### Normal first part of the CFG Scale logic, basically
         diff = cond_stacked - uncond.unsqueeze(1)
-        relative = (diff * weights).sum(1)
+        if weights is not None:
+            diff = diff * weights
+        relative = diff.sum(1)
 
         ### Get the normal result for both mimic and normal scale
         mim_target = uncond + relative * mimicScale

dynthres_unipc.py

@@ -0,0 +1,108 @@
+import gradio as gr
+import torch
+import math
+import traceback
+from modules import shared
+from modules.models.diffusion import uni_pc
+
+######################### UniPC Implementation logic #########################
+
+# The majority of this is straight from modules.models/diffusion/uni_pc/sampler.py
+# Unfortunately that's not an easy middle-injection point, so, just copypasta'd it all
+# It's likely they designed it to intentionally be as difficult to inject into as possible :(
+# (It has hooks but not in useful locations)
+# I stripped the original comments for brevity.
+# Some never-used code (scheduler modes, noise modes, guidance modes) have been removed as well for brevity.
+# The actual impl comes down to just the last line in particular, and the `beforeSample` insert to track step count.
+
+class CustomUniPCSampler(uni_pc.sampler.UniPCSampler):
+    def __init__(self, model, **kwargs):
+        super().__init__(model, *kwargs)
+    @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, **kwargs):
+        if conditioning is not None:
+            if isinstance(conditioning, dict):
+                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}")
+        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 = uni_pc.uni_pc.NoiseScheduleVP('discrete', alphas_cumprod=self.alphas_cumprod)
+        model_type = "v" if self.model.parameterization == "v" else "noise"
+        model_fn = CustomUniPC_model_wrapper(lambda x, t, c: self.model.apply_model(x, t, c), ns, model_type=model_type, guidance_scale=unconditional_guidance_scale, dtData=self.main_class)
+        self.main_class.step = 0
+        def beforeSample(x, t, cond, uncond):
+            self.main_class.step += 1
+            return self.before_sample(x, t, cond, uncond)
+        uni_pc_inst = uni_pc.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=beforeSample, after_sample=self.after_sample, after_update=self.after_update)
+        x = uni_pc_inst.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
+
+def CustomUniPC_model_wrapper(model, noise_schedule, model_type="noise", model_kwargs={}, guidance_scale=1.0, dtData=None):
+    def expand_dims(v, dims):
+        return v[(...,) + (None,)*(dims - 1)]
+    def get_model_input_time(t_continuous):
+        return (t_continuous - 1. / noise_schedule.total_N) * 1000.
+    def noise_pred_fn(x, t_continuous, cond=None):
+        if t_continuous.reshape((-1,)).shape[0] == 1:
+            t_continuous = t_continuous.expand((x.shape[0]))
+        t_input = get_model_input_time(t_continuous)
+        if cond is None:
+            output = model(x, t_input, None, **model_kwargs)
+        else:
+            output = model(x, t_input, cond, **model_kwargs)
+        if model_type == "noise":
+            return output
+        elif model_type == "v":
+            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
+            dims = x.dim()
+            return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
+    def model_fn(x, t_continuous, condition, unconditional_condition):
+        if t_continuous.reshape((-1,)).shape[0] == 1:
+            t_continuous = t_continuous.expand((x.shape[0]))
+        if guidance_scale == 1. or unconditional_condition is None:
+            return noise_pred_fn(x, t_continuous, cond=condition)
+        else:
+            x_in = torch.cat([x] * 2)
+            t_in = torch.cat([t_continuous] * 2)
+            if isinstance(condition, dict):
+                assert isinstance(unconditional_condition, dict)
+                c_in = dict()
+                for k in condition:
+                    if isinstance(condition[k], list):
+                        c_in[k] = [torch.cat([
+                            unconditional_condition[k][i],
+                            condition[k][i]]) for i in range(len(condition[k]))]
+                    else:
+                        c_in[k] = torch.cat([
+                            unconditional_condition[k],
+                            condition[k]])
+            elif isinstance(condition, list):
+                c_in = list()
+                assert isinstance(unconditional_condition, list)
+                for i in range(len(condition)):
+                    c_in.append(torch.cat([unconditional_condition[i], condition[i]]))
+            else:
+                c_in = torch.cat([unconditional_condition, condition])
+            noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
+            #return noise_uncond + guidance_scale * (noise - noise_uncond)
+            return dtData.dynthresh(noise, noise_uncond, guidance_scale, None)
+    return model_fn

scripts/dynamic_thresholding.py

@@ -12,7 +12,7 @@
 
 import gradio as gr
 import torch, traceback
-import dynthres_core
+import dynthres_unipc, dynthres_core
 from modules import scripts, script_callbacks, sd_samplers, sd_samplers_compvis, sd_samplers_kdiffusion, sd_samplers_common
 
 ######################### Data values #########################
@@ -98,13 +98,18 @@ class Script(scripts.Script):
         threshold_percentile *= 0.01
         # Make a placeholder sampler
         sampler = sd_samplers.all_samplers_map[p.sampler_name]
-        def newConstructor(model):
-            dtData = DynThresh(mimic_scale, threshold_percentile, mimic_mode, mimic_scale_min, cfg_mode, cfg_scale_min, power_val, experiment_mode, p.steps)
-            result = sampler.constructor(model)
-            cfg = CustomCFGDenoiser(result.model_wrap_cfg.inner_model, dtData)
-            result.model_wrap_cfg = cfg
-            return result
-        newSampler = sd_samplers_common.SamplerData(fixed_sampler_name, newConstructor, sampler.aliases, sampler.options)
+        dtData = dynthres_core.DynThresh(mimic_scale, threshold_percentile, mimic_mode, mimic_scale_min, cfg_mode, cfg_scale_min, power_val, experiment_mode, p.steps)
+        if p.sampler_name == "UniPC":
+            def uniPCConstructor(model):
+                return CustomVanillaSDSampler(dynthres_unipc.CustomUniPCSampler, model, dtData)
+            newSampler = sd_samplers_common.SamplerData(fixed_sampler_name, uniPCConstructor, sampler.aliases, sampler.options)
+        else:
+            def newConstructor(model):
+                result = sampler.constructor(model)
+                cfg = CustomCFGDenoiser(result.model_wrap_cfg.inner_model, dtData)
+                result.model_wrap_cfg = cfg
+                return result
+            newSampler = sd_samplers_common.SamplerData(fixed_sampler_name, newConstructor, sampler.aliases, sampler.options)
         # Apply for usage
         p.orig_sampler_name = p.sampler_name
         p.sampler_name = fixed_sampler_name
@@ -121,7 +126,12 @@ class Script(scripts.Script):
         del p.orig_sampler_name
         del p.fixed_sampler_name
 
+######################### CompVis Implementation logic #########################
 
+class CustomVanillaSDSampler(sd_samplers_compvis.VanillaStableDiffusionSampler):
+    def __init__(self, constructor, sd_model, dtData):
+        super().__init__(constructor, sd_model)
+        self.sampler.main_class = dtData
 
 ######################### K-Diffusion Implementation logic #########################
 
@@ -135,6 +145,7 @@ class CustomCFGDenoiser(sd_samplers_kdiffusion.CFGDenoiser):
         # conds_list shape is (batch, cond, 2)
         weights = torch.tensor(conds_list, device=uncond.device).select(2, 1)
         weights = weights.reshape(*weights.shape, 1, 1, 1)
+        self.main_class.step = self.step
         return self.main_class.dynthresh(x_out[:-uncond.shape[0]], denoised_uncond, cond_scale, weights)
 
 ######################### XYZ Plot Script Support logic #########################