turn off classifier free guidance if predicting x_start for diffusion prior

* add unrelated embedding metric

* change to torch.roll

Co-authored-by: nousr <z@localhost.com>
Co-authored-by: nousr <>

README.md

@@ -902,7 +902,7 @@ Please note that the script internally passes text_embed and image_embed to the
 ### Usage 
 
 ```bash
-$ pyhon train_diffusion_prior.py
+$ python train_diffusion_prior.py
 ```
 
 The most significant parameters for the script are as follows:
@@ -981,6 +981,8 @@ Once built, images will be saved to the same directory the command is invoked
 - [ ] make sure FILIP works with DALL-E2 from x-clip https://arxiv.org/abs/2111.07783
 - [ ] make sure resnet hyperparameters can be configurable across unet depth (groups and expansion factor)
 - [ ] offer save / load methods on the trainer classes to automatically take care of state dicts for scalers / optimizers / saving versions and checking for breaking changes
+- [ ] offer setting in diffusion prior to split time and image embeddings into multiple tokens, configurable, for more surface area during attention
+- [ ] bring in skip-layer excitatons (from lightweight gan paper) to see if it helps for either decoder of unet or vqgan-vae training
 
 ## Citations
 

dalle2_pytorch/dalle2_pytorch.py

@@ -706,7 +706,7 @@ class DiffusionPriorNetwork(nn.Module):
         **kwargs
     ):
         super().__init__()
-        self.time_embeddings = nn.Embedding(num_timesteps, dim) if exists(num_timesteps) else nn.Sequential(Rearrange('b -> b 1'), MLP(1, dim)) # also offer a continuous version of timestep embeddings, with a 2 layer MLP
+        self.time_embeddings = nn.Embedding(num_timesteps, dim) if exists(num_timesteps) else nn.Sequential(SinusoidalPosEmb(dim), MLP(dim, dim)) # also offer a continuous version of timestep embeddings, with a 2 layer MLP
         self.learned_query = nn.Parameter(torch.randn(dim))
         self.causal_transformer = CausalTransformer(dim = dim, **kwargs)
 
@@ -765,7 +765,7 @@ class DiffusionPriorNetwork(nn.Module):
         # but let's just do it right
 
         if exists(mask):
-            mask = F.pad(mask, (0, 2), value = True) # extend mask for text embedding, noised image embedding, time step embedding, and learned query
+            mask = F.pad(mask, (0, 3), value = True) # extend mask for text embedding, noised image embedding, time step embedding, and learned query
 
         time_embed = self.time_embeddings(diffusion_timesteps)
         time_embed = rearrange(time_embed, 'b d -> b 1 d')
@@ -776,6 +776,7 @@ class DiffusionPriorNetwork(nn.Module):
             text_encodings,
             text_embed,
             time_embed,
+            image_embed,
             learned_queries
         ), dim = -2)
 
@@ -799,13 +800,14 @@ class DiffusionPrior(BaseGaussianDiffusion):
         image_size = None,
         image_channels = 3,
         timesteps = 1000,
-        cond_drop_prob = 0.2,
+        cond_drop_prob = 0.,
         loss_type = "l1",
         predict_x_start = True,
         beta_schedule = "cosine",
         condition_on_text_encodings = True, # the paper suggests this is needed, but you can turn it off for your CLIP preprocessed text embed -> image embed training
         sampling_clamp_l2norm = False,
         training_clamp_l2norm = False,
+        init_image_embed_l2norm = False,
         image_embed_scale = None,           # this is for scaling the l2-normed image embedding, so it is more suitable for gaussian diffusion, as outlined by Katherine (@crowsonkb) https://github.com/lucidrains/DALLE2-pytorch/issues/60#issue-1226116132
         clip_adapter_overrides = dict()
     ):
@@ -832,7 +834,7 @@ class DiffusionPrior(BaseGaussianDiffusion):
         self.image_embed_dim = default(image_embed_dim, lambda: clip.dim_latent)
         self.channels = default(image_channels, lambda: clip.image_channels)
 
-        self.cond_drop_prob = cond_drop_prob
+        self.cond_drop_prob = cond_drop_prob if not predict_x_start else 0.
         self.condition_on_text_encodings = condition_on_text_encodings
 
         # in paper, they do not predict the noise, but predict x0 directly for image embedding, claiming empirically better results. I'll just offer both.
@@ -844,6 +846,7 @@ class DiffusionPrior(BaseGaussianDiffusion):
         # whether to force an l2norm, similar to clipping denoised, when sampling
         self.sampling_clamp_l2norm = sampling_clamp_l2norm
         self.training_clamp_l2norm = training_clamp_l2norm
+        self.init_image_embed_l2norm = init_image_embed_l2norm
 
     def p_mean_variance(self, x, t, text_cond, clip_denoised: bool):
         pred = self.net(x, t, **text_cond)
@@ -878,11 +881,16 @@ class DiffusionPrior(BaseGaussianDiffusion):
         device = self.betas.device
 
         b = shape[0]
-        img = torch.randn(shape, device=device)
+        image_embed = torch.randn(shape, device=device)
+
+        if self.init_image_embed_l2norm:
+            image_embed = l2norm(image_embed) * self.image_embed_scale
 
         for i in tqdm(reversed(range(0, self.num_timesteps)), desc='sampling loop time step', total=self.num_timesteps):
