feat(ml): export clip models to ONNX and host models on Hugging Face (#4700)

* export clip models

* export to hf

refactored export code

* export mclip, general refactoring

cleanup

* updated conda deps

* do transforms with pillow and numpy, add tokenization config to export, general refactoring

* moved conda dockerfile, re-added poetry

* minor fixes

* updated link

* updated tests

* removed `requirements.txt` from workflow

* fixed mimalloc path

* removed torchvision

* cleaner np typing

* review suggestions

* update default model name

* update test

machine-learning/app/conftest.py

@@ -1,5 +1,6 @@
 import json
-from typing import Any, Iterator, TypeAlias
+from pathlib import Path
+from typing import Any, Iterator
 from unittest import mock
 
 import numpy as np
@@ -8,8 +9,7 @@ from fastapi.testclient import TestClient
 from PIL import Image
 
 from .main import app, init_state
-
-ndarray: TypeAlias = np.ndarray[int, np.dtype[np.float32]]
+from .schemas import ndarray_f32
 
 
 @pytest.fixture
@@ -18,13 +18,13 @@ def pil_image() -> Image.Image:
 
 
 @pytest.fixture
-def cv_image(pil_image: Image.Image) -> ndarray:
+def cv_image(pil_image: Image.Image) -> ndarray_f32:
     return np.asarray(pil_image)[:, :, ::-1]  # PIL uses RGB while cv2 uses BGR
 
 
 @pytest.fixture
 def mock_get_model() -> Iterator[mock.Mock]:
-    with mock.patch("app.models.cache.InferenceModel.from_model_type", autospec=True) as mocked:
+    with mock.patch("app.models.cache.from_model_type", autospec=True) as mocked:
         yield mocked
 
 
@@ -37,3 +37,25 @@ def deployed_app() -> TestClient:
 @pytest.fixture(scope="session")
 def responses() -> dict[str, Any]:
     return json.load(open("responses.json", "r"))
+
+
+@pytest.fixture(scope="session")
+def clip_model_cfg() -> dict[str, Any]:
+    return {
+        "embed_dim": 512,
+        "vision_cfg": {"image_size": 224, "layers": 12, "width": 768, "patch_size": 32},
+        "text_cfg": {"context_length": 77, "vocab_size": 49408, "width": 512, "heads": 8, "layers": 12},
+    }
+
+
+@pytest.fixture(scope="session")
+def clip_preprocess_cfg() -> dict[str, Any]:
+    return {
+        "size": [224, 224],
+        "mode": "RGB",
+        "mean": [0.48145466, 0.4578275, 0.40821073],
+        "std": [0.26862954, 0.26130258, 0.27577711],
+        "interpolation": "bicubic",
+        "resize_mode": "shortest",
+        "fill_color": 0,
+    }

machine-learning/app/models/__init__.py

@@ -1,3 +1,25 @@
-from .clip import CLIPEncoder
+from typing import Any
+
+from app.schemas import ModelType
+
+from .base import InferenceModel
+from .clip import MCLIPEncoder, OpenCLIPEncoder, is_mclip, is_openclip
 from .facial_recognition import FaceRecognizer
 from .image_classification import ImageClassifier
+
+
+def from_model_type(model_type: ModelType, model_name: str, **model_kwargs: Any) -> InferenceModel:
+    match model_type:
+        case ModelType.CLIP:
+            if is_openclip(model_name):
+                return OpenCLIPEncoder(model_name, **model_kwargs)
+            elif is_mclip(model_name):
+                return MCLIPEncoder(model_name, **model_kwargs)
+            else:
+                raise ValueError(f"Unknown CLIP model {model_name}")
+        case ModelType.FACIAL_RECOGNITION:
+            return FaceRecognizer(model_name, **model_kwargs)
+        case ModelType.IMAGE_CLASSIFICATION:
+            return ImageClassifier(model_name, **model_kwargs)
+        case _:
+            raise ValueError(f"Unknown model type {model_type}")

machine-learning/app/models/base.py

@@ -25,7 +25,7 @@ class InferenceModel(ABC):
     ) -> None:
         self.model_name = model_name
         self.loaded = False
