adjust unipc img2img parameters

modules/models/diffusion/uni_pc/sampler.py

@@ -7,9 +7,8 @@ try:
 except:
     pass
 
-from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC
+from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC, get_time_steps
 from modules import shared, devices
-from ldm.modules.diffusionmodules.util import extract_into_tensor
 
 
 class UniPCSampler(object):
@@ -21,6 +20,8 @@ class UniPCSampler(object):
         self.after_sample = None
         self.register_buffer('alphas_cumprod', to_torch(model.alphas_cumprod))
 
+        self.noise_schedule = NoiseScheduleVP("discrete", alphas_cumprod=self.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
@@ -33,35 +34,32 @@ class UniPCSampler(object):
         # 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
+        approx_denoise_strength = num_inference_steps / self.inflated_steps
+        self.denoise_steps = max(num_inference_steps, shared.opts.uni_pc_order)
 
-        # 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)
+        init_timestep = max(self.inflated_steps - self.denoise_steps, 0)
 
         # actual number of steps we'll run
-        self.steps = max(
-                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)
+        all_timesteps = get_time_steps(
+            self.noise_schedule,
+            shared.opts.uni_pc_skip_type,
+            self.noise_schedule.T,
+            1./self.noise_schedule.total_N,
+            self.inflated_steps+1,
+            t.device,
+        )
 
-        sample_timesteps = scheduler_timesteps[t_start:]
-        latent_timestep = sample_timesteps[:1].repeat(x0.shape[0])
+        # the rest of the timesteps will be used for denoising
+        self.timesteps = all_timesteps[-(self.denoise_steps+1):]
+
+        latent_timestep = (
+            (   # get the timestep of our first denoise step
+                self.timesteps[:1]
+                # multiply by number of alphas to get int index
+                * self.noise_schedule.total_N
+            ).int() - 1 # minus one for 0-indexed
+        ).repeat(x0.shape[0])
 
         alphas_cumprod = self.alphas_cumprod
         sqrt_alpha_prod = alphas_cumprod[latent_timestep] ** 0.5
@@ -78,16 +76,12 @@ class UniPCSampler(object):
 
     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,
+            self.noise_schedule,
             model_type=model_type,
             guidance_type="classifier-free",
             #condition=conditioning,
@@ -97,7 +91,7 @@ class UniPCSampler(object):
 
         self.uni_pc = UniPC(
                 model_fn,
-                noise_schedule,
+                self.noise_schedule,
                 predict_x0=True,
                 thresholding=False,
                 variant=shared.opts.uni_pc_variant,
@@ -110,12 +104,13 @@ class UniPCSampler(object):
 
         return self.uni_pc.sample(
                 x_latent,
-                steps=self.steps,
+                steps=self.denoise_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,
+                denoise_to_zero=True,
+                timesteps=self.timesteps,
             )
 
     def register_buffer(self, name, attr):
@@ -182,14 +177,12 @@ class UniPCSampler(object):
         else:
             img = x_T
 
-        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,
+            self.noise_schedule,
             model_type=model_type,
             guidance_type="classifier-free",
             #condition=conditioning,
@@ -197,7 +190,7 @@ class UniPCSampler(object):
             guidance_scale=unconditional_guidance_scale,
         )
 
-        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)
+        uni_pc = UniPC(model_fn, self.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)
         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

modules/models/diffusion/uni_pc/uni_pc.py

@@ -371,6 +371,22 @@ def model_wrapper(
     assert guidance_type in ["uncond", "classifier", "classifier-free"]
     return model_fn
 
+def get_time_steps(noise_schedule, skip_type, t_T, t_0, N, device):
+    """Compute the intermediate time steps for sampling.
+    """
+    if skip_type == 'logSNR':
+        lambda_T = noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
+        lambda_0 = noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
+        logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
+        return noise_schedule.inverse_lambda(logSNR_steps)
+    elif skip_type == 'time_uniform':
+        return torch.linspace(t_T, t_0, N + 1).to(device)
+    elif skip_type == 'time_quadratic':
+        t_order = 2
+        t = torch.linspace(t_T**(1. / t_order), t_0**(1. / t_order), N + 1).pow(t_order).to(device)
+        return t
+    else:
+        raise ValueError("Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
 
 class UniPC:
     def __init__(
@@ -459,23 +475,6 @@ class UniPC:
         else:
             return self.noise_prediction_fn(x, t)
 
-    def get_time_steps(self, skip_type, t_T, t_0, N, device):
-        """Compute the intermediate time steps for sampling.
-        """
-        if skip_type == 'logSNR':
-            lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
-            lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
-            logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
-            return self.noise_schedule.inverse_lambda(logSNR_steps)
-        elif skip_type == 'time_uniform':
-            return torch.linspace(t_T, t_0, N + 1).to(device)
-        elif skip_type == 'time_quadratic':
-            t_order = 2
-            t = torch.linspace(t_T**(1. / t_order), t_0**(1. / t_order), N + 1).pow(t_order).to(device)
-            return t
-        else:
-            raise ValueError("Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
-
     def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
         """
         Get the order of each step for sampling by the singlestep DPM-Solver.
@@ -502,9 +501,9 @@ class UniPC:
             raise ValueError("'order' must be '1' or '2' or '3'.")
         if skip_type == 'logSNR':
             # To reproduce the results in DPM-Solver paper
-            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
+            timesteps_outer = get_time_steps(self.noise_schedule, skip_type, t_T, t_0, K, device)
         else:
-            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[torch.cumsum(torch.tensor([0,] + orders), 0).to(device)]
+            timesteps_outer = get_time_steps(self.noise_schedule, skip_type, t_T, t_0, steps, device)[torch.cumsum(torch.tensor([0,] + orders), 0).to(device)]
         return timesteps_outer, orders
 
     def denoise_to_zero_fn(self, x, s):
@@ -748,15 +747,16 @@ class UniPC:
 
     def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',
         method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',
-        atol=0.0078, rtol=0.05, corrector=False,
+        atol=0.0078, rtol=0.05, corrector=False, timesteps=None,
     ):
         t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
         t_T = self.noise_schedule.T if t_start is None else t_start
         device = x.device
         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)
+            if timesteps == None:
+                timesteps = get_time_steps(self.noise_schedule, 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, order {order}")
+            assert steps >= order, "UniPC order must be < sampling steps"
             assert timesteps.shape[0] - 1 == steps
             with Progress(TextColumn('[cyan]{task.description}'), BarColumn(), TaskProgressColumn(), TimeRemainingColumn(), TimeElapsedColumn()) as progress:
                 task = progress.add_task(description="Initializing", total=steps)