allow for configuring expansion factor in convnext

README.md

@@ -34,7 +34,7 @@ diffusion = GaussianDiffusion(
     loss_type = 'l1'    # L1 or L2
 )
 
-training_images = torch.randn(8, 3, 128, 128)
+training_images = torch.randn(8, 3, 128, 128) # your images need to be normalized from a range of -1 to +1
 loss = diffusion(training_images)
 loss.backward()
 # after a lot of training

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)
@@ -192,17 +232,36 @@ class Unet(nn.Module):
     def __init__(
         self,
         dim,
+        init_dim = None,
         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,
+        convnext_mult = 2
     ):
         super().__init__()
+
+        # determine dimensions
+
         self.channels = channels
 
-        dims = [channels, *map(lambda m: dim * m, dim_mults)]
+        init_dim = default(init_dim, dim // 3 * 2)
+        self.init_conv = nn.Conv2d(channels, init_dim, 7, padding = 3)
+
+        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 = partial(ConvNextBlock, mult = convnext_mult)
+        else:
+            block_klass = partial(ResnetBlock, groups = resnet_block_groups)
+
+        # time embeddings
+
         if with_time_emb:
             time_dim = dim * 4
             self.time_mlp = nn.Sequential(
@@ -215,6 +274,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)
@@ -223,41 +284,43 @@ 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),
+                block_klass(dim_in, dim_out, time_emb_dim = time_dim),
-                ConvNextBlock(dim_out, 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),
+                block_klass(dim_out * 2, dim_in, time_emb_dim = time_dim),
-                ConvNextBlock(dim_in, 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)
         )
 
     def forward(self, x, time):
+        x = self.init_conv(x)
+
         t = self.time_mlp(time) if exists(self.time_mlp) else None
 
         h = []
 
-        for convnext, convnext2, attn, downsample in self.downs:
+        for block1, block2, attn, downsample in self.downs:
-            x = convnext(x, t)
+            x = block1(x, t)
-            x = convnext2(x, t)
+            x = block2(x, t)
             x = attn(x)
             h.append(x)
             x = downsample(x)
@@ -266,10 +329,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 = block1(x, t)
-            x = convnext2(x, t)
+            x = block2(x, t)
             x = attn(x)
             x = upsample(x)
 

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.10.0',
+  version = '0.11.1',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',