bvandeusen
0f35a0c484
Consolidated merge of feat/tag-suggestions branch. Original 64-commit history was lost to git-object corruption in a Nextcloud-synced checkout; this single commit captures the equivalent diff. Includes: - pgvector-backed tag suggestion infra (WD14 + SigLIP centroids, ml-worker container, Celery tasks, suggestion service, accept/reject endpoints + modal UI with green/red chip buttons) - Character/fandom integrity: title-case normalization on every write path, fandom-id backfill, maintenance task + settings button, migrations g26041901 + h26041901 to canonicalize legacy rows with case-only duplicate merging - Tag-underscores + modal polish: WD14 name canonicalization at emit + accept + add/bulk-add paths, migration i26041901 for legacy-row rename-or-merge across character/fandom/NULL kinds, suggestion-accept refresh parity via awaited loadTags, persistent chip tint
59 lines
2.3 KiB
Python
59 lines
2.3 KiB
Python
"""SigLIP SO400M image-embedding wrapper (PyTorch CPU)."""
|
|
from __future__ import annotations
|
|
import os
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
# 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)
|