[Feature] Enable non-timm encoders with explicit embed_dim

stable_pretraining/backbone/vit.py

@@ -232,9 +232,9 @@ class MaskedEncoder(nn.Module):
     Wraps a timm ViT model and adds flexible masking via :class:`PatchMasking`.
     Handles all ViT internals: patch embedding, positional embeddings, prefix
     tokens (CLS, registers), and transformer blocks.
-    :param model_or_model_name: timm model name string or pre-instantiated nn.Module
+    :param encoder_name: timm model name string or pre-instantiated nn.Module
     :param masking: PatchMasking instance. If None, no masking is applied.
-    :param pretrained: Load pretrained weights (only when model_or_model_name is str)
+    :param pretrained: Load pretrained weights (only when encoder_name is str)
     :param img_size: Override default image size
     :param patch_size: Override default patch size (will reinitialize patch_embed)
     :param dynamic_img_size: Enable dynamic image size support with pos_embed interpolation
@@ -243,7 +243,7 @@ class MaskedEncoder(nn.Module):
 
         masking = PatchMasking(mask_ratio=0.75, block_size=4)
         encoder = MaskedEncoder(
-            model_or_model_name="vit_base_patch16_224",
+            encoder_name="vit_base_patch16_224",
             masking=masking,
             pretrained=True,
         )
