### Download PyTorch Model Weights

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/vjepa2_demo.ipynb

Example command to download PyTorch model weights using `wget`. Ensure to replace `YOUR_DIR` with your desired local path.

```bash
wget https://dl.fbaipublicfiles.com/vjepa2/vitg-384.pt -P YOUR_DIR
```

--------------------------------

### Download Attentive Probe Weights

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/vjepa2_demo.ipynb

Example command to download pretrained attentive probe weights for classification. Replace `YOUR_DIR` with your desired local path.

```bash
wget https://dl.fbaipublicfiles.com/vjepa2/evals/ssv2-vitg-384-64x2x3.pt -P YOUR_DIR
```

--------------------------------

### Setup VJepa2 Environment

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Commands to create and activate a conda environment, and install the VJepa2 package. For development, use `pip install -e .`.

```bash
conda create -n vjepa2-312 python=3.12
conda activate vjepa2-312
pip install .  # or `pip install -e .` for development mode
```

--------------------------------

### Local Post-training

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Command to run post-training of the action-conditioned model locally, starting from the pretrained VJEPA 2 backbone. Specify the devices to use.

```python
python -m app.main --fname configs/train/vitg16/droid-256px-8f.yaml \
  --devices cuda:0
```

--------------------------------

### VisionTransformer Forward Pass Example

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Example of using the core `VisionTransformer` encoder, which handles both image and video inputs and supports optional masking and intermediate layer outputs.

```python
from src.models.vision_transformer import vit_giant_xformers_rope
import torch

```

--------------------------------

### Load Action-Conditioned V-JEPA 2 Backbone via PyTorch Hub

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Load the action-conditioned V-JEPA 2 backbone and predictor using PyTorch Hub. Ensure PyTorch, timm, and einops are installed. PyTorch with CUDA support is recommended.

```python
import torch

vjepa2_encoder, vjepa2_ac_predictor = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_ac_vit_giant')
```

--------------------------------

### Load Pretrained V-JEPA 2 Models from Huggingface

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Use this snippet to load pretrained V-JEPA 2 models and their corresponding processors from Huggingface. Ensure you have the 'transformers' library installed.

```python
from transformers import AutoVideoProcessor, AutoModel

hf_repo = "facebook/vjepa2-vitg-fpc64-256"
# facebook/vjepa2-vitl-fpc64-256
# facebook/vjepa2-vith-fpc64-256
# facebook/vjepa2-vitg-fpc64-256
# facebook/vjepa2-vitg-fpc64-384

model = AutoModel.from_pretrained(hf_repo)
processor = AutoVideoProcessor.from_pretrained(hf_repo)
```

--------------------------------

### Load V-JEPA 2 Models via PyTorch Hub

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Installs PyTorch, timm, and einops. Loads the V-JEPA 2 preprocessor and various V-JEPA 2 model variants (ViT-large, ViT-huge, ViT-giant, ViT-giant-384) using torch.hub.load. CUDA support for PyTorch is recommended.

```python
import torch

# preprocessor
processor = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_preprocessor')
# models
# V-JEPA 2
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')
```

--------------------------------

### Load V-JEPA 2.1 Models via PyTorch Hub

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Installs PyTorch, timm, and einops. Loads various V-JEPA 2.1 model variants (ViT-base-384, ViT-large-384, ViT-giant-384, ViT-gigantic-384) using torch.hub.load. CUDA support for PyTorch is recommended.

```python
import torch

# V-JEPA 2.1
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')
```

--------------------------------

### Launch V-JEPA 2.1 Pretraining (Bash)

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Use these commands to launch V-JEPA 2.1 pretraining. Choose between local single-GPU debugging or SLURM distributed training. Ensure correct paths and account/QoS information for SLURM.

```bash
python -m app.main \
    --fname configs/train_2_1/vitg16/pretrain-256px-16f.yaml \
    --devices cuda:0
```

```bash
python -m app.main_distributed \
    --fname configs/train_2_1/vitg16/pretrain-256px-16f.yaml \
    --time 6000 \
    --account my_account \
    --qos my_qos
```

```bash
python -m app.main_distributed \
    --fname configs/train_2_1/vitg16/cooldown-256px-64f.yaml \
    --time 6000 \
    --account my_account \
    --qos my_qos
```

--------------------------------

### Set up Python Path

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/energy_landscape_example.ipynb

Ensures the vjepa2 library is accessible by modifying the system path.

```python
import sys
sys.path.insert(0, "..")
```

--------------------------------

### Run Usage Demo Script

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Execute the VJepa2 usage demo script after downloading model checkpoints and updating paths. Assumes GPU availability.

