mule-image/backend/app/services/vision/embed.py

"""
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 export_models.py / bootstrap_models.py.
"""
import logging
from pathlib import Path

import numpy as np
import onnxruntime as ort

from app.config import VisionSettings
from app.services.vision.base import Embedder

logger = logging.getLogger(__name__)

# ── 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,
    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")
    w, h = img.size
    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))

    arr = np.array(img, dtype=np.float32) / 255.0
    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"
        textual_path = model_dir / "textual.onnx"

        from app.services.vision.providers import create_session
        from app.config import settings as app_settings
        providers = app_settings.vision.execution_providers

        logger.info("Loading OpenCLIP visual encoder from %s", visual_path)
        self._visual = create_session(str(visual_path), configured_providers=providers)

        logger.info("Loading OpenCLIP textual encoder from %s", textual_path)
        self._textual = create_session(str(textual_path), configured_providers=providers)

    def embed_image(self, image: np.ndarray) -> np.ndarray:
        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)
        return out.astype(np.float32)

    def embed_text(self, text: str) -> np.ndarray:
        import open_clip
        tokenizer = open_clip.get_tokenizer("ViT-B-32")
        tokens = tokenizer([text]).numpy().astype(np.int64)
        eot_indices = tokens.argmax(axis=-1).astype(np.int64)
        inputs = self._textual.get_inputs()
        out = self._textual.run(None, {
            inputs[0].name: tokens,
            inputs[1].name: eot_indices,
        })[0][0]
        out = out / np.linalg.norm(out)
        return out.astype(np.float32)

    @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"

        from app.services.vision.providers import create_session
        from app.config import settings as app_settings
        providers = app_settings.vision.execution_providers

        logger.info("Loading SigLIP2 visual encoder from %s", visual_path)
        self._visual = create_session(str(visual_path), configured_providers=providers)

        logger.info("Loading SigLIP2 textual encoder from %s", textual_path)
        self._textual = create_session(str(textual_path), configured_providers=providers)

    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