Computes SNR as per https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
(noise_scheduler, timesteps)
| 51 | |
| 52 | |
| 53 | def compute_snr(noise_scheduler, timesteps): |
| 54 | """ |
| 55 | Computes SNR as per |
| 56 | https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849 |
| 57 | """ |
| 58 | alphas_cumprod = noise_scheduler.alphas_cumprod |
| 59 | sqrt_alphas_cumprod = alphas_cumprod**0.5 |
| 60 | sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5 |
| 61 | |
| 62 | # Expand the tensors. |
| 63 | # Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026 |
| 64 | sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[timesteps].float() |
| 65 | while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape): |
| 66 | sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None] |
| 67 | alpha = sqrt_alphas_cumprod.expand(timesteps.shape) |
| 68 | |
| 69 | sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(device=timesteps.device)[timesteps].float() |
| 70 | while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape): |
| 71 | sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None] |
| 72 | sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape) |
| 73 | |
| 74 | # Compute SNR. |
| 75 | snr = (alpha / sigma) ** 2 |
| 76 | return snr |
| 77 | |
| 78 | |
| 79 | def resolve_interpolation_mode(interpolation_type: str): |