-        self._cache_dir = Path(cache_dir) if cache_dir is not None else get_cache_dir(model_name, self.model_type)
+        self._cache_dir = Path(cache_dir) if cache_dir is not None else None
         self.providers = model_kwargs.pop("providers", ["CPUExecutionProvider"])
         #  don't pre-allocate more memory than needed
         self.provider_options = model_kwargs.pop(
@@ -92,7 +92,7 @@ class InferenceModel(ABC):
 
     @property
     def cache_dir(self) -> Path:
-        return self._cache_dir
+        return self._cache_dir if self._cache_dir is not None else get_cache_dir(self.model_name, self.model_type)
 
     @cache_dir.setter
     def cache_dir(self, cache_dir: Path) -> None:

machine-learning/app/models/cache.py

@@ -4,6 +4,8 @@ from aiocache.backends.memory import SimpleMemoryCache
 from aiocache.lock import OptimisticLock
 from aiocache.plugins import BasePlugin, TimingPlugin
 
+from app.models import from_model_type
+
 from ..schemas import ModelType
 from .base import InferenceModel
 
@@ -50,7 +52,7 @@ class ModelCache:
         async with OptimisticLock(self.cache, key) as lock:
             model = await self.cache.get(key)
             if model is None:
-                model = InferenceModel.from_model_type(model_type, model_name, **model_kwargs)
+                model = from_model_type(model_type, model_name, **model_kwargs)
                 await lock.cas(model, ttl=self.ttl)
         return model
 

machine-learning/app/models/clip.py

@@ -1,23 +1,24 @@
-import os
-import zipfile
+import json
+from abc import abstractmethod
+from functools import cached_property
 from io import BytesIO
+from pathlib import Path
 from typing import Any, Literal
 
+import numpy as np
 import onnxruntime as ort
-import torch
-from clip_server.model.clip import BICUBIC, _convert_image_to_rgb
-from clip_server.model.clip_onnx import _MODELS, _S3_BUCKET_V2, CLIPOnnxModel, download_model
-from clip_server.model.pretrained_models import _VISUAL_MODEL_IMAGE_SIZE
-from clip_server.model.tokenization import Tokenizer
+from huggingface_hub import snapshot_download
 from PIL import Image
-from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
+from transformers import AutoTokenizer
+
+from app.config import log
+from app.models.transforms import crop, get_pil_resampling, normalize, resize, to_numpy
+from app.schemas import ModelType, ndarray_f32, ndarray_i32, ndarray_i64
 
-from ..config import log
-from ..schemas import ModelType
 from .base import InferenceModel
 
 
-class CLIPEncoder(InferenceModel):
+class BaseCLIPEncoder(InferenceModel):
     _model_type = ModelType.CLIP
 
     def __init__(
@@ -27,48 +28,29 @@ class CLIPEncoder(InferenceModel):
         mode: Literal["text", "vision"] | None = None,
         **model_kwargs: Any,
     ) -> None:
-        if mode is not None and mode not in ("text", "vision"):
-            raise ValueError(f"Mode must be 'text', 'vision', or omitted; got '{mode}'")
-        if model_name not in _MODELS:
-            raise ValueError(f"Unknown model name {model_name}.")
         self.mode = mode
         super().__init__(model_name, cache_dir, **model_kwargs)
 
-    def _download(self) -> None:
-        models: tuple[tuple[str, str], tuple[str, str]] = _MODELS[self.model_name]
-        text_onnx_path = self.cache_dir / "textual.onnx"
-        vision_onnx_path = self.cache_dir / "visual.onnx"
-
-        if not text_onnx_path.is_file():
-            self._download_model(*models[0])
-
-        if not vision_onnx_path.is_file():
-            self._download_model(*models[1])
-
     def _load(self) -> None:
         if self.mode == "text" or self.mode is None:
             log.debug(f"Loading clip text model '{self.model_name}'")
+
             self.text_model = ort.InferenceSession(
-                self.cache_dir / "textual.onnx",
+                self.textual_path.as_posix(),
                 sess_options=self.sess_options,
                 providers=self.providers,
                 provider_options=self.provider_options,
             )
-            self.text_outputs = [output.name for output in self.text_model.get_outputs()]
-            self.tokenizer = Tokenizer(self.model_name)
 
         if self.mode == "vision" or self.mode is None:
             log.debug(f"Loading clip vision model '{self.model_name}'")
+
             self.vision_model = ort.InferenceSession(
-                self.cache_dir / "visual.onnx",
+                self.visual_path.as_posix(),
                 sess_options=self.sess_options,
                 providers=self.providers,
                 provider_options=self.provider_options,
             )
-            self.vision_outputs = [output.name for output in self.vision_model.get_outputs()]
-
-            image_size = _VISUAL_MODEL_IMAGE_SIZE[CLIPOnnxModel.get_model_name(self.model_name)]
-            self.transform = _transform_pil_image(image_size)
 
     def _predict(self, image_or_text: Image.Image | str) -> list[float]:
         if isinstance(image_or_text, bytes):
@@ -78,55 +60,163 @@ class CLIPEncoder(InferenceModel):
             case Image.Image():
                 if self.mode == "text":
                     raise TypeError("Cannot encode image as text-only model")
-                pixel_values = self.transform(image_or_text)
-                assert isinstance(pixel_values, torch.Tensor)
-                pixel_values = torch.unsqueeze(pixel_values, 0).numpy()
-                outputs = self.vision_model.run(self.vision_outputs, {"pixel_values": pixel_values})
+
+                outputs = self.vision_model.run(None, self.transform(image_or_text))
             case str():
                 if self.mode == "vision":
                     raise TypeError("Cannot encode text as vision-only model")
-                text_inputs: dict[str, torch.Tensor] = self.tokenizer(image_or_text)
-                inputs = {
-                    "input_ids": text_inputs["input_ids"].int().numpy(),
-                    "attention_mask": text_inputs["attention_mask"].int().numpy(),
-                }
-                outputs = self.text_model.run(self.text_outputs, inputs)
+
+                outputs = self.text_model.run(None, self.tokenize(image_or_text))
             case _:
                 raise TypeError(f"Expected Image or str, but got: {type(image_or_text)}")
 
         return outputs[0][0].tolist()
 
-    def _download_model(self, model_name: str, model_md5: str) -> bool:
-        # downloading logic is adapted from clip-server's CLIPOnnxModel class
-        download_model(
-            url=_S3_BUCKET_V2 + model_name,
-            target_folder=self.cache_dir.as_posix(),
-            md5sum=model_md5,
-            with_resume=True,
-        )
-        file = self.cache_dir / model_name.split("/")[1]
-        if file.suffix == ".zip":
-            with zipfile.ZipFile(file, "r") as zip_ref:
-                zip_ref.extractall(self.cache_dir)
-            os.remove(file)
-        return True
+    @abstractmethod
+    def tokenize(self, text: str) -> dict[str, ndarray_i32]:
+        pass
+
+    @abstractmethod
+    def transform(self, image: Image.Image) -> dict[str, ndarray_f32]:
+        pass
+
+    @property
+    def textual_dir(self) -> Path:
+        return self.cache_dir / "textual"
+
+    @property
+    def visual_dir(self) -> Path:
+        return self.cache_dir / "visual"
+
+    @property
+    def model_cfg_path(self) -> Path:
+        return self.cache_dir / "config.json"
+
+    @property
+    def textual_path(self) -> Path:
+        return self.textual_dir / "model.onnx"
+
+    @property
+    def visual_path(self) -> Path:
+        return self.visual_dir / "model.onnx"
+
+    @property
+    def preprocess_cfg_path(self) -> Path:
+        return self.visual_dir / "preprocess_cfg.json"
 
     @property
     def cached(self) -> bool:
-        return (self.cache_dir / "textual.onnx").is_file() and (self.cache_dir / "visual.onnx").is_file()
+        return self.textual_path.is_file() and self.visual_path.is_file()
 
 
-# same as `_transform_blob` without `_blob2image`
-def _transform_pil_image(n_px: int) -> Compose:
-    return Compose(
-        [
-            Resize(n_px, interpolation=BICUBIC),
-            CenterCrop(n_px),
-            _convert_image_to_rgb,
-            ToTensor(),
-            Normalize(
-                (0.48145466, 0.4578275, 0.40821073),
-                (0.26862954, 0.26130258, 0.27577711),
-            ),
-        ]
-    )
+class OpenCLIPEncoder(BaseCLIPEncoder):
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: str | None = None,
+        mode: Literal["text", "vision"] | None = None,
+        **model_kwargs: Any,
+    ) -> None:
+        super().__init__(_clean_model_name(model_name), cache_dir, mode, **model_kwargs)
+
+    def _download(self) -> None:
+        snapshot_download(
+            f"immich-app/{self.model_name}",
+            cache_dir=self.cache_dir,
+            local_dir=self.cache_dir,
+            local_dir_use_symlinks=False,
+        )
+
+    def _load(self) -> None:
+        super()._load()
+
+        self.tokenizer = AutoTokenizer.from_pretrained(self.textual_dir)
+        self.sequence_length = self.model_cfg["text_cfg"]["context_length"]
+
+        self.size = (
+            self.preprocess_cfg["size"][0] if type(self.preprocess_cfg["size"]) == list else self.preprocess_cfg["size"]
+        )
+        self.resampling = get_pil_resampling(self.preprocess_cfg["interpolation"])
+        self.mean = np.array(self.preprocess_cfg["mean"], dtype=np.float32)
+        self.std = np.array(self.preprocess_cfg["std"], dtype=np.float32)
+
+    def tokenize(self, text: str) -> dict[str, ndarray_i32]:
+        input_ids: ndarray_i64 = self.tokenizer(
+            text,
+            max_length=self.sequence_length,
+            return_tensors="np",
+            return_attention_mask=False,
+            padding="max_length",
+            truncation=True,
+        ).input_ids
+        return {"text": input_ids.astype(np.int32)}
+
+    def transform(self, image: Image.Image) -> dict[str, ndarray_f32]:
+        image = resize(image, self.size)
+        image = crop(image, self.size)
+        image_np = to_numpy(image)
+        image_np = normalize(image_np, self.mean, self.std)
+        return {"image": np.expand_dims(image_np.transpose(2, 0, 1), 0)}
+
+    @cached_property
+    def model_cfg(self) -> dict[str, Any]:
+        return json.load(self.model_cfg_path.open())
+
+    @cached_property
+    def preprocess_cfg(self) -> dict[str, Any]:
+        return json.load(self.preprocess_cfg_path.open())
+
+
+class MCLIPEncoder(OpenCLIPEncoder):
+    def tokenize(self, text: str) -> dict[str, ndarray_i32]:
+        tokens: dict[str, ndarray_i64] = self.tokenizer(text, return_tensors="np")
+        return {k: v.astype(np.int32) for k, v in tokens.items()}
+
+
+_OPENCLIP_MODELS = {
+    "RN50__openai",
+    "RN50__yfcc15m",
+    "RN50__cc12m",
+    "RN101__openai",
+    "RN101__yfcc15m",
+    "RN50x4__openai",
+    "RN50x16__openai",
+    "RN50x64__openai",
+    "ViT-B-32__openai",
+    "ViT-B-32__laion2b_e16",
+    "ViT-B-32__laion400m_e31",
+    "ViT-B-32__laion400m_e32",
+    "ViT-B-32__laion2b-s34b-b79k",
+    "ViT-B-16__openai",
+    "ViT-B-16__laion400m_e31",
+    "ViT-B-16__laion400m_e32",
+    "ViT-B-16-plus-240__laion400m_e31",
+    "ViT-B-16-plus-240__laion400m_e32",
+    "ViT-L-14__openai",
+    "ViT-L-14__laion400m_e31",
+    "ViT-L-14__laion400m_e32",
+    "ViT-L-14__laion2b-s32b-b82k",
+    "ViT-L-14-336__openai",
+    "ViT-H-14__laion2b-s32b-b79k",
+    "ViT-g-14__laion2b-s12b-b42k",
+}
+
+
+_MCLIP_MODELS = {
+    "LABSE-Vit-L-14",
+    "XLM-Roberta-Large-Vit-B-32",
+    "XLM-Roberta-Large-Vit-B-16Plus",
+    "XLM-Roberta-Large-Vit-L-14",
+}
+
+
+def _clean_model_name(model_name: str) -> str:
+    return model_name.split("/")[-1].replace("::", "__")
+
+
+def is_openclip(model_name: str) -> bool:
+    return _clean_model_name(model_name) in _OPENCLIP_MODELS
+
+
+def is_mclip(model_name: str) -> bool:
+    return _clean_model_name(model_name) in _MCLIP_MODELS

