Update files

modeling_dpt.py

@@ -8,7 +8,7 @@ from typing import Optional
 
 import torch
 from huggingface_hub import hf_hub_download
-from transformers import AutoModel, PreTrainedModel
+from transformers import AutoConfig, AutoModel, PreTrainedModel
 
 from .configuration_dpt import TIPSv2DPTConfig
 
@@ -61,7 +61,6 @@ class TIPSv2DPTModel(PreTrainedModel):
     _no_split_modules = []
     _supports_cache_class = False
     _tied_weights_keys = []
-    _keys_to_ignore_on_load_unexpected = {"backbone"}
 
     @property
     def all_tied_weights_keys(self):
@@ -75,6 +74,10 @@ class TIPSv2DPTModel(PreTrainedModel):
 
         ppc = tuple(config.post_process_channels)
 
+        backbone_config = AutoConfig.from_pretrained(config.backbone_repo, trust_remote_code=True)
+        backbone = AutoModel.from_config(backbone_config, trust_remote_code=True)
+        self.vision_encoder = backbone.vision_encoder
+
         self.depth_head = dpt_mod.DPTDepthHead(
             input_embed_dim=config.embed_dim, channels=config.channels,
             post_process_channels=ppc, readout_type=config.readout_type,
@@ -90,17 +93,9 @@ class TIPSv2DPTModel(PreTrainedModel):
             post_process_channels=ppc, readout_type=config.readout_type,
             num_classes=config.num_seg_classes,
         )
-        self._backbone = None
-
-    def _get_backbone(self):
-        if self._backbone is None:
-            self._backbone = AutoModel.from_pretrained(self.config.backbone_repo, trust_remote_code=True)
-            self._backbone.to(self.device).eval()
-        return self._backbone.vision_encoder
 
     def _extract_intermediate(self, pixel_values):
-        backbone = self._get_backbone()
-        intermediate = backbone.get_intermediate_layers(
+        intermediate = self.vision_encoder.get_intermediate_layers(
             pixel_values, n=self.config.block_indices,
             reshape=True, return_class_token=True, norm=True,
         )