shared/embedder.py

"""The ONE embedder. Imported by both the indexer (PC) and the server (Space).

Any divergence between index-time and serve-time preprocessing silently
destroys retrieval accuracy, so there must be exactly one implementation.
Vision-only: the server never needs the text tower.
"""

import numpy as np
import torch
from PIL import Image

from transformers import CLIPVisionModelWithProjection, CLIPImageProcessor

try:
    from .version import MODEL_ID
except ImportError:  # allow running as a loose script
    from version import MODEL_ID


class Embedder:
    def __init__(self, model_id: str = MODEL_ID):
        self.model_id = model_id
        self.proc = CLIPImageProcessor.from_pretrained(model_id)
        self.model = CLIPVisionModelWithProjection.from_pretrained(model_id).eval()

    @torch.no_grad()
    def embed(self, images: list[Image.Image]) -> np.ndarray:
        """Return (B, 512) float32, L2-normalized rows. Input PIL images (any mode)."""
        rgb = [im.convert("RGB") for im in images]
        inputs = self.proc(images=rgb, return_tensors="pt")
        emb = self.model(**inputs).image_embeds
        emb = emb / emb.norm(dim=-1, keepdim=True)
        return emb.float().cpu().numpy()