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