kaiko-eva 0.2.2__py3-none-any.whl → 0.3.1__py3-none-any.whl

eva/vision/models/modules/semantic_segmentation.py

@@ -1,4 +1,4 @@
-""""Neural Network Semantic Segmentation Module."""
+"""Neural Network Semantic Segmentation Module."""
 
 import functools
 from typing import Any, Callable, Dict, Iterable, List, Tuple
@@ -12,7 +12,7 @@ from torch.optim import lr_scheduler
 from typing_extensions import override
 
 from eva.core.metrics import structs as metrics_lib
-from eva.core.models.modules import module
+from eva.core.models.modules import SchedulerConfiguration, module
 from eva.core.models.modules.typings import INPUT_BATCH, INPUT_TENSOR_BATCH
 from eva.core.models.modules.utils import batch_postprocess, grad, submodule_state_dict
 from eva.core.utils import parser
@@ -32,7 +32,7 @@ class SemanticSegmentationModule(module.ModelModule):
         lr_multiplier_encoder: float = 0.0,
         inferer: Inferer | None = None,
         optimizer: OptimizerCallable = optim.AdamW,
-        lr_scheduler: LRSchedulerCallable = lr_scheduler.ConstantLR,
+        lr_scheduler: LRSchedulerCallable | SchedulerConfiguration = lr_scheduler.ConstantLR,
         metrics: metrics_lib.MetricsSchema | None = None,
         postprocess: batch_postprocess.BatchPostProcess | None = None,
         save_decoder_only: bool = True,
@@ -116,7 +116,7 @@ class SemanticSegmentationModule(module.ModelModule):
     def forward(
         self,
         tensor: torch.Tensor,
-        to_size: Tuple[int, int],
+        to_size: Tuple[int, ...],
         *args: Any,
         **kwargs: Any,
     ) -> torch.Tensor:
