Upload model

config.json

@@ -18,8 +18,7 @@
   "vision_encoder_config": {
     "anyres_patch_sampling": false,
     "image_aspect_ratio": "pad",
-    "model_type": "xgenmm_vision_encoder",
-    "temporal_encoder_mode": "gttm"
+    "model_type": "xgenmm_vision_encoder"
   },
   "vision_tokenizer_config": {
     "model_type": "xgenmm_vision_tokenizer"

modeling_xgenmm.py

@@ -78,14 +78,15 @@ class XGenMMConfig(PretrainedConfig):
             vision_encoder_config = {
                 "image_aspect_ratio": "pad",
                 "anyres_patch_sampling": False,
-                "temporal_encoder_mode": "gttm",
             }
             logger.info(
                 "vision_encoder_config is None. initializing the XGenMMVisionEncoderConfig with default values."
             )
 
         if vision_tokenizer_config is None:
-            vision_tokenizer_config = {}
+            vision_tokenizer_config = {
+                "temporal_encoder_mode": "gttm",
+            }
             logger.info(
                 "vision_tokenizer_config is None. Initializing the XGenMMVisionTokenizerConfig with default values."
             )
@@ -843,6 +844,64 @@ class TokenTuringMachineUnit(nn.Module):
         return (mem_out_tokens, output_tokens)
 
 
+class GroupedTokenTuringMachine7(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        output_size=32,
+        memory_size_per_group=4,
+        num_layers=4,
+        num_heads=8,
+    ):
+        super().__init__()
+
+        self.ttm_unit = GroupedTokenTuringMachineUnit(
+            dim=dim,
+            process_size=output_size,
+            memory_size_per_group=memory_size_per_group,
+            num_layers=num_layers,
+            num_heads=num_heads)
+
+        self.initial_memory = nn.Parameter(torch.randn(output_size, memory_size_per_group, dim))
+        
+        self.pos_emb = PositionalEncoding1D(dim)
+
+        self.initial_reduction = TokenLearnerAttentionModule(dim=dim, num_target_tokens=output_size)
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.Tensor):
+                shape (b, T, n, D)
+        """
+        b, T, n, D = x.shape
+
+        memory_tokens = repeat(self.initial_memory, "n g d -> b n g d", b=b)
+
+        mean_x = torch.mean(x, dim=-2, keepdim=False)
+        positional_embeddings = self.pos_emb(mean_x)  # (b, T, d)
+
+        for i in range(T):
+            step_tokens = x[:, i, :, :]
+
+            pos = positional_embeddings[:, i, :]
+            pos = pos.unsqueeze(1)
+            step_tokens = step_tokens + pos
+
+            step_tokens = self.initial_reduction(step_tokens)
+
+            # print(memory_tokens.shape)
+            # print(step_tokens.shape)
+
+            memory_tokens = self.ttm_unit(memory_tokens, step_tokens)
+
+        memory_tokens = torch.mean(memory_tokens, dim=-2, keepdim=False)
+        # memory_tokens = torch.amax(memory_tokens, dim=-2, keepdim=False)
+
+        return memory_tokens.unsqueeze(1)
+
+
 class GroupedTokenTuringMachine4(nn.Module):
     def __init__(
         self,
@@ -1000,10 +1059,7 @@ class TokenTuringMachine(nn.Module):
                 pos = pos.unsqueeze(1)
                 step_tokens = step_tokens + pos
 
-            # print(step_tokens.shape)
             memory_tokens, output_tokens = self.ttm_unit(memory_tokens, step_tokens)
-            # print(f'memory_tokens shape: {memory_tokens.shape}')
-            # print(f'output_tokens shape: {output_tokens.shape}')
             output_tokens_list.append(output_tokens)
 
         if self.final_output_only:
@@ -1016,6 +1072,31 @@ class TokenTuringMachine(nn.Module):
             return output_tokens
 
 
+class TokenLearner(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        output_size=128,
+    ):
+        super().__init__()
+
+        self.final_output = TokenLearnerAttentionModule(dim=dim, num_target_tokens=output_size)
+
+    def forward(self, x):
+        """
+        Args:
+            x (torch.Tensor):
+                shape (b, T, n, D)
+        """
+        b, T, n, D = x.shape
+
+        output_tokens = x.view(b, -1, D)
+        output_tokens = self.final_output(output_tokens)
+
+        return output_tokens.unsqueeze(1)
+
+
 def num_params(module, filter_to_trainable=False):
     """Returns the number of parameters in the module, or optionally only the trainable parameters"""
     if filter_to_trainable:
@@ -1094,8 +1175,13 @@ class PerceiverResampler(VisionTokenizer):
         if self.temporal_encoder_mode=='gttm':
             # self.ttm = TokenTuringMachine(dim=dim, memory_size=128, memory_out_mode=True)
             self.temporal_encoder = GroupedTokenTuringMachine(dim=dim, process_size=128, memory_size_per_group=4)
-        elif self.temporal_encoder_mode=='gttm_pool':
+        elif self.temporal_encoder_mode=='gttm4':
             self.temporal_encoder = GroupedTokenTuringMachine4(dim=dim, process_size=128, memory_size_per_group=4, output_size=32)
+        elif self.temporal_encoder_mode=='tokenlearner':
+            self.temporal_encoder = TokenLearner(dim=dim, output_size=32)
+        elif self.temporal_encoder_mode=='gttm7':
+            self.temporal_encoder = GroupedTokenTuringMachine7(dim=dim, memory_size_per_group=4, output_size=32)
+
 
     def forward(self, x, vision_attn_masks):
         """
@@ -1126,7 +1212,7 @@ class PerceiverResampler(VisionTokenizer):
             latents = attn(x, latents, vision_attn_masks) + latents
             latents = ff(latents) + latents
 
-        if self.video_mode is not None:
+        if self.temporal_encoder_mode is not None:
             latents = self.temporal_encoder(latents)
 
         if exists(self.projection):