change: Address review

Author: harshaljanjani

src/transformers/models/omdet_turbo/configuration_omdet_turbo.py

@@ -149,7 +149,8 @@ def __init__(
         # Init timm backbone with hardcoded values for BC
         timm_default_kwargs = {
             "out_indices": [1, 2, 3],
-            "timm_model_kwargs": {"img_size": image_size, "always_partition": True},
+            "img_size": image_size,
+            "always_partition": True,
         }
         backbone_config, kwargs = consolidate_backbone_kwargs_to_config(
             backbone_config=backbone_config,
@@ -160,16 +161,11 @@ def __init__(
             **kwargs,
         )
 
-        # BC: hub configs with backbone_kwargs pass img_size/always_partition as direct attrs
-        if getattr(backbone_config, "model_type", None) == "timm_backbone" and not getattr(
-            backbone_config, "timm_model_kwargs", None
-        ):
-            timm_extra = {}
+        timm_kwargs = {}
+        if getattr(backbone_config, "model_type", None) == "timm_backbone":
             for attr in ("img_size", "always_partition"):
                 if hasattr(backbone_config, attr):
-                    timm_extra[attr] = getattr(backbone_config, attr)
-            if timm_extra:
-                backbone_config.timm_model_kwargs = timm_extra
+                    timm_kwargs[attr] = getattr(backbone_config, attr)
 
         if text_config is None:
             logger.info("`text_config` is `None`. Initializing the config with the default `clip_text_model`")
@@ -222,8 +218,14 @@ def __init__(
         self.eval_size = eval_size
         self.learn_initial_query = learn_initial_query
         self.cache_size = cache_size
+        self.timm_kwargs = timm_kwargs
 
         super().__init__(is_encoder_decoder=is_encoder_decoder, **kwargs)
 
+    def to_dict(self):
+        output = super().to_dict()
+        output.pop("timm_kwargs", None)
+        return output
+
 
 __all__ = ["OmDetTurboConfig"]

src/transformers/models/omdet_turbo/modeling_omdet_turbo.py

@@ -26,7 +26,6 @@
 
 from ... import initialization as init
 from ...activations import ACT2CLS, ACT2FN
-from ...backbone_utils import load_backbone
 from ...integrations import use_kernel_forward_from_hub
 from ...masking_utils import create_bidirectional_mask
 from ...modeling_layers import GradientCheckpointingLayer
@@ -39,7 +38,7 @@
     logging,
     torch_compilable_check,
 )
-from ..auto import AutoModel
+from ..auto import AutoBackbone, AutoModel
 from .configuration_omdet_turbo import OmDetTurboConfig
 
 
@@ -279,7 +278,7 @@ class OmDetTurboVisionBackbone(nn.Module):
     def __init__(self, config: OmDetTurboConfig):
         super().__init__()
         self.apply_layernorm_after_vision_backbone = config.apply_layernorm_after_vision_backbone
-        self.vision_backbone = load_backbone(config)
+        self.vision_backbone = AutoBackbone.from_config(config.backbone_config, **getattr(config, "timm_kwargs", {}))
         self.layer_norms = nn.ModuleList(
             [nn.LayerNorm(in_channel_dim, eps=config.layer_norm_eps) for in_channel_dim in config.encoder_in_channels]
         )

src/transformers/models/timm_backbone/modeling_timm_backbone.py

@@ -49,7 +49,6 @@ def __init__(self, config, **kwargs):
         out_indices = config.out_indices if getattr(config, "out_indices", None) is not None else (-1,)
         pretrained = kwargs.pop("pretrained", False)
         in_chans = kwargs.pop("in_chans", config.num_channels)
-        timm_model_kwargs = getattr(config, "timm_model_kwargs", {})
 
         backbone = timm.create_model(
             config.backbone,
@@ -59,7 +58,6 @@ def __init__(self, config, **kwargs):
             in_chans=in_chans,
             out_indices=out_indices,
             output_stride=config.output_stride,
-            **timm_model_kwargs,
             **kwargs,
         )