Upload diffusion.py

scripts/diffusion.py

@@ -110,6 +110,7 @@ class Diffusion(L.LightningModule):
 
     ############ FORWARD DIFFUSION #########
     def subs_parameterization(self, logits, noised_latents):
+        print(logits.size()) # [bsz x bsz x seq_len]
         logits = logits.float()
         logits[:, :, self.mask_index] += self.neg_infinity
         
@@ -147,7 +148,7 @@ class Diffusion(L.LightningModule):
             x: int torch.Tensor with shape (batch_size, diffusion_model_input_length), input. 
             move_chance: float torch.Tensor with shape (batch_size, 1).
         """
-        #latents = latents.mean(dim=1) # [bsz x seq_len x 1280] --> [bsz x 1280] as per markdown
+        latents = torch.mean(latents, dim=2) # [bsz x seq_len x 1280] --> [bsz x seq_len] as per markdown
         move_indices = torch.rand(* latents.shape, device=latents.device) < move_chance
         noised_latents = torch.where(move_indices, self.mask_index, latents)
         return noised_latents
@@ -172,13 +173,14 @@ class Diffusion(L.LightningModule):
 
         xt = self.q_xt(x0, move_chance)
         model_output = self.forward(xt, unet_conditioning)
+        print(f'model out: {model_output}')
+        print(f'model out dim: {model_output.size()}') # [bsz x bsz x seq_len]
     
         # SUBS parameterization, continuous time.
-        idx = x0.long()
-        print(f'idx: {idx.size()}')
-        print(f'idx min: {idx.min()}')
-        print(f'idx max: {idx.max()}')
-        print(f'model out: {model_output.size()}')
+        idx = torch.mean(x0, dim=2).long()[:, :, None]
+        print(f'idx: {idx}')
+        print(f'idx dim: {idx.size()}') # [bsz x seq_len x 1]
+
         log_p_theta = torch.gather(input=model_output, dim=-1, index=idx).squeeze(-1)
         scale = (dsigma / torch.expm1(sigma))[:, None]
         return - log_p_theta * scale