From 96bb2ff3105964afae337e7064ceb13ead7b904c Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Wed, 8 Jun 2022 23:07:01 -0700 Subject: [PATCH] alpha cosine noise schedule is now working for continuous time gaussian diffusion --- .../continuous_time_gaussian_diffusion.py | 13 +++++++++---- setup.py | 2 +- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py b/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py index 9e50abc..ad95045 100644 --- a/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py +++ b/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py @@ -59,11 +59,14 @@ class MonotonicLinear(nn.Module): # log(snr) that approximates the original linear schedule -def beta_linear_log_snr(t): - return -torch.log(expm1(1e-4 + 10 * (t ** 2))) +def log(t, eps = 1e-20): + return torch.log(t.clamp(min = eps)) -def alpha_cosine_log_snr(t): - raise NotImplementedError +def beta_linear_log_snr(t): + return -log(expm1(1e-4 + 10 * (t ** 2))) + +def alpha_cosine_log_snr(t, s = 0.008): + return -log((torch.cos((t + s) / (1 + s) * torch.pi * 0.5) ** -2) - 1) class learned_noise_schedule(nn.Module): """ described in section H and then I.2 of the supplementary material for variational ddpm paper """ @@ -135,6 +138,8 @@ class ContinuousTimeGaussianDiffusion(nn.Module): if noise_schedule == 'linear': self.log_snr = beta_linear_log_snr + elif noise_schedule == 'cosine': + self.log_snr = alpha_cosine_log_snr elif noise_schedule == 'learned': log_snr_max, log_snr_min = [beta_linear_log_snr(torch.tensor([time])).item() for time in (0., 1.)] diff --git a/setup.py b/setup.py index e301868..c81d446 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ from setuptools import setup, find_packages setup( name = 'denoising-diffusion-pytorch', packages = find_packages(), - version = '0.17.6', + version = '0.17.7', license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang',