Duplicate from epfl-neuroai/vjepa2-encoder-basic

Co-authored-by: david <dawidtang@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,147 @@
+---
+tags:
+- neuroscience
+- fmri
+- video
+- v-jepa
+- pytorch
+library_name: pytorch
+---
+
+# V-JEPA2 Offline Encoder for Video-Evoked BOLD Responses
+
+This repository contains a PyTorch checkpoint for a basic V-JEPA2-based offline encoder trained to predict video-evoked BOLD responses. The encoder is intended for research workflows involving neural response prediction and neural response-guided visual synthesis.
+
+The checkpoint stores decoder weights and metadata for an offline encoder. This repository includes a custom `transformers.AutoModel` wrapper and does not require the original training codebase.
+
+## Input/Output Contract
+
+The intended input is a short video clip corresponding to the training stimulus duration:
+
+- **Input**: one 3-second RGB video clip, represented as a float tensor shaped `[B, T, C, H, W]` with values in `[0, 1]`.
+- **Output**: one vector of predicted z-scored fMRI beta responses per video, shaped `[B, 20484]`.
+- **Temporal dimension**: the output has no time dimension. Each 3-second video maps to a single predicted response vector.
+
+This makes the encoder suitable for scoring or optimizing short generated videos against static target neural-response patterns.
+
+The video-input path resizes frames to `224 x 224` and applies the ImageNet normalization used by the V-JEPA2 training pipeline. If you pass already-normalized V-JEPA2 inputs, call `model.predict_fmri(video, normalize=False)`.
+
+## Loading
+
+This checkpoint can be loaded with `transformers.AutoModel` and `trust_remote_code=True`.
+
+Example:
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "epfl-neuroai/vjepa2-encoder-basic",
+    trust_remote_code=True,
+)
+model.eval()
+
+# Replace this with a preprocessed 3-second video tensor.
+# Shape: [batch, frames, channels, height, width].
+video = torch.zeros(1, 16, 3, 224, 224)
+
+with torch.no_grad():
+    prediction = model.predict_fmri(video)
+
+print(prediction.shape)  # [1, 20484]
+```
+
+For decoder-only debugging, the model can also run from precomputed V-JEPA2 layer features:
+
+```python
+model = AutoModel.from_pretrained(
+    "epfl-neuroai/vjepa2-encoder-basic",
+    trust_remote_code=True,
+    load_vjepa=False,
+)
+
+features = [
+    torch.zeros(1, decoder.mean.shape[1])
+    for decoder in model.decoders
+]
+
+with torch.no_grad():
+    prediction = model.forward_features(features)
+```
+
+## Data
+
+This checkpoint was trained using data from:
+
+- **BOLD Moments Dataset (BMD)**: whole-brain fMRI responses to short naturalistic videos.
+- **Social interaction video fMRI dataset from Emalie McMahon and collaborators**: fMRI responses to naturalistic two-person social action videos.
+
+This repository does not include the underlying fMRI datasets or stimulus videos.
+
+## Files
+
+- `vjepa2_offline_encoder.pth`: PyTorch checkpoint containing decoder weights, decoding-unit selection metadata, feature-extractor configuration, and registered attributes.
+- `config.json`, `configuration_vjepa2_fmri_encoder.py`, `modeling_vjepa2_fmri_encoder.py`: custom Transformers files for `AutoModel` loading.
+- `requirements.txt`: minimal Python dependencies.
+
+## Backbone Source
+
+The V-JEPA2 backbone weights are shipped in this repository as:
+
+```text
+vitl.pt
+```
+
+The loader uses the V-JEPA2 Torch Hub architecture with `pretrained=False`, then loads the local `vitl.pt` weights directly. This avoids relying on the moving `facebookresearch/vjepa2` Torch Hub checkpoint URL while preserving compatibility with the original decoder features. The decoder checkpoint uses canonical `extractor_config["layer_names"]` metadata.
+
+## Citations
+
+If you use this checkpoint, please cite the V-JEPA/V-JEPA 2 backbone papers and source datasets:
+
+```bibtex
+@article{bardes2024revisiting,
+  title={Revisiting Feature Prediction for Learning Visual Representations from Video},
+  author={Bardes, Adrien and Garrido, Quentin and Ponce, Jean and Chen, Xinlei and Rabbat, Michael and LeCun, Yann and Assran, Mahmoud and Ballas, Nicolas},
+  journal={arXiv preprint arXiv:2404.08471},
+  year={2024}
+}
+
+@article{assran2025vjepa2,
+  title={V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning},
+  author={Assran, Mido 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 others},
+  journal={arXiv preprint arXiv:2506.09985},
+  year={2025}
+}
+
+@article{tang2025diverse,
+  title={Diverse perceptual representations across visual pathways emerge from a single objective},
+  author={Tang, Yingtian and Gokce, Abdulkadir and Al-Karkari, Khaled Jedoui and Yamins, Daniel and Schrimpf, Martin},
+  journal={bioRxiv},
+  pages={2025--07},
+  year={2025},
+  publisher={Cold Spring Harbor Laboratory}
+}
+
+@article{lahner2024modeling,
+  title={Modeling short visual events through the BOLD moments video fMRI dataset and metadata},
+  author={Lahner, Benjamin and Dwivedi, Kshitij and Iamshchinina, Polina and Graumann, Monika and Lascelles, Alex and Roig, Gemma and Gifford, Alessandro Thomas and Pan, Bowen and Jin, SouYoung and Ratan Murty, N Apurva and others},
+  journal={Nature communications},
+  volume={15},
+  number={1},
+  pages={6241},
+  year={2024},
+  publisher={Nature Publishing Group UK London}
+}
+
+@article{mcmahon2023hierarchical,
+  title={Hierarchical organization of social action features along the lateral visual pathway},
+  author={McMahon, Emalie and Bonner, Michael F and Isik, Leyla},
+  journal={Current Biology},
+  volume={33},
+  number={23},
+  pages={5035--5047},
+  year={2023},
+  publisher={Elsevier}
+}
+```

config.json

@@ -0,0 +1,19 @@
+{
+  "model_type": "vjepa2_fmri_encoder",
+  "architectures": [
+    "VJEPA2FMRIEncoderModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_vjepa2_fmri_encoder.VJEPA2FMRIEncoderConfig",
+    "AutoModel": "modeling_vjepa2_fmri_encoder.VJEPA2FMRIEncoderModel"
+  },
+  "checkpoint_filename": "vjepa2_offline_encoder.pth",
+  "output_dim": 20484,
+  "input_duration_seconds": 3.0,
+  "input_format": "video_tensor_b_t_c_h_w",
+  "output_description": "z_scored_fmri_betas_no_time_dimension",
+  "vjepa_size": "large",
+  "load_vjepa": true,
+  "image_size": 224,
+  "normalize_input": true
+}

configuration_vjepa2_fmri_encoder.py

@@ -0,0 +1,35 @@
+"""Transformers config for the V-JEPA2 fMRI encoder."""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class VJEPA2FMRIEncoderConfig(PretrainedConfig):
+    model_type = "vjepa2_fmri_encoder"
+
+    def __init__(
+        self,
+        checkpoint_filename: str = "vjepa2_offline_encoder.pth",
+        output_dim: int = 20484,
+        input_duration_seconds: float = 3.0,
+        input_format: str = "video_tensor_b_t_c_h_w",
+        output_description: str = "z_scored_fmri_betas_no_time_dimension",
+        backbone_filename: str = "vitl.pt",
+        vjepa_size: str = "large",
+        load_vjepa: bool = True,
+        image_size: int = 224,
+        normalize_input: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.checkpoint_filename = checkpoint_filename
+        self.output_dim = int(output_dim)
+        self.input_duration_seconds = float(input_duration_seconds)
+        self.input_format = input_format
+        self.output_description = output_description
+        self.backbone_filename = backbone_filename
+        self.vjepa_size = vjepa_size
+        self.load_vjepa = bool(load_vjepa)
+        self.image_size = int(image_size)
+        self.normalize_input = bool(normalize_input)

modeling_vjepa2_fmri_encoder.py

@@ -0,0 +1,299 @@
+"""Custom AutoModel implementation for a basic V-JEPA2 fMRI encoder."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Iterable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import PreTrainedModel
+
+try:
+    from .configuration_vjepa2_fmri_encoder import VJEPA2FMRIEncoderConfig
+except ImportError:
+    from configuration_vjepa2_fmri_encoder import VJEPA2FMRIEncoderConfig
+
+
+class RidgeDecoder(nn.Module):
+    def __init__(self, state_dict: dict[str, torch.Tensor]) -> None:
+        super().__init__()
+        self.register_buffer("mean", state_dict["steps.1.mean"])
+        self.register_buffer("std", state_dict["steps.1.std"])
+        self.register_buffer("coef", state_dict["steps.2.regressor._coef"])
+        self.register_buffer("intercept", state_dict["steps.2.regressor._intercept"])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.reshape(x.shape[0], -1)
+        x = (x - self.mean.to(device=x.device)) / self.std.to(device=x.device)
+        coef = self.coef.to(device=x.device)
+        x = x.to(dtype=coef.dtype)
+        return x @ coef.T + self.intercept.to(device=x.device)
+
+
+class HookedFeatureExtractor:
+    def __init__(self, layer_names: Iterable[str], ret_type: str = "chw", spatial_pool: int = 14) -> None:
+        self.layer_names = list(layer_names)
+        self.ret_type = ret_type
+        self.spatial_pool = int(spatial_pool)
+
+    @staticmethod
+    def _get_layer(model: nn.Module, layer_name: str) -> nn.Module:
+        layer: object = model
+        for part in layer_name.split("."):
+            layer = layer[int(part)] if part.isdigit() else getattr(layer, part)
+        if not isinstance(layer, nn.Module):
+            raise TypeError(f"{layer_name} did not resolve to a torch module")
+        return layer
+
+    @staticmethod
+    def _unwrap_output(output: Any) -> torch.Tensor:
+        if isinstance(output, (list, tuple)):
+            if len(output) == 0:
+                raise ValueError("Received an empty feature tuple.")
+            output = output[0]
+        if not torch.is_tensor(output):
+            raise TypeError(f"Expected tensor feature output, got {type(output)!r}")
+        return output
+
+    def __call__(self, model: nn.Module, videos: torch.Tensor, **model_kwargs) -> list[torch.Tensor]:
+        outputs: dict[str, torch.Tensor] = {}
+        hooks = [
+            self._get_layer(model, name).register_forward_hook(
+                lambda _module, _inputs, output, name=name: outputs.__setitem__(name, self._unwrap_output(output))
+            )
+            for name in self.layer_names
+        ]
+        try:
+            model(videos, **model_kwargs)
+        finally:
+            for hook in hooks:
+                hook.remove()
+        return [self._process_feature(outputs[name]) for name in self.layer_names]
+
+    def _process_feature(self, feature: torch.Tensor) -> torch.Tensor:
+        batch, tokens, channels = feature.shape
+        feature = feature.reshape(batch, -1, 14, 14, channels).permute(0, 1, 4, 2, 3)
+        if self.spatial_pool > 1:
+            batch, frames, channels, height, width = feature.shape
+            new_height = height // self.spatial_pool
+            new_width = width // self.spatial_pool
+            feature = feature.reshape(
+                batch,
+                frames,
+                channels,
+                new_height,
+                self.spatial_pool,
+                new_width,
+                self.spatial_pool,
+            )
+            feature = feature.permute(0, 1, 2, 3, 5, 4, 6).mean(dim=(-2, -1))
+        if self.ret_type == "chw":
+            return feature.mean(dim=1)
+        if self.ret_type == "tchw":
+            return feature
+        raise ValueError(f"Unsupported ret_type: {self.ret_type}")
+
+
+class LocalVJEPA2Backbone(nn.Module):
+    def __init__(self, size: str, image_size: int, normalize_input: bool, checkpoint_path: str) -> None:
+        super().__init__()
+        self.image_size = int(image_size)
+        self.normalize_input = bool(normalize_input)
+        self.register_buffer("image_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 1, 3, 1, 1))
+        self.register_buffer("image_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 1, 3, 1, 1))
+        hub_name = {
+            "large": "vjepa2_vit_large",
+            "huge": "vjepa2_vit_huge",
+            "giant": "vjepa2_vit_giant",
+        }[size]
+        backbone = torch.hub.load("facebookresearch/vjepa2", hub_name, pretrained=False)
+        backbone, predictor = backbone if isinstance(backbone, (list, tuple)) else (backbone, None)
+        state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+        backbone.load_state_dict(_clean_backbone_key(state_dict["target_encoder"]), strict=False)
+        if predictor is not None and "predictor" in state_dict:
+            predictor.load_state_dict(_clean_backbone_key(state_dict["predictor"]), strict=False)
+        self.backbone = backbone
+
+    def forward(self, videos: torch.Tensor, normalize: bool | None = None) -> torch.Tensor:
+        if videos.ndim != 5:
+            raise ValueError(f"Expected video tensor shaped [B, T, C, H, W], got {tuple(videos.shape)}")
+        if videos.shape[2] != 3:
+            raise ValueError(f"Expected RGB video with 3 channels at dim 2, got {videos.shape[2]}")
+
+        videos = videos.float()
+        batch, frames, channels, height, width = videos.shape
+        if height != self.image_size or width != self.image_size:
+            videos = videos.reshape(batch * frames, channels, height, width)
+            videos = F.interpolate(
+                videos,
+                size=(self.image_size, self.image_size),
+                mode="bilinear",
+                align_corners=False,
+            )
+            videos = videos.reshape(batch, frames, channels, self.image_size, self.image_size)
+
+        normalize = self.normalize_input if normalize is None else bool(normalize)
+        if normalize:
+            videos = (videos - self.image_mean.to(device=videos.device, dtype=videos.dtype)) / self.image_std.to(
+                device=videos.device,
+                dtype=videos.dtype,
+            )
+        return self.backbone(videos.permute(0, 2, 1, 3, 4))
+
+
+def _clean_backbone_key(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    cleaned = {}
+    for key, value in state_dict.items():
+        key = key.replace("module.", "").replace("backbone.", "")
+        cleaned[key] = value
+    return cleaned
+
+
+class VJEPA2FMRIEncoderModel(PreTrainedModel):
+    config_class = VJEPA2FMRIEncoderConfig
+    base_model_prefix = "vjepa2_fmri_encoder"
+    main_input_name = "videos"
+
+    def __init__(self, config: VJEPA2FMRIEncoderConfig) -> None:
+        super().__init__(config)
+        self.decoders = nn.ModuleList()
+        self.register_buffer("decoding_units", torch.empty(0, dtype=torch.long))
+        self.extractor: HookedFeatureExtractor | None = None
+        self.vjepa: LocalVJEPA2Backbone | None = None
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str | os.PathLike[str],
+        *model_args: Any,
+        config: VJEPA2FMRIEncoderConfig | None = None,
+        load_vjepa: bool | None = None,
+        vjepa_size: str | None = None,
+        normalize_input: bool | None = None,
+        **kwargs: Any,
+    ) -> "VJEPA2FMRIEncoderModel":
+        if model_args:
+            raise TypeError("Unexpected positional arguments for VJEPA2FMRIEncoderModel.from_pretrained")
+
+        revision = kwargs.pop("revision", None)
+        token = kwargs.pop("token", None)
+        cache_dir = kwargs.pop("cache_dir", None)
+        local_files_only = kwargs.pop("local_files_only", False)
+        for ignored in ("trust_remote_code", "state_dict", "ignore_mismatched_sizes", "adapter_kwargs", "weights_only"):
+            kwargs.pop(ignored, None)
+        if kwargs:
+            raise TypeError(f"Unsupported keyword argument(s): {', '.join(sorted(kwargs))}")
+
+        if config is None:
+            config = VJEPA2FMRIEncoderConfig.from_pretrained(
+                pretrained_model_name_or_path,
+                revision=revision,
+                token=token,
+                cache_dir=cache_dir,
+                local_files_only=local_files_only,
+            )
+
+        checkpoint_path = cls._resolve_file_path(
+            pretrained_model_name_or_path,
+            filename=config.checkpoint_filename,
+            revision=revision,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+
+        model = cls(config)
+        model.decoders = nn.ModuleList([RidgeDecoder(state_dict) for state_dict in checkpoint["decoders_state_dict"]])
+        model.register_buffer("decoding_units", checkpoint["decoding_units"].long())
+        for name, value in checkpoint.get("registered_attrs", {}).items():
+            if torch.is_tensor(value):
+                model.register_buffer(name, value)
+
+        load_vjepa = config.load_vjepa if load_vjepa is None else bool(load_vjepa)
+        vjepa_size = config.vjepa_size if vjepa_size is None else vjepa_size
+        normalize_input = config.normalize_input if normalize_input is None else bool(normalize_input)
+        if load_vjepa:
+            backbone_path = cls._resolve_file_path(
+                pretrained_model_name_or_path,
+                filename=config.backbone_filename,
+                revision=revision,
+                token=token,
+                cache_dir=cache_dir,
+                local_files_only=local_files_only,
+            )
+            extractor_config = checkpoint["extractor_config"]
+            model.extractor = HookedFeatureExtractor(
+                layer_names=cls._resolve_layer_names(extractor_config),
+                ret_type=extractor_config.get("ret_type", "chw"),
+                spatial_pool=extractor_config.get("spatial_pool", 14),
+            )
+            model.vjepa = LocalVJEPA2Backbone(
+                size=vjepa_size,
+                image_size=config.image_size,
+                normalize_input=normalize_input,
+                checkpoint_path=backbone_path,
+            )
+        model.eval()
+        return model
+
+    @staticmethod
+    def _resolve_file_path(
+        pretrained_model_name_or_path: str | os.PathLike[str],
+        *,
+        filename: str,
+        revision: str | None,
+        token: str | bool | None,
+        cache_dir: str | os.PathLike[str] | None,
+        local_files_only: bool,
+    ) -> str:
+        path = Path(pretrained_model_name_or_path)
+        if path.exists():
+            file_path = path / filename if path.is_dir() else path
+            if not file_path.exists():
+                raise FileNotFoundError(f"Missing file: {file_path}")
+            return str(file_path)
+
+        from huggingface_hub import hf_hub_download
+
+        return hf_hub_download(
+            repo_id=str(pretrained_model_name_or_path),
+            filename=filename,
+            repo_type="model",
+            revision=revision,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+
+    @staticmethod
+    def _resolve_layer_names(extractor_config: dict[str, Any]) -> list[str]:
+        layer_names = extractor_config.get("layer_names")
+        if layer_names is None:
+            layer_names = extractor_config.get("loi")
+        if layer_names is None:
+            raise KeyError("extractor_config must contain `layer_names` or `loi`.")
+        return list(layer_names)
+
+    def forward_features(self, features: list[torch.Tensor]) -> torch.Tensor:
+        if len(features) != len(self.decoders):
+            raise ValueError(f"Expected {len(self.decoders)} feature tensors, got {len(features)}")
+        outputs = [decoder(feature) for decoder, feature in zip(self.decoders, features)]
+        output = torch.stack(outputs, dim=-1)
+        index = self.decoding_units.to(output.device).unsqueeze(0).unsqueeze(-1)
+        index = index.expand(output.shape[0], -1, -1)
+        return output.gather(dim=2, index=index).squeeze(-1)
+
+    def forward(self, videos: torch.Tensor, normalize: bool | None = None) -> torch.Tensor:
+        if self.vjepa is None or self.extractor is None:
+            raise RuntimeError("This model was loaded with load_vjepa=False.")
+        features = self.extractor(self.vjepa, videos, normalize=normalize)
+        return self.forward_features(features)
+
+    def predict_fmri(self, videos: torch.Tensor, normalize: bool | None = None) -> torch.Tensor:
+        """Predict z-scored fMRI beta responses for videos shaped [B, T, C, H, W]."""
+
+        return self(videos, normalize=normalize)

requirements.txt

@@ -0,0 +1,3 @@
+torch
+huggingface_hub
+transformers

vitl.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5346856ec9df69487fe72a25bf2632aaa8112df33fb67708e3f7374edc1f7012
+size 5127726842

vjepa2_offline_encoder.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b2ec1499735098a1c97e67c84837241c349809f278d334f8c0e0c7b5ef1fe3b
+size 2125320301