fix: load hash-pinned DINOv2 weights without network

recognization/.gitignore

@@ -19,6 +19,7 @@ ENV/
 
 # Temporary weights
 runs/detect/train/weights/last.pt
+weights/dinov2_vits14_lvd142m.safetensors
 *.log
 stock-*/identity_head.pt
 .demo-token

recognization/code/gallery.py

@@ -38,6 +38,78 @@ DEFAULT_THRESHOLD = 0.55
 # The same line for a fine-tuned head, whose similarity scale is its own.
 DEFAULT_HEAD_THRESHOLD = 0.40
 
+# DINOv2 ViT-S/14, LVD-142M. timm's default source is the Hugging Face copy of
+# these weights (Apache-2.0); pretrained=True fetches it even when YOLO_OFFLINE=1.
+# The pin is the sha256 of that model.safetensors (timm hub snapshot 4610ca14).
+# Not in Git (~84 MB). Place it at weights/dinov2_vits14_lvd142m.safetensors, or
+# point DORAYAKI_DINOV2_WEIGHTS at a copy. See README "DINOv2 weights".
+DINOV2_MODEL = "vit_small_patch14_dinov2.lvd142m"
+DINOV2_SHA256 = "04d27f3400d059fc0cfd7d17dd1909a75bf3ea8fb3eeb48b97cb99e57ee20081"
+DINOV2_FILENAME = "dinov2_vits14_lvd142m.safetensors"
+
+
+def dinov2_weights_path() -> str:
+    override = os.environ.get("DORAYAKI_DINOV2_WEIGHTS")
+    if override:
+        return override
+    here = os.path.dirname(os.path.abspath(__file__))
+    return os.path.join(os.path.dirname(here), "weights", DINOV2_FILENAME)
+
+
+def _sha256(path: str) -> str:
+    import hashlib
+
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def load_dinov2(image_size: int) -> Any:
+    """Build the gallery backbone from the pinned local file. Never downloads.
+
+    A missing or altered file raises FileNotFoundError naming the path and the
+    expected sha256. The head in finetune_identity.py is a Linear(384, 128) on
+    this model's features; a different backbone will not load that head.
+    """
+    import timm
+    from safetensors.torch import load_file
+    from timm.models.vision_transformer import checkpoint_filter_fn
+
+    path = dinov2_weights_path()
+    if not os.path.isfile(path):
+        raise FileNotFoundError(
+            f"DINOv2 weights not found at {path}. "
+            f"Copy model.safetensors (sha256 {DINOV2_SHA256}) from the local "
+            f"Hugging Face cache of timm/{DINOV2_MODEL}, or set DORAYAKI_DINOV2_WEIGHTS. "
+            f"The gallery does not download weights."
+        )
+    digest = _sha256(path)
+    if digest != DINOV2_SHA256:
+        raise FileNotFoundError(
+            f"DINOv2 weights at {path} are sha256 {digest}, expected {DINOV2_SHA256}."
+        )
+    model = timm.create_model(
+        DINOV2_MODEL,
+        pretrained=False,
+        num_classes=0,
+        img_size=int(image_size),
+        dynamic_img_size=True,
+    )
+    # The file is the 518px checkpoint. timm's own filter resizes pos_embed when
+    # the caller asks for another size; pretrained=True would do this and also
+    # download, so it stays off.
+    state = checkpoint_filter_fn(load_file(path), model)
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    unexpected = [key for key in unexpected if not key.startswith("head.")]
+    if missing or unexpected:
+        raise RuntimeError(
+            f"DINOv2 weights at {path} do not match {DINOV2_MODEL} "
+            f"(missing={missing}, unexpected={unexpected})."
+        )
+    return model
+
 
 @dataclass
 class GalleryEntry:
@@ -138,13 +210,7 @@ class ProductGallery:
         device = self.device
         if device is None:
             device = "mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu"
-        model = timm.create_model(
-            "vit_small_patch14_dinov2.lvd142m",
-            pretrained=True,
-            num_classes=0,
-            img_size=self.image_size,
-            dynamic_img_size=True,
-        )
+        model = load_dinov2(self.image_size)
         model = model.eval().to(device)
         self.device = device
         self._model = model

recognization/code/test_gallery.py

@@ -21,7 +21,13 @@ if os.path.exists(os.path.join(project_root, "stock")):
 if current_dir not in sys.path:
     sys.path.insert(0, current_dir)
 
-from gallery import GalleryEntry, ProductGallery  # noqa: E402
+from gallery import (  # noqa: E402
+    DINOV2_SHA256,
+    GalleryEntry,
+    ProductGallery,
+    dinov2_weights_path,
+    load_dinov2,
+)
 
 
 def a_gallery(views, threshold=0.55):
@@ -108,5 +114,35 @@ class EnrollmentTests(unittest.TestCase):
             self.assertEqual(len(gallery._views_for("sku_ab")), 1)
 
 
+
+class OfflineWeightsTests(unittest.TestCase):
+    def test_a_missing_weights_file_names_the_path_and_does_not_fetch(self):
+        previous = os.environ.get("DORAYAKI_DINOV2_WEIGHTS")
+        os.environ["DORAYAKI_DINOV2_WEIGHTS"] = os.path.join(
+            "does-not-exist", "dinov2_vits14_lvd142m.safetensors"
+        )
+        try:
+            with self.assertRaises(FileNotFoundError) as raised:
+                load_dinov2(224)
+        finally:
+            if previous is None:
+                os.environ.pop("DORAYAKI_DINOV2_WEIGHTS", None)
+            else:
+                os.environ["DORAYAKI_DINOV2_WEIGHTS"] = previous
+        message = str(raised.exception)
+        self.assertIn("does-not-exist", message)
+        self.assertIn(DINOV2_SHA256, message)
+        self.assertIn("does not download", message)
+
+    def test_the_default_path_stays_inside_the_repo(self):
+        previous = os.environ.pop("DORAYAKI_DINOV2_WEIGHTS", None)
+        try:
+            self.assertTrue(dinov2_weights_path().endswith(
+                os.path.join("weights", "dinov2_vits14_lvd142m.safetensors")
+            ))
+        finally:
+            if previous is not None:
+                os.environ["DORAYAKI_DINOV2_WEIGHTS"] = previous
+
 if __name__ == "__main__":
     unittest.main()