From 0ffff59ca020bac3ebdff47c4f9bd95dec535d04 Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Sun, 30 Oct 2022 09:32:06 -0700 Subject: [PATCH] add option for random fourier features, given misinterpretation of Katherine's code, thanks to @tmabraham for addressing this in https://github.com/lucidrains/denoising-diffusion-pytorch/issues/112 --- .../continuous_time_gaussian_diffusion.py | 2 +- .../denoising_diffusion_pytorch.py | 17 +++++++++-------- .../elucidated_diffusion.py | 2 +- ..._param_continuous_time_gaussian_diffusion.py | 2 +- setup.py | 2 +- 5 files changed, 13 insertions(+), 12 deletions(-) diff --git a/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py b/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py index 7822d8b..3d455d3 100644 --- a/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py +++ b/denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py @@ -126,7 +126,7 @@ class ContinuousTimeGaussianDiffusion(nn.Module): p2_loss_weight_k = 1 ): super().__init__() - assert model.learned_sinusoidal_cond + assert model.random_or_learned_sinusoidal_cond assert not model.self_condition, 'not supported yet' self.model = model diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 8d4dc2e..50835b7 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -140,15 +140,15 @@ class SinusoidalPosEmb(nn.Module): emb = torch.cat((emb.sin(), emb.cos()), dim=-1) return emb -class LearnedSinusoidalPosEmb(nn.Module): - """ following @crowsonkb 's lead with learned sinusoidal pos emb """ +class RandomOrLearnedSinusoidalPosEmb(nn.Module): + """ following @crowsonkb 's lead with random (learned optional) sinusoidal pos emb """ """ https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8 """ - def __init__(self, dim): + def __init__(self, dim, is_random = False): super().__init__() assert (dim % 2) == 0 half_dim = dim // 2 - self.weights = nn.Parameter(torch.randn(half_dim)) + self.weights = nn.Parameter(torch.randn(half_dim), requires_grad = not is_random) def forward(self, x): x = rearrange(x, 'b -> b 1') @@ -271,6 +271,7 @@ class Unet(nn.Module): resnet_block_groups = 8, learned_variance = False, learned_sinusoidal_cond = False, + random_fourier_features = False, learned_sinusoidal_dim = 16 ): super().__init__() @@ -293,10 +294,10 @@ class Unet(nn.Module): time_dim = dim * 4 - self.learned_sinusoidal_cond = learned_sinusoidal_cond + self.random_or_learned_sinusoidal_cond = learned_sinusoidal_cond or random_fourier_features - if learned_sinusoidal_cond: - sinu_pos_emb = LearnedSinusoidalPosEmb(learned_sinusoidal_dim) + if self.random_or_learned_sinusoidal_cond: + sinu_pos_emb = RandomOrLearnedSinusoidalPosEmb(learned_sinusoidal_dim, random_fourier_features) fourier_dim = learned_sinusoidal_dim + 1 else: sinu_pos_emb = SinusoidalPosEmb(dim) @@ -429,7 +430,7 @@ class GaussianDiffusion(nn.Module): ): super().__init__() assert not (type(self) == GaussianDiffusion and model.channels != model.out_dim) - assert not model.learned_sinusoidal_cond + assert not model.random_or_learned_sinusoidal_cond self.model = model self.channels = self.model.channels diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py index 7ddaf19..2b62b4f 100644 --- a/denoising_diffusion_pytorch/elucidated_diffusion.py +++ b/denoising_diffusion_pytorch/elucidated_diffusion.py @@ -52,7 +52,7 @@ class ElucidatedDiffusion(nn.Module): S_noise = 1.003, ): super().__init__() - assert net.learned_sinusoidal_cond + assert net.random_or_learned_sinusoidal_cond self.self_condition = net.self_condition self.net = net diff --git a/denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py b/denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py index 6e01c73..c42fb84 100644 --- a/denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py +++ b/denoising_diffusion_pytorch/v_param_continuous_time_gaussian_diffusion.py @@ -61,7 +61,7 @@ class VParamContinuousTimeGaussianDiffusion(nn.Module): clip_sample_denoised = True, ): super().__init__() - assert model.learned_sinusoidal_cond + assert model.random_or_learned_sinusoidal_cond assert not model.self_condition, 'not supported yet' self.model = model diff --git a/setup.py b/setup.py index 50e65bd..41bf50d 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.29.0', + version = '0.29.1', license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang',