V-JEPA 2

V-JEPA 2 is an official PyTorch codebase for self-supervised video models designed for understanding, prediction, and planning. It includes three main variants:

1. V-JEPA 2: A self-supervised approach for training video encoders using internet-scale video data. It excels at motion understanding and human action anticipation.
2. V-JEPA 2.1: An improved version focusing on high-quality and temporally consistent dense features. It utilizes a Dense Predictive Loss (masking-based self-supervision where all tokens contribute to the loss), Deep Self-Supervision (applying loss at multiple intermediate representations), and Multi-Modal Tokenizers for images and videos.
3. V-JEPA 2-AC: A latent action-conditioned world model post-trained from V-JEPA 2 using robot trajectory interaction data. It is designed for robot manipulation tasks (like reaching, grasping, and pick-and-place) via planning from image goals, without requiring environment-specific data or task-specific calibration.

You can load pretrained V-JEPA 2 and V-JEPA 2.1 models directly using torch.hub. Ensure you have torch, timm, and einops installed locally. It is strongly recommended to use a PyTorch installation with CUDA support.

Available symbols include the preprocessor and various model sizes for both V-JEPA 2 and V-JEPA 2.1.

import torch

# preprocessor
processor = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_preprocessor')

# V-JEPA 2 models
vjepa2_vit_large = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_large')
vjepa2_vit_huge = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_huge')
vjepa2_vit_giant = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_giant')
vjepa2_vit_giant_384 = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_giant_384')

# V-JEPA 2.1 models
vjepa2_1_vit_base_384 = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_base_384')
vjepa2_1_vit_large_384 = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_large_384')
vjepa2_1_vit_giant_384 = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_giant_384')
vjepa2_1_vit_gigantic_384 = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_gigantic_384')

To load V-JEPA 2 encoder weights into a PyTorch model, use the load_pretrained_vjepa_pt_weights function. This function handles the preprocessing of the state_dict by removing module. and backbone. prefixes to ensure compatibility with your model instance.

Weights can be downloaded via wget from the official Facebook repositories (e.g., https://dl.fbaipublicfiles.com/vjepa2/vitg-384.pt).

def load_pretrained_vjepa_pt_weights(model, pretrained_weights):
    # Load weights of the VJEPA2 encoder
    pretrained_dict = torch.load(pretrained_weights, weights_only=True, map_location="cpu")["encoder"]
    pretrained_dict = {k.replace("module.", ""): v for k, v in pretrained_dict.items()}
    pretrained_dict = {k.replace("backbone.", ""): v for k, v in pretrained_dict.items()}
    msg = model.load_state_dict(pretrained_dict, strict=False)
    print("Pretrained weights found at {} and loaded with msg: {}".format(pretrained_weights, msg))

After extracting patch-wise features from the encoder, you can pass them to an AttentiveClassifier to predict action classes (e.g., for the Something-Something V2 dataset).

1. Initialize AttentiveClassifier with the encoder's embed_dim.
2. Load pretrained probe weights.
3. Pass the patch features to the classifier.
4. Map the output indices to class names using a JSON mapping file (e.g., ssv2_classes.json).

from src.models.attentive_pooler import AttentiveClassifier

# Initialize the classifier
classifier = AttentiveClassifier(
    embed_dim=model_pt.embed_dim, 
    num_heads=16, 
    depth=4, 
    num_classes=174
).cuda().eval()

# Load weights and predict
load_pretrained_vjepa_classifier_weights(classifier, classifier_model_path)

with torch.inference_mode():
    out_classifier = classifier(out_patch_features_pt)
    # out_classifier contains logits for top-k prediction