@@ -174,7 +174,8 @@ class SemanticSegmentationModule(module.ModelModule):
             The batch step output.
         """
         data, targets, metadata = INPUT_TENSOR_BATCH(*batch)
-        predictions = self(data, to_size=targets.shape[-self.spatial_dims :])
+        to_size = targets.shape[-self.spatial_dims :] if self.inferer is None else None
+        predictions = self(data, to_size=to_size)
         loss = self.criterion(predictions, targets)
         return {
             "loss": loss,
@@ -183,11 +184,21 @@ class SemanticSegmentationModule(module.ModelModule):
             "metadata": metadata,
         }
 
-    def _forward_networks(self, tensor: torch.Tensor, to_size: Tuple[int, int]) -> torch.Tensor:
+    def _forward_networks(
+        self, tensor: torch.Tensor, to_size: Tuple[int, ...] | None = None
+    ) -> torch.Tensor:
         """Passes the input tensor through the encoder and decoder."""
-        features = self.encoder(tensor) if self.encoder else tensor
+        if self.encoder:
+            to_size = to_size or tuple(tensor.shape[-self.spatial_dims :])
+            features = self.encoder(tensor)
+        else:
+            if to_size is None:
+                raise ValueError("`to_size` must be provided when no encoder is used.")
+            features = tensor
+
         if isinstance(self.decoder, segmentation.Decoder):
             if not isinstance(features, list):
                 raise ValueError(f"Expected a list of feature map tensors, got {type(features)}.")
             return self.decoder(DecoderInputs(features, to_size, tensor))
+
         return self.decoder(features)

eva/vision/models/networks/backbones/__init__.py

@@ -1,13 +1,12 @@
 """Vision Model Backbones API."""
 
 from eva.vision.models.networks.backbones import pathology, radiology, timm, universal
-from eva.vision.models.networks.backbones.registry import BackboneModelRegistry, register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 __all__ = [
     "radiology",
     "pathology",
     "timm",
     "universal",
-    "BackboneModelRegistry",
-    "register_model",
+    "backbone_registry",
 ]

eva/vision/models/networks/backbones/_utils.py

@@ -36,7 +36,7 @@ def load_hugingface_model(
 
     return models.HuggingFaceModel(
         model_name_or_path=model_name,
-        tensor_transforms=tensor_transforms,
+        transforms=tensor_transforms,
         model_kwargs=model_kwargs,
     )
 

eva/vision/models/networks/backbones/pathology/bioptimus.py

@@ -8,10 +8,10 @@ from torch import nn
 from eva.core.models import transforms
 from eva.vision.models import wrappers
 from eva.vision.models.networks.backbones import _utils
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-@register_model("pathology/bioptimus_h_optimus_0")
+@backbone_registry.register("pathology/bioptimus_h_optimus_0")
 def bioptimus_h_optimus_0(
     dynamic_img_size: bool = True,
     out_indices: int | Tuple[int, ...] | None = None,
@@ -39,7 +39,7 @@ def bioptimus_h_optimus_0(
     )
 
 
-@register_model("pathology/bioptimus_h0_mini")
+@backbone_registry.register("pathology/bioptimus_h0_mini")
 def bioptimus_h0_mini(
     dynamic_img_size: bool = True,
     out_indices: int | Tuple[int, ...] | None = None,
@@ -72,7 +72,7 @@ def bioptimus_h0_mini(
             "mlp_layer": timm.layers.SwiGLUPacked,
             "act_layer": nn.SiLU,
         },
-        tensor_transforms=(
+        transforms=(
             transforms.ExtractCLSFeatures(include_patch_tokens=include_patch_tokens)
             if out_indices is None
             else None

eva/vision/models/networks/backbones/pathology/gigapath.py

@@ -5,10 +5,10 @@ from typing import Tuple
 import timm
 from torch import nn
 
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-@register_model("pathology/prov_gigapath")
+@backbone_registry.register("pathology/prov_gigapath")
 def prov_gigapath(
     dynamic_img_size: bool = True,
     out_indices: int | Tuple[int, ...] | None = None,

eva/vision/models/networks/backbones/pathology/histai.py

@@ -5,10 +5,10 @@ from typing import Tuple
 from torch import nn
 
 from eva.vision.models.networks.backbones import _utils
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-@register_model("pathology/histai_hibou_b")
+@backbone_registry.register("pathology/histai_hibou_b")
 def histai_hibou_b(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
     """Initializes the hibou-B pathology FM by hist.ai (https://huggingface.co/histai/hibou-B).
 
@@ -30,7 +30,7 @@ def histai_hibou_b(out_indices: int | Tuple[int, ...] | None = None) -> nn.Modul
     )
 
 
-@register_model("pathology/histai_hibou_l")
+@backbone_registry.register("pathology/histai_hibou_l")
 def histai_hibou_l(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
     """Initializes the hibou-L pathology FM by hist.ai (https://huggingface.co/histai/hibou-L).
 

eva/vision/models/networks/backbones/pathology/hkust.py

@@ -7,10 +7,10 @@ import timm
 from torch import nn
 
 from eva.core.models.wrappers import _utils
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-@register_model("pathology/hkust_gpfm")
+@backbone_registry.register("pathology/hkust_gpfm")
 def hkust_gpfm(
     dynamic_img_size: bool = True,
     out_indices: int | Tuple[int, ...] | None = None,

eva/vision/models/networks/backbones/pathology/kaiko.py

@@ -6,10 +6,10 @@ import torch
 from torch import nn
 
 from eva.vision.models.networks.backbones import _utils
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-@register_model("pathology/kaiko_midnight_12k")
+@backbone_registry.register("pathology/kaiko_midnight_12k")
 def kaiko_midnight_12k(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
     """Initializes the Midnight-12k pathology FM by kaiko.ai.
 
@@ -26,7 +26,7 @@ def kaiko_midnight_12k(out_indices: int | Tuple[int, ...] | None = None) -> nn.M
     )
 
 
-@register_model("pathology/kaiko_vits16")
+@backbone_registry.register("pathology/kaiko_vits16")
 def kaiko_vits16(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
 ) -> nn.Module:
@@ -49,7 +49,7 @@ def kaiko_vits16(
     )
 
 
