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

modeling_tips.py

@@ -1,34 +1,23 @@
 """TIPSv2 model for HuggingFace — wraps vision and text encoders."""
 
-import importlib
-import os
 from dataclasses import dataclass
-from pathlib import Path
 from typing import List, Optional, Union
 
-import numpy as np
 import torch
-from huggingface_hub import hf_hub_download
 from transformers import PreTrainedModel
+from transformers.utils import cached_file
 
 from .configuration_tips import TIPSv2Config
+from .image_encoder import vit_base, vit_giant2, vit_large, vit_small, vit_so400m
+from .text_encoder import TextEncoder, Tokenizer
 
-_this_dir = Path(__file__).parent
-_sibling_cache = {}
-
-
-def _load_sibling(name, repo_id=None):
-    """Import a sibling .py from the same dir, downloading from HF if needed."""
-    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
+_VISION_FACTORIES = {
+    "vit_small": vit_small,
+    "vit_base": vit_base,
+    "vit_large": vit_large,
+    "vit_so400m": vit_so400m,
+    "vit_giant2": vit_giant2,
+}
 
 
 @dataclass
@@ -75,12 +64,7 @@ class TIPSv2Model(PreTrainedModel):
     def __init__(self, config: TIPSv2Config):
         super().__init__(config)
 
-        repo_id = getattr(config, "_name_or_path", None)
-        ie = _load_sibling("image_encoder", repo_id)
-        te = _load_sibling("text_encoder", repo_id)
-
-        build_fn = getattr(ie, config.vision_fn)
-        self.vision_encoder = build_fn(
+        self.vision_encoder = _VISION_FACTORIES[config.vision_fn](
             img_size=config.img_size,
             patch_size=config.patch_size,
             ffn_layer=config.ffn_layer,
@@ -90,7 +74,7 @@ class TIPSv2Model(PreTrainedModel):
             interpolate_offset=0.0,
         )
 
-        self.text_encoder = te.TextEncoder(
+        self.text_encoder = TextEncoder(
             config={
                 "hidden_size": config.text_hidden_size,
                 "mlp_dim": config.text_mlp_dim,
@@ -101,14 +85,10 @@ class TIPSv2Model(PreTrainedModel):
         )
 
         self._tokenizer = None
-        self._te_mod = te
 
     def _load_tokenizer(self):
-        """Lazy-load the SentencePiece tokenizer."""
-        tok_path = _this_dir / "tokenizer.model"
-        if not tok_path.exists():
-            tok_path = hf_hub_download(self.name_or_path, "tokenizer.model")
-        return self._te_mod.Tokenizer(str(tok_path))
+        """Load the SentencePiece tokenizer shipped with the checkpoint."""
+        return Tokenizer(cached_file(self.name_or_path, "tokenizer.model"))
 
     @torch.no_grad()
     def encode_image(self, pixel_values: torch.Tensor) -> TIPSv2ImageOutput: