Spaces:

Shamima
/

code-search

Runtime error

App Files Files Community

code-search / functions_data.csv

Shamima

Upload 2 files

11b8e37 over 1 year ago

raw history blame contribute delete

No virus

154 kB

	function_name,docstring,function_body,file_path
	diffusion_from_config,,"def diffusion_from_config(config: Dict[str, Any]) ->GaussianDiffusion:
	schedule = config['schedule']
	steps = config['timesteps']
	respace = config.get('respacing', None)
	mean_type = config.get('mean_type', 'epsilon')
	betas = get_named_beta_schedule(schedule, steps)
	channel_scales = config.get('channel_scales', None)
	channel_biases = config.get('channel_biases', None)
	if channel_scales is not None:
	channel_scales = np.array(channel_scales)
	if channel_biases is not None:
	channel_biases = np.array(channel_biases)
	kwargs = dict(betas=betas, model_mean_type=mean_type, model_var_type=
	'learned_range', loss_type='mse', channel_scales=channel_scales,
	channel_biases=channel_biases)
	if respace is None:
	return GaussianDiffusion(**kwargs)
	else:
	return SpacedDiffusion(use_timesteps=space_timesteps(steps, respace
	), **kwargs)
	",point_e\diffusion\configs.py
	get_beta_schedule,"This is the deprecated API for creating beta schedules.

	See get_named_beta_schedule() for the new library of schedules.","def get_beta_schedule(beta_schedule, *, beta_start, beta_end,
	num_diffusion_timesteps):
	""""""""""""
	if beta_schedule == 'linear':
	betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps,
	dtype=np.float64)
	else:
	raise NotImplementedError(beta_schedule)
	assert betas.shape == (num_diffusion_timesteps,)
	return betas
	",point_e\diffusion\gaussian_diffusion.py
	get_named_beta_schedule,"Get a pre-defined beta schedule for the given name.

	The beta schedule library consists of beta schedules which remain similar
	in the limit of num_diffusion_timesteps.
	Beta schedules may be added, but should not be removed or changed once
	they are committed to maintain backwards compatibility.","def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
	""""""""""""
	if schedule_name == 'linear':
	scale = 1000 / num_diffusion_timesteps
	return get_beta_schedule('linear', beta_start=scale * 0.0001,
	beta_end=scale * 0.02, num_diffusion_timesteps=
	num_diffusion_timesteps)
	elif schedule_name == 'cosine':
	return betas_for_alpha_bar(num_diffusion_timesteps, lambda t: math.
	cos((t + 0.008) / 1.008 * math.pi / 2) ** 2)
	else:
	raise NotImplementedError(f'unknown beta schedule: {schedule_name}')
	",point_e\diffusion\gaussian_diffusion.py
	betas_for_alpha_bar,"Create a beta schedule that discretizes the given alpha_t_bar function,
	which defines the cumulative product of (1-beta) over time from t = [0,1].

	:param num_diffusion_timesteps: the number of betas to produce.
	:param alpha_bar: a lambda that takes an argument t from 0 to 1 and
	produces the cumulative product of (1-beta) up to that
	part of the diffusion process.
	:param max_beta: the maximum beta to use; use values lower than 1 to
	prevent singularities.","def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
	""""""""""""
	betas = []
	for i in range(num_diffusion_timesteps):
	t1 = i / num_diffusion_timesteps
	t2 = (i + 1) / num_diffusion_timesteps
	betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
	return np.array(betas)
	",point_e\diffusion\gaussian_diffusion.py
	space_timesteps,"Create a list of timesteps to use from an original diffusion process,
	given the number of timesteps we want to take from equally-sized portions
	of the original process.
	For example, if there's 300 timesteps and the section counts are [10,15,20]
	then the first 100 timesteps are strided to be 10 timesteps, the second 100
	are strided to be 15 timesteps, and the final 100 are strided to be 20.
	:param num_timesteps: the number of diffusion steps in the original
	process to divide up.
	:param section_counts: either a list of numbers, or a string containing
	comma-separated numbers, indicating the step count
	per section. As a special case, use ""ddimN"" where N
	is a number of steps to use the striding from the
	DDIM paper.
	:return: a set of diffusion steps from the original process to use.","def space_timesteps(num_timesteps, section_counts):
	""""""""""""
	if isinstance(section_counts, str):
	if section_counts.startswith('ddim'):
	desired_count = int(section_counts[len('ddim'):])
	for i in range(1, num_timesteps):
	if len(range(0, num_timesteps, i)) == desired_count:
	return set(range(0, num_timesteps, i))
	raise ValueError(
	f'cannot create exactly {num_timesteps} steps with an integer stride'
	)
	elif section_counts.startswith('exact'):
	res = set(int(x) for x in section_counts[len('exact'):].split(','))
	for x in res:
	if x < 0 or x >= num_timesteps:
	raise ValueError(f'timestep out of bounds: {x}')
	return res
	section_counts = [int(x) for x in section_counts.split(',')]
	size_per = num_timesteps // len(section_counts)
	extra = num_timesteps % len(section_counts)
	start_idx = 0
	all_steps = []
	for i, section_count in enumerate(section_counts):
	size = size_per + (1 if i < extra else 0)
	if size < section_count:
	raise ValueError(
	f'cannot divide section of {size} steps into {section_count}')
	if section_count <= 1:
	frac_stride = 1
	else:
	frac_stride = (size - 1) / (section_count - 1)
	cur_idx = 0.0
	taken_steps = []
	for _ in range(section_count):
	taken_steps.append(start_idx + round(cur_idx))
	cur_idx += frac_stride
	all_steps += taken_steps
	start_idx += size
	return set(all_steps)
	",point_e\diffusion\gaussian_diffusion.py
	_extract_into_tensor,"Extract values from a 1-D numpy array for a batch of indices.

	:param arr: the 1-D numpy array.
	:param timesteps: a tensor of indices into the array to extract.
	:param broadcast_shape: a larger shape of K dimensions with the batch
	dimension equal to the length of timesteps.
	:return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.","def _extract_into_tensor(arr, timesteps, broadcast_shape):
	""""""""""""
	res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
	while len(res.shape) < len(broadcast_shape):
	res = res[..., None]
	return res + th.zeros(broadcast_shape, device=timesteps.device)
	",point_e\diffusion\gaussian_diffusion.py
	normal_kl,"Compute the KL divergence between two gaussians.
	Shapes are automatically broadcasted, so batches can be compared to
	scalars, among other use cases.","def normal_kl(mean1, logvar1, mean2, logvar2):
	""""""""""""
	tensor = None
	for obj in (mean1, logvar1, mean2, logvar2):
	if isinstance(obj, th.Tensor):
	tensor = obj
	break
	assert tensor is not None, 'at least one argument must be a Tensor'
	logvar1, logvar2 = [(x if isinstance(x, th.Tensor) else th.tensor(x).to
	(tensor)) for x in (logvar1, logvar2)]
	return 0.5 * (-1.0 + logvar2 - logvar1 + th.exp(logvar1 - logvar2) + (
	mean1 - mean2) ** 2 * th.exp(-logvar2))
	",point_e\diffusion\gaussian_diffusion.py
	approx_standard_normal_cdf,"A fast approximation of the cumulative distribution function of the
	standard normal.","def approx_standard_normal_cdf(x):
	""""""""""""
	return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.
	pow(x, 3))))
	",point_e\diffusion\gaussian_diffusion.py
	discretized_gaussian_log_likelihood,"Compute the log-likelihood of a Gaussian distribution discretizing to a
	given image.
	:param x: the target images. It is assumed that this was uint8 values,
	rescaled to the range [-1, 1].
	:param means: the Gaussian mean Tensor.
	:param log_scales: the Gaussian log stddev Tensor.
	:return: a tensor like x of log probabilities (in nats).","def discretized_gaussian_log_likelihood(x, *, means, log_scales):
	""""""""""""
	assert x.shape == means.shape == log_scales.shape
	centered_x = x - means
	inv_stdv = th.exp(-log_scales)
	plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
	cdf_plus = approx_standard_normal_cdf(plus_in)
	min_in = inv_stdv * (centered_x - 1.0 / 255.0)
	cdf_min = approx_standard_normal_cdf(min_in)
	log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
	log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
	cdf_delta = cdf_plus - cdf_min
	log_probs = th.where(x < -0.999, log_cdf_plus, th.where(x > 0.999,
	log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))))
	assert log_probs.shape == x.shape
	return log_probs
	",point_e\diffusion\gaussian_diffusion.py
	mean_flat,Take the mean over all non-batch dimensions.,"def mean_flat(tensor):
	""""""""""""
	return tensor.flatten(1).mean(1)
	",point_e\diffusion\gaussian_diffusion.py
	__init__,,"def __init__(self, *, betas: Sequence[float], model_mean_type: str,
	model_var_type: str, loss_type: str, discretized_t0: bool=False,
	channel_scales: Optional[np.ndarray]=None, channel_biases: Optional[np.
	ndarray]=None):
	self.model_mean_type = model_mean_type
	self.model_var_type = model_var_type
	self.loss_type = loss_type
	self.discretized_t0 = discretized_t0
	self.channel_scales = channel_scales
	self.channel_biases = channel_biases
	betas = np.array(betas, dtype=np.float64)
	self.betas = betas
	assert len(betas.shape) == 1, 'betas must be 1-D'
	assert (betas > 0).all() and (betas <= 1).all()
	self.num_timesteps = int(betas.shape[0])
	alphas = 1.0 - betas
	self.alphas_cumprod = np.cumprod(alphas, axis=0)
	self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
	self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
	assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
	self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
	self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
	self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
	self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
	self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
	self.posterior_variance = betas * (1.0 - self.alphas_cumprod_prev) / (
	1.0 - self.alphas_cumprod)
	self.posterior_log_variance_clipped = np.log(np.append(self.
	posterior_variance[1], self.posterior_variance[1:]))
	self.posterior_mean_coef1 = betas * np.sqrt(self.alphas_cumprod_prev) / (
	1.0 - self.alphas_cumprod)
	self.posterior_mean_coef2 = (1.0 - self.alphas_cumprod_prev) * np.sqrt(
	alphas) / (1.0 - self.alphas_cumprod)
	",point_e\diffusion\gaussian_diffusion.py
	get_sigmas,,"def get_sigmas(self, t):
	return _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, t.shape)
	",point_e\diffusion\gaussian_diffusion.py
	q_mean_variance,"Get the distribution q(x_t \| x_0).

	:param x_start: the [N x C x ...] tensor of noiseless inputs.
	:param t: the number of diffusion steps (minus 1). Here, 0 means one step.
	:return: A tuple (mean, variance, log_variance), all of x_start's shape.","def q_mean_variance(self, x_start, t):
	""""""""""""
	mean = _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape
	) * x_start
	variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape
	)
	log_variance = _extract_into_tensor(self.log_one_minus_alphas_cumprod,
	t, x_start.shape)
	return mean, variance, log_variance
	",point_e\diffusion\gaussian_diffusion.py
	q_sample,"Diffuse the data for a given number of diffusion steps.

	In other words, sample from q(x_t \| x_0).

	:param x_start: the initial data batch.
	:param t: the number of diffusion steps (minus 1). Here, 0 means one step.
	:param noise: if specified, the split-out normal noise.
	:return: A noisy version of x_start.","def q_sample(self, x_start, t, noise=None):
	""""""""""""
	if noise is None:
	noise = th.randn_like(x_start)
	assert noise.shape == x_start.shape
	return _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape
	) * x_start + _extract_into_tensor(self.
	sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
	",point_e\diffusion\gaussian_diffusion.py
	q_posterior_mean_variance,"Compute the mean and variance of the diffusion posterior:

	q(x_{t-1} \| x_t, x_0)","def q_posterior_mean_variance(self, x_start, x_t, t):
	""""""""""""
	assert x_start.shape == x_t.shape
	posterior_mean = _extract_into_tensor(self.posterior_mean_coef1, t, x_t
	.shape) * x_start + _extract_into_tensor(self.posterior_mean_coef2,
	t, x_t.shape) * x_t
	posterior_variance = _extract_into_tensor(self.posterior_variance, t,
	x_t.shape)
	posterior_log_variance_clipped = _extract_into_tensor(self.
	posterior_log_variance_clipped, t, x_t.shape)
	assert posterior_mean.shape[0] == posterior_variance.shape[0
	] == posterior_log_variance_clipped.shape[0] == x_start.shape[0]
	return posterior_mean, posterior_variance, posterior_log_variance_clipped
	",point_e\diffusion\gaussian_diffusion.py
	p_mean_variance,"Apply the model to get p(x_{t-1} \| x_t), as well as a prediction of
	the initial x, x_0.

	:param model: the model, which takes a signal and a batch of timesteps
	as input.
	:param x: the [N x C x ...] tensor at time t.
	:param t: a 1-D Tensor of timesteps.
	:param clip_denoised: if True, clip the denoised signal into [-1, 1].
	:param denoised_fn: if not None, a function which applies to the
	x_start prediction before it is used to sample. Applies before
	clip_denoised.
	:param model_kwargs: if not None, a dict of extra keyword arguments to
	pass to the model. This can be used for conditioning.
	:return: a dict with the following keys:
	- 'mean': the model mean output.
	- 'variance': the model variance output.
	- 'log_variance': the log of 'variance'.
	- 'pred_xstart': the prediction for x_0.","def p_mean_variance(self, model, x, t, clip_denoised=False, denoised_fn=
	None, model_kwargs=None):
	""""""""""""
	if model_kwargs is None:
	model_kwargs = {}
	B, C = x.shape[:2]
	assert t.shape == (B,)
	model_output = model(x, t, **model_kwargs)
	if isinstance(model_output, tuple):
	model_output, extra = model_output
	else:
	extra = None
	if self.model_var_type in ['learned', 'learned_range']:
	assert model_output.shape == (B, C * 2, *x.shape[2:])
	model_output, model_var_values = th.split(model_output, C, dim=1)
	if self.model_var_type == 'learned':
	model_log_variance = model_var_values
	model_variance = th.exp(model_log_variance)
	else:
	min_log = _extract_into_tensor(self.
	posterior_log_variance_clipped, t, x.shape)
	max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
	frac = (model_var_values + 1) / 2
	model_log_variance = frac * max_log + (1 - frac) * min_log
	model_variance = th.exp(model_log_variance)
	else:
	model_variance, model_log_variance = {'fixed_large': (np.append(
	self.posterior_variance[1], self.betas[1:]), np.log(np.append(
	self.posterior_variance[1], self.betas[1:]))), 'fixed_small': (
	self.posterior_variance, self.posterior_log_variance_clipped)}[self
	.model_var_type]
	model_variance = _extract_into_tensor(model_variance, t, x.shape)
	model_log_variance = _extract_into_tensor(model_log_variance, t, x.
	shape)

	def process_xstart(x):
	if denoised_fn is not None:
	x = denoised_fn(x)
	if clip_denoised:
	return x.clamp(-1, 1)
	return x
	if self.model_mean_type == 'x_prev':
	pred_xstart = process_xstart(self._predict_xstart_from_xprev(x_t=x,
	t=t, xprev=model_output))
	model_mean = model_output
	elif self.model_mean_type in ['x_start', 'epsilon']:
	if self.model_mean_type == 'x_start':
	pred_xstart = process_xstart(model_output)
	else:
	pred_xstart = process_xstart(self._predict_xstart_from_eps(x_t=
	x, t=t, eps=model_output))
	model_mean, _, _ = self.q_posterior_mean_variance(x_start=
	pred_xstart, x_t=x, t=t)
	else:
	raise NotImplementedError(self.model_mean_type)
	assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
	return {'mean': model_mean, 'variance': model_variance, 'log_variance':
	model_log_variance, 'pred_xstart': pred_xstart, 'extra': extra}
	",point_e\diffusion\gaussian_diffusion.py
	_predict_xstart_from_eps,,"def _predict_xstart_from_eps(self, x_t, t, eps):
	assert x_t.shape == eps.shape
	return _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape
	) * x_t - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t,
	x_t.shape) * eps
	",point_e\diffusion\gaussian_diffusion.py
	_predict_xstart_from_xprev,,"def _predict_xstart_from_xprev(self, x_t, t, xprev):
	assert x_t.shape == xprev.shape
	return _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape
	) * xprev - _extract_into_tensor(self.posterior_mean_coef2 / self.
	posterior_mean_coef1, t, x_t.shape) * x_t
	",point_e\diffusion\gaussian_diffusion.py
	_predict_eps_from_xstart,,"def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
	return (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.
	shape) * x_t - pred_xstart) / _extract_into_tensor(self.
	sqrt_recipm1_alphas_cumprod, t, x_t.shape)
	",point_e\diffusion\gaussian_diffusion.py
	condition_mean,"Compute the mean for the previous step, given a function cond_fn that
	computes the gradient of a conditional log probability with respect to
	x. In particular, cond_fn computes grad(log(p(y\|x))), and we want to
	condition on y.

	This uses the conditioning strategy from Sohl-Dickstein et al. (2015).","def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
	""""""""""""
	gradient = cond_fn(x, t, **model_kwargs)
	new_mean = p_mean_var['mean'].float() + p_mean_var['variance'
	] * gradient.float()
	return new_mean
	",point_e\diffusion\gaussian_diffusion.py
	condition_score,"Compute what the p_mean_variance output would have been, should the
	model's score function be conditioned by cond_fn.

	See condition_mean() for details on cond_fn.

	Unlike condition_mean(), this instead uses the conditioning strategy
	from Song et al (2020).","def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
	""""""""""""
	alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
	eps = self._predict_eps_from_xstart(x, t, p_mean_var['pred_xstart'])
	eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x, t, **model_kwargs)
	out = p_mean_var.copy()
	out['pred_xstart'] = self._predict_xstart_from_eps(x, t, eps)
	out['mean'], _, _ = self.q_posterior_mean_variance(x_start=out[
	'pred_xstart'], x_t=x, t=t)
	return out
	",point_e\diffusion\gaussian_diffusion.py
	p_sample,"Sample x_{t-1} from the model at the given timestep.

	:param model: the model to sample from.
	:param x: the current tensor at x_{t-1}.
	:param t: the value of t, starting at 0 for the first diffusion step.
	:param clip_denoised: if True, clip the x_start prediction to [-1, 1].
	:param denoised_fn: if not None, a function which applies to the
	x_start prediction before it is used to sample.
	:param cond_fn: if not None, this is a gradient function that acts
	similarly to the model.
	:param model_kwargs: if not None, a dict of extra keyword arguments to
	pass to the model. This can be used for conditioning.
	:return: a dict containing the following keys:
	- 'sample': a random sample from the model.
	- 'pred_xstart': a prediction of x_0.","def p_sample(self, model, x, t, clip_denoised=False, denoised_fn=None,
	cond_fn=None, model_kwargs=None):
	""""""""""""
	out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,
	denoised_fn=denoised_fn, model_kwargs=model_kwargs)
	noise = th.randn_like(x)
	nonzero_mask = (t != 0).float().view(-1, ([1] (len(x.shape) - 1)))
	if cond_fn is not None:
	out['mean'] = self.condition_mean(cond_fn, out, x, t, model_kwargs=
	model_kwargs)
	sample = out['mean'] + nonzero_mask * th.exp(0.5 * out['log_variance']
	) * noise
	return {'sample': sample, 'pred_xstart': out['pred_xstart']}
	",point_e\diffusion\gaussian_diffusion.py
	p_sample_loop,"Generate samples from the model.

	:param model: the model module.
	:param shape: the shape of the samples, (N, C, H, W).
	:param noise: if specified, the noise from the encoder to sample.
	Should be of the same shape as `shape`.
	:param clip_denoised: if True, clip x_start predictions to [-1, 1].
	:param denoised_fn: if not None, a function which applies to the
	x_start prediction before it is used to sample.
	:param cond_fn: if not None, this is a gradient function that acts
	similarly to the model.
	:param model_kwargs: if not None, a dict of extra keyword arguments to
	pass to the model. This can be used for conditioning.
	:param device: if specified, the device to create the samples on.
	If not specified, use a model parameter's device.
	:param progress: if True, show a tqdm progress bar.
	:return: a non-differentiable batch of samples.","def p_sample_loop(self, model, shape, noise=None, clip_denoised=False,
	denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,
	progress=False, temp=1.0):
	""""""""""""
	final = None
	for sample in self.p_sample_loop_progressive(model, shape, noise=noise,
	clip_denoised=clip_denoised, denoised_fn=denoised_fn, cond_fn=
	cond_fn, model_kwargs=model_kwargs, device=device, progress=
	progress, temp=temp):
	final = sample
	return final['sample']
	",point_e\diffusion\gaussian_diffusion.py
	p_sample_loop_progressive,"Generate samples from the model and yield intermediate samples from
	each timestep of diffusion.

	Arguments are the same as p_sample_loop().
	Returns a generator over dicts, where each dict is the return value of
	p_sample().","def p_sample_loop_progressive(self, model, shape, noise=None, clip_denoised
	=False, denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,
	progress=False, temp=1.0):
	""""""""""""
	if device is None:
	device = next(model.parameters()).device
	assert isinstance(shape, (tuple, list))
	if noise is not None:
	img = noise
	else:
	img = th.randn(shape, device=device) temp
	indices = list(range(self.num_timesteps))[::-1]
	if progress:
	from tqdm.auto import tqdm
	indices = tqdm(indices)
	for i in indices:
	t = th.tensor([i] * shape[0], device=device)
	with th.no_grad():
	out = self.p_sample(model, img, t, clip_denoised=clip_denoised,
	denoised_fn=denoised_fn, cond_fn=cond_fn, model_kwargs=
	model_kwargs)
	yield self.unscale_out_dict(out)
	img = out['sample']
	",point_e\diffusion\gaussian_diffusion.py
	ddim_sample,"Sample x_{t-1} from the model using DDIM.

	Same usage as p_sample().","def ddim_sample(self, model, x, t, clip_denoised=False, denoised_fn=None,
	cond_fn=None, model_kwargs=None, eta=0.0):
	""""""""""""
	out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,
	denoised_fn=denoised_fn, model_kwargs=model_kwargs)
	if cond_fn is not None:
	out = self.condition_score(cond_fn, out, x, t, model_kwargs=
	model_kwargs)
	eps = self._predict_eps_from_xstart(x, t, out['pred_xstart'])
	alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
	alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
	sigma = eta * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar)) * th.sqrt(
	1 - alpha_bar / alpha_bar_prev)
	noise = th.randn_like(x)
	mean_pred = out['pred_xstart'] * th.sqrt(alpha_bar_prev) + th.sqrt(1 -
	alpha_bar_prev - sigma ** 2) * eps
	nonzero_mask = (t != 0).float().view(-1, ([1] (len(x.shape) - 1)))
	sample = mean_pred + nonzero_mask * sigma * noise
	return {'sample': sample, 'pred_xstart': out['pred_xstart']}
	",point_e\diffusion\gaussian_diffusion.py
	ddim_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_fn
	=None, cond_fn=None, model_kwargs=None, eta=0.0):
	""""""""""""
	assert eta == 0.0, 'Reverse ODE only for deterministic path'
	out = self.p_mean_variance(model, x, t, clip_denoised=clip_denoised,
	denoised_fn=denoised_fn, model_kwargs=model_kwargs)
	if cond_fn is not None:
	out = self.condition_score(cond_fn, out, x, t, model_kwargs=
	model_kwargs)
	eps = (_extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) *
	x - out['pred_xstart']) / _extract_into_tensor(self.
	sqrt_recipm1_alphas_cumprod, t, x.shape)
	alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
	mean_pred = out['pred_xstart'] * th.sqrt(alpha_bar_next) + th.sqrt(1 -
	alpha_bar_next) * eps
	return {'sample': mean_pred, 'pred_xstart': out['pred_xstart']}
	",point_e\diffusion\gaussian_diffusion.py
	ddim_sample_loop,"Generate samples from the model using DDIM.

	Same usage as p_sample_loop().","def ddim_sample_loop(self, model, shape, noise=None, clip_denoised=False,
	denoised_fn=None, cond_fn=None, model_kwargs=None, device=None,
	progress=False, eta=0.0, temp=1.0):
	""""""""""""
	final = None
	for sample in self.ddim_sample_loop_progressive(model, shape, noise=
	noise, clip_denoised=clip_denoised, denoised_fn=denoised_fn,
	cond_fn=cond_fn, model_kwargs=model_kwargs, device=device, progress
	=progress, eta=eta, temp=temp):
	final = sample
	return final['sample']
	",point_e\diffusion\gaussian_diffusion.py
	ddim_sample_loop_progressive,"Use DDIM to sample from the model and yield intermediate samples from
	each timestep of DDIM.

	Same usage as p_sample_loop_progressive().","def ddim_sample_loop_progressive(self, model, shape, noise=None,
	clip_denoised=False, denoised_fn=None, cond_fn=None, model_kwargs=None,
	device=None, progress=False, eta=0.0, temp=1.0):
	""""""""""""
	if device is None:
	device = next(model.parameters()).device
	assert isinstance(shape, (tuple, list))
	if noise is not None:
	img = noise
	else:
	img = th.randn(shape, device=device) temp
	indices = list(range(self.num_timesteps))[::-1]
	if progress:
	from tqdm.auto import tqdm
	indices = tqdm(indices)
	for i in indices:
	t = th.tensor([i] * shape[0], device=device)
	with th.no_grad():
	out = self.ddim_sample(model, img, t, clip_denoised=
	clip_denoised, denoised_fn=denoised_fn, cond_fn=cond_fn,
	model_kwargs=model_kwargs, eta=eta)
	yield self.unscale_out_dict(out)
	img = out['sample']
	",point_e\diffusion\gaussian_diffusion.py
	_vb_terms_bpd,"Get a term for the variational lower-bound.

	The resulting units are bits (rather than nats, as one might expect).
	This allows for comparison to other papers.

	:return: a dict with the following keys:
	- 'output': a shape [N] tensor of NLLs or KLs.
	- 'pred_xstart': the x_0 predictions.","def _vb_terms_bpd(self, model, x_start, x_t, t, clip_denoised=False,
	model_kwargs=None):
	""""""""""""
	true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
	x_start=x_start, x_t=x_t, t=t)
	out = self.p_mean_variance(model, x_t, t, clip_denoised=clip_denoised,
	model_kwargs=model_kwargs)
	kl = normal_kl(true_mean, true_log_variance_clipped, out['mean'], out[
	'log_variance'])
	kl = mean_flat(kl) / np.log(2.0)
	decoder_nll = -discretized_gaussian_log_likelihood(x_start, means=out[
	'mean'], log_scales=0.5 * out['log_variance'])
	if not self.discretized_t0:
	decoder_nll = th.zeros_like(decoder_nll)
	assert decoder_nll.shape == x_start.shape
	decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
	output = th.where(t == 0, decoder_nll, kl)
	return {'output': output, 'pred_xstart': out['pred_xstart'], 'extra':
	out['extra']}
	",point_e\diffusion\gaussian_diffusion.py
	training_losses,"Compute training losses for a single timestep.

	:param model: the model to evaluate loss on.
	:param x_start: the [N x C x ...] tensor of inputs.
	:param t: a batch of timestep indices.
	:param model_kwargs: if not None, a dict of extra keyword arguments to
	pass to the model. This can be used for conditioning.
	:param noise: if specified, the specific Gaussian noise to try to remove.
	:return: a dict with the key ""loss"" containing a tensor of shape [N].
	Some mean or variance settings may also have other keys.","def training_losses(self, model, x_start, t, model_kwargs=None, noise=None
	) ->Dict[str, th.Tensor]:
	""""""""""""
	x_start = self.scale_channels(x_start)
	if model_kwargs is None:
	model_kwargs = {}
	if noise is None:
	noise = th.randn_like(x_start)
	x_t = self.q_sample(x_start, t, noise=noise)
	terms = {}
	if self.loss_type == 'kl' or self.loss_type == 'rescaled_kl':
	vb_terms = self._vb_terms_bpd(model=model, x_start=x_start, x_t=x_t,
	t=t, clip_denoised=False, model_kwargs=model_kwargs)
	terms['loss'] = vb_terms['output']
	if self.loss_type == 'rescaled_kl':
	terms['loss'] *= self.num_timesteps
	extra = vb_terms['extra']
	elif self.loss_type == 'mse' or self.loss_type == 'rescaled_mse':
	model_output = model(x_t, t, **model_kwargs)
	if isinstance(model_output, tuple):
	model_output, extra = model_output
	else:
	extra = {}
	if self.model_var_type in ['learned', 'learned_range']:
	B, C = x_t.shape[:2]
	assert model_output.shape == (B, C * 2, *x_t.shape[2:])
	model_output, model_var_values = th.split(model_output, C, dim=1)
	frozen_out = th.cat([model_output.detach(), model_var_values],
	dim=1)
	terms['vb'] = self._vb_terms_bpd(model=lambda *args, r=
	frozen_out: r, x_start=x_start, x_t=x_t, t=t, clip_denoised
	=False)['output']
	if self.loss_type == 'rescaled_mse':
	terms['vb'] *= self.num_timesteps / 1000.0
	target = {'x_prev': self.q_posterior_mean_variance(x_start=x_start,
	x_t=x_t, t=t)[0], 'x_start': x_start, 'epsilon': noise}[self.
	model_mean_type]
	assert model_output.shape == target.shape == x_start.shape
	terms['mse'] = mean_flat((target - model_output) ** 2)
	if 'vb' in terms:
	terms['loss'] = terms['mse'] + terms['vb']
	else:
	terms['loss'] = terms['mse']
	else:
	raise NotImplementedError(self.loss_type)
	if 'losses' in extra:
	terms.update({k: loss for k, (loss, _scale) in extra['losses'].items()}
	)
	for loss, scale in extra['losses'].values():
	terms['loss'] = terms['loss'] + loss * scale
	return terms
	",point_e\diffusion\gaussian_diffusion.py
	_prior_bpd,"Get the prior KL term for the variational lower-bound, measured in
	bits-per-dim.

	This term can't be optimized, as it only depends on the encoder.

	:param x_start: the [N x C x ...] tensor of inputs.
	:return: a batch of [N] KL values (in bits), one per batch element.","def _prior_bpd(self, x_start):
	""""""""""""
	batch_size = x_start.shape[0]
	t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
	qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
	kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0,
	logvar2=0.0)
	return mean_flat(kl_prior) / np.log(2.0)
	",point_e\diffusion\gaussian_diffusion.py
	calc_bpd_loop,"Compute the entire variational lower-bound, measured in bits-per-dim,
	as well as other related quantities.

	:param model: the model to evaluate loss on.
	:param x_start: the [N x C x ...] tensor of inputs.
	:param clip_denoised: if True, clip denoised samples.
	:param model_kwargs: if not None, a dict of extra keyword arguments to
	pass to the model. This can be used for conditioning.

	:return: a dict containing the following keys:
	- total_bpd: the total variational lower-bound, per batch element.
	- prior_bpd: the prior term in the lower-bound.
	- vb: an [N x T] tensor of terms in the lower-bound.
	- xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
	- 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=None
	):
	""""""""""""
	device = x_start.device
	batch_size = x_start.shape[0]
	vb = []
	xstart_mse = []
	mse = []
	for t in list(range(self.num_timesteps))[::-1]:
	t_batch = th.tensor([t] * batch_size, device=device)
	noise = th.randn_like(x_start)
	x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
	with th.no_grad():
	out = self._vb_terms_bpd(model, x_start=x_start, x_t=x_t, t=
	t_batch, clip_denoised=clip_denoised, model_kwargs=model_kwargs
	)
	vb.append(out['output'])
	xstart_mse.append(mean_flat((out['pred_xstart'] - x_start) ** 2))
	eps = self._predict_eps_from_xstart(x_t, t_batch, out['pred_xstart'])
	mse.append(mean_flat((eps - noise) ** 2))
	vb = th.stack(vb, dim=1)
	xstart_mse = th.stack(xstart_mse, dim=1)
	mse = th.stack(mse, dim=1)
	prior_bpd = self._prior_bpd(x_start)
	total_bpd = vb.sum(dim=1) + prior_bpd
	return {'total_bpd': total_bpd, 'prior_bpd': prior_bpd, 'vb': vb,
	'xstart_mse': xstart_mse, 'mse': mse}
	",point_e\diffusion\gaussian_diffusion.py
	scale_channels,,"def scale_channels(self, x: th.Tensor) ->th.Tensor:
	if self.channel_scales is not None:
	x = x * th.from_numpy(self.channel_scales).to(x).reshape([1, -1, *(
	[1] * (len(x.shape) - 2))])
	if self.channel_biases is not None:
	x = x + th.from_numpy(self.channel_biases).to(x).reshape([1, -1, *(
	[1] * (len(x.shape) - 2))])
	return x
	",point_e\diffusion\gaussian_diffusion.py
	unscale_channels,,"def unscale_channels(self, x: th.Tensor) ->th.Tensor:
	if self.channel_biases is not None:
	x = x - th.from_numpy(self.channel_biases).to(x).reshape([1, -1, *(
	[1] * (len(x.shape) - 2))])
	if self.channel_scales is not None:
	x = x / th.from_numpy(self.channel_scales).to(x).reshape([1, -1, *(
	[1] * (len(x.shape) - 2))])
	return x
	",point_e\diffusion\gaussian_diffusion.py
	unscale_out_dict,,"def unscale_out_dict(self, out: Dict[str, Union[th.Tensor, Any]]) ->Dict[
	str, Union[th.Tensor, Any]]:
	return {k: (self.unscale_channels(v) if isinstance(v, th.Tensor) else v
	) for k, v in out.items()}
	",point_e\diffusion\gaussian_diffusion.py
	__init__,,"def __init__(self, use_timesteps: Iterable[int], **kwargs):
	self.use_timesteps = set(use_timesteps)
	self.timestep_map = []
	self.original_num_steps = len(kwargs['betas'])
	base_diffusion = GaussianDiffusion(**kwargs)
	last_alpha_cumprod = 1.0
	new_betas = []
	for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
	if i in self.use_timesteps:
	new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
	last_alpha_cumprod = alpha_cumprod
	self.timestep_map.append(i)
	kwargs['betas'] = np.array(new_betas)
	super().__init__(**kwargs)
	",point_e\diffusion\gaussian_diffusion.py
	p_mean_variance,,"def p_mean_variance(self, model, args, *kwargs):
	return super().p_mean_variance(self._wrap_model(model), args, *kwargs)
	",point_e\diffusion\gaussian_diffusion.py
	training_losses,,"def training_losses(self, model, args, *kwargs):
	return super().training_losses(self._wrap_model(model), args, *kwargs)
	",point_e\diffusion\gaussian_diffusion.py
	condition_mean,,"def condition_mean(self, cond_fn, args, *kwargs):
	return super().condition_mean(self._wrap_model(cond_fn), args, *kwargs)
	",point_e\diffusion\gaussian_diffusion.py
	condition_score,,"def condition_score(self, cond_fn, args, *kwargs):
	return super().condition_score(self._wrap_model(cond_fn), args, *kwargs)
	",point_e\diffusion\gaussian_diffusion.py
	_wrap_model,,"def _wrap_model(self, model):
	if isinstance(model, _WrappedModel):
	return model
	return _WrappedModel(model, self.timestep_map, self.original_num_steps)
	",point_e\diffusion\gaussian_diffusion.py
	__init__,,"def __init__(self, model, timestep_map, original_num_steps):
	self.model = model
	self.timestep_map = timestep_map
	self.original_num_steps = original_num_steps
	",point_e\diffusion\gaussian_diffusion.py
	__call__,,"def __call__(self, x, ts, **kwargs):
	map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
	new_ts = map_tensor[ts]
	return self.model(x, new_ts, **kwargs)
	",point_e\diffusion\gaussian_diffusion.py
	karras_sample,,"def karras_sample(args, *kwargs):
	last = None
	for x in karras_sample_progressive(args, *kwargs):
	last = x['x']
	return last
	",point_e\diffusion\k_diffusion.py
	karras_sample_progressive,,"def karras_sample_progressive(diffusion, model, shape, steps, clip_denoised
	=True, progress=False, model_kwargs=None, device=None, sigma_min=0.002,
	sigma_max=80, rho=7.0, sampler='heun', s_churn=0.0, s_tmin=0.0, s_tmax=
	float('inf'), s_noise=1.0, guidance_scale=0.0):
	sigmas = get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device)
	x_T = th.randn(shape, device=device) sigma_max
	sample_fn = {'heun': sample_heun, 'dpm': sample_dpm, 'ancestral':
	sample_euler_ancestral}[sampler]
	if sampler != 'ancestral':
	sampler_args = dict(s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax,
	s_noise=s_noise)
	else:
	sampler_args = {}
	if isinstance(diffusion, KarrasDenoiser):

	def denoiser(x_t, sigma):
	_, denoised = diffusion.denoise(model, x_t, sigma, **model_kwargs)
	if clip_denoised:
	denoised = denoised.clamp(-1, 1)
	return denoised
	elif isinstance(diffusion, GaussianDiffusion):
	model = GaussianToKarrasDenoiser(model, diffusion)

	def denoiser(x_t, sigma):
	_, denoised = model.denoise(x_t, sigma, clip_denoised=
	clip_denoised, model_kwargs=model_kwargs)
	return denoised
	else:
	raise NotImplementedError
	if guidance_scale != 0 and guidance_scale != 1:

	def guided_denoiser(x_t, sigma):
	x_t = th.cat([x_t, x_t], dim=0)
	sigma = th.cat([sigma, sigma], dim=0)
	x_0 = denoiser(x_t, sigma)
	cond_x_0, uncond_x_0 = th.split(x_0, len(x_0) // 2, dim=0)
	x_0 = uncond_x_0 + guidance_scale * (cond_x_0 - uncond_x_0)
	return x_0
	else:
	guided_denoiser = denoiser
	for obj in sample_fn(guided_denoiser, x_T, sigmas, progress=progress,
	**sampler_args):
	if isinstance(diffusion, GaussianDiffusion):
	yield diffusion.unscale_out_dict(obj)
	else:
	yield obj
	",point_e\diffusion\k_diffusion.py
	get_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'):
	""""""""""""
	ramp = th.linspace(0, 1, n)
	min_inv_rho = sigma_min ** (1 / rho)
	max_inv_rho = sigma_max ** (1 / rho)
	sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
	return append_zero(sigmas).to(device)
	",point_e\diffusion\k_diffusion.py
	to_d,Converts a denoiser output to a Karras ODE derivative.,"def to_d(x, sigma, denoised):
	""""""""""""
	return (x - denoised) / append_dims(sigma, x.ndim)
	",point_e\diffusion\k_diffusion.py
	get_ancestral_step,"Calculates the noise level (sigma_down) to step down to and the amount
	of noise to add (sigma_up) when doing an ancestral sampling step.","def get_ancestral_step(sigma_from, sigma_to):
	""""""""""""
	sigma_up = (sigma_to ** 2 * (sigma_from 2 - sigma_to 2) /
	sigma_from 2) 0.5
	sigma_down = (sigma_to 2 - sigma_up 2) ** 0.5
	return sigma_down, sigma_up
	",point_e\diffusion\k_diffusion.py
	sample_euler_ancestral,Ancestral sampling with Euler method steps.,"@th.no_grad()
	def sample_euler_ancestral(model, x, sigmas, progress=False):
	""""""""""""
	s_in = x.new_ones([x.shape[0]])
	indices = range(len(sigmas) - 1)
	if progress:
	from tqdm.auto import tqdm
	indices = tqdm(indices)
	for i in indices:
	denoised = model(x, sigmas[i] * s_in)
	sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1])
	yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i],
	'pred_xstart': denoised}
	d = to_d(x, sigmas[i], denoised)
	dt = sigma_down - sigmas[i]
	x = x + d * dt
	x = x + th.randn_like(x) * sigma_up
	yield {'x': x, 'pred_xstart': x}
	",point_e\diffusion\k_diffusion.py
	sample_heun,Implements Algorithm 2 (Heun steps) from Karras et al. (2022).,"@th.no_grad()
	def sample_heun(denoiser, x, sigmas, progress=False, s_churn=0.0, s_tmin=
	0.0, s_tmax=float('inf'), s_noise=1.0):
	""""""""""""
	s_in = x.new_ones([x.shape[0]])
	indices = range(len(sigmas) - 1)
	if progress:
	from tqdm.auto import tqdm
	indices = tqdm(indices)
	for i in indices:
	gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1
	) if s_tmin <= sigmas[i] <= s_tmax else 0.0
	eps = th.randn_like(x) * s_noise
	sigma_hat = sigmas[i] * (gamma + 1)
	if gamma > 0:
	x = x + eps * (sigma_hat 2 - sigmas[i] 2) ** 0.5
	denoised = denoiser(x, sigma_hat * s_in)
	d = to_d(x, sigma_hat, denoised)
	yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat,
	'pred_xstart': denoised}
	dt = sigmas[i + 1] - sigma_hat
	if sigmas[i + 1] == 0:
	x = x + d * dt
	else:
	x_2 = x + d * dt
	denoised_2 = denoiser(x_2, sigmas[i + 1] * s_in)
	d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
	d_prime = (d + d_2) / 2
	x = x + d_prime * dt
	yield {'x': x, 'pred_xstart': denoised}
	",point_e\diffusion\k_diffusion.py
	sample_dpm,A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022).,"@th.no_grad()
	def sample_dpm(denoiser, x, sigmas, progress=False, s_churn=0.0, s_tmin=0.0,
	s_tmax=float('inf'), s_noise=1.0):
	""""""""""""
	s_in = x.new_ones([x.shape[0]])
	indices = range(len(sigmas) - 1)
	if progress:
	from tqdm.auto import tqdm
	indices = tqdm(indices)
	for i in indices:
	gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1
	) if s_tmin <= sigmas[i] <= s_tmax else 0.0
	eps = th.randn_like(x) * s_noise
	sigma_hat = sigmas[i] * (gamma + 1)
	if gamma > 0:
	x = x + eps * (sigma_hat 2 - sigmas[i] 2) ** 0.5
	denoised = denoiser(x, sigma_hat * s_in)
	d = to_d(x, sigma_hat, denoised)
	yield {'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat,
	'denoised': denoised}
	sigma_mid = ((sigma_hat (1 / 3) + sigmas[i + 1] (1 / 3)) / 2
	) ** 3
	dt_1 = sigma_mid - sigma_hat
	dt_2 = sigmas[i + 1] - sigma_hat
	x_2 = x + d * dt_1
	denoised_2 = denoiser(x_2, sigma_mid * s_in)
	d_2 = to_d(x_2, sigma_mid, denoised_2)
	x = x + d_2 * dt_2
	yield {'x': x, 'pred_xstart': denoised}
	",point_e\diffusion\k_diffusion.py
	append_dims,Appends dimensions to the end of a tensor until it has target_dims dimensions.,"def append_dims(x, target_dims):
	""""""""""""
	dims_to_append = target_dims - x.ndim
	if dims_to_append < 0:
	raise ValueError(
	f'input has {x.ndim} dims but target_dims is {target_dims}, which is less'
	)
	return x[(...,) + (None,) * dims_to_append]
	",point_e\diffusion\k_diffusion.py
	append_zero,,"def append_zero(x):
	return th.cat([x, x.new_zeros([1])])
	",point_e\diffusion\k_diffusion.py
	__init__,,"def __init__(self, sigma_data: float=0.5):
	self.sigma_data = sigma_data
	",point_e\diffusion\k_diffusion.py
	get_snr,,"def get_snr(self, sigmas):
	return sigmas ** -2
	",point_e\diffusion\k_diffusion.py
	get_sigmas,,"def get_sigmas(self, sigmas):
	return sigmas
	",point_e\diffusion\k_diffusion.py
	get_scalings,,"def get_scalings(self, sigma):
	c_skip = self.sigma_data 2 / (sigma 2 + self.sigma_data ** 2)
	c_out = sigma * self.sigma_data / (sigma 2 + self.sigma_data 2
	) ** 0.5
	c_in = 1 / (sigma 2 + self.sigma_data 2) ** 0.5
	return c_skip, c_out, c_in
	",point_e\diffusion\k_diffusion.py
	training_losses,,"def training_losses(self, model, x_start, sigmas, model_kwargs=None, noise=None
	):
	if model_kwargs is None:
	model_kwargs = {}
	if noise is None:
	noise = th.randn_like(x_start)
	terms = {}
	dims = x_start.ndim
	x_t = x_start + noise * append_dims(sigmas, dims)
	c_skip, c_out, _ = [append_dims(x, dims) for x in self.get_scalings(sigmas)
	]
	model_output, denoised = self.denoise(model, x_t, sigmas, **model_kwargs)
	target = (x_start - c_skip * x_t) / c_out
	terms['mse'] = mean_flat((model_output - target) ** 2)
	terms['xs_mse'] = mean_flat((denoised - x_start) ** 2)
	if 'vb' in terms:
	terms['loss'] = terms['mse'] + terms['vb']
	else:
	terms['loss'] = terms['mse']
	return terms
	",point_e\diffusion\k_diffusion.py
	denoise,,"def denoise(self, model, x_t, sigmas, **model_kwargs):
	c_skip, c_out, c_in = [append_dims(x, x_t.ndim) for x in self.
	get_scalings(sigmas)]
	rescaled_t = 1000 * 0.25 * th.log(sigmas + 1e-44)
	model_output = model(c_in * x_t, rescaled_t, **model_kwargs)
	denoised = c_out * model_output + c_skip * x_t
	return model_output, denoised
	",point_e\diffusion\k_diffusion.py
	__init__,,"def __init__(self, model, diffusion):
	from scipy import interpolate
	self.model = model
	self.diffusion = diffusion
	self.alpha_cumprod_to_t = interpolate.interp1d(diffusion.alphas_cumprod,
	np.arange(0, diffusion.num_timesteps))
	",point_e\diffusion\k_diffusion.py
	sigma_to_t,,"def sigma_to_t(self, sigma):
	alpha_cumprod = 1.0 / (sigma ** 2 + 1)
	if alpha_cumprod > self.diffusion.alphas_cumprod[0]:
	return 0
	elif alpha_cumprod <= self.diffusion.alphas_cumprod[-1]:
	return self.diffusion.num_timesteps - 1
	else:
	return float(self.alpha_cumprod_to_t(alpha_cumprod))
	",point_e\diffusion\k_diffusion.py
	denoise,,"def denoise(self, x_t, sigmas, clip_denoised=True, model_kwargs=None):
	t = th.tensor([self.sigma_to_t(sigma) for sigma in sigmas.cpu().numpy()
	], dtype=th.long, device=sigmas.device)
	c_in = append_dims(1.0 / (sigmas 2 + 1) 0.5, x_t.ndim)
	out = self.diffusion.p_mean_variance(self.model, x_t * c_in, t,
	clip_denoised=clip_denoised, model_kwargs=model_kwargs)
	return None, out['pred_xstart']
	",point_e\diffusion\k_diffusion.py
	__init__,,"def __init__(self, device: torch.device, models: Sequence[nn.Module],
	diffusions: Sequence[GaussianDiffusion], num_points: Sequence[int],
	aux_channels: Sequence[str], model_kwargs_key_filter: Sequence[str]=(
	'*',), guidance_scale: Sequence[float]=(3.0, 3.0), clip_denoised: bool=
	True, use_karras: Sequence[bool]=(True, True), karras_steps: Sequence[
	int]=(64, 64), sigma_min: Sequence[float]=(0.001, 0.001), sigma_max:
	Sequence[float]=(120, 160), s_churn: Sequence[float]=(3, 0)):
	n = len(models)
	assert n > 0
	if n > 1:
	if len(guidance_scale) == 1:
	guidance_scale = list(guidance_scale) + [1.0] * (n - 1)
	if len(use_karras) == 1:
	use_karras = use_karras * n
	if len(karras_steps) == 1:
	karras_steps = karras_steps * n
	if len(sigma_min) == 1:
	sigma_min = sigma_min * n
	if len(sigma_max) == 1:
	sigma_max = sigma_max * n
	if len(s_churn) == 1:
	s_churn = s_churn * n
	if len(model_kwargs_key_filter) == 1:
	model_kwargs_key_filter = model_kwargs_key_filter * n
	if len(model_kwargs_key_filter) == 0:
	model_kwargs_key_filter = [''] n
	assert len(guidance_scale) == n
	assert len(use_karras) == n
	assert len(karras_steps) == n
	assert len(sigma_min) == n
	assert len(sigma_max) == n
	assert len(s_churn) == n
	assert len(model_kwargs_key_filter) == n
	self.device = device
	self.num_points = num_points
	self.aux_channels = aux_channels
	self.model_kwargs_key_filter = model_kwargs_key_filter
	self.guidance_scale = guidance_scale
	self.clip_denoised = clip_denoised
	self.use_karras = use_karras
	self.karras_steps = karras_steps
	self.sigma_min = sigma_min
	self.sigma_max = sigma_max
	self.s_churn = s_churn
	self.models = models
	self.diffusions = diffusions
	",point_e\diffusion\sampler.py
	num_stages,,"@property
	def num_stages(self) ->int:
	return len(self.models)
	",point_e\diffusion\sampler.py
	sample_batch,,"def sample_batch(self, batch_size: int, model_kwargs: Dict[str, Any]
	) ->torch.Tensor:
	samples = None
	for x in self.sample_batch_progressive(batch_size, model_kwargs):
	samples = x
	return samples
	",point_e\diffusion\sampler.py
	sample_batch_progressive,,"def sample_batch_progressive(self, batch_size: int, model_kwargs: Dict[str,
	Any]) ->Iterator[torch.Tensor]:
	samples = None
	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(
	self.models, self.diffusions, self.num_points, self.guidance_scale,
	self.use_karras, self.karras_steps, self.sigma_min, self.sigma_max,
	self.s_churn, self.model_kwargs_key_filter):
	stage_model_kwargs = model_kwargs.copy()
	if stage_key_filter != '*':
	use_keys = set(stage_key_filter.split(','))
	stage_model_kwargs = {k: v for k, v in stage_model_kwargs.items
	() if k in use_keys}
	if samples is not None:
	stage_model_kwargs['low_res'] = samples
	if hasattr(model, 'cached_model_kwargs'):
	stage_model_kwargs = model.cached_model_kwargs(batch_size,
	stage_model_kwargs)
	sample_shape = batch_size, 3 + len(self.aux_channels), stage_num_points
	if stage_guidance_scale != 1 and stage_guidance_scale != 0:
	for k, v in stage_model_kwargs.copy().items():
	stage_model_kwargs[k] = torch.cat([v, torch.zeros_like(v)],
	dim=0)
	if stage_use_karras:
	samples_it = karras_sample_progressive(diffusion=diffusion,
	model=model, shape=sample_shape, steps=stage_karras_steps,
	clip_denoised=self.clip_denoised, model_kwargs=
	stage_model_kwargs, device=self.device, sigma_min=
	stage_sigma_min, sigma_max=stage_sigma_max, s_churn=
	stage_s_churn, guidance_scale=stage_guidance_scale)
	else:
	internal_batch_size = batch_size
	if stage_guidance_scale:
	model = self._uncond_guide_model(model, stage_guidance_scale)
	internal_batch_size *= 2
	samples_it = diffusion.p_sample_loop_progressive(model, shape=(
	internal_batch_size, *sample_shape[1:]), model_kwargs=
	stage_model_kwargs, device=self.device, clip_denoised=self.
	clip_denoised)
	for x in samples_it:
	samples = x['pred_xstart'][:batch_size]
	if 'low_res' in stage_model_kwargs:
	samples = torch.cat([stage_model_kwargs['low_res'][:len(
	samples)], samples], dim=-1)
	yield samples
	",point_e\diffusion\sampler.py
	combine,,"@classmethod
	def combine(cls, *samplers: 'PointCloudSampler') ->'PointCloudSampler':
	assert all(x.device == samplers[0].device for x in samplers[1:])
	assert all(x.aux_channels == samplers[0].aux_channels for x in samplers[1:]
	)
	assert all(x.clip_denoised == samplers[0].clip_denoised for x in
	samplers[1:])
	return cls(device=samplers[0].device, models=[x for y in samplers for x in
	y.models], diffusions=[x for y in samplers for x in y.diffusions],
	num_points=[x for y in samplers for x in y.num_points],
	aux_channels=samplers[0].aux_channels, model_kwargs_key_filter=[x for
	y in samplers for x in y.model_kwargs_key_filter], guidance_scale=[
	x for y in samplers for x in y.guidance_scale], clip_denoised=
	samplers[0].clip_denoised, use_karras=[x for y in samplers for x in
	y.use_karras], karras_steps=[x for y in samplers for x in y.
	karras_steps], sigma_min=[x for y in samplers for x in y.sigma_min],
	sigma_max=[x for y in samplers for x in y.sigma_max], s_churn=[x for
	y in samplers for x in y.s_churn])
	",point_e\diffusion\sampler.py
	_uncond_guide_model,,"def _uncond_guide_model(self, model: Callable[..., torch.Tensor], scale: float
	) ->Callable[..., torch.Tensor]:

	def model_fn(x_t, ts, **kwargs):
	half = x_t[:len(x_t) // 2]
	combined = torch.cat([half, half], dim=0)
	model_out = model(combined, ts, **kwargs)
	eps, rest = model_out[:, :3], model_out[:, 3:]
	cond_eps, uncond_eps = torch.chunk(eps, 2, dim=0)
	half_eps = uncond_eps + scale * (cond_eps - uncond_eps)
	eps = torch.cat([half_eps, half_eps], dim=0)
	return torch.cat([eps, rest], dim=1)
	return model_fn
	",point_e\diffusion\sampler.py
	split_model_output,,"def split_model_output(self, output: torch.Tensor, rescale_colors: bool=False
	) ->Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
	assert len(self.aux_channels) + 3 == output.shape[1
	], 'there must be three spatial channels before aux'
	pos, joined_aux = output[:, :3], output[:, 3:]
	aux = {}
	for i, name in enumerate(self.aux_channels):
	v = joined_aux[:, i]
	if name in {'R', 'G', 'B', 'A'}:
	v = v.clamp(0, 255).round()
	if rescale_colors:
	v = v / 255.0
	aux[name] = v
	return pos, aux
	",point_e\diffusion\sampler.py
	output_to_point_clouds,,"def output_to_point_clouds(self, output: torch.Tensor) ->List[PointCloud]:
	res = []
	for sample in output:
	xyz, aux = self.split_model_output(sample[None], rescale_colors=True)
	res.append(PointCloud(coords=xyz[0].t().cpu().numpy(), channels={k:
	v[0].cpu().numpy() for k, v in aux.items()}))
	return res
	",point_e\diffusion\sampler.py
	with_options,,"def with_options(self, guidance_scale: float, clip_denoised: bool,
	use_karras: Sequence[bool]=(True, True), karras_steps: Sequence[int]=(
	64, 64), sigma_min: Sequence[float]=(0.001, 0.001), sigma_max: Sequence
	[float]=(120, 160), s_churn: Sequence[float]=(3, 0)) ->'PointCloudSampler':
	return PointCloudSampler(device=self.device, models=self.models,
	diffusions=self.diffusions, num_points=self.num_points,
	aux_channels=self.aux_channels, model_kwargs_key_filter=self.
	model_kwargs_key_filter, guidance_scale=guidance_scale,
	clip_denoised=clip_denoised, use_karras=use_karras, karras_steps=
	karras_steps, sigma_min=sigma_min, sigma_max=sigma_max, s_churn=s_churn
	)
	",point_e\diffusion\sampler.py
	get_torch_devices,,"def get_torch_devices() ->List[Union[str, torch.device]]:
	if torch.cuda.is_available():
	return [torch.device(f'cuda:{i}') for i in range(torch.cuda.
	device_count())]
	else:
	return ['cpu']
	",point_e\evals\feature_extractor.py
	normalize_point_clouds,,"def normalize_point_clouds(pc: np.ndarray) ->np.ndarray:
	centroids = np.mean(pc, axis=1, keepdims=True)
	pc = pc - centroids
	m = np.max(np.sqrt(np.sum(pc ** 2, axis=-1, keepdims=True)), axis=1,
	keepdims=True)
	pc = pc / m
	return pc
	",point_e\evals\feature_extractor.py
	supports_predictions,,"@property
	@abstractmethod
	def supports_predictions(self) ->bool:
	pass
	",point_e\evals\feature_extractor.py
	feature_dim,,"@property
	@abstractmethod
	def feature_dim(self) ->int:
	pass
	",point_e\evals\feature_extractor.py
	num_classes,,"@property
	@abstractmethod
	def num_classes(self) ->int:
	pass
	",point_e\evals\feature_extractor.py
	features_and_preds,"For a stream of point cloud batches, compute feature vectors and class
	predictions.

	:param point_clouds: a streamer for a sample batch. Typically, arr_0
	will contain the XYZ coordinates.
	:return: a tuple (features, predictions)
	- features: a [B x feature_dim] array of feature vectors.
	- predictions: a [B x num_classes] array of probabilities.","@abstractmethod
	def features_and_preds(self, streamer: NpzStreamer) ->Tuple[np.ndarray, np.
	ndarray]:
	""""""""""""
	",point_e\evals\feature_extractor.py
	__init__,,"def __init__(self, devices: List[Union[str, torch.device]],
	device_batch_size: int=64, cache_dir: Optional[str]=None):
	state_dict = load_checkpoint('pointnet', device=torch.device('cpu'),
	cache_dir=cache_dir)['model_state_dict']
	self.device_batch_size = device_batch_size
	self.devices = devices
	self.models = []
	for device in devices:
	model = get_model(num_class=40, normal_channel=False, width_mult=2)
	model.load_state_dict(state_dict)
	model.to(device)
	model.eval()
	self.models.append(model)
	",point_e\evals\feature_extractor.py
	supports_predictions,,"@property
	def supports_predictions(self) ->bool:
	return True
	",point_e\evals\feature_extractor.py
	feature_dim,,"@property
	def feature_dim(self) ->int:
	return 256
	",point_e\evals\feature_extractor.py
	num_classes,,"@property
	def num_classes(self) ->int:
	return 40
	",point_e\evals\feature_extractor.py
	features_and_preds,,"def features_and_preds(self, streamer: NpzStreamer) ->Tuple[np.ndarray, np.
	ndarray]:
	batch_size = self.device_batch_size * len(self.devices)
	point_clouds = (x['arr_0'] for x in streamer.stream(batch_size, ['arr_0']))
	output_features = []
	output_predictions = []
	with ThreadPool(len(self.devices)) as pool:
	for batch in point_clouds:
	batch = normalize_point_clouds(batch)
	batches = []
	for i, device in zip(range(0, len(batch), self.
	device_batch_size), self.devices):
	batches.append(torch.from_numpy(batch[i:i + self.
	device_batch_size]).permute(0, 2, 1).to(dtype=torch.
	float32, device=device))

	def compute_features(i_batch):
	i, batch = i_batch
	with torch.no_grad():
	return self.models[i](batch, features=True)
	for logits, _, features in pool.imap(compute_features,
	enumerate(batches)):
	output_features.append(features.cpu().numpy())
	output_predictions.append(logits.exp().cpu().numpy())
	return np.concatenate(output_features, axis=0), np.concatenate(
	output_predictions, axis=0)
	",point_e\evals\feature_extractor.py
	compute_statistics,,"def compute_statistics(feats: np.ndarray) ->FIDStatistics:
	mu = np.mean(feats, axis=0)
	sigma = np.cov(feats, rowvar=False)
	return FIDStatistics(mu, sigma)
	",point_e\evals\fid_is.py
	compute_inception_score,,"def compute_inception_score(preds: np.ndarray, split_size: int=5000) ->float:
	scores = []
	for i in range(0, len(preds), split_size):
	part = preds[i:i + split_size]
	kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0))
	)
	kl = np.mean(np.sum(kl, 1))
	scores.append(np.exp(kl))
	return float(np.mean(scores))
	",point_e\evals\fid_is.py
	__init__,,"def __init__(self, mu: np.ndarray, sigma: np.ndarray):
	self.mu = mu
	self.sigma = sigma
	",point_e\evals\fid_is.py
	frechet_distance,Compute the Frechet distance between two sets of statistics.,"def frechet_distance(self, other, eps=1e-06):
	""""""""""""
	mu1, sigma1 = self.mu, self.sigma
	mu2, sigma2 = other.mu, other.sigma
	mu1 = np.atleast_1d(mu1)
	mu2 = np.atleast_1d(mu2)
	sigma1 = np.atleast_2d(sigma1)
	sigma2 = np.atleast_2d(sigma2)
	assert mu1.shape == mu2.shape, f'Training and test mean vectors have different lengths: {mu1.shape}, {mu2.shape}'
	assert sigma1.shape == sigma2.shape, f'Training and test covariances have different dimensions: {sigma1.shape}, {sigma2.shape}'
	diff = mu1 - mu2
	covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
	if not np.isfinite(covmean).all():
	msg = (
	'fid calculation produces singular product; adding %s to diagonal of cov estimates'
	% eps)
	warnings.warn(msg)
	offset = np.eye(sigma1.shape[0]) * eps
	covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
	if np.iscomplexobj(covmean):
	if not np.allclose(np.diagonal(covmean).imag, 0, atol=0.001):
	m = np.max(np.abs(covmean.imag))
	raise ValueError('Imaginary component {}'.format(m))
	covmean = covmean.real
	tr_covmean = np.trace(covmean)
	return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2
	) - 2 * tr_covmean
	",point_e\evals\fid_is.py
	_npz_paths_and_length,,"def _npz_paths_and_length(glob_path: str) ->Tuple[List[str], Optional[int]]:
	count_match = re.match('^(.)\\[:([0-9])\\]$', glob_path)
	if count_match:
	raw_path = count_match[1]
	max_count = int(count_match[2])
	else:
	raw_path = glob_path
	max_count = None
	paths = sorted(glob.glob(raw_path))
	if not len(paths):
	raise ValueError(f'no paths found matching: {glob_path}')
	return paths, max_count
	",point_e\evals\npz_stream.py
	open_npz_arrays,,"@contextmanager
	def open_npz_arrays(path: str, arr_names: Sequence[str]) ->List[NpzArrayReader
	]:
	if not len(arr_names):
	yield []
	return
	arr_name = arr_names[0]
	with open_array(path, arr_name) as arr_f:
	version = np.lib.format.read_magic(arr_f)
	header = None
	if version == (1, 0):
	header = np.lib.format.read_array_header_1_0(arr_f)
	elif version == (2, 0):
	header = np.lib.format.read_array_header_2_0(arr_f)
	if header is None:
	reader = MemoryNpzArrayReader.load(path, arr_name)
	else:
	shape, fortran, dtype = header
	if fortran or dtype.hasobject:
	reader = MemoryNpzArrayReader.load(path, arr_name)
	else:
	reader = StreamingNpzArrayReader(arr_f, shape, dtype)
	with open_npz_arrays(path, arr_names[1:]) as next_readers:
	yield [reader] + next_readers
	",point_e\evals\npz_stream.py
	_read_bytes,"Copied from: https://github.com/numpy/numpy/blob/fb215c76967739268de71aa4bda55dd1b062bc2e/numpy/lib/format.py#L788-L886

	Read from file-like object until size bytes are read.
	Raises ValueError if not EOF is encountered before size bytes are read.
	Non-blocking objects only supported if they derive from io objects.
	Required as e.g. ZipExtFile in python 2.6 can return less data than
	requested.","def _read_bytes(fp, size, error_template='ran out of data'):
	""""""""""""
	data = bytes()
	while True:
	try:
	r = fp.read(size - len(data))
	data += r
	if len(r) == 0 or len(data) == size:
	break
	except io.BlockingIOError:
	pass
	if len(data) != size:
	msg = 'EOF: reading %s, expected %d bytes got %d'
	raise ValueError(msg % (error_template, size, len(data)))
	else:
	return data
	",point_e\evals\npz_stream.py
	open_array,,"@contextmanager
	def open_array(path: str, arr_name: str):
	with open(path, 'rb') as f:
	with zipfile.ZipFile(f, 'r') as zip_f:
	if f'{arr_name}.npy' not in zip_f.namelist():
	raise ValueError(f'missing {arr_name} in npz file')
	with zip_f.open(f'{arr_name}.npy', 'r') as arr_f:
	yield arr_f
	",point_e\evals\npz_stream.py
	_dict_batch_size,,"def _dict_batch_size(objs: Dict[str, np.ndarray]) ->int:
	return len(next(iter(objs.values())))
	",point_e\evals\npz_stream.py
	infos_from_first_file,,"@classmethod
	def infos_from_first_file(cls, glob_path: str) ->Dict[str, 'NumpyArrayInfo']:
	paths, _ = _npz_paths_and_length(glob_path)
	return cls.infos_from_file(paths[0])
	",point_e\evals\npz_stream.py
	infos_from_file,Extract the info of every array in an npz file.,"@classmethod
	def infos_from_file(cls, npz_path: str) ->Dict[str, 'NumpyArrayInfo']:
	""""""""""""
	if not os.path.exists(npz_path):
	raise FileNotFoundError(f'batch of samples was not found: {npz_path}')
	results = {}
	with open(npz_path, 'rb') as f:
	with zipfile.ZipFile(f, 'r') as zip_f:
	for name in zip_f.namelist():
	if not name.endswith('.npy'):
	continue
	key_name = name[:-len('.npy')]
	with zip_f.open(name, 'r') as arr_f:
	version = np.lib.format.read_magic(arr_f)
	if version == (1, 0):
	header = np.lib.format.read_array_header_1_0(arr_f)
	elif version == (2, 0):
	header = np.lib.format.read_array_header_2_0(arr_f)
	else:
	raise ValueError(
	f'unknown numpy array version: {version}')
	shape, _, dtype = header
	results[key_name] = cls(name=key_name, dtype=dtype,
	shape=shape)
	return results
	",point_e\evals\npz_stream.py
	elem_shape,,"@property
	def elem_shape(self) ->Tuple[int]:
	return self.shape[1:]
	",point_e\evals\npz_stream.py
	validate,,"def validate(self):
	if self.name in {'R', 'G', 'B'}:
	if len(self.shape) != 2:
	raise ValueError(
	f""expecting exactly 2-D shape for '{self.name}' but got: {self.shape}""
	)
	elif self.name == 'arr_0':
	if len(self.shape) < 2:
	raise ValueError(
	f'expecting at least 2-D shape but got: {self.shape}')
	elif len(self.shape) == 3:
	if not np.issubdtype(self.dtype, np.floating):
	raise ValueError(
	f'invalid dtype for audio batch: {self.dtype} (expected float)'
	)
	elif self.dtype != np.uint8:
	raise ValueError(
	f'invalid dtype for image batch: {self.dtype} (expected uint8)'
	)
	",point_e\evals\npz_stream.py
	__init__,,"def __init__(self, glob_path: str):
	self.paths, self.trunc_length = _npz_paths_and_length(glob_path)
	self.infos = NumpyArrayInfo.infos_from_file(self.paths[0])
	",point_e\evals\npz_stream.py
	keys,,"def keys(self) ->List[str]:
	return list(self.infos.keys())
	",point_e\evals\npz_stream.py
	stream,,"def stream(self, batch_size: int, keys: Sequence[str]) ->Iterator[Dict[str,
	np.ndarray]]:
	cur_batch = None
	num_remaining = self.trunc_length
	for path in self.paths:
	if num_remaining is not None and num_remaining <= 0:
	break
	with open_npz_arrays(path, keys) as readers:
	combined_reader = CombinedReader(keys, readers)
	while num_remaining is None or num_remaining > 0:
	read_bs = batch_size
	if cur_batch is not None:
	read_bs -= _dict_batch_size(cur_batch)
	if num_remaining is not None:
	read_bs = min(read_bs, num_remaining)
	batch = combined_reader.read_batch(read_bs)
	if batch is None:
	break
	if num_remaining is not None:
	num_remaining -= _dict_batch_size(batch)
	if cur_batch is None:
	cur_batch = batch
	else:
	cur_batch = {k: np.concatenate([cur_batch[k], v], axis=
	0) for k, v in batch.items()}
	if _dict_batch_size(cur_batch) == batch_size:
	yield cur_batch
	cur_batch = None
	if cur_batch is not None:
	yield cur_batch
	",point_e\evals\npz_stream.py
	read_batch,,"@abstractmethod
	def read_batch(self, batch_size: int) ->Optional[np.ndarray]:
	pass
	",point_e\evals\npz_stream.py
	__init__,,"def __init__(self, arr_f, shape, dtype):
	self.arr_f = arr_f
	self.shape = shape
	self.dtype = dtype
	self.idx = 0
	",point_e\evals\npz_stream.py
	read_batch,,"def read_batch(self, batch_size: int) ->Optional[np.ndarray]:
	if self.idx >= self.shape[0]:
	return None
	bs = min(batch_size, self.shape[0] - self.idx)
	self.idx += bs
	if self.dtype.itemsize == 0:
	return np.ndarray([bs, *self.shape[1:]], dtype=self.dtype)
	read_count = bs * np.prod(self.shape[1:])
	read_size = int(read_count * self.dtype.itemsize)
	data = _read_bytes(self.arr_f, read_size, 'array data')
	return np.frombuffer(data, dtype=self.dtype).reshape([bs, *self.shape[1:]])
	",point_e\evals\npz_stream.py
	__init__,,"def __init__(self, arr):
	self.arr = arr
	self.idx = 0
	",point_e\evals\npz_stream.py
	load,,"@classmethod
	def load(cls, path: str, arr_name: str):
	with open(path, 'rb') as f:
	arr = np.load(f)[arr_name]
	return cls(arr)
	",point_e\evals\npz_stream.py
	read_batch,,"def read_batch(self, batch_size: int) ->Optional[np.ndarray]:
	if self.idx >= self.arr.shape[0]:
	return None
	res = self.arr[self.idx:self.idx + batch_size]
	self.idx += batch_size
	return res
	",point_e\evals\npz_stream.py
	__init__,,"def __init__(self, keys: List[str], readers: List[NpzArrayReader]):
	self.keys = keys
	self.readers = readers
	",point_e\evals\npz_stream.py
	read_batch,,"def read_batch(self, batch_size: int) ->Optional[Dict[str, np.ndarray]]:
	batches = [r.read_batch(batch_size) for r in self.readers]
	any_none = any(x is None for x in batches)
	all_none = all(x is None for x in batches)
	if any_none != all_none:
	raise RuntimeError('different keys had different numbers of elements')
	if any_none:
	return None
	if any(len(x) != len(batches[0]) for x in batches):
	raise RuntimeError('different keys had different numbers of elements')
	return dict(zip(self.keys, batches))
	",point_e\evals\npz_stream.py
	__init__,,"def __init__(self, num_class, normal_channel=True, width_mult=1):
	super(get_model, self).__init__()
	self.width_mult = width_mult
	in_channel = 6 if normal_channel else 3
	self.normal_channel = normal_channel
	self.sa1 = PointNetSetAbstraction(npoint=512, radius=0.2, nsample=32,
	in_channel=in_channel, mlp=[64 * width_mult, 64 * width_mult, 128 *
	width_mult], group_all=False)
	self.sa2 = PointNetSetAbstraction(npoint=128, radius=0.4, nsample=64,
	in_channel=128 * width_mult + 3, mlp=[128 * width_mult, 128 *
	width_mult, 256 * width_mult], group_all=False)
	self.sa3 = PointNetSetAbstraction(npoint=None, radius=None, nsample=
	None, in_channel=256 * width_mult + 3, mlp=[256 * width_mult, 512 *
	width_mult, 1024 * width_mult], group_all=True)
	self.fc1 = nn.Linear(1024 * width_mult, 512 * width_mult)
	self.bn1 = nn.BatchNorm1d(512 * width_mult)
	self.drop1 = nn.Dropout(0.4)
	self.fc2 = nn.Linear(512 * width_mult, 256 * width_mult)
	self.bn2 = nn.BatchNorm1d(256 * width_mult)
	self.drop2 = nn.Dropout(0.4)
	self.fc3 = nn.Linear(256 * width_mult, num_class)
	",point_e\evals\pointnet2_cls_ssg.py
	forward,,"def forward(self, xyz, features=False):
	B, _, _ = xyz.shape
	if self.normal_channel:
	norm = xyz[:, 3:, :]
	xyz = xyz[:, :3, :]
	else:
	norm = None
	l1_xyz, l1_points = self.sa1(xyz, norm)
	l2_xyz, l2_points = self.sa2(l1_xyz, l1_points)
	l3_xyz, l3_points = self.sa3(l2_xyz, l2_points)
	x = l3_points.view(B, 1024 * self.width_mult)
	x = self.drop1(F.relu(self.bn1(self.fc1(x))))
	result_features = self.bn2(self.fc2(x))
	x = self.drop2(F.relu(result_features))
	x = self.fc3(x)
	x = F.log_softmax(x, -1)
	if features:
	return x, l3_points, result_features
	else:
	return x, l3_points
	",point_e\evals\pointnet2_cls_ssg.py
	__init__,,"def __init__(self):
	super(get_loss, self).__init__()
	",point_e\evals\pointnet2_cls_ssg.py
	forward,,"def forward(self, pred, target, trans_feat):
	total_loss = F.nll_loss(pred, target)
	return total_loss
	",point_e\evals\pointnet2_cls_ssg.py
	timeit,,"def timeit(tag, t):
	print('{}: {}s'.format(tag, time() - t))
	return time()
	",point_e\evals\pointnet2_utils.py
	pc_normalize,,"def pc_normalize(pc):
	l = pc.shape[0]
	centroid = np.mean(pc, axis=0)
	pc = pc - centroid
	m = np.max(np.sqrt(np.sum(pc ** 2, axis=1)))
	pc = pc / m
	return pc
	",point_e\evals\pointnet2_utils.py
	square_distance,"Calculate Euclid distance between each two points.

	src^T * dst = xn * xm + yn * ym + zn * zm;
	sum(src^2, dim=-1) = xnxn + ynyn + zn*zn;
	sum(dst^2, dim=-1) = xmxm + ymym + zm*zm;
	dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
	= sum(src2,dim=-1)+sum(dst2,dim=-1)-2src^Tdst

	Input:
	src: source points, [B, N, C]
	dst: target points, [B, M, C]
	Output:
	dist: per-point square distance, [B, N, M]","def square_distance(src, dst):
	""""""""""""
	B, N, _ = src.shape
	_, M, _ = dst.shape
	dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
	dist += torch.sum(src ** 2, -1).view(B, N, 1)
	dist += torch.sum(dst ** 2, -1).view(B, 1, M)
	return dist
	",point_e\evals\pointnet2_utils.py
	index_points,"Input:
	points: input points data, [B, N, C]
	idx: sample index data, [B, S]
	Return:
	new_points:, indexed points data, [B, S, C]","def index_points(points, idx):
	""""""""""""
	device = points.device
	B = points.shape[0]
	view_shape = list(idx.shape)
	view_shape[1:] = [1] * (len(view_shape) - 1)
	repeat_shape = list(idx.shape)
	repeat_shape[0] = 1
	batch_indices = torch.arange(B, dtype=torch.long).to(device).view(
	view_shape).repeat(repeat_shape)
	new_points = points[batch_indices, idx, :]
	return new_points
	",point_e\evals\pointnet2_utils.py
	farthest_point_sample,"Input:
	xyz: pointcloud data, [B, N, 3]
	npoint: number of samples
	Return:
	centroids: sampled pointcloud index, [B, npoint]","def farthest_point_sample(xyz, npoint, deterministic=False):
	""""""""""""
	device = xyz.device
	B, N, C = xyz.shape
	centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
	distance = torch.ones(B, N).to(device) * 10000000000.0
	if deterministic:
	farthest = torch.arange(0, B, dtype=torch.long).to(device)
	else:
	farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
	batch_indices = torch.arange(B, dtype=torch.long).to(device)
	for i in range(npoint):
	centroids[:, i] = farthest
	centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
	dist = torch.sum((xyz - centroid) ** 2, -1)
	mask = dist < distance
	distance[mask] = dist[mask]
	farthest = torch.max(distance, -1)[1]
	return centroids
	",point_e\evals\pointnet2_utils.py
	query_ball_point,"Input:
	radius: local region radius
	nsample: max sample number in local region
	xyz: all points, [B, N, 3]
	new_xyz: query points, [B, S, 3]
	Return:
	group_idx: grouped points index, [B, S, nsample]","def query_ball_point(radius, nsample, xyz, new_xyz):
	""""""""""""
	device = xyz.device
	B, N, C = xyz.shape
	_, S, _ = new_xyz.shape
	group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N
	).repeat([B, S, 1])
	sqrdists = square_distance(new_xyz, xyz)
	group_idx[sqrdists > radius ** 2] = N
	group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
	group_first = group_idx[:, :, 0].view(B, S, 1).repeat([1, 1, nsample])
	mask = group_idx == N
	group_idx[mask] = group_first[mask]
	return group_idx
	",point_e\evals\pointnet2_utils.py
	sample_and_group,"Input:
	npoint:
	radius:
	nsample:
	xyz: input points position data, [B, N, 3]
	points: input points data, [B, N, D]
	Return:
	new_xyz: sampled points position data, [B, npoint, nsample, 3]
	new_points: sampled points data, [B, npoint, nsample, 3+D]","def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False,
	deterministic=False):
	""""""""""""
	B, N, C = xyz.shape
	S = npoint
	fps_idx = farthest_point_sample(xyz, npoint, deterministic=deterministic)
	new_xyz = index_points(xyz, fps_idx)
	idx = query_ball_point(radius, nsample, xyz, new_xyz)
	grouped_xyz = index_points(xyz, idx)
	grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
	if points is not None:
	grouped_points = index_points(points, idx)
	new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1)
	else:
	new_points = grouped_xyz_norm
	if returnfps:
	return new_xyz, new_points, grouped_xyz, fps_idx
	else:
	return new_xyz, new_points
	",point_e\evals\pointnet2_utils.py
	sample_and_group_all,"Input:
	xyz: input points position data, [B, N, 3]
	points: input points data, [B, N, D]
	Return:
	new_xyz: sampled points position data, [B, 1, 3]
	new_points: sampled points data, [B, 1, N, 3+D]","def sample_and_group_all(xyz, points):
	""""""""""""
	device = xyz.device
	B, N, C = xyz.shape
	new_xyz = torch.zeros(B, 1, C).to(device)
	grouped_xyz = xyz.view(B, 1, N, C)
	if points is not None:
	new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
	else:
	new_points = grouped_xyz
	return new_xyz, new_points
	",point_e\evals\pointnet2_utils.py
	__init__,,"def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
	super(PointNetSetAbstraction, self).__init__()
	self.npoint = npoint
	self.radius = radius
	self.nsample = nsample
	self.mlp_convs = nn.ModuleList()
	self.mlp_bns = nn.ModuleList()
	last_channel = in_channel
	for out_channel in mlp:
	self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
	self.mlp_bns.append(nn.BatchNorm2d(out_channel))
	last_channel = out_channel
	self.group_all = group_all
	",point_e\evals\pointnet2_utils.py
	forward,"Input:
	xyz: input points position data, [B, C, N]
	points: input points data, [B, D, N]
	Return:
	new_xyz: sampled points position data, [B, C, S]
	new_points_concat: sample points feature data, [B, D', S]","def forward(self, xyz, points):
	""""""""""""
	xyz = xyz.permute(0, 2, 1)
	if points is not None:
	points = points.permute(0, 2, 1)
	if self.group_all:
	new_xyz, new_points = sample_and_group_all(xyz, points)
	else:
	new_xyz, new_points = sample_and_group(self.npoint, self.radius,
	self.nsample, xyz, points, deterministic=not self.training)
	new_points = new_points.permute(0, 3, 2, 1)
	for i, conv in enumerate(self.mlp_convs):
	bn = self.mlp_bns[i]
	new_points = F.relu(bn(conv(new_points)))
	new_points = torch.max(new_points, 2)[0]
	new_xyz = new_xyz.permute(0, 2, 1)
	return new_xyz, new_points
	",point_e\evals\pointnet2_utils.py
	__init__,,"def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list):
	super(PointNetSetAbstractionMsg, self).__init__()
	self.npoint = npoint
	self.radius_list = radius_list
	self.nsample_list = nsample_list
	self.conv_blocks = nn.ModuleList()
	self.bn_blocks = nn.ModuleList()
	for i in range(len(mlp_list)):
	convs = nn.ModuleList()
	bns = nn.ModuleList()
	last_channel = in_channel + 3
	for out_channel in mlp_list[i]:
	convs.append(nn.Conv2d(last_channel, out_channel, 1))
	bns.append(nn.BatchNorm2d(out_channel))
	last_channel = out_channel
	self.conv_blocks.append(convs)
	self.bn_blocks.append(bns)
	",point_e\evals\pointnet2_utils.py
	forward,"Input:
	xyz: input points position data, [B, C, N]
	points: input points data, [B, D, N]
	Return:
	new_xyz: sampled points position data, [B, C, S]
	new_points_concat: sample points feature data, [B, D', S]","def forward(self, xyz, points):
	""""""""""""
	xyz = xyz.permute(0, 2, 1)
	if points is not None:
	points = points.permute(0, 2, 1)
	B, N, C = xyz.shape
	S = self.npoint
	new_xyz = index_points(xyz, farthest_point_sample(xyz, S, deterministic
	=not self.training))
	new_points_list = []
	for i, radius in enumerate(self.radius_list):
	K = self.nsample_list[i]
	group_idx = query_ball_point(radius, K, xyz, new_xyz)
	grouped_xyz = index_points(xyz, group_idx)
	grouped_xyz -= new_xyz.view(B, S, 1, C)
	if points is not None:
	grouped_points = index_points(points, group_idx)
	grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
	else:
	grouped_points = grouped_xyz
	grouped_points = grouped_points.permute(0, 3, 2, 1)
	for j in range(len(self.conv_blocks[i])):
	conv = self.conv_blocks[i][j]
	bn = self.bn_blocks[i][j]
	grouped_points = F.relu(bn(conv(grouped_points)))
	new_points = torch.max(grouped_points, 2)[0]
	new_points_list.append(new_points)
	new_xyz = new_xyz.permute(0, 2, 1)
	new_points_concat = torch.cat(new_points_list, dim=1)
	return new_xyz, new_points_concat
	",point_e\evals\pointnet2_utils.py
	__init__,,"def __init__(self, in_channel, mlp):
	super(PointNetFeaturePropagation, self).__init__()
	self.mlp_convs = nn.ModuleList()
	self.mlp_bns = nn.ModuleList()
	last_channel = in_channel
	for out_channel in mlp:
	self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
	self.mlp_bns.append(nn.BatchNorm1d(out_channel))
	last_channel = out_channel
	",point_e\evals\pointnet2_utils.py
	forward,"Input:
	xyz1: input points position data, [B, C, N]
	xyz2: sampled input points position data, [B, C, S]
	points1: input points data, [B, D, N]
	points2: input points data, [B, D, S]
	Return:
	new_points: upsampled points data, [B, D', N]","def forward(self, xyz1, xyz2, points1, points2):
	""""""""""""
	xyz1 = xyz1.permute(0, 2, 1)
	xyz2 = xyz2.permute(0, 2, 1)
	points2 = points2.permute(0, 2, 1)
	B, N, C = xyz1.shape
	_, S, _ = xyz2.shape
	if S == 1:
	interpolated_points = points2.repeat(1, N, 1)
	else:
	dists = square_distance(xyz1, xyz2)
	dists, idx = dists.sort(dim=-1)
	dists, idx = dists[:, :, :3], idx[:, :, :3]
	dist_recip = 1.0 / (dists + 1e-08)
	norm = torch.sum(dist_recip, dim=2, keepdim=True)
	weight = dist_recip / norm
	interpolated_points = torch.sum(index_points(points2, idx) * weight
	.view(B, N, 3, 1), dim=2)
	if points1 is not None:
	points1 = points1.permute(0, 2, 1)
	new_points = torch.cat([points1, interpolated_points], dim=-1)
	else:
	new_points = interpolated_points
	new_points = new_points.permute(0, 2, 1)
	for i, conv in enumerate(self.mlp_convs):
	bn = self.mlp_bns[i]
	new_points = F.relu(bn(conv(new_points)))
	return new_points
	",point_e\evals\pointnet2_utils.py
	clear_scene,,"def clear_scene():
	bpy.ops.object.select_all(action='SELECT')
	bpy.ops.object.delete()
	",point_e\evals\scripts\blender_script.py
	clear_lights,,"def clear_lights():
	bpy.ops.object.select_all(action='DESELECT')
	for obj in bpy.context.scene.objects.values():
	if isinstance(obj.data, bpy.types.Light):
	obj.select_set(True)
	bpy.ops.object.delete()
	",point_e\evals\scripts\blender_script.py
	import_model,,"def import_model(path):
	clear_scene()
	_, ext = os.path.splitext(path)
	ext = ext.lower()
	if ext == '.obj':
	bpy.ops.import_scene.obj(filepath=path)
	elif ext in ['.glb', '.gltf']:
	bpy.ops.import_scene.gltf(filepath=path)
	elif ext == '.stl':
	bpy.ops.import_mesh.stl(filepath=path)
	elif ext == '.fbx':
	bpy.ops.import_scene.fbx(filepath=path)
	elif ext == '.dae':
	bpy.ops.wm.collada_import(filepath=path)
	elif ext == '.ply':
	bpy.ops.import_mesh.ply(filepath=path)
	else:
	raise RuntimeError(f'unexpected extension: {ext}')
	",point_e\evals\scripts\blender_script.py
	scene_root_objects,,"def scene_root_objects():
	for obj in bpy.context.scene.objects.values():
	if not obj.parent:
	yield obj
	",point_e\evals\scripts\blender_script.py
	scene_bbox,,"def scene_bbox(single_obj=None, ignore_matrix=False):
	bbox_min = (math.inf,) * 3
	bbox_max = (-math.inf,) * 3
	found = False
	for obj in (scene_meshes() if single_obj is None else [single_obj]):
	found = True
	for coord in obj.bound_box:
	coord = Vector(coord)
	if not ignore_matrix:
	coord = obj.matrix_world @ coord
	bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
	bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
	if not found:
	raise RuntimeError('no objects in scene to compute bounding box for')
	return Vector(bbox_min), Vector(bbox_max)
	",point_e\evals\scripts\blender_script.py
	scene_meshes,,"def scene_meshes():
	for obj in bpy.context.scene.objects.values():
	if isinstance(obj.data, bpy.types.Mesh):
	yield obj
	",point_e\evals\scripts\blender_script.py
	normalize_scene,,"def normalize_scene():
	bbox_min, bbox_max = scene_bbox()
	scale = 1 / max(bbox_max - bbox_min)
	for obj in scene_root_objects():
	obj.scale = obj.scale * scale
	bpy.context.view_layer.update()
	bbox_min, bbox_max = scene_bbox()
	offset = -(bbox_min + bbox_max) / 2
	for obj in scene_root_objects():
	obj.matrix_world.translation += offset
	bpy.ops.object.select_all(action='DESELECT')
	",point_e\evals\scripts\blender_script.py
	create_camera,,"def create_camera():
	camera_data = bpy.data.cameras.new(name='Camera')
	camera_object = bpy.data.objects.new('Camera', camera_data)
	bpy.context.scene.collection.objects.link(camera_object)
	bpy.context.scene.camera = camera_object
	",point_e\evals\scripts\blender_script.py
	set_camera,,"def set_camera(direction, camera_dist=2.0):
	camera_pos = -camera_dist * direction
	bpy.context.scene.camera.location = camera_pos
	rot_quat = direction.to_track_quat('-Z', 'Y')
	bpy.context.scene.camera.rotation_euler = rot_quat.to_euler()
	bpy.context.view_layer.update()
	",point_e\evals\scripts\blender_script.py
	randomize_camera,,"def randomize_camera(camera_dist=2.0):
	direction = random_unit_vector()
	set_camera(direction, camera_dist=camera_dist)
	",point_e\evals\scripts\blender_script.py
	pan_camera,,"def pan_camera(time, axis='Z', camera_dist=2.0, elevation=-0.1):
	angle = time * math.pi * 2
	direction = [-math.cos(angle), -math.sin(angle), -elevation]
	assert axis in ['X', 'Y', 'Z']
	if axis == 'X':
	direction = [direction[2], *direction[:2]]
	elif axis == 'Y':
	direction = [direction[0], -elevation, direction[1]]
	direction = Vector(direction).normalized()
	set_camera(direction, camera_dist=camera_dist)
	",point_e\evals\scripts\blender_script.py
	place_camera,,"def place_camera(time, camera_pose_mode='random', camera_dist_min=2.0,
	camera_dist_max=2.0):
	camera_dist = random.uniform(camera_dist_min, camera_dist_max)
	if camera_pose_mode == 'random':
	randomize_camera(camera_dist=camera_dist)
	elif camera_pose_mode == 'z-circular':
	pan_camera(time, axis='Z', camera_dist=camera_dist)
	elif camera_pose_mode == 'z-circular-elevated':
	pan_camera(time, axis='Z', camera_dist=camera_dist, elevation=
	0.2617993878)
	else:
	raise ValueError(f'Unknown camera pose mode: {camera_pose_mode}')
	",point_e\evals\scripts\blender_script.py
	create_light,,"def create_light(location, energy=1.0, angle=0.5 * math.pi / 180):
	light_data = bpy.data.lights.new(name='Light', type='SUN')
	light_data.energy = energy
	light_data.angle = angle
	light_object = bpy.data.objects.new(name='Light', object_data=light_data)
	direction = -location
	rot_quat = direction.to_track_quat('-Z', 'Y')
	light_object.rotation_euler = rot_quat.to_euler()
	bpy.context.view_layer.update()
	bpy.context.collection.objects.link(light_object)
	light_object.location = location
	",point_e\evals\scripts\blender_script.py
	create_random_lights,,"def create_random_lights(count=4, distance=2.0, energy=1.5):
	clear_lights()
	for _ in range(count):
	create_light(random_unit_vector() * distance, energy=energy)
	",point_e\evals\scripts\blender_script.py
	create_camera_light,,"def create_camera_light():
	clear_lights()
	create_light(bpy.context.scene.camera.location, energy=5.0)
	",point_e\evals\scripts\blender_script.py
	create_uniform_light,,"def create_uniform_light(backend):
	clear_lights()
	pos = Vector(UNIFORM_LIGHT_DIRECTION)
	angle = 0.0092 if backend == 'CYCLES' else math.pi
	create_light(pos, energy=5.0, angle=angle)
	create_light(-pos, energy=5.0, angle=angle)
	",point_e\evals\scripts\blender_script.py
	create_vertex_color_shaders,,"def create_vertex_color_shaders():
	for obj in bpy.context.scene.objects.values():
	if not isinstance(obj.data, bpy.types.Mesh):
	continue
	if len(obj.data.materials):
	continue
	color_keys = (obj.data.vertex_colors or {}).keys()
	if not len(color_keys):
	continue
	mat = bpy.data.materials.new(name='VertexColored')
	mat.use_nodes = True
	bsdf_node = None
	for node in mat.node_tree.nodes:
	if node.type == 'BSDF_PRINCIPLED':
	bsdf_node = node
	assert bsdf_node is not None, 'material has no Principled BSDF node to modify'
	socket_map = {}
	for input in bsdf_node.inputs:
	socket_map[input.name] = input
	socket_map['Specular'].default_value = 0.0
	socket_map['Roughness'].default_value = 1.0
	v_color = mat.node_tree.nodes.new('ShaderNodeVertexColor')
	v_color.layer_name = color_keys[0]
	mat.node_tree.links.new(v_color.outputs[0], socket_map['Base Color'])
	obj.data.materials.append(mat)
	",point_e\evals\scripts\blender_script.py
	create_default_materials,,"def create_default_materials():
	for obj in bpy.context.scene.objects.values():
	if isinstance(obj.data, bpy.types.Mesh):
	if not len(obj.data.materials):
	mat = bpy.data.materials.new(name='DefaultMaterial')
	mat.use_nodes = True
	obj.data.materials.append(mat)
	",point_e\evals\scripts\blender_script.py
	find_materials,,"def find_materials():
	all_materials = set()
	for obj in bpy.context.scene.objects.values():
	if not isinstance(obj.data, bpy.types.Mesh):
	continue
	for mat in obj.data.materials:
	all_materials.add(mat)
	return all_materials
	",point_e\evals\scripts\blender_script.py
	get_socket_value,,"def get_socket_value(tree, socket):
	default = socket.default_value
	if not isinstance(default, float):
	default = list(default)
	for link in tree.links:
	if link.to_socket == socket:
	return link.from_socket, default
	return None, default
	",point_e\evals\scripts\blender_script.py
	clear_socket_input,,"def clear_socket_input(tree, socket):
	for link in list(tree.links):
	if link.to_socket == socket:
	tree.links.remove(link)
	",point_e\evals\scripts\blender_script.py
	set_socket_value,,"def set_socket_value(tree, socket, socket_and_default):
	clear_socket_input(tree, socket)
	old_source_socket, default = socket_and_default
	if isinstance(default, float) and not isinstance(socket.default_value,
	float):
	socket.default_value = [default] * 3 + [1.0]
	else:
	socket.default_value = default
	if old_source_socket is not None:
	tree.links.new(old_source_socket, socket)
	",point_e\evals\scripts\blender_script.py
	setup_nodes,,"def setup_nodes(output_path, capturing_material_alpha: bool=False):
	tree = bpy.context.scene.node_tree
	links = tree.links
	for node in tree.nodes:
	tree.nodes.remove(node)

	def node_op(op: str, *args, clamp=False):
	node = tree.nodes.new(type='CompositorNodeMath')
	node.operation = op
	if clamp:
	node.use_clamp = True
	for i, arg in enumerate(args):
	if isinstance(arg, (int, float)):
	node.inputs[i].default_value = arg
	else:
	links.new(arg, node.inputs[i])
	return node.outputs[0]

	def node_clamp(x, maximum=1.0):
	return node_op('MINIMUM', x, maximum)

	def node_mul(x, y, **kwargs):
	return node_op('MULTIPLY', x, y, **kwargs)
	input_node = tree.nodes.new(type='CompositorNodeRLayers')
	input_node.scene = bpy.context.scene
	input_sockets = {}
	for output in input_node.outputs:
	input_sockets[output.name] = output
	if capturing_material_alpha:
	color_socket = input_sockets['Image']
	else:
	raw_color_socket = input_sockets['Image']
	color_node = tree.nodes.new(type='CompositorNodeConvertColorSpace')
	color_node.from_color_space = 'Linear'
	color_node.to_color_space = 'sRGB'
	tree.links.new(raw_color_socket, color_node.inputs[0])
	color_socket = color_node.outputs[0]
	split_node = tree.nodes.new(type='CompositorNodeSepRGBA')
	tree.links.new(color_socket, split_node.inputs[0])
	for i, channel in (enumerate('rgba') if not capturing_material_alpha else
	[(0, 'MatAlpha')]):
	output_node = tree.nodes.new(type='CompositorNodeOutputFile')
	output_node.base_path = f'{output_path}_{channel}'
	links.new(split_node.outputs[i], output_node.inputs[0])
	if capturing_material_alpha:
	return
	depth_out = node_clamp(node_mul(input_sockets['Depth'], 1 / MAX_DEPTH))
	output_node = tree.nodes.new(type='CompositorNodeOutputFile')
	output_node.base_path = f'{output_path}_depth'
	links.new(depth_out, output_node.inputs[0])
	",point_e\evals\scripts\blender_script.py
	render_scene,,"def render_scene(output_path, fast_mode: bool):
	use_workbench = bpy.context.scene.render.engine == 'BLENDER_WORKBENCH'
	if use_workbench:
	bpy.context.scene.render.engine = 'BLENDER_EEVEE'
	bpy.context.scene.eevee.taa_render_samples = 1
	if fast_mode:
	if bpy.context.scene.render.engine == 'BLENDER_EEVEE':
	bpy.context.scene.eevee.taa_render_samples = 1
	elif bpy.context.scene.render.engine == 'CYCLES':
	bpy.context.scene.cycles.samples = 256
	elif bpy.context.scene.render.engine == 'CYCLES':
	bpy.context.scene.cycles.time_limit = 40
	bpy.context.view_layer.update()
	bpy.context.scene.use_nodes = True
	bpy.context.scene.view_layers['ViewLayer'].use_pass_z = True
	bpy.context.scene.view_settings.view_transform = 'Raw'
	bpy.context.scene.render.film_transparent = True
	bpy.context.scene.render.resolution_x = 512
	bpy.context.scene.render.resolution_y = 512
	bpy.context.scene.render.image_settings.file_format = 'PNG'
	bpy.context.scene.render.image_settings.color_mode = 'BW'
	bpy.context.scene.render.image_settings.color_depth = '16'
	bpy.context.scene.render.filepath = output_path
	setup_nodes(output_path)
	bpy.ops.render.render(write_still=True)
	for channel_name in ['r', 'g', 'b', 'a', 'depth']:
	sub_dir = f'{output_path}_{channel_name}'
	image_path = os.path.join(sub_dir, os.listdir(sub_dir)[0])
	name, ext = os.path.splitext(output_path)
	if channel_name == 'depth' or not use_workbench:
	os.rename(image_path, f'{name}_{channel_name}{ext}')
	else:
	os.remove(image_path)
	os.removedirs(sub_dir)
	if use_workbench:
	bpy.context.scene.use_nodes = False
	bpy.context.scene.render.engine = 'BLENDER_WORKBENCH'
	bpy.context.scene.render.image_settings.color_mode = 'RGBA'
	bpy.context.scene.render.image_settings.color_depth = '8'
	bpy.context.scene.display.shading.color_type = 'TEXTURE'
	bpy.context.scene.display.shading.light = 'FLAT'
	if fast_mode:
	bpy.context.scene.display.render_aa = 'FXAA'
	os.remove(output_path)
	bpy.ops.render.render(write_still=True)
	bpy.context.scene.render.image_settings.color_mode = 'BW'
	bpy.context.scene.render.image_settings.color_depth = '16'
	",point_e\evals\scripts\blender_script.py
	scene_fov,,"def scene_fov():
	x_fov = bpy.context.scene.camera.data.angle_x
	y_fov = bpy.context.scene.camera.data.angle_y
	width = bpy.context.scene.render.resolution_x
	height = bpy.context.scene.render.resolution_y
	if bpy.context.scene.camera.data.angle == x_fov:
	y_fov = 2 * math.atan(math.tan(x_fov / 2) * height / width)
	else:
	x_fov = 2 * math.atan(math.tan(y_fov / 2) * width / height)
	return x_fov, y_fov
	",point_e\evals\scripts\blender_script.py
	write_camera_metadata,,"def write_camera_metadata(path):
	x_fov, y_fov = scene_fov()
	bbox_min, bbox_max = scene_bbox()
	matrix = bpy.context.scene.camera.matrix_world
	with open(path, 'w') as f:
	json.dump(dict(format_version=FORMAT_VERSION, max_depth=MAX_DEPTH,
	bbox=[list(bbox_min), list(bbox_max)], origin=list(matrix.col[3
	])[:3], x_fov=x_fov, y_fov=y_fov, x=list(matrix.col[0])[:3], y=
	list(-matrix.col[1])[:3], z=list(-matrix.col[2])[:3]), f)
	",point_e\evals\scripts\blender_script.py
	save_rendering_dataset,,"def save_rendering_dataset(input_path: str, output_path: str, num_images:
	int, backend: str, light_mode: str, camera_pose: str, camera_dist_min:
	float, camera_dist_max: float, fast_mode: bool):
	assert light_mode in ['random', 'uniform', 'camera']
	assert camera_pose in ['random', 'z-circular', 'z-circular-elevated']
	import_model(input_path)
	bpy.context.scene.render.engine = backend
	normalize_scene()
	if light_mode == 'random':
	create_random_lights()
	elif light_mode == 'uniform':
	create_uniform_light(backend)
	create_camera()
	create_vertex_color_shaders()
	for i in range(num_images):
	t = i / max(num_images - 1, 1)
	place_camera(t, camera_pose_mode=camera_pose, camera_dist_min=
	camera_dist_min, camera_dist_max=camera_dist_max)
	if light_mode == 'camera':
	create_camera_light()
	render_scene(os.path.join(output_path, f'{i:05}.png'), fast_mode=
	fast_mode)
	write_camera_metadata(os.path.join(output_path, f'{i:05}.json'))
	with open(os.path.join(output_path, 'info.json'), 'w') as f:
	info = dict(backend=backend, light_mode=light_mode, fast_mode=
	fast_mode, format_version=FORMAT_VERSION, channels=['R', 'G',
	'B', 'A', 'D'], scale=0.5)
	json.dump(info, f)
	",point_e\evals\scripts\blender_script.py
	main,,"def main():
	try:
	dash_index = sys.argv.index('--')
	except ValueError as exc:
	raise ValueError(""arguments must be preceded by '--'"") from exc
	raw_args = sys.argv[dash_index + 1:]
	parser = argparse.ArgumentParser()
	parser.add_argument('--input_path', required=True, type=str)
	parser.add_argument('--output_path', required=True, type=str)
	parser.add_argument('--num_images', type=int, default=20)
	parser.add_argument('--backend', type=str, default='BLENDER_EEVEE')
	parser.add_argument('--light_mode', type=str, default='uniform')
	parser.add_argument('--camera_pose', type=str, default='random')
	parser.add_argument('--camera_dist_min', type=float, default=2.0)
	parser.add_argument('--camera_dist_max', type=float, default=2.0)
	parser.add_argument('--fast_mode', action='store_true')
	args = parser.parse_args(raw_args)
	save_rendering_dataset(input_path=args.input_path, output_path=args.
	output_path, num_images=args.num_images, backend=args.backend,
	light_mode=args.light_mode, camera_pose=args.camera_pose,
	camera_dist_min=args.camera_dist_min, camera_dist_max=args.
	camera_dist_max, fast_mode=args.fast_mode)
	",point_e\evals\scripts\blender_script.py
	main,,"def main():
	parser = argparse.ArgumentParser()
	parser.add_argument('--cache_dir', type=str, default=None)
	parser.add_argument('batch_1', type=str)
	parser.add_argument('batch_2', type=str)
	args = parser.parse_args()
	print('creating classifier...')
	clf = PointNetClassifier(devices=get_torch_devices(), cache_dir=args.
	cache_dir)
	print('computing first batch activations')
	features_1, _ = clf.features_and_preds(NpzStreamer(args.batch_1))
	stats_1 = compute_statistics(features_1)
	del features_1
	features_2, _ = clf.features_and_preds(NpzStreamer(args.batch_2))
	stats_2 = compute_statistics(features_2)
	del features_2
	print(f'P-FID: {stats_1.frechet_distance(stats_2)}')
	",point_e\evals\scripts\evaluate_pfid.py
	main,,"def main():
	parser = argparse.ArgumentParser()
	parser.add_argument('--cache_dir', type=str, default=None)
	parser.add_argument('batch', type=str)
	args = parser.parse_args()
	print('creating classifier...')
	clf = PointNetClassifier(devices=get_torch_devices(), cache_dir=args.
	cache_dir)
	print('computing batch predictions')
	_, preds = clf.features_and_preds(NpzStreamer(args.batch))
	print(f'P-IS: {compute_inception_score(preds)}')
	",point_e\evals\scripts\evaluate_pis.py
	checkpoint,"Evaluate a function without caching intermediate activations, allowing for
	reduced memory at the expense of extra compute in the backward pass.
	:param func: the function to evaluate.
	:param inputs: the argument sequence to pass to `func`.
	:param params: a sequence of parameters `func` depends on but does not
	explicitly take as arguments.
	:param flag: if False, disable gradient checkpointing.","def checkpoint(func: Callable[..., Union[torch.Tensor, Sequence[torch.
	Tensor]]], inputs: Sequence[torch.Tensor], params: Iterable[torch.
	Tensor], flag: bool):
	""""""""""""
	if flag:
	args = tuple(inputs) + tuple(params)
	return CheckpointFunction.apply(func, len(inputs), *args)
	else:
	return func(*inputs)
	",point_e\models\checkpoint.py
	forward,,"@staticmethod
	def forward(ctx, run_function, length, *args):
	ctx.run_function = run_function
	ctx.input_tensors = list(args[:length])
	ctx.input_params = list(args[length:])
	with torch.no_grad():
	output_tensors = ctx.run_function(*ctx.input_tensors)
	return output_tensors
	",point_e\models\checkpoint.py
	backward,,"@staticmethod
	def backward(ctx, *output_grads):
	ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.
	input_tensors]
	with torch.enable_grad():
	shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
	output_tensors = ctx.run_function(*shallow_copies)
	input_grads = torch.autograd.grad(output_tensors, ctx.input_tensors +
	ctx.input_params, output_grads, allow_unused=True)
	del ctx.input_tensors
	del ctx.input_params
	del output_tensors
	return (None, None) + input_grads
	",point_e\models\checkpoint.py
	model_from_config,,"def model_from_config(config: Dict[str, Any], device: torch.device
	) ->nn.Module:
	config = config.copy()
	name = config.pop('name')
	if name == 'PointDiffusionTransformer':
	return PointDiffusionTransformer(device=device, dtype=torch.float32,
	**config)
	elif name == 'CLIPImagePointDiffusionTransformer':
	return CLIPImagePointDiffusionTransformer(device=device, dtype=
	torch.float32, **config)
	elif name == 'CLIPImageGridPointDiffusionTransformer':
	return CLIPImageGridPointDiffusionTransformer(device=device, dtype=
	torch.float32, **config)
	elif name == 'UpsamplePointDiffusionTransformer':
	return UpsamplePointDiffusionTransformer(device=device, dtype=torch
	.float32, **config)
	elif name == 'CLIPImageGridUpsamplePointDiffusionTransformer':
	return CLIPImageGridUpsamplePointDiffusionTransformer(device=device,
	dtype=torch.float32, **config)
	elif name == 'CrossAttentionPointCloudSDFModel':
	return CrossAttentionPointCloudSDFModel(device=device, dtype=torch.
	float32, **config)
	raise ValueError(f'unknown model name: {name}')
	",point_e\models\configs.py
	default_cache_dir,,"@lru_cache()
	def default_cache_dir() ->str:
	return os.path.join(os.path.abspath(os.getcwd()), 'point_e_model_cache')
	",point_e\models\download.py
	fetch_file_cached,"Download the file at the given URL into a local file and return the path.
	If cache_dir is specified, it will be used to download the files.
	Otherwise, default_cache_dir() is used.","def fetch_file_cached(url: str, progress: bool=True, cache_dir: Optional[
	str]=None, chunk_size: int=4096) ->str:
	""""""""""""
	if cache_dir is None:
	cache_dir = default_cache_dir()
	os.makedirs(cache_dir, exist_ok=True)
	local_path = os.path.join(cache_dir, url.split('/')[-1])
	if os.path.exists(local_path):
	return local_path
	response = requests.get(url, stream=True)
	size = int(response.headers.get('content-length', '0'))
	with FileLock(local_path + '.lock'):
	if progress:
	pbar = tqdm(total=size, unit='iB', unit_scale=True)
	tmp_path = local_path + '.tmp'
	with open(tmp_path, 'wb') as f:
	for chunk in response.iter_content(chunk_size):
	if progress:
	pbar.update(len(chunk))
	f.write(chunk)
	os.rename(tmp_path, local_path)
	if progress:
	pbar.close()
	return local_path
	",point_e\models\download.py
	load_checkpoint,,"def load_checkpoint(checkpoint_name: str, device: torch.device, progress:
	bool=True, cache_dir: Optional[str]=None, chunk_size: int=4096) ->Dict[
	str, torch.Tensor]:
	if checkpoint_name not in MODEL_PATHS:
	raise ValueError(
	f'Unknown checkpoint name {checkpoint_name}. Known names are: {MODEL_PATHS.keys()}.'
	)
	path = fetch_file_cached(MODEL_PATHS[checkpoint_name], progress=
	progress, cache_dir=cache_dir, chunk_size=chunk_size)
	return torch.load(path, map_location=device)
	",point_e\models\download.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_data: int,
	width: int, heads: int, init_scale: float, data_width: Optional[int]=None):
	super().__init__()
	self.n_data = n_data
	self.width = width
	self.heads = heads
	self.data_width = width if data_width is None else data_width
	self.c_q = nn.Linear(width, width, device=device, dtype=dtype)
	self.c_kv = nn.Linear(self.data_width, width * 2, device=device, dtype=
	dtype)
	self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
	self.attention = QKVMultiheadCrossAttention(device=device, dtype=dtype,
	heads=heads, n_data=n_data)
	init_linear(self.c_q, init_scale)
	init_linear(self.c_kv, init_scale)
	init_linear(self.c_proj, init_scale)
	",point_e\models\perceiver.py
	forward,,"def forward(self, x, data):
	x = self.c_q(x)
	data = self.c_kv(data)
	x = checkpoint(self.attention, (x, data), (), True)
	x = self.c_proj(x)
	return x
	",point_e\models\perceiver.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int,
	n_data: int):
	super().__init__()
	self.device = device
	self.dtype = dtype
	self.heads = heads
	self.n_data = n_data
	",point_e\models\perceiver.py
	forward,,"def forward(self, q, kv):
	_, n_ctx, _ = q.shape
	bs, n_data, width = kv.shape
	attn_ch = width // self.heads // 2
	scale = 1 / math.sqrt(math.sqrt(attn_ch))
	q = q.view(bs, n_ctx, self.heads, -1)
	kv = kv.view(bs, n_data, self.heads, -1)
	k, v = torch.split(kv, attn_ch, dim=-1)
	weight = torch.einsum('bthc,bshc->bhts', q * scale, k * scale)
	wdtype = weight.dtype
	weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
	return torch.einsum('bhts,bshc->bthc', weight, v).reshape(bs, n_ctx, -1)
	",point_e\models\perceiver.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_data: int,
	width: int, heads: int, data_width: Optional[int]=None, init_scale:
	float=1.0):
	super().__init__()
	if data_width is None:
	data_width = width
	self.attn = MultiheadCrossAttention(device=device, dtype=dtype, n_data=
	n_data, width=width, heads=heads, data_width=data_width, init_scale
	=init_scale)
	self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
	self.ln_2 = nn.LayerNorm(data_width, device=device, dtype=dtype)
	self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=
	init_scale)
	self.ln_3 = nn.LayerNorm(width, device=device, dtype=dtype)
	",point_e\models\perceiver.py
	forward,,"def forward(self, x: torch.Tensor, data: torch.Tensor):
	x = x + self.attn(self.ln_1(x), self.ln_2(data))
	x = x + self.mlp(self.ln_3(x))
	return x
	",point_e\models\perceiver.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_data: int,
	width: int, layers: int, heads: int, init_scale: float=0.25, data_width:
	Optional[int]=None):
	super().__init__()
	self.width = width
	self.layers = layers
	init_scale = init_scale * math.sqrt(1.0 / width)
	self.resblocks = nn.ModuleList([ResidualCrossAttentionBlock(device=
	device, dtype=dtype, n_data=n_data, width=width, heads=heads,
	init_scale=init_scale, data_width=data_width) for _ in range(layers)])
	",point_e\models\perceiver.py
	forward,,"def forward(self, x: torch.Tensor, data: torch.Tensor):
	for block in self.resblocks:
	x = block(x, data)
	return x
	",point_e\models\perceiver.py
	_image_to_pil,,"def _image_to_pil(obj: Optional[ImageType]) ->Image.Image:
	if obj is None:
	return Image.fromarray(np.zeros([64, 64, 3], dtype=np.uint8))
	if isinstance(obj, np.ndarray):
	return Image.fromarray(obj.astype(np.uint8))
	elif isinstance(obj, torch.Tensor):
	return Image.fromarray(obj.detach().cpu().numpy().astype(np.uint8))
	else:
	return obj
	",point_e\models\pretrained_clip.py
	__init__,,"def __init__(self, device: torch.device, dtype: Optional[torch.dtype]=torch
	.float32, ensure_used_params: bool=True, clip_name: str='ViT-L/14',
	cache_dir: Optional[str]=None):
	super().__init__()
	assert clip_name in ['ViT-L/14', 'ViT-B/32']
	self.device = device
	self.ensure_used_params = ensure_used_params
	import clip
	self.clip_model, self.preprocess = clip.load(clip_name, device=device,
	download_root=cache_dir or default_cache_dir())
	self.clip_name = clip_name
	if dtype is not None:
	self.clip_model.to(dtype)
	self._tokenize = clip.tokenize
	",point_e\models\pretrained_clip.py
	feature_dim,,"@property
	def feature_dim(self) ->int:
	if self.clip_name == 'ViT-L/14':
	return 768
	else:
	return 512
	",point_e\models\pretrained_clip.py
	grid_size,,"@property
	def grid_size(self) ->int:
	if self.clip_name == 'ViT-L/14':
	return 16
	else:
	return 7
	",point_e\models\pretrained_clip.py
	grid_feature_dim,,"@property
	def grid_feature_dim(self) ->int:
	if self.clip_name == 'ViT-L/14':
	return 1024
	else:
	return 768
	",point_e\models\pretrained_clip.py
	forward,"Generate a batch of embeddings from a mixture of images, texts,
	precomputed embeddings, and possibly empty values.

	For each batch element, at most one of images, texts, and embeddings
	should have a non-None value. Embeddings from multiple modalities
	cannot be mixed for a single batch element. If no modality is provided,
	a zero embedding will be used for the batch element.","def forward(self, batch_size: int, images: Optional[Iterable[Optional[
	ImageType]]]=None, texts: Optional[Iterable[Optional[str]]]=None,
	embeddings: Optional[Iterable[Optional[torch.Tensor]]]=None
	) ->torch.Tensor:
	""""""""""""
	image_seq = [None] * batch_size if images is None else list(images)
	text_seq = [None] * batch_size if texts is None else list(texts)
	embedding_seq = [None] * batch_size if embeddings is None else list(
	embeddings)
	assert len(image_seq
	) == batch_size, 'number of images should match batch size'
	assert len(text_seq
	) == batch_size, 'number of texts should match batch size'
	assert len(embedding_seq
	) == batch_size, 'number of embeddings should match batch size'
	if self.ensure_used_params:
	return self._static_multimodal_embed(images=image_seq, texts=
	text_seq, embeddings=embedding_seq)
	result = torch.zeros((batch_size, self.feature_dim), device=self.device)
	index_images = []
	index_texts = []
	for i, (image, text, emb) in enumerate(zip(image_seq, text_seq,
	embedding_seq)):
	assert sum([int(image is not None), int(text is not None), int(emb
	is not None)]
	) < 2, 'only one modality may be non-None per batch element'
	if image is not None:
	index_images.append((i, image))
	elif text is not None:
	index_texts.append((i, text))
	elif emb is not None:
	result[i] = emb.to(result)
	if len(index_images):
	embs = self.embed_images(img for _, img in index_images)
	for (i, _), emb in zip(index_images, embs):
	result[i] = emb.to(result)
	if len(index_texts):
	embs = self.embed_text(text for _, text in index_texts)
	for (i, _), emb in zip(index_texts, embs):
	result[i] = emb.to(result)
	return result
	",point_e\models\pretrained_clip.py
	_static_multimodal_embed,"Like forward(), but always runs all encoders to ensure that
	the forward graph looks the same on every rank.","def _static_multimodal_embed(self, images: List[Optional[ImageType]]=None,
	texts: List[Optional[str]]=None, embeddings: List[Optional[torch.Tensor
	]]=None) ->torch.Tensor:
	""""""""""""
	image_emb = self.embed_images(images)
	text_emb = self.embed_text(t if t else '' for t in texts)
	joined_embs = torch.stack([(emb.to(device=self.device, dtype=torch.
	float32) if emb is not None else torch.zeros(self.feature_dim,
	device=self.device)) for emb in embeddings], dim=0)
	image_flag = torch.tensor([(x is not None) for x in images], device=
	self.device)[:, None].expand_as(image_emb)
	text_flag = torch.tensor([(x is not None) for x in texts], device=self.
	device)[:, None].expand_as(image_emb)
	emb_flag = torch.tensor([(x is not None) for x in embeddings], device=
	self.device)[:, None].expand_as(image_emb)
	return image_flag.float() * image_emb + text_flag.float(
	) * text_emb + emb_flag.float(
	) * joined_embs + self.clip_model.logit_scale * 0
	",point_e\models\pretrained_clip.py
	embed_images,":param xs: N images, stored as numpy arrays, tensors, or PIL images.
	:return: an [N x D] tensor of features.","def embed_images(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
	""""""""""""
	clip_inputs = self.images_to_tensor(xs)
	results = self.clip_model.encode_image(clip_inputs).float()
	return results / torch.linalg.norm(results, dim=-1, keepdim=True)
	",point_e\models\pretrained_clip.py
	embed_text,Embed text prompts as an [N x D] tensor.,"def embed_text(self, prompts: Iterable[str]) ->torch.Tensor:
	""""""""""""
	enc = self.clip_model.encode_text(self._tokenize(list(prompts),
	truncate=True).to(self.device)).float()
	return enc / torch.linalg.norm(enc, dim=-1, keepdim=True)
	",point_e\models\pretrained_clip.py
	embed_images_grid,"Embed images into latent grids.

	:param xs: an iterable of images to embed.
	:return: a tensor of shape [N x C x L], where L = self.grid_size**2.","def embed_images_grid(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
	""""""""""""
	if self.ensure_used_params:
	extra_value = 0.0
	for p in self.parameters():
	extra_value = extra_value + p.mean() * 0.0
	else:
	extra_value = 0.0
	x = self.images_to_tensor(xs).to(self.clip_model.dtype)
	vt = self.clip_model.visual
	x = vt.conv1(x)
	x = x.reshape(x.shape[0], x.shape[1], -1)
	x = x.permute(0, 2, 1)
	x = torch.cat([vt.class_embedding.to(x.dtype) + torch.zeros(x.shape[0],
	1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)
	x = x + vt.positional_embedding.to(x.dtype)
	x = vt.ln_pre(x)
	x = x.permute(1, 0, 2)
	x = vt.transformer(x)
	x = x.permute(1, 2, 0)
	return x[..., 1:].contiguous().float() + extra_value
	",point_e\models\pretrained_clip.py
	images_to_tensor,,"def images_to_tensor(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
	return torch.stack([self.preprocess(_image_to_pil(x)) for x in xs], dim=0
	).to(self.device)
	",point_e\models\pretrained_clip.py
	__init__,,"def __init__(self, device: torch.device, **kwargs):
	self.model = ImageCLIP(device, dtype=None, ensure_used_params=False, **
	kwargs)
	for parameter in self.model.parameters():
	parameter.requires_grad_(False)
	",point_e\models\pretrained_clip.py
	feature_dim,,"@property
	def feature_dim(self) ->int:
	return self.model.feature_dim
	",point_e\models\pretrained_clip.py
	grid_size,,"@property
	def grid_size(self) ->int:
	return self.model.grid_size
	",point_e\models\pretrained_clip.py
	grid_feature_dim,,"@property
	def grid_feature_dim(self) ->int:
	return self.model.grid_feature_dim
	",point_e\models\pretrained_clip.py
	__call__,,"def __call__(self, batch_size: int, images: Optional[Iterable[Optional[
	ImageType]]]=None, texts: Optional[Iterable[Optional[str]]]=None,
	embeddings: Optional[Iterable[Optional[torch.Tensor]]]=None
	) ->torch.Tensor:
	return self.model(batch_size=batch_size, images=images, texts=texts,
	embeddings=embeddings)
	",point_e\models\pretrained_clip.py
	embed_images,,"def embed_images(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
	with torch.no_grad():
	return self.model.embed_images(xs)
	",point_e\models\pretrained_clip.py
	embed_text,,"def embed_text(self, prompts: Iterable[str]) ->torch.Tensor:
	with torch.no_grad():
	return self.model.embed_text(prompts)
	",point_e\models\pretrained_clip.py
	embed_images_grid,,"def embed_images_grid(self, xs: Iterable[Optional[ImageType]]) ->torch.Tensor:
	with torch.no_grad():
	return self.model.embed_images_grid(xs)
	",point_e\models\pretrained_clip.py
	device,Get the device that should be used for input tensors.,"@property
	@abstractmethod
	def device(self) ->torch.device:
	""""""""""""
	",point_e\models\sdf.py
	default_batch_size,"Get a reasonable default number of query points for the model.
	In some cases, this might be the only supported size.","@property
	@abstractmethod
	def default_batch_size(self) ->int:
	""""""""""""
	",point_e\models\sdf.py
	encode_point_clouds,"Encode a batch of point clouds to cache part of the SDF calculation
	done by forward().

	:param point_clouds: a batch of [batch x 3 x N] points.
	:return: a state representing the encoded point cloud batch.","@abstractmethod
	def encode_point_clouds(self, point_clouds: torch.Tensor) ->Dict[str, torch
	.Tensor]:
	""""""""""""
	",point_e\models\sdf.py
	forward,"Predict the SDF at the coordinates x, given a batch of point clouds.

	Either point_clouds or encoded should be passed. Only exactly one of
	these arguments should be None.

	:param x: a [batch x 3 x N'] tensor of query points.
	:param point_clouds: a [batch x 3 x N] batch of point clouds.
	:param encoded: the result of calling encode_point_clouds().
	:return: a [batch x N'] tensor of SDF predictions.","def forward(self, x: torch.Tensor, point_clouds: Optional[torch.Tensor]=
	None, encoded: Optional[Dict[str, torch.Tensor]]=None) ->torch.Tensor:
	""""""""""""
	assert point_clouds is not None or encoded is not None
	assert point_clouds is None or encoded is None
	if point_clouds is not None:
	encoded = self.encode_point_clouds(point_clouds)
	return self.predict_sdf(x, encoded)
	",point_e\models\sdf.py
	predict_sdf,"Predict the SDF at the query points given the encoded point clouds.

	Each query point should be treated independently, only conditioning on
	the point clouds themselves.","@abstractmethod
	def predict_sdf(self, x: torch.Tensor, encoded: Optional[Dict[str, torch.
	Tensor]]) ->torch.Tensor:
	""""""""""""
	",point_e\models\sdf.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
	4096, width: int=512, encoder_layers: int=12, encoder_heads: int=8,
	decoder_layers: int=4, decoder_heads: int=8, init_scale: float=0.25):
	super().__init__()
	self._device = device
	self.n_ctx = n_ctx
	self.encoder_input_proj = nn.Linear(3, width, device=device, dtype=dtype)
	self.encoder = Transformer(device=device, dtype=dtype, n_ctx=n_ctx,
	width=width, layers=encoder_layers, heads=encoder_heads, init_scale
	=init_scale)
	self.decoder_input_proj = nn.Linear(3, width, device=device, dtype=dtype)
	self.decoder = SimplePerceiver(device=device, dtype=dtype, n_data=n_ctx,
	width=width, layers=decoder_layers, heads=decoder_heads, init_scale
	=init_scale)
	self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
	self.output_proj = nn.Linear(width, 1, device=device, dtype=dtype)
	",point_e\models\sdf.py
	device,,"@property
	def device(self) ->torch.device:
	return self._device
	",point_e\models\sdf.py
	default_batch_size,,"@property
	def default_batch_size(self) ->int:
	return self.n_query
	",point_e\models\sdf.py
	encode_point_clouds,,"def encode_point_clouds(self, point_clouds: torch.Tensor) ->Dict[str, torch
	.Tensor]:
	h = self.encoder_input_proj(point_clouds.permute(0, 2, 1))
	h = self.encoder(h)
	return dict(latents=h)
	",point_e\models\sdf.py
	predict_sdf,,"def predict_sdf(self, x: torch.Tensor, encoded: Optional[Dict[str, torch.
	Tensor]]) ->torch.Tensor:
	data = encoded['latents']
	x = self.decoder_input_proj(x.permute(0, 2, 1))
	x = self.decoder(x, data)
	x = self.ln_post(x)
	x = self.output_proj(x)
	return x[..., 0]
	",point_e\models\sdf.py
	init_linear,,"def init_linear(l, stddev):
	nn.init.normal_(l.weight, std=stddev)
	if l.bias is not None:
	nn.init.constant_(l.bias, 0.0)
	",point_e\models\transformer.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int,
	width: int, heads: int, init_scale: float):
	super().__init__()
	self.n_ctx = n_ctx
	self.width = width
	self.heads = heads
	self.c_qkv = nn.Linear(width, width * 3, device=device, dtype=dtype)
	self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
	self.attention = QKVMultiheadAttention(device=device, dtype=dtype,
	heads=heads, n_ctx=n_ctx)
	init_linear(self.c_qkv, init_scale)
	init_linear(self.c_proj, init_scale)
	",point_e\models\transformer.py
	forward,,"def forward(self, x):
	x = self.c_qkv(x)
	x = checkpoint(self.attention, (x,), (), True)
	x = self.c_proj(x)
	return x
	",point_e\models\transformer.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, width: int,
	init_scale: float):
	super().__init__()
	self.width = width
	self.c_fc = nn.Linear(width, width * 4, device=device, dtype=dtype)
	self.c_proj = nn.Linear(width * 4, width, device=device, dtype=dtype)
	self.gelu = nn.GELU()
	init_linear(self.c_fc, init_scale)
	init_linear(self.c_proj, init_scale)
	",point_e\models\transformer.py
	forward,,"def forward(self, x):
	return self.c_proj(self.gelu(self.c_fc(x)))
	",point_e\models\transformer.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int,
	n_ctx: int):
	super().__init__()
	self.device = device
	self.dtype = dtype
	self.heads = heads
	self.n_ctx = n_ctx
	",point_e\models\transformer.py
	forward,,"def forward(self, qkv):
	bs, n_ctx, width = qkv.shape
	attn_ch = width // self.heads // 3
	scale = 1 / math.sqrt(math.sqrt(attn_ch))
	qkv = qkv.view(bs, n_ctx, self.heads, -1)
	q, k, v = torch.split(qkv, attn_ch, dim=-1)
	weight = torch.einsum('bthc,bshc->bhts', q * scale, k * scale)
	wdtype = weight.dtype
	weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
	return torch.einsum('bhts,bshc->bthc', weight, v).reshape(bs, n_ctx, -1)
	",point_e\models\transformer.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int,
	width: int, heads: int, init_scale: float=1.0):
	super().__init__()
	self.attn = MultiheadAttention(device=device, dtype=dtype, n_ctx=n_ctx,
	width=width, heads=heads, init_scale=init_scale)
	self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
	self.mlp = MLP(device=device, dtype=dtype, width=width, init_scale=
	init_scale)
	self.ln_2 = nn.LayerNorm(width, device=device, dtype=dtype)
	",point_e\models\transformer.py
	forward,,"def forward(self, x: torch.Tensor):
	x = x + self.attn(self.ln_1(x))
	x = x + self.mlp(self.ln_2(x))
	return x
	",point_e\models\transformer.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int,
	width: int, layers: int, heads: int, init_scale: float=0.25):
	super().__init__()
	self.n_ctx = n_ctx
	self.width = width
	self.layers = layers
	init_scale = init_scale * math.sqrt(1.0 / width)
	self.resblocks = nn.ModuleList([ResidualAttentionBlock(device=device,
	dtype=dtype, n_ctx=n_ctx, width=width, heads=heads, init_scale=
	init_scale) for _ in range(layers)])
	",point_e\models\transformer.py
	forward,,"def forward(self, x: torch.Tensor):
	for block in self.resblocks:
	x = block(x)
	return x
	",point_e\models\transformer.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype,
	input_channels: int=3, output_channels: int=3, n_ctx: int=1024, width:
	int=512, layers: int=12, heads: int=8, init_scale: float=0.25,
	time_token_cond: bool=False):
	super().__init__()
	self.input_channels = input_channels
	self.output_channels = output_channels
	self.n_ctx = n_ctx
	self.time_token_cond = time_token_cond
	self.time_embed = MLP(device=device, dtype=dtype, width=width,
	init_scale=init_scale * math.sqrt(1.0 / width))
	self.ln_pre = nn.LayerNorm(width, device=device, dtype=dtype)
	self.backbone = Transformer(device=device, dtype=dtype, n_ctx=n_ctx +
	int(time_token_cond), width=width, layers=layers, heads=heads,
	init_scale=init_scale)
	self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
	self.input_proj = nn.Linear(input_channels, width, device=device, dtype
	=dtype)
	self.output_proj = nn.Linear(width, output_channels, device=device,
	dtype=dtype)
	with torch.no_grad():
	self.output_proj.weight.zero_()
	self.output_proj.bias.zero_()
	",point_e\models\transformer.py
	forward,":param x: an [N x C x T] tensor.
	:param t: an [N] tensor.
	:return: an [N x C' x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor):
	""""""""""""
	assert x.shape[-1] == self.n_ctx
	t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
	return self._forward_with_cond(x, [(t_embed, self.time_token_cond)])
	",point_e\models\transformer.py
	_forward_with_cond,,"def _forward_with_cond(self, x: torch.Tensor, cond_as_token: List[Tuple[
	torch.Tensor, bool]]) ->torch.Tensor:
	h = self.input_proj(x.permute(0, 2, 1))
	for emb, as_token in cond_as_token:
	if not as_token:
	h = h + emb[:, None]
	extra_tokens = [(emb[:, None] if len(emb.shape) == 2 else emb) for emb,
	as_token in cond_as_token if as_token]
	if len(extra_tokens):
	h = torch.cat(extra_tokens + [h], dim=1)
	h = self.ln_pre(h)
	h = self.backbone(h)
	h = self.ln_post(h)
	if len(extra_tokens):
	h = h[:, sum(h.shape[1] for h in extra_tokens):]
	h = self.output_proj(h)
	return h.permute(0, 2, 1)
	",point_e\models\transformer.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
	1024, token_cond: bool=False, cond_drop_prob: float=0.0, frozen_clip:
	bool=True, cache_dir: Optional[str]=None, **kwargs):
	super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + int(
	token_cond), **kwargs)
	self.n_ctx = n_ctx
	self.token_cond = token_cond
	self.clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device,
	cache_dir=cache_dir)
	self.clip_embed = nn.Linear(self.clip.feature_dim, self.backbone.width,
	device=device, dtype=dtype)
	self.cond_drop_prob = cond_drop_prob
	",point_e\models\transformer.py
	cached_model_kwargs,,"def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]
	) ->Dict[str, Any]:
	with torch.no_grad():
	return dict(embeddings=self.clip(batch_size, **model_kwargs))
	",point_e\models\transformer.py
	forward,":param x: an [N x C x T] tensor.
	:param t: an [N] tensor.
	:param images: a batch of images to condition on.
	:param texts: a batch of texts to condition on.
	:param embeddings: a batch of CLIP embeddings to condition on.
	:return: an [N x C' x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, images: Optional[
	Iterable[Optional[ImageType]]]=None, texts: Optional[Iterable[Optional[
	str]]]=None, embeddings: Optional[Iterable[Optional[torch.Tensor]]]=None):
	""""""""""""
	assert x.shape[-1] == self.n_ctx
	t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
	clip_out = self.clip(batch_size=len(x), images=images, texts=texts,
	embeddings=embeddings)
	assert len(clip_out.shape) == 2 and clip_out.shape[0] == x.shape[0]
	if self.training:
	mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
	clip_out = clip_out * mask[:, None].to(clip_out)
	clip_out = math.sqrt(clip_out.shape[1]) * clip_out
	clip_embed = self.clip_embed(clip_out)
	cond = [(clip_embed, self.token_cond), (t_embed, self.time_token_cond)]
	return self._forward_with_cond(x, cond)
	",point_e\models\transformer.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
	1024, cond_drop_prob: float=0.0, frozen_clip: bool=True, cache_dir:
	Optional[str]=None, **kwargs):
	clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device,
	cache_dir=cache_dir)
	super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + clip.
	grid_size 2, kwargs)
	self.n_ctx = n_ctx
	self.clip = clip
	self.clip_embed = nn.Sequential(nn.LayerNorm(normalized_shape=(self.
	clip.grid_feature_dim,), device=device, dtype=dtype), nn.Linear(
	self.clip.grid_feature_dim, self.backbone.width, device=device,
	dtype=dtype))
	self.cond_drop_prob = cond_drop_prob
	",point_e\models\transformer.py
	cached_model_kwargs,,"def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]
	) ->Dict[str, Any]:
	_ = batch_size
	with torch.no_grad():
	return dict(embeddings=self.clip.embed_images_grid(model_kwargs[
	'images']))
	",point_e\models\transformer.py
	forward,":param x: an [N x C x T] tensor.
	:param t: an [N] tensor.
	:param images: a batch of images to condition on.
	:param embeddings: a batch of CLIP latent grids to condition on.
	:return: an [N x C' x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, images: Optional[
	Iterable[ImageType]]=None, embeddings: Optional[Iterable[torch.Tensor]]
	=None):
	""""""""""""
	assert images is not None or embeddings is not None, 'must specify images or embeddings'
	assert images is None or embeddings is None, 'cannot specify both images and embeddings'
	assert x.shape[-1] == self.n_ctx
	t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
	if images is not None:
	clip_out = self.clip.embed_images_grid(images)
	else:
	clip_out = embeddings
	if self.training:
	mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
	clip_out = clip_out * mask[:, None, None].to(clip_out)
	clip_out = clip_out.permute(0, 2, 1)
	clip_embed = self.clip_embed(clip_out)
	cond = [(t_embed, self.time_token_cond), (clip_embed, True)]
	return self._forward_with_cond(x, cond)
	",point_e\models\transformer.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype,
	cond_input_channels: Optional[int]=None, cond_ctx: int=1024, n_ctx: int
	=4096 - 1024, channel_scales: Optional[Sequence[float]]=None,
	channel_biases: Optional[Sequence[float]]=None, **kwargs):
	super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + cond_ctx, **
	kwargs)
	self.n_ctx = n_ctx
	self.cond_input_channels = cond_input_channels or self.input_channels
	self.cond_point_proj = nn.Linear(self.cond_input_channels, self.
	backbone.width, device=device, dtype=dtype)
	self.register_buffer('channel_scales', torch.tensor(channel_scales,
	dtype=dtype, device=device) if channel_scales is not None else None)
	self.register_buffer('channel_biases', torch.tensor(channel_biases,
	dtype=dtype, device=device) if channel_biases is not None else None)
	",point_e\models\transformer.py
	forward,":param x: an [N x C1 x T] tensor.
	:param t: an [N] tensor.
	:param low_res: an [N x C2 x T'] tensor of conditioning points.
	:return: an [N x C3 x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, *, low_res: torch.Tensor):
	""""""""""""
	assert x.shape[-1] == self.n_ctx
	t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
	low_res_embed = self._embed_low_res(low_res)
	cond = [(t_embed, self.time_token_cond), (low_res_embed, True)]
	return self._forward_with_cond(x, cond)
	",point_e\models\transformer.py
	_embed_low_res,,"def _embed_low_res(self, x: torch.Tensor) ->torch.Tensor:
	if self.channel_scales is not None:
	x = x * self.channel_scales[None, :, None]
	if self.channel_biases is not None:
	x = x + self.channel_biases[None, :, None]
	return self.cond_point_proj(x.permute(0, 2, 1))
	",point_e\models\transformer.py
	__init__,,"def __init__(self, *, device: torch.device, dtype: torch.dtype, n_ctx: int=
	4096 - 1024, cond_drop_prob: float=0.0, frozen_clip: bool=True,
	cache_dir: Optional[str]=None, **kwargs):
	clip = (FrozenImageCLIP if frozen_clip else ImageCLIP)(device,
	cache_dir=cache_dir)
	super().__init__(device=device, dtype=dtype, n_ctx=n_ctx + clip.
	grid_size 2, kwargs)
	self.n_ctx = n_ctx
	self.clip = clip
	self.clip_embed = nn.Sequential(nn.LayerNorm(normalized_shape=(self.
	clip.grid_feature_dim,), device=device, dtype=dtype), nn.Linear(
	self.clip.grid_feature_dim, self.backbone.width, device=device,
	dtype=dtype))
	self.cond_drop_prob = cond_drop_prob
	",point_e\models\transformer.py
	cached_model_kwargs,,"def cached_model_kwargs(self, batch_size: int, model_kwargs: Dict[str, Any]
	) ->Dict[str, Any]:
	if 'images' not in model_kwargs:
	zero_emb = torch.zeros([batch_size, self.clip.grid_feature_dim,
	self.clip.grid_size ** 2], device=next(self.parameters()).device)
	return dict(embeddings=zero_emb, low_res=model_kwargs['low_res'])
	with torch.no_grad():
	return dict(embeddings=self.clip.embed_images_grid(model_kwargs[
	'images']), low_res=model_kwargs['low_res'])
	",point_e\models\transformer.py
	forward,":param x: an [N x C1 x T] tensor.
	:param t: an [N] tensor.
	:param low_res: an [N x C2 x T'] tensor of conditioning points.
	:param images: a batch of images to condition on.
	:param embeddings: a batch of CLIP latent grids to condition on.
	:return: an [N x C3 x T] tensor.","def forward(self, x: torch.Tensor, t: torch.Tensor, *, low_res: torch.
	Tensor, images: Optional[Iterable[ImageType]]=None, embeddings:
	Optional[Iterable[torch.Tensor]]=None):
	""""""""""""
	assert x.shape[-1] == self.n_ctx
	t_embed = self.time_embed(timestep_embedding(t, self.backbone.width))
	low_res_embed = self._embed_low_res(low_res)
	if images is not None:
	clip_out = self.clip.embed_images_grid(images)
	elif embeddings is not None:
	clip_out = embeddings
	else:
	clip_out = torch.zeros([len(x), self.clip.grid_feature_dim, self.
	clip.grid_size ** 2], dtype=x.dtype, device=x.device)
	if self.training:
	mask = torch.rand(size=[len(x)]) >= self.cond_drop_prob
	clip_out = clip_out * mask[:, None, None].to(clip_out)
	clip_out = clip_out.permute(0, 2, 1)
	clip_embed = self.clip_embed(clip_out)
	cond = [(t_embed, self.time_token_cond), (clip_embed, True), (
	low_res_embed, True)]
	return self._forward_with_cond(x, cond)
	",point_e\models\transformer.py
	timestep_embedding,"Create sinusoidal timestep embeddings.
	:param timesteps: a 1-D Tensor of N indices, one per batch element.
	These may be fractional.
	:param dim: the dimension of the output.
	:param max_period: controls the minimum frequency of the embeddings.
	:return: an [N x dim] Tensor of positional embeddings.","def timestep_embedding(timesteps, dim, max_period=10000):
	""""""""""""
	half = dim // 2
	freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=
	half, dtype=torch.float32) / half).to(device=timesteps.device)
	args = timesteps[:, None].to(timesteps.dtype) * freqs[None]
	embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
	if dim % 2:
	embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1]
	)], dim=-1)
	return embedding
	",point_e\models\util.py
	load,Load the mesh from a .npz file.,"@classmethod
	def load(cls, f: Union[str, BinaryIO]) ->'TriMesh':
	""""""""""""
	if isinstance(f, str):
	with open(f, 'rb') as reader:
	return cls.load(reader)
	else:
	obj = np.load(f)
	keys = list(obj.keys())
	verts = obj['verts']
	faces = obj['faces']
	normals = obj['normals'] if 'normals' in keys else None
	vertex_channels = {}
	face_channels = {}
	for key in keys:
	if key.startswith('v_'):
	vertex_channels[key[2:]] = obj[key]
	elif key.startswith('f_'):
	face_channels[key[2:]] = obj[key]
	return cls(verts=verts, faces=faces, normals=normals,
	vertex_channels=vertex_channels, face_channels=face_channels)
	",point_e\util\mesh.py
	save,Save the mesh to a .npz file.,"def save(self, f: Union[str, BinaryIO]):
	""""""""""""
	if isinstance(f, str):
	with open(f, 'wb') as writer:
	self.save(writer)
	else:
	obj_dict = dict(verts=self.verts, faces=self.faces)
	if self.normals is not None:
	obj_dict['normals'] = self.normals
	for k, v in self.vertex_channels.items():
	obj_dict[f'v_{k}'] = v
	for k, v in self.face_channels.items():
	obj_dict[f'f_{k}'] = v
	np.savez(f, **obj_dict)
	",point_e\util\mesh.py
	has_vertex_colors,,"def has_vertex_colors(self) ->bool:
	return self.vertex_channels is not None and all(x in self.
	vertex_channels for x in 'RGB')
	",point_e\util\mesh.py
	write_ply,,"def write_ply(self, raw_f: BinaryIO):
	write_ply(raw_f, coords=self.verts, rgb=np.stack([self.vertex_channels[
	x] for x in 'RGB'], axis=1) if self.has_vertex_colors() else None,
	faces=self.faces)
	",point_e\util\mesh.py
	marching_cubes_mesh,"Run marching cubes on the SDF predicted from a point cloud to produce a
	mesh representing the 3D surface.

	:param pc: the point cloud to apply marching cubes to.
	:param model: the model to use to predict SDF values.
	:param grid_size: the number of samples along each axis. A total of
	grid_size**3 function evaluations are performed.
	:param side_length: the size of the cube containing the model, which is
	assumed to be centered at the origin.
	:param fill_vertex_channels: if True, use the nearest neighbor of each mesh
	vertex in the point cloud to compute vertex
	data (e.g. colors).","def marching_cubes_mesh(pc: PointCloud, model: PointCloudSDFModel,
	batch_size: int=4096, grid_size: int=128, side_length: float=1.02,
	fill_vertex_channels: bool=True, progress: bool=False) ->TriMesh:
	""""""""""""
	voxel_size = side_length / (grid_size - 1)
	min_coord = -side_length / 2

	def int_coord_to_float(int_coords: torch.Tensor) ->torch.Tensor:
	return int_coords.float() * voxel_size + min_coord
	with torch.no_grad():
	cond = model.encode_point_clouds(torch.from_numpy(pc.coords).
	permute(1, 0).to(model.device)[None])
	indices = range(0, grid_size ** 3, batch_size)
	if progress:
	indices = tqdm(indices)
	volume = []
	for i in indices:
	indices = torch.arange(i, min(i + batch_size, grid_size ** 3), step
	=1, dtype=torch.int64, device=model.device)
	zs = int_coord_to_float(indices % grid_size)
	ys = int_coord_to_float(torch.div(indices, grid_size, rounding_mode
	='trunc') % grid_size)
	xs = int_coord_to_float(torch.div(indices, grid_size ** 2,
	rounding_mode='trunc'))
	coords = torch.stack([xs, ys, zs], dim=0)
	with torch.no_grad():
	volume.append(model(coords[None], encoded=cond)[0])
	volume_np = torch.cat(volume).view(grid_size, grid_size, grid_size).cpu(
	).numpy()
	if np.all(volume_np < 0) or np.all(volume_np > 0):
	volume_np -= np.mean(volume_np)
	verts, faces, normals, _ = skimage.measure.marching_cubes(volume=
	volume_np, level=0, allow_degenerate=False, spacing=(voxel_size,) * 3)
	old_f1 = faces[:, 0].copy()
	faces[:, 0] = faces[:, 1]
	faces[:, 1] = old_f1
	verts += min_coord
	return TriMesh(verts=verts, faces=faces, normals=normals,
	vertex_channels=None if not fill_vertex_channels else
	_nearest_vertex_channels(pc, verts))
	",point_e\util\pc_to_mesh.py
	_nearest_vertex_channels,,"def _nearest_vertex_channels(pc: PointCloud, verts: np.ndarray) ->Dict[str,
	np.ndarray]:
	nearest = pc.nearest_points(verts)
	return {ch: arr[nearest] for ch, arr in pc.channels.items()}
	",point_e\util\pc_to_mesh.py
	plot_point_cloud,"Render a point cloud as a plot to the given image path.

	:param pc: the PointCloud to plot.
	:param image_path: the path to save the image, with a file extension.
	:param color: if True, show the RGB colors from the point cloud.
	:param grid_size: the number of random rotations to render.","def plot_point_cloud(pc: PointCloud, color: bool=True, grid_size: int=1,
	fixed_bounds: Optional[Tuple[Tuple[float, float, float], Tuple[float,
	float, float]]]=((-0.75, -0.75, -0.75), (0.75, 0.75, 0.75))):
	""""""""""""
	fig = plt.figure(figsize=(8, 8))
	for i in range(grid_size):
	for j in range(grid_size):
	ax = fig.add_subplot(grid_size, grid_size, 1 + j + i *
	grid_size, projection='3d')
	color_args = {}
	if color:
	color_args['c'] = np.stack([pc.channels['R'], pc.channels[
	'G'], pc.channels['B']], axis=-1)
	c = pc.coords
	if grid_size > 1:
	theta = np.pi * 2 * (i * grid_size + j) / grid_size ** 2
	rotation = np.array([[np.cos(theta), -np.sin(theta), 0.0],
	[np.sin(theta), np.cos(theta), 0.0], [0.0, 0.0, 1.0]])
	c = c @ rotation
	ax.scatter(c[:, 0], c[:, 1], c[:, 2], **color_args)
	if fixed_bounds is None:
	min_point = c.min(0)
	max_point = c.max(0)
	size = (max_point - min_point).max() / 2
	center = (min_point + max_point) / 2
	ax.set_xlim3d(center[0] - size, center[0] + size)
	ax.set_ylim3d(center[1] - size, center[1] + size)
	ax.set_zlim3d(center[2] - size, center[2] + size)
	else:
	ax.set_xlim3d(fixed_bounds[0][0], fixed_bounds[1][0])
	ax.set_ylim3d(fixed_bounds[0][1], fixed_bounds[1][1])
	ax.set_zlim3d(fixed_bounds[0][2], fixed_bounds[1][2])
	return fig
	",point_e\util\plotting.py
	write_ply,"Write a PLY file for a mesh or a point cloud.

	:param coords: an [N x 3] array of floating point coordinates.
	:param rgb: an [N x 3] array of vertex colors, in the range [0.0, 1.0].
	:param faces: an [N x 3] array of triangles encoded as integer indices.","def write_ply(raw_f: BinaryIO, coords: np.ndarray, rgb: Optional[np.ndarray
	]=None, faces: Optional[np.ndarray]=None):
	""""""""""""
	with buffered_writer(raw_f) as f:
	f.write(b'ply\n')
	f.write(b'format binary_little_endian 1.0\n')
	f.write(bytes(f'element vertex {len(coords)}\n', 'ascii'))
	f.write(b'property float x\n')
	f.write(b'property float y\n')
	f.write(b'property float z\n')
	if rgb is not None:
	f.write(b'property uchar red\n')
	f.write(b'property uchar green\n')
	f.write(b'property uchar blue\n')
	if faces is not None:
	f.write(bytes(f'element face {len(faces)}\n', 'ascii'))
	f.write(b'property list uchar int vertex_index\n')
	f.write(b'end_header\n')
	if rgb is not None:
	rgb = (rgb * 255.499).round().astype(int)
	vertices = [(coord, rgb) for coord, rgb in zip(coords.tolist(
	), rgb.tolist())]
	format = struct.Struct('<3f3B')
	for item in vertices:
	f.write(format.pack(*item))
	else:
	format = struct.Struct('<3f')
	for vertex in coords.tolist():
	f.write(format.pack(*vertex))
	if faces is not None:
	format = struct.Struct('<B3I')
	for tri in faces.tolist():
	f.write(format.pack(len(tri), *tri))
	",point_e\util\ply_util.py
	buffered_writer,,"@contextmanager
	def buffered_writer(raw_f: BinaryIO) ->Iterator[io.BufferedIOBase]:
	if isinstance(raw_f, io.BufferedIOBase):
	yield raw_f
	else:
	f = io.BufferedWriter(raw_f)
	yield f
	f.flush()
	",point_e\util\ply_util.py
	preprocess,,"def preprocess(data, channel):
	if channel in COLORS:
	return np.round(data * 255.0)
	return data
	",point_e\util\point_cloud.py
	load,Load the point cloud from a .npz file.,"@classmethod
	def load(cls, f: Union[str, BinaryIO]) ->'PointCloud':
	""""""""""""
	if isinstance(f, str):
	with open(f, 'rb') as reader:
	return cls.load(reader)
	else:
	obj = np.load(f)
	keys = list(obj.keys())
	return PointCloud(coords=obj['coords'], channels={k: obj[k] for k in
	keys if k != 'coords'})
	",point_e\util\point_cloud.py
	save,Save the point cloud to a .npz file.,"def save(self, f: Union[str, BinaryIO]):
	""""""""""""
	if isinstance(f, str):
	with open(f, 'wb') as writer:
	self.save(writer)
	else:
	np.savez(f, coords=self.coords, **self.channels)
	",point_e\util\point_cloud.py
	write_ply,,"def write_ply(self, raw_f: BinaryIO):
	write_ply(raw_f, coords=self.coords, rgb=np.stack([self.channels[x] for
	x in 'RGB'], axis=1) if all(x in self.channels for x in 'RGB') else
	None)
	",point_e\util\point_cloud.py
	random_sample,"Sample a random subset of this PointCloud.

	:param num_points: maximum number of points to sample.
	:param subsample_kwargs: arguments to self.subsample().
	:return: a reduced PointCloud, or self if num_points is not less than
	the current number of points.","def random_sample(self, num_points: int, **subsample_kwargs) ->'PointCloud':
	""""""""""""
	if len(self.coords) <= num_points:
	return self
	indices = np.random.choice(len(self.coords), size=(num_points,),
	replace=False)
	return self.subsample(indices, **subsample_kwargs)
	",point_e\util\point_cloud.py
	farthest_point_sample,"Sample a subset of the point cloud that is evenly distributed in space.

	First, a random point is selected. Then each successive point is chosen
	such that it is furthest from the currently selected points.

	The time complexity of this operation is O(NM), where N is the original
	number of points and M is the reduced number. Therefore, performance
	can be improved by randomly subsampling points with random_sample()
	before running farthest_point_sample().

	:param num_points: maximum number of points to sample.
	:param init_idx: if specified, the first point to sample.
	:param subsample_kwargs: arguments to self.subsample().
	:return: a reduced PointCloud, or self if num_points is not less than
	the current number of points.","def farthest_point_sample(self, num_points: int, init_idx: Optional[int]=
	None, **subsample_kwargs) ->'PointCloud':
	""""""""""""
	if len(self.coords) <= num_points:
	return self
	init_idx = random.randrange(len(self.coords)
	) if init_idx is None else init_idx
	indices = np.zeros([num_points], dtype=np.int64)
	indices[0] = init_idx
	sq_norms = np.sum(self.coords ** 2, axis=-1)

	def compute_dists(idx: int):
	return sq_norms + sq_norms[idx] - 2 * (self.coords @ self.coords[idx])
	cur_dists = compute_dists(init_idx)
	for i in range(1, num_points):
	idx = np.argmax(cur_dists)
	indices[i] = idx
	cur_dists = np.minimum(cur_dists, compute_dists(idx))
	return self.subsample(indices, **subsample_kwargs)
	",point_e\util\point_cloud.py
	subsample,,"def subsample(self, indices: np.ndarray, average_neighbors: bool=False
	) ->'PointCloud':
	if not average_neighbors:
	return PointCloud(coords=self.coords[indices], channels={k: v[
	indices] for k, v in self.channels.items()})
	new_coords = self.coords[indices]
	neighbor_indices = PointCloud(coords=new_coords, channels={}
	).nearest_points(self.coords)
	neighbor_indices[indices] = np.arange(len(indices))
	new_channels = {}
	for k, v in self.channels.items():
	v_sum = np.zeros_like(v[:len(indices)])
	v_count = np.zeros_like(v[:len(indices)])
	np.add.at(v_sum, neighbor_indices, v)
	np.add.at(v_count, neighbor_indices, 1)
	new_channels[k] = v_sum / v_count
	return PointCloud(coords=new_coords, channels=new_channels)
	",point_e\util\point_cloud.py
	select_channels,,"def select_channels(self, channel_names: List[str]) ->np.ndarray:
	data = np.stack([preprocess(self.channels[name], name) for name in
	channel_names], axis=-1)
	return data
	",point_e\util\point_cloud.py
	nearest_points,"For each point in another set of points, compute the point in this
	pointcloud which is closest.

	:param points: an [N x 3] array of points.
	:param batch_size: the number of neighbor distances to compute at once.
	Smaller values save memory, while larger values may
	make the computation faster.
	:return: an [N] array of indices into self.coords.","def nearest_points(self, points: np.ndarray, batch_size: int=16384
	) ->np.ndarray:
	""""""""""""
	norms = np.sum(self.coords ** 2, axis=-1)
	all_indices = []
	for i in range(0, len(points), batch_size):
	batch = points[i:i + batch_size]
	dists = norms + np.sum(batch ** 2, axis=-1)[:, None] - 2 * (batch @
	self.coords.T)
	all_indices.append(np.argmin(dists, axis=-1))
	return np.concatenate(all_indices, axis=0)
	",point_e\util\point_cloud.py
	combine,,"def combine(self, other: 'PointCloud') ->'PointCloud':
	assert self.channels.keys() == other.channels.keys()
	return PointCloud(coords=np.concatenate([self.coords, other.coords],
	axis=0), channels={k: np.concatenate([v, other.channels[k]], axis=0
	) for k, v in self.channels.items()})
	",point_e\util\point_cloud.py