0.22.3

dalle2_pytorch/dalle2_pytorch.py

@@ -806,6 +806,7 @@ class DiffusionPriorNetwork(nn.Module):
         num_time_embeds = 1,
         num_image_embeds = 1,
         num_text_embeds = 1,
+        attend_all_text_encodings = True,
         **kwargs
     ):
         super().__init__()
@@ -831,6 +832,8 @@ class DiffusionPriorNetwork(nn.Module):
         self.learned_query = nn.Parameter(torch.randn(dim))
         self.causal_transformer = CausalTransformer(dim = dim, **kwargs)
 
+        self.attend_all_text_encodings = attend_all_text_encodings
+
     def forward_with_cond_scale(
         self,
         *args,
@@ -852,7 +855,6 @@ class DiffusionPriorNetwork(nn.Module):
         *,
         text_embed,
         text_encodings = None,
-        mask = None,
         cond_drop_prob = 0.
     ):
         batch, dim, device, dtype = *image_embed.shape, image_embed.device, image_embed.dtype
@@ -871,7 +873,10 @@ class DiffusionPriorNetwork(nn.Module):
         if not exists(text_encodings):
             text_encodings = torch.empty((batch, 0, dim), device = device, dtype = dtype)
 
-        mask = torch.any(text_encodings != 0., dim = -1)
+        if self.attend_all_text_encodings:
+            mask = torch.ones((batch, text_encodings.shape[-2]), device = device, dtype = torch.bool)
+        else:
+            mask = torch.any(text_encodings != 0., dim = -1)
 
         # classifier free guidance
 
@@ -1812,6 +1817,7 @@ class Unet(nn.Module):
         text_tokens = None
 
         if exists(text_encodings) and self.cond_on_text_encodings:
+            assert text_encodings.shape[0] == batch_size, f'the text encodings being passed into the unet does not have the proper batch size - text encoding shape {text_encodings.shape} - required batch size is {batch_size}'
             assert self.text_embed_dim == text_encodings.shape[-1], f'the text encodings you are passing in have a dimension of {text_encodings.shape[-1]}, but the unet was created with text_embed_dim of {self.text_embed_dim}.'
 
             text_mask = torch.any(text_encodings != 0., dim = -1)

dalle2_pytorch/train_configs.py

@@ -133,6 +133,7 @@ class DiffusionPriorNetworkConfig(BaseModel):
     num_time_embeds: int = 1
     num_image_embeds: int = 1
     num_text_embeds: int = 1
+    attend_all_text_encodings: bool = True
     dim_head: int = 64
     heads: int = 8
     ff_mult: int = 4

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.22.0'
+__version__ = '0.22.3'