-@register_model("pathology/kaiko_vits8")
+@backbone_registry.register("pathology/kaiko_vits8")
 def kaiko_vits8(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
 ) -> nn.Module:
@@ -72,7 +72,7 @@ def kaiko_vits8(
     )
 
 
-@register_model("pathology/kaiko_vitb16")
+@backbone_registry.register("pathology/kaiko_vitb16")
 def kaiko_vitb16(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
 ) -> nn.Module:
@@ -95,7 +95,7 @@ def kaiko_vitb16(
     )
 
 
-@register_model("pathology/kaiko_vitb8")
+@backbone_registry.register("pathology/kaiko_vitb8")
 def kaiko_vitb8(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
 ) -> nn.Module:
@@ -118,7 +118,7 @@ def kaiko_vitb8(
     )
 
 
-@register_model("pathology/kaiko_vitl14")
+@backbone_registry.register("pathology/kaiko_vitl14")
 def kaiko_vitl14(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
 ) -> nn.Module:

eva/vision/models/networks/backbones/pathology/lunit.py

@@ -13,14 +13,14 @@ from typing import Tuple
 from torch import nn
 
 from eva.vision.models import wrappers
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 VITS_URL_PREFIX = (
     "https://github.com/lunit-io/benchmark-ssl-pathology/releases/download/pretrained-weights"
 )
 
 
-@register_model("pathology/lunit_vits16")
+@backbone_registry.register("pathology/lunit_vits16")
 def lunit_vits16(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
 ) -> nn.Module:
@@ -44,7 +44,7 @@ def lunit_vits16(
     )
 
 
-@register_model("pathology/lunit_vits8")
+@backbone_registry.register("pathology/lunit_vits8")
 def lunit_vits8(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
 ) -> nn.Module:

eva/vision/models/networks/backbones/pathology/mahmood.py

