From 9a09b2eef30264490a0b6f155f8044a47cfe2a04 Mon Sep 17 00:00:00 2001
From: cool-bigdogs-tshirt
 <131823432+cool-bigdogs-tshirt@users.noreply.github.com>
Date: Thu, 27 Apr 2023 11:24:26 -0500
Subject: [PATCH] attempt at unipc latent upscaling

i should have taken linear algebra before i dropped out...
---
 modules/models/diffusion/uni_pc/sampler.py | 74 ++++++++++++++++++++++
 modules/processing.py                      |  3 +-
 2 files changed, 76 insertions(+), 1 deletion(-)

diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py
index a241c8a7c..41b8c9a5b 100644
--- a/modules/models/diffusion/uni_pc/sampler.py
+++ b/modules/models/diffusion/uni_pc/sampler.py
@@ -4,6 +4,7 @@ import torch
 
 from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC
 from modules import shared, devices
+from ldm.modules.diffusionmodules.util import extract_into_tensor
 
 
 class UniPCSampler(object):
@@ -15,6 +16,79 @@ class UniPCSampler(object):
         self.after_sample = None
         self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod))
 
+    def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True):
+        # persist steps so we can eventually find denoising strength
+        self.inflated_steps = ddim_num_steps
+
+    @torch.no_grad()
+    def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
+        if noise is None:
+            noise = torch.randn_like(x0)
+
+        # first time we have all the info to get the real parameters from the ui
+        hires_steps = t[0] + 1
+        inflated_steps = self.inflated_steps
+        self.denoising_strength = hires_steps/inflated_steps
+
+        adjusted_steps = int(hires_steps * self.denoising_strength)
+        self.steps = max(adjusted_steps, shared.opts.uni_pc_order+1)
+
+        t = torch.full(t.shape, self.steps).to(t.device)
+
+        timesteps = torch.asarray(list(range(
+            t,
+            self.model.num_timesteps,
+            self.model.num_timesteps // hires_steps,
+        ))) + 1
+        alphas = self.model.alphas_cumprod[timesteps]
+        sqrt_one_minus_alphas = torch.sqrt(1. - alphas)
+        a = extract_into_tensor(torch.sqrt(alphas), t, x0.shape) * x0 
+        b = extract_into_tensor(sqrt_one_minus_alphas, t, x0.shape) * noise
+
+        return (a+b)
+
+    def decode(self, x_latent, conditioning, t_start, unconditional_guidance_scale=1.0, unconditional_conditioning=None,
+               use_original_steps=False, callback=None):
+        #print(f'steps {self.steps} denoising {self.denoising_strength}')
+
+        noise_schedule = NoiseScheduleVP("discrete", alphas_cumprod=self.alphas_cumprod)
+
+        # same as in .sample(), i guess
+        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),
+            noise_schedule,
+            model_type=model_type,
+            guidance_type="classifier-free",
+            #condition=conditioning,
+            #unconditional_condition=unconditional_conditioning,
+            guidance_scale=unconditional_guidance_scale,
+        )
+
+        self.uni_pc = UniPC(
+                model_fn,
+                noise_schedule,
+                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,
+            )
+
+        return self.uni_pc.sample(
+                x_latent,
+                steps=self.steps,
+                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,
+                t_start=self.denoising_strength,
+            )
+
     def register_buffer(self, name, attr):
         if type(attr) == torch.Tensor:
             if attr.device != devices.device:
diff --git a/modules/processing.py b/modules/processing.py
index 04379fabe..c3ea4b2b8 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -970,7 +970,8 @@ class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
         shared.state.nextjob()
 
         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
+        if self.sampler_name in ['PLMS']:
+            # PLMS does not support img2img, use fallback instead
             img2img_sampler_name = shared.opts.fallback_sampler
         self.sampler = sd_samplers.create_sampler(img2img_sampler_name, self.sd_model)