```python
# Then update your model paths in vjepa2_demo.py.
pt_model_path = YOUR_DIR/vitg-384.pt
classifier_model_path = YOUR_DIR/ssv2-vitg-384-64x2x3.pt

# Then run the script (assumes your machine has a GPU)
python -m notebooks.vjepa2_demo
```

--------------------------------

### Download Sample Video and Class Mapping

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/vjepa2_demo.ipynb

Downloads a sample video file and a JSON mapping for action recognition classes if they do not already exist locally. This ensures the necessary data is available for running the V-JEPA 2 demo.

```python
sample_video_path = "sample_video.mp4"
# Download the video if not yet downloaded to local path
if not os.path.exists(sample_video_path):
    video_url = "https://huggingface.co/datasets/nateraw/kinetics-mini/resolve/main/val/bowling/-WH-lxmGJVY_000005_000015.mp4"
    command = ["wget", video_url, "-O", sample_video_path]
    subprocess.run(command)
    print("Downloading video")

```

--------------------------------

### Local Pretraining

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Command to launch initial pretraining of a ViT-L model locally. Specify the devices to use.

```python
python -m app.main --fname configs/train/vitl16/pretrain-256px-16f.yaml \
  --devices cuda:0
```

--------------------------------

### Initialize VJEPA2 Model and Transforms

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/energy_landscape_example.ipynb

Loads the VJEPA2-AC model and sets up image transformations for processing.

```python
# Initialize VJEPA 2-AC model
encoder, predictor = torch.hub.load("facebookresearch/vjepa2", "vjepa2_ac_vit_giant")

# Initialize transform
crop_size = 256
tokens_per_frame = int((crop_size // encoder.patch_size) ** 2)
transform = make_transforms(
    random_horizontal_flip=False,
    random_resize_aspect_ratio=(1., 1.),
    random_resize_scale=(1., 1.),
    reprob=0.,
    auto_augment=False,
    motion_shift=False,
    crop_size=crop_size,
)
```

--------------------------------

### Initialize and Load V-JEPA 2 Models

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/vjepa2_demo.ipynb

Initializes HuggingFace and PyTorch V-JEPA 2 models, loads pretrained weights, and builds corresponding preprocessing transforms. Ensure `pt_model_path` is set correctly.

```python
# HuggingFace model repo name
hf_model_name = (
    "facebook/vjepa2-vitg-fpc64-384"  # Replace with your favored model, e.g. facebook/vjepa2-vitg-fpc64-384
)
# Path to local PyTorch weights
pt_model_path = "YOUR_MODEL_PATH"

# Initialize the HuggingFace model, load pretrained weights
model_hf = AutoModel.from_pretrained(hf_model_name)
model_hf.cuda().eval()

# Build HuggingFace preprocessing transform
hf_transform = AutoVideoProcessor.from_pretrained(hf_model_name)
img_size = hf_transform.crop_size["height"]  # E.g. 384, 256, etc.

# Initialize the PyTorch model, load pretrained weights
model_pt = vit_giant_xformers_rope(img_size=(img_size, img_size), num_frames=64)
model_pt.cuda().eval()
load_pretrained_vjepa_pt_weights(model_pt, pt_model_path)

### Can also use torch.hub to load the model
# model_pt, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_giant_384')
# model_pt.cuda().eval()

# Build PyTorch preprocessing transform
pt_video_transform = build_pt_video_transform(img_size=img_size)
```

--------------------------------

### Launch Probe-Based Evaluation Scripts

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Commands for downloading checkpoints, training attentive probes locally or distributedly using SLURM, and performing inference. Ensure configuration files are correctly set for dataset paths and checkpoint locations.

```bash
# ---- Download backbone and probe checkpoints ----
wget https://dl.fbaipublicfiles.com/vjepa2/vitg-384.pt          -P /checkpoints/
wget https://dl.fbaipublicfiles.com/vjepa2/evals/ssv2-vitg-384-64x2x3.pt -P /checkpoints/ssv2/vitg-384/v1/

# ---- Train a new attentive probe locally (2 GPUs) ----
# Edit configs/eval/vitg-384/ssv2.yaml to set:
#   checkpoint: /checkpoints/vitg-384.pt
#   dataset_train: /data/ssv2/train.csv
#   dataset_val: /data/ssv2/val.csv
python -m evals.main \
    --fname configs/eval/vitg-384/ssv2.yaml \
    --devices cuda:0 cuda:1

# ---- Distributed probe training via SLURM ----
python -m evals.main_distributed \
    --fname configs/eval/vitg-384/ssv2.yaml \
    --time 8600 \
    --account my_account \
    --qos my_qos

# ---- Inference from existing checkpoint ----
# Place checkpoint at: [folder]/[eval_name]/[tag]/latest.pt
# (e.g. /checkpoints/ssv2/vitg-384/v1/latest.pt — matches folder/eval/tag in config)
python -m evals.main \
    --fname configs/inference/vitg-384/ssv2.yaml \
    --devices cuda:0

# Expected output (SSv2 top-1 accuracy):
# [epoch 0] val loss: 0.821 | val top-1: 77.3% | val top-5: 95.1%
```

