Add files using upload-large-folder tool

openvla-7b-prismatic/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

openvla-7b-prismatic/README.md

@@ -0,0 +1,41 @@
+---
+library_name: transformers
+tags:
+- robotics
+- vla
+- image-text-to-text
+- multimodal
+- pretraining
+license: mit
+language:
+- en
+pipeline_tag: image-text-to-text
+---
+
+# OpenVLA 7B (Prismatic-Compatible Version)
+
+<b>This is the same model as the [OpenVLA 7B model](https://huggingface.co/openvla/openvla-7b), except that this checkpoint is in a format that is
+compatible with the training script from the original [Prismatic VLMs project codebase](https://github.com/TRI-ML/prismatic-vlms), which the OpenVLA
+team built on top of to develop the OpenVLA model. See details for the OpenVLA 7B model here: https://huggingface.co/openvla/openvla-7b</b>
+
+This Prismatic-compatible checkpoint may be useful if you wish to <b>fully fine-tune</b> OpenVLA (all 7.5 billion parameters) via native PyTorch Fully
+Sharded Data Parallel (FSDP) using the Prismatic VLMs training script. If you instead wish to do Parameter-Efficient Fine-Tuning via LoRA, you
+can use the OpenVLA checkpoint linked above, which is compatible with the Hugging Face `transformers` library. We recommend fine-tuning via LoRA if
+you do not have sufficient compute to fully fine-tune a 7B-parameter model (e.g., multiple A100/H100 GPUs).
+
+## Usage Instructions
+
+See the [OpenVLA GitHub README](https://github.com/openvla/openvla/blob/main/README.md) for instructions on how to use this checkpoint for full fine-tuning.
+
+## Citation
+
+**BibTeX:**
+
+```bibtex
+@article{kim24openvla,
+    title={OpenVLA: An Open-Source Vision-Language-Action Model},
+    author={{Moo Jin} Kim and Karl Pertsch and Siddharth Karamcheti and Ted Xiao and Ashwin Balakrishna and Suraj Nair and Rafael Rafailov and Ethan Foster and Grace Lam and Pannag Sanketi and Quan Vuong and Thomas Kollar and Benjamin Burchfiel and Russ Tedrake and Dorsa Sadigh and Sergey Levine and Percy Liang and Chelsea Finn},
+    journal = {arXiv preprint arXiv:2406.09246},
+    year={2024}
+} 
+```

openvla-7b-prismatic/checkpoints/step-295000-epoch-40-loss=0.2200.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2c2497cd9e0ecced65e54b0771172f22e3ed64d0c0af339e094349715d3b3602
+size 30165309772

openvla-7b-prismatic/config.json

@@ -0,0 +1,41 @@
+{
+  "data_root_dir": "/scr/user/data",
+  "hf_token": ".hf_token",
+  "pretrained_checkpoint": "",
+  "resume_epoch": null,
+  "resume_step": null,
+  "run_id": "prism-dinosiglip-224px+mx-oxe-magic-soup-plus+n8+b32+x7",
+  "run_id_note": null,
+  "run_root_dir": "./runs",
+  "save_interval": 2500,
+  "seed": 7,
+  "stage": "vla-full-train",
+  "trackers": [
+    "jsonl",
+    "wandb"
+  ],
+  "vla": {
+    "base_vlm": "prism-dinosiglip-224px+7b",
+    "data_mix": "oxe_magic_soup_plus_minus",
+    "enable_gradient_checkpointing": true,
+    "enable_mixed_precision_training": true,
+    "epochs": 1000,
+    "expected_world_size": 64,
+    "freeze_vision_backbone": false,
+    "global_batch_size": 2048,
+    "learning_rate": 2e-05,
+    "lr_scheduler_type": "constant",
+    "max_grad_norm": 1.0,
+    "max_steps": null,
+    "per_device_batch_size": 32,
+    "reduce_in_full_precision": true,
+    "shuffle_buffer_size": 256000,
+    "train_strategy": "fsdp-full-shard",
+    "type": "prism-dinosiglip-224px+mx-oxe-magic-soup-plus",
+    "vla_id": "prism-dinosiglip-224px+mx-oxe-magic-soup-plus",
+    "warmup_ratio": 0.0,
+    "weight_decay": 0.0
+  },
+  "wandb_entity": "",
+  "wandb_project": ""
+}

openvla-7b-prismatic/config.yml

@@ -0,0 +1,37 @@
+data_root_dir: /scr/user/data
+hf_token: .hf_token
+pretrained_checkpoint: null
+resume_epoch: null
+resume_step: null
+run_id: prism-dinosiglip-224px+mx-oxe-magic-soup-plus+n8+b32+x7
+run_id_note: null
+run_root_dir: ./runs
+save_interval: 2500
+seed: 7
+stage: vla-full-train
+trackers:
+- jsonl
+- wandb
+vla:
+  base_vlm: prism-dinosiglip-224px+7b
+  data_mix: oxe_magic_soup_plus_minus
+  enable_gradient_checkpointing: true
+  enable_mixed_precision_training: true
+  epochs: 1000
+  expected_world_size: 64
+  freeze_vision_backbone: false
+  global_batch_size: 2048
+  learning_rate: 2.0e-05
+  lr_scheduler_type: constant
+  max_grad_norm: 1.0
+  max_steps: null
+  per_device_batch_size: 32
+  reduce_in_full_precision: true
+  shuffle_buffer_size: 256000
+  train_strategy: fsdp-full-shard
+  type: prism-dinosiglip-224px+mx-oxe-magic-soup-plus
+  vla_id: prism-dinosiglip-224px+mx-oxe-magic-soup-plus
+  warmup_ratio: 0.0
+  weight_decay: 0.0
+wandb_entity: null
+wandb_project: null

openvla-7b-prismatic/dataset_statistics.json

@@ -0,0 +1,3127 @@
+{
+  "fractal20220817_data": {
+    "action": {
+      "mean": [
+        0.006987582892179489,
+        0.006265917327255011,
+        -0.01262515690177679,
+        0.04333311319351196,
+        -0.005756212864071131,
+        0.0009130256366916001,
+        0.5354204773902893
+      ],
+      "std": [
+        0.0692116990685463,
+        0.05970962345600128,
+        0.07353084534406662,
+        0.15610496699810028,
+        0.13164450228214264,
+        0.14593800902366638,
+        0.497110515832901
+      ],
+      "max": [
+        2.9984593391418457,
+        22.09052848815918,
+        2.7507524490356445,
+        1.570636510848999,
+        1.5321086645126343,
+        1.5691522359848022,
+        1.0
+      ],
+      "min": [
+        -2.0204520225524902,
+        -5.497899532318115,
+        -2.031663417816162,
+        -1.569917917251587,
+        -1.569892168045044,
+        -1.570419430732727,
+        0.0
+      ],
+      "q01": [
+        -0.22453527510166169,
+        -0.14820013284683228,
+        -0.231589707583189,
+        -0.3517994859814644,
+        -0.4193011274933815,
+        -0.43643461108207704,
+        0.0
+      ],
+      "q99": [
+        0.17824687153100965,
+        0.14938379630446405,
+        0.21842354819178575,
+        0.5892666035890578,
+        0.35272657424211445,
+        0.44796681255102094,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 3786400,
+    "num_trajectories": 87212
+  },
+  "kuka": {
+    "action": {
+      "mean": [
+        -0.0004668905457947403,
+        0.00040138536132872105,
+        -0.001280792523175478,
+        0.0,
+        0.0,
+        -0.03722453489899635,
+        0.4131543040275574
+      ],
+      "std": [
+        0.02083250693976879,
+        0.02915887162089348,
+        0.06422865390777588,
+        0.0,
+        0.0,
+        0.14224295318126678,
+        0.49086448550224304
+      ],
+      "max": [
+        0.1697135865688324,
+        0.2777623236179352,
+        0.43710532784461975,
+        0.0,
+        0.0,
+        1.9684287309646606,
+        1.0
+      ],
+      "min": [
+        -0.159867063164711,
+        -0.2892282009124756,
+        -0.2795473635196686,
+        0.0,
+        0.0,
+        -1.9875637292861938,
+        0.0
+      ],
+      "q01": [
+        -0.06619441494345665,
+        -0.08713878810405731,
+        -0.15083016991615295,
+        0.0,
+        0.0,
+        -0.5415697038173676,
+        0.0
+      ],
+      "q99": [
+        0.06601839080452929,
+        0.08732476785779003,
+        0.18168179214000715,
+        0.0,
+        0.0,
+        0.2923380345106127,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 2455879,
+    "num_trajectories": 209880
+  },
+  "bridge_orig": {
+    "action": {
+      "mean": [
+        0.0002334194869035855,
+        0.00013004911306779832,
+        -0.00012762474943883717,
+        -0.0001556558854645118,
+        -0.0004039328487124294,
+        0.00023557482927571982,
+        0.5764579176902771
+      ],
+      "std": [
+        0.009765930473804474,
+        0.013689135201275349,
+        0.012667362578213215,
+        0.028534092009067535,
+        0.030637972056865692,
+        0.07691419124603271,
+        0.4973701536655426
+      ],
+      "max": [
+        0.41691166162490845,
+        0.25864794850349426,
+        0.21218234300613403,
+        3.122201919555664,
+        1.8618112802505493,
+        6.280478477478027,
+        1.0
+      ],
+      "min": [
+        -0.4007510244846344,
+        -0.13874775171279907,
+        -0.22553899884223938,
+        -3.2010786533355713,
+        -1.8618112802505493,
+        -6.279075622558594,
+        0.0
+      ],
+      "q01": [
+        -0.02872725307941437,
+        -0.04170349963009357,
+        -0.026093858778476715,
+        -0.08092105075716972,
+        -0.09288699507713317,
+        -0.20718276381492615,
+        0.0
+      ],
+      "q99": [
+        0.028309678435325586,
+        0.040855254605412394,
+        0.040161586627364146,
+        0.08192047759890528,
+        0.07792850524187081,
+        0.20382574498653397,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 2135463,
+    "num_trajectories": 60064
+  },
+  "taco_play": {
+    "action": {
+      "mean": [
+        -0.003845922416076064,
+        0.009671456180512905,
+        0.012780580669641495,
+        -0.005403771996498108,
+        -0.009606587700545788,
+        -0.002480733208358288,
+        0.4263913035392761
+      ],
+      "std": [
+        0.23254038393497467,
+        0.36298269033432007,
+        0.28692901134490967,
+        0.2617705166339874,
+        0.2438892275094986,
+        0.5216503143310547,
+        0.4946896731853485
+      ],
+      "max": [
+        1.4915844202041626,
+        2.1842432022094727,
+        2.6836395263671875,
+        5.035226821899414,
+        2.665864944458008,
+        4.250768661499023,
+        1.0
+      ],
+      "min": [
+        -4.242457866668701,
+        -3.192805051803589,
+        -1.3371467590332031,
+        -4.202683448791504,
+        -2.6722638607025146,
+        -3.3467135429382324,
+        0.0
+      ],
+      "q01": [
+        -0.7106140398979186,
+        -1.056944659948349,
+        -0.5878450274467468,
+        -0.7682853937149048,
+        -0.7180147767066956,
+        -1.5527938604354858,
+        0.0
+      ],
+      "q99": [
+        0.6482916426658629,
+        1.0051310062408447,
+        0.9480248689651489,
+        0.6926478147506714,
+        0.6351067513227462,
+        1.628010264635086,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 237798,
+    "num_trajectories": 3603
+  },
+  "jaco_play": {
+    "action": {
+      "mean": [
+        0.0009658430935814977,
+        -0.00580078037455678,
+        -0.00395062193274498,
+        0.0,
+        0.0,
+        0.0,
+        0.34934908151626587
+      ],
+      "std": [
+        0.12235074490308762,
+        0.09678777307271957,
+        0.11155334860086441,
+        0.0,
+        0.0,
+        0.0,
+        0.4768252968788147
+      ],
+      "max": [
+        0.20000000298023224,
+        0.20000000298023224,
+        0.20000000298023224,
+        0.0,
+        0.0,
+        0.0,
+        1.0
+      ],
+      "min": [
+        -0.20000000298023224,
+        -0.20000000298023224,
+        -0.20000000298023224,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -0.20000000298023224,
+        -0.20000000298023224,
+        -0.20000000298023224,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.20000000298023224,
+        0.20000000298023224,
+        0.20000000298023224,
+        0.0,
+        0.0,
+        0.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 77965,
+    "num_trajectories": 1085
+  },
+  "berkeley_cable_routing": {
+    "action": {
+      "mean": [
+        -0.07139874249696732,
+        0.023609008640050888,
+        0.10241943597793579,
+        0.0,
+        0.0,
+        0.049671024084091187,
+        0.0
+      ],
+      "std": [
+        0.1815500408411026,
+        0.1810990273952484,
+        0.21220779418945312,
+        0.0,
+        0.0,
+        0.3475511968135834,
+        0.0
+      ],
+      "max": [
+        0.9633283019065857,
+        1.0,
+        1.0,
+        0.0,
+        0.0,
+        1.0,
+        0.0
+      ],
+      "min": [
+        -0.9809081554412842,
+        -0.9554349184036255,
+        -0.9994775056838989,
+        0.0,
+        0.0,
+        -1.0,
+        0.0
+      ],
+      "q01": [
+        -0.5534318816661835,
+        -0.4797285574674606,
+        -0.5314934802055359,
+        0.0,
+        0.0,
+        -0.8855219376087189,
+        0.0
+      ],
+      "q99": [
+        0.42652835428714786,
+        0.5000944086909298,
+        0.639823433756829,
+        0.0,
+        0.0,
+        0.984243879914284,
+        0.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 42328,
+    "num_trajectories": 1647
+  },
+  "roboturk": {
+    "action": {
+      "mean": [
+        0.0014448732836171985,
+        -0.0015945249469950795,
+        -0.0011753785656765103,
+        0.0023012510500848293,
+        -0.0009382463176734746,
+        -0.00011485807772260159,
+        0.5746025443077087
+      ],
+      "std": [
+        0.04935386776924133,
+        0.0635455846786499,
+        0.061164740473032,
+        0.09553450345993042,
+        0.08420111238956451,
+        0.06517903506755829,
+        0.49452081322669983
+      ],
+      "max": [
+        0.39124172925949097,
+        0.4601028263568878,
+        0.4870833456516266,
+        1.816888689994812,
+        1.8240282535552979,
+        1.4824820756912231,
+        1.0
+      ],
+      "min": [
+        -0.6546999216079712,
+        -0.6365841031074524,
+        -0.4217723608016968,
+        -1.6695482730865479,
+        -1.8023357391357422,
+        -1.4630827903747559,
+        0.0
+      ],
+      "q01": [
+        -0.1342635464668274,
+        -0.19996687173843383,
+        -0.1482972100377083,
+        -0.20720748245716095,
+        -0.09676413893699647,
+        -0.18075634717941286,
+        0.0
+      ],
+      "q99": [
+        0.14956976801157001,
+        0.1805950567126275,
+        0.18841815620660796,
+        0.21615413755178453,
+        0.09457383215427405,
+        0.18543301910162005,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 187507,
+    "num_trajectories": 1995
+  },
+  "viola": {
+    "action": {
+      "mean": [
+        0.04761844128370285,
+        -0.029204415157437325,
+        0.05586736649274826,
+        -0.002618510741740465,
+        0.006867344491183758,
+        -0.01682133786380291,
+        0.7323777675628662
+      ],
+      "std": [
+        0.39157867431640625,
+        0.4076525568962097,
+        0.40077948570251465,
+        0.10023996233940125,
+        0.0844319611787796,
+        0.10375042259693146,
+        0.44260647892951965
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.375,
+        0.36321428418159485,
+        0.375,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.375,
+        -0.375,
+        -0.375,
+        0.0
+      ],
+      "q01": [
+        -0.9628571271896362,
+        -1.0,
+        -1.0,
+        -0.26249998807907104,
+        -0.21321429312229156,
+        -0.3385714292526245,
+        0.0
+      ],
+      "q99": [
+        0.9114285707473755,
+        0.868571400642395,
+        1.0,
+        0.2817857265472412,
+        0.2239285707473755,
+        0.3557142913341522,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 76324,
+    "num_trajectories": 150
+  },
+  "berkeley_autolab_ur5": {
+    "action": {
+      "mean": [
+        0.0005683620693162084,
+        0.001217700308188796,
+        -0.0005296372692100704,
+        0.00021029810886830091,
+        6.0695128922816366e-05,
+        0.001204986940138042,
+        0.6298308372497559
+      ],
+      "std": [
+        0.0115329809486866,
+        0.007990492507815361,
+        0.009577835910022259,
+        0.009432995691895485,
+        0.016427582129836082,
+        0.011053967289626598,
+        0.48267969489097595
+      ],
+      "max": [
+        0.019999999552965164,
+        0.019999999552965164,
+        0.019999999552965164,
+        0.06666667014360428,
+        0.06666667014360428,
+        0.06666667014360428,
+        1.0
+      ],
+      "min": [
+        -0.019999999552965164,
+        -0.019999999552965164,
+        -0.019999999552965164,
+        -0.06666667014360428,
+        -0.06666667014360428,
+        -0.06666667014360428,
+        0.0
+      ],
+      "q01": [
+        -0.019999999552965164,
+        -0.019999999552965164,
+        -0.019999999552965164,
+        -0.02628571353852749,
+        -0.06666667014360428,
+        -0.03847619146108627,
+        0.0
+      ],
+      "q99": [
+        0.019999999552965164,
+        0.019999999552965164,
+        0.019999999552965164,
+        0.031809523701667786,
+        0.06666667014360428,
+        0.036571428179740906,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 97939,
+    "num_trajectories": 1000
+  },
+  "toto": {
+    "action": {
+      "mean": [
+        0.38542115688323975,
+        0.007769413758069277,
+        0.3632740378379822,
+        -0.6652036905288696,
+        0.1890396922826767,
+        0.03298724442720413,
+        0.0
+      ],
+      "std": [
+        0.12211652100086212,
+        0.19378550350666046,
+        0.10178236663341522,
+        0.5725259184837341,
+        0.29884573817253113,
+        0.3259911835193634,
+        0.0
+      ],
+      "max": [
+        0.6839867234230042,
+        0.4454185664653778,
+        0.7984078526496887,
+        2.120781660079956,
+        1.371164321899414,
+        1.4118704795837402,
+        0.0
+      ],
+      "min": [
+        0.09922284632921219,
+        -0.5180193781852722,
+        0.13791072368621826,
+        -2.635117530822754,
+        -1.0734480619430542,
+        -1.9282547235488892,
+        0.0
+      ],
+      "q01": [
+        0.1756722891330719,
+        -0.3077590811252594,
+        0.235383919775486,
+        -2.0908505964279174,
+        -0.6191593289375306,
+        -0.7488683319091797,
+        0.0
+      ],
+      "q99": [
+        0.6136963081359863,
+        0.33704194784164443,
+        0.6681221985816956,
+        0.7422861719131538,
+        0.7955395007133507,
+        0.740464625358582,
+        0.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 325699,
+    "num_trajectories": 1003
+  },
+  "stanford_hydra_dataset_converted_externally_to_rlds": {
+    "action": {
+      "mean": [
+        0.0007790001109242439,
+        0.00013707754260394722,
+        -0.0002548607881180942,
+        0.0012903271708637476,
+        -0.004751681815832853,
+        0.002692886395379901,
+        0.48855218291282654
+      ],
+      "std": [
+        0.008022161200642586,
+        0.009131459519267082,
+        0.009574338793754578,
+        0.04122216999530792,
+        0.0384303517639637,
+        0.04606688767671585,
+        0.49976691603660583
+      ],
+      "max": [
+        0.02499854564666748,
+        0.02499903365969658,
+        0.024999922141432762,
+        0.24974457919597626,
+        0.24997030198574066,
+        0.24999946355819702,
+        1.0
+      ],
+      "min": [
+        -0.024999044835567474,
+        -0.024999700486660004,
+        -0.02499929815530777,
+        -0.24993225932121277,
+        -0.2499666064977646,
+        -0.2499932497739792,
+        0.0
+      ],
+      "q01": [
+        -0.019992006458342076,
+        -0.02415412735193968,
+        -0.022941758055239916,
+        -0.11085530579090118,
+        -0.12024572037160397,
+        -0.13314770206809043,
+        0.0
+      ],
+      "q99": [
+        0.022886231057345868,
+        0.022358838934451335,
+        0.02410089675337076,
+        0.12370114490389822,
+        0.11323311634361738,
+        0.18474749639630164,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 358234,
+    "num_trajectories": 570
+  },
+  "austin_buds_dataset_converted_externally_to_rlds": {
+    "action": {
+      "mean": [
+        -0.07678354531526566,
+        0.0036849044263362885,
+        0.05644911900162697,
+        0.0,
+        0.0,
+        0.0,
+        0.3510494828224182
+      ],
+      "std": [
+        0.6367740631103516,
+        0.37889179587364197,
+        0.47796326875686646,
+        0.0,
+        0.0,
+        0.0,
+        0.47721168398857117
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.0,
+        0.0,
+        0.0,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -1.0,
+        -0.9599999785423279,
+        -0.8714285492897034,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        0.8600000143051147,
+        1.0,
+        0.0,
+        0.0,
+        0.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 34112,
+    "num_trajectories": 50
+  },
+  "nyu_franka_play_dataset_converted_externally_to_rlds": {
+    "action": {
+      "mean": [
+        0.001021989737637341,
+        -0.00012002651783404872,
+        0.00032894269679673016,
+        0.0015034361276775599,
+        -0.002198522910475731,
+        -0.001663230243138969,
+        0.7230083346366882
+      ],
+      "std": [
+        0.01327415369451046,
+        0.013215910643339157,
+        0.012822109274566174,
+        0.2732451558113098,
+        0.057022541761398315,
+        0.039172880351543427,
+        0.44752755761146545
+      ],
+      "max": [
+        0.06424188613891602,
+        0.07027634978294373,
+        0.06129661202430725,
+        6.281067848205566,
+        0.1967729926109314,
+        0.26377415657043457,
+        1.0
+      ],
+      "min": [
+        -0.05952230095863342,
+        -0.07232445478439331,
+        -0.06730806827545166,
+        -6.278434753417969,
+        -0.21479034423828125,
+        -0.3627619743347168,
+        0.0
+      ],
+      "q01": [
+        -0.03199600875377655,
+        -0.032861671447753905,
+        -0.03368805110454559,
+        -0.12080862045288086,
+        -0.12175218224525451,
+        -0.11370223641395569,
+        0.0
+      ],
+      "q99": [
+        0.03101520001888276,
+        0.0373908892273903,
+        0.03646374464035038,
+        0.11764093399047852,
+        0.1258920183777809,
+        0.09366151213645942,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 44875,
+    "num_trajectories": 456
+  },
+  "furniture_bench_dataset_converted_externally_to_rlds": {
+    "action": {
+      "mean": [
+        0.00014610752987209707,
+        0.0010830952087417245,
+        0.0006224989192560315,
+        -0.003303206292912364,
+        -0.0026880695950239897,
+        0.018242603167891502,
+        0.48854944109916687
+      ],
+      "std": [
+        0.01610708422958851,
+        0.014891477301716805,
+        0.014014219865202904,
+        0.058274295181035995,
+        0.11417088657617569,
+        0.33479776978492737,
+        0.49991825222969055
+      ],
+      "max": [
+        0.10000000149011612,
+        0.10000000149011612,
+        0.10000000149011612,
+        0.8651833534240723,
+        1.0909736156463623,
+        2.863185405731201,
+        1.0
+      ],
+      "min": [
+        -0.10495579987764359,
+        -0.10939455777406693,
+        -0.10000000149011612,
+        -0.971906840801239,
+        -1.0475432872772217,
+        -3.06000018119812,
+        0.0
+      ],
+      "q01": [
+        -0.053988199681043625,
+        -0.05049169331789017,
+        -0.032499241530895236,
+        -0.1953887003660202,
+        -0.41674559473991396,
+        -0.8886768388748169,
+        0.0
+      ],
+      "q99": [
+        0.05414841488003723,
+        0.04965164884924884,
+        0.060055799782276154,
+        0.18231668293476103,
+        0.39867786407470646,
+        0.8772023963928218,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 3948057,
+    "num_trajectories": 5100
+  },
+  "ucsd_kitchen_dataset_converted_externally_to_rlds": {
+    "action": {
+      "mean": [
+        410.37567138671875,
+        116.9518814086914,
+        192.35032653808594,
+        -121.22441864013672,
+        -33.84893035888672,
+        50.016136169433594,
+        0.741813600063324
+      ],
+      "std": [
+        122.81494903564453,
+        108.8009033203125,
+        130.303466796875,
+        116.28205108642578,
+        27.621843338012695,
+        41.02094650268555,
+        0.43763357400894165
+      ],
+      "max": [
+        678.0,
+        400.0,
+        507.0,
+        180.00001525878906,
+        6.000013828277588,
+        116.99998474121094,
+        1.0
+      ],
+      "min": [
+        172.0,
+        -166.0,
+        -99.99999237060547,
+        -180.00001525878906,
+        -89.0,
+        -96.00010681152344,
+        0.0
+      ],
+      "q01": [
+        200.00001052856445,
+        -102.31004211425781,
+        -94.99993370056153,
+        -180.00001525878906,
+        -88.00001525878906,
+        -38.999977111816406,
+        0.0
+      ],
+      "q99": [
+        637.0,
+        368.30999999999995,
+        493.0,
+        180.00001525878906,
+        0.999983012676239,
+        105.00001525878906,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 3970,
+    "num_trajectories": 150
+  },
+  "austin_sailor_dataset_converted_externally_to_rlds": {
+    "action": {
+      "mean": [
+        0.011825348250567913,
+        0.006461074110120535,
+        0.06023626774549484,
+        0.0,
+        0.0,
+        0.0016465914668515325,
+        0.5260950326919556
+      ],
+      "std": [
+        0.46348899602890015,
+        0.41240179538726807,
+        0.411862850189209,
+        0.0,
+        0.0,
+        0.0578610822558403,
+        0.49894046783447266
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.0,
+        0.0,
+        0.375,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        0.0,
+        0.0,
+        -0.375,
+        0.0
+      ],
+      "q01": [
+        -1.0,
+        -0.9828571677207947,
+        -0.6000000238418579,
+        0.0,
+        0.0,
+        -0.17249999940395355,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        0.9457142949104309,
+        1.0,
+        0.0,
+        0.0,
+        0.17892856895923615,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 353094,
+    "num_trajectories": 240
+  },
+  "austin_sirius_dataset_converted_externally_to_rlds": {
+    "action": {
+      "mean": [
+        0.07747682929039001,
+        0.03195561468601227,
+        0.04244732856750488,
+        0.0,
+        0.0,
+        -0.01603456400334835,
+        0.43260177969932556
+      ],
+      "std": [
+        0.3906329572200775,
+        0.2998155355453491,
+        0.2782271206378937,
+        0.0,
+        0.0,
+        0.08120622485876083,
+        0.49528297781944275
+      ],
+      "max": [
+        1.0002285242080688,
+        0.960608720779419,
+        1.105179786682129,
+        0.0,
+        0.0,
+        0.341785728931427,
+        1.0
+      ],
+      "min": [
+        -1.0183025598526,
+        -0.9800000190734863,
+        -0.9774575233459473,
+        0.0,
+        0.0,
+        -0.34607142210006714,
+        0.0
+      ],
+      "q01": [
+        -0.780905865430832,
+        -0.5667179036140442,
+        -0.5254343223571777,
+        0.0,
+        0.0,
+        -0.28495091378688814,
+        0.0
+      ],
+      "q99": [
+        0.9569637751579284,
+        0.6971374487876891,
+        0.8124888157844541,
+        0.0,
+        0.0,
+        0.1971428543329239,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 279939,
+    "num_trajectories": 559
+  },
+  "dlr_edan_shared_control_converted_externally_to_rlds": {
+    "action": {
+      "mean": [
+        0.006647810339927673,
+        -0.0007657372043468058,
+        0.006522852927446365,
+        0.0011679717572405934,
+        -0.006395625416189432,
+        -0.011902998201549053,
+        0.6985887289047241
+      ],
+      "std": [
+        0.021393608301877975,
+        0.01814231649041176,
+        0.03374375030398369,
+        0.01743541844189167,
+        0.03394376486539841,
+        0.04641875624656677,
+        0.4588589072227478
+      ],
+      "max": [
+        0.18991442024707794,
+        0.0739002525806427,
+        0.18064819276332855,
+        0.0866486132144928,
+        0.13464981317520142,
+        0.16910280287265778,
+        1.0
+      ],
+      "min": [
+        -0.10054297000169754,
+        -0.08427435159683228,
+        -0.13533438742160797,
+        -0.17556548118591309,
+        -0.18485672771930695,
+        -0.2680685818195343,
+        0.0
+      ],
+      "q01": [
+        -0.02987122368067503,
+        -0.06013262912631035,
+        -0.08286409199237824,
+        -0.05924444157630205,
+        -0.15986866518855095,
+        -0.15636983573436739,
+        0.0
+      ],
+      "q99": [
+        0.08832092039287087,
+        0.042126184627413736,
+        0.11311905644834042,
+        0.0643695573508739,
+        0.03941855944693088,
+        0.156646853685379,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 8928,
+    "num_trajectories": 104
+  },
+  "iamlab_cmu_pickup_insert_converted_externally_to_rlds": {
+    "action": {
+      "mean": [
+        0.5274372696876526,
+        0.02858201041817665,
+        0.18712575733661652,
+        1.2339589595794678,
+        0.03226623684167862,
+        -1.4199490547180176,
+        0.5550631880760193
+      ],
+      "std": [
+        0.08108345419168472,
+        0.1116757020354271,
+        0.07747554779052734,
+        2.8737246990203857,
+        0.02774704433977604,
+        2.7678682804107666,
+        0.49695101380348206
+      ],
+      "max": [
+        0.6634981632232666,
+        0.23428471386432648,
+        0.4308285415172577,
+        3.1415927410125732,
+        0.13647015392780304,
+        3.141592502593994,
+        1.0
+      ],
+      "min": [
+        0.3071657121181488,
+        -0.29754969477653503,
+        0.06578229367733002,
+        -3.1415927410125732,
+        -0.04584203287959099,
+        -3.141592502593994,
+        0.0
+      ],
+      "q01": [
+        0.3148897051811218,
+        -0.20317550599575043,
+        0.06785467118024827,
+        -3.140952730178833,
+        -0.029743434861302376,
+        -3.141091251373291,
+        0.0
+      ],
+      "q99": [
+        0.6472805738449097,
+        0.20846802592277527,
+        0.36855655312538155,
+        3.1409926891326903,
+        0.11424950212240226,
+        3.1410969257354737,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 146241,
+    "num_trajectories": 631
+  },
+  "utaustin_mutex": {
+    "action": {
+      "mean": [
+        0.06176406890153885,
+        -0.005005486309528351,
+        0.10216785222291946,
+        -0.03314131125807762,
+        0.013895004987716675,
+        -0.011317633092403412,
+        0.5038976669311523
+      ],
+      "std": [
+        0.1875014752149582,
+        0.4468473494052887,
+        0.3792876601219177,
+        0.14097853004932404,
+        0.06453701853752136,
+        0.11765272170305252,
+        0.501045286655426
+      ],
+      "max": [
+        1.0,
+        1.0,
+        1.0,
+        0.375,
+        0.375,
+        0.375,
+        1.0
+      ],
+      "min": [
+        -1.0,
+        -1.0,
+        -1.0,
+        -0.375,
+        -0.375,
+        -0.375,
+        0.0
+      ],
+      "q01": [
+        -0.4285714328289032,
+        -0.9800000190734863,
+        -0.5571428537368774,
+        -0.375,
+        -0.15642857551574707,
+        -0.335357129573822,
+        0.0
+      ],
+      "q99": [
+        0.5914285778999329,
+        0.9714285731315613,
+        1.0,
+        0.3278571367263794,
+        0.207857146859169,
+        0.25607141852378845,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 361883,
+    "num_trajectories": 1500
+  },
+  "berkeley_fanuc_manipulation": {
+    "action": {
+      "mean": [
+        0.0007744057802483439,
+        -0.00031240080716088414,
+        -0.0015001941937953234,
+        -0.0007515158504247665,
+        -0.00015832878125365824,
+        0.00014327642566058785,
+        0.699295699596405
+      ],
+      "std": [
+        0.0034070091787725687,
+        0.0049921851605176926,
+        0.005344334989786148,
+        0.00759894959628582,
+        0.004081866703927517,
+        0.008568956516683102,
+        0.4586937427520752
+      ],
+      "max": [
+        0.009999999776482582,
+        0.009999999776482582,
+        0.009999999776482582,
+        0.03490658476948738,
+        0.03490658476948738,
+        0.03490658476948738,
+        1.0
+      ],
+      "min": [
+        -0.009999999776482582,
+        -0.009999999776482582,
+        -0.009999999776482582,
+        -0.03490658476948738,
+        -0.03490658476948738,
+        -0.03490658476948738,
+        0.0
+      ],
+      "q01": [
+        -0.009999999776482582,
+        -0.009999999776482582,
+        -0.009999999776482582,
+        -0.03490658476948738,
+        0.0,
+        -0.03490658476948738,
+        0.0
+      ],
+      "q99": [
+        0.009999999776482582,
+        0.009999999776482582,
+        0.009999999776482582,
+        0.03490658476948738,
+        0.0,
+        0.03490658476948738,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 62613,
+    "num_trajectories": 415
+  },
+  "cmu_stretch": {
+    "action": {
+      "mean": [
+        0.00036304505192674696,
+        0.0,
+        0.0016466958913952112,
+        0.0,
+        0.0,
+        0.0,
+        0.3987048268318176
+      ],
+      "std": [
+        0.004081828519701958,
+        0.0,
+        0.0037743328139185905,
+        0.0,
+        0.0,
+        0.0,
+        0.48963725566864014
+      ],
+      "max": [
+        0.02338407188653946,
+        0.0,
+        0.023404927924275398,
+        0.0,
+        0.0,
+        0.0,
+        1.0
+      ],
+      "min": [
+        -0.019353797659277916,
+        0.0,
+        -0.02019215188920498,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        -0.011175686959177256,
+        0.0,
+        -0.0032206363626755773,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.014501785952597848,
+        0.0,
+        0.015056106168776728,
+        0.0,
+        0.0,
+        0.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 25016,
+    "num_trajectories": 135
+  },
+  "bc_z": {
+    "action": {
+      "mean": [
+        -0.009958467446267605,
+        0.0008958321413956583,
+        0.004995597992092371,
+        0.00029755113064311445,
+        -0.008735382929444313,
+        -0.030693737789988518,
+        0.8344562649726868
+      ],
+      "std": [
+        0.03053455986082554,
+        0.0231423731893301,
+        0.020641816779971123,
+        0.04155943542718887,
+        0.046427831053733826,
+        0.0769818127155304,
+        0.3610210120677948
+      ],
+      "max": [
+        0.2165454924106598,
+        0.1251407265663147,
+        0.10772687941789627,
+        0.33544227480888367,
+        0.28117990493774414,
+        0.40614867210388184,
+        1.0
+      ],
+      "min": [
+        -0.1677047461271286,
+        -0.14630407094955444,
+        -0.10066790133714676,
+        -0.29421567916870117,
+        -0.32101404666900635,
+        -0.4635624885559082,
+        0.0
+      ],
+      "q01": [
+        -0.09220654994249344,
+        -0.06456145539879798,
+        -0.049121275544166565,
+        -0.11594625547528267,
+        -0.14152548640966414,
+        -0.2251061636209488,
+        0.0
+      ],
+      "q99": [
+        0.07628866866230968,
+        0.058019736707210584,
+        0.052540797740221024,
+        0.11740604028105736,
+        0.11703975558280955,
+        0.16729306846857078,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 6015535,
+    "num_trajectories": 43264
+  },
+  "fmb_dataset": {
+    "action": {
+      "mean": [
+        0.059029702097177505,
+        -0.06476633995771408,
+        -0.09787475317716599,
+        0.004325388930737972,
+        0.00028963794466108084,
+        -0.04457257315516472,
+        0.7336440086364746
+      ],
+      "std": [
+        0.28809213638305664,
+        0.2820415794849396,
+        0.4626740515232086,
+        0.3266514539718628,
+        0.10842999070882797,
+        0.3440099358558655,
+        0.4435282051563263
+      ],
+      "max": [
+        1.399999976158142,
+        1.0,
+        1.399999976158142,
+        1.0,
+        1.0,
+        1.0,
+        1.0
+      ],
+      "min": [
+        -1.399999976158142,
+        -1.399999976158142,
+        -1.0,
+        -1.0,
+        -1.0,
+        -1.0,
+        0.0
+      ],
+      "q01": [
+        -0.8257142901420593,
+        -1.399999976158142,
+        -1.0,
+        -1.0,
+        -0.3028571307659149,
+        -1.0,
+        0.0
+      ],
+      "q99": [
+        1.0,
+        0.5257142782211304,
+        1.0,
+        1.0,
+        0.3400000035762787,
+        1.0,
+        1.0
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 1137459,
+    "num_trajectories": 8612
+  },
+  "dobbe": {
+    "action": {
+      "mean": [
+        -0.0001120665911003016,
+        0.0011229600058868527,
+        -0.00010194431524723768,
+        -7.371398532995954e-05,
+        -0.00067531579406932,
+        -5.6643435527803376e-05,
+        0.6318281888961792
+      ],
+      "std": [
+        0.04264938458800316,
+        0.04428559169173241,
+        0.12224084138870239,
+        0.005388413090258837,
+        0.011246449314057827,
+        0.006287882570177317,
+        0.39732322096824646
+      ],
+      "max": [
+        38.590423583984375,
+        17.932697296142578,
+        4.843764305114746,
+        1.4372116327285767,
+        0.4340403974056244,
+        1.2057193517684937,
+        0.9998947381973267
+      ],
+      "min": [
+        -5.700923442840576,
+        -21.605947494506836,
+        -123.72489929199219,
+        -1.7229845523834229,
+        -0.4998578727245331,
+        -0.8867913484573364,
+        1.4196479014572105e-06
+      ],
+      "q01": [
+        -0.01119564864784479,
+        -0.014266146533191203,
+        -0.0071747214533388615,
+        -0.009444301575422287,
+        -0.03990109823644161,
+        -0.017422311007976532,
+        4.003279136668425e-05
+      ],
+      "q99": [
+        0.01015154086053368,
+        0.017181577533483497,
+        0.007216989761218411,
+        0.010380979906767595,
+        0.03556173853576176,
+        0.018032474815845446,
+        0.9982578039169312
+      ],
+      "mask": [
+        true,
+        true,
+        true,
+        true,
+        true,
+        true,
+        false
+      ]
+    },
+    "proprio": {
+      "mean": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "std": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "max": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "min": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q01": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ],
+      "q99": [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.0
+      ]
+    },
+    "num_transitions": 1139911,
+    "num_trajectories": 5208
+  }
+}

openvla-7b-prismatic/run-metrics.jsonl

@@ -0,0 +1 @@
+{"hparams": {"data_root_dir": "/scr/user/data", "hf_token": ".hf_token", "pretrained_checkpoint": "", "resume_epoch": null, "resume_step": null, "run_id": "prism-dinosiglip-224px+mx-oxe-magic-soup-plus+n8+b32+x7", "run_id_note": null, "run_root_dir": "./runs", "save_interval": 2500, "seed": 7, "stage": "vla-full-train", "trackers": ["jsonl", "wandb"], "vla": {"base_vlm": "prism-dinosiglip-224px+7b", "data_mix": "oxe_magic_soup_plus_minus", "enable_gradient_checkpointing": true, "enable_mixed_precision_training": true, "epochs": 1000, "expected_world_size": 64, "freeze_vision_backbone": false, "global_batch_size": 2048, "learning_rate": 2e-05, "lr_scheduler_type": "constant", "max_grad_norm": 1.0, "max_steps": null, "per_device_batch_size": 32, "reduce_in_full_precision": true, "shuffle_buffer_size": 256000, "train_strategy": "fsdp-full-shard", "type": "prism-dinosiglip-224px+mx-oxe-magic-soup-plus", "vla_id": "prism-dinosiglip-224px+mx-oxe-magic-soup-plus", "warmup_ratio": 0.0, "weight_decay": 0.0}, "wandb_entity": "", "wandb_project": ""}, "run_id": "prism-dinosiglip-224px+mx-oxe-magic-soup-plus+n8+b32+x7"}

openvla-7b/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

openvla-7b/README.md

@@ -0,0 +1,99 @@
+---
+library_name: transformers
+tags:
+- robotics
+- vla
+- image-text-to-text
+- multimodal
+- pretraining
+license: mit
+language:
+- en
+pipeline_tag: image-text-to-text
+---
+
+# OpenVLA 7B
+
+OpenVLA 7B (`openvla-7b`) is an open vision-language-action model trained on 970K robot manipulation episodes from the [Open X-Embodiment](https://robotics-transformer-x.github.io/) dataset. 
+The model takes language instructions and camera images as input and generates robot actions. It supports controlling multiple robots out-of-the-box, and can be quickly adapted for new robot domains via (parameter-efficient) fine-tuning.
+
+All OpenVLA checkpoints, as well as our [training codebase](https://github.com/openvla/openvla) are released under an MIT License.
+
+For full details, please read [our paper](https://arxiv.org/abs/2406.09246) and see [our project page](https://openvla.github.io/).
+
+## Model Summary
+
+- **Developed by:** The OpenVLA team consisting of researchers from Stanford, UC Berkeley, Google Deepmind, and the Toyota Research Institute.
+- **Model type:** Vision-language-action (language, image => robot actions)
+- **Language(s) (NLP):** en
+- **License:** MIT
+- **Finetuned from:** [`prism-dinosiglip-224px`](https://github.com/TRI-ML/prismatic-vlms), a VLM trained from:
+  + **Vision Backbone**: DINOv2 ViT-L/14 and SigLIP ViT-So400M/14
+  + **Language Model**: Llama-2
+- **Pretraining Dataset:** [Open X-Embodiment](https://robotics-transformer-x.github.io/) -- specific component datasets can be found [here](https://github.com/openvla/openvla).
+- **Repository:** [https://github.com/openvla/openvla](https://github.com/openvla/openvla)
+- **Paper:** [OpenVLA: An Open-Source Vision-Language-Action Model](https://arxiv.org/abs/2406.09246)
+- **Project Page & Videos:** [https://openvla.github.io/](https://openvla.github.io/)
+
+## Uses
+
+OpenVLA models take a language instruction and a camera image of a robot workspace as input, and predict (normalized) robot actions consisting of 7-DoF end-effector deltas
+of the form (x, y, z, roll, pitch, yaw, gripper). To execute on an actual robot platform, actions need to be *un-normalized* subject to statistics computed on a per-robot,
+per-dataset basis. See [our repository](https://github.com/openvla/openvla) for more information.
+
+OpenVLA models can be used zero-shot to control robots for specific combinations of embodiments and domains seen in the Open-X pretraining mixture (e.g., for 
+[BridgeV2 environments with a Widow-X robot](https://rail-berkeley.github.io/bridgedata/)). They can also be efficiently *fine-tuned* for new tasks and robot setups
+given minimal demonstration data; [see here](https://github.com/openvla/openvla/blob/main/scripts/finetune.py).
+
+**Out-of-Scope:** OpenVLA models do not zero-shot generalize to new (unseen) robot embodiments, or setups that are not represented in the pretraining mix; in these cases,
+we suggest collecting a dataset of demonstrations on the desired setup, and fine-tuning OpenVLA models instead.
+
+## Getting Started
+
+OpenVLA 7B can be used to control multiple robots for domains represented in the pretraining mixture out-of-the-box. For example,
+here is an example for loading `openvla-7b` for zero-shot instruction following in the [BridgeV2 environments] with a Widow-X robot:
+
+```python
+# Install minimal dependencies (`torch`, `transformers`, `timm`, `tokenizers`, ...)
+# > pip install -r https://raw.githubusercontent.com/openvla/openvla/main/requirements-min.txt
+from transformers import AutoModelForVision2Seq, AutoProcessor
+from PIL import Image
+
+import torch
+
+# Load Processor & VLA
+processor = AutoProcessor.from_pretrained("openvla/openvla-7b", trust_remote_code=True)
+vla = AutoModelForVision2Seq.from_pretrained(
+    "openvla/openvla-7b",
+    attn_implementation="flash_attention_2",  # [Optional] Requires `flash_attn`
+    torch_dtype=torch.bfloat16, 
+    low_cpu_mem_usage=True, 
+    trust_remote_code=True
+).to("cuda:0")
+
+# Grab image input & format prompt
+image: Image.Image = get_from_camera(...)
+prompt = "In: What action should the robot take to {<INSTRUCTION>}?\nOut:"
+
+# Predict Action (7-DoF; un-normalize for BridgeV2)
+inputs = processor(prompt, image).to("cuda:0", dtype=torch.bfloat16)
+action = vla.predict_action(**inputs, unnorm_key="bridge_orig", do_sample=False)
+
+# Execute...
+robot.act(action, ...)
+```
+
+For more examples, including scripts for fine-tuning OpenVLA models on your own robot demonstration datasets, see [our training repository](https://github.com/openvla/openvla).
+
+## Citation
+
+**BibTeX:**
+
+```bibtex
+@article{kim24openvla,
+    title={OpenVLA: An Open-Source Vision-Language-Action Model},
+    author={{Moo Jin} Kim and Karl Pertsch and Siddharth Karamcheti and Ted Xiao and Ashwin Balakrishna and Suraj Nair and Rafael Rafailov and Ethan Foster and Grace Lam and Pannag Sanketi and Quan Vuong and Thomas Kollar and Benjamin Burchfiel and Russ Tedrake and Dorsa Sadigh and Sergey Levine and Percy Liang and Chelsea Finn},
+    journal = {arXiv preprint arXiv:2406.09246},
+    year={2024}
+} 
+```

openvla-7b/added_tokens.json

@@ -0,0 +1,3 @@
+{
+  "<PAD>": 32000
+}

openvla-7b/configuration_prismatic.py

@@ -0,0 +1,140 @@
+"""
+configuration_prismatic.py
+
+HuggingFace-style configuration definition for Prismatic VLMs, inheriting from `transformers.PretrainedConfig`.
+Default configuration specifies `siglip-224px+7b`.
+"""
+
+from typing import Any, Dict, List, Optional
+
+from transformers import PretrainedConfig
+from transformers.models.auto import CONFIG_MAPPING
+
+# === Utilities for Mapping Prismatic names to HF names ===
+# fmt: off
+VISION_BACKBONE_TO_RESOLUTION: Dict[str, List[int]] = {
+    "clip-vit-l": [224], "siglip-vit-so400m": [224], "dinov2-vit-l": [224], "in1k-vit-l": [224],
+
+    "clip-vit-l-336px": [336],
+    "siglip-vit-so400m-384px": [384],
+
+    "dinoclip-vit-l-336px": [336, 336],
+    "dinosiglip-vit-so-224px": [224, 224],
+    "dinosiglip-vit-so-384px": [384, 384],
+}
+VISION_BACKBONE_TO_TIMM_ID: Dict[str, List[str]] = {
+    "clip-vit-l": ["vit_large_patch14_clip_224.openai"],
+    "clip-vit-l-336px": ["vit_large_patch14_clip_336.openai"],
+
+    "dinov2-vit-l": ["vit_large_patch14_reg4_dinov2.lvd142m"],
+    "in1k-vit-l": ["vit_large_patch16_224.augreg_in21k_ft_in1k"],
+
+    "siglip-vit-so400m": ["vit_so400m_patch14_siglip_224"],
+    "siglip-vit-so400m-384px": ["vit_so400m_patch14_siglip_384"],
+
+    "dinoclip-vit-l-336px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_large_patch14_clip_336.openai"],
+    "dinosiglip-vit-so-224px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_224"],
+    "dinosiglip-vit-so-384px": ["vit_large_patch14_reg4_dinov2.lvd142m", "vit_so400m_patch14_siglip_384"],
+}
+TIMM_OVERRIDE_ACT_LAYER: Dict[str, List[Optional[str]]] = {
+    "clip-vit-l": ["quick_gelu"], "clip-vit-l-336px": ["quick_gelu"],
+    "dinov2-vit-l": [None], "in1k-vit-l": [None],
+    "siglip-vit-so400m": [None], "siglip-vit-so400m-384px": [None],
+    "dinoclip-vit-l-336px": [None, "quick_gelu"],
+    "dinosiglip-vit-so-224px": [None, None], "dinosiglip-vit-so-384px": [None, None]
+}
+
+LLM_BACKBONE_TO_HF_PATH = {
+    "llama2-7b-pure": "meta-llama/Llama-2-7b-hf", "llama2-13b-pure": "meta-llama/Llama-2-13b-hf",
+    "llama2-7b-chat": "meta-llama/Llama-2-7b-chat-hf", "llama2-13b-chat": "meta-llama/Llama-2-13b-chat-hf",
+
+    "vicuna-v15-7b": "lmsys/vicuna-7b-v1.5", "vicuna-v15-13b": "lmsys/vicuna-13b-v1.5",
+
+    "mistral-v0.1-7b-pure": "mistralai/Mistral-7B-v0.1",
+    "mistral-v0.1-7b-instruct": "mistralai/Mistral-7B-Instruct-v0.1",
+
+    "phi-2-3b": "microsoft/phi-2",
+}
+LLM_BACKBONE_TO_HF_METACLASS = {
+    "llama2-7b-pure": "llama", "llama2-13b-pure": "llama", "llama2-7b-chat": "llama", "llama2-13b-chat": "llama",
+    "vicuna-v15-7b": "llama", "vicuna-v15-13b": "llama",
+
+    "mistral-v0.1-7b-pure": "mistral", "mistral-v0.1-7b-instruct": "mistral",
+
+    "phi-2-3b": "phi",
+}
+
+VALID_VISION_BACKBONES = set(VISION_BACKBONE_TO_RESOLUTION.keys())
+VALID_LLM_BACKBONES = set(LLM_BACKBONE_TO_HF_PATH)
+# fmt: on
+
+
+class PrismaticConfig(PretrainedConfig):
+    model_type: str = "prismatic"
+    is_composition: bool = False
+
+    def __init__(
+        self,
+        vision_backbone_id: str = "siglip-vit-so400m",
+        llm_backbone_id: str = "vicuna-v15-7b",
+        arch_specifier: str = "no-align+gelu-mlp",
+        use_fused_vision_backbone: Optional[bool] = None,
+        image_resize_strategy: str = "letterbox",
+        text_config: Optional[Dict[str, Any]] = None,
+        llm_max_length: int = 2048,
+        pad_token_id: int = 32000,
+        pad_to_multiple_of: int = 64,
+        output_projector_states: bool = False,
+        **kwargs: str,
+    ) -> None:
+        if vision_backbone_id not in VALID_VISION_BACKBONES:
+            raise ValueError(f"Vision backbone `{vision_backbone_id}` not in {VALID_VISION_BACKBONES = }")
+
+        if llm_backbone_id not in VALID_LLM_BACKBONES:
+            raise ValueError(f"LLM backbone `{llm_backbone_id}` not in {VALID_LLM_BACKBONES = }")
+
+        # Set Prismatic Configuration Fields
+        self.vision_backbone_id = vision_backbone_id
+        self.llm_backbone_id = llm_backbone_id
+        self.arch_specifier = arch_specifier
+        self.output_projector_states = output_projector_states
+
+        # [Contract] All vision backbone parameters are lists =>> supports fused backbones with different preprocessing
+        self.use_fused_vision_backbone = (
+            use_fused_vision_backbone
+            if use_fused_vision_backbone is not None
+            else any(self.vision_backbone_id.startswith(v) for v in ["dinoclip", "dinosiglip"])
+        )
+
+        self.timm_model_ids = VISION_BACKBONE_TO_TIMM_ID[self.vision_backbone_id]
+        self.timm_override_act_layers = TIMM_OVERRIDE_ACT_LAYER[self.vision_backbone_id]
+        self.image_sizes = VISION_BACKBONE_TO_RESOLUTION[self.vision_backbone_id]
+        self.image_resize_strategy = image_resize_strategy
+
+        self.hf_llm_id = LLM_BACKBONE_TO_HF_PATH[self.llm_backbone_id]
+        self.llm_max_length = llm_max_length
+        self.pad_token_id, self.pad_to_multiple_of = pad_token_id, pad_to_multiple_of
+
+        # [IMPORTANT] HF Utilities actually look for a `text_config` field... we need to use that specific naming!
+        self.text_config = (
+            CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]](**text_config)
+            if text_config is not None
+            else CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]]()
+        )
+
+        # Dispatch **kwargs to super() =>> note that `pad_token_id` collides, so we pass it in here as well...
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+
+class OpenVLAConfig(PrismaticConfig):
+    model_type: str = "openvla"
+
+    def __init__(
+        self,
+        norm_stats: Optional[Dict[str, Dict[str, Dict[str, Dict[str, List[float]]]]]] = None,
+        n_action_bins: int = 256,
+        **kwargs: str,
+    ) -> None:
+        self.norm_stats, self.n_action_bins = norm_stats, n_action_bins
+
+        super().__init__(**kwargs)

openvla-7b/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 32000,
+  "transformers_version": "4.40.1"
+}

openvla-7b/model-00001-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10d8636256018712c5e5c823d12e22b5797f99bb721bd123bf6bf2379892be85
+size 6948961960

openvla-7b/model-00002-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2050b14f21d48904d269f48d5a980fecea87cd7b36641d9b0f015e72d1fe216a
+size 6971232040

openvla-7b/model-00003-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea65305a1577f36f721965bf84c8caec0a948ce7ce84d754701637376c531fef
+size 1162406824

openvla-7b/model.safetensors.index.json

@@ -0,0 +1,989 @@
+{
+  "metadata": {
+    "total_size": 15082474368
+  },
+  "weight_map": {
+    "language_model.lm_head.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.embed_tokens.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.0.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.1.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.10.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.11.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.11.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.11.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.11.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.11.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.11.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.11.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.11.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.11.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.12.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.12.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.12.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.12.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.12.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.12.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.13.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.13.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.14.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.15.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.16.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.17.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.18.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.19.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.2.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.2.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.20.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.20.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.21.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.22.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.23.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.23.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.23.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.23.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.23.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.23.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.23.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.23.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.23.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.24.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.24.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.24.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.24.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.24.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.24.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.24.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.24.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.24.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.25.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.25.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.25.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.25.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.25.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.25.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.25.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.25.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.25.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.26.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.26.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.26.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.26.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.26.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.26.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.26.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.26.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.26.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.27.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.27.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.27.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.27.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.27.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.27.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.27.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.27.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.27.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.28.input_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.28.mlp.down_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.28.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.28.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.28.post_attention_layernorm.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.28.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.28.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.28.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.28.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.29.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.mlp.gate_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.29.mlp.up_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.29.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.29.self_attn.k_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.29.self_attn.o_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.29.self_attn.q_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.29.self_attn.v_proj.weight": "model-00002-of-00003.safetensors",
+    "language_model.model.layers.3.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.3.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.30.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.30.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.input_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.mlp.down_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.mlp.gate_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.mlp.up_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.post_attention_layernorm.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.self_attn.k_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.self_attn.o_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.self_attn.q_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.31.self_attn.v_proj.weight": "model-00003-of-00003.safetensors",
+    "language_model.model.layers.4.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.4.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.5.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.6.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.7.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.8.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.input_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.mlp.down_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.mlp.gate_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.mlp.up_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.post_attention_layernorm.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.self_attn.k_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.self_attn.o_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.self_attn.q_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.layers.9.self_attn.v_proj.weight": "model-00001-of-00003.safetensors",
+    "language_model.model.norm.weight": "model-00003-of-00003.safetensors",
+    "projector.fc1.bias": "model-00001-of-00003.safetensors",
+    "projector.fc1.weight": "model-00001-of-00003.safetensors",
+    "projector.fc2.bias": "model-00001-of-00003.safetensors",
+    "projector.fc2.weight": "model-00001-of-00003.safetensors",
+    "projector.fc3.bias": "model-00001-of-00003.safetensors",
+    "projector.fc3.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.0.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.1.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.10.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.11.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.12.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.13.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.14.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.15.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.16.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.17.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.18.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.19.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.2.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.20.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.21.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.22.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.23.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.3.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.4.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.5.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.6.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.7.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.8.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.ls1.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.ls2.scale_factor": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.blocks.9.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.cls_token": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.norm.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.norm.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.patch_embed.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.patch_embed.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.pos_embed": "model-00001-of-00003.safetensors",
+    "vision_backbone.featurizer.reg_token": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.kv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.kv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.latent": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.norm.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.norm.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.q.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.attn_pool.q.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.0.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.1.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.10.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.11.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.12.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.13.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.14.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.15.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.16.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.17.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.18.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.19.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.2.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.20.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.21.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.22.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.23.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.24.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.25.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.26.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.3.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.4.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.5.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.6.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.7.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.8.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.attn.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.attn.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.attn.qkv.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.attn.qkv.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.mlp.fc1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.mlp.fc1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.mlp.fc2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.mlp.fc2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.norm1.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.norm1.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.norm2.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.blocks.9.norm2.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.norm.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.norm.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.patch_embed.proj.bias": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.patch_embed.proj.weight": "model-00001-of-00003.safetensors",
+    "vision_backbone.fused_featurizer.pos_embed": "model-00001-of-00003.safetensors"
+  }
+}

openvla-7b/modeling_prismatic.py

@@ -0,0 +1,1086 @@
+"""
+modeling_prismatic.py
+
+Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions.
+Inherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained,
+but exactly replicate the logic in `prismatic.models.vlms.prismatic.py`.
+"""
+
+import logging
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import timm
+import tokenizers
+import torch
+import torch.nn as nn
+import transformers
+from timm.models.vision_transformer import LayerScale
+from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from prismatic.training.train_utils import (
+    get_current_action_mask,
+    get_next_actions_mask,
+)
+from prismatic.vla.constants import (
+    ACTION_DIM,
+    ACTION_PROPRIO_NORMALIZATION_TYPE,
+    ACTION_TOKEN_BEGIN_IDX,
+    IGNORE_INDEX,
+    NUM_ACTIONS_CHUNK,
+    STOP_INDEX,
+    NormalizationType,
+)
+
+from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
+
+# Set up logger
+logger = logging.getLogger(__name__)
+
+
+# === Utility Functions for Monkey-Patching ===
+def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        result = fn(*args, **kwargs)
+        return result[0] if isinstance(result, tuple) else result
+
+    return wrapper
+
+
+# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
+#   =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
+#   =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
+def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
+    return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
+
+
+def ls_apply_patch(ls_module: LayerScale):
+    ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
+    ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
+    del ls_module.gamma
+
+
+# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
+class PrismaticVisionBackbone(nn.Module):
+    """
+    Vision backbone for Prismatic models that handles image feature extraction.
+
+    Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
+    For fused backbones, features from both models are concatenated along the feature dimension.
+    """
+
+    def __init__(
+        self,
+        use_fused_vision_backbone: bool,
+        image_sizes: List[int],
+        timm_model_ids: List[str],
+        timm_override_act_layers: List[Optional[str]],
+    ) -> None:
+        """
+        Initialize the vision backbone.
+
+        Args:
+            use_fused_vision_backbone: Whether to use two backbones and fuse their features
+            image_sizes: List of image sizes for each backbone
+            timm_model_ids: List of TIMM model IDs to use for each backbone
+            timm_override_act_layers: List of activation layer overrides for each backbone
+        """
+        super().__init__()
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.num_images_in_input = 1  # Default value, can be overridden later
+
+        # Validate number of (fused) vision backbones
+        if len(timm_model_ids) > 2:
+            raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
+
+        # Create primary featurizer
+        self.featurizer = self._create_featurizer(
+            model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
+        )
+        self.embed_dim = self.featurizer.embed_dim
+
+        # Create secondary featurizer if using fused backbone
+        if self.use_fused_vision_backbone:
+            self.fused_featurizer = self._create_featurizer(
+                model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
+            )
+            self.embed_dim += self.fused_featurizer.embed_dim
+
+        # Patch LayerScale modules for HF compatibility
+        self._patch_layer_scales()
+
+    def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
+        """
+        Create a TIMM-based featurizer model with appropriate configurations.
+
+        Args:
+            model_id: The TIMM model ID to load
+            img_size: Input image size for the model
+            act_layer: Override for the activation layer type
+
+        Returns:
+            A configured featurizer model
+        """
+        featurizer = timm.create_model(
+            model_id,
+            pretrained=False,
+            num_classes=0,
+            img_size=img_size,
+            act_layer=act_layer,
+        )
+
+        # Monkey-patch the forward function to extract the second-to-last layer features
+        num_blocks = len(featurizer.blocks)
+        featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
+
+        return featurizer
+
+    def _patch_layer_scales(self) -> None:
+        """
+        Patch all LayerScale modules to be compatible with HF's parameter naming.
+
+        HF Transformers overwrites parameters with names containing 'gamma',
+        so we need to rename and modify the forward method.
+        """
+        # Patch primary featurizer
+        for module in self.featurizer.modules():
+            if isinstance(module, LayerScale):
+                ls_apply_patch(module)
+
+        # Patch secondary featurizer if it exists
+        if self.use_fused_vision_backbone:
+            for module in self.fused_featurizer.modules():
+                if isinstance(module, LayerScale):
+                    ls_apply_patch(module)
+
+    def get_num_patches(self) -> int:
+        """
+        Returns the number of vision patches output by the vision backbone.
+
+        Returns:
+            Number of patches per image
+        """
+        return self.featurizer.patch_embed.num_patches
+
+    def get_num_images_in_input(self) -> int:
+        """
+        Returns the number of input images for the vision backbone.
+
+        Returns:
+            Number of images expected in the input
+        """
+        return self.num_images_in_input
+
+    def set_num_images_in_input(self, num_images_in_input: int) -> None:
+        """
+        Sets the number of input images for the vision backbone.
+
+        Args:
+            num_images_in_input: Number of images to expect in the input
+        """
+        self.num_images_in_input = num_images_in_input
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        Implements the forward pass for the vision backbone.
+
+        If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
+        (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
+
+        Args:
+            pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
+        """
+        if self.num_images_in_input == 1:
+            if not self.use_fused_vision_backbone:
+                return self.featurizer(pixel_values)
+
+            # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
+            img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
+            patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
+
+            return torch.cat([patches, patches_fused], dim=2)
+
+        else:
+            assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
+
+            # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
+            images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
+
+            # Process each image and collect patches
+            all_patches = []
+            for img in images:
+                # Split each image further into two stacks of channels (each with 3 channels)
+                img_regular, img_fused = torch.split(img, [3, 3], dim=1)
+
+                # Get patches from both SigLIP and DINOv2 vision transformers
+                patches = self.featurizer(img_regular)
+                patches_fused = self.fused_featurizer(img_fused)
+
+                # Concatenate SigLIP and DINOv2 patches along the hidden dimension
+                combined_patches = torch.cat([patches, patches_fused], dim=2)
+                all_patches.append(combined_patches)
+
+            # Concatenate all patches along the patch dimension
+            return torch.cat(all_patches, dim=1)
+
+
+# === Prismatic Projector (nn.Module) Definitions ===
+class PrismaticProjector(nn.Module):
+    def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
+        super().__init__()
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.vision_dim, self.llm_dim = vision_dim, llm_dim
+
+        # Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
+        if not self.use_fused_vision_backbone:
+            self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
+            self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+            self.act_fn1 = nn.GELU()
+        else:
+            initial_projection_dim = 4 * vision_dim
+            self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
+            self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
+            self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+            self.act_fn1 = nn.GELU()
+            self.act_fn2 = nn.GELU()
+
+    def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
+        if not self.use_fused_vision_backbone:
+            projected_features = self.fc1(img_patches)
+            projected_features = self.act_fn1(projected_features)
+            projected_features = self.fc2(projected_features)
+        else:
+            projected_features = self.fc1(img_patches)
+            projected_features = self.act_fn1(projected_features)
+            projected_features = self.fc2(projected_features)
+            projected_features = self.act_fn2(projected_features)
+            projected_features = self.fc3(projected_features)
+
+        return projected_features
+
+
+# === Main HF Class Definitions ===
+@dataclass
+class PrismaticCausalLMOutputWithPast(ModelOutput):
+    """Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+    # Additions for VLMs
+    projector_features: Optional[torch.FloatTensor] = None
+
+
+class PrismaticPreTrainedModel(PreTrainedModel):
+    config_class: PretrainedConfig = PrismaticConfig
+    base_model_prefix: str = "model"
+    supports_gradient_checkpointing: bool = True
+
+    _no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
+    _skip_keys_device_placement: str = "past_key_values"
+    _supports_flash_attn_2: bool = True
+
+    def _init_weights(self, module: nn.Module) -> None:
+        # Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
+        #   => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
+        #      https://github.com/TRI-ML/prismatic-vlms
+        std = (
+            self.config.initializer_range
+            if hasattr(self.config, "initializer_range")
+            else self.config.text_config.initializer_range
+        )
+
+        if hasattr(module, "class_embedding"):
+            module.class_embedding.data.normal_(mean=0.0, std=std)
+
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    @property
+    def _supports_sdpa(self) -> bool:
+        """Check LLM supports SDPA Attention"""
+        return self.language_model._supports_sdpa
+
+
+class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
+    def __init__(self, config: PrismaticConfig) -> None:
+        super().__init__(config)
+
+        # [Validation] Lightweight Validate on `config` Fields + Dependency Versions
+        if config.use_fused_vision_backbone is None:
+            raise ValueError("Missing config field `use_fused_vision_backbone`")
+
+        if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
+            raise NotImplementedError(
+                "TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
+                "if you urgently need support for latest TIMM versions."
+            )
+
+        if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
+            logger.warning(
+                f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
+                f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
+                f"there might be inference-time regressions due to dependency changes. If in doubt, please"
+                f"use the above versions."
+            )
+
+        # Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
+        self.vision_backbone = PrismaticVisionBackbone(
+            config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
+        )
+
+        # Create Multimodal Projector
+        self.projector = PrismaticProjector(
+            config.use_fused_vision_backbone,
+            vision_dim=self.vision_backbone.embed_dim,
+            llm_dim=config.text_config.hidden_size,
+        )
+
+        # Instantiate LLM Backbone
+        self.language_model = AutoModelForCausalLM.from_config(
+            config.text_config, attn_implementation=config._attn_implementation
+        )
+        self.vocab_size = config.text_config.vocab_size
+        self.pad_token_id = config.pad_token_id
+        self.llm_dim = config.text_config.hidden_size
+
+        # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
+        self.post_init()
+
+    # === `PreTrainedModel` Boilerplate ===
+    def get_input_embeddings(self) -> nn.Module:
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def get_decoder(self) -> nn.Module:
+        return self.language_model.get_decoder()
+
+    def set_decoder(self, decoder: nn.Module) -> None:
+        self.language_model.set_decoder(decoder)
+
+    def tie_weights(self) -> None:
+        self.language_model.tie_weights()  # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
+
+    def resize_token_embeddings(
+        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
+    ) -> nn.Embedding:
+        updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+
+        # Update config/instance variables
+        self.config.text_config.vocab_size = updated_embeddings.num_embeddings
+        self.vocab_size = updated_embeddings.num_embeddings
+
+        return updated_embeddings
+
+    def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
+        """
+        Replace embeddings in input_embeddings at positions where all_actions_mask is True
+        with embeddings from noisy_action_features, using vectorized operations.
+
+        Args:
+            input_embeddings: Tensor of shape (B, S, D)
+            all_actions_mask: Boolean tensor of shape (B, S)
+            noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
+
+        Returns:
+            Modified input_embeddings tensor
+        """
+        # Clone input to avoid modifying the original tensor
+        new_input_embeddings = input_embeddings.clone()
+
+        # Create a tensor with the same shape of input_embeddings to hold the noisy action features
+        repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
+
+        # Create batch indices for splicing
+        batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
+        batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
+
+        # Get indices where mask is True for each sample
+        masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
+
+        # Move the noisy action features into their correct positions
+        repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
+
+        # Combine original input embeddings and noisy action embeddings using the mask
+        new_input_embeddings = torch.where(
+            all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
+        )
+
+        return new_input_embeddings
+
+    def _process_action_masks(self, labels):
+        """Helper to get action masks from labels"""
+        current_action_mask = get_current_action_mask(labels)
+        next_actions_mask = get_next_actions_mask(labels)
+        all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)
+        return all_actions_mask
+
+    def _process_vision_features(self, pixel_values, language_embeddings=None, use_film=False):
+        """Process vision features with optional FiLM conditioning"""
+        if use_film:
+            # FiLM: Infuse language inputs into visual features
+            patch_features = self.vision_backbone(pixel_values, language_embeddings)  # (bsz, 256 * num_images, D)
+        else:
+            patch_features = self.vision_backbone(pixel_values)  # (bsz, 256 * num_images, D)
+
+        # Project patch embeddings into language embedding space
+        return self.projector(patch_features)
+
+    def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
+        """Process proprioceptive features and append to vision features"""
+        if proprio_projector is not None and proprio is not None:
+            # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
+            # proprio: (bsz, proprio_dim) or (propro_dim,)
+            proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1)  # (bsz, proprio_dim)
+            proprio_features = proprio_projector(proprio)  # (bsz, llm_dim)
+            proprio_features = proprio_features.unsqueeze(dim=1)  # (bsz, 1, llm_dim)
+            # For simplicity, just append proprio token to the end of projected vision patch tokens
+            return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
+        return projected_patch_embeddings
+
+    def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
+        """Build multimodal embeddings and attention mask"""
+        # Update attention mask
+        projected_patch_attention_mask = None
+        if attention_mask is not None:
+            projected_patch_attention_mask = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=True,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+
+        # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
+        multimodal_embeddings = torch.cat(
+            [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
+        )
+
+        multimodal_attention_mask = None
+        if attention_mask is not None:
+            multimodal_attention_mask = torch.cat(
+                [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
+            )
+
+        return multimodal_embeddings, multimodal_attention_mask
+
+    def _build_multimodal_labels(self, labels, projected_patch_embeddings):
+        """Build multimodal labels with IGNORE_INDEX for patch embeddings"""
+        if labels is not None:
+            projected_patch_labels = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=IGNORE_INDEX,
+                dtype=labels.dtype,
+                device=labels.device,
+            )
+            return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1)
+        return None
+
+    # === Core Prismatic VLM `forward()` Logic ===
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_projector_features: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        proprio=None,
+        proprio_projector=None,
+        noisy_actions=None,
+        noisy_action_projector=None,
+        diffusion_timestep_embeddings=None,
+        use_film: bool = False,
+    ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
+        """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        output_projector_features = output_projector_features if output_projector_features is not None else False
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)
+        use_cache = use_cache and not self.training
+
+        # Instantiate Placeholder for Projector Features
+        projected_patch_embeddings = None
+
+        # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===
+        if input_ids.shape[1] == 1:
+            assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!"
+            assert past_key_values is not None, "You must provide `past_key_values` during cached generation!"
+            assert labels is None, "Unexpected key `labels` provided during cached generation!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=past_key_values,
+                inputs_embeds=None,
+                labels=None,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Unimodal Forward ===
+        elif pixel_values is None:
+            assert (input_ids is not None) and (inputs_embeds is None), "Missing `input_ids` in language-only forward!"
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!"
+
+            language_model_output = self.language_model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=None,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
+        # === Handle Multimodal Forward ===
+        elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):
+            assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!"
+
+            # Get input embeddings (from language model embeddings)
+            input_embeddings = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)
+
+            # Extract action masks
+            all_actions_mask = self._process_action_masks(labels)
+
+            # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
+            language_embeddings = input_embeddings[~all_actions_mask].reshape(
+                input_embeddings.shape[0], -1, input_embeddings.shape[2]
+            )  # (B, lang_seq_len, llm_dim)
+
+            # Get visual features
+            projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+
+            # Add proprioceptive state if provided
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+            # [Diffusion] Add diffusion timestep embedding if provided
+            if diffusion_timestep_embeddings is not None:
+                # For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens
+                projected_patch_embeddings = torch.cat(
+                    (projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
+                )
+
+            # Process action embeddings
+            if noisy_actions is not None:
+                # Get mask corresponding to all action tokens
+                all_actions_mask = self._process_action_masks(labels)
+
+                # Reshape noisy actions into individual action tokens
+                # noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)
+                B = noisy_actions.shape[0]
+                noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)
+
+                # Project noisy action tokens into language model embedding space
+                noisy_action_features = noisy_action_projector(noisy_actions)  # (B, chunk_len * action_dim, llm_dim)
+
+                # Replace embeddings of the action tokens with noisy action embeddings
+                input_embeddings = self._replace_input_embeddings(
+                    input_embeddings, all_actions_mask, noisy_action_features
+                )
+            else:
+                # Replace the embeddings of the action tokens with zeros
+                # (Later on, the positional embeddings will be added to them)
+                all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+                input_embeddings = input_embeddings * ~all_actions_mask
+
+            # Build multimodal embeddings & attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Build labels for multimodal sequence if needed
+            multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
+
+            # Dispatch to language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=multimodal_labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+            import pdb; pdb.set_trace()
+
+        # === Otherwise =>> Assume Invalid! ===
+        elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):
+            raise ValueError("Non-homogenous batch of (text, image) input -- forward() does not support mixed batches!")
+
+        else:
+            raise ValueError(
+                "Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\n"
+                f"=> `input_ids` = {input_ids is not None}\n"
+                f"=> `attention_mask` = {attention_mask is not None}\n"
+                f"=> `pixel_values` = {pixel_values is not None}\n"
+                f"=> `labels` = {labels is not None}\n"
+                f"=> `input_embeds` = {inputs_embeds is not None}\n"
+                f"=> `past_key_values` = {past_key_values is not None}\n"
+                f"=> `use_cache` = {use_cache}"
+            )
+
+        # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)
+        if not return_dict:
+            if output_projector_features and (projected_patch_embeddings is not None):
+                return *language_model_output, projected_patch_embeddings
+
+            return language_model_output
+
+        return PrismaticCausalLMOutputWithPast(
+            loss=language_model_output.loss,
+            logits=language_model_output.logits,
+            past_key_values=language_model_output.past_key_values,
+            hidden_states=language_model_output.hidden_states,
+            attentions=language_model_output.attentions,
+            projector_features=projected_patch_embeddings,
+        )
+
+    # === GenerationMixin Methods ===
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs: str,
+    ) -> Dict[str, torch.Tensor]:
+        """Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
+        if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
+            (inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
+        ):
+            raise ValueError("Generation with batch size > 1 is not currently supported!")
+
+        # Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        # If `input_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"input_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        # Make sure `pixel_values` are preserved in `model_inputs`
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "pixel_values": pixel_values,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+            }
+        )
+
+        return model_inputs
+
+    # Defer to Language Model (all handle this differently, with different return types)
+    def _reorder_cache(self, *args, **kwargs) -> Any:
+        return self.language_model._reorder_cache(*args, **kwargs)
+
+
+class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
+    config_class: PretrainedConfig = OpenVLAConfig
+
+    def __init__(self, config: OpenVLAConfig) -> None:
+        super().__init__(config)
+        self.norm_stats = config.norm_stats
+
+        # Compute action bins
+        self.bins = np.linspace(-1, 1, config.n_action_bins)
+        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
+
+        # Compute vocab size for de-tokenization -- revert added "multiple of"
+        self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
+
+    def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
+        """Prepares input for action prediction by adding necessary tokens"""
+        # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
+        placeholder_action_token_ids = (
+            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+        )
+        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)
+
+        # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
+        stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX
+        input_ids = torch.cat([input_ids, stop_token_id], dim=-1)
+
+        # Extend the attention mask to fit the new shape of input
+        # Note: Only batch size == 1 supported right now
+        mask_extension = (
+            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+            .to(attention_mask.device)
+            .to(attention_mask.dtype)
+        )
+        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+        return input_ids, attention_mask
+
+    def _prepare_labels_for_action_prediction(self, labels, input_ids):
+        """Creates labels tensor for action prediction if not provided"""
+        # Extend labels tensor with fake action labels
+        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1
+        labels_extension = (
+            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+            * ARBITRARY_ACTION_TOKEN_IDX
+        )
+        labels = torch.cat([labels, labels_extension], dim=-1)
+
+        # Replace last label token with stop token
+        labels[:, -1] = STOP_INDEX
+
+        return labels
+
+    def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
+        """Unnormalize actions using dataset statistics"""
+        action_norm_stats = self.get_action_stats(unnorm_key)
+
+        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
+        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+        else:
+            raise ValueError("Unsupported action/proprio normalization type detected!")
+
+        actions = np.where(
+            mask,
+            0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
+            normalized_actions,
+        )
+
+        return actions
+
+    def _run_diffusion_prediction(
+        self,
+        input_embeddings,
+        all_actions_mask,
+        noise,
+        action_head,
+        projected_patch_embeddings,
+        labels,
+        attention_mask,
+        NUM_PATCHES,
+        NUM_PROMPT_TOKENS,
+        noisy_action_projector,
+    ):
+        """Run diffusion-based action prediction"""
+        # Clone embedding for reuse in each timestep
+        orig_projected_patch_embeddings = projected_patch_embeddings.clone()
+        curr_noisy_actions = noise
+
+        # Reverse diffusion: Iteratively denoise to generate action prediction
+        for t in action_head.noise_scheduler.timesteps:
+            # Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action
+            # embedding, and diffusion timestep embedding)
+            timesteps = torch.Tensor([t]).to(labels.device)
+            diffusion_timestep_embeddings = (
+                action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)
+            )  # (B, llm_dim)
+            diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1)  # (B, 1, llm_dim)
+
+            # [Diffusion] Replace the embeddings of the action tokens with noisy actions
+            # (Later on, the positional embeddings will be added to them)
+
+            # For simplicity, append diffusion timestep embedding to the end of projected vision tokens
+            projected_patch_embeddings = torch.cat(
+                (orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1
+            )
+
+            # Reshape and project noisy actions into language embedding space
+            B = curr_noisy_actions.shape[0]
+            orig_curr_noisy_actions_shape = curr_noisy_actions.shape
+            curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)
+            noisy_action_features = noisy_action_projector(curr_noisy_actions)
+            curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)
+
+            # Replace action token embeddings with noisy action embeddings
+            input_embeddings = self._replace_input_embeddings(
+                input_embeddings.clone(), all_actions_mask, noisy_action_features
+            )
+
+            # Build multimodal embeddings and attention mask
+            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+                input_embeddings, projected_patch_embeddings, attention_mask
+            )
+
+            # Forward pass through language model
+            language_model_output = self.language_model(
+                input_ids=None,
+                attention_mask=multimodal_attention_mask,
+                position_ids=None,
+                past_key_values=None,
+                inputs_embeds=multimodal_embeddings,
+                labels=None,
+                use_cache=None,
+                output_attentions=False,
+                output_hidden_states=True,
+                return_dict=True,
+            )
+
+            # Extract hidden states for action portion of response
+            last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)
+            actions_hidden_states = last_hidden_states[
+                :,
+                NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+                :,
+            ]  # (B, act_chunk_len, D)
+
+            # Predict noise and update noisy actions: x_t -> x_{t-1}
+            noise_pred = action_head.predict_noise(actions_hidden_states)
+            curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample
+
+        curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+
+        # Return final actions
+        return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states
+
+    def _regression_or_discrete_prediction(
+        self,
+        input_embeddings,
+        all_actions_mask,
+        projected_patch_embeddings,
+        attention_mask,
+        labels,
+        NUM_PATCHES,
+        NUM_PROMPT_TOKENS,
+        action_head=None,
+    ):
+        """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
+        # Zero out action token embeddings
+        all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+        input_embeddings = input_embeddings * ~all_actions_mask
+
+        # Build multimodal embeddings and attention mask
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        )
+
+        # Forward pass through language model
+        language_model_output = self.language_model(
+            input_ids=None,
+            attention_mask=multimodal_attention_mask,
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=multimodal_embeddings,
+            labels=None,
+            use_cache=None,
+            output_attentions=False,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+
+        # Extract hidden states for action tokens
+        last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)
+        actions_hidden_states = last_hidden_states[
+            :,
+            NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+            :,
+        ]  # (B, act_chunk_len, D)
+
+        # Handle different prediction methods
+        if action_head is not None:
+            # L1 regression prediction
+            normalized_actions = action_head.predict_action(actions_hidden_states)
+            normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+            normalized_actions = normalized_actions.float().cpu().detach().numpy()
+        else:
+            # Discrete token-based prediction
+            predicted_action_token_ids = (
+                language_model_output.logits[
+                    :,
+                    NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+                ]
+                .argmax(dim=2)
+                .cpu()
+                .numpy()
+            )
+            discretized_actions = self.vocab_size - predicted_action_token_ids
+            discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
+            normalized_actions = self.bin_centers[discretized_actions]
+            normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+
+        return normalized_actions, actions_hidden_states
+
+    def predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """Predict actions from input sequence, with options for different prediction methods.
+
+        Args:
+            input_ids: Input token ids
+            unnorm_key: Key for unnormalization statistics
+            proprio: Proprioceptive features
+            proprio_projector: Projector for proprioceptive features
+            action_head: Optional head for L1 regression or diffusion-based prediction
+            noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+            use_film: Whether to use FiLM conditioning
+            **kwargs: Additional arguments including pixel_values and attention_mask
+
+        Returns:
+            Tuple of (unnormalized_actions, action_hidden_states)
+        """
+        # If the special empty token ('') does not already appear after the colon (':') token in the prompt
+        # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        # Create fake labels tensor (needed for action mask)
+        labels = input_ids.clone()
+        labels[:] = IGNORE_INDEX
+
+        # Get number of tokens in prompt (excluding the start token)
+        NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token
+
+        # Prepare inputs by adding necessary tokens
+        input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
+
+        # Update labels tensor for action mask computation later
+        labels = self._prepare_labels_for_action_prediction(labels, input_ids)
+
+        # Get input embeddings and action masks
+        input_embeddings = self.get_input_embeddings()(input_ids)
+        all_actions_mask = self._process_action_masks(labels)
+
+        # Extract language embeddings
+        language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            input_embeddings.shape[0], -1, input_embeddings.shape[2]
+        )
+
+        # Process vision features
+        projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+
+        # Add proprioceptive features if provided
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+        # Use diffusion if provided, otherwise use regression or discrete prediction
+        use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
+
+        # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+        if use_proprio:
+            NUM_PATCHES += 1
+        if use_diffusion:
+            NUM_PATCHES += 1
+
+        if use_diffusion:
+            # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion
+            noise = torch.randn(
+                size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype
+            )
+
+            # Run diffusion-based prediction
+            normalized_actions, actions_hidden_states = self._run_diffusion_prediction(
+                input_embeddings,
+                all_actions_mask,
+                noise,
+                action_head,
+                projected_patch_embeddings,
+                labels,
+                attention_mask,
+                NUM_PATCHES,
+                NUM_PROMPT_TOKENS,
+                noisy_action_projector,
+            )
+        else:
+            # Run regression or discrete token-based prediction
+            normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
+                input_embeddings,
+                all_actions_mask,
+                projected_patch_embeddings,
+                attention_mask,
+                labels,
+                NUM_PATCHES,
+                NUM_PROMPT_TOKENS,
+                action_head,
+            )
+
+        # Unnormalize predicted actions
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+
+        return actions, actions_hidden_states
+
+    @staticmethod
+    def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
+        """Validate and resolve the unnormalization key for action statistics"""
+        if unnorm_key is None:
+            assert len(norm_stats) == 1, (
+                f"Your model was trained on more than one dataset, "
+                f"please pass a `unnorm_key` from the following options to choose the statistics "
+                f"used for un-normalizing actions: {norm_stats.keys()}"
+            )
+            unnorm_key = next(iter(norm_stats.keys()))
+
+        assert unnorm_key in norm_stats, (
+            f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
+            f"please choose from: {norm_stats.keys()}"
+        )
+        return unnorm_key
+
+    def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
+        """Get the dimensionality of the policy's action space."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+        return len(self.norm_stats[unnorm_key]["action"]["min"])
+
+    def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
+        """Get all the logged statistics for the given dataset."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+        return self.norm_stats[unnorm_key]["action"]

openvla-7b/modeling_prismatic_rl.py

@@ -0,0 +1,1344 @@
+"""
+modeling_prismatic.py
+Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions.
+Inherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained,
+but exactly replicate the logic in `prismatic.models.vlms.prismatic.py`.
+"""
+
+import logging
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Callable, ClassVar, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import timm
+import tokenizers
+import torch
+import torch.nn as nn
+import transformers
+from timm.models.vision_transformer import LayerScale
+from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from prismatic.training.train_utils import (
+    get_current_action_mask,
+    get_next_actions_mask,
+)
+from prismatic.vla.constants import (
+    ACTION_DIM,
+    ACTION_PROPRIO_NORMALIZATION_TYPE,
+    ACTION_TOKEN_BEGIN_IDX,
+    IGNORE_INDEX,
+    NUM_ACTIONS_CHUNK,
+    STOP_INDEX,
+    NormalizationType,
+)
+
+from .configuration_prismatic import OpenVLAConfig, PrismaticConfig
+
+TRAINING_MODE = "SFT"
+
+# Set up logger
+logger = logging.getLogger(__name__)
+
+
+# === Utility Functions for Monkey-Patching ===
+def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]:
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        result = fn(*args, **kwargs)
+        return result[0] if isinstance(result, tuple) else result
+
+    return wrapper
+
+
+# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.
+#   =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109
+#   =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960
+def _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:
+    return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor
+
+
+def ls_apply_patch(ls_module: LayerScale):
+    ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())
+    ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)
+    del ls_module.gamma
+
+
+# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===
+class PrismaticVisionBackbone(nn.Module):
+    """
+    Vision backbone for Prismatic models that handles image feature extraction.
+    Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.
+    For fused backbones, features from both models are concatenated along the feature dimension.
+    """
+
+    def __init__(
+        self,
+        use_fused_vision_backbone: bool,
+        image_sizes: List[int],
+        timm_model_ids: List[str],
+        timm_override_act_layers: List[Optional[str]],
+    ) -> None:
+        """
+        Initialize the vision backbone.
+        Args:
+            use_fused_vision_backbone: Whether to use two backbones and fuse their features
+            image_sizes: List of image sizes for each backbone
+            timm_model_ids: List of TIMM model IDs to use for each backbone
+            timm_override_act_layers: List of activation layer overrides for each backbone
+        """
+        super().__init__()
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.num_images_in_input = 1  # Default value, can be overridden later
+
+        # Validate number of (fused) vision backbones
+        if len(timm_model_ids) > 2:
+            raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!")
+
+        # Create primary featurizer
+        self.featurizer = self._create_featurizer(
+            model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]
+        )
+        self.embed_dim = self.featurizer.embed_dim
+
+        # Create secondary featurizer if using fused backbone
+        if self.use_fused_vision_backbone:
+            self.fused_featurizer = self._create_featurizer(
+                model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]
+            )
+            self.embed_dim += self.fused_featurizer.embed_dim
+
+        # Patch LayerScale modules for HF compatibility
+        self._patch_layer_scales()
+
+    def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:
+        """
+        Create a TIMM-based featurizer model with appropriate configurations.
+        Args:
+            model_id: The TIMM model ID to load
+            img_size: Input image size for the model
+            act_layer: Override for the activation layer type
+        Returns:
+            A configured featurizer model
+        """
+        featurizer = timm.create_model(
+            model_id,
+            pretrained=False,
+            num_classes=0,
+            img_size=img_size,
+            act_layer=act_layer,
+        )
+
+        # Monkey-patch the forward function to extract the second-to-last layer features
+        num_blocks = len(featurizer.blocks)
+        featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))
+
+        return featurizer
+
+    def _patch_layer_scales(self) -> None:
+        """
+        Patch all LayerScale modules to be compatible with HF's parameter naming.
+        HF Transformers overwrites parameters with names containing 'gamma',
+        so we need to rename and modify the forward method.
+        """
+        # Patch primary featurizer
+        for module in self.featurizer.modules():
+            if isinstance(module, LayerScale):
+                ls_apply_patch(module)
+
+        # Patch secondary featurizer if it exists
+        if self.use_fused_vision_backbone:
+            for module in self.fused_featurizer.modules():
+                if isinstance(module, LayerScale):
+                    ls_apply_patch(module)
+
+    def get_num_patches(self) -> int:
+        """
+        Returns the number of vision patches output by the vision backbone.
+        Returns:
+            Number of patches per image
+        """
+        return self.featurizer.patch_embed.num_patches
+
+    def get_num_images_in_input(self) -> int:
+        """
+        Returns the number of input images for the vision backbone.
+        Returns:
+            Number of images expected in the input
+        """
+        return self.num_images_in_input
+
+    def set_num_images_in_input(self, num_images_in_input: int) -> None:
+        """
+        Sets the number of input images for the vision backbone.
+        Args:
+            num_images_in_input: Number of images to expect in the input
+        """
+        self.num_images_in_input = num_images_in_input
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        Implements the forward pass for the vision backbone.
+        If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features
+        (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).
+        Args:
+            pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).
+        """
+        if self.num_images_in_input == 1:
+            if not self.use_fused_vision_backbone:
+                return self.featurizer(pixel_values)
+
+            # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack
+            img, img_fused = torch.split(pixel_values, [3, 3], dim=1)
+            patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)
+
+            return torch.cat([patches, patches_fused], dim=2)
+
+        else:
+            assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!"
+
+            # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)
+            images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)
+
+            # Process each image and collect patches
+            all_patches = []
+            for img in images:
+                # Split each image further into two stacks of channels (each with 3 channels)
+                img_regular, img_fused = torch.split(img, [3, 3], dim=1)
+
+                # Get patches from both SigLIP and DINOv2 vision transformers
+                patches = self.featurizer(img_regular)
+                patches_fused = self.fused_featurizer(img_fused)
+
+                # Concatenate SigLIP and DINOv2 patches along the hidden dimension
+                combined_patches = torch.cat([patches, patches_fused], dim=2)
+                all_patches.append(combined_patches)
+
+            # Concatenate all patches along the patch dimension
+            return torch.cat(all_patches, dim=1)
+
+
+# === Prismatic Projector (nn.Module) Definitions ===
+class PrismaticProjector(nn.Module):
+    def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:
+        super().__init__()
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.vision_dim, self.llm_dim = vision_dim, llm_dim
+
+        # Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!
+        if not self.use_fused_vision_backbone:
+            self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)
+            self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+            self.act_fn1 = nn.GELU()
+        else:
+            initial_projection_dim = 4 * vision_dim
+            self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)
+            self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)
+            self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)
+            self.act_fn1 = nn.GELU()
+            self.act_fn2 = nn.GELU()
+
+    def forward(self, img_patches: torch.Tensor) -> torch.Tensor:
+        if not self.use_fused_vision_backbone:
+            projected_features = self.fc1(img_patches)
+            projected_features = self.act_fn1(projected_features)
+            projected_features = self.fc2(projected_features)
+        else:
+            projected_features = self.fc1(img_patches)
+            projected_features = self.act_fn1(projected_features)
+            projected_features = self.fc2(projected_features)
+            projected_features = self.act_fn2(projected_features)
+            projected_features = self.fc3(projected_features)
+
+        return projected_features
+
+
+# === Main HF Class Definitions ===
+@dataclass
+class PrismaticCausalLMOutputWithPast(ModelOutput):
+    """Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features."""
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+    # Additions for VLMs
+    projector_features: Optional[torch.FloatTensor] = None
+
+
+class PrismaticPreTrainedModel(PreTrainedModel):
+    config_class: PretrainedConfig = PrismaticConfig
+    base_model_prefix: str = "model"
+    supports_gradient_checkpointing: bool = True
+
+    _no_split_modules: ClassVar[List[str]] = ["PrismaticProjector"]
+    _skip_keys_device_placement: str = "past_key_values"
+    _supports_flash_attn_2: bool = True
+
+    def _init_weights(self, module: nn.Module) -> None:
+        # Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!
+        #   => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at
+        #      https://github.com/TRI-ML/prismatic-vlms
+        std = (
+            self.config.initializer_range
+            if hasattr(self.config, "initializer_range")
+            else self.config.text_config.initializer_range
+        )
+
+        if hasattr(module, "class_embedding"):
+            module.class_embedding.data.normal_(mean=0.0, std=std)
+
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    @property
+    def _supports_sdpa(self) -> bool:
+        """Check LLM supports SDPA Attention"""
+        return self.language_model._supports_sdpa
+
+
+class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
+    def __init__(self, config: PrismaticConfig) -> None:
+        super().__init__(config)
+
+        # [Validation] Lightweight Validate on `config` Fields + Dependency Versions
+        if config.use_fused_vision_backbone is None:
+            raise ValueError("Missing config field `use_fused_vision_backbone`")
+
+        if timm.__version__ not in {"0.9.10", "0.9.11", "0.9.12", "0.9.16"}:
+            raise NotImplementedError(
+                "TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue "
+                "if you urgently need support for latest TIMM versions."
+            )
+
+        if (transformers.__version__ != "4.40.1") or (tokenizers.__version__ != "0.19.1"):
+            logger.warning(
+                f"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got "
+                f"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; "
+                f"there might be inference-time regressions due to dependency changes. If in doubt, please"
+                f"use the above versions."
+            )
+
+        # Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)
+        self.vision_backbone = PrismaticVisionBackbone(
+            config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers
+        )
+
+        # Create Multimodal Projector
+        self.projector = PrismaticProjector(
+            config.use_fused_vision_backbone,
+            vision_dim=self.vision_backbone.embed_dim,
+            llm_dim=config.text_config.hidden_size,
+        )
+
+        # Instantiate LLM Backbone
+        self.language_model = AutoModelForCausalLM.from_config(
+            config.text_config, attn_implementation=config._attn_implementation
+        )
+        self.vocab_size = config.text_config.vocab_size
+        self.pad_token_id = config.pad_token_id
+        self.llm_dim = config.text_config.hidden_size
+
+        # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing
+        self.post_init()
+
+    # === `PreTrainedModel` Boilerplate ===
+    def get_input_embeddings(self) -> nn.Module:
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.language_model.set_input_embeddings(value)
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.language_model.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def get_decoder(self) -> nn.Module:
+        return self.language_model.get_decoder()
+
+    def set_decoder(self, decoder: nn.Module) -> None:
+        self.language_model.set_decoder(decoder)
+
+    def tie_weights(self) -> None:
+        self.language_model.tie_weights()  # Note: `Llama-2` and `Mistral` don't tie weights (no-op)
+
+    def resize_token_embeddings(
+        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
+    ) -> nn.Embedding:
+        updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
+
+        # Update config/instance variables
+        self.config.text_config.vocab_size = updated_embeddings.num_embeddings
+        self.vocab_size = updated_embeddings.num_embeddings
+
+        return updated_embeddings
+
+    def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):
+        """
+        Replace embeddings in input_embeddings at positions where all_actions_mask is True
+        with embeddings from noisy_action_features, using vectorized operations.
+        Args:
+            input_embeddings: Tensor of shape (B, S, D)
+            all_actions_mask: Boolean tensor of shape (B, S)
+            noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample
+        Returns:
+            Modified input_embeddings tensor
+        """
+        # Clone input to avoid modifying the original tensor
+        new_input_embeddings = input_embeddings.clone()
+
+        # Create a tensor with the same shape of input_embeddings to hold the noisy action features
+        repositioned_noisy_action_features = torch.zeros_like(input_embeddings)
+
+        # Create batch indices for splicing
+        batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)
+        batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])
+
+        # Get indices where mask is True for each sample
+        masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])
+
+        # Move the noisy action features into their correct positions
+        repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features
+
+        # Combine original input embeddings and noisy action embeddings using the mask
+        new_input_embeddings = torch.where(
+            all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings
+        )
+
+        return new_input_embeddings
+
+    def _process_action_masks(self, labels):
+        """Helper to get action masks from labels"""
+        current_action_mask = get_current_action_mask(labels)
+        next_actions_mask = get_next_actions_mask(labels)
+        all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)
+        return all_actions_mask
+
+    def _process_vision_features(self, pixel_values, language_embeddings=None, use_film=False):
+        """Process vision features with optional FiLM conditioning"""
+        if use_film:
+            # FiLM: Infuse language inputs into visual features
+            patch_features = self.vision_backbone(pixel_values, language_embeddings)  # (bsz, 256 * num_images, D)
+        else:
+            patch_features = self.vision_backbone(pixel_values)  # (bsz, 256 * num_images, D)
+
+        # Project patch embeddings into language embedding space
+        return self.projector(patch_features)
+
+    def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):
+        """Process proprioceptive features and append to vision features"""
+        if proprio_projector is not None and proprio is not None:
+            # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)
+            # proprio: (bsz, proprio_dim) or (propro_dim,)
+            proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1)  # (bsz, proprio_dim)
+            proprio_features = proprio_projector(proprio)  # (bsz, llm_dim)
+            proprio_features = proprio_features.unsqueeze(dim=1)  # (bsz, 1, llm_dim)
+            # For simplicity, just append proprio token to the end of projected vision patch tokens
+            return torch.cat((projected_patch_embeddings, proprio_features), dim=1)
+        return projected_patch_embeddings
+
+    def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):
+        """Build multimodal embeddings and attention mask"""
+        # Update attention mask
+        projected_patch_attention_mask = None
+        if attention_mask is not None:
+            projected_patch_attention_mask = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=True,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+
+        # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
+        multimodal_embeddings = torch.cat(
+            [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
+        )
+
+        multimodal_attention_mask = None
+        if attention_mask is not None:
+            multimodal_attention_mask = torch.cat(
+                [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
+            )
+
+        return multimodal_embeddings, multimodal_attention_mask
+
+    def _build_multimodal_labels(self, labels, projected_patch_embeddings):
+        """Build multimodal labels with IGNORE_INDEX for patch embeddings"""
+        if labels is not None:
+            projected_patch_labels = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=IGNORE_INDEX,
+                dtype=labels.dtype,
+                device=labels.device,
+            )
+            return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1)
+        return None
+
+    # === GenerationMixin Methods ===
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs: str,
+    ) -> Dict[str, torch.Tensor]:
+        """Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic."""
+        if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (
+            (inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)
+        ):
+            raise ValueError("Generation with batch size > 1 is not currently supported!")
+
+        # Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+
+        # If `input_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"input_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        # Make sure `pixel_values` are preserved in `model_inputs`
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "pixel_values": pixel_values,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+            }
+        )
+        return model_inputs
+
+    # Defer to Language Model (all handle this differently, with different return types)
+    def _reorder_cache(self, *args, **kwargs) -> Any:
+        return self.language_model._reorder_cache(*args, **kwargs)
+    
+    def _prepare_input_for_action_prediction_verl(self, input_ids, attention_mask):
+        """Prepares input for action prediction by adding necessary tokens"""
+        # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
+        placeholder_action_token_ids = (
+            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+        )
+        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)
+
+        # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
+        stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX
+        input_ids = torch.cat([input_ids, stop_token_id], dim=-1)
+
+        # Extend the attention mask to fit the new shape of input
+        # Note: Only batch size == 1 supported right now
+        mask_extension = (
+            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+            .to(attention_mask.device)
+            .to(attention_mask.dtype)
+        )
+        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+        return input_ids, attention_mask
+
+    def _prepare_labels_for_action_prediction_verl(self, labels, input_ids):
+        """Creates labels tensor for action prediction if not provided"""
+        # Extend labels tensor with fake action labels
+        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1
+        labels_extension = (
+            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+            * ARBITRARY_ACTION_TOKEN_IDX
+        )
+        labels = torch.cat([labels, labels_extension], dim=-1)
+
+        # Replace last label token with stop token
+        labels[:, -1] = STOP_INDEX
+
+        return labels
+    
+    def _verl_discrete_compute_logits(
+        self,
+        input_embeddings,
+        all_actions_mask,
+        projected_patch_embeddings,
+        attention_mask,
+        labels,
+        NUM_PATCHES,
+        NUM_PROMPT_TOKENS,
+        action_head=None,
+    ):#contintue!!!!!
+        """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
+        # Zero out action token embeddings
+        all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+        input_embeddings = input_embeddings * ~all_actions_mask
+
+        # Build multimodal embeddings and attention mask
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        )
+
+        # Forward pass through language model
+        language_model_output = self.language_model(
+            input_ids=None,
+            attention_mask=multimodal_attention_mask,
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=multimodal_embeddings,
+            labels=None,
+            use_cache=None,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+        )
+      
+        compute_logits = language_model_output.logits[
+                    :,
+                    NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+                ]
+            
+        return  compute_logits
+    
+    def _forward_sft(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        proprio=None,
+        proprio_projector=None,
+    ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
+        """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance."""
+        # Instantiate Placeholder for Projector Features
+        projected_patch_embeddings = None
+
+        # === Handle Multimodal Forward ===
+        # Get input embeddings (from language model embeddings)
+        input_embeddings = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)
+
+        # Extract action masks
+        all_actions_mask = self._process_action_masks(labels)
+
+        # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)
+        language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            input_embeddings.shape[0], -1, input_embeddings.shape[2]
+        )  # (B, lang_seq_len, llm_dim)
+
+        # Get visual features
+        projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings)
+
+        # Add proprioceptive state if provided
+        projected_patch_embeddings = self._process_proprio_features(
+            projected_patch_embeddings, proprio, proprio_projector
+        )
+
+        # Replace the embeddings of the action tokens with zeros
+        # (Later on, the positional embeddings will be added to them)
+        all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+        input_embeddings = input_embeddings * ~all_actions_mask
+
+        # Build multimodal embeddings & attention mask
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        )
+
+        # Build labels for multimodal sequence if needed
+        multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)
+
+        # Dispatch to language model
+        language_model_output = self.language_model(
+            input_ids=None,
+            attention_mask=multimodal_attention_mask,
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=multimodal_embeddings,
+            labels=multimodal_labels,
+            use_cache=None,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+        )
+
+        return PrismaticCausalLMOutputWithPast(
+            loss=language_model_output.loss,
+            logits=language_model_output.logits,
+            past_key_values=language_model_output.past_key_values,
+            hidden_states=language_model_output.hidden_states,
+            attentions=language_model_output.attentions,
+            projector_features=projected_patch_embeddings,
+        )
+
+    def _forward_rl(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        pixel_values=None,
+        attention_mask=None,
+        labels=None,
+        proprio=None,
+        proprio_projector=None,
+        **kwargs: str,
+    ) :
+        """Predict actions from input sequence, with options for different prediction methods.
+        Args:
+            input_ids: Input token ids
+            unnorm_key: Key for unnormalization statistics
+            proprio: Proprioceptive features
+            proprio_projector: Projector for proprioceptive features
+            action_head: Optional head for L1 regression or diffusion-based prediction
+            noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+            use_film: Whether to use FiLM conditioning
+            **kwargs: Additional arguments including pixel_values and attention_mask
+        Returns:
+            Tuple of (unnormalized_actions, action_hidden_states)
+        """
+        # Create fake labels tensor (needed for action mask)
+        labels = input_ids.clone()
+        labels[:] = IGNORE_INDEX
+
+        # # Get number of tokens in prompt (excluding the start token)
+        NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token
+
+
+        # # Prepare inputs by adding necessary tokens
+        # #input_ids, attention_mask = self._prepare_input_for_action_prediction_verl(input_ids, attention_mask)
+        
+        # #test
+        placeholder_action_token_ids = (
+            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+        )
+        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)
+
+        # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
+        stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX
+        input_ids = torch.cat([input_ids, stop_token_id], dim=-1)
+
+        # Extend the attention mask to fit the new shape of input
+        # Note: Only batch size == 1 supported right now
+        mask_extension = (
+            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+            .to(attention_mask.device)
+            .to(attention_mask.dtype)
+        )
+        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1
+        labels_extension = (
+            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+            * ARBITRARY_ACTION_TOKEN_IDX
+        )
+        labels = torch.cat([labels, labels_extension], dim=-1)
+
+        # # Replace last label token with stop token
+        labels[:, -1] = STOP_INDEX
+
+        
+        # Get input embeddings and action masks
+        
+        #NUM_PROMPT_TOKENS = kwargs["num_prompt_tokens"]
+        
+        input_embeddings = self.get_input_embeddings()(input_ids)
+        
+        newline_positions = labels != IGNORE_INDEX
+
+        # Calculate cumulative sum to identify regions between newlines
+        cumsum = torch.cumsum(newline_positions, dim=1)
+
+        # Create the mask
+        mask = (1 <= cumsum) & (cumsum <= ACTION_DIM)
+
+        # Extract the action part only
+        action_tokens_only_mask = labels > ACTION_TOKEN_BEGIN_IDX
+        current_action_mask = action_tokens_only_mask * mask
+
+        #next_actions_mask = get_next_actions_mask(labels)
+        newline_positions = labels != IGNORE_INDEX
+
+        # Calculate cumulative sum to identify regions between newlines
+        cumsum = torch.cumsum(newline_positions, dim=1)
+
+        # Create the mask
+        mask = cumsum > ACTION_DIM
+
+        # Extract the action part only
+        action_tokens_only_mask = labels > ACTION_TOKEN_BEGIN_IDX
+        next_actions_mask = action_tokens_only_mask * mask
+        
+        all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)
+        
+        # Extract language embeddings
+        language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            input_embeddings.shape[0], -1, input_embeddings.shape[2]
+        )
+
+        # Process vision features
+        #projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+        patch_features = self.vision_backbone(pixel_values)  # (bsz, 256 * num_images, D)
+
+        projected_patch_embeddings = self.projector(patch_features)
+        #test end
+        
+        
+        # Add proprioceptive features if provided
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+        # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+        if use_proprio:
+            NUM_PATCHES += 1
+        
+        all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+        input_embeddings = input_embeddings * ~all_actions_mask
+        projected_patch_attention_mask = None
+        if attention_mask is not None:
+            projected_patch_attention_mask = torch.full(
+                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),
+                fill_value=True,
+                dtype=attention_mask.dtype,
+                device=attention_mask.device,
+            )
+
+        # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)
+        multimodal_embeddings = torch.cat(
+            [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1
+        )
+
+        multimodal_attention_mask = None
+        if attention_mask is not None:
+            multimodal_attention_mask = torch.cat(
+                [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1
+            )
+        # Forward pass through language model
+        language_model_output = self.language_model(
+            input_ids=None,
+            attention_mask=multimodal_attention_mask,
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=multimodal_embeddings,
+            labels=None,
+            use_cache=None,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+        )
+
+    
+        compute_logits = language_model_output.logits[
+                    :,
+                    NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+                ]
+                
+            #test end
+
+        return compute_logits
+    
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        pixel_values: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_projector_features: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        proprio=None,
+        proprio_projector=None,
+        noisy_actions=None,
+        noisy_action_projector=None,
+        diffusion_timestep_embeddings=None,
+        use_film: bool = False,
+    ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:
+        if TRAINING_MODE == "SFT":
+            return self._forward_sft(
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                attention_mask=attention_mask,
+                labels=labels,
+                proprio=proprio,
+                proprio_projector=proprio_projector,
+            )
+        elif TRAINING_MODE == "RL":
+            return self._forward_rl(
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                attention_mask=attention_mask,
+                labels=None,
+                proprio=proprio,
+                proprio_projector=proprio_projector,
+            )
+        else:
+            raise ValueError(f"Unsupported training mode: {TRAINING_MODE}. Supported modes are 'SFT' and 'RL'.")
+
+class OpenVLAForActionPredictionForRL(PrismaticForConditionalGeneration):
+    config_class: PretrainedConfig = OpenVLAConfig
+
+    def __init__(self, config: OpenVLAConfig) -> None:
+        super().__init__(config)
+        self.norm_stats = config.norm_stats
+
+        # Compute action bins
+        self.bins = np.linspace(-1, 1, config.n_action_bins)
+        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
+
+        # Compute vocab size for de-tokenization -- revert added "multiple of"
+        self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of
+
+    def _prepare_input_for_action_prediction(self, input_ids, attention_mask):
+        """Prepares input for action prediction by adding necessary tokens"""
+        # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens
+        placeholder_action_token_ids = (
+            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)
+        )
+        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)
+
+        # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)
+        stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX
+        input_ids = torch.cat([input_ids, stop_token_id], dim=-1)
+
+        # Extend the attention mask to fit the new shape of input
+        # Note: Only batch size == 1 supported right now
+        mask_extension = (
+            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))
+            .to(attention_mask.device)
+            .to(attention_mask.dtype)
+        )
+        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)
+
+        return input_ids, attention_mask
+
+    def _prepare_labels_for_action_prediction(self, labels, input_ids):
+        """Creates labels tensor for action prediction if not provided"""
+        # Extend labels tensor with fake action labels
+        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1
+        labels_extension = (
+            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)
+            * ARBITRARY_ACTION_TOKEN_IDX
+        )
+        labels = torch.cat([labels, labels_extension], dim=-1)
+
+        # Replace last label token with stop token
+        labels[:, -1] = STOP_INDEX
+
+        return labels
+
+    def _unnormalize_actions(self, normalized_actions, unnorm_key=None):
+        """Unnormalize actions using dataset statistics"""
+        action_norm_stats = self.get_action_stats(unnorm_key)
+
+        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"])
+        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:
+            mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool))
+            action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"])
+        else:
+            raise ValueError("Unsupported action/proprio normalization type detected!")
+
+        actions = np.where(
+            mask,
+            0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,
+            normalized_actions,
+        )
+
+        return actions
+    def _regression_or_discrete_prediction(
+        self,
+        input_embeddings,
+        all_actions_mask,
+        projected_patch_embeddings,
+        attention_mask,
+        labels,
+        NUM_PATCHES,
+        NUM_PROMPT_TOKENS,
+        action_head=None,
+    ):
+        """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
+        # Zero out action token embeddings
+        all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+        input_embeddings = input_embeddings * ~all_actions_mask
+
+        # Build multimodal embeddings and attention mask
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        )
+
+        # Forward pass through language model
+        language_model_output = self.language_model(
+            input_ids=None,
+            attention_mask=multimodal_attention_mask,
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=multimodal_embeddings,
+            labels=None,
+            use_cache=None,
+            output_attentions=False,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+
+        # Extract hidden states for action tokens
+        last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)
+        actions_hidden_states = last_hidden_states[
+            :,
+            NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+            :,
+        ]  # (B, act_chunk_len, D)
+
+        # Handle different prediction methods
+        if action_head is not None:
+            # L1 regression prediction
+            normalized_actions = action_head.predict_action(actions_hidden_states)
+            normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+            normalized_actions = normalized_actions.float().cpu().detach().numpy()
+        else:
+            # Discrete token-based prediction
+            predicted_action_token_ids = (
+                language_model_output.logits[
+                    :,
+                    NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,
+                ]
+                .argmax(dim=2)
+                .cpu()
+                .numpy()
+            )
+            discretized_actions = self.vocab_size - predicted_action_token_ids
+            discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
+            normalized_actions = self.bin_centers[discretized_actions]
+            normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+
+        return normalized_actions, actions_hidden_states
+    
+    def _verl_discrete_prediction(
+        self,
+        input_embeddings,
+        all_actions_mask,
+        projected_patch_embeddings,
+        attention_mask,
+        labels,
+        NUM_PATCHES,
+        NUM_PROMPT_TOKENS,
+        action_head=None,
+        do_sample=True,
+        temperature=1,
+    ):
+        """Run L1 regression-based continuous action prediction or discrete action tokens prediction."""
+        # Zero out action token embeddings
+        all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)
+        input_embeddings = input_embeddings * ~all_actions_mask
+
+        # Build multimodal embeddings and attention mask
+        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(
+            input_embeddings, projected_patch_embeddings, attention_mask
+        )
+
+        # Forward pass through language model
+        language_model_output = self.language_model(
+            input_ids=None,
+            attention_mask=multimodal_attention_mask,
+            position_ids=None,
+            past_key_values=None,
+            inputs_embeds=multimodal_embeddings,
+            labels=None,
+            use_cache=None,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+        )
+        NUM_PROMPT_TOKENS = NUM_PROMPT_TOKENS + NUM_PATCHES
+        batch_size = language_model_output.logits.shape[0]
+        device = language_model_output.logits.device
+
+        start_indices = NUM_PROMPT_TOKENS.unsqueeze(1)  # [batch_size, 1]
+        position_offsets = torch.arange(ACTION_DIM * NUM_ACTIONS_CHUNK, device=device).unsqueeze(0)  # [1, seq_length]
+        seq_indices = start_indices + position_offsets  # [batch_size, ACTION_DIM*NUM_ACTIONS_CHUNK]
+        if do_sample == False:
+            #padding + only get last 256 token
+            reponse_ids_logits = language_model_output.logits[
+                torch.arange(batch_size, device=device).unsqueeze(-1),  
+                seq_indices, 
+                :
+            ]
+            start_index = self.vocab_size - 256 
+            response_last256 = reponse_ids_logits[..., -256-64:-64]  # Shape: [batch_size, seq_len, 256]
+            last256_argmax = response_last256.argmax(dim=-1)  # Shape: [batch_size, seq_len]
+            reponse_ids = last256_argmax + start_index  # Shape: [batch_size, seq_len]
+            #padding + only get last 256 token end
+            
+            predicted_action_token_ids = reponse_ids.cpu().numpy()
+                
+        else:
+            assert temperature>0
+            
+            action_logits = language_model_output.logits[
+                torch.arange(batch_size, device=device).unsqueeze(-1),  
+                seq_indices, 
+                :
+            ]
+            
+            #padding + only get last 256 token
+            action_logits_last256 = action_logits[..., -256-64:-64]
+            scaled_logits = action_logits_last256 / temperature
+            probs = torch.softmax(scaled_logits, dim=-1)
+            assert probs.shape[-1] == 256
+            probs_flat = probs.reshape(-1, probs.shape[-1])
+            sampled_indices_flat = torch.multinomial(probs_flat, num_samples=1)
+            original_ids_flat = sampled_indices_flat + (self.vocab_size - 256)
+            reponse_ids = original_ids_flat.view(action_logits.shape[0], -1)
+            #padding + only get last 256 token end
+            
+            predicted_action_token_ids = reponse_ids.cpu().numpy()
+     
+        discretized_actions = self.vocab_size - predicted_action_token_ids
+        discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
+        normalized_actions = self.bin_centers[discretized_actions]
+        #normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)
+        normalized_actions = normalized_actions.reshape(-1, ACTION_DIM)
+
+        return normalized_actions, reponse_ids
+        #return normalized_actions, actions_hidden_states
+
+    def predict_action(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """Predict actions from input sequence, with options for different prediction methods.
+        Args:
+            input_ids: Input token ids
+            unnorm_key: Key for unnormalization statistics
+            proprio: Proprioceptive features
+            proprio_projector: Projector for proprioceptive features
+            action_head: Optional head for L1 regression or diffusion-based prediction
+            noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+            use_film: Whether to use FiLM conditioning
+            **kwargs: Additional arguments including pixel_values and attention_mask
+        Returns:
+            Tuple of (unnormalized_actions, action_hidden_states)
+        """
+        # If the special empty token ('') does not already appear after the colon (':') token in the prompt
+        # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time
+        if not torch.all(input_ids[:, -1] == 29871):
+            input_ids = torch.cat(
+                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1
+            )
+
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+
+        # Create fake labels tensor (needed for action mask)
+        labels = input_ids.clone()
+        labels[:] = IGNORE_INDEX
+
+        # Get number of tokens in prompt (excluding the start token)
+        NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token
+
+        # Prepare inputs by adding necessary tokens
+        input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
+
+        # Update labels tensor for action mask computation later
+        labels = self._prepare_labels_for_action_prediction(labels, input_ids)
+
+        # Get input embeddings and action masks
+        input_embeddings = self.get_input_embeddings()(input_ids)
+        all_actions_mask = self._process_action_masks(labels)
+
+        # Extract language embeddings
+        language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            input_embeddings.shape[0], -1, input_embeddings.shape[2]
+        )
+
+        # Process vision features
+        projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+
+        # Add proprioceptive features if provided
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+        # Use diffusion if provided, otherwise use regression or discrete prediction
+        use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
+
+        # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+        if use_proprio:
+            NUM_PATCHES += 1
+        if use_diffusion:
+            raise ValueError
+        else:
+            # Run regression or discrete token-based prediction
+            normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(
+                input_embeddings,
+                all_actions_mask,
+                projected_patch_embeddings,
+                attention_mask,
+                labels,
+                NUM_PATCHES,
+                NUM_PROMPT_TOKENS,
+                action_head,
+            )
+
+        # Unnormalize predicted actions
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+
+        return actions, actions_hidden_states
+
+    def generate_action_verl(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        unnorm_key: Optional[str] = None,
+        proprio=None,
+        proprio_projector=None,
+        action_head=None,
+        noisy_action_projector=None,
+        use_film: bool = False,
+        **kwargs: str,
+    ) -> np.ndarray:
+        """Predict actions from input sequence, with options for different prediction methods.
+        Args:
+            input_ids: Input token ids
+            unnorm_key: Key for unnormalization statistics
+            proprio: Proprioceptive features
+            proprio_projector: Projector for proprioceptive features
+            action_head: Optional head for L1 regression or diffusion-based prediction
+            noisy_action_projector: Projector for noisy actions in diffusion-based prediction
+            use_film: Whether to use FiLM conditioning
+            **kwargs: Additional arguments including pixel_values and attention_mask
+        Returns:
+            Tuple of (unnormalized_actions, action_hidden_states)
+        """
+
+        pixel_values = kwargs["pixel_values"]
+        attention_mask = kwargs["attention_mask"]
+        do_sample = kwargs["do_sample"]
+        temperature = kwargs["temperature"]
+        
+        # Create fake labels tensor (needed for action mask)
+        labels = input_ids.clone()
+        labels[:] = IGNORE_INDEX
+
+        # Get number of tokens in prompt (excluding the start token)
+        #NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token
+        padding_idx = kwargs["padding_idx"]
+        num_prompt_tokens = input_ids.ne(padding_idx).sum(dim=1) - 1
+
+        # Prepare inputs by adding necessary tokens
+        input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)
+
+        # Update labels tensor for action mask computation later
+        labels = self._prepare_labels_for_action_prediction(labels, input_ids)
+        
+        #here to convert padding from before to last
+        padding_mask = input_ids.ne(padding_idx)
+        assert torch.all(padding_mask==attention_mask.ne(0))
+        #print("in predict_action padding_mask:", padding_mask)
+        padding_mask = padding_mask.int() 
+        sorted_indices = torch.argsort(padding_mask, dim=1, descending=True, stable=True)
+        input_ids = torch.gather(input_ids, 1, sorted_indices)
+        attention_mask = torch.gather(attention_mask, 1, sorted_indices)
+        labels = torch.gather(labels, 1, sorted_indices)
+        assert use_film==False        
+
+        # Get input embeddings and action masks
+        input_embeddings = self.get_input_embeddings()(input_ids)
+        all_actions_mask = self._process_action_masks(labels)
+
+        # Extract language embeddings
+        language_embeddings = input_embeddings[~all_actions_mask].reshape(
+            input_embeddings.shape[0], -1, input_embeddings.shape[2]
+        )
+
+        # Process vision features
+        projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)
+
+        # Add proprioceptive features if provided
+        use_proprio = proprio_projector is not None and proprio is not None
+        if use_proprio:
+            proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype)
+            projected_patch_embeddings = self._process_proprio_features(
+                projected_patch_embeddings, proprio, proprio_projector
+            )
+
+        # Use diffusion if provided, otherwise use regression or discrete prediction
+        use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler")
+
+        # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)
+        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()
+        if use_proprio:
+            NUM_PATCHES += 1
+        if use_diffusion:
+            NUM_PATCHES += 1
+
+        if use_diffusion:
+            raise ValueError
+        else:
+            # Run regression or discrete token-based prediction
+            normalized_actions, reponse_ids = self._verl_discrete_prediction(
+                input_embeddings,
+                all_actions_mask,
+                projected_patch_embeddings,
+                attention_mask,
+                labels,
+                NUM_PATCHES,
+                num_prompt_tokens,
+                action_head,
+                do_sample=do_sample,
+                temperature=temperature,
+            )
+
+        # Unnormalize predicted actions
+        actions = self._unnormalize_actions(normalized_actions, unnorm_key)
+        #verl add!
+        actions = actions.reshape(-1 ,NUM_ACTIONS_CHUNK, ACTION_DIM)
+        #
+        return actions, reponse_ids
+
+    @staticmethod
+    def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
+        """Validate and resolve the unnormalization key for action statistics"""
+        if unnorm_key is None:
+            assert len(norm_stats) == 1, (
+                f"Your model was trained on more than one dataset, "
+                f"please pass a `unnorm_key` from the following options to choose the statistics "
+                f"used for un-normalizing actions: {norm_stats.keys()}"
+            )
+            unnorm_key = next(iter(norm_stats.keys()))
+
+        assert unnorm_key in norm_stats, (
+            f"The `unnorm_key` you chose is not in the set of available dataset statistics, "
+            f"please choose from: {norm_stats.keys()}"
+        )
+        return unnorm_key
+
+    def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:
+        """Get the dimensionality of the policy's action space."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+        return len(self.norm_stats[unnorm_key]["action"]["min"])
+
+    def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
+        """Get all the logged statistics for the given dataset."""
+        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
+        return self.norm_stats[unnorm_key]["action"]

openvla-7b/preprocessor_config.json

@@ -0,0 +1,114 @@
+{
+  "auto_map": {
+    "AutoImageProcessor": "processing_prismatic.PrismaticImageProcessor",
+    "AutoProcessor": "processing_prismatic.PrismaticProcessor"
+  },
+  "image_processor_type": "PrismaticImageProcessor",
+  "image_resize_strategy": "resize-naive",
+  "input_sizes": [
+    [
+      3,
+      224,
+      224
+    ],
+    [
+      3,
+      224,
+      224
+    ]
+  ],
+  "interpolations": [
+    "bicubic",
+    "bicubic"
+  ],
+  "means": [
+    [
+      0.485,
+      0.456,
+      0.406
+    ],
+    [
+      0.5,
+      0.5,
+      0.5
+    ]
+  ],
+  "processor_class": "PrismaticProcessor",
+  "stds": [
+    [
+      0.229,
+      0.224,
+      0.225
+    ],
+    [
+      0.5,
+      0.5,
+      0.5
+    ]
+  ],
+  "tvf_crop_params": [
+    {
+      "output_size": [
+        224,
+        224
+      ]
+    },
+    {
+      "output_size": [
+        224,
+        224
+      ]
+    }
+  ],
+  "tvf_do_letterbox": false,
+  "tvf_letterbox_fill": null,
+  "tvf_normalize_params": [
+    {
+      "inplace": false,
+      "mean": [
+        0.484375,
+        0.455078125,
+        0.40625
+      ],
+      "std": [
+        0.228515625,
+        0.2236328125,
+        0.224609375
+      ]
+    },
+    {
+      "inplace": false,
+      "mean": [
+        0.5,
+        0.5,
+        0.5
+      ],
+      "std": [
+        0.5,
+        0.5,
+        0.5
+      ]
+    }
+  ],
+  "tvf_resize_params": [
+    {
+      "antialias": true,
+      "interpolation": 3,
+      "max_size": null,
+      "size": [
+        224,
+        224
+      ]
+    },
+    {
+      "antialias": true,
+      "interpolation": 3,
+      "max_size": null,
+      "size": [
+        224,
+        224
+      ]
+    }
+  ],
+  "use_fused_vision_backbone": true
+}

openvla-7b/processing_prismatic.py

@@ -0,0 +1,257 @@
+"""
+processing_prismatic.py
+
+HuggingFace-style preprocessor definitions for Prismatic VLMs, inheriting from `ProcessorMixin`. Default configuration
+specifies `siglip-224px+7b`.
+"""
+
+from typing import Any, ClassVar, List, Optional, Tuple, Union
+
+import timm.data
+import torch
+import torchvision.transforms.functional as TVF
+from PIL import Image
+from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
+from transformers import PreTrainedTokenizerBase
+from transformers.image_processing_utils import BatchFeature, ImageProcessingMixin
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
+from transformers.utils import TensorType
+
+
+# === Image Processing ===
+def letterbox_pad_transform(image: Image.Image, padding_fill_value: Tuple[int, int, int]) -> Image.Image:
+    """Given a PIL.Image, pad to square by adding a symmetric border around the height/width."""
+    (w, h), max_wh = image.size, max(image.size)
+    horizontal_pad, vertical_pad = int((max_wh - w) / 2), int((max_wh - h) / 2)
+    padding = (horizontal_pad, vertical_pad, horizontal_pad, vertical_pad)
+
+    return TVF.pad(image, padding, fill=padding_fill_value, padding_mode="constant")
+
+
+class PrismaticImageProcessor(ImageProcessingMixin):
+    model_input_names: ClassVar[List[str]] = ["pixel_values"]
+
+    def __init__(
+        self,
+        use_fused_vision_backbone: bool = False,
+        image_resize_strategy: str = "letterbox",
+        input_sizes: Optional[List[Tuple[int, int, int]]] = None,
+        interpolations: Optional[List[str]] = None,
+        means: Optional[List[Tuple[float, float, float]]] = None,
+        stds: Optional[List[Tuple[float, float, float]]] = None,
+        **kwargs: str,
+    ) -> None:
+        """
+        Initialize a PrismaticImageProcessor as a wrapper around a torchvision transform; this transform will be
+        created by TIMM, and edited to follow our custom `image_resize_strategy` logic.
+
+        @param use_fused_vision_backbone: Boolean indicating single or fused (dual) vision backbone
+        @param image_resize_strategy: Prismatic image resize strategy in < resize-naive | resize-crop | letterbox >
+        @param input_size: [TIMM :: `data_cfg`] Input image size as tuple (channels, width, height)
+        @param interpolation: [TIMM :: `data_cfg`] Interpolation as string (default: "bicubic")
+        @param mean: [TIMM :: `data_cfg`] Normalization mean as float tuple (or two-tuple if `fused_backbone`)
+        @param std: [TIMM :: `data_cfg`] Normalization std as float tuple (or two-tuple if `fused_backbone`)
+        """
+        self.use_fused_vision_backbone = use_fused_vision_backbone
+        self.image_resize_strategy = image_resize_strategy
+
+        # Handle `None` default values
+        input_sizes = [(3, 224, 224)] if input_sizes is None else input_sizes
+        means = [(0.5, 0.5, 0.5)] if means is None else means
+        stds = [(0.5, 0.5, 0.5)] if stds is None else stds
+
+        # TIMM `data_cfg` Parameters
+        self.input_sizes, self.interpolations, self.means, self.stds = input_sizes, interpolations, means, stds
+
+        # Grab torchvision transforms via TIMM =>> need to parse for specific "functional" transform values!
+        self.tvf_resize_params, self.tvf_crop_params, self.tvf_normalize_params = [], [], []
+        self.tvf_do_letterbox, self.tvf_letterbox_fill = False, None
+
+        for idx in range(len(input_sizes)):
+            transform = timm.data.create_transform(
+                input_size=self.input_sizes[idx],
+                interpolation=self.interpolations[idx],
+                mean=self.means[idx],
+                std=self.stds[idx],
+                crop_pct=1.0,  # Set to 1.0 to ignore cropping (initial Resize sets `input_size`)
+                crop_mode="center",  # Default crop mode -- no-op when `crop_pct == 1.0`
+                is_training=False,  # No image augmentations when loading the transform!
+            )
+
+            # [Validation] Ensure appropriate transform structure, expected sizes
+            if not (
+                isinstance(transform, Compose)
+                and (len(transform.transforms) == 4)
+                and isinstance(transform.transforms[0], Resize)
+                and isinstance(transform.transforms[1], CenterCrop)
+                and isinstance(transform.transforms[2], ToTensor)
+                and isinstance(transform.transforms[3], Normalize)
+                and (transform.transforms[0].size == self.input_sizes[idx][-1])
+                and (transform.transforms[1].size == self.input_sizes[idx][-2:])
+            ):
+                raise ValueError(f"Unexpected TIMM image transformation structure/sizes: `{transform}`")
+
+            # HF Image Processors *must* be JSON-serializable; as such, cannot have torchvision. as an attribute.
+            #   => Instead, we're going to parse the transform and call "torchvision.transforms.functional" (`tvf`)
+            resize_t, crop_t, norm_t = transform.transforms[0], transform.transforms[1], transform.transforms[3]
+            self.tvf_resize_params.append(
+                {
+                    "size": resize_t.size,
+                    "interpolation": TVF.pil_modes_mapping[resize_t.interpolation],
+                    "max_size": None,
+                    "antialias": True,
+                }
+            )
+            self.tvf_crop_params.append({"output_size": crop_t.size})
+            self.tvf_normalize_params.append(
+                {
+                    "mean": norm_t.mean.float().numpy().tolist(),
+                    "std": norm_t.std.float().numpy().tolist(),
+                    "inplace": False,
+                }
+            )
+            self.tvf_do_letterbox, self.tvf_letterbox_fill = False, None
+
+            # Handle Prismatic `image_resize_strategy`
+            if self.image_resize_strategy == "resize-naive":
+                self.tvf_resize_params[idx]["size"] = (resize_t.size, resize_t.size)
+            elif self.image_resize_strategy == "letterbox":
+                self.tvf_do_letterbox, self.tvf_letterbox_fill = True, tuple([int(x * 255) for x in self.means[idx]])
+            elif self.image_resize_strategy == "resize-crop":
+                pass
+            else:
+                raise ValueError(f"Image resize strategy `{self.image_resize_strategy}` is not supported!")
+
+        # Dispatch **kwargs to super()
+        super().__init__(**kwargs)
+
+    def apply_transform(self, img: Image.Image) -> torch.Tensor:
+        """Apply `functional` variant of TIMM's Transform = Compose([Resize -> CenterCrop -> ToTensor -> Normalize])"""
+        if self.tvf_do_letterbox:
+            img = letterbox_pad_transform(img, self.tvf_letterbox_fill)
+
+        # [Contract] Fused Backbones expect "channel-stacked" inputs; we'll unpack on the model side!
+        imgs_t = []
+        for idx in range(len(self.input_sizes)):
+            img_idx = TVF.resize(img, **self.tvf_resize_params[idx])
+            img_idx = TVF.center_crop(img_idx, **self.tvf_crop_params[idx])
+            img_idx_t = TVF.to_tensor(img_idx)
+            img_idx_t = TVF.normalize(img_idx_t, **self.tvf_normalize_params[idx])
+            imgs_t.append(img_idx_t)
+
+        # [Contract] `imgs_t` is a list of Tensors of shape [3, input_size, input_size]; stack along dim = 0
+        img_t = torch.vstack(imgs_t)
+
+        return img_t
+
+    def preprocess(
+        self,
+        images: Union[Image.Image, List[Image.Image]],
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        **_: str,
+    ) -> BatchFeature:
+        """
+        Preprocess an image (or batch of images); note that unlike the `transformers :: BaseImageProcessor` we
+        explicitly only handle PIL.Image.Image instances for simplicity.
+
+        @param images: A (batch of) PIL.Image.Image instance(s) to preprocess.
+        @param return_tensors: BatchFeature default Tensor format (e.g., "pt" for torch); if None, returns np.ndarray
+
+        @return: Instance of `transformers :: BatchFeature` with a single key "pixel_values"
+        """
+        if not isinstance(images, list):
+            images = [images]
+
+        # Apply `self.img_transform` to each image (will return list of torch.Tensors); stack into "batched" Tensor
+        pixel_values = torch.stack([self.apply_transform(img.convert("RGB")) for img in images])
+
+        # Return BatchFeature =>> note that for compatibility, constructor expects Dict[str, np.ndarray], so we convert
+        return BatchFeature(data={"pixel_values": pixel_values.float().numpy()}, tensor_type=return_tensors)
+
+    def __call__(self, images: Union[Image.Image, List[Image.Image]], **kwargs) -> BatchFeature:
+        return self.preprocess(images, **kwargs)
+
+
+# === PrismaticProcessor =>> Wraps both ImageProcessor and Tokenizer ===
+#   =>> https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava/processing_llava.py
+class PrismaticProcessor(ProcessorMixin):
+    attributes: ClassVar[List[str]] = ["image_processor", "tokenizer"]
+    image_processor_class: str = "AutoImageProcessor"
+    tokenizer_class: str = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor: Optional[ImageProcessingMixin] = None,
+        tokenizer: Optional[PreTrainedTokenizerBase] = None,
+    ) -> None:
+        super().__init__(image_processor, tokenizer)
+
+    def __call__(
+        self,
+        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]],
+        images: Union[Image.Image, List[Image.Image]],
+        padding: Union[bool, str, PaddingStrategy] = False,
+        truncation: Optional[Union[bool, str, TruncationStrategy]] = None,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
+    ) -> BatchFeature:
+        """
+        Preprocess a given (batch) of text/images for a Prismatic VLM; forwards text to the underlying LLM's tokenizer,
+        forwards images to PrismaticImageProcessor.
+
+        @param text: The (batch) of text to encode; must be a string or list of strings.
+        @param images: A (batch of) PIL.Image.Image instance(s) to preprocess.
+        @param padding: Sequence padding strategy (if multiple specified) in < True = "longest" | "max_length" | False >
+        @param truncation: Truncation strategy for the output sequences; requires `max_length` to be specified
+        @param max_length: Maximum length (in tokens) to truncate
+        @param return_tensors: Type of return tensors (usually "pt" or TensorType.PYTORCH)
+
+        @return: BatchFeature with keys for `input_ids`, `attention_mask` and `pixel_values`.
+        """
+        pixel_values = self.image_processor(images, return_tensors=return_tensors)["pixel_values"]
+        text_inputs = self.tokenizer(
+            text, return_tensors=return_tensors, padding=padding, truncation=truncation, max_length=max_length
+        )
+
+        # [Validate] Need same number of images and text inputs!
+        if pixel_values.shape[0] != text_inputs.input_ids.shape[0]:
+            raise ValueError("Batch is malformed; expected same number of images and text inputs!")
+
+        return BatchFeature(data={**text_inputs, "pixel_values": pixel_values})
+
+    # === Tokenizer Dispatch Utilities =>> check `PreTrainedTokenizerBase` for documentation ===
+    def batch_decode(
+        self,
+        sequences: Union[List[int], List[List[int]], torch.Tensor, Any],  # `Any` = np.ndarray | tf.Tensor
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: Optional[bool] = None,
+        **kwargs: str,
+    ) -> List[str]:
+        return self.tokenizer.batch_decode(
+            sequences=sequences,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+    def decode(
+        self,
+        token_ids: Union[int, List[int], torch.Tensor, Any],  # `Any` = np.ndarray | tf.Tensor
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: Optional[bool] = None,
+        **kwargs: str,
+    ) -> str:
+        return self.tokenizer.decode(
+            token_ids=token_ids,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+    @property
+    def model_input_names(self) -> List[str]:
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+
+        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))

openvla-7b/processor_config.json

@@ -0,0 +1,6 @@
+{
+  "auto_map": {
+    "AutoProcessor": "processing_prismatic.PrismaticProcessor"
+  },
+  "processor_class": "PrismaticProcessor"
+}

openvla-7b/special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<PAD>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

openvla-7b/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347
+size 499723

openvla-7b/tokenizer_config.json

@@ -0,0 +1,53 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32000": {
+      "content": "<PAD>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoProcessor": "processing_prismatic.PrismaticProcessor"
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "</s>",
+  "legacy": false,
+  "model_max_length": 2048,
+  "pad_token": "<PAD>",
+  "padding_side": "right",
+  "processor_class": "PrismaticProcessor",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

prismatic/config.json

@@ -0,0 +1,59 @@
+{
+  "dataset": {
+    "align_stage_components": [
+      "download/llava-laion-cc-sbu-558k/chat.json",
+      "download/llava-laion-cc-sbu-558k"
+    ],
+    "dataset_id": "llava-lvis4v-lrv",
+    "dataset_root_dir": "/opt/ml/input/data/training/skaramcheti/datasets/prismatic-vlms",
+    "finetune_stage_components": [
+      "download/llava-v1.5-instruct/llava_v1_5_lvis4v_lrv_mix1231k.json",
+      "download/llava-v1.5-instruct"
+    ],
+    "type": "llava-lvis4v-lrv"
+  },
+  "hf_token": ".hf_token",
+  "model": {
+    "align_epochs": 1,
+    "align_global_batch_size": 256,
+    "align_learning_rate": 0.001,
+    "align_lr_scheduler_type": "linear-warmup+cosine-decay",
+    "align_max_grad_norm": 1.0,
+    "align_max_steps": null,
+    "align_per_device_batch_size": 16,
+    "align_train_strategy": "fsdp-shard-grad-op",
+    "align_warmup_ratio": 0.03,
+    "align_weight_decay": 0.0,
+    "arch_specifier": "no-align+fused-gelu-mlp",
+    "enable_gradient_checkpointing": true,
+    "enable_mixed_precision_training": true,
+    "finetune_epochs": 2,
+    "finetune_global_batch_size": 128,
+    "finetune_learning_rate": 2e-05,
+    "finetune_lr_scheduler_type": "linear-warmup+cosine-decay",
+    "finetune_max_grad_norm": 1.0,
+    "finetune_max_steps": null,
+    "finetune_per_device_batch_size": 16,
+    "finetune_train_strategy": "fsdp-full-shard",
+    "finetune_warmup_ratio": 0.03,
+    "finetune_weight_decay": 0.1,
+    "image_resize_strategy": "resize-naive",
+    "llm_backbone_id": "llama2-7b-pure",
+    "llm_max_length": 2048,
+    "model_id": "prism-dinosiglip-224px+7b",
+    "reduce_in_full_precision": false,
+    "type": "prism-dinosiglip-224px+7b",
+    "vision_backbone_id": "dinosiglip-vit-so-224px"
+  },
+  "pretrained_checkpoint": null,
+  "run_id": "llava-lvis4v-lrv+prism-dinosiglip-224px+7b+stage-finetune+x7",
+  "run_root_dir": "/opt/ml/input/data/training/x-prismatic-vlms/runs",
+  "seed": 7,
+  "stage": "finetune",
+  "trackers": [
+    "jsonl",
+    "wandb"
+  ],
+  "wandb_entity": "stanford-voltron",
+  "wandb_project": "onyx-vlms"
+}

prismatic/config.yaml

@@ -0,0 +1,52 @@
+dataset:
+  align_stage_components:
+  - download/llava-laion-cc-sbu-558k/chat.json
+  - download/llava-laion-cc-sbu-558k
+  dataset_id: llava-lvis4v-lrv
+  dataset_root_dir: /opt/ml/input/data/training/skaramcheti/datasets/prismatic-vlms
+  finetune_stage_components:
+  - download/llava-v1.5-instruct/llava_v1_5_lvis4v_lrv_mix1231k.json
+  - download/llava-v1.5-instruct
+  type: llava-lvis4v-lrv
+hf_token: .hf_token
+model:
+  align_epochs: 1
+  align_global_batch_size: 256
+  align_learning_rate: 0.001
+  align_lr_scheduler_type: linear-warmup+cosine-decay
+  align_max_grad_norm: 1.0
+  align_max_steps: null
+  align_per_device_batch_size: 16
+  align_train_strategy: fsdp-shard-grad-op
+  align_warmup_ratio: 0.03
+  align_weight_decay: 0.0
+  arch_specifier: no-align+fused-gelu-mlp
+  enable_gradient_checkpointing: true
+  enable_mixed_precision_training: true
+  finetune_epochs: 2
+  finetune_global_batch_size: 128
+  finetune_learning_rate: 2.0e-05
+  finetune_lr_scheduler_type: linear-warmup+cosine-decay
+  finetune_max_grad_norm: 1.0
+  finetune_max_steps: null
+  finetune_per_device_batch_size: 16
+  finetune_train_strategy: fsdp-full-shard
+  finetune_warmup_ratio: 0.03
+  finetune_weight_decay: 0.1
+  image_resize_strategy: resize-naive
+  llm_backbone_id: llama2-7b-pure
+  llm_max_length: 2048
+  model_id: prism-dinosiglip-224px+7b
+  reduce_in_full_precision: false
+  type: prism-dinosiglip-224px+7b
+  vision_backbone_id: dinosiglip-vit-so-224px
+pretrained_checkpoint: null
+run_id: llava-lvis4v-lrv+prism-dinosiglip-224px+7b+stage-finetune+x7
+run_root_dir: /opt/ml/input/data/training/x-prismatic-vlms/runs
+seed: 7
+stage: finetune
+trackers:
+- jsonl
+- wandb
+wandb_entity: stanford-voltron
+wandb_project: onyx-vlms

prismatic/latest-checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b0281cb5c9e37d08e1a9bde840522518b02bad916b96058ee9bb99e0956e0159
+size 13620752058

prismatic/run-metrics.jsonl

@@ -0,0 +1 @@
+{"hparams": {"dataset": {"align_stage_components": ["download/llava-laion-cc-sbu-558k/chat.json", "download/llava-laion-cc-sbu-558k"], "dataset_id": "llava-lvis4v-lrv", "dataset_root_dir": "/opt/ml/input/data/training/skaramcheti/datasets/prismatic-vlms", "finetune_stage_components": ["download/llava-v1.5-instruct/llava_v1_5_lvis4v_lrv_mix1231k.json", "download/llava-v1.5-instruct"], "type": "llava-lvis4v-lrv"}, "hf_token": ".hf_token", "model": {"align_epochs": 1, "align_global_batch_size": 256, "align_learning_rate": 0.001, "align_lr_scheduler_type": "linear-warmup+cosine-decay", "align_max_grad_norm": 1.0, "align_max_steps": null, "align_per_device_batch_size": 16, "align_train_strategy": "fsdp-shard-grad-op", "align_warmup_ratio": 0.03, "align_weight_decay": 0.0, "arch_specifier": "no-align+fused-gelu-mlp", "enable_gradient_checkpointing": true, "enable_mixed_precision_training": true, "finetune_epochs": 2, "finetune_global_batch_size": 128, "finetune_learning_rate": 2e-05, "finetune_lr_scheduler_type": "linear-warmup+cosine-decay", "finetune_max_grad_norm": 1.0, "finetune_max_steps": null, "finetune_per_device_batch_size": 16, "finetune_train_strategy": "fsdp-full-shard", "finetune_warmup_ratio": 0.03, "finetune_weight_decay": 0.1, "image_resize_strategy": "resize-naive", "llm_backbone_id": "llama2-7b-pure", "llm_max_length": 2048, "model_id": "prism-dinosiglip-224px+7b", "reduce_in_full_precision": false, "type": "prism-dinosiglip-224px+7b", "vision_backbone_id": "dinosiglip-vit-so-224px"}, "pretrained_checkpoint": null, "run_id": "llava-lvis4v-lrv+prism-dinosiglip-224px+7b+stage-finetune+x7", "run_root_dir": "/opt/ml/input/data/training/x-prismatic-vlms/runs", "seed": 7, "stage": "finetune", "trackers": ["jsonl", "wandb"], "wandb_entity": "stanford-voltron", "wandb_project": "onyx-vlms"}, "run_id": "llava-lvis4v-lrv+prism-dinosiglip-224px+7b+stage-finetune+x7"}