Motif Vision Encoder release

Co-authored-by: gkalstn0 <gkalstn0@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,38 @@
+*.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
+*.gif filter=lfs diff=lfs merge=lfs -text
+assets/architecture.png filter=lfs diff=lfs merge=lfs -text
+assets/dense_attention_comparison.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Motif Technologies
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,139 @@
+---
+license: mit
+library_name: transformers
+pipeline_tag: image-feature-extraction
+tags:
+  - motif
+  - vision-transformer
+  - self-supervised
+  - image-feature-extraction
+  - video
+  - custom_code
+---
+
+# Motif Vision Encoder
+
+Motif Vision Encoder is a unified image + video self-supervised vision encoder on a ViT
+backbone. A single 3D-convolutional tokenizer ingests both modalities — an image is a
+1-frame clip (`T=1`), a video is `T>1` — so the same weights produce dense patch-level
+features and a language-aligned global (CLS) representation.
+
+Trained on **~1/3 the data of DINOv3** (0.5B vs 1.7B samples), it still reaches competitive
+performance across image and video benchmarks — and leads on DAVIS video tracking.
+
+<p align="center">
+  <img src="assets/haaland_attn_blk20.gif" width="480" alt="Point tracking on a video clip: Motif vs V-JEPA 2.1"/>
+</p>
+<p align="center"><em>Point tracking on a video clip (top: Motif, bottom: V-JEPA 2.1) — a query point propagated across frames by patch-feature cosine similarity. Motif tracks the subject more reliably than V-JEPA 2.1.</em></p>
+
+- **Architecture**: ViT-7B (embed 4096 / depth 40 / heads 32), patch 16, 3D axial RoPE
+  (`base=100`), SwiGLU FFN, LayerScale, per-head QK-norm, gated attention, 4 register tokens.
+- **Tokenizer**: `Conv3d(kernel=stride=(tubelet, patch, patch))` — image `(B,3,H,W)` → `T=1`,
+  video `(B,T,3,H,W)`. Token layout `[CLS] + [register × 4] + [patch × N]`.
+
+
+## Usage
+
+The model ships a self-contained `modeling_motif_vision_encoder.py`, so it loads with `trust_remote_code=True`.
+
+### Image
+
+```python
+import torch
+from transformers import AutoImageProcessor, AutoModel
+from transformers.image_utils import load_image
+
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = load_image(url)
+
+repo = "Motif-Technologies/Motif-Vision-Encoder"
+processor = AutoImageProcessor.from_pretrained(repo)
+model = AutoModel.from_pretrained(repo, trust_remote_code=True, dtype=torch.bfloat16).to("cuda").eval()
+
+inputs = processor(images=image, return_tensors="pt").to(model.device, torch.bfloat16)
+with torch.inference_mode():
+    outputs = model(**inputs)
+
+outputs.last_hidden_state   # (1, 1 + 4 + N, 4096)  CLS + registers + patch tokens
+outputs.pooler_output       # (1, 4096)             global (CLS) representation
+
+patch_tokens = outputs.last_hidden_state[:, 5:, :]   # (1, N, 4096), N = (H/16)*(W/16)
+```
+
+The processor resizes the shorter side to 512, center-crops to 512×512, and normalizes with
+ImageNet mean/std (BICUBIC). `H`/`W` must be multiples of 16.
+
+### Video
+
+An image is a 1-frame clip; a video is the same call with a `(B, T, 3, H, W)` tensor. Apply the
+same per-frame transform (resize → center-crop → ImageNet norm) and stack over time:
+
+```python
+import torch
+
+video = torch.randn(1, 8, 3, 256, 256, device="cuda", dtype=torch.bfloat16)  # (B, T, 3, H, W)
+with torch.inference_mode():
+    outputs = model(pixel_values=video)
+```
+
+## Model details
+
+<p align="center">
+  <img src="assets/architecture.png" width="820" alt="Motif Vision Encoder architecture: image and video inputs, patch embedding, 40-block transformer stack, and transformer block internals"/>
+</p>
+
+| | |
+|---|---|
+| Backbone | ViT-7B, patch 16, embed 4096, depth 40, heads 32, SwiGLU |
+| Register tokens | 4 |
+| Position encoding | 3D axial RoPE (T,H,W), `base=100.0` |
+| Video tokenizer | 3D Conv, tubelet size 2 |
+| Precision | bf16 weights |
+| Training | DINO + iBOT + KoLeo self-distillation, Gram anchoring, contrastive caption alignment |
+| Training data | ~0.47B samples — 448.6M images (96%) + 18.5M video clips (4%) |
+
+Outputs (`BaseModelOutputWithPooling`): `last_hidden_state` `(B, 1+4+N, 4096)`,
+`pooler_output` `(B, 4096)`.
+
+## Evaluation
+
+Compared against the strongest publicly reported self-supervised / vision backbones. Higher is
+better for every column. Best comparable value per column in bold, second best <u>underlined</u>.
+
+DAVIS S/M/L follow the DINOv3 protocol (J&F-mean at video short side 420/480, 840/960,
+1260/1440 px). V-JEPA 2.1 is not part of the DINOv3 Table 5 tracking benchmark, so only its
+single-resolution (S) figure is available.
+
+| Model | Training<br>data | DAVIS S<br>J&F ↑ | DAVIS M<br>J&F ↑ | DAVIS L<br>J&F ↑ | ImageNet-1K<br>lin. probe ↑ | ADE20K<br>mIoU ↑ | K400 ↑ |
+|---|---|---|---|---|---|---|---|
+| **Motif Vision Encoder** | 0.5B | **74.0** | **80.5** | **83.5** | 87.4 | <u>52.0</u> | 87.4 |
+| DINOv3 | 1.7B | <u>71.1</u> | <u>79.7</u> | <u>83.3</u> | 88.4 | **55.9** | <u>87.8</u> |
+| Web-DINO | 2B | 57.2 | 65.8 | 69.5 | 85.9 | 42.7 | 86.8 |
+| PEcore | 5.4B | 48.2 | 53.1 | 49.8 | **89.3** | 38.9 | **87.9** |
+| SigLIP2 | 10B | 56.1 | 62.3 | 62.9 | <u>89.1</u> | 45.4 | 86.9 |
+| OpenCLIP | 2B | – | – | – | – | – | – |
+| V-JEPA 2.1 | 0.022B | 69.0 | – | – | 85.5 | 47.9 | 87.7 |
+
+Protocol: DINOv3-style linear/attentive probes for image tasks; V-JEPA 2-style protocol for
+video. Baseline DAVIS / ADE20K / K400 figures are taken from the DINOv3 technical report's
+unified evaluation (Tab. 3, 5, 6) and ImageNet from Tab. 7; OpenCLIP is not in that report and
+its benchmarks are not reported under a comparable protocol.
+
+Motif is state of the art on DAVIS video tracking at every resolution (74.0 / 80.5 / 83.5 J&F)
+and stays competitive on the other image and video benchmarks, using roughly 1/3 of DINOv3's
+training data (~0.5B samples).
+
+<p align="center">
+  <img src="assets/davis_mask_propagation.gif" width="820" alt="Mask propagation: ground truth vs DINOv3 vs Motif"/>
+</p>
+
+<p align="center">
+  <img src="assets/dense_attention_comparison.png" width="820" alt="Dense attention and feature-similarity comparison across Motif, DINOv3, V-JEPA 2.1, and SigLIP2"/>
+</p>
+<p align="center"><em>Dense features on a single image (768px). Columns: query point, CLS attention, query-point attention, patch-feature cosine similarity. Motif and DINOv3 keep attention and similarity tightly localized on the queried object, while V-JEPA 2.1 and SigLIP2 are noticeably noisier.</em></p>
+
+
+## License
+
+Released under the **MIT License** (see `LICENSE`). The model was trained on data governed by the
+respective dataset licenses; downstream users are responsible for compliance with those terms.

config.json

@@ -0,0 +1,36 @@
+{
+  "architectures": [
+    "MotifVisionModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_motif_vision_encoder.MotifVisionConfig",
+    "AutoModel": "modeling_motif_vision_encoder.MotifVisionModel"
+  },
+  "depth": 40,
+  "drop_path_rate": 0.0,
+  "dtype": "bfloat16",
+  "embed_dim": 4096,
+  "ffn_bias": true,
+  "ffn_layer": "swiglu64",
+  "ffn_ratio": 3.0,
+  "gated_attention": "elementwise",
+  "img_size": 512,
+  "in_chans": 3,
+  "layerscale_init": 1e-05,
+  "mask_k_bias": true,
+  "model_type": "motif_vision",
+  "n_storage_tokens": 4,
+  "norm_layer": "layernormbf16",
+  "num_frames": 1,
+  "num_heads": 32,
+  "patch_size": 16,
+  "pos_embed_rope_base": 100.0,
+  "pos_embed_rope_rescale_coords": 2.0,
+  "proj_bias": true,
+  "qk_norm": true,
+  "qkv_bias": false,
+  "transformers_version": "5.8.1",
+  "tubelet_size": 2,
+  "untie_cls_and_patch_norms": false,
+  "untie_global_and_local_cls_norm": true
+}

