Add training and inference source code

2ea2835 verified about 1 month ago

22.5 kB

"""V-JEPA2 sparse and intermediate-dense video encoder adapter.
The historical sparse API remains available.  The dense API additionally
returns the complete post-residual hidden state of an intermediate encoder
block, before V-JEPA2's final LayerNorm and predictor.  It is intended for
feature-distribution losses whose gradient must replay through preprocessing
and a frozen encoder prefix.
"""

from __future__ import annotations

from pathlib import Path

from typing import Dict, Optional, Sequence, Tuple

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

from video_fd.encoders.base_sparse_video_encoder import (
    SparseVideoEncoderAdapter,
    differentiable_resize_crop,
)


class VJEPA2Adapter(SparseVideoEncoderAdapter):
    """Hugging Face V-JEPA2 adapter with an optional memory-pruned dense mode.
    Dense-mode configuration keys:
    * ``dense_layer_index``: zero-based complete encoder block (default 15)
    * ``dense_num_frames``: contiguous input frames (default 64)
    * ``resize_short_side``: differentiable resize short edge (default 292)
    * ``dense_only``: discard predictor, final norm, and blocks after the dense
      target once the checkpoint is loaded
    ``model`` is an injection seam for tiny unit tests.  Production callers
    normally omit it and load ``model_id`` with ``from_pretrained``.
    """

    def __init__(self, config, model: Optional[nn.Module] = None):
        config = dict(config)
        config.setdefault("name", "vjepa2")
        config.setdefault("input_size", config.get("crop_size", 256))
        config.setdefault("input_frames", 16)
        config.setdefault("norm_mean", (0.485, 0.456, 0.406))
        config.setdefault("norm_std", (0.229, 0.224, 0.225))
        super().__init__(config)
        self.model_id = str(config.get("model_id", "assets/models/vjepa2-vitl-fpc64-256"))
        self.allow_temporal_pos_interpolation = bool(config.get("allow_temporal_pos_interpolation", True))
        self.tubelet_size = int(config.get("tubelet_size", 2))
        self.patch_size = int(config.get("patch_size", 16))
        self.dense_layer_index = int(config.get("dense_layer_index", 15))
        self.dense_num_frames = int(config.get("dense_num_frames", 64))
        self.resize_short_side = int(
            config.get("resize_short_side", config.get("dense_resize_short_side", 292))
        )
        self.dense_only = bool(config.get("dense_only", False))
        if self.dense_num_frames <= 0:
            raise ValueError(
                f"dense_num_frames must be positive, got {self.dense_num_frames}"
            )
        if self.resize_short_side < self.input_size:
            raise ValueError(
                "resize_short_side must be >= input_size, got "
                f"{self.resize_short_side} < {self.input_size}"
            )
        self._model = None
        self._dense_pruned = False
        if model is None:
            self._load()
        else:
            self._adopt_model(model)

    def _load(self):
        model_dir = Path(self.model_id).expanduser()
        if not model_dir.is_dir():
            raise FileNotFoundError(f"Prepare the local V-JEPA2 model directory: {model_dir}")
        self.model_id = str(model_dir.resolve())
        try:
            from transformers import VJEPA2Model  # type: ignore
        except Exception as exc:  # pragma: no cover
            raise ImportError(
                "V-JEPA2 requires `transformers` with VJEPA2Model. Original error: %r" % (exc,)
            ) from exc
        model = VJEPA2Model.from_pretrained(self.model_id, local_files_only=True)
        self._adopt_model(model)

    def _adopt_model(self, model: nn.Module) -> None:
        """Install a loaded model, derive its geometry, optionally prune it."""
        cfg = getattr(model, "config", None)
        if cfg is not None:
            self.patch_size = int(getattr(cfg, "patch_size", self.patch_size))
            self.tubelet_size = int(getattr(cfg, "tubelet_size", self.tubelet_size))
            self.input_size = int(getattr(cfg, "crop_size", getattr(cfg, "image_size", self.input_size)))
            cfg_frames = getattr(cfg, "frames_per_clip", None)
            if cfg_frames is not None and int(cfg_frames) != self.input_frames and not self.allow_temporal_pos_interpolation:
                print(f"[vjepa2] checkpoint frames_per_clip={cfg_frames} != {self.input_frames}; "
                      "temporal interpolation disabled -> will fall back to global pooling.")
        self._model = model.to(torch.float32)
        self._validate_dense_layer(self.dense_layer_index)
        if self.dense_only:
            self.prune_for_dense_prefix(self.dense_layer_index)
        self.freeze()

    def _grid_shape(self, n_tokens: int) -> Optional[Tuple[int, int, int]]:
        Tp = self.input_frames // self.tubelet_size
        Hp = self.input_size // self.patch_size
        Wp = self.input_size // self.patch_size
        return (Tp, Hp, Wp) if Tp * Hp * Wp == n_tokens else None

    def _encode_tokens(self, pixel_values: torch.Tensor) -> torch.Tensor:
        if self._dense_pruned:
            raise RuntimeError(
                "forward_sparse is unavailable after prune_for_dense_prefix(); "
                "construct VJEPA2Adapter with dense_only=False to retain it"
            )
        # transformers VJEPA2 expects pixel_values_videos [B, T, 3, H, W]
        # Predictor output is unused by the sparse encoder and can be skipped
        # without changing encoder.last_hidden_state.
        out = self._model(pixel_values_videos=pixel_values, skip_predictor=True)
        tokens = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
        return tokens  # [B, N, D]

    # -- intermediate dense feature extraction -----------------------------
    @property
    def dense_feature_dim(self) -> int:
        cfg = getattr(self._model, "config", None)
        if cfg is None or not hasattr(cfg, "hidden_size"):
            raise RuntimeError("loaded V-JEPA2 model does not expose config.hidden_size")
        return int(cfg.hidden_size)

    @property
    def dense_raw_tokens(self) -> int:
        return (
            (self.dense_num_frames // self.tubelet_size)
            * (self.input_size // self.patch_size)
            * (self.input_size // self.patch_size)
        )

    @property
    def retained_dense_blocks(self) -> int:
        return len(self._model.encoder.layer)

    @staticmethod
    def _integer_index_tensor(
        values: torch.Tensor,
        *,
        name: str,
        device: torch.device,
    ) -> torch.Tensor:
        indices = torch.as_tensor(values)
        if (
            indices.dtype == torch.bool
            or indices.is_floating_point()
            or indices.is_complex()
        ):
            raise TypeError(f"{name} must contain only integers")
        return indices.to(device=device, dtype=torch.long)

    def _temporal_window_indices(
        self,
        video_clip_01: torch.Tensor,
        *,
        temporal_window_starts: Optional[torch.Tensor],
        temporal_window_indices: Optional[torch.Tensor],
    ) -> Optional[torch.Tensor]:
        """Validate and canonicalize optional per-video contiguous windows."""
        if (
            temporal_window_starts is not None
            and temporal_window_indices is not None
        ):
            raise ValueError(
                "provide either temporal_window_starts or "
                "temporal_window_indices, not both"
            )
        if (
            temporal_window_starts is None
            and temporal_window_indices is None
        ):
            return None

        batch_size = int(video_clip_01.shape[0])
        total_frames = int(video_clip_01.shape[1])
        max_start = total_frames - self.dense_num_frames
        device = video_clip_01.device

        if temporal_window_starts is not None:
            starts = self._integer_index_tensor(
                temporal_window_starts,
                name="temporal_window_starts",
                device=device,
            )
            if starts.ndim != 1 or int(starts.shape[0]) != batch_size:
                raise ValueError(
                    "temporal_window_starts must have shape [B], got "
                    f"{tuple(starts.shape)} for B={batch_size}"
                )
            if starts.numel() and (
                int(starts.min()) < 0 or int(starts.max()) > max_start
            ):
                raise ValueError(
                    "temporal_window_starts must be in "
                    f"[0,{max_start}]"
                )
            offsets = torch.arange(
                self.dense_num_frames,
                dtype=torch.long,
                device=device,
            )
            return starts[:, None] + offsets[None, :]

        indices = self._integer_index_tensor(
            temporal_window_indices,
            name="temporal_window_indices",
            device=device,
        )
        expected_shape = (batch_size, self.dense_num_frames)
        if tuple(indices.shape) != expected_shape:
            raise ValueError(
                "temporal_window_indices must have shape "
                f"[B,{self.dense_num_frames}], got {tuple(indices.shape)} "
                f"for B={batch_size}"
            )
        if indices.numel():
            starts = indices[:, 0]
            if int(starts.min()) < 0 or int(starts.max()) > max_start:
                raise ValueError(
                    "temporal_window_indices must describe windows starting "
                    f"in [0,{max_start}]"
                )
            offsets = torch.arange(
                self.dense_num_frames,
                dtype=torch.long,
                device=device,
            )
            expected = starts[:, None] + offsets[None, :]
            if not bool(torch.equal(indices, expected)):
                raise ValueError(
                    "each temporal_window_indices row must be a contiguous "
                    "strictly increasing window"
                )
        return indices

    def preprocess_dense(
        self,
        video_clip_01: torch.Tensor,
        temporal_window_starts: Optional[torch.Tensor] = None,
        temporal_window_indices: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Official differentiable V-JEPA2 preprocessing.
        The input is ``[B,T,3,H,W]`` in ``[0,1]``.  By default a centered
        contiguous ``dense_num_frames`` window is selected.  Callers may
        instead supply explicit integer starts with shape ``[B]`` or explicit
        contiguous indices with shape ``[B,dense_num_frames]``.  Sampling is
        intentionally external: this adapter never consumes ambient RNG
        state.  The short side is then resized to 292 (configurable), a 256
        center crop is taken, and ImageNet normalization is applied.  The
        returned layout remains ``[B,T,3,256,256]``.
        """
        if video_clip_01.ndim != 5:
            raise ValueError(
                "video_clip_01 must have shape [B,T,3,H,W], "
                f"got {tuple(video_clip_01.shape)}"
            )
        if video_clip_01.shape[2] != 3:
            raise ValueError(
                f"video_clip_01 must have 3 channels, got {video_clip_01.shape[2]}"
            )
        total_frames = int(video_clip_01.shape[1])
        if total_frames < self.dense_num_frames:
            raise ValueError(
                f"V-JEPA2 dense mode needs at least {self.dense_num_frames} frames, "
                f"got {total_frames}"
            )
        frame_indices = self._temporal_window_indices(
            video_clip_01,
            temporal_window_starts=temporal_window_starts,
            temporal_window_indices=temporal_window_indices,
        )
        if frame_indices is None:
            start = (total_frames - self.dense_num_frames) // 2
            clip = video_clip_01.narrow(1, start, self.dense_num_frames)
        else:
            gather_indices = frame_indices[:, :, None, None, None].expand(
                -1,
                -1,
                int(video_clip_01.shape[2]),
                int(video_clip_01.shape[3]),
                int(video_clip_01.shape[4]),
            )
            clip = torch.gather(video_clip_01, 1, gather_indices)
        x = differentiable_resize_crop(
            clip,
            self.input_size,
            crop_mode="center",
            resize_short_side=self.resize_short_side,
            interpolation="bilinear",
            round_center_crop=True,
        )
        B, T, C, H, W = x.shape
        flat = x.reshape(B * T, C, H, W)
        flat = (flat - self.norm_mean.to(flat.dtype)) / self.norm_std.to(flat.dtype)
        return flat.reshape(B, T, C, H, W)

    def forward_dense(
        self,
        video_clip_01: torch.Tensor,
        layer_index: Optional[int] = None,
        token_indices: Optional[torch.Tensor] = None,
        gradient_checkpointing: bool = False,
        temporal_window_starts: Optional[torch.Tensor] = None,
        temporal_window_indices: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Return a raw intermediate dense token matrix from an input video.
        No final encoder LayerNorm, predictor, pooling, or flattening is applied.
        Gradients are retained with respect to ``video_clip_01`` even though all
        V-JEPA2 parameters are frozen.
        """
        pixel_values = self.preprocess_dense(
            video_clip_01,
            temporal_window_starts=temporal_window_starts,
            temporal_window_indices=temporal_window_indices,
        )
        return self.encode_dense_preprocessed(
            pixel_values,
            layer_index=layer_index,
            token_indices=token_indices,
            gradient_checkpointing=gradient_checkpointing,
        )

    def forward_dense_layers(
        self,
        video_clip_01: torch.Tensor,
        layer_indices: Sequence[int],
        token_indices: Optional[torch.Tensor] = None,
        gradient_checkpointing: bool = False,
        temporal_window_starts: Optional[torch.Tensor] = None,
        temporal_window_indices: Optional[torch.Tensor] = None,
    ) -> Dict[int, torch.Tensor]:
        """Return several raw intermediate token matrices in one prefix pass.
        The encoder is executed exactly once through ``max(layer_indices)``.
        Each returned value is the complete post-residual output of its
        zero-based block, before V-JEPA2's final LayerNorm.  The dictionary
        follows the caller-provided layer order.
        """
        pixel_values = self.preprocess_dense(
            video_clip_01,
            temporal_window_starts=temporal_window_starts,
            temporal_window_indices=temporal_window_indices,
        )
        return self.encode_dense_preprocessed_layers(
            pixel_values,
            layer_indices=layer_indices,
            token_indices=token_indices,
            gradient_checkpointing=gradient_checkpointing,
        )

    def encode_dense_preprocessed(
        self,
        pixel_values: torch.Tensor,
        layer_index: Optional[int] = None,
        token_indices: Optional[torch.Tensor] = None,
        gradient_checkpointing: bool = False,
    ) -> torch.Tensor:
        """Run embeddings and complete blocks ``0..layer_index``.
        ``pixel_values`` must already have V-JEPA2 layout and normalization.
        This lower-level entrypoint lets replay code avoid preprocessing twice.
        """
        target_layer = self.dense_layer_index if layer_index is None else int(layer_index)
        return self.encode_dense_preprocessed_layers(
            pixel_values,
            layer_indices=(target_layer,),
            token_indices=token_indices,
            gradient_checkpointing=gradient_checkpointing,
        )[target_layer]

    def encode_dense_preprocessed_layers(
        self,
        pixel_values: torch.Tensor,
        layer_indices: Sequence[int],
        token_indices: Optional[torch.Tensor] = None,
        gradient_checkpointing: bool = False,
    ) -> Dict[int, torch.Tensor]:
        """Run one encoder prefix and return all requested block outputs.
        ``pixel_values`` must already have V-JEPA2 layout and normalization.
        Blocks are evaluated only once, through the largest requested
        zero-based layer.  Token selection happens as each requested layer is
        captured so that later blocks do not retain unneeded token outputs.
        """
        if pixel_values.ndim != 5:
            raise ValueError(
                "pixel_values must have shape [B,T,3,H,W], "
                f"got {tuple(pixel_values.shape)}"
            )
        if pixel_values.shape[2] != 3:
            raise ValueError(f"pixel_values must have 3 channels, got {pixel_values.shape[2]}")
        if pixel_values.shape[1] != self.dense_num_frames:
            raise ValueError(
                f"pixel_values must contain {self.dense_num_frames} frames, "
                f"got {pixel_values.shape[1]}"
            )
        if tuple(pixel_values.shape[-2:]) != (self.input_size, self.input_size):
            raise ValueError(
                f"pixel_values must be {self.input_size}x{self.input_size}, "
                f"got {tuple(pixel_values.shape[-2:])}"
            )

        try:
            requested_layers = tuple(int(layer) for layer in layer_indices)
        except TypeError as exc:
            raise TypeError("layer_indices must be a sequence of integers") from exc
        if not requested_layers:
            raise ValueError("layer_indices must not be empty")
        if len(set(requested_layers)) != len(requested_layers):
            raise ValueError("layer_indices must contain unique indices")
        for target_layer in requested_layers:
            self._validate_dense_layer(target_layer)
        max_target_layer = max(requested_layers)
        if max_target_layer >= self.retained_dense_blocks:
            raise ValueError(
                f"layer_index {max_target_layer} was pruned; retained blocks are "
                f"0..{self.retained_dense_blocks - 1}"
            )

        encoder = self._model.encoder
        hidden_states = encoder.embeddings(pixel_values)
        expected_tokens = self.dense_raw_tokens
        if hidden_states.shape[1] != expected_tokens:
            raise RuntimeError(
                "unexpected V-JEPA2 dense token grid: expected "
                f"{expected_tokens}, got {hidden_states.shape[1]}"
            )
        self._raw_dim = int(hidden_states.shape[-1])

        indices = None
        if token_indices is not None:
            indices = torch.as_tensor(token_indices, device=hidden_states.device)
            if indices.ndim != 1:
                raise ValueError(
                    f"token_indices must be 1-D, got shape {tuple(indices.shape)}"
                )
            if indices.numel() == 0:
                raise ValueError("token_indices must not be empty")
            if indices.dtype != torch.long:
                indices = indices.to(torch.long)
            if int(indices.min()) < 0 or int(indices.max()) >= hidden_states.shape[1]:
                raise IndexError(
                    f"token_indices must be in [0,{hidden_states.shape[1] - 1}]"
                )

        requested_set = set(requested_layers)
        captured: Dict[int, torch.Tensor] = {}
        for block_index in range(max_target_layer + 1):
            block = encoder.layer[block_index]

            def run_block(x: torch.Tensor, module=block) -> torch.Tensor:
                return module(
                    x,
                    position_mask=None,
                    head_mask=None,
                    output_attentions=False,
                )[0]

            if (
                gradient_checkpointing
                and torch.is_grad_enabled()
                and hidden_states.requires_grad
            ):
                hidden_states = checkpoint(
                    run_block,
                    hidden_states,
                    use_reentrant=False,
                )
            else:
                hidden_states = run_block(hidden_states)

            if block_index in requested_set:
                # This is deliberately before encoder.layernorm:
                # HF hidden_states[i+1] is the complete post-residual result
                # of zero-based block i.
                captured[block_index] = (
                    hidden_states
                    if indices is None
                    else hidden_states.index_select(1, indices)
                )

        return {layer: captured[layer] for layer in requested_layers}

    def prune_for_dense_prefix(self, layer_index: Optional[int] = None):
        """Irreversibly retain only the modules required by dense-prefix mode."""
        target_layer = self.dense_layer_index if layer_index is None else int(layer_index)
        self._validate_dense_layer(target_layer)
        if self._dense_pruned and target_layer >= len(self._model.encoder.layer):
            raise ValueError(
                f"cannot restore pruned layer {target_layer}; only "
                f"{len(self._model.encoder.layer)} blocks remain"
            )
        self._model.encoder.layer = nn.ModuleList(
            list(self._model.encoder.layer[: target_layer + 1])
        )
        # Dense prefix never executes either module.  Replacing them with None
        # unregisters their parameters immediately, which matters before FSDP.
        self._model.encoder.layernorm = None
        self._model.predictor = None
        self._dense_pruned = True
        self.dense_only = True
        self.dense_layer_index = target_layer
        return self

    def _validate_dense_layer(self, layer_index: int) -> None:
        if layer_index < 0:
            raise ValueError(f"dense layer index must be non-negative, got {layer_index}")
        encoder = getattr(self._model, "encoder", None)
        layers = getattr(encoder, "layer", None)
        if layers is None:
            raise RuntimeError("loaded V-JEPA2 model does not expose encoder.layer")
        configured_layers = int(
            getattr(getattr(self._model, "config", None), "num_hidden_layers", len(layers))
        )
        if layer_index >= configured_layers:
            raise ValueError(
                f"dense layer index {layer_index} out of range for "
                f"{configured_layers} encoder blocks"
            )

	"""V-JEPA2 sparse and intermediate-dense video encoder adapter.

	The historical sparse API remains available. The dense API additionally
	returns the complete post-residual hidden state of an intermediate encoder
	block, before V-JEPA2's final LayerNorm and predictor. It is intended for
	feature-distribution losses whose gradient must replay through preprocessing
	and a frozen encoder prefix.
	"""

	from __future__ import annotations

	from pathlib import Path

	from typing import Dict, Optional, Sequence, Tuple

	import torch
	import torch.nn as nn
	from torch.utils.checkpoint import checkpoint

	from video_fd.encoders.base_sparse_video_encoder import (
	SparseVideoEncoderAdapter,
	differentiable_resize_crop,
	)


	class VJEPA2Adapter(SparseVideoEncoderAdapter):
	"""Hugging Face V-JEPA2 adapter with an optional memory-pruned dense mode.

	Dense-mode configuration keys:

	* ``dense_layer_index``: zero-based complete encoder block (default 15)
	* ``dense_num_frames``: contiguous input frames (default 64)
	* ``resize_short_side``: differentiable resize short edge (default 292)
	* ``dense_only``: discard predictor, final norm, and blocks after the dense
	target once the checkpoint is loaded

	``model`` is an injection seam for tiny unit tests. Production callers
	normally omit it and load ``model_id`` with ``from_pretrained``.
	"""

	def __init__(self, config, model: Optional[nn.Module] = None):
	config = dict(config)
	config.setdefault("name", "vjepa2")
	config.setdefault("input_size", config.get("crop_size", 256))
	config.setdefault("input_frames", 16)
	config.setdefault("norm_mean", (0.485, 0.456, 0.406))
	config.setdefault("norm_std", (0.229, 0.224, 0.225))
	super().__init__(config)
	self.model_id = str(config.get("model_id", "assets/models/vjepa2-vitl-fpc64-256"))
	self.allow_temporal_pos_interpolation = bool(config.get("allow_temporal_pos_interpolation", True))
	self.tubelet_size = int(config.get("tubelet_size", 2))
	self.patch_size = int(config.get("patch_size", 16))
	self.dense_layer_index = int(config.get("dense_layer_index", 15))
	self.dense_num_frames = int(config.get("dense_num_frames", 64))
	self.resize_short_side = int(
	config.get("resize_short_side", config.get("dense_resize_short_side", 292))
	)
	self.dense_only = bool(config.get("dense_only", False))
	if self.dense_num_frames <= 0:
	raise ValueError(
	f"dense_num_frames must be positive, got {self.dense_num_frames}"
	)
	if self.resize_short_side < self.input_size:
	raise ValueError(
	"resize_short_side must be >= input_size, got "
	f"{self.resize_short_side} < {self.input_size}"
	)
	self._model = None
	self._dense_pruned = False
	if model is None:
	self._load()
	else:
	self._adopt_model(model)

	def _load(self):
	model_dir = Path(self.model_id).expanduser()
	if not model_dir.is_dir():
	raise FileNotFoundError(f"Prepare the local V-JEPA2 model directory: {model_dir}")
	self.model_id = str(model_dir.resolve())
	try:
	from transformers import VJEPA2Model # type: ignore
	except Exception as exc: # pragma: no cover
	raise ImportError(
	"V-JEPA2 requires `transformers` with VJEPA2Model. Original error: %r" % (exc,)
	) from exc
	model = VJEPA2Model.from_pretrained(self.model_id, local_files_only=True)
	self._adopt_model(model)

	def _adopt_model(self, model: nn.Module) -> None:
	"""Install a loaded model, derive its geometry, optionally prune it."""
	cfg = getattr(model, "config", None)
	if cfg is not None:
	self.patch_size = int(getattr(cfg, "patch_size", self.patch_size))
	self.tubelet_size = int(getattr(cfg, "tubelet_size", self.tubelet_size))
	self.input_size = int(getattr(cfg, "crop_size", getattr(cfg, "image_size", self.input_size)))
	cfg_frames = getattr(cfg, "frames_per_clip", None)
	if cfg_frames is not None and int(cfg_frames) != self.input_frames and not self.allow_temporal_pos_interpolation:
	print(f"[vjepa2] checkpoint frames_per_clip={cfg_frames} != {self.input_frames}; "
	"temporal interpolation disabled -> will fall back to global pooling.")
	self._model = model.to(torch.float32)
	self._validate_dense_layer(self.dense_layer_index)
	if self.dense_only:
	self.prune_for_dense_prefix(self.dense_layer_index)
	self.freeze()

	def _grid_shape(self, n_tokens: int) -> Optional[Tuple[int, int, int]]:
	Tp = self.input_frames // self.tubelet_size
	Hp = self.input_size // self.patch_size
	Wp = self.input_size // self.patch_size
	return (Tp, Hp, Wp) if Tp * Hp * Wp == n_tokens else None

	def _encode_tokens(self, pixel_values: torch.Tensor) -> torch.Tensor:
	if self._dense_pruned:
	raise RuntimeError(
	"forward_sparse is unavailable after prune_for_dense_prefix(); "
	"construct VJEPA2Adapter with dense_only=False to retain it"
	)
	# transformers VJEPA2 expects pixel_values_videos [B, T, 3, H, W]
	# Predictor output is unused by the sparse encoder and can be skipped
	# without changing encoder.last_hidden_state.
	out = self._model(pixel_values_videos=pixel_values, skip_predictor=True)
	tokens = out.last_hidden_state if hasattr(out, "last_hidden_state") else out[0]
	return tokens # [B, N, D]

	# -- intermediate dense feature extraction -----------------------------
	@property
	def dense_feature_dim(self) -> int:
	cfg = getattr(self._model, "config", None)
	if cfg is None or not hasattr(cfg, "hidden_size"):
	raise RuntimeError("loaded V-JEPA2 model does not expose config.hidden_size")
	return int(cfg.hidden_size)

	@property
	def dense_raw_tokens(self) -> int:
	return (
	(self.dense_num_frames // self.tubelet_size)
	* (self.input_size // self.patch_size)
	* (self.input_size // self.patch_size)
	)

	@property
	def retained_dense_blocks(self) -> int:
	return len(self._model.encoder.layer)

	@staticmethod
	def _integer_index_tensor(
	values: torch.Tensor,
	*,
	name: str,
	device: torch.device,
	) -> torch.Tensor:
	indices = torch.as_tensor(values)
	if (
	indices.dtype == torch.bool
	or indices.is_floating_point()
	or indices.is_complex()
	):
	raise TypeError(f"{name} must contain only integers")
	return indices.to(device=device, dtype=torch.long)

	def _temporal_window_indices(
	self,
	video_clip_01: torch.Tensor,
	*,
	temporal_window_starts: Optional[torch.Tensor],
	temporal_window_indices: Optional[torch.Tensor],
	) -> Optional[torch.Tensor]:
	"""Validate and canonicalize optional per-video contiguous windows."""
	if (
	temporal_window_starts is not None
	and temporal_window_indices is not None
	):
	raise ValueError(
	"provide either temporal_window_starts or "
	"temporal_window_indices, not both"
	)
	if (
	temporal_window_starts is None
	and temporal_window_indices is None
	):
	return None

	batch_size = int(video_clip_01.shape[0])
	total_frames = int(video_clip_01.shape[1])
	max_start = total_frames - self.dense_num_frames
	device = video_clip_01.device

	if temporal_window_starts is not None:
	starts = self._integer_index_tensor(
	temporal_window_starts,
	name="temporal_window_starts",
	device=device,
	)
	if starts.ndim != 1 or int(starts.shape[0]) != batch_size:
	raise ValueError(
	"temporal_window_starts must have shape [B], got "
	f"{tuple(starts.shape)} for B={batch_size}"
	)
	if starts.numel() and (
	int(starts.min()) < 0 or int(starts.max()) > max_start
	):
	raise ValueError(
	"temporal_window_starts must be in "
	f"[0,{max_start}]"
	)
	offsets = torch.arange(
	self.dense_num_frames,
	dtype=torch.long,
	device=device,
	)
	return starts[:, None] + offsets[None, :]

	indices = self._integer_index_tensor(
	temporal_window_indices,
	name="temporal_window_indices",
	device=device,
	)
	expected_shape = (batch_size, self.dense_num_frames)
	if tuple(indices.shape) != expected_shape:
	raise ValueError(
	"temporal_window_indices must have shape "
	f"[B,{self.dense_num_frames}], got {tuple(indices.shape)} "
	f"for B={batch_size}"
	)
	if indices.numel():
	starts = indices[:, 0]
	if int(starts.min()) < 0 or int(starts.max()) > max_start:
	raise ValueError(
	"temporal_window_indices must describe windows starting "
	f"in [0,{max_start}]"
	)
	offsets = torch.arange(
	self.dense_num_frames,
	dtype=torch.long,
	device=device,
	)
	expected = starts[:, None] + offsets[None, :]
	if not bool(torch.equal(indices, expected)):
	raise ValueError(
	"each temporal_window_indices row must be a contiguous "
	"strictly increasing window"
	)
	return indices

	def preprocess_dense(
	self,
	video_clip_01: torch.Tensor,
	temporal_window_starts: Optional[torch.Tensor] = None,
	temporal_window_indices: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""Official differentiable V-JEPA2 preprocessing.

	The input is ``[B,T,3,H,W]`` in ``[0,1]``. By default a centered
	contiguous ``dense_num_frames`` window is selected. Callers may
	instead supply explicit integer starts with shape ``[B]`` or explicit
	contiguous indices with shape ``[B,dense_num_frames]``. Sampling is
	intentionally external: this adapter never consumes ambient RNG
	state. The short side is then resized to 292 (configurable), a 256
	center crop is taken, and ImageNet normalization is applied. The
	returned layout remains ``[B,T,3,256,256]``.
	"""
	if video_clip_01.ndim != 5:
	raise ValueError(
	"video_clip_01 must have shape [B,T,3,H,W], "
	f"got {tuple(video_clip_01.shape)}"
	)
	if video_clip_01.shape[2] != 3:
	raise ValueError(
	f"video_clip_01 must have 3 channels, got {video_clip_01.shape[2]}"
	)
	total_frames = int(video_clip_01.shape[1])
	if total_frames < self.dense_num_frames:
	raise ValueError(
	f"V-JEPA2 dense mode needs at least {self.dense_num_frames} frames, "
	f"got {total_frames}"
	)
	frame_indices = self._temporal_window_indices(
	video_clip_01,
	temporal_window_starts=temporal_window_starts,
	temporal_window_indices=temporal_window_indices,
	)
	if frame_indices is None:
	start = (total_frames - self.dense_num_frames) // 2
	clip = video_clip_01.narrow(1, start, self.dense_num_frames)
	else:
	gather_indices = frame_indices[:, :, None, None, None].expand(
	-1,
	-1,
	int(video_clip_01.shape[2]),
	int(video_clip_01.shape[3]),
	int(video_clip_01.shape[4]),
	)
	clip = torch.gather(video_clip_01, 1, gather_indices)
	x = differentiable_resize_crop(
	clip,
	self.input_size,
	crop_mode="center",
	resize_short_side=self.resize_short_side,
	interpolation="bilinear",
	round_center_crop=True,
	)
	B, T, C, H, W = x.shape
	flat = x.reshape(B * T, C, H, W)
	flat = (flat - self.norm_mean.to(flat.dtype)) / self.norm_std.to(flat.dtype)
	return flat.reshape(B, T, C, H, W)

	def forward_dense(
	self,
	video_clip_01: torch.Tensor,
	layer_index: Optional[int] = None,
	token_indices: Optional[torch.Tensor] = None,
	gradient_checkpointing: bool = False,
	temporal_window_starts: Optional[torch.Tensor] = None,
	temporal_window_indices: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""Return a raw intermediate dense token matrix from an input video.

	No final encoder LayerNorm, predictor, pooling, or flattening is applied.
	Gradients are retained with respect to ``video_clip_01`` even though all
	V-JEPA2 parameters are frozen.
	"""
	pixel_values = self.preprocess_dense(
	video_clip_01,
	temporal_window_starts=temporal_window_starts,
	temporal_window_indices=temporal_window_indices,
	)
	return self.encode_dense_preprocessed(
	pixel_values,
	layer_index=layer_index,
	token_indices=token_indices,
	gradient_checkpointing=gradient_checkpointing,
	)

	def forward_dense_layers(
	self,
	video_clip_01: torch.Tensor,
	layer_indices: Sequence[int],
	token_indices: Optional[torch.Tensor] = None,
	gradient_checkpointing: bool = False,
	temporal_window_starts: Optional[torch.Tensor] = None,
	temporal_window_indices: Optional[torch.Tensor] = None,
	) -> Dict[int, torch.Tensor]:
	"""Return several raw intermediate token matrices in one prefix pass.

	The encoder is executed exactly once through ``max(layer_indices)``.
	Each returned value is the complete post-residual output of its
	zero-based block, before V-JEPA2's final LayerNorm. The dictionary
	follows the caller-provided layer order.
	"""
	pixel_values = self.preprocess_dense(
	video_clip_01,
	temporal_window_starts=temporal_window_starts,
	temporal_window_indices=temporal_window_indices,
	)
	return self.encode_dense_preprocessed_layers(
	pixel_values,
	layer_indices=layer_indices,
	token_indices=token_indices,
	gradient_checkpointing=gradient_checkpointing,
	)

	def encode_dense_preprocessed(
	self,
	pixel_values: torch.Tensor,
	layer_index: Optional[int] = None,
	token_indices: Optional[torch.Tensor] = None,
	gradient_checkpointing: bool = False,
	) -> torch.Tensor:
	"""Run embeddings and complete blocks ``0..layer_index``.

	``pixel_values`` must already have V-JEPA2 layout and normalization.
	This lower-level entrypoint lets replay code avoid preprocessing twice.
	"""
	target_layer = self.dense_layer_index if layer_index is None else int(layer_index)
	return self.encode_dense_preprocessed_layers(
	pixel_values,
	layer_indices=(target_layer,),
	token_indices=token_indices,
	gradient_checkpointing=gradient_checkpointing,
	)[target_layer]

	def encode_dense_preprocessed_layers(
	self,
	pixel_values: torch.Tensor,
	layer_indices: Sequence[int],
	token_indices: Optional[torch.Tensor] = None,
	gradient_checkpointing: bool = False,
	) -> Dict[int, torch.Tensor]:
	"""Run one encoder prefix and return all requested block outputs.

	``pixel_values`` must already have V-JEPA2 layout and normalization.
	Blocks are evaluated only once, through the largest requested
	zero-based layer. Token selection happens as each requested layer is
	captured so that later blocks do not retain unneeded token outputs.
	"""
	if pixel_values.ndim != 5:
	raise ValueError(
	"pixel_values must have shape [B,T,3,H,W], "
	f"got {tuple(pixel_values.shape)}"
	)
	if pixel_values.shape[2] != 3:
	raise ValueError(f"pixel_values must have 3 channels, got {pixel_values.shape[2]}")
	if pixel_values.shape[1] != self.dense_num_frames:
	raise ValueError(
	f"pixel_values must contain {self.dense_num_frames} frames, "
	f"got {pixel_values.shape[1]}"
	)
	if tuple(pixel_values.shape[-2:]) != (self.input_size, self.input_size):
	raise ValueError(
	f"pixel_values must be {self.input_size}x{self.input_size}, "
	f"got {tuple(pixel_values.shape[-2:])}"
	)

	try:
	requested_layers = tuple(int(layer) for layer in layer_indices)
	except TypeError as exc:
	raise TypeError("layer_indices must be a sequence of integers") from exc
	if not requested_layers:
	raise ValueError("layer_indices must not be empty")
	if len(set(requested_layers)) != len(requested_layers):
	raise ValueError("layer_indices must contain unique indices")
	for target_layer in requested_layers:
	self._validate_dense_layer(target_layer)
	max_target_layer = max(requested_layers)
	if max_target_layer >= self.retained_dense_blocks:
	raise ValueError(
	f"layer_index {max_target_layer} was pruned; retained blocks are "
	f"0..{self.retained_dense_blocks - 1}"
	)

	encoder = self._model.encoder
	hidden_states = encoder.embeddings(pixel_values)
	expected_tokens = self.dense_raw_tokens
	if hidden_states.shape[1] != expected_tokens:
	raise RuntimeError(
	"unexpected V-JEPA2 dense token grid: expected "
	f"{expected_tokens}, got {hidden_states.shape[1]}"
	)
	self._raw_dim = int(hidden_states.shape[-1])

	indices = None
	if token_indices is not None:
	indices = torch.as_tensor(token_indices, device=hidden_states.device)
	if indices.ndim != 1:
	raise ValueError(
	f"token_indices must be 1-D, got shape {tuple(indices.shape)}"
	)
	if indices.numel() == 0:
	raise ValueError("token_indices must not be empty")
	if indices.dtype != torch.long:
	indices = indices.to(torch.long)
	if int(indices.min()) < 0 or int(indices.max()) >= hidden_states.shape[1]:
	raise IndexError(
	f"token_indices must be in [0,{hidden_states.shape[1] - 1}]"
	)

	requested_set = set(requested_layers)
	captured: Dict[int, torch.Tensor] = {}
	for block_index in range(max_target_layer + 1):
	block = encoder.layer[block_index]

	def run_block(x: torch.Tensor, module=block) -> torch.Tensor:
	return module(
	x,
	position_mask=None,
	head_mask=None,
	output_attentions=False,
	)[0]

	if (
	gradient_checkpointing
	and torch.is_grad_enabled()
	and hidden_states.requires_grad
	):
	hidden_states = checkpoint(
	run_block,
	hidden_states,
	use_reentrant=False,
	)
	else:
	hidden_states = run_block(hidden_states)

	if block_index in requested_set:
	# This is deliberately before encoder.layernorm:
	# HF hidden_states[i+1] is the complete post-residual result
	# of zero-based block i.
	captured[block_index] = (
	hidden_states
	if indices is None
	else hidden_states.index_select(1, indices)
	)

	return {layer: captured[layer] for layer in requested_layers}

	def prune_for_dense_prefix(self, layer_index: Optional[int] = None):
	"""Irreversibly retain only the modules required by dense-prefix mode."""
	target_layer = self.dense_layer_index if layer_index is None else int(layer_index)
	self._validate_dense_layer(target_layer)
	if self._dense_pruned and target_layer >= len(self._model.encoder.layer):
	raise ValueError(
	f"cannot restore pruned layer {target_layer}; only "
	f"{len(self._model.encoder.layer)} blocks remain"
	)
	self._model.encoder.layer = nn.ModuleList(
	list(self._model.encoder.layer[: target_layer + 1])
	)
	# Dense prefix never executes either module. Replacing them with None
	# unregisters their parameters immediately, which matters before FSDP.
	self._model.encoder.layernorm = None
	self._model.predictor = None
	self._dense_pruned = True
	self.dense_only = True
	self.dense_layer_index = target_layer
	return self

	def _validate_dense_layer(self, layer_index: int) -> None:
	if layer_index < 0:
	raise ValueError(f"dense layer index must be non-negative, got {layer_index}")
	encoder = getattr(self._model, "encoder", None)
	layers = getattr(encoder, "layer", None)
	if layers is None:
	raise RuntimeError("loaded V-JEPA2 model does not expose encoder.layer")
	configured_layers = int(
	getattr(getattr(self._model, "config", None), "num_hidden_layers", len(layers))
	)
	if layer_index >= configured_layers:
	raise ValueError(
	f"dense layer index {layer_index} out of range for "
	f"{configured_layers} encoder blocks"
	)