Add load-time DINOv3 key remap hook (compat transformers 4.56 & 5.6.x)

modeling_page.py

@@ -695,6 +695,44 @@ class PaGEBackbone(nn.Module):
         self.embed_dim = int(dinov3_config.hidden_size)
         # CLS(1) + num_register_tokens
         self._num_front = 1 + int(getattr(dinov3_config, "num_register_tokens", 0) or 0)
+        # DINOv3ViTModel's internal naming differs across transformers versions:
+        #   4.56.x  -> layer stack flattened:  model.layer.N.*        (model.<dinov3 top-level>)
+        #   5.6.x   -> layer stack nested:     model.model.layer.N.*  (extra `.model` wrapper)
+        # A single safetensors file must load under both, so remap incoming keys at load time.
+        self._register_load_state_dict_pre_hook(self._remap_dinov3_keys)
+
+    @staticmethod
+    def _dinov3_has_nested_layer(dinov3_module) -> bool:
+        """True if this transformers version nests the layer stack under an inner `.model`."""
+        inner = getattr(dinov3_module, "model", None)
+        if not isinstance(inner, nn.Module):
+            return False
+        return any(k.startswith("model.layer.") for k in inner.state_dict().keys())
+
+    def _remap_dinov3_keys(self, state_dict, prefix, *args, **kwargs):
+        """Normalize DINOv3 backbone keys (embeddings / layer / norm / rope_embeddings)
+        from whichever convention the checkpoint uses into the one this transformers
+        version expects."""
+        nested = self._dinov3_has_nested_layer(self.model)
+        model_pref = prefix + "model."
+        new = {}
+        for k in list(state_dict.keys()):
+            if not k.startswith(model_pref):
+                continue
+            rest = k[len(model_pref):]  # after "<prefix>model."
+            # rest is one of: "embeddings...", "norm...", "rope_embeddings...", "layer...",
+            #                  or "model.layer..." (nested-conv checkpoint under a flat version, etc.)
+            if rest.startswith("model.layer."):
+                core = rest[len("model."):]  # -> "layer..."
+            else:
+                core = rest  # "layer..." / "embeddings..." / "norm..." / "rope_embeddings..."
+            if nested and core.startswith("layer."):
+                target = model_pref + "model." + core
+            else:
+                target = model_pref + core
+            if target != k:
+                new[target] = state_dict.pop(k)
+        state_dict.update(new)
 
     def _get_patch_tokens(self, x: torch.Tensor) -> torch.Tensor:
         out = self.model(pixel_values=x, return_dict=True)