vasanth009's picture
Add MLX fp16 V-JEPA2 ViT-g encoder (from facebook/vjepa2-vitg-fpc64-256)
fa5afa4 verified
|
Raw
History Blame Contribute Delete
2.38 kB
---
license: cc-by-nc-4.0
library_name: mlx
tags:
- mlx
- vjepa2
- video
- apple-silicon
- tribe
- fmri
base_model: facebook/vjepa2-vitg-fpc64-256
pipeline_tag: feature-extraction
---
# V-JEPA2 ViT-g β†’ MLX (TRIBE vision backbone)
MLX (fp16) conversion of **[`facebook/vjepa2-vitg-fpc64-256`](https://huggingface.co/facebook/vjepa2-vitg-fpc64-256)** for fast video feature extraction on **Apple Silicon**.
Used by the TRIBE v2 Mac fork: [vasanthsreeram/tribev2](https://github.com/vasanthsreeram/tribev2) with the brain encoding head [`facebook/tribev2`](https://huggingface.co/facebook/tribev2).
## Files
| File | Description |
|------|-------------|
| `model.safetensors` | Encoder weights (fp16, MLX layout), ~1.9 GB |
| `config.json` | ViT-g encoder config (`hidden=1408`, 40 layers, …) |
## How this was converted
1. Load HF `AutoModel` for `facebook/vjepa2-vitg-fpc64-256` in **fp32**.
2. Map tensors with [vjepa2-mlx](https://github.com/xocialize/vjepa2-mlx) `convert_state_dict` (Conv3d layout, QKV, etc.).
3. Keep **`encoder.*` only** (predictor dropped β€” not used by TRIBE).
4. Cast to **fp16**, save safetensors with `format=mlx` metadata.
5. Config: `mlp_ratio=48/11` to match HF (not integer 4).
Full write-up: [docs/MLX_CONVERSION.md](https://github.com/vasanthsreeram/tribev2/blob/main/docs/MLX_CONVERSION.md) in the code fork.
## Validation
Against official torch TRIBE vision-only preds on a 1 s clip (same config):
- **Cosine β‰ˆ 0.96**, Pearson **r β‰ˆ 0.96**
- Encoder-only random input vs HF: cosine β‰ˆ **0.9994**
- Runtime: ~**11Γ—** faster than torch CPU encode on the test Mac
## Usage with TRIBE fork
```bash
hf download vasanth009/vjepa2-vitg-fpc64-256-mlx --local-dir mlx_weights/V-JEPA2-vitg-fpc64-256
# then in the tribev2 fork:
python demo_data/run_efficient.py your.mp4 --open
```
```python
from tribev2.mlx_vjepa import encode_clip_frames, install_mlx_video_hooks
# requires vjepa2_mlx package for the encoder graph
```
## License & attribution
- Weights derived from Meta V-JEPA2 β€” **CC-BY-NC-4.0** (non-commercial), same spirit as upstream.
- Cite the [TRIBE paper](https://arxiv.org/abs/2605.04326) if you use the brain model.
- Conversion utilities based on [vjepa2-mlx](https://github.com/xocialize/vjepa2-mlx).
**Not affiliated with Meta.** This is a community conversion for Apple Silicon inference.