Upload TrajectoryVLA

config.json

@@ -1,64 +1,37 @@
 {
+  "arch_specifier": "no-align+gelu-mlp",
+  "architectures": [
+    "TrajectoryVLA"
+  ],
   "auto_map": {
-    "AutoConfig": "prismatic_config.TrajectoryVLAConfig"
+    "AutoModelForVision2Seq": "prismatic_model.TrajectoryVLA"
   },
-  "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"
+  "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"
   },
-  "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"
+  "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"
 }

prismatic_model.py

@@ -26,7 +26,7 @@ import torch
 import torch.nn as nn
 import transformers
 from timm.models.vision_transformer import LayerScale
-from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
+from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel, AutoModel, AutoConfig
 from transformers.modeling_outputs import ModelOutput
 import collections
 import math
@@ -40,6 +40,8 @@ from PIL import Image
 from pathlib import Path
 from torch.amp.autocast_mode import autocast  # Corrected import for latest PyTorch
 from scipy.spatial.transform import Rotation as R
+# import automodel
+
 ht_token_path = Path(".hf_token")
 HF_TOKEN  = ht_token_path.read_text().strip() if isinstance(ht_token_path, Path) else hf_token_path
 
@@ -256,7 +258,10 @@ class LLMBackbone(nn.Module):
 
         return tokenizer
 
+# @AutoModel.register(PrismaticConfig)
 class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
+    # model_type: ClassVar[str] = "prismatic"
+    config_class: PretrainedConfig = PrismaticConfig
     def __init__(self, config: PrismaticConfig) -> None:
         super().__init__(config)
         # [Validation] Lightweight Validate on `config` Fields + Dependency Versions
@@ -773,7 +778,7 @@ class TimestepProjModule(nn.Module):
 
 
 # class Waypointer(nn.Module):
-
+# @AutoModel.register(TrajectoryVLAConfig)
 class TrajectoryVLA(PrismaticForConditionalGeneration):
 
 
@@ -781,6 +786,7 @@ class TrajectoryVLA(PrismaticForConditionalGeneration):
 
     def __init__(self, config: TrajectoryVLAConfig) -> None:
         super().__init__(config.prismatic_config)
+
         self.control_tokenizer = WaypointTokenizer(self.llm_backbone.tokenizer)
         self.timestep_proj = TimestepProjModule(
             config.timestep_proj_config,
@@ -1008,7 +1014,10 @@ def read_pt(pt_path):
 # control_target = read_pt('/work/nikolay_nikolov/debug/inference/control_target.pt')
 
 if __name__ == "__main__":
-
+    AutoConfig.register("prismatic",PrismaticConfig)
+    AutoConfig.register("trajectoryvla",TrajectoryVLAConfig)
+    AutoModel.register('prismatic',PrismaticForConditionalGeneration)
+    AutoModel.register('trajectoryvla',TrajectoryVLA)
     prismatic_config_dict = {
         "vision_backbone_id":"dinosiglip-vit-so-224px",
         "llm_backbone_id":"llama2-7b-pure",