diff --git a/modules/models/diffusion/uni_pc/sampler.py b/modules/models/diffusion/uni_pc/sampler.py index 953e786db..e3f8a6651 100644 --- a/modules/models/diffusion/uni_pc/sampler.py +++ b/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 diff --git a/modules/models/diffusion/uni_pc/uni_pc.py b/modules/models/diffusion/uni_pc/uni_pc.py index 895fc58c3..d86572f69 100644 --- a/modules/models/diffusion/uni_pc/uni_pc.py +++ b/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)