Repository: james77777778/keras-image-models (language: python)

Issue:
Improve test coverage



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index adbc733..d74b769 100644
--- a/README.md
+++ b/README.md
@@ -29,9 +29,9 @@ pip install keras kimm
 [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/14WxYgVjlwCIO9MwqPYW-dskbTL2UHsVN?usp=sharing)
 
 ```python
-import cv2
 import keras
 from keras import ops
+from keras import utils
 from keras.applications.imagenet_utils import decode_predictions
 
 import kimm
@@ -43,15 +43,15 @@ print(kimm.list_models())
 print(kimm.list_models("efficientnet", weights="imagenet"))  # fuzzy search
 
 # Initialize the model with pretrained weights
-model = kimm.models.EfficientNetV2B0()
-image_size = model._default_size
+model = kimm.models.VisionTransformerTiny16()
+image_size = (model._default_size, model._default_size)
 
 # Load an image as the model input
 image_path = keras.utils.get_file(
     "african_elephant.jpg", "https://i.imgur.com/Bvro0YD.png"
 )
-image = cv2.imread(image_path)
-image = cv2.resize(image, (image_size, image_size))
+image = utils.load_img(image_path, target_size=image_size)
+image = utils.img_to_array(image)
 x = ops.convert_to_tensor(image)
 x = ops.expand_dims(x, axis=0)
 
@@ -62,9 +62,9 @@ print("Predicted:", decode_predictions(preds, top=3)[0])
 
 ```bash
 ['ConvMixer1024D20', 'ConvMixer1536D20', 'ConvMixer736D32', 'ConvNeXtAtto', ...]
