Upload config

config.json

@@ -1,37 +1,64 @@
 {
-  "arch_specifier": "no-align+gelu-mlp",
-  "architectures": [
-    "TrajectoryVLA"
-  ],
   "auto_map": {
-    "AutoModelForVision2Seq": "prismatic_model.TrajectoryVLA"
+    "AutoConfig": "prismatic_config.TrajectoryVLAConfig"
   },
-  "hf_llm_id": "meta-llama/Llama-2-7b-hf",
-  "image_resize_strategy": "letterbox",
-  "image_sizes": [
-    224,
-    224
-  ],
-  "llm_backbone_id": "llama2-7b-pure",
-  "llm_max_length": 2048,
-  "model_type": "prismatic",
-  "output_projector_states": false,
-  "pad_to_multiple_of": 64,
-  "pad_token_id": 32000,
-  "return_dict": false,
-  "text_config": {
-    "model_type": "llama"
+  "cheat": false,
+  "model_type": "trajectoryvla",
+  "num_timesteps": 6,
+  "prismatic_config": {
+    "architectures": [
+      "TrajectoryVLA"
+    ],
+    "auto_map": {
+      "AutoModelForVision2Seq": "prismatic_model.TrajectoryVLA"
+    },
+    "model_type": "prismatic",
+    "return_dict": false,
+    "torch_dtype": "bfloat16"
   },
-  "timm_model_ids": [
-    "vit_large_patch14_reg4_dinov2.lvd142m",
-    "vit_so400m_patch14_siglip_224"
-  ],
-  "timm_override_act_layers": [
-    null,
-    null
-  ],
-  "torch_dtype": "bfloat16",
-  "transformers_version": "4.44.2",
-  "use_fused_vision_backbone": true,
-  "vision_backbone_id": "dinosiglip-vit-so-224px"
+  "rotation_components": 9,
+  "seperate_control_proj": true,
+  "timestep_proj_config": {
+    "num_tokens": 3,
+    "pos_embed_scale": 8,
+    "proj_layers": [
+      128,
+      512,
+      1024
+    ],
+    "time_delta_sec": 0.1
+  },
+  "token_proj_config": {
+    "control_tokens_layers": [
+      4096,
+      2048,
+      1024
+    ],
+    "image_tokens_mode": "vit",
+    "llm_image_tokens_layers": [],
+    "vit_tokens_layers": [
+      2176,
+      1024
+    ]
+  },
+  "token_size": 1024,
+  "transformer_config": {
+    "decoder_block_config": {
+      "dropout": 0.0,
+      "feature_size": 1024,
+      "head_dim": 64,
+      "num_heads": 16
+    },
+    "encoder_block_config": {
+      "feature_size": 1024,
+      "head_dim": 64,
+      "num_heads": 16
+    },
+    "num_blocks": 2,
+    "pos_embed_config": {
+      "embedding_dim": 1024,
+      "num_embeddings": 300
+    }
+  },
+  "transformers_version": "4.44.2"
 }

prismatic_config.py

@@ -176,7 +176,8 @@ class TrajectoryVLAConfig(PretrainedConfig):
         # **kwargs: str,
     ):
 
-        super().__init__(**prismatic_config)
+        # super().__init__(**prismatic_config)
+        super().__init__()
         self.prismatic_config = PrismaticConfig(**prismatic_config)
 
         self.token_size = token_size
@@ -197,29 +198,6 @@ class TrajectoryVLAConfig(PretrainedConfig):
     @property
     def num_timestep_tokens(self) -> int:
         return self.timestep_proj_config['num_tokens']
-# class WaypointerConfig(ConfigurableModuleConfig):
-#     token_size: int = 1024  # Timestep token size
-
-#     cheat: bool  # If True, cheat and use action tokens; Works only with OpenVLA checkpoint
-
-#     timestep_proj_config: AutoConfig  # Timestep tokens
-#     token_proj_config: TokenProjectorConfig  # LLM output tokens projection and packing
-#     transformer_config: AutoConfig  # Transformer config
-
-#     # Output configurations
-#     num_timesteps: int = 20  # Number of prediction time steps
-#     rotation_components: int = 3  # Number of rotation componens: euler -> 3, quaternion -> 4, rotmat -> 9
-#     separate_control_proj: bool = True  # If True, project control components separately
-
-#     @property
-#     def control_components(self) -> int:
-#         # Number of control dimensions: 3 translation, N rotation, 1 gripper
-#         return 3 + self.rotation_components + 1
-
-#     @property
-#     def num_timestep_tokens(self) -> int:
-#         return self.timestep_proj_config.num_tokens
-
 
 class OpenVLAConfig(PrismaticConfig):
     model_type: str = "openvla"
@@ -237,73 +215,95 @@ class OpenVLAConfig(PrismaticConfig):
 if  __name__ == "__main__" :
     # yaml_file = 'barrel/pipes/vlams/configs/waypoints/waypointer_multistep_fractal.yaml'
 