@@ -8,10 +8,10 @@ from torch import nn
 
 from eva.vision.models import wrappers
 from eva.vision.models.networks.backbones import _utils
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-@register_model("pathology/mahmood_uni")
+@backbone_registry.register("pathology/mahmood_uni")
 def mahmood_uni(
     dynamic_img_size: bool = True,
     out_indices: int | Tuple[int, ...] | None = None,
@@ -41,7 +41,7 @@ def mahmood_uni(
     )
 
 
-@register_model("pathology/mahmood_uni2_h")
+@backbone_registry.register("pathology/mahmood_uni2_h")
 def mahmood_uni2_h(
     dynamic_img_size: bool = True,
     out_indices: int | Tuple[int, ...] | None = None,

eva/vision/models/networks/backbones/pathology/owkin.py

@@ -5,10 +5,10 @@ from typing import Tuple
 from torch import nn
 
 from eva.vision.models.networks.backbones import _utils
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-@register_model("pathology/owkin_phikon")
+@backbone_registry.register("pathology/owkin_phikon")
 def owkin_phikon(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
     """Initializes the phikon pathology FM by owkin (https://huggingface.co/owkin/phikon).
 
@@ -22,7 +22,7 @@ def owkin_phikon(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
     return _utils.load_hugingface_model(model_name="owkin/phikon", out_indices=out_indices)
 
 
-@register_model("pathology/owkin_phikon_v2")
+@backbone_registry.register("pathology/owkin_phikon_v2")
 def owkin_phikon_v2(out_indices: int | Tuple[int, ...] | None = None) -> nn.Module:
     """Initializes the phikon-v2 pathology FM by owkin (https://huggingface.co/owkin/phikon-v2).
 

eva/vision/models/networks/backbones/pathology/paige.py

@@ -11,10 +11,10 @@ import torch.nn as nn
 from eva.core.models import transforms
 from eva.vision.models import wrappers
 from eva.vision.models.networks.backbones import _utils
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-@register_model("pathology/paige_virchow2")
+@backbone_registry.register("pathology/paige_virchow2")
 def paige_virchow2(
     dynamic_img_size: bool = True,
     out_indices: int | Tuple[int, ...] | None = None,
@@ -43,7 +43,7 @@ def paige_virchow2(
             "mlp_layer": timm.layers.SwiGLUPacked,
             "act_layer": nn.SiLU,
         },
-        tensor_transforms=(
+        transforms=(
             transforms.ExtractCLSFeatures(include_patch_tokens=include_patch_tokens)
             if out_indices is None
             else None

eva/vision/models/networks/backbones/radiology/swin_unetr.py

@@ -9,10 +9,10 @@ from monai.networks.nets import swin_unetr
 from monai.utils import misc
 from torch import nn
 
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-@register_model("radiology/swin_unetr_encoder")
+@backbone_registry.register("radiology/swin_unetr_encoder")
 class SwinUNETREncoder(nn.Module):
     """Swin transformer encoder based on UNETR [0].
 

eva/vision/models/networks/backbones/radiology/voco.py

@@ -4,10 +4,10 @@ from typing_extensions import override
 
 from eva.core.models.wrappers import _utils
 from eva.vision.models.networks.backbones.radiology import swin_unetr
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-class _VoCo(swin_unetr.SwinUNETREncoder):
+class _VoCo(swin_unetr.SwinUNETREncoder):  # type: ignore
     """Base class for the VoCo self-supervised encoders."""
 
     _checkpoint: str
@@ -39,7 +39,7 @@ class _VoCo(swin_unetr.SwinUNETREncoder):
         self.load_state_dict(state_dict)
 
 
-@register_model("radiology/voco_b")
+@backbone_registry.register("radiology/voco_b")
 class VoCoB(_VoCo):
     """VoCo Self-supervised pre-trained B model."""
 
@@ -51,7 +51,7 @@ class VoCoB(_VoCo):
         super().__init__(feature_size=48, out_indices=out_indices)
 
 
-@register_model("radiology/voco_l")
+@backbone_registry.register("radiology/voco_l")
 class VoCoL(_VoCo):
     """VoCo Self-supervised pre-trained L model."""
 
@@ -63,7 +63,7 @@ class VoCoL(_VoCo):
         super().__init__(feature_size=96, out_indices=out_indices)
 
 
-@register_model("radiology/voco_h")
+@backbone_registry.register("radiology/voco_h")
 class VoCoH(_VoCo):
     """VoCo Self-supervised pre-trained H model."""
 

eva/vision/models/networks/backbones/registry.py

@@ -1,47 +1,5 @@
 """Backbone Model Registry."""
 
-from typing import Any, Callable, Dict, List
+from eva.core.utils.registry import Registry
 
-import torch.nn as nn
-
-
-class BackboneModelRegistry:
-    """A model registry for accessing backbone models by name."""
-
-    _registry: Dict[str, Callable[..., nn.Module]] = {}
-
-    @classmethod
-    def register(cls, name: str) -> Callable:
-        """Decorator to register a new model."""
-
-        def decorator(model_fn: Callable[..., nn.Module]) -> Callable[..., nn.Module]:
-            if name in cls._registry:
-                raise ValueError(f"Model {name} is already registered.")
-            cls._registry[name] = model_fn
-            return model_fn
-
-        return decorator
-
-    @classmethod
-    def get(cls, model_name: str) -> Callable[..., nn.Module]:
-        """Gets a model function from the registry."""
-        if model_name not in cls._registry:
-            raise ValueError(f"Model {model_name} not found in the registry.")
-        return cls._registry[model_name]
-
-    @classmethod
-    def load_model(cls, model_name: str, model_kwargs: Dict[str, Any] | None = None) -> nn.Module:
-        """Loads & initializes a model class from the registry."""
-        model_fn = cls.get(model_name)
-        return model_fn(**(model_kwargs or {}))
-
-    @classmethod
-    def list_models(cls) -> List[str]:
-        """List all models in the registry."""
-        register_models = [name for name in cls._registry.keys() if not name.startswith("timm")]
-        return register_models + ["timm/<model_name>"]
-
-
-def register_model(name: str) -> Callable:
-    """Simple decorator to register a model."""
-    return BackboneModelRegistry.register(name)
+backbone_registry = Registry()

eva/vision/models/networks/backbones/timm/backbones.py

@@ -8,7 +8,7 @@ from loguru import logger
 from torch import nn
 
 from eva.vision.models import wrappers
-from eva.vision.models.networks.backbones.registry import BackboneModelRegistry
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
 def timm_model(
@@ -46,7 +46,7 @@ def timm_model(
     )
 
 
-BackboneModelRegistry._registry.update(
+backbone_registry._registry.update(
     {
         f"timm/{model_name}": functools.partial(timm_model, model_name=model_name)
         for model_name in timm.list_models()

eva/vision/models/networks/backbones/universal/__init__.py

@@ -1,8 +1,15 @@
 """Universal Vision Model Backbones API."""
 
 from eva.vision.models.networks.backbones.universal.vit import (
+    vit_base_patch16_224_dino_1chan,
     vit_small_patch16_224_dino,
+    vit_small_patch16_224_dino_1chan,
     vit_small_patch16_224_random,
 )
 
-__all__ = ["vit_small_patch16_224_dino", "vit_small_patch16_224_random"]
+__all__ = [
+    "vit_small_patch16_224_dino",
+    "vit_small_patch16_224_random",
+    "vit_small_patch16_224_dino_1chan",
+    "vit_base_patch16_224_dino_1chan",
+]

eva/vision/models/networks/backbones/universal/vit.py

@@ -5,10 +5,10 @@ from typing import Tuple
 import timm
 from torch import nn
 
-from eva.vision.models.networks.backbones.registry import register_model
+from eva.vision.models.networks.backbones.registry import backbone_registry
 
 
-@register_model("universal/vit_small_patch16_224_random")
+@backbone_registry.register("universal/vit_small_patch16_224_random")
 def vit_small_patch16_224_random(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
 ) -> nn.Module:
@@ -31,7 +31,7 @@ def vit_small_patch16_224_random(
     )
 
 
-@register_model("universal/vit_small_patch16_224_dino")
+@backbone_registry.register("universal/vit_small_patch16_224_dino")
 def vit_small_patch16_224_dino(
     dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
 ) -> nn.Module:
@@ -52,3 +52,53 @@ def vit_small_patch16_224_dino(
         out_indices=out_indices,
         dynamic_img_size=dynamic_img_size,
     )
+
+
+@backbone_registry.register("universal/vit_small_patch16_224_dino_1chan")
+def vit_small_patch16_224_dino_1chan(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes a ViTS-16 baseline model pretrained w/ DINO for single-channel images.
+
+    Args:
+        dynamic_img_size: Support different input image sizes by allowing to change
+            the grid size (interpolate abs and/or ROPE pos) in the forward pass.
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The torch ViTS-16 based foundation model.
+    """
+    return timm.create_model(
+        model_name="vit_small_patch16_224.dino",
+        in_chans=1,
+        num_classes=0,
+        pretrained=True,
+        features_only=out_indices is not None,
+        out_indices=out_indices,
+        dynamic_img_size=dynamic_img_size,
+    )
+
+
+@backbone_registry.register("universal/vit_base_patch16_224_dino_1chan")
+def vit_base_patch16_224_dino_1chan(
+    dynamic_img_size: bool = True, out_indices: int | Tuple[int, ...] | None = None
+) -> nn.Module:
+    """Initializes a ViTB-16 baseline model pretrained w/ DINO for single-channel images.
+
+    Args:
+        dynamic_img_size: Support different input image sizes by allowing to change
+            the grid size (interpolate abs and/or ROPE pos) in the forward pass.
+        out_indices: Whether and which multi-level patch embeddings to return.
+
+    Returns:
+        The torch ViTB-16 based foundation model.
+    """
+    return timm.create_model(
+        model_name="vit_base_patch16_224.dino",
+        in_chans=1,
+        num_classes=0,
+        pretrained=True,
+        features_only=out_indices is not None,
+        out_indices=out_indices,
+        dynamic_img_size=dynamic_img_size,
+    )

eva/vision/models/networks/decoders/segmentation/decoder2d.py

@@ -115,4 +115,4 @@ class Decoder2D(base.Decoder):
         if self._combine_features:
             features = self._forward_features(features)
         logits = self._forward_head(features)
-        return self._upscale(logits, image_size)
+        return self._upscale(logits, image_size)  # type: ignore

eva/vision/models/networks/decoders/segmentation/linear.py

@@ -117,4 +117,4 @@ class LinearDecoder(base.Decoder):
         """
         patch_embeddings = self._forward_features(decoder_inputs.features)
         logits = self._forward_head(patch_embeddings)
-        return self._cls_seg(logits, decoder_inputs.image_size)
+        return self._cls_seg(logits, decoder_inputs.image_size)  # type: ignore

eva/vision/models/networks/decoders/segmentation/semantic/common.py

@@ -1,7 +1,7 @@
 """Common semantic segmentation decoders.
 
-This module contains implementations of different types of decoder models 
-used in semantic segmentation. These decoders convert the high-level features 
+This module contains implementations of different types of decoder models
+used in semantic segmentation. These decoders convert the high-level features
 output by an encoder into pixel-wise predictions for segmentation tasks.
 """
 

eva/vision/models/networks/decoders/segmentation/typings.py

@@ -11,7 +11,7 @@ class DecoderInputs(NamedTuple):
     features: List[torch.Tensor]
     """List of image features generated by the encoder from the original images."""
 
-    image_size: Tuple[int, int]
+    image_size: Tuple[int, ...]
     """Size of the original input images to be used for upsampling."""
 
     images: torch.Tensor | None = None

eva/vision/models/wrappers/from_registry.py

@@ -2,18 +2,20 @@
 
 from typing import Any, Callable, Dict
 
+import torch
 from typing_extensions import override
 
-from eva.core.models import wrappers
-from eva.vision.models.networks.backbones import BackboneModelRegistry
+from eva.core.models.wrappers import base
+from eva.core.utils import factory
+from eva.vision.models.networks.backbones import backbone_registry
 
 
-class ModelFromRegistry(wrappers.BaseModel):
+class ModelFromRegistry(base.BaseModel[torch.Tensor, torch.Tensor]):
     """Wrapper class for vision backbone models.
 
     This class can be used by load backbones available in eva's
     model registry by name. New backbones can be registered by using
-    the `@register_model(model_name)` decorator.
+    the `@backbone_registry.register(model_name)` decorator.
     """
 
     def __init__(
@@ -21,7 +23,7 @@ class ModelFromRegistry(wrappers.BaseModel):
         model_name: str,
         model_kwargs: Dict[str, Any] | None = None,
         model_extra_kwargs: Dict[str, Any] | None = None,
-        tensor_transforms: Callable | None = None,
+        transforms: Callable | None = None,
     ) -> None:
         """Initializes the model.
 
@@ -29,10 +31,10 @@ class ModelFromRegistry(wrappers.BaseModel):
             model_name: The name of the model to load.
             model_kwargs: The arguments used for instantiating the model.
             model_extra_kwargs: Extra arguments used for instantiating the model.
-            tensor_transforms: The transforms to apply to the output tensor
+            transforms: The transforms to apply to the output tensor
                 produced by the model.
         """
-        super().__init__(tensor_transforms=tensor_transforms)
+        super().__init__(transforms=transforms)
 
         self._model_name = model_name
         self._model_kwargs = model_kwargs or {}
@@ -42,7 +44,10 @@ class ModelFromRegistry(wrappers.BaseModel):
 
     @override
     def load_model(self) -> None:
-        self._model = BackboneModelRegistry.load_model(
-            self._model_name, self._model_kwargs | self._model_extra_kwargs
+        self._model = factory.ModuleFactory(
+            registry=backbone_registry,
+            name=self._model_name,
+            init_args=self._model_kwargs | self._model_extra_kwargs,
         )
+
         ModelFromRegistry.__name__ = self._model_name