--------------------------------

### Create Video Preprocessing Pipelines with `make_transforms`

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Generates video transformation pipelines for training or evaluation. Use `training=False` for evaluation (center crop, no augmentation) and `training=True` for training (random augmentations).

```python
from evals.video_classification_frozen.utils import make_transforms
import torch
import numpy as np

# Evaluation transform (center crop, no augmentation)
eval_transform = make_transforms(
    training=False,
    crop_size=384,
    normalize=((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
)

# Multi-view eval transform (e.g., 3 spatial crops for inference ensembling)
multiview_transform = make_transforms(
    training=False,
    crop_size=256,
    num_views_per_clip=3,
)

# Training transform (random resize-crop + horizontal flip + optional RandAugment)
train_transform = make_transforms(
    training=True,
    random_horizontal_flip=True,
    random_resize_scale=(0.3, 1.0),
    random_resize_aspect_ratio=(0.75, 1.333),
    auto_augment=True,
    motion_shift=False,
    crop_size=256,
)

# Apply to a raw video clip  (T x H x W x C  uint8 numpy array)
frames_np = np.random.randint(0, 255, (16, 360, 640, 3), dtype=np.uint8)
frames_list = [frames_np[i] for i in range(len(frames_np))]

clips = eval_transform(frames_list)      # returns list of tensors
print(clips[0].shape)                    # torch.Size([3, 16, 384, 384])  (C, T, H, W)
```

--------------------------------

### Distributed Pretraining

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Command to launch initial pretraining of a ViT-L model distributedly. Requires specifying time, account, and QoS.

```python
python -m app.main_distributed \
  --fname configs/train/vitl16/pretrain-256px-16f.yaml
  --time 6000
  --account my_account --qos=my_qos
```

--------------------------------

### End-to-End Inference Demo (Python)

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

This Python script demonstrates loading V-JEPA 2.1 models (PyTorch and HuggingFace), processing a video clip, and extracting features for classification. It includes downloading a sample video and comparing features from both model implementations.

```python
# python -m notebooks.vjepa2_demo
import torch, torch.nn.functional as F, json, subprocess, os
import numpy as np
from decord import VideoReader
from transformers import AutoModel, AutoVideoProcessor
from src.models.vision_transformer import vit_giant_xformers_rope
from src.models.attentive_pooler import AttentiveClassifier
import src.datasets.utils.video.transforms as video_transforms
import src.datasets.utils.video.volume_transforms as volume_transforms

MEAN, STD = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
IMG_SIZE = 384

# Download sample video
if not os.path.exists("sample_video.mp4"):
    subprocess.run(["wget",
        "https://huggingface.co/datasets/nateraw/kinetics-mini/resolve/main/val/bowling/-WH-lxmGJVY_000005_000015.mp4",
        "-O", "sample_video.mp4"])

# Load video (64 frames)
vr = VideoReader("sample_video.mp4")
frames = vr.get_batch(np.arange(0, 128, 2)).asnumpy()          # (64, H, W, 3)
video_tensor = torch.from_numpy(frames).permute(0, 3, 1, 2)    # (64, 3, H, W)

# --- PyTorch model ---
pt_transform = video_transforms.Compose([
    video_transforms.Resize(int(256.0 / 224 * IMG_SIZE), interpolation="bilinear"),
    video_transforms.CenterCrop(size=(IMG_SIZE, IMG_SIZE)),
    volume_transforms.ClipToTensor(),
    video_transforms.Normalize(mean=MEAN, std=STD),
])
encoder_pt = vit_giant_xformers_rope(img_size=(IMG_SIZE, IMG_SIZE), num_frames=64).cuda().eval()
ckpt = torch.load("vitg-384.pt", map_location="cpu")
encoder_state = {k.replace("module.","").replace("backbone.",""): v
                 for k, v in ckpt["encoder"].items()}
encoder_pt.load_state_dict(encoder_state, strict=False)

x_pt = pt_transform(video_tensor).cuda().unsqueeze(0)           # (1, 3, 64, 384, 384)
with torch.inference_mode():
    features_pt = encoder_pt(x_pt)                              # (1, N_patches, 1408)

# --- HuggingFace model ---
model_hf = AutoModel.from_pretrained("facebook/vjepa2-vitg-fpc64-384").cuda().eval()
processor_hf = AutoVideoProcessor.from_pretrained("facebook/vjepa2-vitg-fpc64-384")
x_hf = processor_hf(video_tensor, return_tensors="pt")["pixel_values_videos"].cuda()
with torch.inference_mode():
    features_hf = model_hf.get_vision_features(x_hf)           # (1, N_patches, 1408)

print(f"Feature match: {torch.allclose(features_pt, features_hf, atol=1e-3)}")
```