-['EfficientNetB0', 'EfficientNetB1', 'EfficientNetB2', 'EfficientNetB3', ...]
-1/1 ━━━━━━━━━━━━━━━━━━━━ 11s 11s/step
-Predicted: [('n02504458', 'African_elephant', 0.90578836), ('n01871265', 'tusker', 0.024864597), ('n02504013', 'Indian_elephant', 0.01161992)]
+['VisionTransformerBase16', 'VisionTransformerBase32', 'VisionTransformerSmall16', ...]
+1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 1s/step
+Predicted: [('n02504458', 'African_elephant', 0.6895825), ('n01871265', 'tusker', 0.17934209), ('n02504013', 'Indian_elephant', 0.12927249)]
 ```
 
 ### An end-to-end example: fine-tuning an image classification model on a cats vs. dogs dataset
diff --git a/kimm/__init__.py b/kimm/__init__.py
index 74309d2..57a48e1 100644
--- a/kimm/__init__.py
+++ b/kimm/__init__.py
@@ -1,4 +1,4 @@
 from kimm import models  # force to add models to the registry
 from kimm.utils.model_registry import list_models
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"
diff --git a/kimm/models/base_model.py b/kimm/models/base_model.py
index d1ad061..06baa88 100644
--- a/kimm/models/base_model.py
+++ b/kimm/models/base_model.py
@@ -1,4 +1,3 @@
-import abc
 import pathlib
 import typing
 import urllib.parse
@@ -12,6 +11,12 @@
 
 
 class BaseModel(models.Model):
+    default_origin = (
+        "https://github.com/james77777778/kimm/releases/download/0.1.0/"
+    )
+    available_feature_keys = []
+    available_weights = []
+
     def __init__(
         self,
         inputs,
@@ -183,12 +188,6 @@ def load_pretrained_weights(self, weights_url: typing.Optional[str] = None):
             )
             self.load_weights(weights_path)
 
-    @staticmethod
-    @abc.abstractmethod
-    def available_feature_keys():
-        # TODO: add docstring
-        raise NotImplementedError
-
     def get_config(self):
         # Don't chain to super here. The default `get_config()` for functional
         # models is nested and cannot be passed to BaseModel.
@@ -215,6 +214,19 @@ def get_config(self):
     def fix_config(self, config: typing.Dict):
         return config
 
-    @property
-    def default_origin(self):
-        return "https://github.com/james77777778/kimm/releases/download/0.1.0/"
+    def get_weights_url(self, weights):
+        if weights is None:
+            return None
+
+        for _weights, _origin, _file_name in self.available_weights:
+            if weights == _weights:
+                return f"{_origin}/{_file_name}"
+
+        # Failed to find the weights
+        _available_weights_name = [
+            _weights for _weights, _ in self.available_weights
+        ]
+        raise ValueError(
+            f"Available weights are {_available_weights_name}. "
+            f"Received weights={weights}"
+        )
diff --git a/kimm/models/convmixer.py b/kimm/models/convmixer.py
index fe2212b..d3d455a 100644
--- a/kimm/models/convmixer.py
+++ b/kimm/models/convmixer.py
@@ -52,6 +52,8 @@ def __init__(
         activation: str = "relu",
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         input_tensor = kwargs.pop("input_tensor", None)
         self.set_properties(kwargs)
         inputs = self.determine_input_tensor(
@@ -100,10 +102,6 @@ def __init__(
         self.kernel_size = kernel_size
         self.activation = activation
 
-    @staticmethod
-    def available_feature_keys():
-        raise NotImplementedError
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -136,6 +134,15 @@ def fix_config(self, config):
 
 
 class ConvMixer736D32(ConvMixer):
+    available_feature_keys = ["STEM", *[f"BLOCK{i}" for i in range(32)]]
+    available_weights = [
+        (
+            "imagenet",
+            ConvMixer.default_origin,
+            "convmixer736d32_convmixer_768_32.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -151,9 +158,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convmixer736d32_convmixer_768_32.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             32,
             768,
@@ -173,14 +177,17 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(32)])
-        return feature_keys
-
 
 class ConvMixer1024D20(ConvMixer):
+    available_feature_keys = ["STEM", *[f"BLOCK{i}" for i in range(20)]]
+    available_weights = [
+        (
+            "imagenet",
+            ConvMixer.default_origin,
+            "convmixer1024d20_convmixer_1024_20_ks9_p14.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -196,9 +203,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convmixer1024d20_convmixer_1024_20_ks9_p14.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             20,
             1024,
@@ -218,14 +222,17 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(20)])
-        return feature_keys
-
 
 class ConvMixer1536D20(ConvMixer):
+    available_feature_keys = ["STEM", *[f"BLOCK{i}" for i in range(20)]]
+    available_weights = [
+        (
+            "imagenet",
+            ConvMixer.default_origin,
+            "convmixer1536d20_convmixer_1536_20.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -241,9 +248,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convmixer1536d20_convmixer_1536_20.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             20,
             1536,
@@ -263,12 +267,6 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(20)])
-        return feature_keys
-
 
 add_model_to_registry(ConvMixer736D32, "imagenet")
 add_model_to_registry(ConvMixer1024D20, "imagenet")
diff --git a/kimm/models/convnext.py b/kimm/models/convnext.py
index 06fd880..f0edf67 100644
--- a/kimm/models/convnext.py
+++ b/kimm/models/convnext.py
@@ -120,6 +120,11 @@ def apply_convnext_stage(
 
 @keras.saving.register_keras_serializable(package="kimm")
 class ConvNeXt(BaseModel):
+    available_feature_keys = [
+        "STEM_S4",
+        *[f"BLOCK{i}_S{2**(i+2)}" for i in range(4)],
+    ]
+
     def __init__(
         self,
         depths: typing.Sequence[int] = [3, 3, 9, 3],
@@ -130,6 +135,8 @@ def __init__(
         use_conv_mlp: bool = False,
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         input_tensor = kwargs.pop("input_tensor", None)
         self.set_properties(kwargs)
         inputs = self.determine_input_tensor(
@@ -197,12 +204,6 @@ def build_top(self, inputs, classes, classifier_activation, dropout_rate):
         )(x)
         return x
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S4"]
-        feature_keys.extend([f"BLOCK{i}_S{2**(i+2)}" for i in range(4)])
-        return feature_keys
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -237,6 +238,14 @@ def fix_config(self, config):
 
 
 class ConvNeXtAtto(ConvNeXt):
+    available_weights = [
+        (
+            "imagenet",
+            ConvNeXt.default_origin,
+            "convnextatto_convnext_atto.d2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -252,9 +261,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convnextatto_convnext_atto.d2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             (2, 2, 6, 2),
             (40, 80, 160, 320),
@@ -277,6 +283,14 @@ def __init__(
 
 
 class ConvNeXtFemto(ConvNeXt):
+    available_weights = [
+        (
+            "imagenet",
+            ConvNeXt.default_origin,
+            "convnextfemto_convnext_femto.d1_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -292,9 +306,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convnextfemto_convnext_femto.d1_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             (2, 2, 6, 2),
             (48, 96, 192, 384),
@@ -317,6 +328,14 @@ def __init__(
 
 
 class ConvNeXtPico(ConvNeXt):
+    available_weights = [
+        (
+            "imagenet",
+            ConvNeXt.default_origin,
+            "convnextpico_convnext_pico.d1_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -332,9 +351,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convnextpico_convnext_pico.d1_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             (2, 2, 6, 2),
             (64, 128, 256, 512),
@@ -357,6 +373,14 @@ def __init__(
 
 
 class ConvNeXtNano(ConvNeXt):
+    available_weights = [
+        (
+            "imagenet",
+            ConvNeXt.default_origin,
+            "convnextnano_convnext_nano.in12k_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -372,9 +396,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convnextnano_convnext_nano.in12k_ft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             (2, 2, 8, 2),
             (80, 160, 320, 640),
@@ -397,6 +418,14 @@ def __init__(
 
 
 class ConvNeXtTiny(ConvNeXt):
+    available_weights = [
+        (
+            "imagenet",
+            ConvNeXt.default_origin,
+            "convnexttiny_convnext_tiny.in12k_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -412,9 +441,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convnexttiny_convnext_tiny.in12k_ft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             (3, 3, 9, 3),
             (96, 192, 384, 768),
@@ -437,6 +463,14 @@ def __init__(
 
 
 class ConvNeXtSmall(ConvNeXt):
+    available_weights = [
+        (
+            "imagenet",
+            ConvNeXt.default_origin,
+            "convnextsmall_convnext_small.in12k_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -452,9 +486,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convnextsmall_convnext_small.in12k_ft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             (3, 3, 27, 3),
             (96, 192, 384, 768),
@@ -477,6 +508,14 @@ def __init__(
 
 
 class ConvNeXtBase(ConvNeXt):
+    available_weights = [
+        (
+            "imagenet",
+            ConvNeXt.default_origin,
+            "convnextbase_convnext_base.fb_in22k_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -492,9 +531,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convnextbase_convnext_base.fb_in22k_ft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             (3, 3, 27, 3),
             (128, 256, 512, 1024),
@@ -517,6 +553,14 @@ def __init__(
 
 
 class ConvNeXtLarge(ConvNeXt):
+    available_weights = [
+        (
+            "imagenet",
+            ConvNeXt.default_origin,
+            "convnextlarge_convnext_large.fb_in22k_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -532,9 +576,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "convnextlarge_convnext_large.fb_in22k_ft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             (3, 3, 27, 3),
             (192, 384, 768, 1536),
@@ -557,6 +598,8 @@ def __init__(
 
 
 class ConvNeXtXLarge(ConvNeXt):
+    available_weights = []
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
diff --git a/kimm/models/densenet.py b/kimm/models/densenet.py
index 71d5e69..3993cc3 100644
--- a/kimm/models/densenet.py
+++ b/kimm/models/densenet.py
@@ -66,12 +66,19 @@ def apply_dense_transition_block(
 
 @keras.saving.register_keras_serializable(package="kimm")
 class DenseNet(BaseModel):
+    available_feature_keys = [
+        "STEM_S4",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(4), [8, 16, 32, 32])],
+    ]
+
     def __init__(
         self,
         growth_rate: float = 32,
         num_blocks: typing.Sequence[int] = [6, 12, 24, 16],
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         input_tensor = kwargs.pop("input_tensor", None)
         self.set_properties(kwargs)
         inputs = self.determine_input_tensor(
@@ -133,14 +140,6 @@ def __init__(
         self.growth_rate = growth_rate
         self.num_blocks = num_blocks
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S4"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(4), [8, 16, 32, 32])]
-        )
-        return feature_keys
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -161,6 +160,14 @@ def fix_config(self, config: typing.Dict):
 
 
 class DenseNet121(DenseNet):
+    available_weights = [
+        (
+            "imagenet",
+            DenseNet.default_origin,
+            "densenet121_densenet121.ra_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -176,9 +183,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "densenet121_densenet121.ra_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             32,
             [6, 12, 24, 16],
@@ -198,6 +202,14 @@ def __init__(
 
 
 class DenseNet161(DenseNet):
+    available_weights = [
+        (
+            "imagenet",
+            DenseNet.default_origin,
+            "densenet161_densenet161.tv_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -213,9 +225,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "densenet161_densenet161.tv_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             48,
             [6, 12, 36, 24],
@@ -235,6 +244,14 @@ def __init__(
 
 
 class DenseNet169(DenseNet):
+    available_weights = [
+        (
+            "imagenet",
+            DenseNet.default_origin,
+            "densenet169_densenet169.tv_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -250,9 +267,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "densenet169_densenet169.tv_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             32,
             [6, 12, 32, 32],
@@ -272,6 +286,14 @@ def __init__(
 
 
 class DenseNet201(DenseNet):
+    available_weights = [
+        (
+            "imagenet",
+            DenseNet.default_origin,
+            "densenet201_densenet201.tv_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -287,9 +309,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "densenet201_densenet201.tv_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             32,
             [6, 12, 48, 32],
diff --git a/kimm/models/efficientnet.py b/kimm/models/efficientnet.py
index d328c3f..ab6f0b4 100644
--- a/kimm/models/efficientnet.py
+++ b/kimm/models/efficientnet.py
@@ -125,6 +125,16 @@ def apply_edge_residual_block(
 
 @keras.saving.register_keras_serializable(package="kimm")
 class EfficientNet(BaseModel):
+    # for: v1, v1_lite, v2_m, v2_l, v2_xl, tinynet
+    # not for: v2_s, v2_base
+    available_feature_keys = [
+        "STEM_S2",
+        *[
+            f"BLOCK{i}_S{j}"
+            for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
+        ],
+    ]
+
     def __init__(
         self,
         width: float = 1.0,
@@ -137,6 +147,8 @@ def __init__(
         config: str = "v1",
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         _available_configs = [
             "v1",
             "v1_lite",
@@ -270,19 +282,6 @@ def __init__(
         self.activation = activation
         self.config = config
 
-    @staticmethod
-    def available_feature_keys():
-        # for: v1, v1_lite, v2_m, v2_l, v2_xl, tinynet
-        # not for: v2_s, v2_base
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [
-                f"BLOCK{i}_S{j}"
-                for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
-            ]
-        )
-        return feature_keys
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -320,6 +319,14 @@ def fix_config(self, config: typing.Dict):
 
 
 class EfficientNetB0(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetb0_tf_efficientnet_b0.ns_jft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -336,9 +343,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetb0_tf_efficientnet_b0.ns_jft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -367,6 +371,14 @@ def __init__(
 
 
 class EfficientNetB1(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetb1_tf_efficientnet_b1.ns_jft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -383,9 +395,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetb1_tf_efficientnet_b1.ns_jft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -414,6 +423,14 @@ def __init__(
 
 
 class EfficientNetB2(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetb2_tf_efficientnet_b2.ns_jft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -430,9 +447,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetb2_tf_efficientnet_b2.ns_jft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.1,
@@ -461,6 +475,14 @@ def __init__(
 
 
 class EfficientNetB3(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetb3_tf_efficientnet_b3.ns_jft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -477,9 +499,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetb3_tf_efficientnet_b3.ns_jft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.2,
@@ -508,6 +527,14 @@ def __init__(
 
 
 class EfficientNetB4(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetb4_tf_efficientnet_b4.ns_jft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -524,9 +551,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetb4_tf_efficientnet_b4.ns_jft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.4,
@@ -555,6 +579,14 @@ def __init__(
 
 
 class EfficientNetB5(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetb5_tf_efficientnet_b5.ns_jft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -571,9 +603,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetb5_tf_efficientnet_b5.ns_jft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.6,
@@ -602,6 +631,14 @@ def __init__(
 
 
 class EfficientNetB6(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetb6_tf_efficientnet_b6.ns_jft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -618,9 +655,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetb6_tf_efficientnet_b6.ns_jft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.8,
@@ -649,6 +683,14 @@ def __init__(
 
 
 class EfficientNetB7(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetb7_tf_efficientnet_b7.ns_jft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -665,9 +707,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetb7_tf_efficientnet_b7.ns_jft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             2.0,
@@ -696,6 +735,14 @@ def __init__(
 
 
 class EfficientNetLiteB0(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetliteb0_tf_efficientnet_lite0.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -712,9 +759,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetliteb0_tf_efficientnet_lite0.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -743,6 +787,14 @@ def __init__(
 
 
 class EfficientNetLiteB1(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetliteb1_tf_efficientnet_lite1.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -759,9 +811,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetliteb1_tf_efficientnet_lite1.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -790,6 +839,14 @@ def __init__(
 
 
 class EfficientNetLiteB2(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetliteb2_tf_efficientnet_lite2.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -806,9 +863,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetliteb2_tf_efficientnet_lite2.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.1,
@@ -837,6 +891,14 @@ def __init__(
 
 
 class EfficientNetLiteB3(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetliteb3_tf_efficientnet_lite3.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -853,9 +915,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetliteb3_tf_efficientnet_lite3.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.2,
@@ -884,6 +943,14 @@ def __init__(
 
 
 class EfficientNetLiteB4(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetliteb4_tf_efficientnet_lite4.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -900,9 +967,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetliteb4_tf_efficientnet_lite4.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.4,
@@ -931,6 +995,18 @@ def __init__(
 
 
 class EfficientNetV2S(EfficientNet):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetv2s_tf_efficientnetv2_s.in21k_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -947,11 +1023,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = (
-                "efficientnetv2s_tf_efficientnetv2_s.in21k_ft_in1k.keras"
-            )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -978,16 +1049,16 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
-        )
-        return feature_keys
-
 
 class EfficientNetV2M(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetv2m_tf_efficientnetv2_m.in21k_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1004,11 +1075,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = (
-                "efficientnetv2m_tf_efficientnetv2_m.in21k_ft_in1k.keras"
-            )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -1037,6 +1103,14 @@ def __init__(
 
 
 class EfficientNetV2L(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetv2l_tf_efficientnetv2_l.in21k_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1053,11 +1127,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = (
-                "efficientnetv2l_tf_efficientnetv2_l.in21k_ft_in1k.keras"
-            )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -1086,6 +1155,14 @@ def __init__(
 
 
 class EfficientNetV2XL(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetv2xl_tf_efficientnetv2_xl.in21k_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1102,11 +1179,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = (
-                "efficientnetv2xl_tf_efficientnetv2_xl.in21k_ft_in1k.keras"
-            )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -1135,6 +1207,18 @@ def __init__(
 
 
 class EfficientNetV2B0(EfficientNet):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetv2b0_tf_efficientnetv2_b0.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1151,9 +1235,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetv2b0_tf_efficientnetv2_b0.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -1180,16 +1261,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
-        )
-        return feature_keys
-
 
 class EfficientNetV2B1(EfficientNet):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetv2b1_tf_efficientnetv2_b1.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1206,9 +1291,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetv2b1_tf_efficientnetv2_b1.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -1235,16 +1317,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
-        )
-        return feature_keys
-
 
 class EfficientNetV2B2(EfficientNet):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetv2b2_tf_efficientnetv2_b2.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1261,9 +1347,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetv2b2_tf_efficientnetv2_b2.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.1,
@@ -1291,16 +1374,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
-        )
-        return feature_keys
-
 
 class EfficientNetV2B3(EfficientNet):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "efficientnetv2b3_tf_efficientnetv2_b3.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1317,9 +1404,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "efficientnetv2b3_tf_efficientnetv2_b3.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.2,
@@ -1347,16 +1431,16 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
-        )
-        return feature_keys
-
 
 class TinyNetA(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "tinyneta_tinynet_a.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1373,9 +1457,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "tinyneta_tinynet_a.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.0,
             1.2,
@@ -1402,6 +1483,14 @@ def __init__(
 
 
 class TinyNetB(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "tinynetb_tinynet_b.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1418,9 +1507,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "tinynetb_tinynet_b.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             0.75,
             1.1,
@@ -1447,6 +1533,14 @@ def __init__(
 
 
 class TinyNetC(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "tinynetc_tinynet_c.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1463,9 +1557,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "tinynetc_tinynet_c.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             0.54,
             0.85,
@@ -1492,6 +1583,14 @@ def __init__(
 
 
 class TinyNetD(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "tinynetd_tinynet_d.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1508,9 +1607,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "tinynetd_tinynet_d.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             0.54,
             0.695,
@@ -1537,6 +1633,14 @@ def __init__(
 
 
 class TinyNetE(EfficientNet):
+    available_weights = [
+        (
+            "imagenet",
+            EfficientNet.default_origin,
+            "tinynete_tinynet_e.in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1553,9 +1657,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "tinynete_tinynet_e.in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             0.51,
             0.6,
diff --git a/kimm/models/ghostnet.py b/kimm/models/ghostnet.py
index 6dd923f..b932829 100644
--- a/kimm/models/ghostnet.py
+++ b/kimm/models/ghostnet.py
@@ -230,6 +230,14 @@ def apply_ghost_bottleneck(
 
 @keras.saving.register_keras_serializable(package="kimm")
 class GhostNet(BaseModel):
+    available_feature_keys = [
+        "STEM_S2",
+        *[
+            f"BLOCK{i}_S{j}"
+            for i, j in zip(range(9), [2, 4, 4, 8, 8, 16, 16, 32, 32])
+        ],
+    ]
+
     def __init__(
         self,
         width: float = 1.0,
@@ -237,6 +245,8 @@ def __init__(
         version: typing.Literal["v1", "v2"] = "v1",
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         _available_configs = ["default"]
         if config == "default":
             _config = DEFAULT_CONFIG
@@ -332,17 +342,6 @@ def build_top(self, inputs, classes, classifier_activation, dropout_rate):
         )(x)
         return x
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [
-                f"BLOCK{i}_S{j}"
-                for i, j in zip(range(9), [2, 4, 4, 8, 8, 16, 16, 32, 32])
-            ]
-        )
-        return feature_keys
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -367,6 +366,8 @@ def fix_config(self, config):
 
 
 class GhostNet050(GhostNet):
+    available_weights = []
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -402,6 +403,14 @@ def __init__(
 
 
 class GhostNet100(GhostNet):
+    available_weights = [
+        (
+            "imagenet",
+            GhostNet.default_origin,
+            "ghostnet100_ghostnet_100.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -418,9 +427,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "ghostnet100_ghostnet_100.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.0,
             config,
@@ -440,6 +446,8 @@ def __init__(
 
 
 class GhostNet130(GhostNet):
+    available_weights = []
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -475,6 +483,14 @@ def __init__(
 
 
 class GhostNet100V2(GhostNet):
+    available_weights = [
+        (
+            "imagenet",
+            GhostNet.default_origin,
+            "ghostnet100v2_ghostnetv2_100.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -491,9 +507,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "ghostnet100v2_ghostnetv2_100.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.0,
             config,
@@ -513,6 +526,14 @@ def __init__(
 
 
 class GhostNet130V2(GhostNet):
+    available_weights = [
+        (
+            "imagenet",
+            GhostNet.default_origin,
+            "ghostnet130v2_ghostnetv2_130.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -529,9 +550,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "ghostnet130v2_ghostnetv2_130.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.3,
             config,
@@ -551,6 +569,14 @@ def __init__(
 
 
 class GhostNet160V2(GhostNet):
+    available_weights = [
+        (
+            "imagenet",
+            GhostNet.default_origin,
+            "ghostnet160v2_ghostnetv2_160.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -567,9 +593,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "ghostnet160v2_ghostnetv2_160.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.6,
             config,
diff --git a/kimm/models/inception_v3.py b/kimm/models/inception_v3.py
index 50fee21..55a3043 100644
--- a/kimm/models/inception_v3.py
+++ b/kimm/models/inception_v3.py
@@ -203,7 +203,14 @@ def apply_inception_aux_block(inputs, classes, name="inception_aux_block"):
 
 @keras.saving.register_keras_serializable(package="kimm")
 class InceptionV3Base(BaseModel):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(4), [4, 8, 16, 32])],
+    ]
+
     def __init__(self, has_aux_logits: bool = False, **kwargs):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         input_tensor = kwargs.pop("input_tensor", None)
         self.set_properties(kwargs, 299)
         inputs = self.determine_input_tensor(
@@ -264,14 +271,6 @@ def __init__(self, has_aux_logits: bool = False, **kwargs):
         # All references to `self` below this line
         self.has_aux_logits = has_aux_logits
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(4), [4, 8, 16, 32])]
-        )
-        return feature_keys
-
     def get_config(self):
         config = super().get_config()
         config.update({"has_aux_logits": self.has_aux_logits})
@@ -282,6 +281,19 @@ def fix_config(self, config: typing.Dict):
 
 
 class InceptionV3(InceptionV3Base):
+    available_weights = [
+        (
+            "imagenet_aux_logits",
+            InceptionV3Base.default_origin,
+            "inceptionv3_inception_v3.gluon_in1k_aux_logits.keras",
+        ),
+        (
+            "imagenet_no_aux_logits",
+            InceptionV3Base.default_origin,
+            "inceptionv3_inception_v3.gluon_in1k_no_aux_logits.keras",
+        ),
+    ]
+
     def __init__(
         self,
         has_aux_logits: bool = False,
@@ -300,14 +312,9 @@ def __init__(
         kwargs = self.fix_config(kwargs)
         if weights == "imagenet":
             if has_aux_logits:
-                file_name = (
-                    "inceptionv3_inception_v3.gluon_in1k_aux_logits.keras"
-                )
+                weights = f"{weights}_aux_logits"
             else:
-                file_name = (
-                    "inceptionv3_inception_v3.gluon_in1k_no_aux_logits.keras"
-                )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
+                weights = f"{weights}_no_aux_logits"
         super().__init__(
             has_aux_logits,
             input_tensor=input_tensor,
diff --git a/kimm/models/mobilenet_v2.py b/kimm/models/mobilenet_v2.py
index 2af2426..f5150f1 100644
--- a/kimm/models/mobilenet_v2.py
+++ b/kimm/models/mobilenet_v2.py
@@ -24,6 +24,14 @@
 
 @keras.saving.register_keras_serializable(package="kimm")
 class MobileNetV2(BaseModel):
+    available_feature_keys = [
+        "STEM_S2",
+        *[
+            f"BLOCK{i}_S{j}"
+            for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
+        ],
+    ]
+
     def __init__(
         self,
         width: float = 1.0,
@@ -32,6 +40,8 @@ def __init__(
         config: typing.Literal["default"] = "default",
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         _available_configs = ["default"]
         if config == "default":
             _config = DEFAULT_CONFIG
@@ -111,17 +121,6 @@ def __init__(
         self.fix_stem_and_head_channels = fix_stem_and_head_channels
         self.config = config
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [
-                f"BLOCK{i}_S{j}"
-                for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
-            ]
-        )
-        return feature_keys
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -152,6 +151,14 @@ def fix_config(self, config):
 
 
 class MobileNet050V2(MobileNetV2):
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV2.default_origin,
+            "mobilenet050v2_mobilenetv2_050.lamb_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -168,9 +175,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "mobilenet050v2_mobilenetv2_050.lamb_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             0.5,
             1.0,
@@ -191,6 +195,14 @@ def __init__(
 
 
 class MobileNet100V2(MobileNetV2):
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV2.default_origin,
+            "mobilenet100v2_mobilenetv2_100.ra_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -207,9 +219,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "mobilenet100v2_mobilenetv2_100.ra_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.0,
             1.0,
@@ -230,6 +239,14 @@ def __init__(
 
 
 class MobileNet110V2(MobileNetV2):
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV2.default_origin,
+            "mobilenet110v2_mobilenetv2_110d.ra_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -246,9 +263,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "mobilenet110v2_mobilenetv2_110d.ra_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.1,
             1.2,
@@ -269,6 +283,14 @@ def __init__(
 
 
 class MobileNet120V2(MobileNetV2):
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV2.default_origin,
+            "mobilenet120v2_mobilenetv2_120d.ra_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -285,9 +307,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "mobilenet120v2_mobilenetv2_120d.ra_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.2,
             1.4,
@@ -308,6 +327,14 @@ def __init__(
 
 
 class MobileNet140V2(MobileNetV2):
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV2.default_origin,
+            "mobilenet140v2_mobilenetv2_140.ra_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -324,9 +351,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "mobilenet140v2_mobilenetv2_140.ra_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.4,
             1.0,
diff --git a/kimm/models/mobilenet_v3.py b/kimm/models/mobilenet_v3.py
index 92b812d..5d61396 100644
--- a/kimm/models/mobilenet_v3.py
+++ b/kimm/models/mobilenet_v3.py
@@ -90,6 +90,8 @@ def __init__(
         minimal: bool = False,
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         _available_configs = ["small", "large", "lcnet"]
         if config == "small":
             _config = DEFAULT_SMALL_CONFIG
@@ -270,10 +272,6 @@ def build_top(
         )(x)
         return x
 
-    @staticmethod
-    def available_feature_keys():
-        raise NotImplementedError()
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -306,6 +304,18 @@ def fix_config(self, config):
 
 
 class MobileNet050V3Small(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [4, 8, 16, 16, 32, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV3.default_origin,
+            "mobilenet050v3small_mobilenetv3_small_050.lamb_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -322,11 +332,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = (
-                "mobilenet050v3small_mobilenetv3_small_050.lamb_in1k.keras"
-            )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             0.5,
             1.0,
@@ -345,16 +350,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [4, 8, 16, 16, 32, 32])]
-        )
-        return feature_keys
-
 
 class MobileNet075V3Small(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [4, 8, 16, 16, 32, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV3.default_origin,
+            "mobilenet075v3small_mobilenetv3_small_075.lamb_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -371,11 +380,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = (
-                "mobilenet075v3small_mobilenetv3_small_075.lamb_in1k.keras"
-            )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             0.75,
             1.0,
@@ -394,16 +398,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [4, 8, 16, 16, 32, 32])]
-        )
-        return feature_keys
-
 
 class MobileNet100V3Small(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [4, 8, 16, 16, 32, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV3.default_origin,
+            "mobilenet100v3small_mobilenetv3_small_100.lamb_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -420,11 +428,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = (
-                "mobilenet100v3small_mobilenetv3_small_100.lamb_in1k.keras"
-            )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.0,
             1.0,
@@ -443,16 +446,23 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [4, 8, 16, 16, 32, 32])]
-        )
-        return feature_keys
-
 
 class MobileNet100V3SmallMinimal(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [4, 8, 16, 16, 32, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV3.default_origin,
+            (
+                "mobilenet100v3smallminimal_"
+                "tf_mobilenetv3_small_minimal_100.in1k.keras"
+            ),
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -469,12 +479,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = (
-                "mobilenet100v3smallminimal_"
-                "tf_mobilenetv3_small_minimal_100.in1k.keras"
-            )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -497,16 +501,26 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [4, 8, 16, 16, 32, 32])]
-        )
-        return feature_keys
-
 
 class MobileNet100V3Large(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[
+            f"BLOCK{i}_S{j}"
+            for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
+        ],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV3.default_origin,
+            (
+                "mobilenet100v3large_"
+                "mobilenetv3_large_100.miil_in21k_ft_in1k.keras"
+            ),
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -523,12 +537,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = (
-                "mobilenet100v3large_"
-                "mobilenetv3_large_100.miil_in21k_ft_in1k.keras"
-            )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             1.0,
             1.0,
@@ -557,19 +565,26 @@ def build_preprocessing(self, inputs, mode="imagenet"):
         else:
             return super().build_preprocessing(inputs, mode)
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [
-                f"BLOCK{i}_S{j}"
-                for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
-            ]
-        )
-        return feature_keys
-
 
 class MobileNet100V3LargeMinimal(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[
+            f"BLOCK{i}_S{j}"
+            for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
+        ],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV3.default_origin,
+            (
+                "mobilenet100v3largeminimal_"
+                "tf_mobilenetv3_large_minimal_100.in1k.keras"
+            ),
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -586,12 +601,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = (
-                "mobilenet100v3largeminimal_"
-                "tf_mobilenetv3_large_minimal_100.in1k.keras"
-            )
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -614,19 +623,14 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [
-                f"BLOCK{i}_S{j}"
-                for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
-            ]
-        )
-        return feature_keys
-
 
 class LCNet035(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])],
+    ]
+    available_weights = []
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -662,16 +666,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
-        )
-        return feature_keys
-
 
 class LCNet050(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV3.default_origin,
+            "lcnet050_lcnet_050.ra2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -688,9 +696,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "lcnet050_lcnet_050.ra2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             0.5,
@@ -710,16 +715,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
-        )
-        return feature_keys
-
 
 class LCNet075(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV3.default_origin,
+            "lcnet075_lcnet_075.ra2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -736,9 +745,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "lcnet075_lcnet_075.ra2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             0.75,
@@ -758,16 +764,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
-        )
-        return feature_keys
-
 
 class LCNet100(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            MobileNetV3.default_origin,
+            "lcnet100_lcnet_100.ra2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -784,9 +794,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "lcnet100_lcnet_100.ra2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         # default to TF configuration (bn_epsilon=1e-3 and padding="same")
         super().__init__(
             1.0,
@@ -806,16 +813,17 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(6), [2, 4, 8, 16, 16, 32])]
-        )
-        return feature_keys
-
 
 class LCNet150(MobileNetV3):
+    available_feature_keys = [
+        "STEM_S2",
+        *[
+            f"BLOCK{i}_S{j}"
+            for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
+        ],
+    ]
+    available_weights = []
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -851,17 +859,6 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [
-                f"BLOCK{i}_S{j}"
-                for i, j in zip(range(7), [2, 4, 8, 16, 16, 32, 32])
-            ]
-        )
-        return feature_keys
-
 
 add_model_to_registry(MobileNet050V3Small, "imagenet")
 add_model_to_registry(MobileNet075V3Small, "imagenet")
diff --git a/kimm/models/mobilevit.py b/kimm/models/mobilevit.py
index 882e272..c065d5d 100644
--- a/kimm/models/mobilevit.py
+++ b/kimm/models/mobilevit.py
@@ -162,6 +162,11 @@ def apply_mobilevit_block(
 
 @keras.saving.register_keras_serializable(package="kimm")
 class MobileViT(BaseModel):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(5), [2, 4, 8, 16, 32])],
+    ]
+
     def __init__(
         self,
         stem_channels: int = 16,
@@ -170,6 +175,8 @@ def __init__(
         config: str = "v1_s",
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         _available_configs = ["v1_s", "v1_xs", "v1_xss"]
         if config == "v1_s":
             _config = DEFAULT_V1_S_CONFIG
@@ -258,14 +265,6 @@ def __init__(
         self.activation = activation
         self.config = config
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(5), [2, 4, 8, 16, 32])]
-        )
-        return feature_keys
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -291,6 +290,14 @@ def fix_config(self, config):
 
 
 class MobileViTS(MobileViT):
+    available_weights = [
+        (
+            "imagenet",
+            MobileViT.default_origin,
+            "mobilevits_mobilevit_s.cvnets_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -307,9 +314,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "mobilevits_mobilevit_s.cvnets_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             16,
             640,
@@ -330,6 +334,14 @@ def __init__(
 
 
 class MobileViTXS(MobileViT):
+    available_weights = [
+        (
+            "imagenet",
+            MobileViT.default_origin,
+            "mobilevitxs_mobilevit_xs.cvnets_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -346,9 +358,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "mobilevitxs_mobilevit_xs.cvnets_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             16,
             384,
@@ -369,6 +378,14 @@ def __init__(
 
 
 class MobileViTXXS(MobileViT):
+    available_weights = [
+        (
+            "imagenet",
+            MobileViT.default_origin,
+            "mobilevitxxs_mobilevit_xxs.cvnets_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -385,9 +402,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "mobilevitxxs_mobilevit_xxs.cvnets_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             16,
             320,
diff --git a/kimm/models/regnet.py b/kimm/models/regnet.py
index a67d23e..6ab337e 100644
--- a/kimm/models/regnet.py
+++ b/kimm/models/regnet.py
@@ -143,6 +143,11 @@ def apply_bottleneck_block(
 
 @keras.saving.register_keras_serializable(package="kimm")
 class RegNet(BaseModel):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(4), [4, 8, 16, 32])],
+    ]
+
     def __init__(
         self,
         w0: int = 80,
@@ -153,6 +158,8 @@ def __init__(
         se_ratio: float = 0.0,
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         per_stage_config = _generate_regnet(w0, wa, wm, group_size, depth)
 
         input_tensor = kwargs.pop("input_tensor", None)
@@ -202,14 +209,6 @@ def __init__(
         self.depth = depth
         self.se_ratio = se_ratio
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(4), [4, 8, 16, 32])]
-        )
-        return feature_keys
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -237,6 +236,14 @@ def fix_config(self, config):
 
 
 class RegNetX002(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx002_regnetx_002.pycls_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -252,9 +259,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx002_regnetx_002.pycls_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             24,
             36.44,
@@ -276,6 +280,14 @@ def __init__(
 
 
 class RegNetY002(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety002_regnety_002.pycls_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -291,9 +303,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety002_regnety_002.pycls_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             24,
             36.44,
@@ -316,6 +325,14 @@ def __init__(
 
 
 class RegNetX004(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx004_regnetx_004.pycls_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -331,9 +348,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx004_regnetx_004.pycls_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             24,
             24.48,
@@ -355,6 +369,14 @@ def __init__(
 
 
 class RegNetY004(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety004_regnety_004.tv2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -370,9 +392,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety004_regnety_004.tv2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             48,
             27.89,
@@ -395,6 +414,14 @@ def __init__(
 
 
 class RegNetX006(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx006_regnetx_006.pycls_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -410,9 +437,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx006_regnetx_006.pycls_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             48,
             36.97,
@@ -434,6 +458,14 @@ def __init__(
 
 
 class RegNetY006(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety006_regnety_006.pycls_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -449,9 +481,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety006_regnety_006.pycls_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             48,
             32.54,
@@ -474,6 +503,14 @@ def __init__(
 
 
 class RegNetX008(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx008_regnetx_008.tv2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -489,9 +526,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx008_regnetx_008.tv2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             56,
             35.73,
@@ -513,6 +547,14 @@ def __init__(
 
 
 class RegNetY008(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety008_regnety_008.pycls_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -528,9 +570,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety008_regnety_008.pycls_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             56,
             38.84,
@@ -553,6 +592,14 @@ def __init__(
 
 
 class RegNetX016(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx016_regnetx_016.tv2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -568,9 +615,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx016_regnetx_016.tv2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             80,
             34.01,
@@ -592,6 +636,14 @@ def __init__(
 
 
 class RegNetY016(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety016_regnety_016.tv2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -607,9 +659,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety016_regnety_016.tv2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             48,
             20.71,
@@ -632,6 +681,14 @@ def __init__(
 
 
 class RegNetX032(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx032_regnetx_032.tv2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -647,9 +704,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx032_regnetx_032.tv2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             88,
             26.31,
@@ -671,6 +725,14 @@ def __init__(
 
 
 class RegNetY032(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety032_regnety_032.ra_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -686,9 +748,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety032_regnety_032.ra_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             80,
             42.63,
@@ -711,6 +770,14 @@ def __init__(
 
 
 class RegNetX040(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx040_regnetx_040.pycls_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -726,9 +793,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx040_regnetx_040.pycls_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             96,
             38.65,
@@ -750,6 +814,14 @@ def __init__(
 
 
 class RegNetY040(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety040_regnety_040.ra3_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -765,9 +837,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety040_regnety_040.ra3_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             96,
             31.41,
@@ -790,6 +859,14 @@ def __init__(
 
 
 class RegNetX064(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx064_regnetx_064.pycls_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -805,9 +882,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx064_regnetx_064.pycls_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             184,
             60.83,
@@ -829,6 +903,14 @@ def __init__(
 
 
 class RegNetY064(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety064_regnety_064.ra3_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -844,9 +926,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety064_regnety_064.ra3_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             112,
             33.22,
@@ -869,6 +948,14 @@ def __init__(
 
 
 class RegNetX080(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx080_regnetx_080.tv2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -884,9 +971,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx080_regnetx_080.tv2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             80,
             49.56,
@@ -908,6 +992,14 @@ def __init__(
 
 
 class RegNetY080(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety080_regnety_080.ra3_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -923,9 +1015,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety080_regnety_080.ra3_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             192,
             76.82,
@@ -948,6 +1037,14 @@ def __init__(
 
 
 class RegNetX120(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx120_regnetx_120.pycls_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -963,9 +1060,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx120_regnetx_120.pycls_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             168,
             73.36,
@@ -987,6 +1081,14 @@ def __init__(
 
 
 class RegNetY120(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety120_regnety_120.sw_in12k_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1002,9 +1104,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety120_regnety_120.sw_in12k_ft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             168,
             73.36,
@@ -1027,6 +1126,14 @@ def __init__(
 
 
 class RegNetX160(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx160_regnetx_160.tv2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1042,9 +1149,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx160_regnetx_160.tv2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             216,
             55.59,
@@ -1066,6 +1170,14 @@ def __init__(
 
 
 class RegNetY160(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety160_regnety_160.swag_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1081,9 +1193,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety160_regnety_160.swag_ft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             200,
             106.23,
@@ -1106,6 +1215,14 @@ def __init__(
 
 
 class RegNetX320(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnetx320_regnetx_320.tv2_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1121,9 +1238,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnetx320_regnetx_320.tv2_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             320,
             69.86,
@@ -1145,6 +1259,14 @@ def __init__(
 
 
 class RegNetY320(RegNet):
+    available_weights = [
+        (
+            "imagenet",
+            RegNet.default_origin,
+            "regnety320_regnety_320.swag_ft_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -1160,9 +1282,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "regnety320_regnety_320.swag_ft_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             232,
             115.89,
diff --git a/kimm/models/resnet.py b/kimm/models/resnet.py
index 43fe944..298dea2 100644
--- a/kimm/models/resnet.py
+++ b/kimm/models/resnet.py
@@ -105,12 +105,19 @@ def apply_bottleneck_block(
 
 @keras.saving.register_keras_serializable(package="kimm")
 class ResNet(BaseModel):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(4), [4, 8, 16, 32])],
+    ]
+
     def __init__(
         self,
         block_fn: typing.Literal["basic", "bottleneck"],
         num_blocks: typing.Sequence[int],
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         if block_fn not in ("basic", "bottleneck"):
             raise ValueError(
                 "`block_fn` must be one of ('basic', 'bottelneck'). "
@@ -172,14 +179,6 @@ def __init__(
         self.block_fn = block_fn
         self.num_blocks = num_blocks
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(4), [4, 8, 16, 32])]
-        )
-        return feature_keys
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -200,6 +199,14 @@ def fix_config(self, config):
 
 
 class ResNet18(ResNet):
+    available_weights = [
+        (
+            "imagenet",
+            ResNet.default_origin,
+            "resnet18_resnet18.a1_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -215,9 +222,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "resnet18_resnet18.a1_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             "basic",
             [2, 2, 2, 2],
@@ -236,6 +240,14 @@ def __init__(
 
 
 class ResNet34(ResNet):
+    available_weights = [
+        (
+            "imagenet",
+            ResNet.default_origin,
+            "resnet34_resnet34.a1_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -251,9 +263,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "resnet34_resnet34.a1_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             "basic",
             [3, 4, 6, 3],
@@ -272,6 +281,14 @@ def __init__(
 
 
 class ResNet50(ResNet):
+    available_weights = [
+        (
+            "imagenet",
+            ResNet.default_origin,
+            "resnet50_resnet50.a1_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -287,9 +304,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "resnet50_resnet50.a1_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             "bottleneck",
             [3, 4, 6, 3],
@@ -308,6 +322,14 @@ def __init__(
 
 
 class ResNet101(ResNet):
+    available_weights = [
+        (
+            "imagenet",
+            ResNet.default_origin,
+            "resnet101_resnet101.a1_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -323,9 +345,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "resnet101_resnet101.a1_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             "bottleneck",
             [3, 4, 23, 3],
@@ -344,6 +363,14 @@ def __init__(
 
 
 class ResNet152(ResNet):
+    available_weights = [
+        (
+            "imagenet",
+            ResNet.default_origin,
+            "resnet152_resnet152.a1_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -359,9 +386,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "resnet152_resnet152.a1_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             "bottleneck",
             [3, 8, 36, 3],
diff --git a/kimm/models/vgg.py b/kimm/models/vgg.py
index 77775d5..069426e 100644
--- a/kimm/models/vgg.py
+++ b/kimm/models/vgg.py
@@ -108,7 +108,13 @@ def apply_conv_mlp_layer(
 
 @keras.saving.register_keras_serializable(package="kimm")
 class VGG(BaseModel):
+    available_feature_keys = [
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(6), [1, 2, 4, 8, 16, 32])],
+    ]
+
     def __init__(self, config: typing.Union[str, typing.List], **kwargs):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         _available_configs = ["vgg11", "vgg13", "vgg16", "vgg19"]
         if config == "vgg11":
             _config = DEFAULT_VGG11_CONFIG
@@ -178,12 +184,6 @@ def __init__(self, config: typing.Union[str, typing.List], **kwargs):
         # All references to `self` below this line
         self.config = config
 
-    @staticmethod
-    def available_feature_keys():
-        return [
-            f"BLOCK{i}_S{j}" for i, j in zip(range(6), [1, 2, 4, 8, 16, 32])
-        ]
-
     def get_config(self):
         config = super().get_config()
         config.update({"config": self.config})
@@ -202,6 +202,14 @@ def fix_config(self, config: typing.Dict):
 
 
 class VGG11(VGG):
+    available_weights = [
+        (
+            "imagenet",
+            VGG.default_origin,
+            "vgg11_vgg11_bn.tv_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -217,9 +225,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "vgg11_vgg11_bn.tv_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             "vgg11",
             input_tensor=input_tensor,
@@ -237,6 +242,14 @@ def __init__(
 
 
 class VGG13(VGG):
+    available_weights = [
+        (
+            "imagenet",
+            VGG.default_origin,
+            "vgg13_vgg13_bn.tv_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -252,9 +265,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "vgg13_vgg13_bn.tv_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             "vgg13",
             input_tensor=input_tensor,
@@ -272,6 +282,14 @@ def __init__(
 
 
 class VGG16(VGG):
+    available_weights = [
+        (
+            "imagenet",
+            VGG.default_origin,
+            "vgg16_vgg16_bn.tv_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -287,9 +305,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "vgg16_vgg16_bn.tv_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             "vgg16",
             input_tensor=input_tensor,
@@ -307,6 +322,14 @@ def __init__(
 
 
 class VGG19(VGG):
+    available_weights = [
+        (
+            "imagenet",
+            VGG.default_origin,
+            "vgg19_vgg19_bn.tv_in1k.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -322,9 +345,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "vgg19_vgg19_bn.tv_in1k.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             "vgg19",
             input_tensor=input_tensor,
diff --git a/kimm/models/vision_transformer.py b/kimm/models/vision_transformer.py
index 5da247d..82f757b 100644
--- a/kimm/models/vision_transformer.py
+++ b/kimm/models/vision_transformer.py
@@ -23,6 +23,8 @@ def __init__(
         pos_dropout_rate: float = 0.0,
         **kwargs,
     ):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         input_tensor = kwargs.pop("input_tensor", None)
         self.set_properties(kwargs, 384)
         if self._pooling is not None:
@@ -100,10 +102,6 @@ def build_top(self, inputs, classes, classifier_activation, dropout_rate):
         )(x)
         return x
 
-    @staticmethod
-    def available_feature_keys():
-        raise NotImplementedError()
-
     def get_config(self):
         config = super().get_config()
         config.update(
@@ -142,6 +140,18 @@ def fix_config(self, config):
 
 
 class VisionTransformerTiny16(VisionTransformer):
+    available_feature_keys = [
+        "EMBEDDING",
+        *[f"BLOCK{i}" for i in range(12)],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            VisionTransformer.default_origin,
+            "visiontransformertiny16_vit_tiny_patch16_384.keras",
+        )
+    ]
+
     def __init__(
         self,
         mlp_ratio: float = 4.0,
@@ -161,9 +171,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "visiontransformertiny16_vit_tiny_patch16_384.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             16,
             192,
@@ -186,14 +193,14 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["EMBEDDING"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(12)])
-        return feature_keys
-
 
 class VisionTransformerTiny32(VisionTransformer):
+    available_feature_keys = [
+        "EMBEDDING",
+        *[f"BLOCK{i}" for i in range(12)],
+    ]
+    available_weights = []
+
     def __init__(
         self,
         mlp_ratio: float = 4.0,
@@ -235,14 +242,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["EMBEDDING"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(12)])
-        return feature_keys
-
 
 class VisionTransformerSmall16(VisionTransformer):
+    available_feature_keys = [
+        "EMBEDDING",
+        *[f"BLOCK{i}" for i in range(12)],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            VisionTransformer.default_origin,
+            "visiontransformersmall16_vit_small_patch16_384.keras",
+        )
+    ]
+
     def __init__(
         self,
         mlp_ratio: float = 4.0,
@@ -262,9 +275,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "visiontransformersmall16_vit_small_patch16_384.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             16,
             384,
@@ -287,14 +297,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["EMBEDDING"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(12)])
-        return feature_keys
-
 
 class VisionTransformerSmall32(VisionTransformer):
+    available_feature_keys = [
+        "EMBEDDING",
+        *[f"BLOCK{i}" for i in range(12)],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            VisionTransformer.default_origin,
+            "visiontransformersmall32_vit_small_patch32_384.keras",
+        )
+    ]
+
     def __init__(
         self,
         mlp_ratio: float = 4.0,
@@ -314,9 +330,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "visiontransformersmall32_vit_small_patch32_384.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             32,
             384,
@@ -339,14 +352,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["EMBEDDING"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(12)])
-        return feature_keys
-
 
 class VisionTransformerBase16(VisionTransformer):
+    available_feature_keys = [
+        "EMBEDDING",
+        *[f"BLOCK{i}" for i in range(12)],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            VisionTransformer.default_origin,
+            "visiontransformerbase16_vit_base_patch16_384.keras",
+        )
+    ]
+
     def __init__(
         self,
         mlp_ratio: float = 4.0,
@@ -366,9 +385,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "visiontransformerbase16_vit_base_patch16_384.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             16,
             768,
@@ -391,14 +407,20 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["EMBEDDING"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(12)])
-        return feature_keys
-
 
 class VisionTransformerBase32(VisionTransformer):
+    available_feature_keys = [
+        "EMBEDDING",
+        *[f"BLOCK{i}" for i in range(12)],
+    ]
+    available_weights = [
+        (
+            "imagenet",
+            VisionTransformer.default_origin,
+            "visiontransformerbase32_vit_base_patch32_384.keras",
+        )
+    ]
+
     def __init__(
         self,
         mlp_ratio: float = 4.0,
@@ -418,9 +440,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "visiontransformerbase32_vit_base_patch32_384.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             32,
             768,
@@ -443,14 +462,14 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["EMBEDDING"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(12)])
-        return feature_keys
-
 
 class VisionTransformerLarge16(VisionTransformer):
+    available_feature_keys = [
+        "EMBEDDING",
+        *[f"BLOCK{i}" for i in range(24)],
+    ]
+    available_weights = []
+
     def __init__(
         self,
         mlp_ratio: float = 4.0,
@@ -492,14 +511,14 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["EMBEDDING"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(24)])
-        return feature_keys
-
 
 class VisionTransformerLarge32(VisionTransformer):
+    available_feature_keys = [
+        "EMBEDDING",
+        *[f"BLOCK{i}" for i in range(24)],
+    ]
+    available_weights = []
+
     def __init__(
         self,
         mlp_ratio: float = 4.0,
@@ -541,12 +560,6 @@ def __init__(
             **kwargs,
         )
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["EMBEDDING"]
-        feature_keys.extend([f"BLOCK{i}" for i in range(24)])
-        return feature_keys
-
 
 add_model_to_registry(VisionTransformerTiny16, "imagenet")
 add_model_to_registry(VisionTransformerTiny32)
diff --git a/kimm/models/xception.py b/kimm/models/xception.py
index 21fcaa6..b9f743d 100644
--- a/kimm/models/xception.py
+++ b/kimm/models/xception.py
@@ -66,7 +66,14 @@ def apply_xception_block(
 
 @keras.saving.register_keras_serializable(package="kimm")
 class XceptionBase(BaseModel):
+    available_feature_keys = [
+        "STEM_S2",
+        *[f"BLOCK{i}_S{j}" for i, j in zip(range(4), [4, 8, 16, 32])],
+    ]
+
     def __init__(self, **kwargs):
+        kwargs["weights_url"] = self.get_weights_url(kwargs["weights"])
+
         input_tensor = kwargs.pop("input_tensor", None)
         self.set_properties(kwargs)
         inputs = self.determine_input_tensor(
@@ -133,14 +140,6 @@ def __init__(self, **kwargs):
 
         # All references to `self` below this line
 
-    @staticmethod
-    def available_feature_keys():
-        feature_keys = ["STEM_S2"]
-        feature_keys.extend(
-            [f"BLOCK{i}_S{j}" for i, j in zip(range(4), [4, 8, 16, 32])]
-        )
-        return feature_keys
-
     def get_config(self):
         return super().get_config()
 
@@ -154,6 +153,14 @@ def fix_config(self, config: typing.Dict):
 
 
 class Xception(XceptionBase):
+    available_weights = [
+        (
+            "imagenet",
+            XceptionBase.default_origin,
+            "xception.keras",
+        )
+    ]
+
     def __init__(
         self,
         input_tensor: keras.KerasTensor = None,
@@ -169,9 +176,6 @@ def __init__(
         **kwargs,
     ):
         kwargs = self.fix_config(kwargs)
-        if weights == "imagenet":
-            file_name = "xception.keras"
-            kwargs["weights_url"] = f"{self.default_origin}/{file_name}"
         super().__init__(
             input_tensor=input_tensor,
             input_shape=input_shape,
diff --git a/kimm/utils/model_registry.py b/kimm/utils/model_registry.py
index a80b93f..5f5fc40 100644
--- a/kimm/utils/model_registry.py
+++ b/kimm/utils/model_registry.py
@@ -48,7 +48,7 @@ def add_model_to_registry(model_cls, weights: typing.Optional[str] = None):
     feature_keys = []
     if issubclass(model_cls, BaseModel):
         feature_extractor = True
-        feature_keys = model_cls.available_feature_keys()
+        feature_keys = model_cls.available_feature_keys
     for info in MODEL_REGISTRY:
         if info["name"] == model_cls.__name__:
             warnings.warn(


Test patch:
diff --git a/kimm/models/models_test.py b/kimm/models/models_test.py
index 2e3bc41..c5bfddf 100644
--- a/kimm/models/models_test.py
+++ b/kimm/models/models_test.py
@@ -379,7 +379,7 @@ def test_model_feature_extractor(
 
         self.assertIsInstance(y, dict)
         self.assertContainsSubset(
-            model_class.available_feature_keys(), list(y.keys())
+            model_class.available_feature_keys, list(y.keys())
         )
         for feature_info in features:
             name, shape = feature_info
diff --git a/kimm/utils/model_registry_test.py b/kimm/utils/model_registry_test.py
index 7e3d592..93a1047 100644
--- a/kimm/utils/model_registry_test.py
+++ b/kimm/utils/model_registry_test.py
@@ -13,9 +13,7 @@ class DummyModel(models.Model):
 
 
 class DummyFeatureExtractor(BaseModel):
-    @staticmethod
-    def available_feature_keys():
-        return ["A", "B", "C"]
+    available_feature_keys = ["A", "B", "C"]
 
 
 class ModelRegistryTest(testing.TestCase):