feat(fc2b): add SigLIP embedder wrapper
Direct port of ImageRepo's siglip.py. Lazy torch/transformers import so the web container can import the module (for enqueue logic) without the torch cost. EMBED_DIM=1152 asserted against the schema's Vector(1152) columns. Real inference runs in the local integration suite. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,62 @@
|
|||||||
|
"""SigLIP SO400M image-embedding wrapper (PyTorch CPU).
|
||||||
|
|
||||||
|
Direct port of ImageRepo's siglip.py. 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
|
||||||
|
|
||||||
|
MODEL_NAME = os.environ.get(
|
||||||
|
"SIGLIP_MODEL_NAME", "google/siglip-so400m-patch14-384"
|
||||||
|
)
|
||||||
|
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:
|
||||||
|
def __init__(self, model_dir: Path | None = None):
|
||||||
|
self._model_dir = model_dir or _LOCAL_DIR
|
||||||
|
self._model = None
|
||||||
|
self._processor = None
|
||||||
|
self._torch = None
|
||||||
|
|
||||||
|
def load(self) -> None:
|
||||||
|
if self._model is not None:
|
||||||
|
return
|
||||||
|
import torch
|
||||||
|
from transformers import AutoModel, AutoProcessor
|
||||||
|
|
||||||
|
self._torch = torch
|
||||||
|
self._processor = AutoProcessor.from_pretrained(str(self._model_dir))
|
||||||
|
self._model = AutoModel.from_pretrained(str(self._model_dir))
|
||||||
|
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() -> Embedder:
|
||||||
|
global _default_embedder
|
||||||
|
if _default_embedder is None:
|
||||||
|
_default_embedder = Embedder()
|
||||||
|
return _default_embedder
|
||||||
@@ -0,0 +1,30 @@
|
|||||||
|
"""Embedder unit tests. The SigLIP model is a multi-GB download not present
|
||||||
|
in CI, so these test only the import-safety contract and the singleton.
|
||||||
|
Real embedding runs in the local integration suite.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from backend.app.services.ml.embedder import (
|
||||||
|
EMBED_DIM,
|
||||||
|
MODEL_VERSION,
|
||||||
|
Embedder,
|
||||||
|
get_embedder,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_embed_dim_matches_schema():
|
||||||
|
# image_record.siglip_embedding and tag_reference_embedding.embedding
|
||||||
|
# are both Vector(1152). This must stay in sync.
|
||||||
|
assert EMBED_DIM == 1152
|
||||||
|
|
||||||
|
|
||||||
|
def test_model_version_default():
|
||||||
|
assert MODEL_VERSION == "siglip-so400m-patch14-384"
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_embedder_singleton():
|
||||||
|
assert get_embedder() is get_embedder()
|
||||||
|
|
||||||
|
|
||||||
|
def test_not_loaded_until_called():
|
||||||
|
e = Embedder()
|
||||||
|
assert e._model is None # not loaded until .load()
|
||||||
Reference in New Issue
Block a user