update laion

app.py

@@ -17,7 +17,7 @@ for k, v in ckpt.items():
     k = k[len('image_encoder.model.'):]
     new_dict.update({k: v})
 
-model = mae_vit_base_patch16(uni_dim=768, less_u=True)
+model = mae_vit_base_patch16(uni_dim=768, uni_heads=12, less_u=True)
 
 msg = model.load_state_dict(new_dict, strict=False)
 print(msg)

model.py

@@ -143,6 +143,7 @@ class ParallelTransformerBlock(nn.Module):
 
         attn_inner_dim = dim_head * heads
         ff_inner_dim = dim * ff_mult
+        # import ipdb; ipdb.set_trace()
         self.fused_dims = (attn_inner_dim, dim_head, dim_head, (ff_inner_dim * 2))
 
         self.heads = heads
@@ -431,7 +432,7 @@ class MaskedAutoencoderViT(nn.Module):
         # NOTE: +1 for mask token used by MLM objective
         # self.token_emb = nn.Embedding(len(self.tokenizer.vocab) + 1, uni_dim)
 
-        self.token_emb = nn.Embedding(len(self.tokenizer.vocab), uni_dim)
+        self.token_emb = nn.Embedding(len(self.tokenizer.vocab) + 1, uni_dim)
         self.text_cls_token = nn.Parameter(torch.randn(uni_dim))
 
         self.embed_dim = embed_dim
@@ -528,7 +529,7 @@ class MaskedAutoencoderViT(nn.Module):
         # self.text_mask_token = nn.Parameter(torch.randn(embed_dim))
         self.mask_token_id = len(self.tokenizer.vocab)
 
-        # self.text_position_embed = nn.Parameter(torch.zeros(1, text_length, embed_dim), requires_grad=False)
+        self.text_position_embed = nn.Parameter(torch.zeros(1, text_length, embed_dim), requires_grad=False)
         self.text_length = text_length
 
         self.latent_projector_layer = projector_layer