--------------------------------

### Import Libraries and Define Utilities

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/vjepa2_demo.ipynb

Imports necessary libraries for V-JEPA 2 model loading, video processing, and inference. Includes utility functions for loading pretrained weights and building video transforms.

```python
import json
import os
import subprocess

import numpy as np
import torch
import torch.nn.functional as F
from decord import VideoReader
from transformers import AutoVideoProcessor, AutoModel

import src.datasets.utils.video.transforms as video_transforms
import src.datasets.utils.video.volume_transforms as volume_transforms
from src.models.attentive_pooler import AttentiveClassifier
from src.models.vision_transformer import vit_giant_xformers_rope

IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)

def load_pretrained_vjepa_pt_weights(model, pretrained_weights):
    # Load weights of the VJEPA2 encoder
    # The PyTorch state_dict is already preprocessed to have the right key names
    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))


def load_pretrained_vjepa_classifier_weights(model, pretrained_weights):
    # Load weights of the VJEPA2 classifier
    # The PyTorch state_dict is already preprocessed to have the right key names
    pretrained_dict = torch.load(pretrained_weights, weights_only=True, map_location="cpu")["classifiers"][0]
    pretrained_dict = {k.replace("module.", ""): 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))


def build_pt_video_transform(img_size):
    short_side_size = int(256.0 / 224 * img_size)
    # Eval transform has no random cropping nor flip
    eval_transform = video_transforms.Compose(
        [
            video_transforms.Resize(short_side_size, interpolation="bilinear"),
            video_transforms.CenterCrop(size=(img_size, img_size)),
            volume_transforms.ClipToTensor(),
            video_transforms.Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD),
        ]
    )
    return eval_transform


def get_video():
    vr = VideoReader("sample_video.mp4")
    # choosing some frames here, you can define more complex sampling strategy
    frame_idx = np.arange(0, 128, 2)
    video = vr.get_batch(frame_idx).asnumpy()
    return video


def forward_vjepa_video(model_hf, model_pt, hf_transform, pt_transform):
    # Run a sample inference with VJEPA
    with torch.inference_mode():
        # Read and pre-process the image
        video = get_video()  # T x H x W x C
        video = torch.from_numpy(video).permute(0, 3, 1, 2)  # T x C x H x W
        x_pt = pt_transform(video).cuda().unsqueeze(0)
        x_hf = hf_transform(video, return_tensors="pt")["pixel_values_videos"].to("cuda")
        # Extract the patch-wise features from the last layer
        out_patch_features_pt = model_pt(x_pt)
        out_patch_features_hf = model_hf.get_vision_features(x_hf)

    return out_patch_features_hf, out_patch_features_pt


def get_vjepa_video_classification_results(classifier, out_patch_features_pt):
    SOMETHING_SOMETHING_V2_CLASSES = json.load(open("ssv2_classes.json", "r"))

    with torch.inference_mode():
        out_classifier = classifier(out_patch_features_pt)

    print(f"Classifier output shape: {out_classifier.shape}")

    print("Top 5 predicted class names:")
    top5_indices = out_classifier.topk(5).indices[0]
    top5_probs = F.softmax(out_classifier.topk(5).values[0]) * 100.0  # convert to percentage
    for idx, prob in zip(top5_indices, top5_probs):
        str_idx = str(idx.item())
        print(f"{SOMETHING_SOMETHING_V2_CLASSES[str_idx]} ({prob}%)")

    return

```

--------------------------------

### V-JEPA 2.1 Pretraining Configuration (YAML)

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

This YAML configuration file specifies key parameters for V-JEPA 2.1 pretraining, including model architecture, dense predictive loss, and multi-modal data handling.