model-00001-of-00004.safetensors

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

model-00002-of-00004.safetensors

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

model-00003-of-00004.safetensors

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

model-00004-of-00004.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,780 @@
+{
+  "metadata": {
+    "total_parameters": 7390451712,
+    "total_size": 14780903552
+  },
+  "weight_map": {
+    "backbone.blocks.0.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.ls2.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.mlp.w2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.mlp.w2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.mlp.w3.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.mlp.w3.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.0.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.ls2.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.mlp.w2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.mlp.w2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.mlp.w3.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.mlp.w3.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.1.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.10.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.10.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.10.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.10.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.10.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.10.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.11.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.mlp.w1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.mlp.w1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.11.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.mlp.w1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.mlp.w1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.12.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.mlp.w1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.mlp.w1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.13.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.mlp.w1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.mlp.w1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.14.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.mlp.w1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.mlp.w1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.15.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.mlp.w1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.mlp.w1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.16.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.mlp.w1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.mlp.w1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.17.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.mlp.w1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.mlp.w1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.18.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.mlp.w1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.mlp.w1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.19.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.2.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.ls2.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.mlp.w2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.mlp.w2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.mlp.w3.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.mlp.w3.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.2.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.20.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.ls2.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.mlp.w1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.mlp.w1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.mlp.w2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.mlp.w2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.mlp.w3.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.mlp.w3.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.20.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.attn.gate_proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.attn.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.attn.proj.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.attn.proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.attn.qkv.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.ls1.gamma": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.21.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.21.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.21.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.21.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.21.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.21.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.21.norm1.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.norm1.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.norm2.bias": "model-00002-of-00004.safetensors",
+    "backbone.blocks.21.norm2.weight": "model-00002-of-00004.safetensors",
+    "backbone.blocks.22.attn.gate_proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.attn.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.attn.proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.attn.proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.ls1.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.norm2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.22.norm2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.attn.gate_proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.attn.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.attn.proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.attn.proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.ls1.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.norm2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.23.norm2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.attn.gate_proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.attn.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.attn.proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.attn.proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.ls1.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.norm2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.24.norm2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.attn.gate_proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.attn.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.attn.proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.attn.proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.ls1.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.norm2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.25.norm2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.attn.gate_proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.attn.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.attn.proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.attn.proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.ls1.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.norm2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.26.norm2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.attn.gate_proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.attn.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.attn.proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.attn.proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.ls1.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.norm2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.27.norm2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.attn.gate_proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.attn.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.attn.proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.attn.proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.ls1.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.norm2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.28.norm2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.attn.gate_proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.attn.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.attn.proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.attn.proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.ls1.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.norm2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.29.norm2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.3.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.ls2.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.mlp.w2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.mlp.w2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.mlp.w3.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.mlp.w3.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.3.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.30.attn.gate_proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.attn.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.attn.proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.attn.proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.ls1.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.norm2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.30.norm2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.attn.gate_proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.attn.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.attn.proj.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.attn.proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.ls1.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.ls2.gamma": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.mlp.w1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.mlp.w1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.mlp.w2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.mlp.w2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.mlp.w3.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.mlp.w3.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.norm2.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.31.norm2.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.32.attn.gate_proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.attn.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.attn.qkv.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.32.ls1.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.ls2.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.mlp.w1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.mlp.w1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.mlp.w2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.mlp.w2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.mlp.w3.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.mlp.w3.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.norm1.bias": "model-00003-of-00004.safetensors",
+    "backbone.blocks.32.norm1.weight": "model-00003-of-00004.safetensors",
+    "backbone.blocks.32.norm2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.32.norm2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.attn.gate_proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.attn.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.ls1.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.ls2.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.mlp.w1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.mlp.w1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.mlp.w2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.mlp.w2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.mlp.w3.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.mlp.w3.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.norm1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.norm1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.norm2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.33.norm2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.attn.gate_proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.attn.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.ls1.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.ls2.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.mlp.w1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.mlp.w1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.mlp.w2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.mlp.w2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.mlp.w3.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.mlp.w3.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.norm1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.norm1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.norm2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.34.norm2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.attn.gate_proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.attn.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.ls1.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.ls2.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.mlp.w1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.mlp.w1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.mlp.w2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.mlp.w2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.mlp.w3.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.mlp.w3.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.norm1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.norm1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.norm2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.35.norm2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.attn.gate_proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.attn.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.ls1.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.ls2.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.mlp.w1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.mlp.w1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.mlp.w2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.mlp.w2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.mlp.w3.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.mlp.w3.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.norm1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.norm1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.norm2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.36.norm2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.attn.gate_proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.attn.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.ls1.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.ls2.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.mlp.w1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.mlp.w1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.mlp.w2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.mlp.w2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.mlp.w3.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.mlp.w3.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.norm1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.norm1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.norm2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.37.norm2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.attn.gate_proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.attn.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.ls1.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.ls2.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.mlp.w1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.mlp.w1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.mlp.w2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.mlp.w2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.mlp.w3.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.mlp.w3.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.norm1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.norm1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.norm2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.38.norm2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.attn.gate_proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.attn.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.attn.proj.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.attn.proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.attn.qkv.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.ls1.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.ls2.gamma": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.mlp.w1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.mlp.w1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.mlp.w2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.mlp.w2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.mlp.w3.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.mlp.w3.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.norm1.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.norm1.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.norm2.bias": "model-00004-of-00004.safetensors",
+    "backbone.blocks.39.norm2.weight": "model-00004-of-00004.safetensors",
+    "backbone.blocks.4.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.ls2.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.mlp.w2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.mlp.w2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.mlp.w3.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.mlp.w3.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.4.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.ls2.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.mlp.w2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.mlp.w2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.mlp.w3.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.mlp.w3.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.5.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.ls2.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.mlp.w2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.mlp.w2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.mlp.w3.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.mlp.w3.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.6.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.ls2.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.mlp.w2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.mlp.w2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.mlp.w3.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.mlp.w3.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.7.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.ls2.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.mlp.w2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.mlp.w2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.mlp.w3.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.mlp.w3.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.8.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.attn.gate_proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.attn.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.attn.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.attn.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.attn.qkv.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.ls1.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.ls2.gamma": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.mlp.w1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.mlp.w1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.mlp.w2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.mlp.w2.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.mlp.w3.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.mlp.w3.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.norm1.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.norm1.weight": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.norm2.bias": "model-00001-of-00004.safetensors",
+    "backbone.blocks.9.norm2.weight": "model-00001-of-00004.safetensors",
+    "backbone.cls_token": "model-00001-of-00004.safetensors",
+    "backbone.local_cls_norm.bias": "model-00004-of-00004.safetensors",
+    "backbone.local_cls_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.mask_token": "model-00001-of-00004.safetensors",
+    "backbone.norm.bias": "model-00004-of-00004.safetensors",
+    "backbone.norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.patch_embed.proj.bias": "model-00001-of-00004.safetensors",
+    "backbone.patch_embed.proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.rope_embed.periods_h": "model-00001-of-00004.safetensors",
+    "backbone.rope_embed.periods_t": "model-00001-of-00004.safetensors",
+    "backbone.rope_embed.periods_w": "model-00001-of-00004.safetensors",
+    "backbone.storage_tokens": "model-00001-of-00004.safetensors"
+  }
+}

modeling_motif_vision_encoder.py

