add UniAnimate (see notes)

@ d7814fa44a0a1154524b92fce0e3133a2604d333

https://github.com/ali-vilab/UniAnimate

animation/UniAnimate/tools/modules/diffusions/__init__.py

@@ -0,0 +1 @@
+from .diffusion_ddim import *

animation/UniAnimate/tools/modules/diffusions/diffusion_gauss.py

@@ -0,0 +1,498 @@
+"""
+GaussianDiffusion wraps operators for denoising diffusion models, including the
+diffusion and denoising processes, as well as the loss evaluation.
+"""
+import torch
+import torchsde
+import random
+from tqdm.auto import trange
+
+
+__all__ = ['GaussianDiffusion']
+
+
+def _i(tensor, t, x):
+    """
+    Index tensor using t and format the output according to x.
+    """
+    shape = (x.size(0), ) + (1, ) * (x.ndim - 1)
+    return tensor[t.to(tensor.device)].view(shape).to(x.device)
+
+
+class BatchedBrownianTree:
+    """
+    A wrapper around torchsde.BrownianTree that enables batches of entropy.
+    """
+    def __init__(self, x, t0, t1, seed=None, **kwargs):
+        t0, t1, self.sign = self.sort(t0, t1)
+        w0 = kwargs.get('w0', torch.zeros_like(x))
+        if seed is None:
+            seed = torch.randint(0, 2 ** 63 - 1, []).item()
+        self.batched = True
+        try:
+            assert len(seed) == x.shape[0]
+            w0 = w0[0]
+        except TypeError:
+            seed = [seed]
+            self.batched = False
+        self.trees = [torchsde.BrownianTree(
+            t0, w0, t1, entropy=s, **kwargs
+        ) for s in seed]
+    
+    @staticmethod
+    def sort(a, b):
+        return (a, b, 1) if a < b else (b, a, -1)
+
+    def __call__(self, t0, t1):
+        t0, t1, sign = self.sort(t0, t1)
+        w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
+        return w if self.batched else w[0]
+
+
+class BrownianTreeNoiseSampler:
+    """
+    A noise sampler backed by a torchsde.BrownianTree.
+
+    Args:
+        x (Tensor): The tensor whose shape, device and dtype to use to generate
+            random samples.
+        sigma_min (float): The low end of the valid interval.
+        sigma_max (float): The high end of the valid interval.
+        seed (int or List[int]): The random seed. If a list of seeds is
+            supplied instead of a single integer, then the noise sampler will
+            use one BrownianTree per batch item, each with its own seed.
+        transform (callable): A function that maps sigma to the sampler's
+            internal timestep.
+    """
+    def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x):
+        self.transform = transform
+        t0 = self.transform(torch.as_tensor(sigma_min))
+        t1 = self.transform(torch.as_tensor(sigma_max))
+        self.tree = BatchedBrownianTree(x, t0, t1, seed)
+    
+    def __call__(self, sigma, sigma_next):
+        t0 = self.transform(torch.as_tensor(sigma))
+        t1 = self.transform(torch.as_tensor(sigma_next))
+        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
+
+
+def get_scalings(sigma):
+    c_out = -sigma
+    c_in = 1 / (sigma ** 2 + 1. ** 2) ** 0.5
+    return c_out, c_in
+
+
+@torch.no_grad()
+def sample_dpmpp_2m_sde(
+    noise,
+    model,
+    sigmas,
+    eta=1.,
+    s_noise=1.,
+    solver_type='midpoint',
+    show_progress=True
+):
+    """
+    DPM-Solver++ (2M) SDE.
+    """
+    assert solver_type in {'heun', 'midpoint'}
+
+    x = noise * sigmas[0]
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas[sigmas < float('inf')].max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max)
+    old_denoised = None
+    h_last = None
+
+    for i in trange(len(sigmas) - 1, disable=not show_progress):
+        if sigmas[i] == float('inf'):
+            # Euler method
+            denoised = model(noise, sigmas[i])
+            x = denoised + sigmas[i + 1] * noise
+        else:
+            _, c_in = get_scalings(sigmas[i])
+            denoised = model(x * c_in, sigmas[i])
+            if sigmas[i + 1] == 0:
+                # Denoising step
+                x = denoised
+            else:
+                # DPM-Solver++(2M) SDE
+                t, s = -sigmas[i].log(), -sigmas[i + 1].log()
+                h = s - t
+                eta_h = eta * h
+
+                x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + \
+                    (-h - eta_h).expm1().neg() * denoised
+
+                if old_denoised is not None:
+                    r = h_last / h
+                    if solver_type == 'heun':
+                        x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * \
+                            (1 / r) * (denoised - old_denoised)
+                    elif solver_type == 'midpoint':
+                        x = x + 0.5 * (-h - eta_h).expm1().neg() * \
+                            (1 / r) * (denoised - old_denoised)
+
+                x = x + noise_sampler(
+                    sigmas[i],
+                    sigmas[i + 1]
+                ) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
+
+            old_denoised = denoised
+            h_last = h
+    return x
+
+
+class GaussianDiffusion(object):
+
+    def __init__(self, sigmas, prediction_type='eps'):
+        assert prediction_type in {'x0', 'eps', 'v'}
+        self.sigmas = sigmas.float()                        # noise coefficients
+        self.alphas = torch.sqrt(1 - sigmas ** 2).float()   # signal coefficients
+        self.num_timesteps = len(sigmas)
+        self.prediction_type = prediction_type
+
+    def diffuse(self, x0, t, noise=None):
+        """
+        Add Gaussian noise to signal x0 according to:
+        q(x_t | x_0) = N(x_t | alpha_t x_0, sigma_t^2 I).
+        """
+        noise = torch.randn_like(x0) if noise is None else noise
+        xt = _i(self.alphas, t, x0) * x0 + _i(self.sigmas, t, x0) * noise
+        return xt
+    
+    def denoise(
+        self,
+        xt,
+        t,
+        s,
+        model,
+        model_kwargs={},
+        guide_scale=None,
+        guide_rescale=None,
+        clamp=None,
+        percentile=None
+    ):
+        """
+        Apply one step of denoising from the posterior distribution q(x_s | x_t, x0).
+        Since x0 is not available, estimate the denoising results using the learned
+        distribution p(x_s | x_t, \hat{x}_0 == f(x_t)).
+        """
+        s = t - 1 if s is None else s
+
+        # hyperparams
+        sigmas = _i(self.sigmas, t, xt)
+        alphas = _i(self.alphas, t, xt)
+        alphas_s = _i(self.alphas, s.clamp(0), xt)
+        alphas_s[s < 0] = 1.
+        sigmas_s = torch.sqrt(1 - alphas_s ** 2)
+
+        # precompute variables
+        betas = 1 - (alphas / alphas_s) ** 2
+        coef1 = betas * alphas_s / sigmas ** 2
+        coef2 = (alphas * sigmas_s ** 2) / (alphas_s * sigmas ** 2)
+        var = betas * (sigmas_s / sigmas) ** 2
+        log_var = torch.log(var).clamp_(-20, 20)
+
+        # prediction
+        if guide_scale is None:
+            assert isinstance(model_kwargs, dict)
+            out = model(xt, t=t, **model_kwargs)
+        else:
+            # classifier-free guidance (arXiv:2207.12598)
+            # model_kwargs[0]: conditional kwargs
+            # model_kwargs[1]: non-conditional kwargs
+            assert isinstance(model_kwargs, list) and len(model_kwargs) == 2
+            y_out = model(xt, t=t, **model_kwargs[0])
+            if guide_scale == 1.:
+                out = y_out
+            else:
+                u_out = model(xt, t=t, **model_kwargs[1])
+                out = u_out + guide_scale * (y_out - u_out)
+
+                # rescale the output according to arXiv:2305.08891
+                if guide_rescale is not None:
+                    assert guide_rescale >= 0 and guide_rescale <= 1
+                    ratio = (y_out.flatten(1).std(dim=1) / (
+                        out.flatten(1).std(dim=1) + 1e-12
+                    )).view((-1, ) + (1, ) * (y_out.ndim - 1))
+                    out *= guide_rescale * ratio + (1 - guide_rescale) * 1.0
+        
+        # compute x0
+        if self.prediction_type == 'x0':
+            x0 = out
+        elif self.prediction_type == 'eps':
+            x0 = (xt - sigmas * out) / alphas
+        elif self.prediction_type == 'v':
+            x0 = alphas * xt - sigmas * out
+        else:
+            raise NotImplementedError(
+                f'prediction_type {self.prediction_type} not implemented'
+            )
+        
+        # restrict the range of x0
+        if percentile is not None:
+            # NOTE: percentile should only be used when data is within range [-1, 1]
+            assert percentile > 0 and percentile <= 1
+            s = torch.quantile(x0.flatten(1).abs(), percentile, dim=1)
+            s = s.clamp_(1.0).view((-1, ) + (1, ) * (xt.ndim - 1))
+            x0 = torch.min(s, torch.max(-s, x0)) / s
+        elif clamp is not None:
+            x0 = x0.clamp(-clamp, clamp)
+        
+        # recompute eps using the restricted x0
+        eps = (xt - alphas * x0) / sigmas
+
+        # compute mu (mean of posterior distribution) using the restricted x0
+        mu = coef1 * x0 + coef2 * xt
+        return mu, var, log_var, x0, eps
+
+    @torch.no_grad()
+    def sample(
+        self,
+        noise,
+        model,
+        model_kwargs={},
+        condition_fn=None,
+        guide_scale=None,
+        guide_rescale=None,
+        clamp=None,
+        percentile=None,
+        solver='euler_a',
+        steps=20,
+        t_max=None,
+        t_min=None,
+        discretization=None,
+        discard_penultimate_step=None,
+        return_intermediate=None,
+        show_progress=False,
+        seed=-1,
+        **kwargs
+    ):
+        # sanity check
+        assert isinstance(steps, (int, torch.LongTensor))
+        assert t_max is None or (t_max > 0 and t_max <= self.num_timesteps - 1)
+        assert t_min is None or (t_min >= 0 and t_min < self.num_timesteps - 1)
+        assert discretization in (None, 'leading', 'linspace', 'trailing')
+        assert discard_penultimate_step in (None, True, False)
+        assert return_intermediate in (None, 'x0', 'xt')
+
+        # function of diffusion solver
+        solver_fn = {
+            # 'heun': sample_heun,
+            'dpmpp_2m_sde': sample_dpmpp_2m_sde
+        }[solver]
+
+        # options
+        schedule = 'karras' if 'karras' in solver else None
+        discretization = discretization or 'linspace'
+        seed = seed if seed >= 0 else random.randint(0, 2 ** 31)
+        if isinstance(steps, torch.LongTensor):
+            discard_penultimate_step = False
+        if discard_penultimate_step is None:
+            discard_penultimate_step = True if solver in (
+                'dpm2',
+                'dpm2_ancestral',
+                'dpmpp_2m_sde',
+                'dpm2_karras',
+                'dpm2_ancestral_karras',
+                'dpmpp_2m_sde_karras'
+            ) else False
+        
+        # function for denoising xt to get x0
+        intermediates = []
+        def model_fn(xt, sigma):
+            # denoising
+            t = self._sigma_to_t(sigma).repeat(len(xt)).round().long()
+            x0 = self.denoise(
+                xt, t, None, model, model_kwargs, guide_scale, guide_rescale, clamp,
+                percentile
+            )[-2]
+
+            # collect intermediate outputs
+            if return_intermediate == 'xt':
+                intermediates.append(xt)
+            elif return_intermediate == 'x0':
+                intermediates.append(x0)
+            return x0
+        
+        # get timesteps
+        if isinstance(steps, int):
+            steps += 1 if discard_penultimate_step else 0
+            t_max = self.num_timesteps - 1 if t_max is None else t_max
+            t_min = 0 if t_min is None else t_min
+
+            # discretize timesteps
+            if discretization == 'leading':
+                steps = torch.arange(
+                    t_min, t_max + 1, (t_max - t_min + 1) / steps
+                ).flip(0)
+            elif discretization == 'linspace':
+                steps = torch.linspace(t_max, t_min, steps)
+            elif discretization == 'trailing':
+                steps = torch.arange(t_max, t_min - 1, -((t_max - t_min + 1) / steps))
+            else:
+                raise NotImplementedError(
+                    f'{discretization} discretization not implemented'
+                )
+            steps = steps.clamp_(t_min, t_max)
+        steps = torch.as_tensor(steps, dtype=torch.float32, device=noise.device)
+
+        # get sigmas
+        sigmas = self._t_to_sigma(steps)
+        sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
+        if schedule == 'karras':
+            if sigmas[0] == float('inf'):
+                sigmas = karras_schedule(
+                    n=len(steps) - 1,
+                    sigma_min=sigmas[sigmas > 0].min().item(),
+                    sigma_max=sigmas[sigmas < float('inf')].max().item(),
+                    rho=7.
+                ).to(sigmas)
+                sigmas = torch.cat([
+                    sigmas.new_tensor([float('inf')]), sigmas, sigmas.new_zeros([1])
+                ])
+            else:
+                sigmas = karras_schedule(
+                    n=len(steps),
+                    sigma_min=sigmas[sigmas > 0].min().item(),
+                    sigma_max=sigmas.max().item(),
+                    rho=7.
+                ).to(sigmas)
+                sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
+        if discard_penultimate_step:
+            sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
+        
+        # sampling
+        x0 = solver_fn(
+            noise,
+            model_fn,
+            sigmas,
+            show_progress=show_progress,
+            **kwargs
+        )
+        return (x0, intermediates) if return_intermediate is not None else x0
+    
+    @torch.no_grad()
+    def ddim_reverse_sample(
+        self,
+        xt,
+        t,
+        model,
+        model_kwargs={},
+        clamp=None,
+        percentile=None,
+        guide_scale=None,
+        guide_rescale=None,
+        ddim_timesteps=20,
+        reverse_steps=600
+        ):
+        r"""Sample from p(x_{t+1} | x_t) using DDIM reverse ODE (deterministic).
+        """
+        stride = reverse_steps // ddim_timesteps
+
+        # predict distribution of p(x_{t-1} | x_t)
+        _, _, _, x0, eps = self.denoise(
+                xt, t, None, model, model_kwargs, guide_scale, guide_rescale, clamp,
+                percentile
+            )
+        # derive variables
+        s = (t + stride).clamp(0, reverse_steps-1)
+        # hyperparams
+        sigmas = _i(self.sigmas, t, xt)
+        alphas = _i(self.alphas, t, xt)
+        alphas_s = _i(self.alphas, s.clamp(0), xt)
+        alphas_s[s < 0] = 1.
+        sigmas_s = torch.sqrt(1 - alphas_s ** 2)
+        
+        # reverse sample
+        mu = alphas_s * x0 + sigmas_s * eps
+        return mu, x0
+    
+    @torch.no_grad()
+    def ddim_reverse_sample_loop(
+        self,
+        x0,
+        model,
+        model_kwargs={},
+        clamp=None,
+        percentile=None,
+        guide_scale=None,
+        guide_rescale=None,
+        ddim_timesteps=20,
+        reverse_steps=600
+        ):
+        # prepare input
+        b = x0.size(0)
+        xt = x0
+
+        # reconstruction steps
+        steps = torch.arange(0, reverse_steps, reverse_steps // ddim_timesteps)
+        for step in steps:
+            t = torch.full((b, ), step, dtype=torch.long, device=xt.device)
+            xt, _ = self.ddim_reverse_sample(xt, t, model, model_kwargs, clamp, percentile, guide_scale, guide_rescale, ddim_timesteps, reverse_steps)
+        return xt
+    
+    def _sigma_to_t(self, sigma):
+        if sigma == float('inf'):
+            t = torch.full_like(sigma, len(self.sigmas) - 1)
+        else:
+            log_sigmas = torch.sqrt(
+                self.sigmas ** 2 / (1 - self.sigmas ** 2)
+            ).log().to(sigma)
+            log_sigma = sigma.log()
+            dists = log_sigma - log_sigmas[:, None]
+            low_idx = dists.ge(0).cumsum(dim=0).argmax(dim=0).clamp(
+                max=log_sigmas.shape[0] - 2
+            )
+            high_idx = low_idx + 1
+            low, high = log_sigmas[low_idx], log_sigmas[high_idx]
+            w = (low - log_sigma) / (low - high)
+            w = w.clamp(0, 1)
+            t = (1 - w) * low_idx + w * high_idx
+            t = t.view(sigma.shape)
+        if t.ndim == 0:
+            t = t.unsqueeze(0)
+        return t
+
+    def _t_to_sigma(self, t):
+        t = t.float()
+        low_idx, high_idx, w = t.floor().long(), t.ceil().long(), t.frac()
+        log_sigmas = torch.sqrt(self.sigmas ** 2 / (1 - self.sigmas ** 2)).log().to(t)
+        log_sigma = (1 - w) * log_sigmas[low_idx] + w * log_sigmas[high_idx]
+        log_sigma[torch.isnan(log_sigma) | torch.isinf(log_sigma)] = float('inf')
+        return log_sigma.exp()
+    
+    def prev_step(self, model_out, t, xt, inference_steps=50):
+        prev_t = t - self.num_timesteps // inference_steps
+
+        sigmas = _i(self.sigmas, t, xt)
+        alphas = _i(self.alphas, t, xt)
+        alphas_prev = _i(self.alphas, prev_t.clamp(0), xt)
+        alphas_prev[prev_t < 0] = 1.
+        sigmas_prev = torch.sqrt(1 - alphas_prev ** 2)
+        
+        x0 = alphas * xt - sigmas * model_out
+        eps = (xt - alphas * x0) / sigmas
+        prev_sample = alphas_prev * x0 + sigmas_prev * eps
+        return prev_sample
+    
+    def next_step(self, model_out, t, xt, inference_steps=50):
+        t, next_t = min(t - self.num_timesteps // inference_steps, 999), t
+
+        sigmas = _i(self.sigmas, t, xt)
+        alphas = _i(self.alphas, t, xt)
+        alphas_next = _i(self.alphas, next_t.clamp(0), xt)
+        alphas_next[next_t < 0] = 1.
+        sigmas_next = torch.sqrt(1 - alphas_next ** 2)
+        
+        x0 = alphas * xt - sigmas * model_out
+        eps = (xt - alphas * x0) / sigmas
+        next_sample = alphas_next * x0 + sigmas_next * eps
+        return next_sample
+    
+    def get_noise_pred_single(self, xt, t, model, model_kwargs):
+        assert isinstance(model_kwargs, dict)
+        out = model(xt, t=t, **model_kwargs)
+        return out
+    
+    

animation/UniAnimate/tools/modules/diffusions/losses.py

@@ -0,0 +1,28 @@
+import torch
+import math
+
+__all__ = ['kl_divergence', 'discretized_gaussian_log_likelihood']
+
+def kl_divergence(mu1, logvar1, mu2, logvar2):
+    return 0.5 * (-1.0 + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + ((mu1 - mu2) ** 2) * torch.exp(-logvar2))
+
+def standard_normal_cdf(x):
+    r"""A fast approximation of the cumulative distribution function of the standard normal.
+    """
+    return 0.5 * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+
+def discretized_gaussian_log_likelihood(x0, mean, log_scale):
+    assert x0.shape == mean.shape == log_scale.shape
+    cx = x0 - mean
+    inv_stdv = torch.exp(-log_scale)
+    cdf_plus = standard_normal_cdf(inv_stdv * (cx + 1.0 / 255.0))
+    cdf_min = standard_normal_cdf(inv_stdv * (cx - 1.0 / 255.0))
+    log_cdf_plus = torch.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = torch.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = torch.where(
+        x0 < -0.999,
+        log_cdf_plus,
+        torch.where(x0 > 0.999, log_one_minus_cdf_min, torch.log(cdf_delta.clamp(min=1e-12))))
+    assert log_probs.shape == x0.shape
+    return log_probs

animation/UniAnimate/tools/modules/diffusions/schedules.py

@@ -0,0 +1,166 @@
+import math
+import torch
+
+
+def beta_schedule(schedule='cosine',
+                   num_timesteps=1000,
+                   zero_terminal_snr=False,
+                   **kwargs):
+    # compute betas
+    betas = {
+        # 'logsnr_cosine_interp': logsnr_cosine_interp_schedule,
+        'linear': linear_schedule,
+        'linear_sd': linear_sd_schedule,
+        'quadratic': quadratic_schedule,
+        'cosine': cosine_schedule
+    }[schedule](num_timesteps, **kwargs)
+
+    if zero_terminal_snr and abs(betas.max() - 1.0) > 0.0001:
+        betas = rescale_zero_terminal_snr(betas)
+
+    return betas
+
+
+def sigma_schedule(schedule='cosine',
+                   num_timesteps=1000,
+                   zero_terminal_snr=False,
+                   **kwargs):
+    # compute betas
+    betas = {
+        'logsnr_cosine_interp': logsnr_cosine_interp_schedule,
+        'linear': linear_schedule,
+        'linear_sd': linear_sd_schedule,
+        'quadratic': quadratic_schedule,
+        'cosine': cosine_schedule
+    }[schedule](num_timesteps, **kwargs)
+    if schedule == 'logsnr_cosine_interp':
+        sigma = betas
+    else:
+        sigma = betas_to_sigmas(betas)
+    if zero_terminal_snr and abs(sigma.max() - 1.0) > 0.0001:
+        sigma = rescale_zero_terminal_snr(sigma)
+
+    return sigma
+
+
+def linear_schedule(num_timesteps, init_beta, last_beta,  **kwargs):
+    scale = 1000.0 / num_timesteps
+    init_beta = init_beta or scale * 0.0001
+    ast_beta = last_beta or scale * 0.02
+    return torch.linspace(init_beta, last_beta, num_timesteps, dtype=torch.float64)
+
+def logsnr_cosine_interp_schedule(
+        num_timesteps,
+        scale_min=2,
+        scale_max=4,
+        logsnr_min=-15,
+        logsnr_max=15,
+        **kwargs):
+    return logsnrs_to_sigmas(
+        _logsnr_cosine_interp(num_timesteps, logsnr_min, logsnr_max, scale_min, scale_max))
+
+def linear_sd_schedule(num_timesteps, init_beta, last_beta,  **kwargs):
+    return torch.linspace(init_beta ** 0.5, last_beta ** 0.5, num_timesteps, dtype=torch.float64) ** 2
+
+
+def quadratic_schedule(num_timesteps, init_beta, last_beta,  **kwargs):
+    init_beta = init_beta or 0.0015
+    last_beta = last_beta or 0.0195
+    return torch.linspace(init_beta ** 0.5, last_beta ** 0.5, num_timesteps, dtype=torch.float64) ** 2
+
+
+def cosine_schedule(num_timesteps, cosine_s=0.008, **kwargs):
+    betas = []
+    for step in range(num_timesteps):
+        t1 = step / num_timesteps
+        t2 = (step + 1) / num_timesteps
+        fn = lambda u: math.cos((u + cosine_s) / (1 + cosine_s) * math.pi / 2) ** 2
+        betas.append(min(1.0 - fn(t2) / fn(t1), 0.999))
+    return torch.tensor(betas, dtype=torch.float64)
+
+
+# def cosine_schedule(n, cosine_s=0.008, **kwargs):
+#     ramp = torch.linspace(0, 1, n + 1)
+#     square_alphas = torch.cos((ramp + cosine_s) / (1 + cosine_s) * torch.pi / 2) ** 2
+#     betas = (1 - square_alphas[1:] / square_alphas[:-1]).clamp(max=0.999)
+#     return betas_to_sigmas(betas)
+
+
+def betas_to_sigmas(betas):
+    return torch.sqrt(1 - torch.cumprod(1 - betas, dim=0))
+
+
+def sigmas_to_betas(sigmas):
+    square_alphas = 1 - sigmas**2
+    betas = 1 - torch.cat(
+        [square_alphas[:1], square_alphas[1:] / square_alphas[:-1]])
+    return betas
+
+
+
+def sigmas_to_logsnrs(sigmas):
+    square_sigmas = sigmas**2
+    return torch.log(square_sigmas / (1 - square_sigmas))
+
+
+def _logsnr_cosine(n, logsnr_min=-15, logsnr_max=15):
+    t_min = math.atan(math.exp(-0.5 * logsnr_min))
+    t_max = math.atan(math.exp(-0.5 * logsnr_max))
+    t = torch.linspace(1, 0, n)
+    logsnrs = -2 * torch.log(torch.tan(t_min + t * (t_max - t_min)))
+    return logsnrs
+
+
+def _logsnr_cosine_shifted(n, logsnr_min=-15, logsnr_max=15, scale=2):
+    logsnrs = _logsnr_cosine(n, logsnr_min, logsnr_max)
+    logsnrs += 2 * math.log(1 / scale)
+    return logsnrs
+
+def karras_schedule(n, sigma_min=0.002, sigma_max=80.0, rho=7.0):
+    ramp = torch.linspace(1, 0, 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
+    sigmas = torch.sqrt(sigmas**2 / (1 + sigmas**2))
+    return sigmas
+
+def _logsnr_cosine_interp(n,
+                          logsnr_min=-15,
+                          logsnr_max=15,
+                          scale_min=2,
+                          scale_max=4):
+    t = torch.linspace(1, 0, n)
+    logsnrs_min = _logsnr_cosine_shifted(n, logsnr_min, logsnr_max, scale_min)
+    logsnrs_max = _logsnr_cosine_shifted(n, logsnr_min, logsnr_max, scale_max)
+    logsnrs = t * logsnrs_min + (1 - t) * logsnrs_max
+    return logsnrs
+
+
+def logsnrs_to_sigmas(logsnrs):
+    return torch.sqrt(torch.sigmoid(-logsnrs))
+
+
+def rescale_zero_terminal_snr(betas):
+    """
+    Rescale Schedule to Zero Terminal SNR
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1 - betas
+    alphas_bar = alphas.cumprod(0)
+    alphas_bar_sqrt = alphas_bar.sqrt()
+
+    # Store old values. 8 alphas_bar_sqrt_0 = a
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+    # Shift so last timestep is zero.
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+    # Scale so first timestep is back to old value.
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+    # Convert alphas_bar_sqrt to betas
+    alphas_bar = alphas_bar_sqrt ** 2
+    alphas = alphas_bar[1:] / alphas_bar[:-1]
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+    return betas
+