```yaml
# Key fields in configs/train_2_1/vitg16/pretrain-256px-16f.yaml
app: vjepa_2_1
model:
  model_name: vit_giant_xformers
  pred_depth: 24
  pred_embed_dim: 384
  use_rope: true
  modality_embedding: true       # separate embedding for image vs. video tokens
loss:
  predict_all: true              # Dense Predictive Loss — all tokens contribute
  weight_distance_loss: true     # weight loss by spatial distance from context
data:
  batch_size: 24
  crop_size: 256
  dataset_fpcs: [16]
img_data:                        # Multi-Modal: mix in ImageNet images as single-frame videos
  batch_size: 72
  crop_size: 256
  dataset_fpcs: [1]
optimization:
  epochs: 1000
  lr: 0.0006
  ema: [0.99925, 0.99925]        # EMA coefficient range for target encoder
```

--------------------------------

### Initialize World Model for MPC

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/energy_landscape_example.ipynb

Initializes the WorldModel for MPC planning. Adjust `cem_steps` and `samples` for optimization accuracy. Ensure the device is set correctly (e.g., 'cpu' or 'cuda').

```python
from utils.world_model_wrapper import WorldModel

world_model = WorldModel(
    encoder=encoder,
    predictor=predictor,
    tokens_per_frame=tokens_per_frame,
    transform=transform,
    # Doing very few CEM iterations with very few samples just to run efficiently on CPU...
    # ... increase cem_steps and samples for more accurate optimization of energy landscape
    mpc_args={
        "rollout": 2,
        "samples": 25,
        "topk": 10,
        "cem_steps": 2,
        "momentum_mean": 0.15,
        "momentum_mean_gripper": 0.15,
        "momentum_std": 0.75,
        "momentum_std_gripper": 0.15,
        "maxnorm": 0.075,
        "verbose": True
    },
    normalize_reps=True,
    device="cpu"
)
```

--------------------------------

### Distributed Post-training

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Command to run post-training of the action-conditioned model distributedly. Requires specifying time, account, and QoS.

```python
python -m app.main_distributed \
  --fname configs/train/vitg16/droid-256px-8f.yaml
  --time 6000
  --account my_account --qos=my_qos
```

--------------------------------

### Download Model Checkpoints

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Download necessary model checkpoints for V-JEPA 2 inference. Update the `YOUR_DIR` placeholder with your desired directory.

```bash
wget https://dl.fbaipublicfiles.com/vjepa2/vitg-384.pt -P YOUR_DIR
wget https://dl.fbaipublicfiles.com/vjepa2/evals/ssv2-vitg-384-64x2x3.pt -P YOUR_DIR
```

--------------------------------

### Compare HuggingFace and PyTorch Model Outputs

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/vjepa2_demo.ipynb

Performs inference on a video using both HuggingFace and PyTorch V-JEPA 2 models to extract patch-wise features. Compares the output shapes and values to verify model equivalence.

```python
# Inference on video to get the patch-wise features
out_patch_features_hf, out_patch_features_pt = forward_vjepa_video(
    model_hf, model_pt, hf_transform, pt_video_transform
)

print(
    f"""
    Inference results on video:
    HuggingFace output shape: {out_patch_features_hf.shape}
    PyTorch output shape:     {out_patch_features_pt.shape}
    Absolute difference sum:  {torch.abs(out_patch_features_pt - out_patch_features_hf).sum():.6f}
    Close: {torch.allclose(out_patch_features_pt, out_patch_features_hf, atol=1e-3, rtol=1e-3)}
    """
)
```

--------------------------------

### Simulate DataLoader Batch

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Demonstrates how to create a fake batch for testing the collator function, simulating the output of a DataLoader.

```python
import torch

# Simulate a batch from DataLoader — each sample is (video_tensor, label, clip_indices)
fake_batch = [
    (torch.randn(3, 16, 256, 256), 0, [[list(range(16))]]),
    (torch.randn(3, 16, 256, 256), 1, [[list(range(16))]]),
]

# Assuming 'collator' is defined elsewhere and imported
# collations = collator(fake_batch)
# for collated_batch, masks_enc_list, masks_pred_list in collations:
#     videos = collated_batch[0]           # (B, C, T, H, W)  = (2, 3, 16, 256, 256)
#     for masks_enc, masks_pred in zip(masks_enc_list, masks_pred_list):
#         print(f"Encoder mask: {masks_enc.shape}   Predictor mask: {masks_pred.shape}")
        # e.g. Encoder mask: torch.Size([2, 512])   Predictor mask: torch.Size([2, 256])
```

--------------------------------

### Download SSV2 Classes

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/vjepa2_demo.ipynb

Downloads the Something-Something V2 class labels if they are not already present. Uses `wget` to fetch the JSON file.

