"""SAMPLING ONLY.""" import torch from .uni_pc import NoiseScheduleVP, model_wrapper, UniPC, get_time_steps from modules import shared, devices class UniPCSampler(object): def __init__(self, model, **kwargs): super().__init__() self.model = model to_torch = lambda x: x.clone().detach().to(torch.float32).to(model.device) self.before_sample = None 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 @devices.inference_context() 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 / self.inflated_steps self.denoise_steps = max(num_inference_steps, shared.opts.schedulers_solver_order) max(self.inflated_steps - self.denoise_steps, 0) # actual number of steps we'll run 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, ) # 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 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): # 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), self.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, 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, ) return self.uni_pc.sample( x_latent, steps=self.denoise_steps, skip_type=shared.opts.uni_pc_skip_type, method="multistep", order=shared.opts.schedulers_solver_order, lower_order_final=shared.opts.schedulers_use_loworder, denoise_to_zero=True, timesteps=self.timesteps, ) def register_buffer(self, name, attr): if type(attr) == torch.Tensor: if attr.device != devices.device: attr = attr.to(devices.device) setattr(self, name, attr) def set_hooks(self, before_sample, after_sample, after_update): self.before_sample = before_sample self.after_sample = after_sample self.after_update = after_update @devices.inference_context() 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, # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ... **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: shared.log.warning(f"UniPC: got {cbs} conditionings but batch-size is {batch_size}") elif isinstance(conditioning, list): for ctmp in conditioning: if ctmp.shape[0] != batch_size: shared.log.warning(f"UniPC: Got {cbs} conditionings but batch-size is {batch_size}") else: if conditioning.shape[0] != batch_size: shared.log.warning(f"UniPC: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}") # sampling 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 # 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), self.noise_schedule, model_type=model_type, guidance_type="classifier-free", #condition=conditioning, #unconditional_condition=unconditional_conditioning, guidance_scale=unconditional_guidance_scale, ) 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.schedulers_solver_order, lower_order_final=shared.opts.schedulers_use_loworder) return x.to(device), None