-    prismatic_config = PrismaticConfig()
-    print(prismatic_config)
+    # prismatic_config = PrismaticConfig()
+    # print(prismatic_config)
 
     prismatic_config_dict = {
         "vision_backbone_id":"dinosiglip-vit-so-224px",
-        # "llm_backbone_id":"llama2-7b-pure",meta-llama/Llama-2-7b-hf
-        "llm_backbone_id": "meta-llama/Llama-2-7b-hf",
-
+        "llm_backbone_id":"llama2-7b-pure",
         "arch_specifier": "no-align+gelu-mlp", ## TODO: check
-        "use_fused_vision_backbone" :None, ## TODO: check
+        "use_fused_vision_backbone" :True, ## TODO: check
         "image_resize_strategy" : "letterbox",
         "text_config" : None,
         "llm_max_length"  : 2048,
         "pad_token_id" :32000,
         "pad_to_multiple_of" : 64,
         "output_projector_states" : False,
+        "return_dict": False,
     }
+
     token_proj_config = {
         "vit_tokens_layers": [2176, 1024],
         "control_tokens_layers": [4096, 2048, 1024],
         "image_tokens_mode": 'vit',
+        'llm_image_tokens_layers': []
     }
     timestep_proj_config = {
-        "pos_embed_scale": 1.0,
-        "proj_layers": [1024],
+        "pos_embed_scale": 8,
+        "proj_layers": [128,512,1024],
         "time_delta_sec": 0.1,
         "num_tokens":3
     }
-
-    TrajectoryVlaConfig = {
+    pos_embed_config = {
+        "num_embeddings": 300,
+        "embedding_dim": 1024
+    }
+    encoder_block_config = {
+        "feature_size": 1024,
+        "head_dim": 64,
+        "num_heads": 16
+    }
+    decoder_block_config = {
+        "feature_size": 1024,
+        "head_dim": 64,
+        "num_heads": 16,
+        "dropout": 0.0
+    }
+    transformer_config = {
+        "pos_embed_config": pos_embed_config,
+        "encoder_block_config": encoder_block_config,
+        "decoder_block_config": decoder_block_config,
+        "num_blocks": 2
+    }
+    TrajectoryVlaConfig_config = {
         "prismatic_config":prismatic_config_dict,
         "token_size": 1024,
         "cheat": False,
-        "num_timesteps": 20,
-        "rotation_components": 3,
+        "num_timesteps": 6,
+        "rotation_components": 9,
         "seperate_control_proj": True,
-        "timestep_proj_config": {},
-        "token_proj_config": {},
-        "transformer_config": {},
+        "timestep_proj_config": timestep_proj_config,
+        "token_proj_config": token_proj_config,
+        "transformer_config": transformer_config,
+        "num_timestep_tokens": 3,
     }
 
-    TrajectoryVLAConfig = TrajectoryVLAConfig( **TrajectoryVlaConfig)
+    TrajectoryVLAConfig = TrajectoryVLAConfig( **TrajectoryVlaConfig_config)
     print(TrajectoryVLAConfig)
 
-class WaypointTokenizer:
-    """
-    Wraps base LLM/VLM tokenizer and overloads least used token as a control token
+# class WaypointTokenizer:
+#     """
+#     Wraps base LLM/VLM tokenizer and overloads least used token as a control token
 
-    NOTE: By default, assumes a BPE-style tokenizer akin to the LlamaTokenizer,
-        where *the least used tokens* appear at the end of the vocabulary!
+#     NOTE: By default, assumes a BPE-style tokenizer akin to the LlamaTokenizer,
+#         where *the least used tokens* appear at the end of the vocabulary!
 
-    TODO: Adding new token vs overloading? When I call `tokenizer.add_token()` vocab stays the same
-    """
+#     TODO: Adding new token vs overloading? When I call `tokenizer.add_token()` vocab stays the same
+#     """
 
-    def __init__(self, tokenizer: transformers.PreTrainedTokenizerBase, num_tokens: int = 10) -> None:
-        self.tokenizer = tokenizer
-        self.num_tokens = num_tokens
+#     def __init__(self, tokenizer: transformers.PreTrainedTokenizerBase, num_tokens: int = 10) -> None:
+#         self.tokenizer = tokenizer
+#         self.num_tokens = num_tokens
 
-    def __call__(self, *_) -> str:
-        """Get the text token for control"""
-        return self.tokenizer.decode(self.control_token_ids)
+#     def __call__(self, *_) -> str:
+#         """Get the text token for control"""
+#         return self.tokenizer.decode(self.control_token_ids)
 
-    @property
-    def control_token_ids(self) -> np.ndarray:
-        # Assumes we're overwriting the final tokens of the vocabulary (least used tokens)
-        return np.arange(self.num_tokens) + int(self.tokenizer.vocab_size - self.num_tokens)
+#     @property
+#     def control_token_ids(self) -> np.ndarray:
+#         # Assumes we're overwriting the final tokens of the vocabulary (least used tokens)
+#         return np.arange(self.num_tokens) + int(self.tokenizer.vocab_size - self.num_tokens)
 
-    @property
-    def num_control_tokens(self) -> int:
-        return self.num_tokens
+#     @property
+#     def num_control_tokens(self) -> int:
+#         return self.num_tokens