```python
ssv2_classes_path = "ssv2_classes.json"
if not os.path.exists(ssv2_classes_path):
    command = [
        "wget",
        "https://huggingface.co/datasets/huggingface/label-files/resolve/d79675f2d50a7b1ecf98923d42c30526a51818e2/"
        "something-something-v2-id2label.json",
        "-O",
        "ssv2_classes.json",
    ]
    subprocess.run(command)
    print("Downloading SSV2 classes")
```

--------------------------------

### Initialize and Load Attentive Classifier

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/vjepa2_demo.ipynb

Initializes an AttentiveClassifier and loads pretrained weights for action classification. Ensure `classifier_model_path` points to the downloaded weights.

```python
# Initialize the classifier
classifier_model_path = "YOUR_ATTENTIVE_PROBE_PATH"
classifier = (
    AttentiveClassifier(embed_dim=model_pt.embed_dim, num_heads=16, depth=4, num_classes=174).cuda().eval()
)
load_pretrained_vjepa_classifier_weights(classifier, classifier_model_path)

# Get classification results
get_vjepa_video_classification_results(classifier, out_patch_features_pt)
```

--------------------------------

### Load V-JEPA 2 Models via PyTorch Hub

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

This snippet demonstrates how to load the V-JEPA 2 preprocessor and various V-JEPA 2 model variants (ViT-Large, ViT-Huge, ViT-Giant) using PyTorch Hub.

```APIDOC
## Load V-JEPA 2 Models via PyTorch Hub

### Description
Load the V-JEPA 2 preprocessor and different V-JEPA 2 model variants using PyTorch Hub.

### Method
`torch.hub.load(repository, model_name)`

### Parameters
- **repository** (string): 'facebookresearch/vjepa2'
- **model_name** (string): The name of the model to load (e.g., 'vjepa2_preprocessor', 'vjepa2_vit_large').

### Request Example
```python
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')
```
```

--------------------------------

### Load Pretrained V-JEPA 2 Models via PyTorch Hub

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Load V-JEPA 2 and V-JEPA 2.1 encoders, predictors, and preprocessors using torch.hub.load. You can skip downloading pretrained weights by setting pretrained=False.

```python
import torch

# --- Preprocessor (returns a VideoTransform callable) ---
processor = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_preprocessor')
# processor(list_of_frames) -> list of tensors [C, T, H, W]

# --- V-JEPA 2 encoders (256px) ---
encoder_l, predictor_l = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_large')
encoder_h, predictor_h = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_huge')
encoder_g, predictor_g = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_giant')
encoder_g384, _        = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_giant_384')

# --- V-JEPA 2.1 encoders (384px, latest) ---
enc_b, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_base_384')     # 80M params
enc_l, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_large_384')    # 300M params
enc_g, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_giant_384')    # 1B params
enc_G, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_gigantic_384') # 2B params

# --- V-JEPA 2-AC action-conditioned model ---
encoder_ac, predictor_ac = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_ac_vit_giant')

# --- Skip downloading pretrained weights (untrained model) ---
encoder_l_fresh, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_large', pretrained=False)
```

--------------------------------

### Load and Prepare Robot Trajectory

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/energy_landscape_example.ipynb

Loads a robot trajectory, optionally reverses it, and converts it into PyTorch tensors.

```python
# Load robot trajectory

play_in_reverse = False  # Use this FLAG to try loading the trajectory backwards, and see how the energy landscape changes

trajectory = np.load("franka_example_traj.npz")
np_clips = trajectory["observations"]
np_states = trajectory["states"]
if play_in_reverse:
    np_clips = trajectory["observations"][:, ::-1].copy()
    np_states = trajectory["states"][:, ::-1].copy()
np_actions = np.expand_dims(poses_to_diff(np_states[0, 0], np_states[0, 1]), axis=(0, 1))

# Convert trajectory to torch tensors
clips = transform(np_clips[0]).unsqueeze(0)
states = torch.tensor(np_states)
actions = torch.tensor(np_actions)
print(f"clips: {clips.shape}; states: {states.shape}; actions: {actions.shape}")
```

--------------------------------

### Local Probe-based Evaluation

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Command to run probe-based evaluation locally. Specify the GPUs to use with `--devices`. This command launches Something-Something v2 video classification.

```python
python -m evals.main --fname configs/eval/vitl16/ssv2.yaml \
  --devices cuda:0 cuda:1
```

--------------------------------

### PyTorch Hub — Load Pretrained Backbones

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Load pretrained V-JEPA 2 and V-JEPA 2.1 encoders, predictors, and preprocessors directly from PyTorch Hub.