@@ -0,0 +1,1468 @@
+# Copyright (c) Motif Technologies.
+# Self-contained inference model for Motif Vision Encoder (image + video).
+# Auto-assembled from the training repo's inference path; NO training code.
+"""Motif Vision Encoder — unified image/video ViT backbone (inference-only).
+
+Usage:
+    from transformers import AutoModel
+    import torch
+    model = AutoModel.from_pretrained("Motif-Technologies/motif-vision-encoder",
+                                      trust_remote_code=True).eval()
+    # image: (B, 3, H, W)   video: (B, T, 3, H, W)  (H,W multiples of 16)
+    out = model(pixel_values=torch.randn(1, 3, 224, 224))
+    out.last_hidden_state            # (B, 1+num_register+N, D)
+    out.pooler_output                # (B, D)  CLS token
+"""
+import logging
+import math
+from functools import partial
+from typing import Any, Callable, Literal
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+
+
+
+# ---- utils ----
+
+
+def cat_keep_shapes(x_list: list[Tensor]) -> tuple[Tensor, list[tuple[int]], list[int]]:
+    """Concatenate list of tensors while preserving their shapes for later reconstruction."""
+    shapes = [x.shape for x in x_list]
+    num_tokens = [x.select(dim=-1, index=0).numel() for x in x_list]
+    flattened = torch.cat([x.flatten(0, -2) for x in x_list])
+    return flattened, shapes, num_tokens
+
+
+
+def uncat_with_shapes(flattened: Tensor, shapes: list[tuple[int]], num_tokens: list[int]) -> list[Tensor]:
+    """Reverse of cat_keep_shapes: split and reshape flattened tensor back to original shapes."""
+    outputs_splitted = torch.split_with_sizes(flattened, num_tokens, dim=0)
+    shapes_adjusted = [shape[:-1] + torch.Size([flattened.shape[-1]]) for shape in shapes]
+    outputs_reshaped = [o.reshape(shape) for o, shape in zip(outputs_splitted, shapes_adjusted)]
+    return outputs_reshaped
+
+
+
+# ---- rms_norm ----
+
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization.
+
+    A simpler alternative to LayerNorm that normalizes by RMS without centering.
+
+    Args:
+        dim: Number of features.
+        eps: Small constant for numerical stability.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        eps: float = 1e-6,
+        device: torch.device | str | None = None,
+    ) -> None:
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim, device=device))
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Apply RMS normalization."""
+        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return x / rms * self.weight
+
+
+# ---- layer_scale ----
+
+
+class LayerScale(nn.Module):
+    """Per-channel scaling that allows gradual incorporation of each layer's contribution.
+
+    Initializes to a small value (e.g., 1e-5) so that early in training, each layer's
+    contribution is nearly zero, stabilizing deep network training.
+
+    Args:
+        dim: Number of channels.
+        init_values: Initial value for all channels.
+        inplace: Whether to apply scaling in-place.
+        device: Device for parameter allocation.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        init_values: float | Tensor = 1e-5,
+        inplace: bool = False,
+        device: torch.device | None = None,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(torch.empty(dim, device=device))
+        self.init_values = init_values
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Apply per-channel scaling."""
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+# ---- patch_embed ----
+
+
+class PatchEmbed(nn.Module):
+    """Video (5D) or Image (4D) to Patch Embedding via 3D Convolution.
+
+    Handles both modalities through a single Conv3d projection:
+    - Image (B, C, H, W): unsqueeze temporal dim -> (B, C, 1, H, W) -> Conv3d
+    - Video (B, T, C, H, W): transpose -> (B, C, T, H, W) -> Conv3d
+
+    Output: (B, N_total, embed_dim)
+        N_total = (T // tubelet_size) * (H // patch_size) * (W // patch_size)
+
+    Args:
+        img_size: Input image size (used for reference only).
+        patch_size: Spatial patch size in pixels.
+        in_chans: Number of input channels.
+        embed_dim: Output embedding dimension.
+        tubelet_size: Temporal patch size (number of frames per temporal token).
+        flatten_embedding: Whether to flatten spatial dimensions.
+    """
+
+    def __init__(
+        self,
+        img_size: int = 224,
+        patch_size: int = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        tubelet_size: int = 1,
+        flatten_embedding: bool = True,
+    ) -> None:
+        super().__init__()
+        self.img_size = img_size
+        self.patch_size = (patch_size, patch_size) if isinstance(patch_size, int) else patch_size
+        self.tubelet_size = tubelet_size
+        self.flatten_embedding = flatten_embedding
+        self.in_chans = in_chans
+
+        # 3D Convolution: kernel and stride = (tubelet_size, patch_h, patch_w)
+        self.proj = nn.Conv3d(
+            in_chans,
+            embed_dim,
+            kernel_size=(tubelet_size, self.patch_size[0], self.patch_size[1]),
+            stride=(tubelet_size, self.patch_size[0], self.patch_size[1]),
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Tokenize input images or videos.
+
+        Args:
+            x: Input tensor.
+               Image: (B, C, H, W) or Video: (B, T, C, H, W)
+
+        Returns:
+            Patch tokens of shape (B, N_total, embed_dim).
+        """
+        if x.ndim == 4:
+            # Image: (B, C, H, W) -> (B, C, 1, H, W)
+            x = x.unsqueeze(2)
+            # If tubelet_size > 1, repeat the single frame to match kernel size
+            if self.tubelet_size > 1:
+                x = x.expand(-1, -1, self.tubelet_size, -1, -1)
+        elif x.ndim == 5:
+            # Video: (B, T, C, H, W) -> (B, C, T, H, W)
+            x = x.transpose(1, 2)
+
+        # Conv3d Projection -> (B, embed_dim, T', H', W')
+        x = self.proj(x)
+
+        if self.flatten_embedding:
+            # Flatten spatial+temporal: (B, embed_dim, N) -> (B, N, embed_dim)
+            x = x.flatten(2).transpose(1, 2)
+
+        return x
+
+
+# ---- rope ----
+
+
+class RopePositionEmbedding3D(nn.Module):
+    """Full 3D axial RoPE with independent T/H/W frequency bands.
+
+    Unlike the original Motif implementation which simply repeats 2D spatial angles
+    across temporal frames (making temporal positions indistinguishable), this
+    implementation partitions the head dimension into three axis groups:
+
+        D_head = D_T + D_H + D_W   (no spare dimensions)
+
+    By default, the split is spatial-heavy for SSL (spatial quality is priority):
+        D_T = D_head // 4           (25% temporal)
+        D_H = (D_head - D_T) // 2   (37.5% height)
+        D_W = D_head - D_T - D_H    (37.5% width)
+        e.g., D_head=64 → T=16, H=24, W=24
+
+    This can be overridden via ``fhw_dim=(D_T, D_H, D_W)`` for full control.
+
+    For images (T=1): t=0 for all tokens, making temporal angles constant
+    and the output is equivalent to spatial-only RoPE.
+
+    Args:
+        embed_dim: Total embedding dimension.
+        num_heads: Number of attention heads.
+        fhw_dim: Optional explicit (D_T, D_H, D_W) partition. Each must be even.
+            If None, uses the spatial-heavy default described above.
+        base: Frequency base (100.0 for spatial vision convention).
+        min_period: Minimum period (alternative to base).
+        max_period: Maximum period (alternative to base).
+        normalize_coords: How to normalize coordinates.
+        shift_coords: Random shift range during training.
+        jitter_coords: Random jitter multiplier during training.
+        rescale_coords: Random rescale multiplier during training.
+        dtype: Data type for computation.
+        device: Device for parameter allocation.
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        *,
+        num_heads: int,
+        fhw_dim: tuple[int, int, int] | None = None,
+        base: float | None = 100.0,
+        min_period: float | None = None,
+        max_period: float | None = None,
+        normalize_coords: Literal["min", "max", "separate"] = "separate",
+        shift_coords: float | None = None,
+        jitter_coords: float | None = None,
+        rescale_coords: float | None = None,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+    ) -> None:
+        super().__init__()
+        both_periods = min_period is not None and max_period is not None
+        if (base is None and not both_periods) or (base is not None and both_periods):
+            raise ValueError("Either `base` or `min_period`+`max_period` must be provided.")
+
+        D_head = embed_dim // num_heads
+        self.base = base
+        self.min_period = min_period
+        self.max_period = max_period
+        self.D_head = D_head
+        self.normalize_coords = normalize_coords
+        self.shift_coords = shift_coords
+        self.jitter_coords = jitter_coords
+        self.rescale_coords = rescale_coords
+
+        # Partition head dimension into 3 groups: T, H, W (no spare)
+        if fhw_dim is not None:
+            self.D_T, self.D_H, self.D_W = fhw_dim
+            assert self.D_T + self.D_H + self.D_W == D_head, (
+                f"fhw_dim must sum to D_head={D_head}, got {sum(fhw_dim)}"
+            )
+        else:
+            # Default: spatial-heavy split (SSL prioritizes spatial quality)
+            self.D_T = D_head // 4                        # 25% temporal
+            self.D_H = (D_head - self.D_T) // 2           # 37.5% height
+            self.D_W = D_head - self.D_T - self.D_H       # 37.5% width
+        assert self.D_T % 2 == 0 and self.D_H % 2 == 0 and self.D_W % 2 == 0, (
+            f"All axis dims must be even, got T={self.D_T}, H={self.D_H}, W={self.D_W}"
+        )
+
+        self.dtype = dtype
+        # Separate period buffers for each axis (n_freqs = D_axis // 2)
+        self.register_buffer(
+            "periods_t",
+            torch.empty(self.D_T // 2, device=device, dtype=dtype),
+            persistent=True,
+        )
+        self.register_buffer(
+            "periods_h",
+            torch.empty(self.D_H // 2, device=device, dtype=dtype),
+            persistent=True,
+        )
+        self.register_buffer(
+            "periods_w",
+            torch.empty(self.D_W // 2, device=device, dtype=dtype),
+            persistent=True,
+        )
+        self._init_weights()
+
+    def forward(self, *, T: int = 1, H: int, W: int) -> tuple[Tensor, Tensor]:
+        """Compute 3D axial RoPE sin/cos for (T, H, W) grid.
+
+        The head dimension is partitioned as [D_T | D_H | D_W]:
+        - D_T: temporal frequency bands (angles vary with t)
+        - D_H: height frequency bands (angles vary with h)
+        - D_W: width frequency bands (angles vary with w)
+
+        For images (T=1), all tokens get t=0, so temporal angles are constant
+        and the output is equivalent to spatial-only RoPE.
+
+        Args:
+            T: Number of temporal positions (T_grid = num_frames // tubelet_size).
+            H: Height in patches.
+            W: Width in patches.
+
+        Returns:
+            Tuple of (sin, cos), each of shape (T*H*W, D_head).
+        """
+        device = self.periods_t.device
+        dtype = self.dtype
+        dd = {"device": device, "dtype": dtype}
+
+        # 1. Compute normalized coordinates for each axis
+        if T > 1:
+            coords_t = torch.arange(0.5, T, **dd) / T  # [T]
+        else:
+            coords_t = torch.tensor([0.5], **dd)  # [1] - constant for images
+
+        coords_h, coords_w = self._compute_spatial_coords(H, W, **dd)
+
+        # Shift to [-1, +1] range
+        coords_t = 2.0 * coords_t - 1.0  # [T]
+        coords_h = 2.0 * coords_h - 1.0  # [H]
+        coords_w = 2.0 * coords_w - 1.0  # [W]
+
+        # Apply training-time augmentations to spatial coords only
+        if self.training:
+            coords_h, coords_w = self._augment_spatial_coords(coords_h, coords_w, dd)
+
+        # 2. Compute raw angles for each axis (n_freqs = D_axis // 2)
+        angles_t = 2 * math.pi * coords_t[:, None] / self.periods_t[None, :]  # [T, D_T//2]
+        angles_h = 2 * math.pi * coords_h[:, None] / self.periods_h[None, :]  # [H, D_H//2]
+        angles_w = 2 * math.pi * coords_w[:, None] / self.periods_w[None, :]  # [W, D_W//2]
+
+        # 3. Build full 3D grid: create (T*H*W, D_head) angle tensor
+        t_idx, h_idx, w_idx = torch.meshgrid(
+            torch.arange(T, device=device),
+            torch.arange(H, device=device),
+            torch.arange(W, device=device),
+            indexing="ij",
+        )
+        t_idx = t_idx.flatten()  # [T*H*W]
+        h_idx = h_idx.flatten()  # [T*H*W]
+        w_idx = w_idx.flatten()  # [T*H*W]
+
+        # Gather per-token raw angles and concatenate to D_head//2
+        token_angles_t = angles_t[t_idx]  # [T*H*W, D_T//2]
+        token_angles_h = angles_h[h_idx]  # [T*H*W, D_H//2]
+        token_angles_w = angles_w[w_idx]  # [T*H*W, D_W//2]
+        angles_half = torch.cat([token_angles_t, token_angles_h, token_angles_w], dim=-1)  # [T*H*W, D_head//2]
+
+        # tile(2) on full concat — matches Motif 2D RoPE pattern
+        # This ensures rotate_half pairs (dim i ↔ dim i+D//2) have identical angles,
+        # making the rotation orthogonal (preserves dot products in attention).
+        angles = angles_half.tile(2)  # [T*H*W, D_head]
+
+        cos = torch.cos(angles)
+        sin = torch.sin(angles)
+
+        return (sin, cos)
+
+    def _compute_spatial_coords(self, H: int, W: int, **dd) -> tuple[Tensor, Tensor]:
+        """Compute normalized spatial coordinates."""
+        if self.normalize_coords == "max":
+            max_HW = max(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / max_HW
+            coords_w = torch.arange(0.5, W, **dd) / max_HW
+        elif self.normalize_coords == "min":
+            min_HW = min(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / min_HW
+            coords_w = torch.arange(0.5, W, **dd) / min_HW
+        elif self.normalize_coords == "separate":
+            coords_h = torch.arange(0.5, H, **dd) / H
+            coords_w = torch.arange(0.5, W, **dd) / W
+        else:
+            raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}")
+        return coords_h, coords_w
+
+    def _augment_spatial_coords(
+        self,
+        coords_h: Tensor,
+        coords_w: Tensor,
+        dd: dict,
+    ) -> tuple[Tensor, Tensor]:
+        """Apply training-time coordinate augmentations to spatial coords."""
+        if self.shift_coords is not None:
+            shift = torch.empty(2, **dd).uniform_(-self.shift_coords, self.shift_coords)
+            coords_h = coords_h + shift[0]
+            coords_w = coords_w + shift[1]
+        if self.jitter_coords is not None:
+            jitter_max = np.log(self.jitter_coords)
+            jitter = torch.empty(2, **dd).uniform_(-jitter_max, jitter_max).exp()
+            coords_h = coords_h * jitter[0]
+            coords_w = coords_w * jitter[1]
+        if self.rescale_coords is not None:
+            rescale_max = np.log(self.rescale_coords)
+            rescale = torch.empty(1, **dd).uniform_(-rescale_max, rescale_max).exp()
+            coords_h = coords_h * rescale
+            coords_w = coords_w * rescale
+        return coords_h, coords_w
+
+    def _compute_periods(self, n_freqs: int, device: torch.device, dtype: torch.dtype | None) -> Tensor:
+        """Compute frequency periods for a single axis.
+
+        Args:
+            n_freqs: Number of frequency bands (D_axis // 2).
+            device: Device for tensor allocation.
+            dtype: Data type for computation.
+
+        Returns:
+            Tensor of shape (n_freqs,) with logarithmically spaced periods.
+        """
+        if self.base is not None:
+            return self.base ** (
+                2 * torch.arange(n_freqs, device=device, dtype=dtype) / (2 * n_freqs)
+            )
+        else:
+            base = self.max_period / self.min_period
+            exponents = torch.linspace(0, 1, n_freqs, device=device, dtype=dtype)
+            periods = base**exponents
+            periods = periods / base
+            return periods * self.max_period
+
+    def _init_weights(self) -> None:
+        """Initialize frequency periods for all three axes.
+
+        Each axis gets its own frequency schedule based on its dimension size:
+            periods[i] = base^(2i / D_axis)
+
+        This produces logarithmically spaced periods from 1.0 to base,
+        with more frequencies for axes with more allocated dimensions.
+        """
+        device = self.periods_t.device
+        dtype = self.dtype
+
+        self.periods_t.data = self._compute_periods(self.D_T // 2, device, dtype)
+        self.periods_h.data = self._compute_periods(self.D_H // 2, device, dtype)
+        self.periods_w.data = self._compute_periods(self.D_W // 2, device, dtype)
+
+
+# ---- attention ----
+
+
+def rope_rotate_half(x: Tensor) -> Tensor:
+    """Rotate half of the dimensions: [-x2, x1] from [x1, x2].
+
+    Args:
+        x: Input tensor of shape (..., D).
+
+    Returns:
+        Rotated tensor of shape (..., D).
+    """
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def rope_apply(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
+    """Apply rotary position embedding to input tensor.
+
+    Args:
+        x: Input tensor of shape (..., D).
+        sin: Sine angles of shape (..., D).
+        cos: Cosine angles of shape (..., D).
+
+    Returns:
+        Rotated tensor of shape (..., D).
+    """
+    return (x * cos) + (rope_rotate_half(x) * sin)
+
+
+class LinearKMaskedBias(nn.Linear):
+    """Linear layer with masked bias for the K component of QKV.
+
+    Zeroes out the bias for the K component (middle third of output)
+    to avoid interference with RoPE positional encoding.
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        o = self.out_features
+        assert o % 3 == 0
+        if self.bias is not None:
+            self.register_buffer("bias_mask", torch.full_like(self.bias, fill_value=math.nan))
+
+    def forward(self, input: Tensor) -> Tensor:
+        """Forward pass with masked bias."""
+        masked_bias = self.bias * self.bias_mask.to(self.bias.dtype) if self.bias is not None else None
+        return F.linear(input, self.weight, masked_bias)
+
+
+class SelfAttention(nn.Module):
+    """Multi-head self-attention with RoPE support.
+
+    Uses torch.nn.functional.scaled_dot_product_attention for FlashAttention
+    compatibility. RoPE is applied to Q and K on patch tokens only (not CLS/register).
+
+    Args:
+        dim: Model dimension.
+        num_heads: Number of attention heads.
+        qkv_bias: Whether to use bias in QKV projection.
+        proj_bias: Whether to use bias in output projection.
+        attn_drop: Attention dropout probability.
+        proj_drop: Output projection dropout probability.
+        mask_k_bias: Whether to mask K bias (for RoPE compatibility).
+        device: Device for parameter allocation.
+        gated_attention: Gated attention variant. None disables gating,
+            "headwise" applies a per-head scalar gate, "elementwise" applies
+            a per-element gate. Gate scores are query-dependent (derived from
+            input) and applied as sigmoid after SDPA.
+            Reference: https://arxiv.org/abs/2505.06708
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        mask_k_bias: bool = False,
+        device: str | None = None,
+        gated_attention: str | None = None,
+        qk_norm: bool = False,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim**-0.5
+
+        linear_class = LinearKMaskedBias if mask_k_bias else nn.Linear
+        self.qkv = linear_class(dim, dim * 3, bias=qkv_bias, device=device)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim, bias=proj_bias, device=device)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.qk_norm = qk_norm
+        if qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, device=device)
+            self.k_norm = RMSNorm(self.head_dim, device=device)
+
+        self.gated_attention = gated_attention
+        if gated_attention == "headwise":
+            self.gate_proj = nn.Linear(dim, num_heads, bias=True, device=device)
+        elif gated_attention == "elementwise":
+            self.gate_proj = nn.Linear(dim, dim, bias=True, device=device)
+        elif gated_attention is not None:
+            raise ValueError(f"Unknown gated_attention mode: {gated_attention!r}. Use 'headwise' or 'elementwise'.")
+
+    def apply_rope(
+        self,
+        q: Tensor,
+        k: Tensor,
+        rope: tuple[Tensor, Tensor],
+    ) -> tuple[Tensor, Tensor]:
+        """Apply RoPE to query and key tensors.
+
+        RoPE is applied only to patch tokens (prefix tokens like CLS and register
+        are excluded based on the difference between sequence length and rope length).
+
+        Args:
+            q: Query tensor of shape (B, heads, N, D_head).
+            k: Key tensor of shape (B, heads, N, D_head).
+            rope: Tuple of (sin, cos), each of shape (N_patches, D_head).
+
+        Returns:
+            Tuple of rotated (q, k) tensors.
+        """
+        q_dtype = q.dtype
+        k_dtype = k.dtype
+        sin, cos = rope
+        rope_dtype = sin.dtype
+        q = q.to(dtype=rope_dtype)
+        k = k.to(dtype=rope_dtype)
+        N = q.shape[-2]
+        prefix = N - sin.shape[-2]
+        assert prefix >= 0
+        q_prefix = q[:, :, :prefix, :]
+        q = rope_apply(q[:, :, prefix:, :], sin, cos)
+        q = torch.cat((q_prefix, q), dim=-2)
+        k_prefix = k[:, :, :prefix, :]
+        k = rope_apply(k[:, :, prefix:, :], sin, cos)
+        k = torch.cat((k_prefix, k), dim=-2)
+        q = q.to(dtype=q_dtype)
+        k = k.to(dtype=k_dtype)
+        return q, k
+
+    def forward(self, x: Tensor, attn_bias: Tensor | None = None, rope: Tensor | None = None) -> Tensor:
+        """Forward pass for single tensor input.
+
+        Args:
+            x: Input tensor of shape (B, N, D).
+            attn_bias: Unused (kept for interface compatibility).
+            rope: Optional RoPE (sin, cos) tuple.
+
+        Returns:
+            Output tensor of shape (B, N, D).
+        """
+        gate_score = self._compute_gate(x) if self.gated_attention else None
+        qkv = self.qkv(x)
+        attn_v = self.compute_attention(qkv=qkv, attn_bias=attn_bias, rope=rope, gate_score=gate_score)
+        x = self.proj(attn_v)
+        x = self.proj_drop(x)
+        return x
+
+    def forward_list(
+        self,
+        x_list: list[Tensor],
+        attn_bias: Tensor | None = None,
+        rope_list: list[tuple[Tensor, Tensor]] | None = None,
+    ) -> list[Tensor]:
+        """Forward pass for list of tensors (multi-crop efficiency).
+
+        Concatenates inputs for a single QKV projection, then splits for per-crop
+        attention computation (needed because different crops have different RoPE).
+
+        Args:
+            x_list: List of input tensors.
+            attn_bias: Unused.
+            rope_list: List of RoPE (sin, cos) tuples, one per input.
+
+        Returns:
+            List of output tensors.
+        """
+        assert len(x_list) == len(rope_list)
+        x_flat, shapes, num_tokens = cat_keep_shapes(x_list)
+        qkv_flat = self.qkv(x_flat)
+        qkv_list = uncat_with_shapes(qkv_flat, shapes, num_tokens)
+
+        if self.gated_attention:
+            gate_flat = self._compute_gate(x_flat)
+            gate_list = uncat_with_shapes(gate_flat, shapes, num_tokens)
+        else:
+            gate_list = [None] * len(x_list)
+
+        att_out = []
+        for qkv, _, rope, gate_score in zip(qkv_list, shapes, rope_list, gate_list):
+            att_out.append(self.compute_attention(qkv, attn_bias=attn_bias, rope=rope, gate_score=gate_score))
+        x_flat, shapes, num_tokens = cat_keep_shapes(att_out)
+        x_flat = self.proj(x_flat)
+        return uncat_with_shapes(x_flat, shapes, num_tokens)
+
+    def _compute_gate(self, x: Tensor) -> Tensor:
+        """Compute raw gate scores from input.
+
+        Returns the raw projection without reshaping so that the output keeps
+        the same number of leading dimensions as ``x``.  This is critical for
+        ``forward_list`` where ``uncat_with_shapes`` must split a 2-D flat
+        tensor back to per-crop 3-D tensors — adding extra dims here would
+        break that reshape.  The per-head unflatten happens later inside
+        ``compute_attention`` where B and N are known.
+
+        Args:
+            x: Input tensor of shape (..., D). Supports both 2D (flat) and 3D (batched).
+
+        Returns:
+            Raw gate projection. Headwise: (..., num_heads). Elementwise: (..., D).
+        """
+        return self.gate_proj(x)
+
+    def compute_attention(
+        self,
+        qkv: Tensor,
+        attn_bias: Tensor | None = None,
+        rope: tuple[Tensor, Tensor] | None = None,
+        gate_score: Tensor | None = None,
+    ) -> Tensor:
+        """Compute scaled dot-product attention.
+
+        Args:
+            qkv: Combined QKV tensor of shape (B, N, 3*D).
+            attn_bias: Unused.
+            rope: Optional RoPE (sin, cos) tuple.
+            gate_score: Optional gate tensor from _compute_gate.
+
+        Returns:
+            Attention output of shape (B, N, D).
+        """
+        assert attn_bias is None
+        B, N, _ = qkv.shape
+        C = self.qkv.in_features
+
+        qkv = qkv.reshape(B, N, 3, self.num_heads, self.head_dim)
+        q, k, v = torch.unbind(qkv, 2)
+        q, k, v = [t.transpose(1, 2) for t in [q, k, v]]
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        if rope is not None:
+            q, k = self.apply_rope(q, k, rope)
+        x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+        x = x.transpose(1, 2)  # (B, N, num_heads, head_dim)
+        if gate_score is not None:
+            # _compute_gate returns raw projection: (..., num_heads) or (..., D).
+            # Reshape to (B, N, num_heads, 1) or (B, N, num_heads, head_dim) here.
+            if self.gated_attention == "headwise":
+                gate_score = gate_score.unflatten(-1, (self.num_heads, 1))
+            else:  # elementwise
+                gate_score = gate_score.unflatten(-1, (self.num_heads, self.head_dim))
+            x = x * torch.sigmoid(gate_score)
+        return x.reshape([B, N, C])
+
+
+# ---- ffn ----
+
+
+class ListForwardMixin:
+    """Mixin providing forward_list for efficient multi-crop processing."""
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Forward pass for a single tensor."""
+        raise NotImplementedError
+
+    def forward_list(self, x_list: list[Tensor]) -> list[Tensor]:
+        """Forward pass for a list of tensors, concatenated for efficiency."""
+        x_flat, shapes, num_tokens = cat_keep_shapes(x_list)
+        x_flat = self.forward(x_flat)
+        return uncat_with_shapes(x_flat, shapes, num_tokens)
+
+
+class SwiGLUFFN(nn.Module, ListForwardMixin):
+    """SwiGLU Feed-Forward Network: w3(silu(w1(x)) * w2(x)).
+
+    Used for larger ViT models (SO400M+) due to better gradient flow.
+    Hidden dimension is aligned to a multiple of `align_to` for GPU efficiency.
+
+    Args:
+        in_features: Input dimension.
+        hidden_features: Hidden dimension before alignment.
+        out_features: Output dimension (default: same as in_features).
+        act_layer: Unused (SwiGLU has built-in SiLU activation).
+        drop: Unused (no dropout in SwiGLU).
+        bias: Whether to use bias in linear layers.
+        align_to: Align hidden dimension to this multiple.
+        device: Device for parameter allocation.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int | None = None,
+        out_features: int | None = None,
+        act_layer: Callable[..., nn.Module] | None = None,
+        drop: float = 0.0,
+        bias: bool = True,
+        align_to: int = 8,
+        device: str | None = None,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        d = int(hidden_features * 2 / 3)
+        swiglu_hidden_features = d + (-d % align_to)
+        self.w1 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device)
+        self.w2 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device)
+        self.w3 = nn.Linear(swiglu_hidden_features, out_features, bias=bias, device=device)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Forward pass: w3(silu(w1(x)) * w2(x))."""
+        x1 = self.w1(x)
+        x2 = self.w2(x)
+        hidden = F.silu(x1) * x2
+        return self.w3(hidden)
+
+
+# ---- block ----
+
+
+class SelfAttentionBlock(nn.Module):
+    """Pre-norm transformer block: Norm -> Attention -> LayerScale -> Residual (x2).
+
+    Supports both single-tensor and list-of-tensors forward for efficient multi-crop
+    processing.
+
+    Args:
+        dim: Model dimension.
+        num_heads: Number of attention heads.
+        ffn_ratio: FFN hidden dimension ratio.
+        qkv_bias: Whether to use bias in QKV projection.
+        proj_bias: Whether to use bias in output projection.
+        ffn_bias: Whether to use bias in FFN layers.
+        drop: Dropout probability.
+        attn_drop: Attention dropout probability.
+        init_values: LayerScale initial values (None disables LayerScale).
+        drop_path: Stochastic depth drop probability.
+        act_layer: Activation function class.
+        norm_layer: Normalization layer class.
+        attn_class: Attention class.
+        ffn_layer: FFN class.
+        mask_k_bias: Whether to mask K bias.
+        device: Device for parameter allocation.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        ffn_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        init_values: float | None = None,
+        drop_path: float = 0.0,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        attn_class: Callable[..., nn.Module] = SelfAttention,
+        ffn_layer: Callable[..., nn.Module] = SwiGLUFFN,
+        mask_k_bias: bool = False,
+        device: str | None = None,
+        gated_attention: str | None = None,
+        qk_norm: bool = False,
+    ) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = attn_class(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+            mask_k_bias=mask_k_bias,
+            device=device,
+            gated_attention=gated_attention,
+            qk_norm=qk_norm,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values, device=device) if init_values else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * ffn_ratio)
+        self.mlp = ffn_layer(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+            bias=ffn_bias,
+            device=device,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values, device=device) if init_values else nn.Identity()
+
+    def _forward_list(self, x_list: list[Tensor], rope_list: list | None = None) -> list[Tensor]:
+        """Forward pass for a list of tensors (one per crop), each with its own RoPE.
+
+        Pre-norm residual: Norm -> Attention -> LayerScale -> Residual, twice (attn, ffn).
+        """
+        x_out = []
+        for x, rope in zip(x_list, rope_list):
+            x_attn = x + self.ls1(self.attn(self.norm1(x), rope=rope))
+            x_ffn_item = x_attn + self.ls2(self.mlp(self.norm2(x_attn)))
+            x_out.append(x_ffn_item)
+        return x_out
+
+    def forward(
+        self,
+        x_or_x_list: Tensor | list[Tensor],
+        rope_or_rope_list: tuple | list | None = None,
+    ) -> Tensor | list[Tensor]:
+        """Forward pass accepting either a single tensor or list of tensors.
+
+        Args:
+            x_or_x_list: Single tensor (B, N, D) or list of tensors.
+            rope_or_rope_list: Single RoPE tuple or list of RoPE tuples.
+
+        Returns:
+            Output tensor(s) matching input format.
+        """
+        if isinstance(x_or_x_list, Tensor):
+            return self._forward_list([x_or_x_list], rope_list=[rope_or_rope_list])[0]
+        elif isinstance(x_or_x_list, list):
+            if rope_or_rope_list is None:
+                rope_or_rope_list = [None for _ in x_or_x_list]
+            return self._forward_list(x_or_x_list, rope_list=rope_or_rope_list)
+        else:
+            raise AssertionError(f"Unexpected input type: {type(x_or_x_list)}")
+
+
+# ---- vision_transformer ----
+
+
+logger = logging.getLogger("motif")
+
+ffn_layer_dict: dict[str, type] = {
+    "swiglu": SwiGLUFFN,
+    "swiglu32": partial(SwiGLUFFN, align_to=32),
+    "swiglu64": partial(SwiGLUFFN, align_to=64),
+    "swiglu128": partial(SwiGLUFFN, align_to=128),
+}
+
+norm_layer_dict: dict[str, type] = {
+    "layernorm": partial(nn.LayerNorm, eps=1e-6),
+    "layernormbf16": partial(nn.LayerNorm, eps=1e-5),
+    "rmsnorm": RMSNorm,
+}
+
+dtype_dict: dict[str, torch.dtype] = {
+    "fp32": torch.float32,
+    "fp16": torch.float16,
+    "bf16": torch.bfloat16,
+}
+
+
+class MotifVisionTransformer(nn.Module):
+    """Vision Transformer backbone with 3D RoPE for unified image/video processing.
+
+    Key features:
+    - PatchEmbed (Conv3d) for unified image/video tokenization
+    - Full 3D axial RoPE positional encoding (T/H/W)
+    - CLS token + register (storage) tokens
+    - LayerScale
+    - MLP or SwiGLU FFN variants
+
+    Token sequence layout: [CLS] + [Register x n_storage_tokens] + [Patch x N_total]
+
+    Args:
+        img_size: Input image size.
+        patch_size: Spatial patch size.
+        in_chans: Number of input channels.
+        embed_dim: Embedding dimension.
+        depth: Number of transformer blocks.
+        num_heads: Number of attention heads.
+        ffn_ratio: FFN hidden dimension ratio.
+        qkv_bias: Whether to use bias in QKV.
+        drop_path_rate: Stochastic depth rate.
+        layerscale_init: LayerScale initial value (None to disable).
+        norm_layer: Normalization layer name.
+        ffn_layer: FFN layer name.
+        ffn_bias: Whether to use bias in FFN.
+        proj_bias: Whether to use bias in attention output projection.
+        n_storage_tokens: Number of register tokens.
+        mask_k_bias: Whether to mask K bias in attention.
+        untie_cls_and_patch_norms: Use separate norms for CLS and patch tokens.
+        untie_global_and_local_cls_norm: Use separate norm for local CLS tokens.
+        device: Device for parameter allocation.
+        num_frames: Number of input video frames.
+        tubelet_size: Temporal patch size for Conv3d.
+        pos_embed_rope_base: RoPE frequency base (100.0 for vision).
+        gated_attention: Gated attention variant (None, "headwise", "elementwise").
+            See https://arxiv.org/abs/2505.06708.
+    """
+
+    def __init__(
+        self,
+        *,
+        img_size: int = 224,
+        patch_size: int = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        ffn_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        drop_path_rate: float = 0.0,
+        layerscale_init: float | None = None,
+        norm_layer: str = "layernorm",
+        ffn_layer: str = "mlp",
+        ffn_bias: bool = True,
+        proj_bias: bool = True,
+        n_storage_tokens: int = 0,
+        mask_k_bias: bool = False,
+        untie_cls_and_patch_norms: bool = False,
+        untie_global_and_local_cls_norm: bool = False,
+        device: Any | None = None,
+        num_frames: int = 1,
+        tubelet_size: int = 1,
+        pos_embed_rope_base: float = 100.0,
+        pos_embed_rope_rescale_coords: float | None = None,
+        pos_embed_rope_shift_coords: float | None = None,
+        pos_embed_rope_jitter_coords: float | None = None,
+        pos_embed_rope_fhw_dim: tuple[int, int, int] | None = None,
+        gated_attention: str | None = None,
+        qk_norm: bool = False,
+        **ignored_kwargs,
+    ) -> None:
+        super().__init__()
+        if len(ignored_kwargs) > 0:
+            logger.warning(f"Ignored kwargs: {ignored_kwargs}")
+
+        norm_layer_cls = norm_layer_dict[norm_layer]
+
+        self.num_features = self.embed_dim = embed_dim
+        self.n_blocks = depth
+        self.num_heads = num_heads
+        self.patch_size = patch_size
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            tubelet_size=tubelet_size,
+            flatten_embedding=True,
+        )
+
+        self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim, device=device))
+        self.n_storage_tokens = n_storage_tokens
+        if self.n_storage_tokens > 0:
+            self.storage_tokens = nn.Parameter(torch.empty(1, n_storage_tokens, embed_dim, device=device))
+
+        # Convert 0.0 to None for backward compat (0.0 means disabled)
+        _rescale = pos_embed_rope_rescale_coords if pos_embed_rope_rescale_coords else None
+        _shift = pos_embed_rope_shift_coords if pos_embed_rope_shift_coords else None
+        _jitter = pos_embed_rope_jitter_coords if pos_embed_rope_jitter_coords else None
+        self.rope_embed = RopePositionEmbedding3D(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            fhw_dim=pos_embed_rope_fhw_dim,
+            base=pos_embed_rope_base,
+            rescale_coords=_rescale,
+            shift_coords=_shift,
+            jitter_coords=_jitter,
+        )
+
+        logger.info(f"using {ffn_layer} layer as FFN")
+        ffn_layer_cls = ffn_layer_dict[ffn_layer]
+        ffn_ratio_sequence = [ffn_ratio] * depth
+
+        blocks_list = [
+            SelfAttentionBlock(
+                dim=embed_dim,
+                num_heads=num_heads,
+                ffn_ratio=ffn_ratio_sequence[i],
+                qkv_bias=qkv_bias,
+                proj_bias=proj_bias,
+                ffn_bias=ffn_bias,
+                drop_path=drop_path_rate,
+                norm_layer=norm_layer_cls,
+                act_layer=nn.GELU,
+                ffn_layer=ffn_layer_cls,
+                init_values=layerscale_init,
+                mask_k_bias=mask_k_bias,
+                device=device,
+                gated_attention=gated_attention,
+                qk_norm=qk_norm,
+            )
+            for i in range(depth)
+        ]
+
+        self.chunked_blocks = False
+        self.blocks = nn.ModuleList(blocks_list)
+
+        self.norm = norm_layer_cls(embed_dim)
+
+        self.untie_cls_and_patch_norms = untie_cls_and_patch_norms
+        if untie_cls_and_patch_norms:
+            self.cls_norm = norm_layer_cls(embed_dim)
+        else:
+            self.cls_norm = None
+
+        self.untie_global_and_local_cls_norm = untie_global_and_local_cls_norm
+        if untie_global_and_local_cls_norm:
+            self.local_cls_norm = norm_layer_cls(embed_dim)
+        else:
+            self.local_cls_norm = None
+        self.head = nn.Identity()
+        self.mask_token = nn.Parameter(torch.empty(1, embed_dim, device=device))
+
+    def prepare_tokens_with_masks(
+        self,
+        x: Tensor,
+        masks: Tensor | None = None,
+    ) -> tuple[Tensor, tuple[int, int, int]]:
+        """Tokenize input and assemble token sequence with CLS + register + patches.
+
+        Args:
+            x: Input tensor. Image: (B, C, H, W) or Video: (B, T, C, H, W).
+            masks: Boolean mask of shape (B, N_spatial) indicating which patches to mask.
+
+        Returns:
+            Tuple of:
+                - Token sequence: (B, 1 + n_storage + N_total, embed_dim)
+                - Grid dimensions: (T_grid, H_grid, W_grid)
+        """
+        if x.ndim == 5:
+            B, T, C, H, W = x.shape
+            # Video: Conv3d kernel=stride=tubelet downsamples raw T to T // tubelet.
+            T_grid = T // self.patch_embed.tubelet_size
+        else:
+            B, C, H, W = x.shape
+            # Image (4D): PatchEmbed expands raw T=1 to tubelet then Conv3d(stride=tubelet)
+            # produces a single temporal token (output T_out = (tubelet - tubelet)/tubelet + 1 = 1).
+            # vjepa2 vision_transformer.py:171-173 passes T=1 (no division) for the same reason.
+            T_grid = 1
+
+        x = self.patch_embed(x)  # (B, N_total, D)
+
+        # Grid dimensions for RoPE computation
+        H_grid = H // self.patch_embed.patch_size[0]
+        W_grid = W // self.patch_embed.patch_size[1]
+
+        if masks is not None:
+            # Expand spatial mask to spatio-temporal if needed (tube masking)
+            if masks.shape[1] != x.shape[1]:
+                ratio = x.shape[1] // masks.shape[1]
+                masks = masks.unsqueeze(1).repeat(1, ratio, 1).flatten(1)
+
+            # Replace masked positions with mask_token
+            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+            cls_token = self.cls_token
+        else:
+            # Include mask_token in computation graph even when not masking
+            cls_token = self.cls_token + 0 * self.mask_token
+
+        if self.n_storage_tokens > 0:
+            storage_tokens = self.storage_tokens
+        else:
+            storage_tokens = torch.empty(
+                1, 0, cls_token.shape[-1],
+                dtype=cls_token.dtype, device=cls_token.device,
+            )
+
+        x = torch.cat(
+            [
+                cls_token.expand(B, -1, -1),
+                storage_tokens.expand(B, -1, -1),
+                x,
+            ],
+            dim=1,
+        )
+
+        return x, (T_grid, H_grid, W_grid)
+
+    def forward_features_list(
+        self,
+        x_list: list[Tensor],
+        masks_list: list[Tensor | None],
+    ) -> list[dict[str, Tensor]]:
+        """Forward pass for a list of inputs (multi-crop).
+
+        Args:
+            x_list: List of input tensors (global crops, local crops).
+            masks_list: List of corresponding masks (None for unmasked).
+
+        Returns:
+            List of output dictionaries, one per input, containing:
+                - x_norm_clstoken: Normalized CLS token (B, D)
+                - x_storage_tokens: Normalized register tokens (B, n_storage, D)
+                - x_norm_patchtokens: Normalized patch tokens (B, N, D)
+                - x_prenorm: Pre-normalization features (B, 1+n_storage+N, D)
+                - masks: Original masks
+        """
+        x = []
+        rope_params = []
+        for t_x, t_masks in zip(x_list, masks_list):
+            t2_x, grid_tuple = self.prepare_tokens_with_masks(t_x, t_masks)
+            x.append(t2_x)
+            rope_params.append(grid_tuple)
+
+        # Pre-compute RoPE sin/cos once — identical across all blocks.
+        # Hoisting this out of the loop avoids breaking FSDP2's forward prefetch
+        # chain (rope_embed is part of the outer FSDP unit, calling it between
+        # block forwards disrupts the prefetch scheduling).
+        if self.rope_embed is not None:
+            rope_sincos = [self.rope_embed(T=t, H=h, W=w) for t, h, w in rope_params]
+        else:
+            rope_sincos = [None for _ in rope_params]
+
+        for _, blk in enumerate(self.blocks):
+            x = blk(x, rope_sincos)
+
+        all_x = x
+        output = []
+        for idx, (x, masks) in enumerate(zip(all_x, masks_list)):
+            if self.untie_cls_and_patch_norms or self.untie_global_and_local_cls_norm:
+                if self.untie_global_and_local_cls_norm and self.training and idx == 1:
+                    x_norm_cls_reg = self.local_cls_norm(x[:, : self.n_storage_tokens + 1])
+                elif self.untie_cls_and_patch_norms:
+                    x_norm_cls_reg = self.cls_norm(x[:, : self.n_storage_tokens + 1])
+                else:
+                    x_norm_cls_reg = self.norm(x[:, : self.n_storage_tokens + 1])
+                x_norm_patch = self.norm(x[:, self.n_storage_tokens + 1 :])
+            else:
+                x_norm = self.norm(x)
+                x_norm_cls_reg = x_norm[:, : self.n_storage_tokens + 1]
+                x_norm_patch = x_norm[:, self.n_storage_tokens + 1 :]
+            output.append(
+                {
+                    "x_norm_clstoken": x_norm_cls_reg[:, 0],
+                    "x_storage_tokens": x_norm_cls_reg[:, 1:],
+                    "x_norm_patchtokens": x_norm_patch,
+                    "x_prenorm": x,
+                    "masks": masks,
+                }
+            )
+        return output
+
+    def forward_features(
+        self,
+        x: Tensor | list[Tensor],
+        masks: Tensor | list[Tensor | None] | None = None,
+    ) -> dict[str, Tensor] | list[dict[str, Tensor]]:
+        """Forward pass for single or multiple inputs.
+
+        Args:
+            x: Single tensor or list of tensors.
+            masks: Single mask or list of masks.
+
+        Returns:
+            Output dict (single input) or list of output dicts (multiple inputs).
+        """
+        if isinstance(x, torch.Tensor):
+            return self.forward_features_list([x], [masks])[0]
+        else:
+            return self.forward_features_list(x, masks)
+
+    def _get_intermediate_layers_not_chunked(
+        self,
+        x: Tensor,
+        n: int | list[int] = 1,
+    ) -> list[Tensor]:
+        """Run forward pass and collect intermediate block outputs.
+
+        Args:
+            x: Input tensor (B, C, H, W) or (B, T, C, H, W).
+            n: If int, return last n layers. If list, return specific layer indices.
+
+        Returns:
+            List of intermediate outputs, each (B, 1+n_storage+N, D).
+        """
+        x, grid_tuple = self.prepare_tokens_with_masks(x, masks=None)
+        T, H, W = grid_tuple
+
+        output, total_block_len = [], len(self.blocks)
+        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+
+        if self.rope_embed is not None:
+            rope_sincos = self.rope_embed(T=T, H=H, W=W)
+        else:
+            rope_sincos = None
+
+        for i, blk in enumerate(self.blocks):
+            x = blk([x], [rope_sincos])[0]
+            if i in blocks_to_take:
+                output.append(x)
+
+        assert len(output) == len(blocks_to_take), (
+            f"only {len(output)} / {len(blocks_to_take)} blocks found"
+        )
+        return output
+
+    def get_intermediate_layers(
+        self,
+        x: Tensor,
+        n: int | list[int] = 1,
+        reshape: bool = False,
+        return_class_token: bool = False,
+        norm: bool = True,
+    ) -> tuple[Tensor, ...]:
+        """Extract intermediate layer outputs for downstream evaluation.
+
+        This method is critical for dense prediction tasks (segmentation, depth)
+        that need multi-scale features from different transformer blocks.
+
+        Args:
+            x: Input image (B, C, H, W) or video (B, T, C, H, W).
+            n: If int, return outputs from last n layers.
+               If list[int], return outputs from specific layer indices.
+            reshape: If True, reshape patch tokens to spatial form (B, D, H_grid, W_grid).
+            return_class_token: If True, return (patch_tokens, cls_token) tuples.
+            norm: If True, apply final LayerNorm to outputs.
+
+        Returns:
+            If return_class_token is False:
+                Tuple of patch token tensors, one per requested layer.
+                Each tensor is (B, N, D) or (B, D, H_grid, W_grid) if reshape=True.
+            If return_class_token is True:
+                Tuple of (patch_tokens, cls_token) pairs.
+        """
+        # Determine spatial dims for reshape
+        if x.ndim == 5:
+            B, T_in, C, H, W = x.shape
+        else:
+            B, C, H, W = x.shape
+            T_in = 1
+
+        outputs = self._get_intermediate_layers_not_chunked(x, n)
+
+        if norm:
+            outputs_normed = []
+            for out in outputs:
+                if self.untie_cls_and_patch_norms:
+                    x_norm_cls_reg = self.cls_norm(out[:, : self.n_storage_tokens + 1])
+                    x_norm_patch = self.norm(out[:, self.n_storage_tokens + 1 :])
+                    outputs_normed.append(torch.cat((x_norm_cls_reg, x_norm_patch), dim=1))
+                else:
+                    outputs_normed.append(self.norm(out))
+            outputs = outputs_normed
+
+        class_tokens = [out[:, 0] for out in outputs]
+        outputs = [out[:, self.n_storage_tokens + 1 :] for out in outputs]
+
+        if reshape:
+            # Image (T_in=1): PatchEmbed expands to tubelet then Conv3d(stride=tubelet) → T_out=1.
+            # Video: Conv3d downsamples T_in → T_in // tubelet. Matches prepare_tokens_with_masks
+            # and vjepa2 vision_transformer.py:171-177.
+            if x.ndim == 5:
+                T_grid = T_in // self.patch_embed.tubelet_size
+            else:
+                T_grid = 1
+            H_grid = H // self.patch_size
+            W_grid = W // self.patch_size
+            if T_grid > 1:
+                # Video: reshape to (B, D, T_grid, H_grid, W_grid)
+                outputs = [
+                    out.reshape(B, T_grid, H_grid, W_grid, -1).permute(0, 4, 1, 2, 3).contiguous()
+                    for out in outputs
+                ]
+            else:
+                # Image: reshape to (B, D, H_grid, W_grid)
+                outputs = [
+                    out.reshape(B, H_grid, W_grid, -1).permute(0, 3, 1, 2).contiguous()
+                    for out in outputs
+                ]
+
+        if return_class_token:
+            return tuple(zip(outputs, class_tokens))
+        return tuple(outputs)
+
+    def forward(
+        self,
+        *args,
+        is_training: bool = False,
+        **kwargs,
+    ) -> dict[str, Tensor] | list[dict[str, Tensor]] | Tensor:
+        """High-level forward: training returns feature dict, inference returns CLS logits.
+
+        Args:
+            is_training: If True, return full feature dictionary.
+
+        Returns:
+            Feature dict(s) if training, CLS token logits if inference.
+        """
+        ret = self.forward_features(*args, **kwargs)
+        if is_training:
+            return ret
+        else:
+            return self.head(ret["x_norm_clstoken"])
+
+
+# ============================================================================
+# HuggingFace transformers wrapper (inference-only)
+# ============================================================================
+class MotifVisionConfig(PretrainedConfig):
+    """Config for the Motif Vision Encoder backbone (image + video)."""
+
+    model_type = "motif_vision"
+
+    def __init__(
+        self,
+        img_size: int = 224,
+        patch_size: int = 16,
+        in_chans: int = 3,
+        embed_dim: int = 4096,
+        depth: int = 40,
+        num_heads: int = 32,
+        ffn_ratio: float = 3.0,
+        qkv_bias: bool = False,
+        drop_path_rate: float = 0.0,
+        layerscale_init: float | None = 1.0e-5,
+        norm_layer: str = "layernormbf16",
+        ffn_layer: str = "swiglu64",
+        ffn_bias: bool = True,
+        proj_bias: bool = True,
+        n_storage_tokens: int = 4,
+        mask_k_bias: bool = True,
+        untie_cls_and_patch_norms: bool = False,
+        untie_global_and_local_cls_norm: bool = True,
+        num_frames: int = 1,
+        tubelet_size: int = 2,
+        pos_embed_rope_base: float = 100.0,
+        pos_embed_rope_rescale_coords: float | None = 2.0,
+        gated_attention: str | None = "elementwise",
+        qk_norm: bool = True,
+        **kwargs,
+    ):
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.num_heads = num_heads
+        self.ffn_ratio = ffn_ratio
+        self.qkv_bias = qkv_bias
+        self.drop_path_rate = drop_path_rate
+        self.layerscale_init = layerscale_init
+        self.norm_layer = norm_layer
+        self.ffn_layer = ffn_layer
+        self.ffn_bias = ffn_bias
+        self.proj_bias = proj_bias
+        self.n_storage_tokens = n_storage_tokens
+        self.mask_k_bias = mask_k_bias
+        self.untie_cls_and_patch_norms = untie_cls_and_patch_norms
+        self.untie_global_and_local_cls_norm = untie_global_and_local_cls_norm
+        self.num_frames = num_frames
+        self.tubelet_size = tubelet_size
+        self.pos_embed_rope_base = pos_embed_rope_base
+        self.pos_embed_rope_rescale_coords = pos_embed_rope_rescale_coords
+        self.gated_attention = gated_attention
+        self.qk_norm = qk_norm
+        super().__init__(**kwargs)
+
+
+class MotifVisionModel(PreTrainedModel):
+    """Motif Vision Encoder for HF `AutoModel` (inference). Returns dense + CLS features."""
+
+    config_class = MotifVisionConfig
+    base_model_prefix = "motif"
+    main_input_name = "pixel_values"
+    _no_split_modules = ["SelfAttentionBlock"]
+    supports_gradient_checkpointing = False
+
+    def __init__(self, config: MotifVisionConfig):
+        super().__init__(config)
+        self.backbone = MotifVisionTransformer(
+            img_size=config.img_size,
+            patch_size=config.patch_size,
+            in_chans=config.in_chans,
+            embed_dim=config.embed_dim,
+            depth=config.depth,
+            num_heads=config.num_heads,
+            ffn_ratio=config.ffn_ratio,
+            qkv_bias=config.qkv_bias,
+            drop_path_rate=config.drop_path_rate,
+            layerscale_init=config.layerscale_init,
+            norm_layer=config.norm_layer,
+            ffn_layer=config.ffn_layer,
+            ffn_bias=config.ffn_bias,
+            proj_bias=config.proj_bias,
+            n_storage_tokens=config.n_storage_tokens,
+            mask_k_bias=config.mask_k_bias,
+            untie_cls_and_patch_norms=config.untie_cls_and_patch_norms,
+            untie_global_and_local_cls_norm=config.untie_global_and_local_cls_norm,
+            num_frames=config.num_frames,
+            tubelet_size=config.tubelet_size,
+            pos_embed_rope_base=config.pos_embed_rope_base,
+            pos_embed_rope_rescale_coords=config.pos_embed_rope_rescale_coords,
+            gated_attention=config.gated_attention,
+            qk_norm=config.qk_norm,
+        )
+        self.post_init()
+
+    @torch.no_grad()
+    def forward(self, pixel_values: Tensor, return_dict: bool = True, **kwargs):
+        """pixel_values: image (B,3,H,W) or video (B,T,3,H,W). H,W multiples of patch_size."""
+        out = self.backbone.forward_features(pixel_values)
+        cls = out["x_norm_clstoken"]
+        reg = out["x_storage_tokens"]
+        patch = out["x_norm_patchtokens"]
+        last_hidden = torch.cat([cls.unsqueeze(1), reg, patch], dim=1)
+        if not return_dict:
+            return (last_hidden, cls)
+        return BaseModelOutputWithPooling(last_hidden_state=last_hidden, pooler_output=cls)
+
+
+AutoConfig_registered = False
+try:
+    from transformers import AutoConfig, AutoModel
+    AutoConfig.register("motif_vision", MotifVisionConfig)
+    AutoModel.register(MotifVisionConfig, MotifVisionModel)
+    AutoConfig_registered = True
+except Exception:
+    pass

preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "image_processor_type": "BitImageProcessor",
+  "do_resize": true,
+  "size": {
+    "shortest_edge": 512
+  },
+  "do_center_crop": true,
+  "crop_size": {
+    "height": 512,
+    "width": 512
+  },
+  "do_rescale": true,
+  "rescale_factor": 0.00392156862745098,
+  "do_normalize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "resample": 3,
+  "do_convert_rgb": true
+}