bvandeusen
4daa3f2790
Make the SigLIP embedder an operator choice (drop-in to SigLIP 2:
google/siglip2-so400m-patch16-512 is a verified 1152-d model at 512px → no
schema change, better small-cue fidelity). A swap = set model + re-embed +
retrain, all operator-driven; the GPU agent does the re-embed so it's fast.
- settings: embedder_model_name is now a setting (migration 0065) alongside the
existing embedder_model_version; both editable + validated (non-empty) in the
ml admin API. The server embedder loads by HF name (AutoImageProcessor/Model,
model-agnostic), preferring the pre-downloaded local dir for the default so
existing deploys don't re-download; rebuilds on a name change.
- agent: new 'embed' job = whole-image SigLIP embedding (mean-pool video frames)
under the lease-announced model → POST /jobs/submit_embedding writes
image_record.siglip_embedding + siglip_model_version. The lease now announces
the model FROM THE SETTING (not a constant).
- re-embed routing: enqueue_gpu_backfill('embed') selects unembedded + stale-
version images; 'siglip' now re-embeds concept crops whose version != current
(so a swap re-triggers crops, not just the never-embedded back-catalogue). The
CPU ml-worker backfill no longer re-embeds on a version mismatch (it can't
churn the library at 512px) — the GPU agent owns version re-embeds. Daily
'embed' + 'siglip' beats self-heal.
- scoring: score_image only bags embeddings in the CURRENT model's space (whole-
image gated by siglip_model_version, concept regions by embedding_version) so a
mid-swap stale vector isn't scored by new-space heads; legacy NULL = current.
- UI: GpuAgentCard "Embedding model (advanced)" — edit name/version, Save, and
"Re-embed library (GPU)" (queues embed + siglip); points at SigLIP 2.
Tests: lease announces model + submit_embedding round-trip; enqueue 'embed'
selects stale/unembedded; stale-version excluded from scoring; embedder model
settable + empty rejected; siglip gate updated to current-version concept.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
95 lines
3.7 KiB
Python
95 lines
3.7 KiB
Python
"""SigLIP SO400M image-embedding wrapper (PyTorch CPU).
|
||
|
||
torch/transformers are imported lazily inside load() so this module can be
|
||
imported in the web container (which never runs inference) without paying the
|
||
torch import cost.
|
||
"""
|
||
|
||
import os
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
from PIL import Image, ImageFile
|
||
|
||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||
|
||
# Cap torch's intra-op threads so each ml-worker replica is a bounded core
|
||
# consumer on a shared node (torch otherwise uses all cores). Keep
|
||
# N_replicas × this within the cores allotted to ML to avoid oversubscription.
|
||
_INTRA_OP_THREADS = 4
|
||
|
||
DEFAULT_MODEL_NAME = os.environ.get(
|
||
"SIGLIP_MODEL_NAME", "google/siglip-so400m-patch14-384"
|
||
)
|
||
# Back-compat alias (api/gpu imported this name as the fallback embedder id).
|
||
MODEL_NAME = DEFAULT_MODEL_NAME
|
||
MODEL_VERSION = os.environ.get(
|
||
"SIGLIP_MODEL_VERSION", "siglip-so400m-patch14-384"
|
||
)
|
||
EMBED_DIM = 1152
|
||
_LOCAL_DIR = Path(os.environ.get("ML_MODEL_DIR", "/models")) / "siglip"
|
||
|
||
|
||
class Embedder:
|
||
"""Loads whatever SigLIP-family model it's given by HF NAME. For the default
|
||
model it prefers the pre-downloaded local dir (no re-download on existing
|
||
deploys); any other name resolves as an HF repo id (downloaded + cached on
|
||
first use), so an operator model swap (#1190) just works server-side."""
|
||
|
||
def __init__(self, model_name: str | None = None, model_dir: Path | None = None):
|
||
self.model_name = model_name or DEFAULT_MODEL_NAME
|
||
self._explicit_dir = model_dir
|
||
self._model = None
|
||
self._processor = None
|
||
self._torch = None
|
||
|
||
def _source(self) -> str:
|
||
if self._explicit_dir is not None:
|
||
return str(self._explicit_dir)
|
||
if self.model_name == DEFAULT_MODEL_NAME and _LOCAL_DIR.exists():
|
||
return str(_LOCAL_DIR)
|
||
return self.model_name
|
||
|
||
def load(self) -> None:
|
||
if self._model is not None:
|
||
return
|
||
import torch
|
||
from transformers import AutoImageProcessor, AutoModel
|
||
|
||
self._torch = torch
|
||
# Bound torch's CPU thread pool (see _INTRA_OP_THREADS) so each replica
|
||
# stays a predictable core consumer on a shared node.
|
||
torch.set_num_threads(_INTRA_OP_THREADS)
|
||
# IMAGE inference only — AutoImageProcessor loads just the image side
|
||
# (preprocessor_config.json), skipping the SigLIP tokenizer + its
|
||
# sentencepiece dep (operator hit that ImportError 2026-05-25). Works
|
||
# for any SigLIP-family model, keeping the embedder model-agnostic.
|
||
src = self._source()
|
||
self._processor = AutoImageProcessor.from_pretrained(src)
|
||
self._model = AutoModel.from_pretrained(src)
|
||
self._model.eval()
|
||
|
||
def infer(self, image_path: Path) -> np.ndarray:
|
||
"""Return a 1152-dim float32 embedding (SigLIP MAP-pooled output)."""
|
||
self.load()
|
||
img = Image.open(image_path).convert("RGB")
|
||
with self._torch.no_grad():
|
||
inputs = self._processor(images=img, return_tensors="pt")
|
||
out = self._model.get_image_features(**inputs)
|
||
pooled = out.pooler_output if hasattr(out, "pooler_output") else out
|
||
return pooled[0].numpy().astype(np.float32)
|
||
|
||
|
||
_default_embedder: Embedder | None = None
|
||
|
||
|
||
def get_embedder(model_name: str | None = None) -> Embedder:
|
||
"""Cached embedder for `model_name` (default if None). Rebuilds the singleton
|
||
when the requested name changes, so an operator model swap takes effect
|
||
without restarting the worker."""
|
||
global _default_embedder
|
||
name = model_name or DEFAULT_MODEL_NAME
|
||
if _default_embedder is None or _default_embedder.model_name != name:
|
||
_default_embedder = Embedder(model_name=name)
|
||
return _default_embedder
|