Spaces:
Sleeping
Sleeping
File size: 3,795 Bytes
a1dd5ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """DINOv3 — the teacher's substrate. One backbone, three granularities.
Frames (every frame -> scene series), track crops (identity descriptors)
and, later, whatever the FDNN student distills from. Chosen over the
CLIP family for everything appearance-shaped because it has NO text
tower: it was never trained to collapse instances into nameable
categories, which is precisely the failure that made a red and a green
pepper the same object under text-aligned encoders. The instance-level
numbers agree (+10.9 GAP retrieval over DINOv2, which already beat CLIP
at instance top-1).
Gated weights: the user accepted Meta's license on HF (2026-08-01).
No resizing cleverness: everything goes to IMG px square through the
processor, because identity compares crops to crops and frames to
frames — never one to the other — so a shared canonical size keeps
each comparison internally consistent.
"""
from __future__ import annotations
import os
import numpy as np
_M = {}
# ConvNeXt-Tiny, not ViT-S, and the reason is numerical, not taste:
# the DINOv3 ViT variants OVERFLOW FP16 on MPS - every embedding NaN,
# and a NaN loses every >= comparison silently, so the failure mode was
# not an error but "12,583 tracks, 12,583 objects, zero recurrence".
# ViT-S works in fp32 at ~2x its fp16 time, which makes it slower than
# ConvNeXt-Tiny in fp16 - and ConvNeXt measured AUC 0.9994 on the
# corpus's proven pairs (vs yolo26n-reid 0.9047). Faster AND stronger,
# so the 768-d (vs 384) table cost is accepted.
MID = os.environ.get("ELIDEDB_DINOV3",
"facebook/dinov3-convnext-tiny-pretrain-lvd1689m")
IMG = int(os.environ.get("ELIDEDB_DINOV3_SZ", "224"))
BATCH = int(os.environ.get("ELIDEDB_DINOV3_BATCH", "64"))
def _load(mid=None):
mid = mid or MID
if _M.get("mid") == mid:
return _M
import torch
from transformers import AutoImageProcessor, AutoModel
from .device import pick
dev, dtype = pick()
if "vit" in mid.rsplit("/", 1)[-1]:
dtype = torch.float32 # fp16 overflow, see above
_M.clear()
_M["proc"] = AutoImageProcessor.from_pretrained(mid)
_M["model"] = AutoModel.from_pretrained(
mid, dtype=dtype, low_cpu_mem_usage=True).to(dev).eval()
_M["dev"], _M["dtype"], _M["mid"] = dev, dtype, mid
_M["torch"] = torch
return _M
def embed(images, batch=None, mid=None):
"""(N, D) unit-norm float32 embeddings for a list of HWC uint8 arrays.
Pooled (CLS) output — the global instance/scene vector, the thing
retrieval compares. Patch tokens exist but are a different product
for a different consumer; nothing here should quietly average them.
"""
m = _load(mid)
torch = m["torch"]
out = []
batch = batch or BATCH
for i in range(0, len(images), batch):
chunk = [np.ascontiguousarray(x[..., :3]) for x in
images[i:i + batch]]
px = m["proc"](images=chunk, return_tensors="pt",
size={"height": IMG, "width": IMG})
px = px["pixel_values"].to(m["dev"], m["dtype"])
with torch.no_grad():
r = m["model"](pixel_values=px)
v = r.pooler_output.float()
v = v / v.norm(dim=-1, keepdim=True).clamp_min(1e-8)
v = v.cpu().numpy().astype(np.float32)
# FAIL LOUD. The ViT fp16 overflow produced NaN vectors that
# scored 0.0 AUC without a single exception - every comparison
# quietly False. Garbage must stop the run, not grade it.
if np.isnan(v).any():
raise FloatingPointError(
f"{_M['mid']}: NaN embeddings (dtype {_M['dtype']}) - "
"use fp32 for ViT variants")
out.append(v)
return (np.concatenate(out) if out
else np.zeros((0, 384), np.float32))
|