machine-learning/app/models/facial_recognition.py

@@ -9,7 +9,8 @@ from insightface.model_zoo import ArcFaceONNX, RetinaFace
 from insightface.utils.face_align import norm_crop
 from insightface.utils.storage import BASE_REPO_URL, download_file
 
-from ..schemas import ModelType
+from app.schemas import ModelType, ndarray_f32
+
 from .base import InferenceModel
 
 
@@ -68,7 +69,7 @@ class FaceRecognizer(InferenceModel):
         )
         self.rec_model.prepare(ctx_id=0)
 
-    def _predict(self, image: np.ndarray[int, np.dtype[Any]] | bytes) -> list[dict[str, Any]]:
+    def _predict(self, image: ndarray_f32 | bytes) -> list[dict[str, Any]]:
         if isinstance(image, bytes):
             image = cv2.imdecode(np.frombuffer(image, np.uint8), cv2.IMREAD_COLOR)
         bboxes, kpss = self.det_model.detect(image)

machine-learning/app/models/transforms.py

@@ -0,0 +1,35 @@
+import numpy as np
+from PIL import Image
+
+from app.schemas import ndarray_f32
+
+_PIL_RESAMPLING_METHODS = {resampling.name.lower(): resampling for resampling in Image.Resampling}
+
+
+def resize(img: Image.Image, size: int) -> Image.Image:
+    if img.width < img.height:
+        return img.resize((size, int((img.height / img.width) * size)), resample=Image.BICUBIC)
+    else:
+        return img.resize((int((img.width / img.height) * size), size), resample=Image.BICUBIC)
+
+
+# https://stackoverflow.com/a/60883103
+def crop(img: Image.Image, size: int) -> Image.Image:
+    left = int((img.size[0] / 2) - (size / 2))
+    upper = int((img.size[1] / 2) - (size / 2))
+    right = left + size
+    lower = upper + size
+
+    return img.crop((left, upper, right, lower))
+
+
+def to_numpy(img: Image.Image) -> ndarray_f32:
+    return np.asarray(img.convert("RGB")).astype(np.float32) / 255.0
+
+
+def normalize(img: ndarray_f32, mean: float | ndarray_f32, std: float | ndarray_f32) -> ndarray_f32:
+    return (img - mean) / std
+
+
+def get_pil_resampling(resample: str) -> Image.Resampling:
+    return _PIL_RESAMPLING_METHODS[resample.lower()]

