diff --git a/backend/alembic/versions/0010_embeddings_768d_siglip2.py b/backend/alembic/versions/0010_embeddings_768d_siglip2.py new file mode 100644 index 0000000..0a9152c --- /dev/null +++ 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) + """) diff --git a/backend/app/config.py b/backend/app/config.py index d8308ef..7665a26 100644 --- 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): diff --git a/backend/app/models/embeddings.py b/backend/app/models/embeddings.py index 4368f0e..cf17aa0 100644 --- 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()) diff --git a/backend/app/services/vision/bootstrap_models.py b/backend/app/services/vision/bootstrap_models.py index bb9bf77..427071e 100644 --- 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( diff --git a/backend/app/services/vision/classify.py b/backend/app/services/vision/classify.py index 6de4424..4bb5a00 100644 --- 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 diff --git a/backend/app/services/vision/embed.py b/backend/app/services/vision/embed.py index b5357fa..8681bd4 100644 --- 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 diff --git a/backend/app/services/vision/export_models.py b/backend/app/services/vision/export_models.py index 9408e38..1818df3 100644 --- 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.") diff --git a/backend/app/services/vision/onnx_backend.py b/backend/app/services/vision/onnx_backend.py index 125a160..1cad8b1 100644 --- 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