Spaces:
Running
Running
| """ | |
| Simplified from https://github.com/openai/guided-diffusion/blob/main/guided_diffusion/gaussian_diffusion.py. | |
| """ | |
| import math | |
| import numpy as np | |
| import torch as th | |
| def _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, warmup_frac): | |
| betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) | |
| warmup_time = int(num_diffusion_timesteps * warmup_frac) | |
| betas[:warmup_time] = np.linspace(beta_start, beta_end, warmup_time, dtype=np.float64) | |
| return betas | |
| def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps): | |
| """ | |
| This is the deprecated API for creating beta schedules. | |
| See get_named_beta_schedule() for the new library of schedules. | |
| """ | |
| if beta_schedule == "quad": | |
| betas = ( | |
| np.linspace( | |
| beta_start ** 0.5, | |
| beta_end ** 0.5, | |
| num_diffusion_timesteps, | |
| dtype=np.float64, | |
| ) | |
| ** 2 | |
| ) | |
| elif beta_schedule == "linear": | |
| betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64) | |
| elif beta_schedule == "warmup10": | |
| betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.1) | |
| elif beta_schedule == "warmup50": | |
| betas = _warmup_beta(beta_start, beta_end, num_diffusion_timesteps, 0.5) | |
| elif beta_schedule == "const": | |
| betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64) | |
| elif beta_schedule == "jsd": # 1/T, 1/(T-1), 1/(T-2), ..., 1 | |
| betas = 1.0 / np.linspace( | |
| num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64 | |
| ) | |
| else: | |
| raise NotImplementedError(beta_schedule) | |
| assert betas.shape == (num_diffusion_timesteps,) | |
| return betas | |
| def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): | |
| """ | |
| 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. | |
| """ | |
| if schedule_name == "linear": | |
| # Linear schedule from Ho et al, extended to work for any number of | |
| # diffusion steps. | |
| 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 == "squaredcos_cap_v2": | |
| 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}") | |
| def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): | |
| """ | |
| 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. | |
| """ | |
| 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) | |
| class GaussianDiffusion: | |
| """ | |
| Utilities for training and sampling diffusion models. | |
| Original ported from this codebase: | |
| https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42 | |
| :param betas: a 1-D numpy array of betas for each diffusion timestep, | |
| starting at T and going to 1. | |
| """ | |
| def __init__( | |
| self, | |
| *, | |
| betas, | |
| ): | |
| # Use float64 for accuracy. | |
| 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,) | |
| # calculations for diffusion q(x_t | x_{t-1}) and others | |
| 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) | |
| # calculations for posterior q(x_{t-1} | x_t, x_0) | |
| self.posterior_variance = ( | |
| betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod) | |
| ) | |
| # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain | |
| 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) | |
| ) | |
| def q_mean_variance(self, x_start, t): | |
| """ | |
| 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. | |
| """ | |
| 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 | |
| def q_sample(self, x_start, t, noise=None): | |
| """ | |
| 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. | |
| """ | |
| 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 | |
| ) | |
| def q_posterior_mean_variance(self, x_start, x_t, t): | |
| """ | |
| Compute the mean and variance of the diffusion posterior: | |
| q(x_{t-1} | x_t, x_0) | |
| """ | |
| 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 | |
| def p_mean_variance(self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None): | |
| """ | |
| 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. | |
| """ | |
| 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 | |
| assert model_output.shape == (B, C * 2, *x.shape[2:]) | |
| model_output, model_var_values = th.split(model_output, C, dim=1) | |
| 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) | |
| # The model_var_values is [-1, 1] for [min_var, max_var]. | |
| frac = (model_var_values + 1) / 2 | |
| model_log_variance = frac * max_log + (1 - frac) * min_log | |
| model_variance = th.exp(model_log_variance) | |
| def process_xstart(x): | |
| if denoised_fn is not None: | |
| x = denoised_fn(x) | |
| if clip_denoised: | |
| return x.clamp(-1, 1) | |
| return x | |
| 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) | |
| 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, | |
| } | |
| 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 | |
| ) | |
| 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) | |
| def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): | |
| """ | |
| 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). | |
| """ | |
| gradient = cond_fn(x, t, **model_kwargs) | |
| new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() | |
| return new_mean | |
| def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): | |
| """ | |
| 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). | |
| """ | |
| 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 | |
| def p_sample( | |
| self, | |
| model, | |
| x, | |
| t, | |
| clip_denoised=True, | |
| denoised_fn=None, | |
| cond_fn=None, | |
| model_kwargs=None, | |
| ): | |
| """ | |
| 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. | |
| """ | |
| 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))) | |
| ) # no noise when t == 0 | |
| 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"]} | |
| def p_sample_loop( | |
| self, | |
| model, | |
| shape, | |
| noise=None, | |
| clip_denoised=True, | |
| denoised_fn=None, | |
| cond_fn=None, | |
| model_kwargs=None, | |
| device=None, | |
| progress=False, | |
| ): | |
| """ | |
| 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. | |
| """ | |
| 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, | |
| ): | |
| final = sample | |
| return final["sample"] | |
| def p_sample_loop_progressive( | |
| self, | |
| model, | |
| shape, | |
| noise=None, | |
| clip_denoised=True, | |
| denoised_fn=None, | |
| cond_fn=None, | |
| model_kwargs=None, | |
| device=None, | |
| progress=False, | |
| ): | |
| """ | |
| 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(). | |
| """ | |
| 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) | |
| indices = list(range(self.num_timesteps))[::-1] | |
| if progress: | |
| # Lazy import so that we don't depend on tqdm. | |
| 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 out | |
| img = out["sample"] | |
| def ddim_sample( | |
| self, | |
| model, | |
| x, | |
| t, | |
| clip_denoised=True, | |
| denoised_fn=None, | |
| cond_fn=None, | |
| model_kwargs=None, | |
| eta=0.0, | |
| ): | |
| """ | |
| Sample x_{t-1} from the model using DDIM. | |
| Same usage as p_sample(). | |
| """ | |
| 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) | |
| # Usually our model outputs epsilon, but we re-derive it | |
| # in case we used x_start or x_prev prediction. | |
| 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) | |
| ) | |
| # Equation 12. | |
| 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))) | |
| ) # no noise when t == 0 | |
| sample = mean_pred + nonzero_mask * sigma * noise | |
| return {"sample": sample, "pred_xstart": out["pred_xstart"]} | |
| def ddim_reverse_sample( | |
| self, | |
| model, | |
| x, | |
| t, | |
| clip_denoised=True, | |
| denoised_fn=None, | |
| cond_fn=None, | |
| model_kwargs=None, | |
| eta=0.0, | |
| ): | |
| """ | |
| Sample x_{t+1} from the model using DDIM reverse ODE. | |
| """ | |
| 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) | |
| # Usually our model outputs epsilon, but we re-derive it | |
| # in case we used x_start or x_prev prediction. | |
| 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) | |
| # Equation 12. reversed | |
| 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"]} | |
| def ddim_sample_loop( | |
| self, | |
| model, | |
| shape, | |
| noise=None, | |
| clip_denoised=True, | |
| denoised_fn=None, | |
| cond_fn=None, | |
| model_kwargs=None, | |
| device=None, | |
| progress=False, | |
| eta=0.0, | |
| ): | |
| """ | |
| Generate samples from the model using DDIM. | |
| Same usage as p_sample_loop(). | |
| """ | |
| 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, | |
| ): | |
| final = sample | |
| return final["sample"] | |
| def ddim_sample_loop_progressive( | |
| self, | |
| model, | |
| shape, | |
| noise=None, | |
| clip_denoised=True, | |
| denoised_fn=None, | |
| cond_fn=None, | |
| model_kwargs=None, | |
| device=None, | |
| progress=False, | |
| eta=0.0, | |
| ): | |
| """ | |
| Use DDIM to sample from the model and yield intermediate samples from | |
| each timestep of DDIM. | |
| Same usage as p_sample_loop_progressive(). | |
| """ | |
| 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) | |
| indices = list(range(self.num_timesteps))[::-1] | |
| if progress: | |
| # Lazy import so that we don't depend on tqdm. | |
| 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 out | |
| img = out["sample"] | |
| def _extract_into_tensor(arr, timesteps, broadcast_shape): | |
| """ | |
| 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. | |
| """ | |
| 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) | |