Accept loi extractor config alias

modeling_vjepa2_fmri_encoder.py

@@ -38,8 +38,6 @@ class HookedFeatureExtractor:
         self.layer_names = list(layer_names)
         self.ret_type = ret_type
         self.spatial_pool = int(spatial_pool)
-        self.outputs: dict[str, torch.Tensor] = {}
-        self.hooks = []
 
     @staticmethod
     def _get_layer(model: nn.Module, layer_name: str) -> nn.Module:
@@ -50,24 +48,33 @@ class HookedFeatureExtractor:
             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]:
-        self.outputs = {}
-        self.hooks = [
+        outputs: dict[str, torch.Tensor] = {}
+        hooks = [
             self._get_layer(model, name).register_forward_hook(
-                lambda _module, _inputs, output, name=name: self.outputs.__setitem__(name, output)
+                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 self.hooks:
+            for hook in hooks:
                 hook.remove()
-            self.hooks = []
-        return [self._process_feature(self.outputs[name]) for name in self.layer_names]
+        return [self._process_feature(outputs[name]) for name in self.layer_names]
 
     def _process_feature(self, feature: torch.Tensor) -> torch.Tensor:
-        batch, _thw, channels = feature.shape
+        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
@@ -90,8 +97,8 @@ class HookedFeatureExtractor:
         raise ValueError(f"Unsupported ret_type: {self.ret_type}")
 
 
-class VJEPA2Backbone(nn.Module):
-    def __init__(self, size: str, image_size: int, normalize_input: bool) -> None:
+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)
@@ -102,8 +109,13 @@ class VJEPA2Backbone(nn.Module):
             "huge": "vjepa2_vit_huge",
             "giant": "vjepa2_vit_giant",
         }[size]
-        backbone = torch.hub.load("facebookresearch/vjepa2", hub_name, pretrained=True)
-        self.backbone = backbone[0] if isinstance(backbone, (list, tuple)) else backbone
+        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:
@@ -132,6 +144,14 @@ class VJEPA2Backbone(nn.Module):
         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"
@@ -142,7 +162,7 @@ class VJEPA2FMRIEncoderModel(PreTrainedModel):
         self.decoders = nn.ModuleList()
         self.register_buffer("decoding_units", torch.empty(0, dtype=torch.long))
         self.extractor: HookedFeatureExtractor | None = None
-        self.vjepa: VJEPA2Backbone | None = None
+        self.vjepa: LocalVJEPA2Backbone | None = None
 
     @classmethod
     def from_pretrained(
@@ -176,7 +196,7 @@ class VJEPA2FMRIEncoderModel(PreTrainedModel):
                 local_files_only=local_files_only,
             )
 
-        checkpoint_path = cls._resolve_checkpoint_path(
+        checkpoint_path = cls._resolve_file_path(
             pretrained_model_name_or_path,
             filename=config.checkpoint_filename,
             revision=revision,
@@ -197,22 +217,31 @@ class VJEPA2FMRIEncoderModel(PreTrainedModel):
         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=extractor_config["layer_names"],
+                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 = VJEPA2Backbone(
+            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_checkpoint_path(
+    def _resolve_file_path(
         pretrained_model_name_or_path: str | os.PathLike[str],
         *,
         filename: str,
@@ -223,10 +252,10 @@ class VJEPA2FMRIEncoderModel(PreTrainedModel):
     ) -> str:
         path = Path(pretrained_model_name_or_path)
         if path.exists():
-            checkpoint_path = path / filename if path.is_dir() else path
-            if not checkpoint_path.exists():
-                raise FileNotFoundError(f"Missing checkpoint file: {checkpoint_path}")
-            return str(checkpoint_path)
+            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
 
@@ -240,6 +269,15 @@ class VJEPA2FMRIEncoderModel(PreTrainedModel):
             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)}")