Merge pull request #568 from cool-bigdogs-tshirt/bigdogs/unipc-latent-upscale

unipc latent upscale

modules/models/diffusion/uni_pc/sampler.py

@@ -1,9 +1,11 @@
 """SAMPLING ONLY."""
 
+import numpy as np
 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 +17,103 @@ 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
+        # value from the hires steps slider:
+        num_inference_steps = t[0] + 1
+        # (num_inference_steps // denoising_strength):
+        inflated_steps = self.inflated_steps
+        # not exact:
+        self.denoising_strength = num_inference_steps/inflated_steps
+
+        # values used for timesteps that generate noise in diffusers repo
+        init_timestep = min(
+                int(num_inference_steps * self.denoising_strength),
+                num_inference_steps,
+            )
+        t_start = max(num_inference_steps - init_timestep, 0)
+
+        # actual number of steps we'll run
+        self.steps = max(
+                num_inference_steps - init_timestep,
+                shared.opts.uni_pc_order+1,
+            )
+
+        scheduler_timesteps = np.linspace(
+                0,
+                self.model.num_timesteps-1,
+                num_inference_steps + 1,
+            ).round()[::-1][:-1].copy().astype(np.int64)
+        _, unique_indices = np.unique(scheduler_timesteps, return_index=True)
+        scheduler_timesteps = scheduler_timesteps[np.sort(unique_indices)]
+        scheduler_timesteps = torch.from_numpy(scheduler_timesteps).to(t.device)
+
+        sample_timesteps = scheduler_timesteps[t_start:]
+        latent_timestep = sample_timesteps[:1].repeat(x0.shape[0])
+
+        alphas_cumprod = self.alphas_cumprod
+        sqrt_alpha_prod = alphas_cumprod[latent_timestep] ** 0.5
+        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
+        while len(sqrt_alpha_prod.shape) < len(x0.shape):
+            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
+
+        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[latent_timestep]) ** 0.5
+        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
+        while len(sqrt_one_minus_alpha_prod.shape) < len(x0.shape):
+            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
+
+        return (sqrt_alpha_prod * x0 + sqrt_one_minus_alpha_prod * noise)
+
+    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:

modules/processing.py

@@ -970,8 +970,10 @@ 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
-            img2img_sampler_name = shared.opts.fallback_sampler
+        force_latent_upscaler = shared.opts.data.get('xyz_fallback_sampler')
+        if self.sampler_name in ['PLMS'] or force_latent_upscaler is not None:
+            # PLMS does not support img2img, use fallback instead
+            img2img_sampler_name = force_latent_upscaler or shared.opts.fallback_sampler
         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]

scripts/xyz_grid.py

@@ -128,6 +128,14 @@ def apply_styles(p: StableDiffusionProcessingTxt2Img, x: str, _):
     p.styles.extend(x.split(','))
 
 
+def apply_fallback(p, x, xs):
+    sampler_name = sd_samplers.samplers_map.get(x.lower(), None)
+    if sampler_name is None:
+        raise RuntimeError(f"Unknown sampler: {x}")
+
+    opts.data["xyz_fallback_sampler"] = sampler_name
+
+
 def apply_uni_pc_order(p, x, xs):
     opts.data["uni_pc_order"] = min(x, p.steps - 1)
 
@@ -220,6 +228,7 @@ axis_options = [
     AxisOption("Clip skip", int, apply_clip_skip),
     AxisOption("Denoising", float, apply_field("denoising_strength")),
     AxisOptionTxt2Img("Hires upscaler", str, apply_field("hr_upscaler"), choices=lambda: [*shared.latent_upscale_modes, *[x.name for x in shared.sd_upscalers]]),
+    AxisOptionTxt2Img("Fallback latent upscaler sampler", str, apply_fallback, format_value=format_value, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]),
     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)),