make kernel size and sigma for gaussian blur for cascading DDPM overridable at forward. also make sure unets are wrapped in a modulelist so that at sample time, blurring does not happen

README.md

@@ -197,10 +197,10 @@ clip = CLIP(
     dim_image = 512,
     dim_latent = 512,
     num_text_tokens = 49408,
-    text_enc_depth = 6,
+    text_enc_depth = 1,
     text_seq_len = 256,
     text_heads = 8,
-    visual_enc_depth = 6,
+    visual_enc_depth = 1,
     visual_image_size = 256,
     visual_patch_size = 32,
     visual_heads = 8
@@ -209,28 +209,28 @@ clip = CLIP(
 # 2 unets for the decoder (a la cascading DDPM)
 
 unet1 = Unet(
-    dim = 32,
+    dim = 16,
     image_embed_dim = 512,
-    cond_dim = 128,
     channels = 3,
     dim_mults = (1, 2, 4, 8)
 ).cuda()
 
 unet2 = Unet(
-    dim = 32,
+    dim = 16,
     image_embed_dim = 512,
+    lowres_cond = True,         # subsequent unets must have this turned on (and first unet must have this turned off)
     cond_dim = 128,
     channels = 3,
     dim_mults = (1, 2, 4, 8, 16)
 ).cuda()
 
-# decoder, which contains the unet(s) and clip
+# decoder, which contains the unet and clip
 
 decoder = Decoder(
     clip = clip,
     unet = (unet1, unet2),            # insert both unets in order of low resolution to highest resolution (you can have as many stages as you want here)
-    image_sizes = (256, 512),         # resolutions, 256 for first unet, 512 for second. these must be unique and in ascending order (matches with the unets passed in)
-    timesteps = 1000,
+    image_sizes = (256, 512),         # resolutions, 256 for first unet, 512 for second
+    timesteps = 100,
     cond_drop_prob = 0.2
 ).cuda()
 
@@ -257,7 +257,7 @@ mock_image_embed = torch.randn(1, 512).cuda()
 images = decoder.sample(mock_image_embed) # (1, 3, 512, 512)
 ```
 
-Finally, to generate the DALL-E2 images from text. Insert the trained `DiffusionPrior` as well as the `Decoder` (which wraps `CLIP`, the causal transformer, and unet(s))
+Finally, to generate the DALL-E2 images from text. Insert the trained `DiffusionPrior` as well as the `Decoder` (which both contains `CLIP`, a unet, and a causal transformer)
 
 ```python
 from dalle2_pytorch import DALLE2
@@ -349,7 +349,8 @@ unet2 = Unet(
     image_embed_dim = 512,
     cond_dim = 128,
     channels = 3,
-    dim_mults = (1, 2, 4, 8, 16)
+    dim_mults = (1, 2, 4, 8, 16),
+    lowres_cond = True
 ).cuda()
 
 decoder = Decoder(
@@ -409,10 +410,14 @@ Offer training wrappers
 - [x] augment unet so that it can also be conditioned on text encodings (although in paper they hinted this didn't make much a difference)
 - [x] figure out all the current bag of tricks needed to make DDPMs great (starting with the blur trick mentioned in paper)
 - [x] build the cascading ddpm by having Decoder class manage multiple unets at different resolutions
+- [ ] use an image resolution cutoff and do cross attention conditioning only if resources allow, and MLP + sum conditioning on rest
+- [ ] make unet more configurable
+- [ ] figure out some factory methods to make cascading unet instantiations less error-prone
 - [ ] offload unets not being trained on to CPU for memory efficiency (for training each resolution unets separately)
-- [ ] build out latent diffusion architecture in separate file, as it is not faithful to dalle-2 (but offer it as as setting)
-- [ ] become an expert with unets, cleanup unet code, make it fully configurable, add efficient attention (conditional on resolution), port all learnings over to https://github.com/lucidrains/x-unet
 - [ ] train on a toy task, offer in colab
+- [ ] add attention to unet - apply some personal tricks with efficient attention - use the sparse attention mechanism from https://github.com/lucidrains/vit-pytorch#maxvit
+- [ ] build out latent diffusion architecture in separate file, as it is not faithful to dalle-2 (but offer it as as setting)
+- [ ] consider U2-net for decoder https://arxiv.org/abs/2005.09007 (also in separate file as experimental) build out https://github.com/lucidrains/x-unet
 
 ## Citations
 

dalle2_pytorch/dalle2_pytorch.py

@@ -12,6 +12,7 @@ from einops.layers.torch import Rearrange
 from einops_exts import rearrange_many, repeat_many, check_shape
 from einops_exts.torch import EinopsToAndFrom
 
+from kornia.filters.gaussian import GaussianBlur2d
 from kornia.filters import gaussian_blur2d
 
 from dalle2_pytorch.tokenizer import tokenizer
@@ -816,11 +817,6 @@ class Unet(nn.Module):
         attend_at_middle = True, # whether to have a layer of attention at the bottleneck (can turn off for higher resolution in cascading DDPM, before bringing in efficient attention)
     ):
         super().__init__()
-        # save locals to take care of some hyperparameters for cascading DDPM
-
-        self._locals = locals()
-        del self._locals['self']
-        del self._locals['__class__']
 
         # for eventual cascading diffusion
 
@@ -901,15 +897,6 @@ class Unet(nn.Module):
             nn.Conv2d(dim, out_dim, 1)
         )
 
-    # if the current settings for the unet are not correct
-    # for cascading DDPM, then reinit the unet with the right settings
-    def force_lowres_cond(self, lowres_cond):
-        if lowres_cond == self.lowres_cond:
-            return self
-
-        updated_kwargs = {**self._locals, 'lowres_cond': lowres_cond}
-        return self.__class__(**updated_kwargs)
-
     def forward_with_cond_scale(
         self,
         *args,
@@ -1035,17 +1022,7 @@ class Decoder(nn.Module):
         self.clip_image_size = clip.image_size
         self.channels = clip.image_channels
 
-        # automatically take care of ensuring that first unet is unconditional
-        # while the rest of the unets are conditioned on the low resolution image produced by previous unet
-
-        self.unets = nn.ModuleList([])
-        for ind, one_unet in enumerate(cast_tuple(unet)):
-            is_first = ind == 0
-            one_unet = one_unet.force_lowres_cond(not is_first)
-            self.unets.append(one_unet)
-
-        # unet image sizes
-
+        self.unets = nn.ModuleList(unet)
         image_sizes = default(image_sizes, (clip.image_size,))
         image_sizes = tuple(sorted(set(image_sizes)))
 
@@ -1214,7 +1191,7 @@ class Decoder(nn.Module):
 
         return img
 
-    def forward(self, image, text = None, image_embed = None, text_encodings = None, unet_number = None):
+    def forward(self, image, text = None, unet_number = None):
         assert not (len(self.unets) > 1 and not exists(unet_number)), f'you must specify which unet you want trained, from a range of 1 to {len(self.unets)}, if you are training cascading DDPM (multiple unets)'
         unet_number = default(unet_number, 1)
         assert 1 <= unet_number <= len(self.unets)
@@ -1230,10 +1207,8 @@ class Decoder(nn.Module):
 
         times = torch.randint(0, self.num_timesteps, (b,), device = device, dtype = torch.long)
 
-        if not exists(image_embed):
-            image_embed = self.get_image_embed(image)
-
-        text_encodings = self.get_text_encodings(text) if exists(text) and not exists(text_encodings) else None
+        image_embed = self.get_image_embed(image)
+        text_encodings = self.get_text_encodings(text) if exists(text) else None
 
         lowres_cond_img = image if index > 0 else None
         ddpm_image = resize_image_to(image, target_image_size)

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.0.24',
+  version = '0.0.21',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',