diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 5236788..766e01b 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -109,6 +109,39 @@ class PreNorm(nn.Module): # building block modules +class Block(nn.Module): + def __init__(self, dim, dim_out, groups = 8): + super().__init__() + self.block = nn.Sequential( + nn.Conv2d(dim, dim_out, 3, padding = 1), + nn.GroupNorm(groups, dim_out), + nn.SiLU() + ) + def forward(self, x): + return self.block(x) + +class ResnetBlock(nn.Module): + def __init__(self, dim, dim_out, *, time_emb_dim = None, groups = 8): + super().__init__() + self.mlp = nn.Sequential( + nn.SiLU(), + nn.Linear(time_emb_dim, dim_out) + ) if exists(time_emb_dim) else None + + self.block1 = Block(dim, dim_out) + self.block2 = Block(dim_out, dim_out) + self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity() + + def forward(self, x, time_emb = None): + h = self.block1(x) + + if exists(self.mlp) and exists(time_emb): + time_emb = self.mlp(time_emb) + h = rearrange(time_emb, 'b c -> b c 1 1') + h + + h = self.block2(h) + return h + self.res_conv(x) + class ConvNextBlock(nn.Module): """ https://arxiv.org/abs/2201.03545 """ @@ -125,6 +158,7 @@ class ConvNextBlock(nn.Module): LayerNorm(dim) if norm else nn.Identity(), nn.Conv2d(dim, dim_out * mult, 3, padding = 1), nn.GELU(), + LayerNorm(dim_out * mult), nn.Conv2d(dim_out * mult, dim_out, 3, padding = 1) ) @@ -133,7 +167,7 @@ class ConvNextBlock(nn.Module): def forward(self, x, time_emb = None): h = self.ds_conv(x) - if exists(self.mlp): + if exists(self.mlp) and exists(time_emb): assert exists(time_emb), 'time emb must be passed in' condition = self.mlp(time_emb) h = h + rearrange(condition, 'b c -> b c 1 1') @@ -148,15 +182,21 @@ class LinearAttention(nn.Module): self.heads = heads hidden_dim = dim_head * heads self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False) - self.to_out = nn.Conv2d(hidden_dim, dim, 1) + + self.to_out = nn.Sequential( + nn.Conv2d(hidden_dim, dim, 1), + LayerNorm(dim) + ) def forward(self, x): b, c, h, w = x.shape qkv = self.to_qkv(x).chunk(3, dim = 1) q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> b h c (x y)', h = self.heads), qkv) - q = q * self.scale + q = q.softmax(dim = -2) k = k.softmax(dim = -1) + + q = q * self.scale context = torch.einsum('b h d n, b h e n -> b h d e', k, v) out = torch.einsum('b h d e, b h d n -> b h e n', context, q) @@ -196,9 +236,14 @@ class Unet(nn.Module): out_dim = None, dim_mults=(1, 2, 4, 8), channels = 3, - with_time_emb = True + with_time_emb = True, + use_convnext = False, + resnet_block_groups = 8 ): super().__init__() + + # determine dimensions + self.channels = channels init_dim = default(init_dim, dim // 3 * 2) @@ -207,6 +252,15 @@ class Unet(nn.Module): dims = [init_dim, *map(lambda m: dim * m, dim_mults)] in_out = list(zip(dims[:-1], dims[1:])) + # resnet or convnext + + if use_convnext: + block_klass = ConvNextBlock + else: + block_klass = partial(ResnetBlock, groups = resnet_block_groups) + + # time embeddings + if with_time_emb: time_dim = dim * 4 self.time_mlp = nn.Sequential( @@ -219,6 +273,8 @@ class Unet(nn.Module): time_dim = None self.time_mlp = None + # layers + self.downs = nn.ModuleList([]) self.ups = nn.ModuleList([]) num_resolutions = len(in_out) @@ -227,30 +283,30 @@ class Unet(nn.Module): is_last = ind >= (num_resolutions - 1) self.downs.append(nn.ModuleList([ - ConvNextBlock(dim_in, dim_out, time_emb_dim = time_dim, norm = ind != 0), - ConvNextBlock(dim_out, dim_out, time_emb_dim = time_dim), + block_klass(dim_in, dim_out, time_emb_dim = time_dim), + block_klass(dim_out, dim_out, time_emb_dim = time_dim), Residual(PreNorm(dim_out, LinearAttention(dim_out))), Downsample(dim_out) if not is_last else nn.Identity() ])) mid_dim = dims[-1] - self.mid_block1 = ConvNextBlock(mid_dim, mid_dim, time_emb_dim = time_dim) + self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim) self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim))) - self.mid_block2 = ConvNextBlock(mid_dim, mid_dim, time_emb_dim = time_dim) + self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim = time_dim) for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): is_last = ind >= (num_resolutions - 1) self.ups.append(nn.ModuleList([ - ConvNextBlock(dim_out * 2, dim_in, time_emb_dim = time_dim), - ConvNextBlock(dim_in, dim_in, time_emb_dim = time_dim), + block_klass(dim_out * 2, dim_in, time_emb_dim = time_dim), + block_klass(dim_in, dim_in, time_emb_dim = time_dim), Residual(PreNorm(dim_in, LinearAttention(dim_in))), Upsample(dim_in) if not is_last else nn.Identity() ])) out_dim = default(out_dim, channels) self.final_conv = nn.Sequential( - ConvNextBlock(dim, dim), + block_klass(dim, dim), nn.Conv2d(dim, out_dim, 1) ) @@ -261,9 +317,9 @@ class Unet(nn.Module): h = [] - for convnext, convnext2, attn, downsample in self.downs: - x = convnext(x, t) - x = convnext2(x, t) + for block1, block2, attn, downsample in self.downs: + x = block1(x, t) + x = block2(x, t) x = attn(x) h.append(x) x = downsample(x) @@ -272,10 +328,10 @@ class Unet(nn.Module): x = self.mid_attn(x) x = self.mid_block2(x, t) - for convnext, convnext2, attn, upsample in self.ups: + for block1, block2, attn, upsample in self.ups: x = torch.cat((x, h.pop()), dim=1) - x = convnext(x, t) - x = convnext2(x, t) + x = block1(x, t) + x = block2(x, t) x = attn(x) x = upsample(x) diff --git a/setup.py b/setup.py index ea1fb61..3002c57 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.10.1', + version = '0.11.0', license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang',