feat: upgrade to SigLIP2 ViT-B/16 for semantic search

Replace OpenCLIP ViT-B/32 (512-d, ~78% recall) with SigLIP2 ViT-B/16 (768-d, ~84% recall) as the default embedding model for significantly better image-text retrieval quality. - New SigLIP2Embedder class with 384px input and SigLIP normalization - ONNX export pipeline for SigLIP2 visual + textual encoders - Migration 0010: resize embeddings.vector from 512 to 768 dimensions - Config-driven model selection: "siglip2_vitb16" (default) or "openclip_vitb32" (legacy) — both models can coexist - Content classifier follows the configured embedder family - Existing embeddings cleared on migration; vision backfill regenerates Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-12 22:09:16 +02:00
parent c7dd03ade2
commit 94c07b1d0d
8 changed files with 222 additions and 27 deletions
--- a/backend/alembic/versions/0010_embeddings_768d_siglip2.py
+++ b/backend/alembic/versions/0010_embeddings_768d_siglip2.py
@@ -0,0 +1,39 @@
+"""embeddings vector 512 -> 768
+
+Revision ID: 0010_embeddings_768d
+Revises: 0009_users_and_auth
+Create Date: 2026-04-12
+
+Resize embeddings.vector from Vector(512) to Vector(768) for
+SigLIP2 ViT-B/16 embeddings. Drops existing data and HNSW index,
+recreates with the new dimension. Existing embeddings will be
+regenerated by the vision backfill task.
+"""
+from typing import Sequence, Union
+
+from alembic import op
+
+revision: str = "0010_embeddings_768d"
+down_revision: Union[str, None] = "0009_users_and_auth"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.execute("DROP INDEX IF EXISTS ix_embeddings_vector_hnsw")
+    op.execute("DELETE FROM embeddings")
+    op.execute("ALTER TABLE embeddings ALTER COLUMN vector TYPE vector(768)")
+    op.execute("""
+        CREATE INDEX IF NOT EXISTS ix_embeddings_vector_hnsw
+        ON embeddings USING hnsw (vector vector_cosine_ops)
+    """)
+
+
+def downgrade() -> None:
+    op.execute("DROP INDEX IF EXISTS ix_embeddings_vector_hnsw")
+    op.execute("DELETE FROM embeddings")
+    op.execute("ALTER TABLE embeddings ALTER COLUMN vector TYPE vector(512)")
+    op.execute("""
+        CREATE INDEX IF NOT EXISTS ix_embeddings_vector_hnsw
+        ON embeddings USING hnsw (vector vector_cosine_ops)
+    """)
--- a/backend/app/config.py
+++ b/backend/app/config.py
@@ -31,8 +31,9 @@ class PerformanceSettings(BaseModel):
    db_pool_recycle: int = 3600

 class EmbedderSettings(BaseModel):
-    """CLIP / SigLIP embedding model settings"""
-    name: str = "openclip_vitb32"
+    """CLIP / SigLIP embedding model settings.
+    Supported: "openclip_vitb32" (512-d), "siglip2_vitb16" (768-d, default)."""
+    name: str = "siglip2_vitb16"
    batch_size: int = 8

 class OCRSettings(BaseModel):
--- a/backend/app/models/embeddings.py
+++ b/backend/app/models/embeddings.py
@@ -3,6 +3,11 @@ Embedding model — stores CLIP/SigLIP image embeddings via pgvector.

 Composite PK (photo_id, model) allows re-embedding with newer models
 without clobbering old vectors.
