Fix for linear beta schedules to ensure that the SNR at the end of the diffusion chain is 0 (zero terminal SNR), according to the paper Lin et al. (2023). Common Diffusion Noise Schedules and Sample Steps are Flawed. arXiv. http://arxiv.org/abs/2305.08891.
(betas)
| 59 | |
| 60 | |
| 61 | def enforce_zero_terminal_snr(betas): |
| 62 | """ |
| 63 | Fix for linear beta schedules to ensure that the SNR at the end of the |
| 64 | diffusion chain is 0 (zero terminal SNR), according to the paper |
| 65 | |
| 66 | Lin et al. (2023). Common Diffusion Noise Schedules and Sample Steps are |
| 67 | Flawed. arXiv. http://arxiv.org/abs/2305.08891. |
| 68 | """ |
| 69 | # Convert betas to alphas_bar_sqrt |
| 70 | alphas = 1 - betas |
| 71 | alphas_bar = alphas.cumprod(0) |
| 72 | alphas_bar_sqrt = alphas_bar.sqrt() |
| 73 | |
| 74 | # Store old values |
| 75 | alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() |
| 76 | alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() |
| 77 | |
| 78 | # Shift so last timestep is zero |
| 79 | alphas_bar_sqrt -= alphas_bar_sqrt_T |
| 80 | # Scale so first timestep is back to old value |
| 81 | alphas_bar_sqrt *= alphas_bar_sqrt_0 / ( alphas_bar_sqrt_0 - alphas_bar_sqrt_T) |
| 82 | |
| 83 | # Convert alphas_bar_sqrt to betas |
| 84 | alphas_bar = alphas_bar_sqrt ** 2 |
| 85 | alphas = alphas_bar[1:] / alphas_bar[:-1] |
| 86 | alphas = torch.cat([alphas_bar[0:1], alphas]) |
| 87 | betas = 1 - alphas |
| 88 | return betas |
| 89 | |
| 90 | |
| 91 | class GaussianDiffusion(nn.Module): |