Shamima/code-search
0
1function_name,docstring,function_body,file_path2diffusion_from_config,,"def diffusion_from_config(config: Dict[str, Any]) ->GaussianDiffusion:3 schedule = config['schedule']4 steps = config['timesteps']5 respace = config.get('respacing', None)6 mean_type = config.get('mean_type', 'epsilon')7 betas = get_named_beta_schedule(schedule, steps)8 channel_scales = config.get('channel_scales', None)9 channel_biases = config.get('channel_biases', None)10 if channel_scales is not None:11 channel_scales = np.array(channel_scales)12 if channel_biases is not None:13 channel_biases = np.array(channel_biases)14 kwargs = dict(betas=betas, model_mean_type=mean_type, model_var_type=15 'learned_range', loss_type='mse', channel_scales=channel_scales,16 channel_biases=channel_biases)17 if respace is None:18 return GaussianDiffusion(**kwargs)19 else:20 return SpacedDiffusion(use_timesteps=space_timesteps(steps, respace21 ), **kwargs)22",point_e\diffusion\configs.py23get_beta_schedule,"This is the deprecated API for creating beta schedules.24 25See get_named_beta_schedule() for the new library of schedules.","def get_beta_schedule(beta_schedule, *, beta_start, beta_end,26 num_diffusion_timesteps):27 """"""""""""28 if beta_schedule == 'linear':29 betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps,30 dtype=np.float64)31 else:32 raise NotImplementedError(beta_schedule)33 assert betas.shape == (num_diffusion_timesteps,)34 return betas35",point_e\diffusion\gaussian_diffusion.py36get_named_beta_schedule,"Get a pre-defined beta schedule for the given name.37 38The beta schedule library consists of beta schedules which remain similar39in the limit of num_diffusion_timesteps.40Beta schedules may be added, but should not be removed or changed once41they are committed to maintain backwards compatibility.","def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):42 """"""""""""43 if schedule_name == 'linear':44 scale = 1000 / num_diffusion_timesteps45 return get_beta_schedule('linear', beta_start=scale * 0.0001,46 beta_end=scale * 0.02, num_diffusion_timesteps=47 num_diffusion_timesteps)48 elif schedule_name == 'cosine':49 return betas_for_alpha_bar(num_diffusion_timesteps, lambda t: math.50 cos((t + 0.008) / 1.008 * math.pi / 2) ** 2)51 else:52 raise NotImplementedError(f'unknown beta schedule: {schedule_name}')53",point_e\diffusion\gaussian_diffusion.py54betas_for_alpha_bar,"Create a beta schedule that discretizes the given alpha_t_bar function,55which defines the cumulative product of (1-beta) over time from t = [0,1].56 57:param num_diffusion_timesteps: the number of betas to produce.58:param alpha_bar: a lambda that takes an argument t from 0 to 1 and59 produces the cumulative product of (1-beta) up to that60 part of the diffusion process.61:param max_beta: the maximum beta to use; use values lower than 1 to62 prevent singularities.","def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):63 """"""""""""64 betas = []65 for i in range(num_diffusion_timesteps):66 t1 = i / num_diffusion_timesteps67 t2 = (i + 1) / num_diffusion_timesteps68 betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))69 return np.array(betas)70",point_e\diffusion\gaussian_diffusion.py71space_timesteps,"Create a list of timesteps to use from an original diffusion process,72given the number of timesteps we want to take from equally-sized portions73of the original process.74For example, if there's 300 timesteps and the section counts are [10,15,20]75then the first 100 timesteps are strided to be 10 timesteps, the second 10076are strided to be 15 timesteps, and the final 100 are strided to be 20.77:param num_timesteps: the number of diffusion steps in the original78 process to divide up.79:param section_counts: either a list of numbers, or a string containing80 comma-separated numbers, indicating the step count81 per section. As a special case, use ""ddimN"" where N82 is a number of steps to use the striding from the83 DDIM paper.84:return: a set of diffusion steps from the original process to use.","def space_timesteps(num_timesteps, section_counts):85 """"""""""""86 if isinstance(section_counts, str):87 if section_counts.startswith('ddim'):88 desired_count = int(section_counts[len('ddim'):])89 for i in range(1, num_timesteps):90 if len(range(0, num_timesteps, i)) == desired_count:91 return set(range(0, num_timesteps, i))92 raise ValueError(93 f'cannot create exactly {num_timesteps} steps with an integer stride'94 )95 elif section_counts.startswith('exact'):96 res = set(int(x) for x in section_counts[len('exact'):].split(','))97 for x in res:98 if x < 0 or x >= num_timesteps:99 raise ValueError(f'timestep out of bounds: {x}')100 return res101 section_counts = [int(x) for x in section_counts.split(',')]102 size_per = num_timesteps // len(section_counts)103 extra = num_timesteps % len(section_counts)104 start_idx = 0105 all_steps = []106 for i, section_count in enumerate(section_counts):107 size = size_per + (1 if i < extra else 0)108 if size < section_count:109 raise ValueError(110 f'cannot divide section of {size} steps into {section_count}')111 if section_count <= 1:112 frac_stride = 1113 else:114 frac_stride = (size - 1) / (section_count - 1)115 cur_idx = 0.0116 taken_steps = []117 for _ in range(section_count):118 taken_steps.append(start_idx + round(cur_idx))119 cur_idx += frac_stride120 all_steps += taken_steps121 start_idx += size122 return set(all_steps)123",point_e\diffusion\gaussian_diffusion.py124_extract_into_tensor,"Extract values from a 1-D numpy array for a batch of indices.125 126:param arr: the 1-D numpy array.127:param timesteps: a tensor of indices into the array to extract.128:param broadcast_shape: a larger shape of K dimensions with the batch129 dimension equal to the length of timesteps.130:return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.","def _extract_into_tensor(arr, timesteps, broadcast_shape):131 """"""""""""132 res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()133 while len(res.shape) < len(broadcast_shape):134 res = res[..., None]135 return res + th.zeros(broadcast_shape, device=timesteps.device)136",point_e\diffusion\gaussian_diffusion.py137normal_kl,"Compute the KL divergence between two gaussians.138Shapes are automatically broadcasted, so batches can be compared to139scalars, among other use cases.","def normal_kl(mean1, logvar1, mean2, logvar2):140 """"""""""""141 tensor = None142 for obj in (mean1, logvar1, mean2, logvar2):143 if isinstance(obj, th.Tensor):144 tensor = obj145 break146 assert tensor is not None, 'at least one argument must be a Tensor'147 logvar1, logvar2 = [(x if isinstance(x, th.Tensor) else th.tensor(x).to148 (tensor)) for x in (logvar1, logvar2)]149 return 0.5 * (-1.0 + logvar2 - logvar1 + th.exp(logvar1 - logvar2) + (150 mean1 - mean2) ** 2 * th.exp(-logvar2))151",point_e\diffusion\gaussian_diffusion.py152approx_standard_normal_cdf,"A fast approximation of the cumulative distribution function of the153standard normal.","def approx_standard_normal_cdf(x):154 """"""""""""155 return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.156 pow(x, 3))))157",point_e\diffusion\gaussian_diffusion.py158discretized_gaussian_log_likelihood,"Compute the log-likelihood of a Gaussian distribution discretizing to a159given image.160:param x: the target images. It is assumed that this was uint8 values,161 rescaled to the range [-1, 1].162:param means: the Gaussian mean Tensor.163:param log_scales: the Gaussian log stddev Tensor.164:return: a tensor like x of log probabilities (in nats).","def discretized_gaussian_log_likelihood(x, *, means, log_scales):165 """"""""""""166 assert x.shape == means.shape == log_scales.shape167 centered_x = x - means168 inv_stdv = th.exp(-log_scales)169 plus_in = inv_stdv * (centered_x + 1.0 / 255.0)170 cdf_plus = approx_standard_normal_cdf(plus_in)171 min_in = inv_stdv * (centered_x - 1.0 / 255.0)172 cdf_min = approx_standard_normal_cdf(min_in)173 log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))174 log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))175 cdf_delta = cdf_plus - cdf_min176 log_probs = th.where(x < -0.999, log_cdf_plus, th.where(x > 0.999,177 log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))))178 assert log_probs.shape == x.shape179 return log_probs180",point_e\diffusion\gaussian_diffusion.py181mean_flat,Take the mean over all non-batch dimensions.,"def mean_flat(tensor):182 """"""""""""183 return tensor.flatten(1).mean(1)184",point_e\diffusion\gaussian_diffusion.py185__init__,,"def __init__(self, *, betas: Sequence[float], model_mean_type: str,186 model_var_type: str, loss_type: str, discretized_t0: bool=False,187 channel_scales: Optional[np.ndarray]=None, channel_biases: Optional[np.188 ndarray]=None):189 self.model_mean_type = model_mean_type190 self.model_var_type = model_var_type191 self.loss_type = loss_type192 self.discretized_t0 = discretized_t0193 self.channel_scales = channel_scales194 self.channel_biases = channel_biases195 betas = np.array(betas, dtype=np.float64)196 self.betas = betas197 assert len(betas.shape) == 1, 'betas must be 1-D'198 assert (betas > 0).all() and (betas <= 1).all()199 self.num_timesteps = int(betas.shape[0])200 alphas = 1.0 - betas201 self.alphas_cumprod = np.cumprod(alphas, axis=0)202 self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])203 self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)204 assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)205 self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)206 self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)207 self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)208 self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)209 self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)210 self.posterior_variance = betas * (1.0 - self.alphas_cumprod_prev) / (211 1.0 - self.alphas_cumprod)212 self.posterior_log_variance_clipped = np.log(np.append(self.213 posterior_variance[1], self.posterior_variance[1:]))214 self.posterior_mean_coef1 = betas * np.sqrt(self.alphas_cumprod_prev) / (215 1.0 - self.alphas_cumprod)216 self.posterior_mean_coef2 = (1.0 - self.alphas_cumprod_prev) * np.sqrt(217 alphas) / (1.0 - self.alphas_cumprod)218",point_e\diffusion\gaussian_diffusion.py219get_sigmas,,"def get_sigmas(self, t):220 return _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, t.shape)221",point_e\diffusion\gaussian_diffusion.py222q_mean_variance,"Get the distribution q(x_t | x_0).223 224:param x_start: the [N x C x ...] tensor of noiseless inputs.225:param t: the number of diffusion steps (minus 1). Here, 0 means one step.226:return: A tuple (mean, variance, log_variance), all of x_start's shape.","def q_mean_variance(self, x_start, t):227 """"""""""""228 mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape229 ) * x_start230 variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape231 )232 log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod,233 t, x_start.shape)234 return mean, variance, log_variance235",point_e\diffusion\gaussian_diffusion.py236q_sample,"Diffuse the data for a given number of diffusion steps.237 238In other words, sample from q(x_t | x_0).239 240:param x_start: the initial data batch.241:param t: the number of diffusion steps (minus 1). Here, 0 means one step.242:param noise: if specified, the split-out normal noise.243:return: A noisy version of x_start.","def q_sample(self, x_start, t, noise=None):244 """"""""""""245 if noise is None:246 noise = th.randn_like(x_start)247 assert noise.shape == x_start.shape248 return _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape249 ) * x_start + _extract_into_tensor(self.250 sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise251",point_e\diffusion\gaussian_diffusion.py252q_posterior_mean_variance,"Compute the mean and variance of the diffusion posterior:253 254 q(x_{t-1} | x_t, x_0)","def q_posterior_mean_variance(self, x_start, x_t, t):255 """"""""""""256 assert x_start.shape == x_t.shape257 posterior_mean = _extract_into_tensor(self.posterior_mean_coef1, t, x_t258 .shape) * x_start + _extract_into_tensor(self.posterior_mean_coef2,259 t, x_t.shape) * x_t260 posterior_variance = _extract_into_tensor(self.posterior_variance, t,261 x_t.shape)262 posterior_log_variance_clipped = _extract_into_tensor(self.263 posterior_log_variance_clipped, t, x_t.shape)264 assert posterior_mean.shape[0] == posterior_variance.shape[0265 ] == posterior_log_variance_clipped.shape[0] == x_start.shape[0]266 return posterior_mean, posterior_variance, posterior_log_variance_clipped267",point_e\diffusion\gaussian_diffusion.py268p_mean_variance,"Apply the model to get p(x_{t-1} | x_t), as well as a prediction of269the initial x, x_0.270 271:param model: the model, which takes a signal and a batch of timesteps272 as input.273:param x: the [N x C x ...] tensor at time t.274:param t: a 1-D Tensor of timesteps.275:param clip_denoised: if True, clip the denoised signal into [-1, 1].276:param denoised_fn: if not None, a function which applies to the277 x_start prediction before it is used to sample. Applies before278 clip_denoised.279:param model_kwargs: if not None, a dict of extra keyword arguments to280 pass to the model. This can be used for conditioning.281:return: a dict with the following keys:282 - 'mean': the model mean output.283 - 'variance': the model variance output.284 - 'log_variance': the log of 'variance'.285 - 'pred_xstart': the prediction for x_0.","def p_mean_variance(self, model, x, t, clip_denoised=False, denoised_fn=286 None, model_kwargs=None):287 """"""""""""288 if model_kwargs is None:289 model_kwargs = {}290 B, C = x.shape[:2]291 assert t.shape == (B,)292 model_output = model(x, t, **model_kwargs)293 if isinstance(model_output, tuple):294 model_output, extra = model_output295 else:296 extra = None297 if self.model_var_type in ['learned', 'learned_range']:298 assert model_output.shape == (B, C * 2, *x.shape[2:])299 model_output, model_var_values = th.split(model_output, C, dim=1)300 if self.model_var_type == 'learned':301 model_log_variance = model_var_values302 model_variance = th.exp(model_log_variance)303 else:304 min_log = _extract_into_tensor(self.305 posterior_log_variance_clipped, t, x.shape)306 max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)307 frac = (model_var_values + 1) / 2308 model_log_variance = frac * max_log + (1 - frac) * min_log309 model_variance = th.exp(model_log_variance)310 else:311 model_variance, model_log_variance = {'fixed_large': (np.append(312 self.posterior_variance[1], self.betas[1:]), np.log(np.append(313 self.posterior_variance[1], self.betas[1:]))), 'fixed_small': (314 self.posterior_variance, self.posterior_log_variance_clipped)}[self315 .model_var_type]316 model_variance = _extract_into_tensor(model_variance, t, x.shape)317 model_log_variance = _extract_into_tensor(model_log_variance, t, x.318 shape)319 320 def process_xstart(x):321 if denoised_fn is not None:322 x = denoised_fn(x)323 if clip_denoised:324 return x.clamp(-1, 1)325 return x326 if self.model_mean_type == 'x_prev':327 pred_xstart = process_xstart(self._predict_xstart_from_xprev(x_t=x,328 t=t, xprev=model_output))329 model_mean = model_output330 elif self.model_mean_type in ['x_start', 'epsilon']:331 if self.model_mean_type == 'x_start':332 pred_xstart = process_xstart(model_output)333 else:334 pred_xstart = process_xstart(self._predict_xstart_from_eps(x_t=335 x, t=t, eps=model_output))336 model_mean, _, _ = self.q_posterior_mean_variance(x_start=337 pred_xstart, x_t=x, t=t)338 else:339 raise NotImplementedError(self.model_mean_type)340 assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape341 return {'mean': model_mean, 'variance': model_variance, 'log_variance':342 model_log_variance, 'pred_xstart': pred_xstart, 'extra': extra}343",point_e\diffusion\gaussian_diffusion.py344_predict_xstart_from_eps,,"def _predict_xstart_from_eps(self, x_t, t, eps):345 assert x_t.shape == eps.shape346 return _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape347 ) * x_t - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t,348 x_t.shape) * eps349",point_e\diffusion\gaussian_diffusion.py350_predict_xstart_from_xprev,,"def _predict_xstart_from_xprev(self, x_t, t, xprev):351 assert x_t.shape == xprev.shape352 return _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape353 ) * xprev - _extract_into_tensor(self.posterior_mean_coef2 / self.354 posterior_mean_coef1, t, x_t.shape) * x_t355",point_e\diffusion\gaussian_diffusion.py356_predict_eps_from_xstart,,"def _predict_eps_from_xstart(self, x_t, t, pred_xstart):357 return (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.358 shape) * x_t - pred_xstart) / _extract_into_tensor(self.359 sqrt_recipm1_alphas_cumprod, t, x_t.shape)360",point_e\diffusion\gaussian_diffusion.py361condition_mean,"Compute the mean for the previous step, given a function cond_fn that362computes the gradient of a conditional log probability with respect to363x. In particular, cond_fn computes grad(log(p(y|x))), and we want to364condition on y.365 366This uses the conditioning strategy from Sohl-Dickstein et al. (2015).","def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):367 """"""""""""368 gradient = cond_fn(x, t, **model_kwargs)369 new_mean = p_mean_var['mean'].float() + p_mean_var['variance'370 ] * gradient.float()371 return new_mean372",point_e\diffusion\gaussian_diffusion.py373condition_score,"Compute what the p_mean_variance output would have been, should the374model's score function be conditioned by cond_fn.375 376See condition_mean() for details on cond_fn.377 378Unlike condition_mean(), this instead uses the conditioning strategy379from Song et al (2020).","def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):380 """"""""""""381 alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)382 eps = self._predict_eps_from_xstart(x, t, p_mean_var['pred_xstart'])383 eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs)384 out = p_mean_var.copy()385 out['pred_xstart'] = self._predict_xstart_from_eps(x, t, eps)386 out['mean'], _, _ = self.q_posterior_mean_variance(x_start=out[387 'pred_xstart'], x_t=x, t=t)388 return out389",point_e\diffusion\gaussian_diffusion.py390p_sample,"Sample x_{t-1} from the model at the given timestep.391 392:param model: the model to sample from.393:param x: the current tensor at x_{t-1}.394:param t: the value of t, starting at 0 for the first diffusion step.395:param clip_denoised: if True, clip the x_start prediction to [-1, 1].396:param denoised_fn: if not None, a function which applies to the397 x_start prediction before it is used to sample.398:param cond_fn: if not None, this is a gradient function that acts399 similarly to the model.400:param model_kwargs: if not None, a dict of extra keyword arguments to401 pass to the model. This can be used for conditioning.402:return: a dict containing the following keys:403 - 'sample': a random sample from the model.404 - 'pred_xstart': a prediction of x_0.","def p_sample(self, model, x, t, clip_denoised=False, denoised_fn=None,405 cond_fn=None, model_kwargs=None):406 """"""""""""407 out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,408 denoised_fn=denoised_fn, model_kwargs=model_kwargs)409 noise = th.randn_like(x)410 nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))411 if cond_fn is not None:412 out['mean'] = self.condition_mean(cond_fn, out, x, t, model_kwargs=413 model_kwargs)414 sample = out['mean'] + nonzero_mask * th.exp(0.5 * out['log_variance']415 ) * noise416 return {'sample': sample, 'pred_xstart': out['pred_xstart']}417",point_e\diffusion\gaussian_diffusion.py418p_sample_loop,"Generate samples from the model.419 420:param model: the model module.421:param shape: the shape of the samples, (N, C, H, W).422:param noise: if specified, the noise from the encoder to sample.423 Should be of the same shape as `shape`.424:param clip_denoised: if True, clip x_start predictions to [-1, 1].425:param denoised_fn: if not None, a function which applies to the426 x_start prediction before it is used to sample.427:param cond_fn: if not None, this is a gradient function that acts428 similarly to the model.429:param model_kwargs: if not None, a dict of extra keyword arguments to430 pass to the model. This can be used for conditioning.431:param device: if specified, the device to create the samples on.432 If not specified, use a model parameter's device.433:param progress: if True, show a tqdm progress bar.434:return: a non-differentiable batch of samples.","def p_sample_loop(self, model, shape, noise=None, clip_denoised=False,435 denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,436 progress=False, temp=1.0):437 """"""""""""438 final = None439 for sample in self.p_sample_loop_progressive(model, shape, noise=noise,440 clip_denoised=clip_denoised, denoised_fn=denoised_fn, cond_fn=441 cond_fn, model_kwargs=model_kwargs, device=device, progress=442 progress, temp=temp):443 final = sample444 return final['sample']445",point_e\diffusion\gaussian_diffusion.py446p_sample_loop_progressive,"Generate samples from the model and yield intermediate samples from447each timestep of diffusion.448 449Arguments are the same as p_sample_loop().450Returns a generator over dicts, where each dict is the return value of451p_sample().","def p_sample_loop_progressive(self, model, shape, noise=None, clip_denoised452 =False, denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,453 progress=False, temp=1.0):454 """"""""""""455 if device is None:456 device = next(model.parameters()).device457 assert isinstance(shape, (tuple, list))458 if noise is not None:459 img = noise460 else:461 img = th.randn(*shape, device=device) * temp462 indices = list(range(self.num_timesteps))[::-1]463 if progress:464 from tqdm.auto import tqdm465 indices = tqdm(indices)466 for i in indices:467 t = th.tensor([i] * shape[0], device=device)468 with th.no_grad():469 out = self.p_sample(model, img, t, clip_denoised=clip_denoised,470 denoised_fn=denoised_fn, cond_fn=cond_fn, model_kwargs=471 model_kwargs)472 yield self.unscale_out_dict(out)473 img = out['sample']474",point_e\diffusion\gaussian_diffusion.py475ddim_sample,"Sample x_{t-1} from the model using DDIM.476 477Same usage as p_sample().","def ddim_sample(self, model, x, t, clip_denoised=False, denoised_fn=None,478 cond_fn=None, model_kwargs=None, eta=0.0):479 """"""""""""480 out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,481 denoised_fn=denoised_fn, model_kwargs=model_kwargs)482 if cond_fn is not None:483 out = self.condition_score(cond_fn, out, x, t, model_kwargs=484 model_kwargs)485 eps = self._predict_eps_from_xstart(x, t, out['pred_xstart'])486 alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)487 alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)488 sigma = eta * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) * th.sqrt(489 1 - alpha_bar / alpha_bar_prev)490 noise = th.randn_like(x)491 mean_pred = out['pred_xstart'] * th.sqrt(alpha_bar_prev) + th.sqrt(1 -492 alpha_bar_prev - sigma ** 2) * eps493 nonzero_mask = (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))494 sample = mean_pred + nonzero_mask * sigma * noise495 return {'sample': sample, 'pred_xstart': out['pred_xstart']}496",point_e\diffusion\gaussian_diffusion.py497ddim_reverse_sample,Sample x_{t+1} from the model using DDIM reverse ODE.,"def ddim_reverse_sample(self, model, x, t, clip_denoised=False, denoised_fn498 =None, cond_fn=None, model_kwargs=None, eta=0.0):499 """"""""""""500 assert eta == 0.0, 'Reverse ODE only for deterministic path'501 out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,502 denoised_fn=denoised_fn, model_kwargs=model_kwargs)503 if cond_fn is not None:504 out = self.condition_score(cond_fn, out, x, t, model_kwargs=505 model_kwargs)506 eps = (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) *507 x - out['pred_xstart']) / _extract_into_tensor(self.508 sqrt_recipm1_alphas_cumprod, t, x.shape)509 alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)510 mean_pred = out['pred_xstart'] * th.sqrt(alpha_bar_next) + th.sqrt(1 -511 alpha_bar_next) * eps512 return {'sample': mean_pred, 'pred_xstart': out['pred_xstart']}513",point_e\diffusion\gaussian_diffusion.py514ddim_sample_loop,"Generate samples from the model using DDIM.515 516Same usage as p_sample_loop().","def ddim_sample_loop(self, model, shape, noise=None, clip_denoised=False,517 denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,518 progress=False, eta=0.0, temp=1.0):519 """"""""""""520 final = None521 for sample in self.ddim_sample_loop_progressive(model, shape, noise=522 noise, clip_denoised=clip_denoised, denoised_fn=denoised_fn,523 cond_fn=cond_fn, model_kwargs=model_kwargs, device=device, progress524 =progress, eta=eta, temp=temp):525 final = sample526 return final['sample']527",point_e\diffusion\gaussian_diffusion.py528ddim_sample_loop_progressive,"Use DDIM to sample from the model and yield intermediate samples from529each timestep of DDIM.530 531Same usage as p_sample_loop_progressive().","def ddim_sample_loop_progressive(self, model, shape, noise=None,532 clip_denoised=False, denoised_fn=None, cond_fn=None, model_kwargs=None,533 device=None, progress=False, eta=0.0, temp=1.0):534 """"""""""""535 if device is None:536 device = next(model.parameters()).device537 assert isinstance(shape, (tuple, list))538 if noise is not None:539 img = noise540 else:541 img = th.randn(*shape, device=device) * temp542 indices = list(range(self.num_timesteps))[::-1]543 if progress:544 from tqdm.auto import tqdm545 indices = tqdm(indices)546 for i in indices:547 t = th.tensor([i] * shape[0], device=device)548 with th.no_grad():549 out = self.ddim_sample(model, img, t, clip_denoised=550 clip_denoised, denoised_fn=denoised_fn, cond_fn=cond_fn,551 model_kwargs=model_kwargs, eta=eta)552 yield self.unscale_out_dict(out)553 img = out['sample']554",point_e\diffusion\gaussian_diffusion.py555_vb_terms_bpd,"Get a term for the variational lower-bound.556 557The resulting units are bits (rather than nats, as one might expect).558This allows for comparison to other papers.559 560:return: a dict with the following keys:561 - 'output': a shape [N] tensor of NLLs or KLs.562 - 'pred_xstart': the x_0 predictions.","def _vb_terms_bpd(self, model, x_start, x_t, t, clip_denoised=False,563 model_kwargs=None):564 """"""""""""565 true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(566 x_start=x_start, x_t=x_t, t=t)567 out = self.p_mean_variance(model, x_t, t, clip_denoised=clip_denoised,568 model_kwargs=model_kwargs)569 kl = normal_kl(true_mean, true_log_variance_clipped, out['mean'], out[570 'log_variance'])571 kl = mean_flat(kl) / np.log(2.0)572 decoder_nll = -discretized_gaussian_log_likelihood(x_start, means=out[573 'mean'], log_scales=0.5 * out['log_variance'])574 if not self.discretized_t0:575 decoder_nll = th.zeros_like(decoder_nll)576 assert decoder_nll.shape == x_start.shape577 decoder_nll = mean_flat(decoder_nll) / np.log(2.0)578 output = th.where(t == 0, decoder_nll, kl)579 return {'output': output, 'pred_xstart': out['pred_xstart'], 'extra':580 out['extra']}581",point_e\diffusion\gaussian_diffusion.py582training_losses,"Compute training losses for a single timestep.583 584:param model: the model to evaluate loss on.585:param x_start: the [N x C x ...] tensor of inputs.586:param t: a batch of timestep indices.587:param model_kwargs: if not None, a dict of extra keyword arguments to588 pass to the model. This can be used for conditioning.589:param noise: if specified, the specific Gaussian noise to try to remove.590:return: a dict with the key ""loss"" containing a tensor of shape [N].591 Some mean or variance settings may also have other keys.","def training_losses(self, model, x_start, t, model_kwargs=None, noise=None592 ) ->Dict[str, th.Tensor]:593 """"""""""""594 x_start = self.scale_channels(x_start)595 if model_kwargs is None:596 model_kwargs = {}597 if noise is None:598 noise = th.randn_like(x_start)599 x_t = self.q_sample(x_start, t, noise=noise)600 terms = {}601 if self.loss_type == 'kl' or self.loss_type == 'rescaled_kl':602 vb_terms = self._vb_terms_bpd(model=model, x_start=x_start, x_t=x_t,603 t=t, clip_denoised=False, model_kwargs=model_kwargs)604 terms['loss'] = vb_terms['output']605 if self.loss_type == 'rescaled_kl':606 terms['loss'] *= self.num_timesteps607 extra = vb_terms['extra']608 elif self.loss_type == 'mse' or self.loss_type == 'rescaled_mse':609 model_output = model(x_t, t, **model_kwargs)610 if isinstance(model_output, tuple):611 model_output, extra = model_output612 else:613 extra = {}614 if self.model_var_type in ['learned', 'learned_range']:615 B, C = x_t.shape[:2]616 assert model_output.shape == (B, C * 2, *x_t.shape[2:])617 model_output, model_var_values = th.split(model_output, C, dim=1)618 frozen_out = th.cat([model_output.detach(), model_var_values],619 dim=1)620 terms['vb'] = self._vb_terms_bpd(model=lambda *args, r=621 frozen_out: r, x_start=x_start, x_t=x_t, t=t, clip_denoised622 =False)['output']623 if self.loss_type == 'rescaled_mse':624 terms['vb'] *= self.num_timesteps / 1000.0625 target = {'x_prev': self.q_posterior_mean_variance(x_start=x_start,626 x_t=x_t, t=t)[0], 'x_start': x_start, 'epsilon': noise}[self.627 model_mean_type]628 assert model_output.shape == target.shape == x_start.shape629 terms['mse'] = mean_flat((target - model_output) ** 2)630 if 'vb' in terms:631 terms['loss'] = terms['mse'] + terms['vb']632 else:633 terms['loss'] = terms['mse']634 else:635 raise NotImplementedError(self.loss_type)636 if 'losses' in extra:637 terms.update({k: loss for k, (loss, _scale) in extra['losses'].items()}638 )639 for loss, scale in extra['losses'].values():640 terms['loss'] = terms['loss'] + loss * scale641 return terms642",point_e\diffusion\gaussian_diffusion.py643_prior_bpd,"Get the prior KL term for the variational lower-bound, measured in644bits-per-dim.645 646This term can't be optimized, as it only depends on the encoder.647 648:param x_start: the [N x C x ...] tensor of inputs.649:return: a batch of [N] KL values (in bits), one per batch element.","def _prior_bpd(self, x_start):650 """"""""""""651 batch_size = x_start.shape[0]652 t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)653 qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)654 kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0,655 logvar2=0.0)656 return mean_flat(kl_prior) / np.log(2.0)657",point_e\diffusion\gaussian_diffusion.py658calc_bpd_loop,"Compute the entire variational lower-bound, measured in bits-per-dim,659as well as other related quantities.660 661:param model: the model to evaluate loss on.662:param x_start: the [N x C x ...] tensor of inputs.663:param clip_denoised: if True, clip denoised samples.664:param model_kwargs: if not None, a dict of extra keyword arguments to665 pass to the model. This can be used for conditioning.666 667:return: a dict containing the following keys:668 - total_bpd: the total variational lower-bound, per batch element.669 - prior_bpd: the prior term in the lower-bound.670 - vb: an [N x T] tensor of terms in the lower-bound.671 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.672 - mse: an [N x T] tensor of epsilon MSEs for each timestep.","def calc_bpd_loop(self, model, x_start, clip_denoised=False, model_kwargs=None673 ):674 """"""""""""675 device = x_start.device676 batch_size = x_start.shape[0]677 vb = []678 xstart_mse = []679 mse = []680 for t in list(range(self.num_timesteps))[::-1]:681 t_batch = th.tensor([t] * batch_size, device=device)682 noise = th.randn_like(x_start)683 x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)684 with th.no_grad():685 out = self._vb_terms_bpd(model, x_start=x_start, x_t=x_t, t=686 t_batch, clip_denoised=clip_denoised, model_kwargs=model_kwargs687 )688 vb.append(out['output'])689 xstart_mse.append(mean_flat((out['pred_xstart'] - x_start) ** 2))690 eps = self._predict_eps_from_xstart(x_t, t_batch, out['pred_xstart'])691 mse.append(mean_flat((eps - noise) ** 2))692 vb = th.stack(vb, dim=1)693 xstart_mse = th.stack(xstart_mse, dim=1)694 mse = th.stack(mse, dim=1)695 prior_bpd = self._prior_bpd(x_start)696 total_bpd = vb.sum(dim=1) + prior_bpd697 return {'total_bpd': total_bpd, 'prior_bpd': prior_bpd, 'vb': vb,698 'xstart_mse': xstart_mse, 'mse': mse}699",point_e\diffusion\gaussian_diffusion.py700scale_channels,,"def scale_channels(self, x: th.Tensor) ->th.Tensor:701 if self.channel_scales is not None:702 x = x * th.from_numpy(self.channel_scales).to(x).reshape([1, -1, *(703 [1] * (len(x.shape) - 2))])704 if self.channel_biases is not None:705 x = x + th.from_numpy(self.channel_biases).to(x).reshape([1, -1, *(706 [1] * (len(x.shape) - 2))])707 return x708",point_e\diffusion\gaussian_diffusion.py709unscale_channels,,"def unscale_channels(self, x: th.Tensor) ->th.Tensor:710 if self.channel_biases is not None:711 x = x - th.from_numpy(self.channel_biases).to(x).reshape([1, -1, *(712 [1] * (len(x.shape) - 2))])713 if self.channel_scales is not None:714 x = x / th.from_numpy(self.channel_scales).to(x).reshape([1, -1, *(715 [1] * (len(x.shape) - 2))])716 return x717",point_e\diffusion\gaussian_diffusion.py718unscale_out_dict,,"def unscale_out_dict(self, out: Dict[str, Union[th.Tensor, Any]]) ->Dict[719 str, Union[th.Tensor, Any]]:720 return {k: (self.unscale_channels(v) if isinstance(v, th.Tensor) else v721 ) for k, v in out.items()}722",point_e\diffusion\gaussian_diffusion.py723__init__,,"def __init__(self, use_timesteps: Iterable[int], **kwargs):724 self.use_timesteps = set(use_timesteps)725 self.timestep_map = []726 self.original_num_steps = len(kwargs['betas'])727 base_diffusion = GaussianDiffusion(**kwargs)728 last_alpha_cumprod = 1.0729 new_betas = []730 for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):731 if i in self.use_timesteps:732 new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)733 last_alpha_cumprod = alpha_cumprod734 self.timestep_map.append(i)735 kwargs['betas'] = np.array(new_betas)736 super().__init__(**kwargs)737",point_e\diffusion\gaussian_diffusion.py738p_mean_variance,,"def p_mean_variance(self, model, *args, **kwargs):739 return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)740",point_e\diffusion\gaussian_diffusion.py741training_losses,,"def training_losses(self, model, *args, **kwargs):742 return super().training_losses(self._wrap_model(model), *args, **kwargs)743",point_e\diffusion\gaussian_diffusion.py744condition_mean,,"def condition_mean(self, cond_fn, *args, **kwargs):745 return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)746",point_e\diffusion\gaussian_diffusion.py747condition_score,,"def condition_score(self, cond_fn, *args, **kwargs):748 return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)749",point_e\diffusion\gaussian_diffusion.py750_wrap_model,,"def _wrap_model(self, model):751 if isinstance(model, _WrappedModel):752 return model753 return _WrappedModel(model, self.timestep_map, self.original_num_steps)754",point_e\diffusion\gaussian_diffusion.py755__init__,,"def __init__(self, model, timestep_map, original_num_steps):756 self.model = model757 self.timestep_map = timestep_map758 self.original_num_steps = original_num_steps759",point_e\diffusion\gaussian_diffusion.py760__call__,,"def __call__(self, x, ts, **kwargs):761 map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)762 new_ts = map_tensor[ts]763 return self.model(x, new_ts, **kwargs)764",point_e\diffusion\gaussian_diffusion.py765karras_sample,,"def karras_sample(*args, **kwargs):766 last = None767 for x in karras_sample_progressive(*args, **kwargs):768 last = x['x']769 return last770",point_e\diffusion\k_diffusion.py771karras_sample_progressive,,"def karras_sample_progressive(diffusion, model, shape, steps, clip_denoised772 =True, progress=False, model_kwargs=None, device=None, sigma_min=0.002,773 sigma_max=80, rho=7.0, sampler='heun', s_churn=0.0, s_tmin=0.0, s_tmax=774 float('inf'), s_noise=1.0, guidance_scale=0.0):775 sigmas = get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device)776 x_T = th.randn(*shape, device=device) * sigma_max777 sample_fn = {'heun': sample_heun, 'dpm': sample_dpm, 'ancestral':778 sample_euler_ancestral}[sampler]779 if sampler != 'ancestral':780 sampler_args = dict(s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax,781 s_noise=s_noise)782 else:783 sampler_args = {}784 if isinstance(diffusion, KarrasDenoiser):785 786 def denoiser(x_t, sigma):787 _, denoised = diffusion.denoise(model, x_t, sigma, **model_kwargs)788 if clip_denoised:789 denoised = denoised.clamp(-1, 1)790 return denoised791 elif isinstance(diffusion, GaussianDiffusion):792 model = GaussianToKarrasDenoiser(model, diffusion)793 794 def denoiser(x_t, sigma):795 _, denoised = model.denoise(x_t, sigma, clip_denoised=796 clip_denoised, model_kwargs=model_kwargs)797 return denoised798 else:799 raise NotImplementedError800 if guidance_scale != 0 and guidance_scale != 1:801 802 def guided_denoiser(x_t, sigma):803 x_t = th.cat([x_t, x_t], dim=0)804 sigma = th.cat([sigma, sigma], dim=0)805 x_0 = denoiser(x_t, sigma)806 cond_x_0, uncond_x_0 = th.split(x_0, len(x_0) // 2, dim=0)807 x_0 = uncond_x_0 + guidance_scale * (cond_x_0 - uncond_x_0)808 return x_0809 else:810 guided_denoiser = denoiser811 for obj in sample_fn(guided_denoiser, x_T, sigmas, progress=progress,812 **sampler_args):813 if isinstance(diffusion, GaussianDiffusion):814 yield diffusion.unscale_out_dict(obj)815 else:816 yield obj817",point_e\diffusion\k_diffusion.py818get_sigmas_karras,Constructs the noise schedule of Karras et al. (2022).,"def get_sigmas_karras(n, sigma_min, sigma_max, rho=7.0, device='cpu'):819 """"""""""""820 ramp = th.linspace(0, 1, n)821 min_inv_rho = sigma_min ** (1 / rho)822 max_inv_rho = sigma_max ** (1 / rho)823 sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho824 return append_zero(sigmas).to(device)825",point_e\diffusion\k_diffusion.py826to_d,Converts a denoiser output to a Karras ODE derivative.,"def to_d(x, sigma, denoised):827 """"""""""""828 return (x - denoised) / append_dims(sigma, x.ndim)829",point_e\diffusion\k_diffusion.py830get_ancestral_step,"Calculates the noise level (sigma_down) to step down to and the amount831of noise to add (sigma_up) when doing an ancestral sampling step.","def get_ancestral_step(sigma_from, sigma_to):832 """"""""""""833 sigma_up = (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / 834 sigma_from ** 2) ** 0.5835 sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5836 return sigma_down, sigma_up837",point_e\diffusion\k_diffusion.py838sample_euler_ancestral,Ancestral sampling with Euler method steps.,"@th.no_grad()839def sample_euler_ancestral(model, x, sigmas, progress=False):840 """"""""""""841 s_in = x.new_ones([x.shape[0]])842 indices = range(len(sigmas) - 1)843 if progress:844 from tqdm.auto import tqdm845 indices = tqdm(indices)846 for i in indices:847 denoised = model(x, sigmas[i] * s_in)848 sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1])849 yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i],850 'pred_xstart': denoised}851 d = to_d(x, sigmas[i], denoised)852 dt = sigma_down - sigmas[i]853 x = x + d * dt854 x = x + th.randn_like(x) * sigma_up855 yield {'x': x, 'pred_xstart': x}856",point_e\diffusion\k_diffusion.py857sample_heun,Implements Algorithm 2 (Heun steps) from Karras et al. (2022).,"@th.no_grad()858def sample_heun(denoiser, x, sigmas, progress=False, s_churn=0.0, s_tmin=859 0.0, s_tmax=float('inf'), s_noise=1.0):860 """"""""""""861 s_in = x.new_ones([x.shape[0]])862 indices = range(len(sigmas) - 1)863 if progress:864 from tqdm.auto import tqdm865 indices = tqdm(indices)866 for i in indices:867 gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1868 ) if s_tmin <= sigmas[i] <= s_tmax else 0.0869 eps = th.randn_like(x) * s_noise870 sigma_hat = sigmas[i] * (gamma + 1)871 if gamma > 0:872 x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5873 denoised = denoiser(x, sigma_hat * s_in)874 d = to_d(x, sigma_hat, denoised)875 yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat,876 'pred_xstart': denoised}877 dt = sigmas[i + 1] - sigma_hat878 if sigmas[i + 1] == 0:879 x = x + d * dt880 else:881 x_2 = x + d * dt882 denoised_2 = denoiser(x_2, sigmas[i + 1] * s_in)883 d_2 = to_d(x_2, sigmas[i + 1], denoised_2)884 d_prime = (d + d_2) / 2885 x = x + d_prime * dt886 yield {'x': x, 'pred_xstart': denoised}887",point_e\diffusion\k_diffusion.py888sample_dpm,A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022).,"@th.no_grad()889def sample_dpm(denoiser, x, sigmas, progress=False, s_churn=0.0, s_tmin=0.0,890 s_tmax=float('inf'), s_noise=1.0):891 """"""""""""892 s_in = x.new_ones([x.shape[0]])893 indices = range(len(sigmas) - 1)894 if progress:895 from tqdm.auto import tqdm896 indices = tqdm(indices)897 for i in indices:898 gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1899 ) if s_tmin <= sigmas[i] <= s_tmax else 0.0900 eps = th.randn_like(x) * s_noise901 sigma_hat = sigmas[i] * (gamma + 1)902 if gamma > 0:903 x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5904 denoised = denoiser(x, sigma_hat * s_in)905 d = to_d(x, sigma_hat, denoised)906 yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat,907 'denoised': denoised}908 sigma_mid = ((sigma_hat ** (1 / 3) + sigmas[i + 1] ** (1 / 3)) / 2909 ) ** 3910 dt_1 = sigma_mid - sigma_hat911 dt_2 = sigmas[i + 1] - sigma_hat912 x_2 = x + d * dt_1913 denoised_2 = denoiser(x_2, sigma_mid * s_in)914 d_2 = to_d(x_2, sigma_mid, denoised_2)915 x = x + d_2 * dt_2916 yield {'x': x, 'pred_xstart': denoised}917",point_e\diffusion\k_diffusion.py918append_dims,Appends dimensions to the end of a tensor until it has target_dims dimensions.,"def append_dims(x, target_dims):919 """"""""""""920 dims_to_append = target_dims - x.ndim921 if dims_to_append < 0:922 raise ValueError(923 f'input has {x.ndim} dims but target_dims is {target_dims}, which is less'924 )925 return x[(...,) + (None,) * dims_to_append]926",point_e\diffusion\k_diffusion.py927append_zero,,"def append_zero(x):928 return th.cat([x, x.new_zeros([1])])929",point_e\diffusion\k_diffusion.py930__init__,,"def __init__(self, sigma_data: float=0.5):931 self.sigma_data = sigma_data932",point_e\diffusion\k_diffusion.py933get_snr,,"def get_snr(self, sigmas):934 return sigmas ** -2935",point_e\diffusion\k_diffusion.py936get_sigmas,,"def get_sigmas(self, sigmas):937 return sigmas938",point_e\diffusion\k_diffusion.py939get_scalings,,"def get_scalings(self, sigma):940 c_skip = self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2)941 c_out = sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2942 ) ** 0.5943 c_in = 1 / (sigma ** 2 + self.sigma_data ** 2) ** 0.5944 return c_skip, c_out, c_in945",point_e\diffusion\k_diffusion.py946training_losses,,"def training_losses(self, model, x_start, sigmas, model_kwargs=None, noise=None947 ):948 if model_kwargs is None:949 model_kwargs = {}950 if noise is None:951 noise = th.randn_like(x_start)952 terms = {}953 dims = x_start.ndim954 x_t = x_start + noise * append_dims(sigmas, dims)955 c_skip, c_out, _ = [append_dims(x, dims) for x in self.get_scalings(sigmas)956 ]957 model_output, denoised = self.denoise(model, x_t, sigmas, **model_kwargs)958 target = (x_start - c_skip * x_t) / c_out959 terms['mse'] = mean_flat((model_output - target) ** 2)960 terms['xs_mse'] = mean_flat((denoised - x_start) ** 2)961 if 'vb' in terms:962 terms['loss'] = terms['mse'] + terms['vb']963 else:964 terms['loss'] = terms['mse']965 return terms966",point_e\diffusion\k_diffusion.py967denoise,,"def denoise(self, model, x_t, sigmas, **model_kwargs):968 c_skip, c_out, c_in = [append_dims(x, x_t.ndim) for x in self.969 get_scalings(sigmas)]970 rescaled_t = 1000 * 0.25 * th.log(sigmas + 1e-44)971 model_output = model(c_in * x_t, rescaled_t, **model_kwargs)972 denoised = c_out * model_output + c_skip * x_t973 return model_output, denoised974",point_e\diffusion\k_diffusion.py975__init__,,"def __init__(self, model, diffusion):976 from scipy import interpolate977 self.model = model978 self.diffusion = diffusion979 self.alpha_cumprod_to_t = interpolate.interp1d(diffusion.alphas_cumprod,980 np.arange(0, diffusion.num_timesteps))981",point_e\diffusion\k_diffusion.py982sigma_to_t,,"def sigma_to_t(self, sigma):983 alpha_cumprod = 1.0 / (sigma ** 2 + 1)984 if alpha_cumprod > self.diffusion.alphas_cumprod[0]:985 return 0986 elif alpha_cumprod <= self.diffusion.alphas_cumprod[-1]:987 return self.diffusion.num_timesteps - 1988 else:989 return float(self.alpha_cumprod_to_t(alpha_cumprod))990",point_e\diffusion\k_diffusion.py991denoise,,"def denoise(self, x_t, sigmas, clip_denoised=True, model_kwargs=None):992 t = th.tensor([self.sigma_to_t(sigma) for sigma in sigmas.cpu().numpy()993 ], dtype=th.long, device=sigmas.device)994 c_in = append_dims(1.0 / (sigmas ** 2 + 1) ** 0.5, x_t.ndim)995 out = self.diffusion.p_mean_variance(self.model, x_t * c_in, t,996 clip_denoised=clip_denoised, model_kwargs=model_kwargs)997 return None, out['pred_xstart']998",point_e\diffusion\k_diffusion.py999__init__,,"def __init__(self, device: torch.device, models: Sequence[nn.Module],1000 diffusions: Sequence[GaussianDiffusion], num_points: Sequence[int],1001 aux_channels: Sequence[str], model_kwargs_key_filter: Sequence[str]=(1002 '*',), guidance_scale: Sequence[float]=(3.0, 3.0), clip_denoised: bool=1003 True, use_karras: Sequence[bool]=(True, True), karras_steps: Sequence[1004 int]=(64, 64), sigma_min: Sequence[float]=(0.001, 0.001), sigma_max:1005 Sequence[float]=(120, 160), s_churn: Sequence[float]=(3, 0)):1006 n = len(models)1007 assert n > 01008 if n > 1:1009 if len(guidance_scale) == 1:1010 guidance_scale = list(guidance_scale) + [1.0] * (n - 1)1011 if len(use_karras) == 1:1012 use_karras = use_karras * n1013 if len(karras_steps) == 1:1014 karras_steps = karras_steps * n1015 if len(sigma_min) == 1:1016 sigma_min = sigma_min * n1017 if len(sigma_max) == 1:1018 sigma_max = sigma_max * n1019 if len(s_churn) == 1:1020 s_churn = s_churn * n1021 if len(model_kwargs_key_filter) == 1:1022 model_kwargs_key_filter = model_kwargs_key_filter * n1023 if len(model_kwargs_key_filter) == 0:1024 model_kwargs_key_filter = ['*'] * n1025 assert len(guidance_scale) == n1026 assert len(use_karras) == n1027 assert len(karras_steps) == n1028 assert len(sigma_min) == n1029 assert len(sigma_max) == n1030 assert len(s_churn) == n1031 assert len(model_kwargs_key_filter) == n1032 self.device = device1033 self.num_points = num_points1034 self.aux_channels = aux_channels1035 self.model_kwargs_key_filter = model_kwargs_key_filter1036 self.guidance_scale = guidance_scale1037 self.clip_denoised = clip_denoised1038 self.use_karras = use_karras1039 self.karras_steps = karras_steps1040 self.sigma_min = sigma_min1041 self.sigma_max = sigma_max1042 self.s_churn = s_churn1043 self.models = models1044 self.diffusions = diffusions1045",point_e\diffusion\sampler.py1046num_stages,,"@property1047def num_stages(self) ->int:1048 return len(self.models)1049",point_e\diffusion\sampler.py1050sample_batch,,"def sample_batch(self, batch_size: int, model_kwargs: Dict[str, Any]1051 ) ->torch.Tensor:1052 samples = None1053 for x in self.sample_batch_progressive(batch_size, model_kwargs):1054 samples = x1055 return samples1056",point_e\diffusion\sampler.py1057sample_batch_progressive,,"def sample_batch_progressive(self, batch_size: int, model_kwargs: Dict[str,1058 Any]) ->Iterator[torch.Tensor]:1059 samples = None1060 for model, diffusion, stage_num_points, stage_guidance_scale, stage_use_karras, stage_karras_steps, stage_sigma_min, stage_sigma_max, stage_s_churn, stage_key_filter in zip(1061 self.models, self.diffusions, self.num_points, self.guidance_scale,1062 self.use_karras, self.karras_steps, self.sigma_min, self.sigma_max,1063 self.s_churn, self.model_kwargs_key_filter):1064 stage_model_kwargs = model_kwargs.copy()1065 if stage_key_filter != '*':1066 use_keys = set(stage_key_filter.split(','))1067 stage_model_kwargs = {k: v for k, v in stage_model_kwargs.items1068 () if k in use_keys}1069 if samples is not None:1070 stage_model_kwargs['low_res'] = samples1071 if hasattr(model, 'cached_model_kwargs'):1072 stage_model_kwargs = model.cached_model_kwargs(batch_size,1073 stage_model_kwargs)1074 sample_shape = batch_size, 3 + len(self.aux_channels), stage_num_points1075 if stage_guidance_scale != 1 and stage_guidance_scale != 0:1076 for k, v in stage_model_kwargs.copy().items():1077 stage_model_kwargs[k] = torch.cat([v, torch.zeros_like(v)],1078 dim=0)1079 if stage_use_karras:1080 samples_it = karras_sample_progressive(diffusion=diffusion,1081 model=model, shape=sample_shape, steps=stage_karras_steps,1082 clip_denoised=self.clip_denoised, model_kwargs=1083 stage_model_kwargs, device=self.device, sigma_min=1084 stage_sigma_min, sigma_max=stage_sigma_max, s_churn=1085 stage_s_churn, guidance_scale=stage_guidance_scale)1086 else:1087 internal_batch_size = batch_size1088 if stage_guidance_scale:1089 model = self._uncond_guide_model(model, stage_guidance_scale)1090 internal_batch_size *= 21091 samples_it = diffusion.p_sample_loop_progressive(model, shape=(1092 internal_batch_size, *sample_shape[1:]), model_kwargs=1093 stage_model_kwargs, device=self.device, clip_denoised=self.1094 clip_denoised)1095 for x in samples_it:1096 samples = x['pred_xstart'][:batch_size]1097 if 'low_res' in stage_model_kwargs:1098 samples = torch.cat([stage_model_kwargs['low_res'][:len(1099 samples)], samples], dim=-1)1100 yield samples1101",point_e\diffusion\sampler.py1102combine,,"@classmethod1103def combine(cls, *samplers: 'PointCloudSampler') ->'PointCloudSampler':1104 assert all(x.device == samplers[0].device for x in samplers[1:])1105 assert all(x.aux_channels == samplers[0].aux_channels for x in samplers[1:]1106 )1107 assert all(x.clip_denoised == samplers[0].clip_denoised for x in1108 samplers[1:])1109 return cls(device=samplers[0].device, models=[x for y in samplers for x in1110 y.models], diffusions=[x for y in samplers for x in y.diffusions],1111 num_points=[x for y in samplers for x in y.num_points],1112 aux_channels=samplers[0].aux_channels, model_kwargs_key_filter=[x for1113 y in samplers for x in y.model_kwargs_key_filter], guidance_scale=[1114 x for y in samplers for x in y.guidance_scale], clip_denoised=1115 samplers[0].clip_denoised, use_karras=[x for y in samplers for x in1116 y.use_karras], karras_steps=[x for y in samplers for x in y.1117 karras_steps], sigma_min=[x for y in samplers for x in y.sigma_min],1118 sigma_max=[x for y in samplers for x in y.sigma_max], s_churn=[x for1119 y in samplers for x in y.s_churn])1120",point_e\diffusion\sampler.py1121_uncond_guide_model,,"def _uncond_guide_model(self, model: Callable[..., torch.Tensor], scale: float1122 ) ->Callable[..., torch.Tensor]:1123 1124 def model_fn(x_t, ts, **kwargs):1125 half = x_t[:len(x_t) // 2]1126 combined = torch.cat([half, half], dim=0)1127 model_out = model(combined, ts, **kwargs)1128 eps, rest = model_out[:, :3], model_out[:, 3:]1129 cond_eps, uncond_eps = torch.chunk(eps, 2, dim=0)1130 half_eps = uncond_eps + scale * (cond_eps - uncond_eps)1131 eps = torch.cat([half_eps, half_eps], dim=0)1132 return torch.cat([eps, rest], dim=1)1133 return model_fn1134",point_e\diffusion\sampler.py1135split_model_output,,"def split_model_output(self, output: torch.Tensor, rescale_colors: bool=False1136 ) ->Tuple[torch.Tensor, Dict[str, torch.Tensor]]:1137 assert len(self.aux_channels) + 3 == output.shape[11138 ], 'there must be three spatial channels before aux'1139 pos, joined_aux = output[:, :3], output[:, 3:]1140 aux = {}1141 for i, name in enumerate(self.aux_channels):1142 v = joined_aux[:, i]1143 if name in {'R', 'G', 'B', 'A'}:1144 v = v.clamp(0, 255).round()1145 if rescale_colors:1146 v = v / 255.01147 aux[name] = v1148 return pos, aux1149",point_e\diffusion\sampler.py1150output_to_point_clouds,,"def output_to_point_clouds(self, output: torch.Tensor) ->List[PointCloud]:1151 res = []1152 for sample in output:1153 xyz, aux = self.split_model_output(sample[None], rescale_colors=True)1154 res.append(PointCloud(coords=xyz[0].t().cpu().numpy(), channels={k:1155 v[0].cpu().numpy() for k, v in aux.items()}))1156 return res1157",point_e\diffusion\sampler.py1158with_options,,"def with_options(self, guidance_scale: float, clip_denoised: bool,1159 use_karras: Sequence[bool]=(True, True), karras_steps: Sequence[int]=(1160 64, 64), sigma_min: Sequence[float]=(0.001, 0.001), sigma_max: Sequence1161 [float]=(120, 160), s_churn: Sequence[float]=(3, 0)) ->'PointCloudSampler':1162 return PointCloudSampler(device=self.device, models=self.models,1163 diffusions=self.diffusions, num_points=self.num_points,1164 aux_channels=self.aux_channels, model_kwargs_key_filter=self.1165 model_kwargs_key_filter, guidance_scale=guidance_scale,1166 clip_denoised=clip_denoised, use_karras=use_karras, karras_steps=1167 karras_steps, sigma_min=sigma_min, sigma_max=sigma_max, s_churn=s_churn1168 )1169",point_e\diffusion\sampler.py1170get_torch_devices,,"def get_torch_devices() ->List[Union[str, torch.device]]:1171 if torch.cuda.is_available():1172 return [torch.device(f'cuda:{i}') for i in range(torch.cuda.1173 device_count())]1174 else:1175 return ['cpu']1176",point_e\evals\feature_extractor.py1177normalize_point_clouds,,"def normalize_point_clouds(pc: np.ndarray) ->np.ndarray:1178 centroids = np.mean(pc, axis=1, keepdims=True)1179 pc = pc - centroids1180 m = np.max(np.sqrt(np.sum(pc ** 2, axis=-1, keepdims=True)), axis=1,1181 keepdims=True)1182 pc = pc / m1183 return pc1184",point_e\evals\feature_extractor.py1185supports_predictions,,"@property1186@abstractmethod1187def supports_predictions(self) ->bool:1188 pass1189",point_e\evals\feature_extractor.py1190feature_dim,,"@property1191@abstractmethod1192def feature_dim(self) ->int:1193 pass1194",point_e\evals\feature_extractor.py1195num_classes,,"@property1196@abstractmethod1197def num_classes(self) ->int:1198 pass1199",point_e\evals\feature_extractor.py1200features_and_preds,"For a stream of point cloud batches, compute feature vectors and class