-            img = self.p_sample(img, torch.full((b,), i, device = device, dtype = torch.long), text_cond = text_cond)
-        return img
+            times = torch.full((b,), i, device = device, dtype = torch.long)
+            image_embed = self.p_sample(image_embed, times, text_cond = text_cond)
+
+        return image_embed
 
     def p_losses(self, image_embed, times, text_cond, noise = None):
         noise = default(noise, lambda: torch.randn_like(image_embed))

setup.py

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

train_diffusion_prior.py

@@ -46,28 +46,60 @@ def save_model(save_path, state_dict):
     print("====================================== Saving checkpoint ======================================")
     torch.save(state_dict, save_path+'/'+str(time.time())+'_saved_model.pth')
 
-def report_cosine_sims(diffusion_prior,image_reader,text_reader,train_set_size,val_set_size,NUM_TEST_EMBEDDINGS,device):
+
+def report_cosine_sims(diffusion_prior, image_reader, text_reader, train_set_size, val_set_size, NUM_TEST_EMBEDDINGS, device):
     cos = nn.CosineSimilarity(dim=1, eps=1e-6)
 
     tstart = train_set_size+val_set_size
     tend = train_set_size+val_set_size+NUM_TEST_EMBEDDINGS
 
-    for embt, embi in zip(text_reader(batch_size = NUM_TEST_EMBEDDINGS, start=tstart, end = tend),image_reader(batch_size = NUM_TEST_EMBEDDINGS, start=tstart, end = tend)):
+    for embt, embi in zip(text_reader(batch_size=NUM_TEST_EMBEDDINGS, start=tstart, end=tend), image_reader(batch_size=NUM_TEST_EMBEDDINGS, start=tstart, end=tend)):
+        # make a copy of the text embeddings for shuffling
         text_embed = torch.tensor(embt[0]).to(device)
-        text_embed = text_embed /  text_embed.norm(dim=1, keepdim=True)
-        test_text_cond = dict(text_embed = text_embed)
+        text_embed_shuffled = text_embed.clone()
 
+        # roll the text embeddings to simulate "unrelated" captions
+        rolled_idx = torch.roll(torch.arange(NUM_TEST_EMBEDDINGS), 1)
+        text_embed_shuffled = text_embed_shuffled[rolled_idx]
+        text_embed_shuffled = text_embed_shuffled / \
+            text_embed_shuffled.norm(dim=1, keepdim=True)
+        test_text_shuffled_cond = dict(text_embed=text_embed_shuffled)
+
+        # prepare the text embedding
+        text_embed = text_embed / text_embed.norm(dim=1, keepdim=True)
+        test_text_cond = dict(text_embed=text_embed)
+
+        # prepare image embeddings
         test_image_embeddings = torch.tensor(embi[0]).to(device)
-        test_image_embeddings = test_image_embeddings /  test_image_embeddings.norm(dim=1, keepdim=True)
+        test_image_embeddings = test_image_embeddings / \
+            test_image_embeddings.norm(dim=1, keepdim=True)
 
-        predicted_image_embeddings = diffusion_prior.p_sample_loop((NUM_TEST_EMBEDDINGS, 768), text_cond = test_text_cond)
-        predicted_image_embeddings = predicted_image_embeddings / predicted_image_embeddings.norm(dim=1, keepdim=True)
+        # predict on the unshuffled text embeddings
+        predicted_image_embeddings = diffusion_prior.p_sample_loop(
+            (NUM_TEST_EMBEDDINGS, 768), text_cond=test_text_cond)
+        predicted_image_embeddings = predicted_image_embeddings / \
+            predicted_image_embeddings.norm(dim=1, keepdim=True)
 
-        original_similarity = cos(text_embed,test_image_embeddings).cpu().numpy()
-        predicted_similarity = cos(text_embed,predicted_image_embeddings).cpu().numpy()
+        # predict on the shuffled embeddings
+        predicted_unrelated_embeddings = diffusion_prior.p_sample_loop(
+            (NUM_TEST_EMBEDDINGS, 768), text_cond=test_text_shuffled_cond)
+        predicted_unrelated_embeddings = predicted_unrelated_embeddings / \
+            predicted_unrelated_embeddings.norm(dim=1, keepdim=True)
 
-        wandb.log({"CosineSimilarity(text_embed,image_embed)": np.mean(original_similarity)})
-        wandb.log({"CosineSimilarity(text_embed,predicted_image_embed)":np.mean(predicted_similarity)})
+        # calculate similarities
+        original_similarity = cos(
+            text_embed, test_image_embeddings).cpu().numpy()
+        predicted_similarity = cos(
+            text_embed, predicted_image_embeddings).cpu().numpy()
+        unrelated_similarity = cos(
+            text_embed, predicted_unrelated_embeddings).cpu().numpy()
+
+        wandb.log(
+            {"CosineSimilarity(text_embed,image_embed)": np.mean(original_similarity)})
+        wandb.log({"CosineSimilarity(text_embed,predicted_image_embed)": np.mean(
+            predicted_similarity)})
+        wandb.log({"CosineSimilarity(text_embed,predicted_unrelated_embed)": np.mean(
+            unrelated_similarity)})
 
     return np.mean(predicted_similarity - original_similarity)