fix cosine bbeta schedule, thanks to @Zhengxinyang

README.md

@@ -14,7 +14,7 @@ It may also explore an extension of using <a href="https://huggingface.co/spaces
 
 Please join <a href="https://discord.gg/xBPBXfcFHd"><img alt="Join us on Discord" src="https://img.shields.io/discord/823813159592001537?color=5865F2&logo=discord&logoColor=white"></a> if you are interested in helping out with the replication
 
-There was enough interest for a Jax version. It will be completed after the Pytorch version shows signs of life on my toy tasks. <a href="https://github.com/lucidrains/dalle2-jax">Placeholder repository</a>
+There was enough interest for a Jax version. It will be completed after the Pytorch version shows signs of life on my toy tasks. <a href="https://github.com/lucidrains/dalle2-jax">Placeholder repository</a>. I will also eventually extend this to <a href="https://github.com/lucidrains/dalle2-video">text to video</a>, once the repository is in a good place.
 
 ## Install
 
@@ -197,10 +197,10 @@ clip = CLIP(
     dim_image = 512,
     dim_latent = 512,
     num_text_tokens = 49408,
-    text_enc_depth = 1,
+    text_enc_depth = 6,
     text_seq_len = 256,
     text_heads = 8,
-    visual_enc_depth = 1,
+    visual_enc_depth = 6,
     visual_image_size = 256,
     visual_patch_size = 32,
     visual_heads = 8
@@ -209,14 +209,15 @@ clip = CLIP(
 # 2 unets for the decoder (a la cascading DDPM)
 
 unet1 = Unet(
-    dim = 16,
+    dim = 32,
     image_embed_dim = 512,
+    cond_dim = 128,
     channels = 3,
     dim_mults = (1, 2, 4, 8)
 ).cuda()
 
 unet2 = Unet(
-    dim = 16,
+    dim = 32,
     image_embed_dim = 512,
     cond_dim = 128,
     channels = 3,
@@ -228,8 +229,8 @@ unet2 = Unet(
 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
-    timesteps = 100,
+    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,
     cond_drop_prob = 0.2
 ).cuda()
 
@@ -256,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 both contains `CLIP`, a unet, and a causal transformer)
+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))
 
 ```python
 from dalle2_pytorch import DALLE2
@@ -408,15 +409,12 @@ 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
+- [x] add efficient attention in unet
+- [ ] be able to finely customize what to condition on (text, image embed) for specific unet in the cascade (super resolution ddpms near the end may not need too much conditioning)
 - [ ] offload unets not being trained on to CPU for memory efficiency (for training each resolution unets separately)
-- [ ] become an expert with unets, port 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
+- [ ] become an expert with unets, cleanup unet code, make it fully configurable, port all learnings over to https://github.com/lucidrains/x-unet
+- [ ] train on a toy task, offer in colab
 
 ## Citations
 
@@ -465,4 +463,12 @@ Offer training wrappers
 }
 ```
 
+```bibtex
+@inproceedings{Tu2022MaxViTMV,
+    title   = {MaxViT: Multi-Axis Vision Transformer},
+    author  = {Zhe-Wei Tu and Hossein Talebi and Han Zhang and Feng Yang and Peyman Milanfar and Alan Conrad Bovik and Yinxiao Li},
+    year    = {2022}
+}
+```
+
 *Creating noise from data is easy; creating data from noise is generative modeling.* - Yang Song's <a href="https://arxiv.org/abs/2011.13456">paper</a>

dalle2_pytorch/dalle2_pytorch.py

