Fix RoPE inv_freq + version-safe backbone reload in from_pretrained (compat transformers 5.x)

modeling_page.py

@@ -97,8 +97,15 @@ class Axial2dRotaryEmbedding(nn.Module):
             raise ValueError(f"`dim` must be a positive multiple of 4, got {dim}.")
         self.dim = dim
         self.axis_dim = dim // 2
-        inv_freq = 1.0 / (base ** (torch.arange(0, self.axis_dim, 2, dtype=torch.float32) / self.axis_dim))
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.base = base
+        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)
+
+    def _compute_inv_freq(self):
+        return 1.0 / (self.base ** (torch.arange(0, self.axis_dim, 2, dtype=torch.float32) / self.axis_dim))
+
+    def reset_inv_freq(self):
+        """Recompute inv_freq (a non-persistent buffer that meta-init in from_pretrained can corrupt)."""
+        self.inv_freq = self._compute_inv_freq()
 
     def _axis_cos_sin(self, coords, *, device, dtype):
         inv_freq = self.inv_freq.to(device=device, dtype=torch.float32)
@@ -444,8 +451,14 @@ class Axial2dCrossRotaryEmbedding(nn.Module):
             raise ValueError(f"dim must be a positive multiple of 4, got {dim}.")
         self.dim = dim
         self.axis_dim = dim // 2
-        inv_freq = 1.0 / (base ** (torch.arange(0, self.axis_dim, 2, dtype=torch.float32) / self.axis_dim))
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.base = base
+        self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False)
+
+    def _compute_inv_freq(self):
+        return 1.0 / (self.base ** (torch.arange(0, self.axis_dim, 2, dtype=torch.float32) / self.axis_dim))
+
+    def reset_inv_freq(self):
+        self.inv_freq = self._compute_inv_freq()
 
     def _axis_cos_sin(self, coords, *, out_dtype):
         if coords.ndim != 2:
@@ -703,11 +716,13 @@ class PaGEBackbone(nn.Module):
 
     @staticmethod
     def _dinov3_has_nested_layer(dinov3_module) -> bool:
-        """True if this transformers version nests the layer stack under an inner `.model`."""
+        """True if this transformers version nests the layer stack under an inner `.model`
+        (transformers >= 5.x). In 4.56.x the layers are flattened onto the DINOv3ViTModel itself."""
         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())
+        # inner's own keys are relative to it: "layer.0.*" when nested, never "embeddings.*".
+        return any(k.startswith("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)
@@ -878,6 +893,74 @@ class PaGEPreTrainedModel(PreTrainedModel):
             nn.init.zeros_(module.bias)
             nn.init.ones_(module.weight)
 
+    # ------------------------------------------------------------------ #
+    #  Version-safe loading                                              #
+    # ------------------------------------------------------------------ #
+    # The DINOv3 backbones are built from transformers' built-in DINOv3ViTModel, whose
+    # internal parameter naming changed between transformers 4.56.x (layers flattened:
+    # `model.layer.N`) and 5.x (layers nested: `model.model.layer.N`). The checkpoints store
+    # one convention; loading under the other leaves the backbone randomly initialized.
+    # `from_pretrained` in transformers >= 5 bypasses `nn.Module._load_state_dict_pre_hook`,
+    # so the remap hook on PaGEBackbone is not invoked by it. We therefore reload the backbone
+    # weights ourselves through `nn.Module.load_state_dict` (which *does* fire the hook).
+    @staticmethod
+    def _collect_safetensors(path_or_repo, **kwargs):
+        """Return the full state_dict from a local dir or a HF repo id."""
+        import os as _os
+        import glob as _glob
+        from safetensors.torch import load_file as _load_file
+        state = {}
+        if _os.path.isdir(path_or_repo):
+            index = _os.path.join(path_or_repo, "model.safetensors.index.json")
+            if _os.path.isfile(index):
+                import json as _json
+                wm = _json.load(open(index))["weight_map"]
+                files = sorted(set(wm.values()))
+            else:
+                files = ["model.safetensors"]
+            for f in files:
+                state.update(_load_file(_os.path.join(path_or_repo, f)))
+        else:
+            from huggingface_hub import hf_hub_download
+            import json as _json
+            repo_id = path_or_repo
+            try:
+                idx_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors.index.json")
+                wm = _json.load(open(idx_path))["weight_map"]
+                files = sorted(set(wm.values()))
+            except Exception:
+                files = ["model.safetensors"]
+            for f in files:
+                p = hf_hub_download(repo_id=repo_id, filename=f)
+                state.update(_load_file(p))
+        return state
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        model = super().from_pretrained(*args, **kwargs)
+        # Reload DINOv3 backbone weights version-safely (the remap hook fires here).
+        try:
+            path_or_repo = args[0] if args else kwargs.get("pretrained_model_name_or_path")
+            full_state = cls._collect_safetensors(path_or_repo)
+            for branch in ("scene_branch_backbone", "head_branch_backbone"):
+                if not hasattr(model, branch):
+                    continue
+                bb = getattr(model, branch)
+                bb_state = {k[len(branch) + 1:]: v for k, v in full_state.items()
+                            if k.startswith(branch + ".")}
+                if bb_state:
+                    bb.load_state_dict(bb_state, strict=False)
+            # Recompute RoPE inv_freq buffers: they are non-persistent and transformers' meta-init
+            # during from_pretrained leaves them as garbage, which would corrupt axial RoPE.
+            for m in model.modules():
+                if hasattr(m, "reset_inv_freq") and callable(m.reset_inv_freq):
+                    m.reset_inv_freq()
+        except Exception as e:  # pragma: no cover
+            import warnings
+            warnings.warn(f"PaGE: version-safe backbone reload skipped ({e!r}). "
+                          "Backbone weights may be random if your transformers version mismatches the checkpoint.")
+        return model
+
 
 class PaGEModel(PaGEPreTrainedModel):
     """