fix: refuse unavailable gallery and pin head trainer

This commit is contained in:
kiyo9w 2026-09-28 14:41:57 +09:00
parent d5a7a56dff
commit ac315eb48e
4 changed files with 77 additions and 6 deletions

View file

@ -132,12 +132,14 @@ def enrollment_report_text(report: dict) -> str:
def build_model(image_size: int = 518):
import timm
import torch
from gallery import load_dinov2
device = "mps" if torch.backends.mps.is_available() else "cpu"
model = timm.create_model("vit_small_patch14_dinov2.lvd142m", pretrained=True,
num_classes=0, img_size=image_size, dynamic_img_size=True)
# Same hash-pinned local file as inference. pretrained=True would fetch a
# different copy even with YOLO_OFFLINE=1.
model = load_dinov2(image_size)
return model.eval().to(device), device