app/services/embedder.py

from functools import lru_cache

import numpy as np
import torch
from PIL import Image

from app.config import get_settings


@lru_cache(maxsize=1)
def _processor_and_model():
    from transformers import AutoImageProcessor, AutoModel

    settings = get_settings()
    processor = AutoImageProcessor.from_pretrained(settings.embedder_model)
    model = AutoModel.from_pretrained(settings.embedder_model)
    model.eval()
    return processor, model


@torch.inference_mode()
def embed(image: Image.Image) -> np.ndarray:
    """Run DINOv2 on a (cropped) PIL image. Returns an L2-normalized 384-d float32 vector."""
    processor, model = _processor_and_model()
    inputs = processor(images=image.convert("RGB"), return_tensors="pt")
    outputs = model(**inputs)
    # last_hidden_state: [batch, seq, dim]; index 0 is the [CLS] token for DINOv2.
    cls = outputs.last_hidden_state[:, 0, :].squeeze(0).cpu().numpy().astype(np.float32)
    norm = np.linalg.norm(cls)
    if norm == 0:
        return cls
    return cls / norm