src/image_embed.py

"""
image_embed.py — CLIP embeddings via fastembed (ONNX / onnxruntime, NO torch).

CLIP puts images and text in ONE shared vector space, so the same collection can
be searched by an uploaded photo (vision encoder) OR by text (text encoder).
fastembed runs on onnxruntime, whose signed DLLs pass Windows App Control —
unlike torch-based CLIP.

Models (512-dim, aligned):
  images: Qdrant/clip-ViT-B-32-vision
  text  : Qdrant/clip-ViT-B-32-text
"""

IMG_MODEL = "Qdrant/clip-ViT-B-32-vision"
TXT_MODEL = "Qdrant/clip-ViT-B-32-text"

_img = None
_txt = None


def get_image_model():
    global _img
    if _img is None:
        from fastembed import ImageEmbedding
        _img = ImageEmbedding(model_name=IMG_MODEL)
    return _img


def get_text_model():
    global _txt
    if _txt is None:
        from fastembed import TextEmbedding
        _txt = TextEmbedding(model_name=TXT_MODEL)
    return _txt


def embed_images(paths):
    """List of image file paths -> list of 512-dim vectors (as python lists)."""
    return [v.tolist() for v in get_image_model().embed(paths)]


def embed_texts(texts):
    """List of strings -> list of 512-dim CLIP-text vectors (as python lists)."""
    return [v.tolist() for v in get_text_model().embed(texts)]


def embed_image_dataurl(data_url):
    """Embed a base64 data-URL image (from the web upload). Returns one vector."""
    import base64
    import os
    import tempfile
    if "," in data_url:
        data_url = data_url.split(",", 1)[1]
    raw = base64.b64decode(data_url)
    fd, path = tempfile.mkstemp(suffix=".jpg")
    try:
        with os.fdopen(fd, "wb") as fh:
            fh.write(raw)
        return embed_images([path])[0]
    finally:
        try:
            os.remove(path)
        except OSError:
            pass