remove convnext blocks, they do not work well, validated in video diffusion repository

README.md

@@ -99,14 +99,3 @@ Samples and model checkpoints will be logged to `./results` periodically
     note    = {under review}
 }
 ```
-
-```bibtex
-@misc{liu2022convnet,
-    title   = {A ConvNet for the 2020s},
-    author  = {Zhuang Liu and Hanzi Mao and Chao-Yuan Wu and Christoph Feichtenhofer and Trevor Darrell and Saining Xie},
-    year    = {2022},
-    eprint  = {2201.03545},
-    archivePrefix = {arXiv},
-    primaryClass = {cs.CV}
-}
-```

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -128,8 +128,8 @@ class ResnetBlock(nn.Module):
             nn.Linear(time_emb_dim, dim_out)
         ) if exists(time_emb_dim) else None
 
-        self.block1 = Block(dim, dim_out)
+        self.block1 = Block(dim, dim_out, groups = groups)
-        self.block2 = Block(dim_out, dim_out)
+        self.block2 = Block(dim_out, dim_out, groups = groups)
         self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
 
     def forward(self, x, time_emb = None):
@@ -142,39 +142,6 @@ class ResnetBlock(nn.Module):
         h = self.block2(h)
         return h + self.res_conv(x)
 
-class ConvNextBlock(nn.Module):
-    """ https://arxiv.org/abs/2201.03545 """
-
-    def __init__(self, dim, dim_out, *, time_emb_dim = None, mult = 2, norm = True):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.GELU(),
-            nn.Linear(time_emb_dim, dim)
-        ) if exists(time_emb_dim) else None
-
-        self.ds_conv = nn.Conv2d(dim, dim, 7, padding = 3, groups = dim)
-
-        self.net = nn.Sequential(
-            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)
-        )
-
-        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.ds_conv(x)
-
-        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')
-
-        h = self.net(h)
-        return h + self.res_conv(x)
-
 class LinearAttention(nn.Module):
     def __init__(self, dim, heads = 4, dim_head = 32):
         super().__init__()
@@ -237,9 +204,7 @@ class Unet(nn.Module):
         dim_mults=(1, 2, 4, 8),
         channels = 3,
         with_time_emb = True,
-        use_convnext = False,
+        resnet_block_groups = 8
-        resnet_block_groups = 8,
-        convnext_mult = 2
     ):
         super().__init__()
 
@@ -253,12 +218,7 @@ 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
+        block_klass = partial(ResnetBlock, groups = resnet_block_groups)
-
-        if use_convnext:
-            block_klass = partial(ConvNextBlock, mult = convnext_mult)
-        else:
-            block_klass = partial(ResnetBlock, groups = resnet_block_groups)
 
         # time embeddings
 

setup.py

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