mnmly/vjepa2.1-vit-base-384

V-JEPA 2.1 ViT-B/16 (384px) encoder, safetensors for MLX

The V-JEPA 2.1 ViT-B/16 video encoder from Meta FAIR (facebookresearch/vjepa2), repackaged as plain safetensors for the Swift / MLX port mlx-swift-vjepa2. Distilled from ViT-G.

Unofficial. This is a format conversion of Meta's released checkpoint, not a Meta release. The tensors are exactly the encoder that torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_base_384') loads.

Source checkpoint	vjepa2_1_vitb_dist_vitG_384.pt → ema_encoder
Upstream code	facebookresearch/vjepa2 @ 204698b (app/vjepa_2_1/models)
Parameters	86,833,152
Architecture	width 768, depth 12, 12 heads (64-dim), MLP 3072, patch 16, tubelet 2, 3-axis RoPE
Input	384×384 crops, ImageNet normalization; 1 frame → image path, 2+ frames → video path
Output	[N, 768] patch features (N = T/2 · 24 · 24), plus 4-level hierarchical features [N, 3072]

Files

File	dtype	Size	SHA-256
model.safetensors	float32 (bit-identical to the source checkpoint)	347 MB	e89509a1d64dbd117b17d5c32a1d0b1892d3698d75249e4a001cb490ca7e4017
model-fp16.safetensors	float16 (round-to-nearest from fp32)	174 MB	60de5fc3e6bdcb026b00c7e5d4c0f485f634dc42bb2397296ec45341a62de7e8
predictor.safetensors	float32 predictor (bit-identical to the checkpoint's predictor)	92 MB	7260d75185c1ea7da0980deac01b9f8068f144c42a3378ec71fe842fb2469655
predictor-fp16.safetensors	float16 predictor	46 MB	44dd38b28b2154c895b22befbf0f448c2101e033df1e9db661a730602ae323f2
config.json	architecture, preprocessing, provenance

The predictor (22,874,752 parameters) is the one Meta released in the same checkpoint. It predicts the features of hidden tokens from the encoder's features of the visible ones; it outputs 1664-wide features (the ViT-G teacher's feature space, which ViT-B was distilled from).

Keys and tensor layouts follow upstream's state dicts after _clean_backbone_key. The Conv3d patch embeddings stay in PyTorch layout (O, I, kT, kH, kW); the Swift loader transposes them at load time. That means you can also load the files into upstream's PyTorch VisionTransformer / VisionTransformerPredictor with load_state_dict(..., strict=True).

Usage

Swift (mlx-swift-vjepa2)

import VJEPA2

// Downloads model-fp16.safetensors (or model.safetensors for .float32) from this repo and caches it.
let session = try await VJEPA2Session.load(pretrained: .base, dtype: .float16)
let frames = try await VJEPA2VideoReader.frames(url: clipURL, count: 16, stride: 2)
let encoding = session.encode(frames)        // encoding.features: N × 768 floats

Predict hidden patches (the JEPA path) by also loading the predictor:

let session = try await VJEPA2Session.load(pretrained: .base, dtype: .float16, predictor: true)
let mask = VJEPA2TokenMask.block(in: session.tokenGrid(frameCount: frames.count), rows: 6 ..< 18, columns: 6 ..< 18)
let prediction = try session.predict(frames, mask: mask)   // mask.target.count × 1664 floats

CLI: vjepa2-cli embed --model base --video clip.mp4, or vjepa2-cli predict --model base --video clip.mp4. Or pick the model in the V-JEPA 2 Viewer app.

PyTorch (upstream code)

import torch
from safetensors.torch import load_file
from app.vjepa_2_1.models import vision_transformer as vit  # from facebookresearch/vjepa2

encoder = vit.vit_base(
    img_size=(384, 384), patch_size=16, num_frames=64, tubelet_size=2, use_sdpa=True, use_SiLU=False,
    wide_SiLU=True, uniform_power=False, use_rope=True, img_temporal_dim_size=1, interpolate_rope=True)
encoder.load_state_dict(load_file("model.safetensors"), strict=True)

Parity

The weights are checked against upstream's own PyTorch code. The upstream reference runs in fp32 on CPU, on a real 16-frame clip, a 4-frame clip, and a single image. Values are cosine similarity:

Exact path (MLX CPU stream), every saved block	GPU fp32 per stage (min)	End-to-end fp32: video / hierarchical / image	End-to-end fp16
1.000000	0.999888	0.999920 / 0.999944 / 0.999971	0.999933

The predictor is checked the same way, on masks from upstream's own mask generator: every predictor block matches at cosine ≥ 0.9999997 on the exact path, and pixels → masked encoder → predictor matches at ≥ 0.99998 (fp32) / ≥ 0.99997 (fp16) on the GPU.

Preprocessing (an OpenCV-exact bilinear resize, center crop, and normalization) reproduces upstream's tensor bit for bit. The fixture generator and the test suite ship with mlx-swift-vjepa2.

License

The weights and code are released by Meta under the MIT License. That license is reproduced verbatim in LICENSE, with the copyright notice "Copyright (c) Meta Platforms, Inc. and affiliates.". This repackaging adds no restrictions. If you redistribute these files, keep the license and copyright notice with them.

Citation

@article{murlabadia2026vjepa2_1,
  title={V-JEPA 2.1: Unlocking Dense Features in Video Self-Supervised Learning},
  author={Mur-Labadia, Lorenzo and Muckley, Matthew and Bar, Amir and Assran, Mahmoud and
Sinha, Koustuv and Rabbat, Michael and LeCun, Yann and Ballas, Nicolas and Bardes, Adrien},
  journal={arXiv preprint arXiv:2603.14482},
  year={2026}
}

@article{assran2025vjepa2,
  title={V-JEPA~2: Self-Supervised Video Models Enable Understanding, Prediction and Planning},
  author={Assran, Mahmoud and Bardes, Adrien and Fan, David and Garrido, Quentin and Howes, Russell and
Komeili, Mojtaba and Muckley, Matthew and Rizvi, Ammar and Roberts, Claire and Sinha, Koustuv and Zholus, Artem and
Arnaud, Sergio and Gejji, Abha and Martin, Ada and Robert Hogan, Francois and Dugas, Daniel and
Bojanowski, Piotr and Khalidov, Vasil and Labatut, Patrick and Massa, Francisco and Szafraniec, Marc and
Krishnakumar, Kapil and Li, Yong and Ma, Xiaodong and Chandar, Sarath and Meier, Franziska and LeCun, Yann and
Rabbat, Michael and Ballas, Nicolas},
  journal={arXiv preprint arXiv:2506.09985},
  year={2025}
}