```APIDOC
## PyTorch Hub — Load Pretrained Backbones

The `hubconf.py` entrypoint exposes all pretrained encoders and the preprocessor via `torch.hub.load`. For V-JEPA 2 the function returns `(encoder, predictor)` tuples; for the preprocessor it returns a callable video transform pipeline.

```python
import torch

# --- Preprocessor (returns a VideoTransform callable) ---
processor = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_preprocessor')
# processor(list_of_frames) -> list of tensors [C, T, H, W]

# --- V-JEPA 2 encoders (256px) ---
encoder_l, predictor_l = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_large')
encoder_h, predictor_h = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_huge')
encoder_g, predictor_g = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_giant')
encoder_g384, _        = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_giant_384')

# --- V-JEPA 2.1 encoders (384px, latest) ---
enc_b, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_base_384')     # 80M params
enc_l, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_large_384')    # 300M params
enc_g, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_giant_384')    # 1B params
enc_G, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_1_vit_gigantic_384') # 2B params

# --- V-JEPA 2-AC action-conditioned model ---
encoder_ac, predictor_ac = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_ac_vit_giant')

# --- Skip downloading pretrained weights (untrained model) ---
encoder_l_fresh, _ = torch.hub.load('facebookresearch/vjepa2', 'vjepa2_vit_large', pretrained=False)
```
```

--------------------------------

### Load and Use SSv2 Classification Probe

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Loads a pretrained classification probe, infers logits from features, and prints the top 5 predicted classes with their probabilities. Ensure 'ssv2_classes.json' is available.

```python
classifier = AttentiveClassifier(embed_dim=1408, num_heads=16, depth=4, num_classes=174).cuda().eval()
probe_ckpt = torch.load("ssv2-vitg-384-64x2x3.pt", map_location="cpu")
classifier.load_state_dict({k.replace("module.",""): v
                             for k, v in probe_ckpt["classifiers"][0].items()}, strict=False)

with torch.inference_mode():
    logits = classifier(features_pt)                             # (1, 174)

classes = json.load(open("ssv2_classes.json"))
top5_idx = logits.topk(5).indices[0]
top5_pct = F.softmax(logits.topk(5).values[0]) * 100
for idx, pct in zip(top5_idx, top5_pct):
    print(f"  {classes[str(idx.item())]:50s} {pct:.1f}%")
```

--------------------------------

### Print Computed Actions

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/energy_landscape_example.ipynb

Prints the computed actions (x, y, z) obtained from the MPC planning process, formatted to two decimal places.

```python
print(f"Actions returned by planning with CEM (x,y,z) = ({actions[0, 0]:.2f},{actions[0, 1]:.2f} {actions[0, 2]:.2f})")
```

--------------------------------

### Import Libraries

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/energy_landscape_example.ipynb

Imports necessary libraries for numerical operations, plotting, and PyTorch functionalities.

```python
import numpy as np
import matplotlib.pyplot as plt

import torch
from torch.nn import functional as F

from app.vjepa_droid.transforms import make_transforms
from utils.mpc_utils import (
    compute_new_pose,
    poses_to_diff
)
```

--------------------------------

### Load Pretrained V-JEPA 2 Models via HuggingFace

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Load V-JEPA 2 models and processors from HuggingFace repositories. The `get_vision_features` method returns patch-level feature tensors.

```python
from transformers import AutoModel, AutoVideoProcessor
import torch

# Available repos:
# facebook/vjepa2-vitl-fpc64-256
# facebook/vjepa2-vith-fpc64-256
# facebook/vjepa2-vitg-fpc64-256
# facebook/vjepa2-vitg-fpc64-384
hf_repo = "facebook/vjepa2-vitg-fpc64-384"

model = AutoModel.from_pretrained(hf_repo).cuda().eval()
processor = AutoVideoProcessor.from_pretrained(hf_repo)

# Load a raw video tensor  T x C x H x W  (uint8)
import numpy as np
from decord import VideoReader
vr = VideoReader("sample_video.mp4")
frames = vr.get_batch(np.arange(0, 128, 2)).asnumpy()          # (64, H, W, 3)
video_tensor = torch.from_numpy(frames).permute(0, 3, 1, 2)    # (64, 3, H, W)

inputs = processor(video_tensor, return_tensors="pt")           # dict with "pixel_values_videos"
pixel_values = inputs["pixel_values_videos"].to("cuda")

with torch.inference_mode():
    patch_features = model.get_vision_features(pixel_values)    # (1, N_patches, embed_dim)

print(patch_features.shape)  # e.g. torch.Size([1, 1536, 1408]) for ViT-g/16 @ 384px / 64f
```