machine-learning/app/schemas.py

@@ -1,5 +1,7 @@
 from enum import StrEnum
+from typing import TypeAlias
 
+import numpy as np
 from pydantic import BaseModel
 
 
@@ -31,3 +33,8 @@ class ModelType(StrEnum):
     IMAGE_CLASSIFICATION = "image-classification"
     CLIP = "clip"
     FACIAL_RECOGNITION = "facial-recognition"
+
+
+ndarray_f32: TypeAlias = np.ndarray[int, np.dtype[np.float32]]
+ndarray_i64: TypeAlias = np.ndarray[int, np.dtype[np.int64]]
+ndarray_i32: TypeAlias = np.ndarray[int, np.dtype[np.int32]]

machine-learning/app/test_main.py

@@ -1,7 +1,8 @@
 import json
 import pickle
 from io import BytesIO
-from typing import Any, TypeAlias
+from pathlib import Path
+from typing import Any, Callable
 from unittest import mock
 
 import cv2
@@ -14,13 +15,11 @@ from pytest_mock import MockerFixture
 from .config import settings
 from .models.base import PicklableSessionOptions
 from .models.cache import ModelCache
-from .models.clip import CLIPEncoder
+from .models.clip import OpenCLIPEncoder
 from .models.facial_recognition import FaceRecognizer
 from .models.image_classification import ImageClassifier
 from .schemas import ModelType
 
