bvandeusen
8af9f12544
PIL's strict load was raising OSError on images missing the trailing end-of-stream marker (e.g. '6 bytes not processed' on a JPEG without its FF D9 EOI), failing the entire ML task for an image the model could otherwise score fine. Set ImageFile.LOAD_TRUNCATED_IMAGES = True in both ML modules so a minutely-corrupt tail doesn't block tagging or embedding. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
63 lines
2.5 KiB
Python
63 lines
2.5 KiB
Python
"""SigLIP SO400M image-embedding wrapper (PyTorch CPU)."""
|
|
from __future__ import annotations
|
|
import os
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageFile
|
|
|
|
# Mirror wd14.py: tolerate minutely-truncated source images so embedding
|
|
# doesn't fail on the same images WD14 successfully tagged.
|
|
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
|
|
|
# Defer torch/transformers imports to lazily-loaded functions to allow
|
|
# importing MODEL_VERSION in non-ML-worker contexts (e.g., web container
|
|
# centroid-recompute enqueue logic).
|
|
_torch = None
|
|
_AutoModel = None
|
|
_AutoProcessor = None
|
|
|
|
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')
|
|
# Model files live flat under this directory (written by scripts/download_models.py via
|
|
# snapshot_download(local_dir=...)). We point from_pretrained at the local path directly
|
|
# so transformers bypasses its HF cache layout and doesn't need network access at load time.
|
|
_LOCAL_DIR = os.path.join(os.environ.get('ML_MODEL_DIR', '/models'), 'siglip')
|
|
|
|
_model = None
|
|
_processor = None
|
|
|
|
|
|
# NOT thread-safe. Must run in the ml-worker container with --concurrency=1.
|
|
def _load() -> None:
|
|
global _model, _processor, _torch, _AutoModel, _AutoProcessor
|
|
if _model is not None:
|
|
return
|
|
# Lazy import torch/transformers so this module can be imported in contexts
|
|
# where they're not available (e.g., web container for centroid-recompute enqueue).
|
|
import torch
|
|
from transformers import AutoModel, AutoProcessor
|
|
_torch = torch
|
|
_AutoModel = AutoModel
|
|
_AutoProcessor = AutoProcessor
|
|
|
|
_processor = _AutoProcessor.from_pretrained(_LOCAL_DIR)
|
|
_model = _AutoModel.from_pretrained(_LOCAL_DIR)
|
|
_model.eval()
|
|
|
|
|
|
def infer(image_path: str) -> np.ndarray:
|
|
"""Return a 1152-dim float32 numpy embedding for the image.
|
|
|
|
SigLIP uses a MAP-pooled vision head — the pooled output is the retrieval-ready
|
|
embedding the model was trained to produce. `get_image_features` on transformers
|
|
>= 4.45 returns a BaseModelOutputWithPooling, so pull `.pooler_output` explicitly
|
|
rather than relying on the first-field fallback from indexing.
|
|
"""
|
|
_load()
|
|
img = Image.open(image_path).convert('RGB')
|
|
with _torch.no_grad():
|
|
inputs = _processor(images=img, return_tensors='pt')
|
|
out = _model.get_image_features(**inputs)
|
|
pooled = out.pooler_output if hasattr(out, 'pooler_output') else out # (1, 1152)
|
|
return pooled[0].numpy().astype(np.float32)
|