Anonymous-123/ImageNet-Editing
1
1"""2This code started out as a PyTorch port of Ho et al's diffusion models:3https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py4 5Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules.6"""7 8import enum9import math10 11import numpy as np12import torch as th13 14from .nn import mean_flat15from .losses import normal_kl, discretized_gaussian_log_likelihood16 17import pdb18 19 20def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):21 """22 Get a pre-defined beta schedule for the given name.23 24 The beta schedule library consists of beta schedules which remain similar25 in the limit of num_diffusion_timesteps.26 Beta schedules may be added, but should not be removed or changed once27 they are committed to maintain backwards compatibility.28 """29 if schedule_name == "linear":30 # Linear schedule from Ho et al, extended to work for any number of31 # diffusion steps.32 scale = 1000 / num_diffusion_timesteps33 beta_start = scale * 0.000134 beta_end = scale * 0.0235 return np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)36 elif schedule_name == "cosine":37 return betas_for_alpha_bar(38 num_diffusion_timesteps, lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,39 )40 else:41 raise NotImplementedError(f"unknown beta schedule: {schedule_name}")42 43 44def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):45 """46 Create a beta schedule that discretizes the given alpha_t_bar function,47 which defines the cumulative product of (1-beta) over time from t = [0,1].48 49 :param num_diffusion_timesteps: the number of betas to produce.50 :param alpha_bar: a lambda that takes an argument t from 0 to 1 and51 produces the cumulative product of (1-beta) up to that52 part of the diffusion process.53 :param max_beta: the maximum beta to use; use values lower than 1 to54 prevent singularities.55 """56 betas = []57 for i in range(num_diffusion_timesteps):58 t1 = i / num_diffusion_timesteps59 t2 = (i + 1) / num_diffusion_timesteps60 betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))61 return np.array(betas)62 63 64class ModelMeanType(enum.Enum):65 """66 Which type of output the model predicts.67 """68 69 PREVIOUS_X = enum.auto() # the model predicts x_{t-1}70 START_X = enum.auto() # the model predicts x_071 EPSILON = enum.auto() # the model predicts epsilon72 73 74class ModelVarType(enum.Enum):75 """76 What is used as the model's output variance.77 78 The LEARNED_RANGE option has been added to allow the model to predict79 values between FIXED_SMALL and FIXED_LARGE, making its job easier.80 """81 82 LEARNED = enum.auto()83 FIXED_SMALL = enum.auto()84 FIXED_LARGE = enum.auto()85 LEARNED_RANGE = enum.auto()86 87 88class LossType(enum.Enum):89 MSE = enum.auto() # use raw MSE loss (and KL when learning variances)90 RESCALED_MSE = enum.auto() # use raw MSE loss (with RESCALED_KL when learning variances)91 KL = enum.auto() # use the variational lower-bound92 RESCALED_KL = enum.auto() # like KL, but rescale to estimate the full VLB93 94 def is_vb(self):95 return self == LossType.KL or self == LossType.RESCALED_KL96 97 98class GaussianDiffusion:99 """100 Utilities for training and sampling diffusion models.101 102 Ported directly from here, and then adapted over time to further experimentation.103 https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42104 105 :param betas: a 1-D numpy array of betas for each diffusion timestep,106 starting at T and going to 1.107 :param model_mean_type: a ModelMeanType determining what the model outputs.108 :param model_var_type: a ModelVarType determining how variance is output.109 :param loss_type: a LossType determining the loss function to use.110 :param rescale_timesteps: if True, pass floating point timesteps into the111 model so that they are always scaled like in the112 original paper (0 to 1000).113 """114 115 def __init__(116 self, *, betas, model_mean_type, model_var_type, loss_type, rescale_timesteps=False,117 ):118 self.model_mean_type = model_mean_type119 self.model_var_type = model_var_type120 self.loss_type = loss_type121 self.rescale_timesteps = rescale_timesteps122 123 # Use float64 for accuracy.124 betas = np.array(betas, dtype=np.float64)125 self.betas = betas126 assert len(betas.shape) == 1, "betas must be 1-D"127 assert (betas > 0).all() and (betas <= 1).all()128 129 self.num_timesteps = int(betas.shape[0])130 131 alphas = 1.0 - betas132 self.alphas_cumprod = np.cumprod(alphas, axis=0)133 self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])134 self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)135 assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)136 137 # calculations for diffusion q(x_t | x_{t-1}) and others138 self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)139 self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)140 self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)141 self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)142 self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)143 144 # calculations for posterior q(x_{t-1} | x_t, x_0)145 self.posterior_variance = (146 betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)147 )148 # log calculation clipped because the posterior variance is 0 at the149 # beginning of the diffusion chain.150 self.posterior_log_variance_clipped = np.log(151 np.append(self.posterior_variance[1], self.posterior_variance[1:])152 )153 self.posterior_mean_coef1 = (154 betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)155 )156 self.posterior_mean_coef2 = (157 (1.0 - self.alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - self.alphas_cumprod)158 )159 160 def q_mean_variance(self, x_start, t):161 """162 Get the distribution q(x_t | x_0).163 164 :param x_start: the [N x C x ...] tensor of noiseless inputs.165 :param t: the number of diffusion steps (minus 1). Here, 0 means one step.166 :return: A tuple (mean, variance, log_variance), all of x_start's shape.167 """168 mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start169 variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)170 log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)171 return mean, variance, log_variance172 173 def q_sample(self, x_start, t, noise=None):174 """175 Diffuse the data for a given number of diffusion steps.176 177 In other words, sample from q(x_t | x_0).178 179 :param x_start: the initial data batch.180 :param t: the number of diffusion steps (minus 1). Here, 0 means one step.181 :param noise: if specified, the split-out normal noise.182 :return: A noisy version of x_start.183 """184 if noise is None:185 noise = th.randn_like(x_start)186 assert noise.shape == x_start.shape187 return (188 _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start189 + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise190 )191 192 def q_posterior_mean_variance(self, x_start, x_t, t):193 """194 Compute the mean and variance of the diffusion posterior:195 196 q(x_{t-1} | x_t, x_0)197 198 """199 assert x_start.shape == x_t.shape200 posterior_mean = (201 _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start202 + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t203 )204 posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)205 posterior_log_variance_clipped = _extract_into_tensor(206 self.posterior_log_variance_clipped, t, x_t.shape207 )208 assert (209 posterior_mean.shape[0]210 == posterior_variance.shape[0]211 == posterior_log_variance_clipped.shape[0]212 == x_start.shape[0]213 )214 return posterior_mean, posterior_variance, posterior_log_variance_clipped215 216 def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None):217 """218 Apply the model to get p(x_{t-1} | x_t), as well as a prediction of219 the initial x, x_0.220 221 :param model: the model, which takes a signal and a batch of timesteps222 as input.223 :param x: the [N x C x ...] tensor at time t.224 :param t: a 1-D Tensor of timesteps.225 :param clip_denoised: if True, clip the denoised signal into [-1, 1].226 :param denoised_fn: if not None, a function which applies to the227 x_start prediction before it is used to sample. Applies before228 clip_denoised.229 :param model_kwargs: if not None, a dict of extra keyword arguments to230 pass to the model. This can be used for conditioning.231 :return: a dict with the following keys:232 - 'mean': the model mean output.233 - 'variance': the model variance output.234 - 'log_variance': the log of 'variance'.235 - 'pred_xstart': the prediction for x_0.236 """237 if model_kwargs is None:238 model_kwargs = {}239 240 B, C = x.shape[:2]241 assert t.shape == (B,)242 model_output = model(x, self._scale_timesteps(t), **model_kwargs)243 244 if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:245 assert model_output.shape == (B, C * 2, *x.shape[2:])246 model_output, model_var_values = th.split(model_output, C, dim=1)247 if self.model_var_type == ModelVarType.LEARNED:248 model_log_variance = model_var_values249 model_variance = th.exp(model_log_variance)250 else:251 min_log = _extract_into_tensor(self.posterior_log_variance_clipped, t, x.shape)252 max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)253 # The model_var_values is [-1, 1] for [min_var, max_var].254 frac = (model_var_values + 1) / 2255 model_log_variance = frac * max_log + (1 - frac) * min_log256 model_variance = th.exp(model_log_variance)257 else:258 model_variance, model_log_variance = {259 # for fixedlarge, we set the initial (log-)variance like so260 # to get a better decoder log likelihood.261 ModelVarType.FIXED_LARGE: (262 np.append(self.posterior_variance[1], self.betas[1:]),263 np.log(np.append(self.posterior_variance[1], self.betas[1:])),264 ),265 ModelVarType.FIXED_SMALL: (266 self.posterior_variance,267 self.posterior_log_variance_clipped,268 ),269 }[self.model_var_type]270 model_variance = _extract_into_tensor(model_variance, t, x.shape)271 model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)272 273 def process_xstart(x):274 if denoised_fn is not None:275 x = denoised_fn(x)276 if clip_denoised:277 return x.clamp(-1, 1)278 return x279 280 if self.model_mean_type == ModelMeanType.PREVIOUS_X:281 pred_xstart = process_xstart(282 self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)283 )284 model_mean = model_output285 elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:286 if self.model_mean_type == ModelMeanType.START_X:287 pred_xstart = process_xstart(model_output)288 else:289 pred_xstart = process_xstart(290 self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)291 )292 model_mean, _, _ = self.q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t)293 else:294 raise NotImplementedError(self.model_mean_type)295 296 assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape297 return {298 "mean": model_mean,299 "variance": model_variance,300 "log_variance": model_log_variance,301 "pred_xstart": pred_xstart,302 }303 304 def _predict_xstart_from_eps(self, x_t, t, eps):305 assert x_t.shape == eps.shape306 return (307 _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t308 - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps309 )310 311 def _predict_xstart_from_xprev(self, x_t, t, xprev):312 assert x_t.shape == xprev.shape313 return ( # (xprev - coef2*x_t) / coef1314 _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev315 - _extract_into_tensor(316 self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape317 )318 * x_t319 )320 321 def _predict_eps_from_xstart(self, x_t, t, pred_xstart):322 return (323 _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart324 ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)325 326 def _scale_timesteps(self, t):327 if self.rescale_timesteps:328 return t.float() * (1000.0 / self.num_timesteps)329 return t330 331 def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):332 """333 Compute the mean for the previous step, given a function cond_fn that334 computes the gradient of a conditional log probability with respect to335 x. In particular, cond_fn computes grad(log(p(y|x))), and we want to336 condition on y.337 338 This uses the conditioning strategy from Sohl-Dickstein et al. (2015).339 """340 gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs)341 new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()342 return new_mean343 344 def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):345 """346 Compute what the p_mean_variance output would have been, should the347 model's score function be conditioned by cond_fn.348 349 See condition_mean() for details on cond_fn.350 351 Unlike condition_mean(), this instead uses the conditioning strategy352 from Song et al (2020).353 """354 alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)355 356 eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])357 eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, self._scale_timesteps(t), **model_kwargs)358 359 out = p_mean_var.copy()360 out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)361 out["mean"], _, _ = self.q_posterior_mean_variance(x_start=out["pred_xstart"], x_t=x, t=t)362 return out363 364 def p_sample(365 self, model, x, t, clip_denoised=True, denoised_fn=None, cond_fn=None, model_kwargs=None,366 ):367 """368 Sample x_{t-1} from the model at the given timestep.369 370 :param model: the model to sample from.371 :param x: the current tensor at x_{t-1}.372 :param t: the value of t, starting at 0 for the first diffusion step.373 :param clip_denoised: if True, clip the x_start prediction to [-1, 1].374 :param denoised_fn: if not None, a function which applies to the375 x_start prediction before it is used to sample.376 :param cond_fn: if not None, this is a gradient function that acts377 similarly to the model.378 :param model_kwargs: if not None, a dict of extra keyword arguments to379 pass to the model. This can be used for conditioning.380 :return: a dict containing the following keys:381 - 'sample': a random sample from the model.382 - 'pred_xstart': a prediction of x_0.383 """384 out = self.p_mean_variance(385 model,386 x,387 t,388 clip_denoised=clip_denoised,389 denoised_fn=denoised_fn,390 model_kwargs=model_kwargs,391 )392 noise = th.randn_like(x)393 nonzero_mask = (394 (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))395 ) # no noise when t == 0396 if cond_fn is not None:397 out["mean"] = self.condition_mean(cond_fn, out, x, t, model_kwargs=model_kwargs)398 sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise399 return {"sample": sample, "pred_xstart": out["pred_xstart"]}400 401 def p_sample_loop(402 self,403 model,404 shape,405 noise=None,406 clip_denoised=True,407 denoised_fn=None,408 cond_fn=None,409 model_kwargs=None,410 device=None,411 progress=False,412 skip_timesteps=0,413 init_image=None,414 randomize_class=False,415 ):416 """417 Generate samples from the model.418 419 :param model: the model module.420 :param shape: the shape of the samples, (N, C, H, W).421 :param noise: if specified, the noise from the encoder to sample.422 Should be of the same shape as `shape`.423 :param clip_denoised: if True, clip x_start predictions to [-1, 1].424 :param denoised_fn: if not None, a function which applies to the425 x_start prediction before it is used to sample.426 :param cond_fn: if not None, this is a gradient function that acts427 similarly to the model.428 :param model_kwargs: if not None, a dict of extra keyword arguments to429 pass to the model. This can be used for conditioning.430 :param device: if specified, the device to create the samples on.431 If not specified, use a model parameter's device.432 :param progress: if True, show a tqdm progress bar.433 :return: a non-differentiable batch of samples.434 """435 final = None436 for sample in self.p_sample_loop_progressive(437 model,438 shape,439 noise=noise,440 clip_denoised=clip_denoised,441 denoised_fn=denoised_fn,442 cond_fn=cond_fn,443 model_kwargs=model_kwargs,444 device=device,445 progress=progress,446 skip_timesteps=skip_timesteps,447 init_image=init_image,448 randomize_class=randomize_class,449 ):450 final = sample451 return final["sample"]452 453 def p_sample_loop_progressive(454 self,455 model,456 shape,457 noise=None,458 clip_denoised=True,459 denoised_fn=None,460 cond_fn=None,461 model_kwargs=None,462 device=None,463 progress=False,464 skip_timesteps=0,465 init_image=None,466 postprocess_fn=None,467 randomize_class=False,468 ):469 """470 Generate samples from the model and yield intermediate samples from471 each timestep of diffusion.472 473 Arguments are the same as p_sample_loop().474 Returns a generator over dicts, where each dict is the return value of475 p_sample().476 """477 # if device is None:478 # device = next(model.parameters()).device479 assert isinstance(shape, (tuple, list))480 if noise is not None:481 img = noise482 '''483 img_guidance = noise.to(device)484 t_batch = th.tensor([int(t0*self.num_timesteps)-1]*len(img_guidance), device=device)485 img = self.q_sample(img_guidance, t_batch)486 indices = list(range(int(t0*self.num_timesteps)))[::-1]487 '''488 else:489 img = th.randn(*shape, device=device)490 491 # pdb.set_trace()492 if skip_timesteps and init_image is None:493 init_image = th.zeros_like(img)494 495 indices = list(range(self.num_timesteps - skip_timesteps))[::-1]496 497 batch_size = shape[0]498 init_image_batch = th.tile(init_image, dims=(batch_size, 1, 1, 1))499 img = self.q_sample(500 x_start=init_image_batch,501 t=th.tensor(indices[0], dtype=th.long, device=device),502 noise=img,503 )504 505 if progress:506 # Lazy import so that we don't depend on tqdm.507 from tqdm.auto import tqdm508 509 indices = tqdm(indices)510 511 for i in indices:512 t = th.tensor([i] * shape[0], device=device)513 if randomize_class and "y" in model_kwargs:514 model_kwargs["y"] = th.randint(515 low=0,516 high=model.num_classes,517 size=model_kwargs["y"].shape,518 device=model_kwargs["y"].device,519 )520 with th.no_grad():521 out = self.p_sample(522 model,523 img,524 t,525 clip_denoised=clip_denoised,526 denoised_fn=denoised_fn,527 cond_fn=cond_fn,528 model_kwargs=model_kwargs,529 )530 if postprocess_fn is not None:531 out = postprocess_fn(out, t)532 533 yield out534 img = out["sample"]535 536 def ddim_sample(537 self,538 model,539 x,540 t,541 clip_denoised=True,542 denoised_fn=None,543 cond_fn=None,544 model_kwargs=None,545 eta=0.0,546 ):547 """548 Sample x_{t-1} from the model using DDIM.549 550 Same usage as p_sample().551 """552 out = self.p_mean_variance(553 model,554 x,555 t,556 clip_denoised=clip_denoised,557 denoised_fn=denoised_fn,558 model_kwargs=model_kwargs,559 )560 if cond_fn is not None:561 out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)562 563 # Usually our model outputs epsilon, but we re-derive it564 # in case we used x_start or x_prev prediction.565 eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])566 567 alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)568 alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)569 sigma = (570 eta571 * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))572 * th.sqrt(1 - alpha_bar / alpha_bar_prev)573 )574 # Equation 12.575 noise = th.randn_like(x)576 mean_pred = (577 out["pred_xstart"] * th.sqrt(alpha_bar_prev)578 + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps579 )580 nonzero_mask = (581 (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))582 ) # no noise when t == 0583 sample = mean_pred + nonzero_mask * sigma * noise584 return {"sample": sample, "pred_xstart": out["pred_xstart"]}585 586 def ddim_reverse_sample(587 self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None, eta=0.0,588 ):589 """590 Sample x_{t+1} from the model using DDIM reverse ODE.591 """592 assert eta == 0.0, "Reverse ODE only for deterministic path"593 out = self.p_mean_variance(594 model,595 x,596 t,597 clip_denoised=clip_denoised,598 denoised_fn=denoised_fn,599 model_kwargs=model_kwargs,600 )601 # Usually our model outputs epsilon, but we re-derive it602 # in case we used x_start or x_prev prediction.603 eps = (604 _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x605 - out["pred_xstart"]606 ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)607 alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)608 609 # Equation 12. reversed610 mean_pred = out["pred_xstart"] * th.sqrt(alpha_bar_next) + th.sqrt(1 - alpha_bar_next) * eps611 612 return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}613 614 def ddim_sample_loop(615 self,616 model,617 shape,618 noise=None,619 clip_denoised=True,620 denoised_fn=None,621 cond_fn=None,622 model_kwargs=None,623 device=None,624 progress=False,625 eta=0.0,626 skip_timesteps=0,627 init_image=None,628 randomize_class=False,629 ):630 """631 Generate samples from the model using DDIM.632 633 Same usage as p_sample_loop().634 """635 final = None636 for sample in self.ddim_sample_loop_progressive(637 model,638 shape,639 noise=noise,640 clip_denoised=clip_denoised,641 denoised_fn=denoised_fn,642 cond_fn=cond_fn,643 model_kwargs=model_kwargs,644 device=device,645 progress=progress,646 eta=eta,647 skip_timesteps=skip_timesteps,648 init_image=init_image,649 randomize_class=randomize_class,650 ):651 final = sample652 return final["sample"]653 654 def ddim_sample_loop_progressive(655 self,656 model,657 shape,658 noise=None,659 clip_denoised=True,660 denoised_fn=None,661 cond_fn=None,662 model_kwargs=None,663 device=None,664 progress=False,665 eta=0.0,666 skip_timesteps=0,667 init_image=None,668 postprocess_fn=None,669 randomize_class=False,670 ):671 """672 Use DDIM to sample from the model and yield intermediate samples from673 each timestep of DDIM.674 675 Same usage as p_sample_loop_progressive().676 """677 if device is None:678 device = next(model.parameters()).device679 assert isinstance(shape, (tuple, list))680 if noise is not None:681 img = noise682 else:683 img = th.randn(*shape, device=device)684 685 if skip_timesteps and init_image is None:686 init_image = th.zeros_like(img)687 688 indices = list(range(self.num_timesteps - skip_timesteps))[::-1]689 690 if init_image is not None:691 my_t = th.ones([shape[0]], device=device, dtype=th.long) * indices[0]692 batch_size = shape[0]693 init_image_batch = th.tile(init_image, dims=(batch_size, 1, 1, 1))694 img = self.q_sample(init_image_batch, my_t, img)695 696 if progress:697 # Lazy import so that we don't depend on tqdm.698 from tqdm.auto import tqdm699 700 indices = tqdm(indices)701 702 for i in indices:703 t = th.tensor([i] * shape[0], device=device)704 if randomize_class and "y" in model_kwargs:705 model_kwargs["y"] = th.randint(706 low=0,707 high=model.num_classes,708 size=model_kwargs["y"].shape,709 device=model_kwargs["y"].device,710 )711 with th.no_grad():712 out = self.ddim_sample(713 model,714 img,715 t,716 clip_denoised=clip_denoised,717 denoised_fn=denoised_fn,718 cond_fn=cond_fn,719 model_kwargs=model_kwargs,720 eta=eta,721 )722 723 if postprocess_fn is not None:724 out = postprocess_fn(out, t)725 726 yield out727 img = out["sample"]728 729 def _vb_terms_bpd(self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None):730 """731 Get a term for the variational lower-bound.732 733 The resulting units are bits (rather than nats, as one might expect).734 This allows for comparison to other papers.735 736 :return: a dict with the following keys:737 - 'output': a shape [N] tensor of NLLs or KLs.738 - 'pred_xstart': the x_0 predictions.739 """740 true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(741 x_start=x_start, x_t=x_t, t=t742 )743 out = self.p_mean_variance(744 model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs745 )746 kl = normal_kl(true_mean, true_log_variance_clipped, out["mean"], out["log_variance"])747 kl = mean_flat(kl) / np.log(2.0)748 749 decoder_nll = -discretized_gaussian_log_likelihood(750 x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]751 )752 assert decoder_nll.shape == x_start.shape753 decoder_nll = mean_flat(decoder_nll) / np.log(2.0)754 755 # At the first timestep return the decoder NLL,756 # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))757 output = th.where((t == 0), decoder_nll, kl)758 return {"output": output, "pred_xstart": out["pred_xstart"]}759 760 def training_losses(self, model, x_start, t, model_kwargs=None, noise=None):761 """762 Compute training losses for a single timestep.763 764 :param model: the model to evaluate loss on.765 :param x_start: the [N x C x ...] tensor of inputs.766 :param t: a batch of timestep indices.767 :param model_kwargs: if not None, a dict of extra keyword arguments to768 pass to the model. This can be used for conditioning.769 :param noise: if specified, the specific Gaussian noise to try to remove.770 :return: a dict with the key "loss" containing a tensor of shape [N].771 Some mean or variance settings may also have other keys.772 """773 if model_kwargs is None:774 model_kwargs = {}775 if noise is None:776 noise = th.randn_like(x_start)777 x_t = self.q_sample(x_start, t, noise=noise)778 779 terms = {}780 781 if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:782 terms["loss"] = self._vb_terms_bpd(783 model=model,784 x_start=x_start,785 x_t=x_t,786 t=t,787 clip_denoised=False,788 model_kwargs=model_kwargs,789 )["output"]790 if self.loss_type == LossType.RESCALED_KL:791 terms["loss"] *= self.num_timesteps792 elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:793 model_output = model(x_t, self._scale_timesteps(t), **model_kwargs)794 795 if self.model_var_type in [796 ModelVarType.LEARNED,797 ModelVarType.LEARNED_RANGE,798 ]:799 B, C = x_t.shape[:2]800 assert model_output.shape == (B, C * 2, *x_t.shape[2:])801 model_output, model_var_values = th.split(model_output, C, dim=1)802 # Learn the variance using the variational bound, but don't let803 # it affect our mean prediction.804 frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)805 terms["vb"] = self._vb_terms_bpd(806 model=lambda *args, r=frozen_out: r,807 x_start=x_start,808 x_t=x_t,809 t=t,810 clip_denoised=False,811 )["output"]812 if self.loss_type == LossType.RESCALED_MSE:813 # Divide by 1000 for equivalence with initial implementation.814 # Without a factor of 1/1000, the VB term hurts the MSE term.815 terms["vb"] *= self.num_timesteps / 1000.0816 817 target = {818 ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(819 x_start=x_start, x_t=x_t, t=t820 )[0],821 ModelMeanType.START_X: x_start,822 ModelMeanType.EPSILON: noise,823 }[self.model_mean_type]824 assert model_output.shape == target.shape == x_start.shape825 terms["mse"] = mean_flat((target - model_output) ** 2)826 if "vb" in terms:827 terms["loss"] = terms["mse"] + terms["vb"]828 else:829 terms["loss"] = terms["mse"]830 else:831 raise NotImplementedError(self.loss_type)832 833 return terms834 835 def _prior_bpd(self, x_start):836 """837 Get the prior KL term for the variational lower-bound, measured in838 bits-per-dim.839 840 This term can't be optimized, as it only depends on the encoder.841 842 :param x_start: the [N x C x ...] tensor of inputs.843 :return: a batch of [N] KL values (in bits), one per batch element.844 """845 batch_size = x_start.shape[0]846 t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)847 qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)848 kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)849 return mean_flat(kl_prior) / np.log(2.0)850 851 def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):852 """853 Compute the entire variational lower-bound, measured in bits-per-dim,854 as well as other related quantities.855 856 :param model: the model to evaluate loss on.857 :param x_start: the [N x C x ...] tensor of inputs.858 :param clip_denoised: if True, clip denoised samples.859 :param model_kwargs: if not None, a dict of extra keyword arguments to860 pass to the model. This can be used for conditioning.861 862 :return: a dict containing the following keys:863 - total_bpd: the total variational lower-bound, per batch element.864 - prior_bpd: the prior term in the lower-bound.865 - vb: an [N x T] tensor of terms in the lower-bound.866 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.867 - mse: an [N x T] tensor of epsilon MSEs for each timestep.868 """869 device = x_start.device870 batch_size = x_start.shape[0]871 872 vb = []873 xstart_mse = []874 mse = []875 for t in list(range(self.num_timesteps))[::-1]:876 t_batch = th.tensor([t] * batch_size, device=device)877 noise = th.randn_like(x_start)878 x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)879 # Calculate VLB term at the current timestep880 with th.no_grad():881 out = self._vb_terms_bpd(882 model,883 x_start=x_start,884 x_t=x_t,885 t=t_batch,886 clip_denoised=clip_denoised,887 model_kwargs=model_kwargs,888 )889 vb.append(out["output"])890 xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))891 eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])892 mse.append(mean_flat((eps - noise) ** 2))893 894 vb = th.stack(vb, dim=1)895 xstart_mse = th.stack(xstart_mse, dim=1)896 mse = th.stack(mse, dim=1)897 898 prior_bpd = self._prior_bpd(x_start)899 total_bpd = vb.sum(dim=1) + prior_bpd900 return {901 "total_bpd": total_bpd,902 "prior_bpd": prior_bpd,903 "vb": vb,904 "xstart_mse": xstart_mse,905 "mse": mse,906 }907 908 909def _extract_into_tensor(arr, timesteps, broadcast_shape):910 """911 Extract values from a 1-D numpy array for a batch of indices.912 913 :param arr: the 1-D numpy array.914 :param timesteps: a tensor of indices into the array to extract.915 :param broadcast_shape: a larger shape of K dimensions with the batch916 dimension equal to the length of timesteps.917 :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.918 """919 res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()920 while len(res.shape) < len(broadcast_shape):921 res = res[..., None]922 return res.expand(broadcast_shape)923