-ndarray: TypeAlias = np.ndarray[int, np.dtype[np.float32]]
-
 
 class TestImageClassifier:
     classifier_preds = [
@@ -56,30 +55,50 @@ class TestImageClassifier:
 
 class TestCLIP:
     embedding = np.random.rand(512).astype(np.float32)
+    cache_dir = Path("test_cache")
 
-    def test_basic_image(self, pil_image: Image.Image, mocker: MockerFixture) -> None:
-        mocker.patch.object(CLIPEncoder, "download")
+    def test_basic_image(
+        self,
+        pil_image: Image.Image,
+        mocker: MockerFixture,
+        clip_model_cfg: dict[str, Any],
+        clip_preprocess_cfg: Callable[[Path], dict[str, Any]],
+    ) -> None:
+        mocker.patch.object(OpenCLIPEncoder, "download")
+        mocker.patch.object(OpenCLIPEncoder, "model_cfg", clip_model_cfg)
+        mocker.patch.object(OpenCLIPEncoder, "preprocess_cfg", clip_preprocess_cfg)
+        mocker.patch("app.models.clip.AutoTokenizer.from_pretrained", autospec=True)
         mocked = mocker.patch("app.models.clip.ort.InferenceSession", autospec=True)
         mocked.return_value.run.return_value = [[self.embedding]]
-        clip_encoder = CLIPEncoder("ViT-B-32::openai", cache_dir="test_cache", mode="vision")
-        assert clip_encoder.mode == "vision"
+
+        clip_encoder = OpenCLIPEncoder("ViT-B-32::openai", cache_dir="test_cache", mode="vision")
         embedding = clip_encoder.predict(pil_image)
 
+        assert clip_encoder.mode == "vision"
         assert isinstance(embedding, list)
-        assert len(embedding) == 512
+        assert len(embedding) == clip_model_cfg["embed_dim"]
         assert all([isinstance(num, float) for num in embedding])
         clip_encoder.vision_model.run.assert_called_once()
 
-    def test_basic_text(self, mocker: MockerFixture) -> None:
-        mocker.patch.object(CLIPEncoder, "download")
+    def test_basic_text(
+        self,
+        mocker: MockerFixture,
+        clip_model_cfg: dict[str, Any],
+        clip_preprocess_cfg: Callable[[Path], dict[str, Any]],
+    ) -> None:
+        mocker.patch.object(OpenCLIPEncoder, "download")
+        mocker.patch.object(OpenCLIPEncoder, "model_cfg", clip_model_cfg)
+        mocker.patch.object(OpenCLIPEncoder, "preprocess_cfg", clip_preprocess_cfg)
+        mocker.patch("app.models.clip.AutoTokenizer.from_pretrained", autospec=True)
         mocked = mocker.patch("app.models.clip.ort.InferenceSession", autospec=True)
         mocked.return_value.run.return_value = [[self.embedding]]
-        clip_encoder = CLIPEncoder("ViT-B-32::openai", cache_dir="test_cache", mode="text")
-        assert clip_encoder.mode == "text"
+
+        clip_encoder = OpenCLIPEncoder("ViT-B-32::openai", cache_dir="test_cache", mode="text")
         embedding = clip_encoder.predict("test search query")
 
+        assert clip_encoder.mode == "text"
         assert isinstance(embedding, list)
-        assert len(embedding) == 512
+        assert len(embedding) == clip_model_cfg["embed_dim"]
         assert all([isinstance(num, float) for num in embedding])
         clip_encoder.text_model.run.assert_called_once()