--------------------------------

### Distributed Probe-based Evaluation

Source: https://github.com/facebookresearch/vjepa2/blob/main/README.md

Command to run probe-based evaluation distributed via SLURM. Requires specifying time, account, and QoS.

```python
python -m evals.main_distributed \
  --fname configs/eval/vitl/ssv2.yaml  \
  --time 8600 \
  --account my_account --qos=my_qos
```

--------------------------------

### Infer Next Action using MPC with CEM

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/energy_landscape_example.ipynb

Infers the next action using the initialized world model and MPC with CEM. This process requires current states, goal states, and the target forward pass. The output actions are converted to a NumPy array.

```python
with torch.no_grad():
    h = forward_target(clips)
    z_n, z_goal = h[:, :tokens_per_frame], h[:, -tokens_per_frame:]
    s_n = states[:, :1]
    print(f"Starting planning using Cross-Entropy Method...")
    actions = world_model.infer_next_action(z_n, s_n, z_goal).cpu().numpy()
```

--------------------------------

### Visualize Trajectory Frames

Source: https://github.com/facebookresearch/vjepa2/blob/main/notebooks/energy_landscape_example.ipynb

Displays the video frames from the loaded robot trajectory.

```python
# Visualize loaded video frames from traj

T = len(np_clips[0])
plt.figure(figsize=(20, 3))
_ = plt.imshow(np.transpose(np_clips[0], (1, 0, 2, 3)).reshape(256, 256 * T, 3))
```

--------------------------------

### Initialize VisionTransformerPredictorAC

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Initializes an action-conditioned predictor with a frozen encoder. Used for autoregressively predicting future video frames in latent space given robot actions.

```python
from src.models.ac_predictor import vit_ac_predictor
from src.models.vision_transformer import vit_giant_xformers
import torch

encoder = vit_giant_xformers(patch_size=16, img_size=(256,256), num_frames=8, use_rope=True).cuda()
ac_predictor = vit_ac_predictor(
    img_size=(256, 256),
    patch_size=16,
    num_frames=8,
    tubelet_size=2,
    embed_dim=encoder.embed_dim,     # 1408
    predictor_embed_dim=1024,
    depth=24,
    num_heads=16,
    action_embed_dim=7,              # 7-DoF robot arm (Franka)
    is_frame_causal=True,
    use_rope=True,
).cuda()

B, T = 2, 4                   # batch, temporal tubelet steps = num_frames // tubelet_size
video = torch.randn(B, 3, 8, 256, 256).cuda()
ctx_tokens = encoder(video)   # (B, T*H*W, 1408)  = (2, 4*16*16, 1408) = (2, 1024, 1408)

actions = torch.randn(B, T, 7).cuda()   # (B, T, 7) joint delta actions
states  = torch.randn(B, T, 7).cuda()   # (B, T, 7) joint states

future_pred = ac_predictor(ctx_tokens, actions, states)   # (B, T*H*W, 1408)
# Used for planning: compute energy E = ||future_pred - target_encoder_output||^2
# and find actions minimizing E via gradient-based MPC
```

--------------------------------

### Configure MaskCollator for Pretraining

Source: https://context7.com/facebookresearch/vjepa2/llms.txt

Initializes MaskCollator with multiple spatiotemporal mask configurations. Supports mixed batches with different frames per clip for JEPA pretraining.

```python
from src.masks.multiseq_multiblock3d import MaskCollator
import torch

# Two mask configs: fine-grained (many small blocks) + coarse (few large blocks)
cfgs_mask = [
    {
        "spatial_scale": (0.15, 0.15),
        "temporal_scale": (1.0, 1.0),
        "aspect_ratio": (0.75, 1.5),
        "num_blocks": 8,
        "max_temporal_keep": 1.0,
        "max_keep": None,
        "full_complement": False,
        "pred_full_complement": False,
        "inv_block": False,
    },
    {
        "spatial_scale": (0.7, 0.7),
        "temporal_scale": (1.0, 1.0),
        "aspect_ratio": (0.75, 1.5),
        "num_blocks": 2,
        "max_temporal_keep": 1.0,
        "max_keep": None,
        "full_complement": False,
        "pred_full_complement": False,
        "inv_block": False,
    },
]

collator = MaskCollator(
    cfgs_mask=cfgs_mask,
    dataset_fpcs=[16],           # frames per clip in your dataset
    crop_size=(256, 256),
    patch_size=(16, 16),
    tubelet_size=2,
)
```