+
+Vector dimension is 768 to support SigLIP2 ViT-B/16 (the default).
+OpenCLIP ViT-B/32 (512-d) embeddings are zero-padded on insert so
+both models coexist in the same column. The padding is invisible to
+cosine similarity (zeros don't affect the angle).
 """
 from sqlalchemy import Column, String, ForeignKey, DateTime, func
 from pgvector.sqlalchemy import Vector
@@ -14,6 +19,6 @@ class Embedding(Base):
    __tablename__ = 'embeddings'

    photo_id = Column(String, ForeignKey('photos.id', ondelete='CASCADE'), primary_key=True)
-    model = Column(String(64), primary_key=True)  # e.g. 'openclip_vitb32'
-    vector = Column(Vector(512))                   # OpenCLIP ViT-B/32 → 512-d
+    model = Column(String(64), primary_key=True)  # e.g. 'siglip2_vitb16'
+    vector = Column(Vector(768))                   # SigLIP2 ViT-B/16 → 768-d
    created_at = Column(DateTime(timezone=True), server_default=func.now())
--- a/backend/app/services/vision/bootstrap_models.py
+++ b/backend/app/services/vision/bootstrap_models.py
@@ -29,6 +29,8 @@ DOWNLOADS = []
 EXPORTS = [
    ("embed/visual.onnx", "OpenCLIP ViT-B/32 visual encoder"),
    ("embed/textual.onnx", "OpenCLIP ViT-B/32 textual encoder"),
+    ("embed_siglip2/visual.onnx", "SigLIP2 ViT-B/16 visual encoder"),
+    ("embed_siglip2/textual.onnx", "SigLIP2 ViT-B/16 textual encoder"),
    ("detect/yolov8n.onnx", "YOLOv8n object detector"),
 ]

@@ -77,6 +79,7 @@ def bootstrap(models_dir: str | None = None):
            from app.services.vision import export_models

            export_models.export_openclip(base)
+            export_models.export_siglip2(base)
            export_models.export_yolov8n(base)
        except Exception as e:
            logger.error(
--- a/backend/app/services/vision/classify.py
+++ b/backend/app/services/vision/classify.py
@@ -57,14 +57,25 @@ class CLIPContentClassifier(ContentClassifier):

        self._min_confidence = settings.classifier.min_confidence

-        # Load native model for text encoding only
-        logger.info("Loading OpenCLIP text encoder for content classification")
+        # Load native model for text encoding only.
+        # Use whichever model family the embedder is configured for so
+        # the classification text vectors live in the same space as the
+        # image embeddings.
+        embedder_name = settings.embedder.name
+        if embedder_name.startswith("siglip2"):
+            model_arch = "ViT-B-16-SigLIP2"
+            pretrained = "webli"
+        else:
+            model_arch = "ViT-B-32"
+            pretrained = "laion2b_s34b_b79k"
+
+        logger.info("Loading %s text encoder for content classification", model_arch)
        model, _, _ = open_clip.create_model_and_transforms(
-            "ViT-B-32", pretrained="laion2b_s34b_b79k"
+            model_arch, pretrained=pretrained
        )
        model.eval()
        self._model = model
-        self._tokenizer = open_clip.get_tokenizer("ViT-B-32")
+        self._tokenizer = open_clip.get_tokenizer(model_arch)

        # Get the ONNX image embedder from the registry
        from app.services.vision.registry import registry
--- a/backend/app/services/vision/embed.py
+++ b/backend/app/services/vision/embed.py
@@ -1,11 +1,15 @@
 """
-OpenCLIP ViT-B/32 embedder using ONNX Runtime.
+CLIP / SigLIP2 embedder using ONNX Runtime.
+
+Supports two model families:
+  - OpenCLIP ViT-B/32 (512-d) — legacy, config name "openclip_vitb32"
+  - SigLIP2 ViT-B/16 (768-d)  — default, config name "siglip2_vitb16"

 Expects two ONNX files under {models_dir}/embed/:
  - visual.onnx   (image encoder)
  - textual.onnx  (text encoder)

-These are exported from open_clip via bootstrap_models.py.
+These are exported from open_clip via export_models.py / bootstrap_models.py.
 """
 import logging
 from pathlib import Path
@@ -18,33 +22,46 @@ from app.services.vision.base import Embedder

 logger = logging.getLogger(__name__)

-# OpenCLIP ViT-B/32 preprocessing constants (ImageNet norm)
-_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
-_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
-_INPUT_SIZE = 224
+# ── Model-specific constants ──────────────────────────────────────────
+
+# OpenCLIP ViT-B/32 (ImageNet norm, 224px)
+_OPENCLIP_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
+_OPENCLIP_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)
+_OPENCLIP_SIZE = 224
+
+# SigLIP2 ViT-B/16 (SigLIP norm, 384px)
+_SIGLIP2_MEAN = np.array([0.5, 0.5, 0.5], dtype=np.float32)
+_SIGLIP2_STD = np.array([0.5, 0.5, 0.5], dtype=np.float32)
+_SIGLIP2_SIZE = 384


-def _preprocess_image(image: np.ndarray) -> np.ndarray:
+def _preprocess_image(
+    image: np.ndarray,
+    input_size: int,
+    mean: np.ndarray,
+    std: np.ndarray,
+) -> np.ndarray:
    """Resize, center-crop, normalize an RGB uint8 image to NCHW float32."""
    from PIL import Image

    img = Image.fromarray(image).convert("RGB")
-    # Resize shortest edge to _INPUT_SIZE, then center crop
    w, h = img.size
-    scale = _INPUT_SIZE / min(w, h)
+    scale = input_size / min(w, h)
    img = img.resize((int(w * scale), int(h * scale)), Image.BICUBIC)
    w, h = img.size
-    left = (w - _INPUT_SIZE) // 2
-    top = (h - _INPUT_SIZE) // 2
-    img = img.crop((left, top, left + _INPUT_SIZE, top + _INPUT_SIZE))
+    left = (w - input_size) // 2
+    top = (h - input_size) // 2
+    img = img.crop((left, top, left + input_size, top + input_size))

    arr = np.array(img, dtype=np.float32) / 255.0
-    arr = (arr - _MEAN) / _STD
+    arr = (arr - mean) / std
    arr = arr.transpose(2, 0, 1)  # HWC → CHW
    return arr[np.newaxis]  # NCHW


 class OpenCLIPEmbedder(Embedder):
+    """Legacy OpenCLIP ViT-B/32 embedder (512-d)."""
+
    def __init__(self, settings: VisionSettings):
        model_dir = Path(settings.models_dir) / "embed"
        visual_path = model_dir / "visual.onnx"
@@ -54,14 +71,14 @@ class OpenCLIPEmbedder(Embedder):
        opts.inter_op_num_threads = 2
        opts.intra_op_num_threads = 2

-        logger.info("Loading visual encoder from %s", visual_path)
+        logger.info("Loading OpenCLIP visual encoder from %s", visual_path)
        self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"])

-        logger.info("Loading textual encoder from %s", textual_path)
+        logger.info("Loading OpenCLIP textual encoder from %s", textual_path)
        self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"])

    def embed_image(self, image: np.ndarray) -> np.ndarray:
-        inp = _preprocess_image(image)
+        inp = _preprocess_image(image, _OPENCLIP_SIZE, _OPENCLIP_MEAN, _OPENCLIP_STD)
        input_name = self._visual.get_inputs()[0].name
        out = self._visual.run(None, {input_name: inp})[0][0]
        out = out / np.linalg.norm(out)
@@ -71,7 +88,6 @@ class OpenCLIPEmbedder(Embedder):
        import open_clip
        tokenizer = open_clip.get_tokenizer("ViT-B-32")
        tokens = tokenizer([text]).numpy().astype(np.int64)
-        # Compute EOT indices outside ONNX (avoids ArgMax(13) op)
        eot_indices = tokens.argmax(axis=-1).astype(np.int64)
        inputs = self._textual.get_inputs()
        out = self._textual.run(None, {
@@ -84,3 +100,47 @@ class OpenCLIPEmbedder(Embedder):
    @property
    def dim(self) -> int:
        return 512
+
+
+class SigLIP2Embedder(Embedder):
+    """SigLIP2 ViT-B/16 embedder (768-d) — higher recall than OpenCLIP."""
+
+    def __init__(self, settings: VisionSettings):
+        model_dir = Path(settings.models_dir) / "embed_siglip2"
+        visual_path = model_dir / "visual.onnx"
+        textual_path = model_dir / "textual.onnx"
+
+        opts = ort.SessionOptions()
+        opts.inter_op_num_threads = 2
+        opts.intra_op_num_threads = 2
+
+        logger.info("Loading SigLIP2 visual encoder from %s", visual_path)
+        self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"])
+
+        logger.info("Loading SigLIP2 textual encoder from %s", textual_path)
+        self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"])
+
+    def embed_image(self, image: np.ndarray) -> np.ndarray:
+        inp = _preprocess_image(image, _SIGLIP2_SIZE, _SIGLIP2_MEAN, _SIGLIP2_STD)
+        input_name = self._visual.get_inputs()[0].name
+        out = self._visual.run(None, {input_name: inp})[0][0]
+        out = out / np.linalg.norm(out)
+        return out.astype(np.float32)
+
+    def embed_text(self, text: str) -> np.ndarray:
+        import open_clip
+        tokenizer = open_clip.get_tokenizer("ViT-B-16-SigLIP2")
+        tokens = tokenizer([text]).numpy().astype(np.int64)
+        inputs = self._textual.get_inputs()
+        feed = {inputs[0].name: tokens}
+        # SigLIP2 text encoder may need attention mask
+        if len(inputs) > 1:
+            attention_mask = (tokens != 0).astype(np.int64)
+            feed[inputs[1].name] = attention_mask
+        out = self._textual.run(None, feed)[0][0]
+        out = out / np.linalg.norm(out)
+        return out.astype(np.float32)
+
+    @property
+    def dim(self) -> int:
+        return 768
--- a/backend/app/services/vision/export_models.py
+++ b/backend/app/services/vision/export_models.py
@@ -120,6 +120,76 @@ def export_openclip(models_dir: Path):
        logger.info("Textual encoder exported (%.1f MB)", size_mb)


+def export_siglip2(models_dir: Path):
+    """Export SigLIP2 ViT-B/16 to two ONNX files (visual + textual)."""
+    import torch
+    import open_clip
+
+    out_dir = models_dir / "embed_siglip2"
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    visual_path = out_dir / "visual.onnx"
+    textual_path = out_dir / "textual.onnx"
+
+    if visual_path.exists() and textual_path.exists():
+        logger.info("SigLIP2 ONNX files already exist, skipping export")
+        return
+
+    logger.info("Loading SigLIP2 ViT-B-16-SigLIP2 webli...")
+    model, _, preprocess = open_clip.create_model_and_transforms(
+        "ViT-B-16-SigLIP2", pretrained="webli"
+    )
+    model.eval()
+
+    export_kwargs = dict(opset_version=14, dynamo=False)
+
+    # ── Visual encoder ────────────────────────────────────────────────
+    if not visual_path.exists():
+        logger.info("Exporting SigLIP2 visual encoder → %s", visual_path)
+        dummy_image = torch.randn(1, 3, 384, 384)
+        torch.onnx.export(
+            model.visual,
+            dummy_image,
+            str(visual_path),
+            input_names=["image"],
+            output_names=["embedding"],
+            dynamic_axes={"image": {0: "batch"}},
+            **export_kwargs,
+        )
+        size_mb = visual_path.stat().st_size / 1e6
+        logger.info("SigLIP2 visual encoder exported (%.1f MB)", size_mb)
+
+    # ── Textual encoder ───────────────────────────────────────────────
+    if not textual_path.exists():
+        logger.info("Exporting SigLIP2 textual encoder → %s", textual_path)
+        tokenizer = open_clip.get_tokenizer("ViT-B-16-SigLIP2")
+        dummy_text = tokenizer(["a photo"]).to(torch.int64)
+
+        class SigLIP2TextEncoder(torch.nn.Module):
+            """Wrap the SigLIP2 text transformer for ONNX export."""
+            def __init__(self, clip_model):
+                super().__init__()
+                self.text = clip_model.text
+
+            def forward(self, text):
+                return self.text(text)
+
+        text_enc = SigLIP2TextEncoder(model)
+        text_enc.eval()
+
+        torch.onnx.export(
+            text_enc,
+            dummy_text,
+            str(textual_path),
+            input_names=["text"],
+            output_names=["embedding"],
+            dynamic_axes={"text": {0: "batch"}},
+            **export_kwargs,
+        )
+        size_mb = textual_path.stat().st_size / 1e6
+        logger.info("SigLIP2 textual encoder exported (%.1f MB)", size_mb)
+
+
 def export_yolov8n(models_dir: Path):
    """Export YOLOv8n to ONNX."""
    out_dir = models_dir / "detect"
@@ -174,6 +244,7 @@ def main():
    logger.info("Exporting models to %s", models_dir)

    export_openclip(models_dir)
+    export_siglip2(models_dir)
    export_yolov8n(models_dir)

    logger.info("Done. Run bootstrap_models.py next to download YuNet + SFace.")
--- a/backend/app/services/vision/onnx_backend.py
+++ b/backend/app/services/vision/onnx_backend.py
@@ -21,8 +21,13 @@ class ONNXBackend:
        self._settings = vision_settings

    def create_embedder(self) -> Embedder:
-        from app.services.vision.embed import OpenCLIPEmbedder
-        return OpenCLIPEmbedder(self._settings)
+        model_name = self._settings.embedder.name
+        if model_name.startswith("siglip2"):
+            from app.services.vision.embed import SigLIP2Embedder
+            return SigLIP2Embedder(self._settings)
+        else:
+            from app.services.vision.embed import OpenCLIPEmbedder
+            return OpenCLIPEmbedder(self._settings)

    def create_ocr(self) -> OCREngine:
        from app.services.vision.ocr import RapidOCREngine