Create a wrapper class for the forward SDE (VP type). *** Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t. We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for
(
self,
schedule='discrete',
betas=None,
alphas_cumprod=None,
continuous_beta_0=0.1,
continuous_beta_1=20.,
dtype=torch.float32,
)
| 5 | |
| 6 | class NoiseScheduleVP: |
| 7 | def __init__( |
| 8 | self, |
| 9 | schedule='discrete', |
| 10 | betas=None, |
| 11 | alphas_cumprod=None, |
| 12 | continuous_beta_0=0.1, |
| 13 | continuous_beta_1=20., |
| 14 | dtype=torch.float32, |
| 15 | ): |
| 16 | """Create a wrapper class for the forward SDE (VP type). |
| 17 | |
| 18 | *** |
| 19 | Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t. |
| 20 | We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images. |
| 21 | *** |
| 22 | |
| 23 | The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ). |
| 24 | We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper). |
| 25 | Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have: |
| 26 | |
| 27 | log_alpha_t = self.marginal_log_mean_coeff(t) |
| 28 | sigma_t = self.marginal_std(t) |
| 29 | lambda_t = self.marginal_lambda(t) |
| 30 | |
| 31 | Moreover, as lambda(t) is an invertible function, we also support its inverse function: |
| 32 | |
| 33 | t = self.inverse_lambda(lambda_t) |
| 34 | |
| 35 | =============================================================== |
| 36 | |
| 37 | We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]). |
| 38 | |
| 39 | 1. For discrete-time DPMs: |
| 40 | |
| 41 | For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by: |
| 42 | t_i = (i + 1) / N |
| 43 | e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1. |
| 44 | We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3. |
| 45 | |
| 46 | Args: |
| 47 | betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details) |
| 48 | alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details) |
| 49 | |
| 50 | Note that we always have alphas_cumprod = cumprod(1 - betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`. |
| 51 | |
| 52 | **Important**: Please pay special attention for the args for `alphas_cumprod`: |
| 53 | The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that |
| 54 | q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ). |
| 55 | Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have |
| 56 | alpha_{t_n} = \sqrt{\hat{alpha_n}}, |
| 57 | and |
| 58 | log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}). |
| 59 | |
| 60 | |
| 61 | 2. For continuous-time DPMs: |
| 62 | |
| 63 | We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise |
| 64 | schedule are the default settings in DDPM and improved-DDPM: |