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
metadata
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 for fast video feature extraction on Apple Silicon.

Used by the TRIBE v2 Mac fork: vasanthsreeram/tribev2 with the brain encoding head 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 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 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

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
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 if you use the brain model.
  • Conversion utilities based on vjepa2-mlx.

Not affiliated with Meta. This is a community conversion for Apple Silicon inference.