Use relative imports for sibling modules (fixes local loading, save_pretrained, pickling)

modeling_dpt.py

@@ -1,32 +1,13 @@
 """TIPSv2 DPT dense prediction model for HuggingFace."""
 
-import importlib
-import os
 from dataclasses import dataclass
-from pathlib import Path
 from typing import Optional
 
 import torch
-from huggingface_hub import hf_hub_download
 from transformers import AutoConfig, AutoModel, PreTrainedModel
 
 from .configuration_dpt import TIPSv2DPTConfig
-
-_this_dir = Path(__file__).parent
-_sibling_cache = {}
-
-
-def _load_sibling(name, repo_id=None):
-    if name in _sibling_cache:
-        return _sibling_cache[name]
-    path = _this_dir / f"{name}.py"
-    if not path.exists() and repo_id:
-        path = Path(hf_hub_download(repo_id, f"{name}.py"))
-    spec = importlib.util.spec_from_file_location(name, str(path))
-    mod = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(mod)
-    _sibling_cache[name] = mod
-    return mod
+from .dpt_head import DPTDepthHead, DPTNormalsHead, DPTSegmentationHead
 
 
 @dataclass
@@ -69,26 +50,23 @@ class TIPSv2DPTModel(PreTrainedModel):
     def __init__(self, config: TIPSv2DPTConfig):
         super().__init__(config)
 
-        repo_id = getattr(config, "_name_or_path", None)
-        dpt_mod = _load_sibling("dpt_head", repo_id)
-
         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(
+        self.depth_head = DPTDepthHead(
             input_embed_dim=config.embed_dim, channels=config.channels,
             post_process_channels=ppc, readout_type=config.readout_type,
             num_depth_bins=config.num_depth_bins,
             min_depth=config.min_depth, max_depth=config.max_depth,
         )
-        self.normals_head = dpt_mod.DPTNormalsHead(
+        self.normals_head = DPTNormalsHead(
             input_embed_dim=config.embed_dim, channels=config.channels,
             post_process_channels=ppc, readout_type=config.readout_type,
         )
-        self.segmentation_head = dpt_mod.DPTSegmentationHead(
+        self.segmentation_head = DPTSegmentationHead(
             input_embed_dim=config.embed_dim, channels=config.channels,
             post_process_channels=ppc, readout_type=config.readout_type,
             num_classes=config.num_seg_classes,