Add V-JEPA 2.1 inference support
Author: davevanveen
Pull request: yes

# What does this PR do?

Adds inference support and checkpoint conversion for **Meta V-JEPA 2.1** pretrained backbones within the existing `vjepa2` model family.

V-JEPA 2.1 was released by Meta on 2026-03-16 with four pretrained encoders at 384 resolution: ViT-B (80M), ViT-L (300M), ViT-g (1B), and ViT-G (2B). This PR extends the existing `VJEPA2Model` with backward-compatible config fields and modeling changes to support loading these checkpoints.

### Changes

**Config** (`configuration_vjepa2.py`): 8 new fields with backward-compatible defaults that preserve V-JEPA 2.0 behavior:
- `use_rope_interleave`, `use_modality_embeddings`, `interpolate_rope`, `return_all_tokens`, `img_temporal_dim_size`
- Predictor-only: `teacher_embed_dim`, `n_output_distillation`, `hierarchical_layers`

**Modeling** (`modeling_vjepa2.py`):
- Corrected RoPE implementation (`repeat_interleave` vs `repeat`) with config toggle
- Learnable modality embeddings (image/video) for both encoder and predictor, with top-level modality routing
- Separate image patch embedding (`tubelet_size=1`) when `img_temporal_dim_size` is set
- Hierarchical feature extraction with per-layer `norms_block` in encoder
- Smart encoder output: concatenated features when predictor needs them (`n_output_distillation > 1`), single-norm output for `get_vision_features()` / `skip_predictor=True`
- RoPE position interpolation for flexible resolution (encoder only, matches Meta's behavior)
- Multi-layer predictor embed (Linear+GELU+Linear) for `n_output_distillation > 1`
- Predictor context token projection (`proj_context`) for `return_all_tokens`
- Teacher embedding dimension-based output projection sizing
- `VJEPA2ForVideoClassification` guard against unsupported hierarchical configs

**Converter** (`convert_vjepa2_to_hf.py`):
- Four new 2.1 model variants with correct architecture parameters
- Key remappings for all new layers (modality embeddings, norms_block, patch_embed_img, proj_context)
- Checkpoint key handling (`ema_encoder` for distilled models, `target_encoder` for self-supervised)
- Updated test function to handle 2.1 tuple predictor returns

**Tests** (`test_modeling_vjepa2.py`):
- Config defaults test (verifies 2.0 backward compatibility)
- Fast forward pass with shape assertions for `n_output_distillation=1` (distilled path)
- Fast forward pass with shape assertions for `n_output_distillation=4` (multi-layer path)
- All 90 existing tests pass unchanged

**Docs** (`vjepa2.md`): V-JEPA 2.1 section with checkpoint table and architecture notes.

### Verification

Verified end-to-end against Meta's reference implementation on two checkpoints covering both architecture paths:

**ViT-B/384** (80M, `n_output_distillation=1`, distilled):
- All weight keys match (strict load)
- Encoder: max diff **0.0001**
- Predictor: max diff **0.008**

**ViT-g/384** (1B, `n_output_distillation=4`, self-supervised):
- All weight keys match (strict load)
- Encoder: max diff **0.004**
- Predictor target: max diff **0.0002**
- Predictor context: max diff **0.002**

All diffs are within SDPA floating-point precision.

### 2.1 checkpoint summary

| Model | Params | Distilled | `n_output_distillation` | `teacher_embed_dim` |
|-------|--------|-----------|------------------------|---------------------|
| ViT-B/16, 384 | 80M | Yes (ViT-G) | 1 | 1664 |
| ViT-L/16, 384 | 300M | Yes (ViT-G) | 1 | 1664 |
| ViT-g/16, 384 | 1B | No | 4 | — |
| ViT-G/16, 384 | 2B | No | 4 | — |

Fixes https://github.com/huggingface/transformers/issues/45496

## Code Agent Policy

- [x] I confirm that this is not a pure code agent PR.

AI assistance (Claude Code) was used for implementation. I reviewed and verified all changes, ran tests, and validated checkpoint conversion end-to-end on both architecture paths (ViT-B and ViT-g).

## Before submitting

- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? https://github.com/huggingface/transformers/issues/45496
- [x] Did you make sure to update the documentation with your changes?
- [x] Did you write any new necessary tests?

## Who can review?

@yonigozlan — vision models reviewer and original V-JEPA 2 contributor
@molbap — vision models reviewer

Comments:
Hello @davevanveen ! Thanks a lot for opening this PR. Sorry for the delay in reviewing this, we have been very busy with lots of new models to review and recent core changes. I'll try to get to this PR next week!

Appreciate the note, @yonigozlan! No rush from me; I imagine y'all are pretty slammed. Thanks for all the work you do. Cheers.

Thank you for the feedback, @yonigozlan! I appreciate your attention to detail. Happy to iterate further.

Pushed a refactor addressing the API simplification comments. Included inline replies above for changes worth flagging. Summary below: 

- replaced `img_temporal_dim_size` with `use_image_patch_embedder` (boolean) and simplified the single-frame check
- added `tubelet_size` arg on `VJEPA2PatchEmbeddings3D` instead of the config-copy workaround
- replaced predictor-specific `is_predictor` with an explicit `interpolate_rope` flag
- renamed `n_output_distillation` to `num_distillation_outputs` (default 1) and `return_all_tokens` to `use_context_projection`
- added `VJEPA2EncoderOutput.hierarchical_hidden_state` and dropped the `return_hierarchical` argument
- added `VJEPA2PredictorOutput.context_predictions` and stopped overloading `last_hidden_state`
- extended `VJEPA2WithMaskedInputPredictorOutput` and `VJEPA2WithMaskedInputModelOutput` with the new fields, top-level forward chooses predictor input from hierarchical state when present
- simplified distillation naming (`distillation_layers` / `distillation_norms`) and projection-dim logic
- removed the classification-head guard
- updated converter validation to use the new public output fields

Validated the converter on `vit_base_2_1_384` and `vit_giant_2_1_384` — `allclose` passes against Meta's checkpoints for both encoder and predictor.

Hi @yonigozlan - Thanks again for the detailed review! I addressed all requested changes three weeks ago. When you have a chance, I'd appreciate another look. Happy to address any loose ends. Cheers.

Hey @davevanveen ! Me and/or @guarin will follow-up on reviewing the PR soon, sorry for the delay :hugs: 

Hi @molbap, appreciate your review! Great suggestions; all are addressed above. Happy to continue the discussion if anything looks off.

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: vjepa2