@@ -256,7 +256,7 @@ class MaskedEncoder(nn.Module):
 
     def __init__(
         self,
-        model_or_model_name: Union[str, nn.Module] = "vit_base_patch16_224",
+        encoder_name: Union[str, nn.Module] = "vit_base_patch16_224",
         masking: Optional[PatchMasking] = None,
         pretrained: bool = False,
         img_size: Optional[Union[int, Tuple[int, int]]] = None,
@@ -268,7 +268,7 @@ def __init__(
         self.dynamic_img_size = dynamic_img_size
         self.masking = masking
         # === Load or use provided encoder ===
-        if isinstance(model_or_model_name, str):
+        if isinstance(encoder_name, str):
             create_kwargs = {
                 "pretrained": pretrained,
                 "num_classes": 0,
@@ -286,15 +286,15 @@ def __init__(
             if norm_layer is not None:
                 create_kwargs["norm_layer"] = norm_layer
 
-            self.vit = timm.create_model(model_or_model_name, **create_kwargs)
+            self.vit = timm.create_model(encoder_name, **create_kwargs)
         else:
             logger.warning(
                 "MaskedEncoder received a pre-instantiated nn.Module. "
                 "Internals assume a timm ViT model with attributes such as "
                 "patch_embed, pos_embed, cls_token, blocks, norm, etc. "
                 "If you pass a non-timm module, unexpected errors may occur."
             )
-            self.vit = model_or_model_name
+            self.vit = encoder_name
             if patch_size is not None:
                 self._rebuild_patch_embed(patch_size, img_size)
             # Remove classification head if present

stable_pretraining/methods/barlow_twins.py

@@ -57,6 +57,9 @@ class BarlowTwins(Module):
         (default ``5.1e-3`` from the paper).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -66,6 +69,7 @@ def __init__(
         lambd: float = 5.1e-3,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -76,11 +80,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = self.backbone.embed_dim
         else:
             self.backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
 
         self.projector = _build_barlow_projector(embed_dim, list(projector_dims))

stable_pretraining/methods/byol.py

@@ -61,6 +61,9 @@ class BYOL(Module):
     :param ema_decay_end: Final EMA decay (default 1.0).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
 
     Note:
         Use :class:`~stable_pretraining.callbacks.TeacherStudentCallback`
@@ -76,6 +79,7 @@ def __init__(
         ema_decay_end: float = 1.0,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -86,11 +90,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = base_backbone.embed_dim
         else:
             base_backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = base_backbone(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
 
         if len(projector_dims) != 2 or len(predictor_dims) != 2:

stable_pretraining/methods/dino.py

@@ -103,6 +103,9 @@ class DINO(Module):
     :param ema_decay_end: Final EMA (default 1.0).
     :param encoder_kwargs: Extra kwargs forwarded to ``timm.create_model``.
     :param pretrained: Load pretrained timm weights for the encoder.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
 
     Example::
 
@@ -130,6 +133,7 @@ def __init__(
         ema_decay_end: float = 1.0,
         encoder_kwargs: Optional[dict] = None,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -139,11 +143,18 @@ def __init__(
             kw = dict(num_classes=0, pretrained=pretrained)
             kw.update(encoder_kwargs or {})
             base_backbone = timm.create_model(encoder_name, **kw)
+            embed_dim = base_backbone.embed_dim
         else:
             base_backbone = encoder_name
-
-        with torch.no_grad():
-            embed_dim = _to_cls(base_backbone(torch.zeros(1, 3, 224, 224))).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
         self.n_prototypes = n_prototypes
 

stable_pretraining/methods/ijepa.py

@@ -16,7 +16,7 @@
 
     # Create model
     model = IJEPA(
-        model_or_model_name="vit_base_patch16_224",
+        encoder_name="vit_base_patch16_224",
         predictor_embed_dim=384,
         predictor_depth=6,
         num_targets=4,
@@ -82,7 +82,7 @@ class IJEPA(Module):
     The context encoder is wrapped with :class:`TeacherStudentWrapper`, enabling
     automatic EMA updates via :class:`TeacherStudentCallback`.
 
-    :param model_or_model_name: timm model name string or pre-instantiated nn.Module
+    :param encoder_name: timm model name string or pre-instantiated nn.Module
     :param predictor_embed_dim: Predictor hidden dimension (default: 384)
     :param predictor_depth: Number of predictor blocks (default: 6)
     :param num_targets: Number of target blocks to sample (default: 4)
@@ -141,7 +141,7 @@ def configure_optimizers(self):
 
     def __init__(
         self,
-        model_or_model_name: Union[str, nn.Module] = "vit_base_patch16_224",
+        encoder_name: Union[str, nn.Module] = "vit_base_patch16_224",
         predictor_embed_dim: int = 384,
         predictor_depth: int = 6,
         num_targets: int = 4,
@@ -156,7 +156,7 @@ def __init__(
 
         # Encoder with EMA wrapper (enables TeacherStudentCallback)
         base_encoder = MaskedEncoder(
-            model_or_model_name,
+            encoder_name,
             masking=None,
             pretrained=pretrained,
         )

stable_pretraining/methods/lejepa.py

@@ -160,14 +160,17 @@ class LeJEPA(Module):
     The SIGReg term is a sliced goodness-of-fit test that pushes
     projected embeddings toward an isotropic Gaussian, averaged over views.
 
-    :param encoder_name: timm model name (e.g., ``"vit_base_patch16_224"``)
+    :param encoder_name: timm model name string or pre-instantiated nn.Module
     :param projector: Optional projection head.  When ``None``, a 3-layer
         BN+ReLU MLP (``embed_dim → 2048 → 2048 → 512``) is created.
     :param n_slices: Random projection directions for the goodness-of-fit test (default: 1024)
     :param t_max: EP integration upper bound (default: 3.0)
     :param n_points: EP quadrature nodes (default: 17)
     :param lamb: SIGReg weight λ (default: 0.02)
     :param pretrained: Load pretrained timm weights
+    :param embed_dim: Output dimension of the custom encoder. Required when
+    ``encoder_name`` is an ``nn.Module``; ignored when a timm model name
+    string is passed (inferred automatically via ``backbone.embed_dim``).
 
     Example::
 
@@ -216,18 +219,21 @@ def __init__(
         lamb: float = 0.02,
         pretrained: bool = False,
         drop_path_rate: float = 0.1,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
-
-        self.backbone = timm.create_model(
-            encoder_name,
-            pretrained=pretrained,
-            num_classes=0,
-            **({"dynamic_img_size": True} if "vit" in encoder_name else {}),
-            drop_path_rate=drop_path_rate,
-        )
-
-        embed_dim = self.backbone.embed_dim
+        if isinstance(encoder_name, str):
+            self.backbone = timm.create_model(
+                encoder_name,
+                pretrained=pretrained,
+                num_classes=0,
+                **({"dynamic_img_size": True} if "vit" in encoder_name else {}),
+                drop_path_rate=drop_path_rate,
+            )
+            embed_dim = self.backbone.embed_dim
+        else:
+            self.backbone = encoder_name
+            embed_dim = embed_dim
 
         if projector is None:
             projector = nn.Sequential(

stable_pretraining/methods/mae.py

@@ -63,7 +63,7 @@ class MAE(Module):
         - **Decoder**: Lightweight transformer reconstructing masked patches
         - **Target**: Normalized pixel values of masked patches
 
-    :param model_or_model_name: timm model name string or pre-instantiated nn.Module
+    :param encoder_name: timm model name string or pre-instantiated nn.Module
     :param decoder_embed_dim: Decoder hidden dimension (default: 512)
     :param decoder_depth: Number of decoder blocks (default: 8)
     :param decoder_num_heads: Decoder attention heads (default: 16)
@@ -105,7 +105,7 @@ def configure_optimizers(self):
 
     def __init__(
         self,
-        model_or_model_name: Union[str, nn.Module] = "vit_base_patch16_224",
+        encoder_name: Union[str, nn.Module] = "vit_base_patch16_224",
         decoder_embed_dim: int = 512,
         decoder_depth: int = 8,
         decoder_num_heads: int = 16,
@@ -124,7 +124,7 @@ def __init__(
         else:
             self.masking = PatchMasking(mask_ratio=mask_ratio, block_size=block_size)
         self.encoder = MaskedEncoder(
-            model_or_model_name, masking=self.masking, pretrained=pretrained
+            encoder_name, masking=self.masking, pretrained=pretrained
         )
 
         embed_dim = self.encoder.embed_dim

stable_pretraining/methods/mocov2.py

@@ -52,6 +52,9 @@ class MoCov2(Module):
     :param ema_decay_end: Final momentum (default 0.999).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -64,6 +67,7 @@ def __init__(
         ema_decay_end: float = 0.999,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
         if isinstance(encoder_name, str):
@@ -73,11 +77,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = base.embed_dim
         else:
             base = encoder_name
-
-        with torch.no_grad():
-            embed_dim = base(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
         self.temperature = temperature
 

stable_pretraining/methods/mocov3.py

@@ -79,6 +79,9 @@ class MoCov3(Module):
     :param ema_decay_end: Final EMA (default 1.0).
     :param low_resolution: Adapt first conv for low-res input.
     :param pretrained: Load pretrained timm weights.
+    :param embed_dim: Backbone output dimension. Inferred automatically for timm
+        models; must be provided explicitly for custom encoders that do not expose
+        an ``.embed_dim`` attribute.
     """
 
     def __init__(
@@ -91,6 +94,7 @@ def __init__(
         ema_decay_end: float = 1.0,
         low_resolution: bool = False,
         pretrained: bool = False,
+        embed_dim: Optional[int] = None,
     ):
         super().__init__()
 
@@ -101,11 +105,18 @@ def __init__(
                 low_resolution=low_resolution,
                 pretrained=pretrained,
             )
+            embed_dim = base.embed_dim
         else:
             base = encoder_name
-
-        with torch.no_grad():
-            embed_dim = base(torch.zeros(1, 3, 224, 224)).shape[-1]
+            if embed_dim is None:
+                if hasattr(encoder_name, "embed_dim"):
+                    embed_dim = encoder_name.embed_dim
+                else:
+                    raise ValueError(
+                        "embed_dim must be provided when the encoder does not expose "
+                        "an .embed_dim attribute. timm models expose this automatically; "
+                        "for custom encoders, pass embed_dim explicitly."
+                    )
         self.embed_dim = embed_dim
         self.temperature = temperature