bvandeusen
60a9c9e6ef
GPU enablement (#872) cancelled — not worth the Pascal-specific build for a modest CPU→GPU win on an old P4. Remove the dead GPU code (device.py, the CUDA provider branch in tagger, the .to('cuda') path in embedder) so nothing carries it forward. Instead, bound CPU inference threads by default so the ml-worker is a predictable core consumer on a SHARED node — the intended scaling model is multiple worker replicas (each --concurrency=1, each its own cgroup limit), not one big container. ONNX Runtime and torch otherwise size their thread pools to ALL host cores, so each replica would grab every core and oversubscribe / starve the co-located DB+web. Cap both to _INTRA_OP_THREADS=4 (matches the prior per-worker cpus:4 unit): run N replicas where N×4 stays within the cores allotted to ML. - tagger: ort.SessionOptions().intra_op_num_threads = 4 (CPUExecutionProvider). - embedder: torch.set_num_threads(4). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
82 lines
3.0 KiB
Python
82 lines
3.0 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
|
||
|
||
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, SiglipImageProcessor
|
||
|
||
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)
|
||
# FC's embedder only does IMAGE inference — never text. AutoProcessor
|
||
# loads the full processor including SiglipTokenizer, which requires
|
||
# the sentencepiece library at import time even if we never call it.
|
||
# SiglipImageProcessor loads ONLY preprocessor_config.json (image
|
||
# side) and skips the tokenizer config entirely. Operator hit the
|
||
# ImportError 2026-05-25 once the ml-worker started actually running
|
||
# tag_and_embed; switching to the image-only loader avoids the
|
||
# tokenizer dep without adding ~30 MB of unused C++ build to the
|
||
# lean ml-worker image.
|
||
self._processor = SiglipImageProcessor.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
|