@@ -105,8 +105,8 @@ def cosine_beta_schedule(timesteps, s = 0.008):
     as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
     """
     steps = timesteps + 1
-    x = torch.linspace(0, steps, steps)
-    alphas_cumprod = torch.cos(((x / steps) + s) / (1 + s) * torch.pi * 0.5) ** 2
+    x = torch.linspace(0, timesteps, steps)
+    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
     alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
     betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
     return torch.clip(betas, 0, 0.999)
@@ -798,6 +798,20 @@ class CrossAttention(nn.Module):
         out = rearrange(out, 'b h n d -> b n (h d)')
         return self.to_out(out)
 
+class GridAttention(nn.Module):
+    def __init__(self, *args, window_size = 8, **kwargs):
+        super().__init__()
+        self.window_size = window_size
+        self.attn = Attention(*args, **kwargs)
+
+    def forward(self, x):
+        h, w = x.shape[-2:]
+        wsz = self.window_size
+        x = rearrange(x, 'b c (w1 h) (w2 w) -> (b h w) (w1 w2) c', w1 = wsz, w2 = wsz)
+        out = self.attn(x)
+        out = rearrange(out, '(b h w) (w1 w2) c -> b c (w1 h) (w2 w)', w1 = wsz, w2 = wsz, h = h // wsz, w = w // wsz)
+        return out
+
 class Unet(nn.Module):
     def __init__(
         self,
@@ -813,6 +827,8 @@ class Unet(nn.Module):
         lowres_cond_upsample_mode = 'bilinear',
         blur_sigma = 0.1,
         blur_kernel_size = 3,
+        sparse_attn = False,
+        sparse_attn_window = 8,  # window size for sparse attention
         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__()
@@ -875,6 +891,7 @@ class Unet(nn.Module):
 
             self.downs.append(nn.ModuleList([
                 ConvNextBlock(dim_in, dim_out, norm = ind != 0),
+                Residual(GridAttention(dim_out, window_size = sparse_attn_window)) if sparse_attn else nn.Identity(),
                 ConvNextBlock(dim_out, dim_out, cond_dim = layer_cond_dim),
                 Downsample(dim_out) if not is_last else nn.Identity()
             ]))
@@ -891,6 +908,7 @@ class Unet(nn.Module):
 
             self.ups.append(nn.ModuleList([
                 ConvNextBlock(dim_out * 2, dim_in, cond_dim = layer_cond_dim),
+                Residual(GridAttention(dim_in, window_size = sparse_attn_window)) if sparse_attn else nn.Identity(),
                 ConvNextBlock(dim_in, dim_in, cond_dim = layer_cond_dim),
                 Upsample(dim_in)
             ]))
@@ -995,8 +1013,9 @@ class Unet(nn.Module):
 
         hiddens = []
 
-        for convnext, convnext2, downsample in self.downs:
+        for convnext, sparse_attn, convnext2, downsample in self.downs:
             x = convnext(x, c)
+            x = sparse_attn(x)
             x = convnext2(x, c)
             hiddens.append(x)
             x = downsample(x)
@@ -1008,9 +1027,10 @@ class Unet(nn.Module):
 
         x = self.mid_block2(x, mid_c)
 
-        for convnext, convnext2, upsample in self.ups:
+        for convnext, sparse_attn, convnext2, upsample in self.ups:
             x = torch.cat((x, hiddens.pop()), dim=1)
             x = convnext(x, c)
+            x = sparse_attn(x)
             x = convnext2(x, c)
             x = upsample(x)
 
@@ -1214,7 +1234,7 @@ class Decoder(nn.Module):
 
         return img
 
-    def forward(self, image, text = None, unet_number = None):
+    def forward(self, image, text = None, image_embed = None, text_encodings = 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,8 +1250,10 @@ class Decoder(nn.Module):
 
         times = torch.randint(0, self.num_timesteps, (b,), device = device, dtype = torch.long)
 
-        image_embed = self.get_image_embed(image)
-        text_encodings = self.get_text_encodings(text) if exists(text) else None
+        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
 
         lowres_cond_img = image if index > 0 else None
         ddpm_image = resize_image_to(image, target_image_size)

dalle2_pytorch/latent_diffusion.py

@@ -0,0 +1,12 @@
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+
+from einops import rearrange
+
+class LatentDiffusion(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x

setup.py

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