diff --git a/docs/docs/install/environment-variables.md b/docs/docs/install/environment-variables.md index c10a858ed..ccc5e6a90 100644 --- a/docs/docs/install/environment-variables.md +++ b/docs/docs/install/environment-variables.md @@ -154,33 +154,33 @@ Redis (Sentinel) URL example JSON before encoding: ## Machine Learning -| Variable | Description | Default | Containers | -| :---------------------------------------------------------- | :----------------------------------------------------------------------------------------------------------------------------------------------------------- | :--------------------------: | :--------------- | -| `MACHINE_LEARNING_MODEL_TTL` | Inactivity time (s) before a model is unloaded (disabled if \<= 0) | `300` | machine learning | -| `MACHINE_LEARNING_MODEL_TTL_POLL_S` | Interval (s) between checks for the model TTL (disabled if \<= 0) | `10` | machine learning | -| `MACHINE_LEARNING_CACHE_FOLDER` | Directory where models are downloaded | `/cache` | machine learning | -| `MACHINE_LEARNING_REQUEST_THREADS`\*1 | Thread count of the request thread pool (disabled if \<= 0) | number of CPU cores | machine learning | -| `MACHINE_LEARNING_MODEL_INTER_OP_THREADS` | Number of parallel model operations | `1` | machine learning | -| `MACHINE_LEARNING_MODEL_INTRA_OP_THREADS` | Number of threads for each model operation | `2` | machine learning | -| `MACHINE_LEARNING_WORKERS`\*2 | Number of worker processes to spawn | `1` | machine learning | -| `MACHINE_LEARNING_HTTP_KEEPALIVE_TIMEOUT_S`\*3 | HTTP Keep-alive time in seconds | `2` | machine learning | -| `MACHINE_LEARNING_WORKER_TIMEOUT` | Maximum time (s) of unresponsiveness before a worker is killed | `300` (`900` if using ROCm) | machine learning | -| `MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL` | Comma-separated list of (textual) CLIP model(s) to preload and cache | | machine learning | -| `MACHINE_LEARNING_PRELOAD__CLIP__VISUAL` | Comma-separated list of (visual) CLIP model(s) to preload and cache | | machine learning | -| `MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__RECOGNITION` | Comma-separated list of (recognition) facial recognition model(s) to preload and cache | | machine learning | -| `MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__DETECTION` | Comma-separated list of (detection) facial recognition model(s) to preload and cache | | machine learning | -| `MACHINE_LEARNING_PRELOAD__OCR__RECOGNITION` | Comma-separated list of (recognition) OCR model(s) to preload and cache | | machine learning | -| `MACHINE_LEARNING_PRELOAD__OCR__DETECTION` | Comma-separated list of (detection) OCR model(s) to preload and cache | | machine learning | -| `MACHINE_LEARNING_ANN` | Enable ARM-NN hardware acceleration if supported | `True` | machine learning | -| `MACHINE_LEARNING_ANN_FP16_TURBO` | Execute operations in FP16 precision: increasing speed, reducing precision (applies only to ARM-NN) | `False` | machine learning | -| `MACHINE_LEARNING_ANN_TUNING_LEVEL` | ARM-NN GPU tuning level (1: rapid, 2: normal, 3: exhaustive) | `2` | machine learning | -| `MACHINE_LEARNING_DEVICE_IDS`\*4 | Device IDs to use in multi-GPU environments | `0` | machine learning | -| `MACHINE_LEARNING_MAX_BATCH_SIZE__FACIAL_RECOGNITION` | Set the maximum number of faces that will be processed at once by the facial recognition model | None (`1` if using OpenVINO) | machine learning | -| `MACHINE_LEARNING_MAX_BATCH_SIZE__OCR` | Set the maximum number of boxes that will be processed at once by the OCR model | `6` | machine learning | -| `MACHINE_LEARNING_RKNN` | Enable RKNN hardware acceleration if supported | `True` | machine learning | -| `MACHINE_LEARNING_RKNN_THREADS` | How many threads of RKNN runtime should be spun up while inferencing. | `1` | machine learning | -| `MACHINE_LEARNING_MODEL_ARENA` | Pre-allocates CPU memory to avoid memory fragmentation | true | machine learning | -| `MACHINE_LEARNING_OPENVINO_PRECISION` | If set to FP16, uses half-precision floating-point operations for faster inference with reduced accuracy (one of [`FP16`, `FP32`], applies only to OpenVINO) | `FP32` | machine learning | +| Variable | Description | Default | Containers | +| :---------------------------------------------------------- | :-------------------------------------------------------------------------------------------------- | :-------------------------: | :--------------- | +| `MACHINE_LEARNING_MODEL_TTL` | Inactivity time (s) before a model is unloaded (disabled if \<= 0) | `300` | machine learning | +| `MACHINE_LEARNING_MODEL_TTL_POLL_S` | Interval (s) between checks for the model TTL (disabled if \<= 0) | `10` | machine learning | +| `MACHINE_LEARNING_CACHE_FOLDER` | Directory where models are downloaded | `/cache` | machine learning | +| `MACHINE_LEARNING_REQUEST_THREADS`\*1 | Thread count of the request thread pool (disabled if \<= 0) | number of CPU cores | machine learning | +| `MACHINE_LEARNING_MODEL_INTER_OP_THREADS` | Number of parallel model operations | `1` | machine learning | +| `MACHINE_LEARNING_MODEL_INTRA_OP_THREADS` | Number of threads for each model operation | `2` | machine learning | +| `MACHINE_LEARNING_WORKERS`\*2 | Number of worker processes to spawn | `1` | machine learning | +| `MACHINE_LEARNING_HTTP_KEEPALIVE_TIMEOUT_S`\*3 | HTTP Keep-alive time in seconds | `2` | machine learning | +| `MACHINE_LEARNING_WORKER_TIMEOUT` | Maximum time (s) of unresponsiveness before a worker is killed | `300` (`900` if using ROCm) | machine learning | +| `MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL` | Comma-separated list of (textual) CLIP model(s) to preload and cache | | machine learning | +| `MACHINE_LEARNING_PRELOAD__CLIP__VISUAL` | Comma-separated list of (visual) CLIP model(s) to preload and cache | | machine learning | +| `MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__RECOGNITION` | Comma-separated list of (recognition) facial recognition model(s) to preload and cache | | machine learning | +| `MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__DETECTION` | Comma-separated list of (detection) facial recognition model(s) to preload and cache | | machine learning | +| `MACHINE_LEARNING_PRELOAD__OCR__RECOGNITION` | Comma-separated list of (recognition) OCR model(s) to preload and cache | | machine learning | +| `MACHINE_LEARNING_PRELOAD__OCR__DETECTION` | Comma-separated list of (detection) OCR model(s) to preload and cache | | machine learning | +| `MACHINE_LEARNING_PRELOAD__OCR__MAX_RESOLUTION` | Resolution to preload the OCR detection model(s) for (must match the server's OCR resolution) | `736` | machine learning | +| `MACHINE_LEARNING_ANN` | Enable ARM-NN hardware acceleration if supported | `True` | machine learning | +| `MACHINE_LEARNING_ANN_FP16_TURBO` | Execute operations in FP16 precision: increasing speed, reducing precision (applies only to ARM-NN) | `False` | machine learning | +| `MACHINE_LEARNING_ANN_TUNING_LEVEL` | ARM-NN GPU tuning level (1: rapid, 2: normal, 3: exhaustive) | `2` | machine learning | +| `MACHINE_LEARNING_DEVICE_IDS`\*4 | Device IDs to use in multi-GPU environments | `0` | machine learning | +| `MACHINE_LEARNING_MAX_BATCH_SIZE__FACIAL_RECOGNITION` | Set the maximum number of faces that will be processed at once by the facial recognition model | `4` | machine learning | +| `MACHINE_LEARNING_MAX_BATCH_SIZE__OCR` | Set the maximum number of boxes that will be processed at once by the OCR model | `6` | machine learning | +| `MACHINE_LEARNING_RKNN` | Enable RKNN hardware acceleration if supported | `True` | machine learning | +| `MACHINE_LEARNING_RKNN_THREADS` | How many threads of RKNN runtime should be spun up while inferencing. | `1` | machine learning | +| `MACHINE_LEARNING_MODEL_ARENA` | Pre-allocates CPU memory to avoid memory fragmentation | true | machine learning | \*1: It is recommended to begin with this parameter when changing the concurrency levels of the machine learning service and then tune the other ones. diff --git a/machine-learning/Dockerfile b/machine-learning/Dockerfile index b1ca31c88..25e367b14 100644 --- a/machine-learning/Dockerfile +++ b/machine-learning/Dockerfile @@ -31,7 +31,7 @@ ENV PYTHONDONTWRITEBYTECODE=1 \ PYTHONUNBUFFERED=1 \ VIRTUAL_ENV=/opt/venv -RUN apt-get update && apt-get install -y --no-install-recommends g++ +RUN apt-get update && apt-get install -y --no-install-recommends g++ git COPY --from=ghcr.io/astral-sh/uv:0.8.15@sha256:a5727064a0de127bdb7c9d3c1383f3a9ac307d9f2d8a391edc7896c54289ced0 /uv /uvx /bin/ RUN --mount=type=cache,target=/root/.cache/uv \ diff --git a/machine-learning/conftest.py b/machine-learning/conftest.py index a968d6906..34de2cbbb 100644 --- a/machine-learning/conftest.py +++ b/machine-learning/conftest.py @@ -1,4 +1,5 @@ import json +from pathlib import Path from types import SimpleNamespace from typing import Any, Callable, Iterator from unittest import mock @@ -10,6 +11,15 @@ from PIL import Image from immich_ml.config import log from immich_ml.main import app from immich_ml.models.base import InferenceModel +from immich_ml.schemas import Shape +from immich_ml.sessions.ort import Device + +TEST_ASSETS = Path(__file__).parent.parent / "e2e/test-assets" + + +@pytest.fixture(scope="session") +def asset() -> Callable[[str], bytes]: + return lambda path: (TEST_ASSETS / path).read_bytes() @pytest.fixture @@ -20,6 +30,7 @@ def pil_image() -> Image.Image: @pytest.fixture def mock_get_model() -> Iterator[mock.Mock]: with mock.patch("immich_ml.models.cache.get_model_class") as mocked: + mocked.return_value.graph_options = () yield mocked @@ -103,6 +114,15 @@ def providers(request: pytest.FixtureRequest) -> Iterator[mock.Mock]: yield providers +@pytest.fixture(autouse=True) +def gpus() -> Iterator[None]: + with ( + mock.patch("immich_ml.sessions.ort._intel_gpu", return_value=Device("12.71.4-128eu", "26.22.38646.4")), + mock.patch("immich_ml.sessions.ort._amd_gpu", return_value=Device("gfx1100", "7.2.0")), + ): + yield + + @pytest.fixture(scope="function") def ort_pybind() -> Iterator[mock.Mock]: with mock.patch("immich_ml.sessions.ort.ort.capi._pybind_state") as mocked: @@ -120,15 +140,29 @@ def ov_device_ids(request: pytest.FixtureRequest, ort_pybind: mock.Mock) -> Iter @pytest.fixture(scope="function") def ort_session() -> Iterator[mock.Mock]: - with mock.patch("immich_ml.sessions.ort.ort.InferenceSession") as mocked: + # a graph is opened as it is: preparing one takes a child process and real files + with ( + mock.patch("immich_ml.sessions.ort.ort.InferenceSession") as mocked, + mock.patch("immich_ml.sessions.ort.prepared", side_effect=lambda spec: spec.model_path), + ): yield mocked @pytest.fixture(scope="function") def stub_session() -> Callable[..., mock.Mock]: - def _make(shape: tuple[Any, ...], outputs: Any = None, name: str = "input.1") -> mock.Mock: + def _make( + shape: tuple[Any, ...], + outputs: Any = None, + name: str = "input.1", + normalizes_input: bool = False, + shapes: tuple[Shape, ...] = (Shape(batch=1),), + ) -> mock.Mock: session = mock.Mock() session.get_inputs.return_value = [SimpleNamespace(name=name, shape=shape)] + session.normalizes_input = normalizes_input + session.shapes = shapes + session.batches = tuple(sorted({shape.batch for shape in shapes}, reverse=True)) + session.for_shape.return_value = session # a stub is the session and the one graph in it if outputs is not None: session.run.return_value = outputs return session @@ -139,12 +173,16 @@ def stub_session() -> Callable[..., mock.Mock]: @pytest.fixture(scope="function") def ann_session() -> Iterator[mock.Mock]: with mock.patch("immich_ml.sessions.ann.Ann") as mocked: + mocked.return_value.input_shapes.__getitem__.return_value = [(1, 3, 112, 112)] # armnn compiles one batch yield mocked @pytest.fixture(scope="function") def rknn_session() -> Iterator[mock.Mock]: with mock.patch("immich_ml.sessions.rknn.RknnPoolExecutor") as mocked: + mocked.return_value.custom_string = "" # an unstamped binary, i.e. one that wants normalized input + compiled = SimpleNamespace(name="input", shape=(1, 224, 224, 3)) # the compiler reports NHWC either way + mocked.return_value.inputs = [compiled, compiled] yield mocked @@ -188,6 +226,7 @@ def exception() -> Iterator[mock.Mock]: @pytest.fixture(scope="function") def snapshot_download() -> Iterator[mock.Mock]: + # a download that leaves the model cached, as one that found the artifact does with ( mock.patch("immich_ml.models.base.snapshot_download") as mocked, mock.patch.object(InferenceModel, "cached", new_callable=mock.PropertyMock, side_effect=lambda: mocked.called), diff --git a/machine-learning/immich_ml/config.py b/machine-learning/immich_ml/config.py index 747d1a498..2a5f47cb1 100644 --- a/machine-learning/immich_ml/config.py +++ b/machine-learning/immich_ml/config.py @@ -13,7 +13,7 @@ from rich.logging import RichHandler from uvicorn import Server from uvicorn.workers import UvicornWorker -from .schemas import ModelOrganization, ModelPrecision +from .schemas import ModelOrganization class ClipSettings(BaseModel): @@ -29,6 +29,7 @@ class FacialRecognitionSettings(BaseModel): class OcrSettings(BaseModel): recognition: str | None = None detection: str | None = None + max_resolution: int = 736 # must match the server's OCR setting, or the preloaded model will not be used class PreloadModelData(BaseModel): @@ -38,8 +39,8 @@ class PreloadModelData(BaseModel): class MaxBatchSize(BaseModel): - facial_recognition: int | None = None - ocr: int | None = None + facial_recognition: int = 4 + ocr: int = 6 def default_worker_timeout() -> int: @@ -71,9 +72,7 @@ class Settings(BaseSettings): rknn: bool = True rknn_threads: int = 1 preload: PreloadModelData | None = None - max_batch_size: MaxBatchSize | None = None - openvino_precision: ModelPrecision = ModelPrecision.FP32 - rocm_precision: ModelPrecision = ModelPrecision.FP32 + max_batch_size: MaxBatchSize = MaxBatchSize() model_organization: ModelOrganization = ModelOrganization.APP model_revision: str = "main" diff --git a/machine-learning/immich_ml/main.py b/machine-learning/immich_ml/main.py index fdad03eae..8f27cb0d9 100644 --- a/machine-learning/immich_ml/main.py +++ b/machine-learning/immich_ml/main.py @@ -91,6 +91,7 @@ async def preload_models(preload: PreloadModelData) -> None: model_name = model_name.strip() model = model_cache.get(model_name, model_type, model_task, **options) await load(model) + await attempt(model, model.build, MODEL_FILE_ERRORS) if preload.clip.textual is not None: await load_models(preload.clip.textual, ModelType.TEXTUAL, ModelTask.SEARCH) @@ -117,6 +118,7 @@ async def preload_models(preload: PreloadModelData) -> None: preload.ocr.detection, ModelType.DETECTION, ModelTask.OCR, + maxResolution=preload.ocr.max_resolution, ) if preload.ocr.recognition is not None: diff --git a/machine-learning/immich_ml/models/base.py b/machine-learning/immich_ml/models/base.py index a54a80485..2e999c859 100644 --- a/machine-learning/immich_ml/models/base.py +++ b/machine-learning/immich_ml/models/base.py @@ -15,6 +15,7 @@ from immich_ml.sessions.ort import OrtSession from ..config import clean_name, log, settings from ..schemas import ModelFormat, ModelIdentity, ModelSession, ModelTask, ModelType from ..sessions.ann import AnnSession +from ..sessions.policy import ShapePolicy _IGNORED_PATTERNS: dict[ModelFormat, list[str]] = { ModelFormat.ONNX: ["*.armnn", "*.rknn"], @@ -26,6 +27,10 @@ _IGNORED_PATTERNS: dict[ModelFormat, list[str]] = { class InferenceModel(ABC): depends: ClassVar[list[ModelIdentity]] identity: ClassVar[ModelIdentity] + # options a graph is built for, so one instance cannot serve two values of them + graph_options: ClassVar[tuple[str, ...]] = () + # the shapes this model feeds, which pick the artifact and tell the engine what it may pin + shape_policy: ShapePolicy = ShapePolicy() threads: ClassVar[int] = 2 # how many threads one run takes for CPU inference def __init__( @@ -78,6 +83,10 @@ class InferenceModel(ABC): del self.session self.loaded = False + def build(self) -> None: + self.load() + self.session.warm() + def predict(self, *inputs: Any, **model_kwargs: Any) -> Any: self.load() return self._predict(*inputs, **model_kwargs) @@ -115,23 +124,19 @@ class InferenceModel(ABC): case ModelFormat.ARMNN: return AnnSession(self.model_path) case ModelFormat.ONNX: - return OrtSession(self.model_path, threads=self.threads) + return OrtSession(self.model_path, self.shape_policy, threads=self.threads) case ModelFormat.RKNN: return rknn.RknnSession(self.model_path) - def model_path_for_format(self, model_format: ModelFormat) -> Path: - model_path_prefix = rknn.model_prefix if model_format == ModelFormat.RKNN else None - if model_path_prefix: - return self.model_dir / model_path_prefix / f"model.{model_format}" - return self.model_dir / f"model.{model_format}" - @property def model_dir(self) -> Path: return self.cache_dir / self.model_type.value @property def model_path(self) -> Path: - return self.model_path_for_format(self.model_format) + if self.model_format == ModelFormat.RKNN: + return rknn.model_path(self.model_dir, self.shape_policy.label) + return self.model_dir / f"model.{self.model_format}" @property def model_task(self) -> ModelTask: @@ -151,7 +156,8 @@ class InferenceModel(ABC): @property def _cache_dir_default(self) -> Path: - return settings.cache_folder / self.model_task.value / self.model_name + cache_dir = settings.cache_folder / self.model_task.value / self.model_name + return cache_dir if settings.legacy_models else cache_dir / settings.model_revision @property def cached(self) -> bool: diff --git a/machine-learning/immich_ml/models/cache.py b/machine-learning/immich_ml/models/cache.py index 1a8dad2eb..a433411b9 100644 --- a/machine-learning/immich_ml/models/cache.py +++ b/machine-learning/immich_ml/models/cache.py @@ -17,7 +17,7 @@ class ModelCache: self, model_name: str, model_type: ModelType, model_task: ModelTask, ttl: int | None = None, **options: Any ) -> InferenceModel: model_cls = get_model_class(model_name, model_type, model_task) - key = (model_name, model_type, model_task) + key = (model_name, model_type, model_task, *(options.get(option) for option in model_cls.graph_options)) if key not in self._models: self._models[key] = model_cls(model_name, **options) elif key not in self._expiries: diff --git a/machine-learning/immich_ml/models/clip/textual.py b/machine-learning/immich_ml/models/clip/textual.py index 003d30e56..d5ad3400f 100644 --- a/machine-learning/immich_ml/models/clip/textual.py +++ b/machine-learning/immich_ml/models/clip/textual.py @@ -15,18 +15,33 @@ from immich_ml.models.transforms import clean_text, serialize_np_array from immich_ml.schemas import ModelSession, ModelTask, ModelType +def _mapped(model_path: Path, name: str) -> NDArray[Any]: + import onnx + + initializer = next(i for i in onnx.load(model_path, load_external_data=False).graph.initializer if i.name == name) + data = {entry.key: entry.value for entry in initializer.external_data} + dtype = onnx.helper.tensor_dtype_to_np_dtype(initializer.data_type) + offset, shape = int(data.get("offset", 0)), tuple(initializer.dims) + return np.memmap(model_path.parent / data["location"], dtype=dtype, mode="r", offset=offset, shape=shape) + + class BaseCLIPTextualEncoder(InferenceModel): depends = [] identity = (ModelType.TEXTUAL, ModelTask.SEARCH) threads = 4 # lower search latency def _predict(self, inputs: str, language: str | None = None) -> str: - tokens = self.tokenize(inputs, language=language) - res: NDArray[np.float32] = self.session.run(None, tokens)[0][0] + tokens: dict[str, NDArray[Any]] = self.tokenize(inputs, language=language) + graph = self.session.for_shape(self.shape_policy.dims[0]) + if self.embedding is not None: # a graph that left its token table to the host + tokens["token_embeds"] = self.embedding[next(iter(tokens.values()))] + res: NDArray[np.float32] = graph.run(None, {node.name: tokens[node.name] for node in graph.get_inputs()})[0][0] return serialize_np_array(res) def _load(self) -> ModelSession: session = super()._load() + table = session.for_shape(self.shape_policy.dims[0]).get_metadata().get("embedding") + self.embedding = _mapped(self.model_dir / "model.onnx", table) if table is not None else None log.debug(f"Loading tokenizer for CLIP model '{self.model_name}'") self.tokenizer = self._load_tokenizer() tokenizer_kwargs: dict[str, Any] | None = self.text_cfg.get("tokenizer_kwargs") diff --git a/machine-learning/immich_ml/models/clip/visual.py b/machine-learning/immich_ml/models/clip/visual.py index 37362fc53..fe627b738 100644 --- a/machine-learning/immich_ml/models/clip/visual.py +++ b/machine-learning/immich_ml/models/clip/visual.py @@ -19,7 +19,7 @@ from immich_ml.models.transforms import ( serialize_np_array, to_numpy, ) -from immich_ml.schemas import ModelSession, ModelTask, ModelType +from immich_ml.schemas import ModelGraph, ModelInput, ModelSession, ModelTask, ModelType class BaseCLIPVisualEncoder(InferenceModel): @@ -28,11 +28,12 @@ class BaseCLIPVisualEncoder(InferenceModel): def _predict(self, inputs: Image.Image | bytes) -> str: image = decode_pil(inputs) - res: NDArray[np.float32] = self.session.run(None, self.transform(image))[0][0] + session = self.session.for_shape(self.shape_policy.dims[0]) + res: NDArray[np.float32] = session.run(None, self.transform(session, image))[0][0] return serialize_np_array(res) @abstractmethod - def transform(self, image: Image.Image) -> dict[str, NDArray[np.float32]]: + def transform(self, session: ModelGraph, image: Image.Image) -> ModelInput: pass @property @@ -69,11 +70,12 @@ class OpenClipVisualEncoder(BaseCLIPVisualEncoder): return super()._load() - def transform(self, image: Image.Image) -> dict[str, NDArray[np.float32]]: + def transform(self, session: ModelGraph, image: Image.Image) -> ModelInput: image = self._resize(image) - 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)} + if session.normalizes_input: + rgb: NDArray[np.uint8] = np.asarray(image if image.mode == "RGB" else image.convert("RGB")) + return {"image": rgb[None]} + return {"image": normalize(to_numpy(image), self.mean, self.std).transpose(2, 0, 1)[None]} def _resize(self, image: Image.Image) -> Image.Image: match self.preprocess_cfg.get("resize_mode", "shortest"): # open_clip's default for older configs diff --git a/machine-learning/immich_ml/models/facial_recognition/_ops.py b/machine-learning/immich_ml/models/facial_recognition/_ops.py index 6bdf5f5c3..1fddb5a4c 100644 --- a/machine-learning/immich_ml/models/facial_recognition/_ops.py +++ b/machine-learning/immich_ml/models/facial_recognition/_ops.py @@ -8,11 +8,11 @@ from functools import lru_cache import cv2 import numpy as np +from immich_model.constants import FACE_DETECTION_SIZE as DET_SIZE from numpy.typing import NDArray from immich_ml.models.transforms import ensure_dims -DET_SIZE = 640 ALIGNED_SIZE = 112 # the FPN levels for which the fused detector emits a (scores, boxes, kps) triple, @@ -69,6 +69,5 @@ def umeyama(src: NDArray[np.float32], dst: NDArray[np.float32]) -> NDArray[np.fl return np.hstack([scale * rotation, translation[:, None]], dtype=np.float32) -def align_face(image: NDArray[np.uint8], kps: NDArray[np.float32]) -> NDArray[np.float32]: - matrix = umeyama(kps, ARCFACE_DST) - return cv2.warpAffine(image, matrix, (ALIGNED_SIZE, ALIGNED_SIZE)).astype(np.float32) +def align_face(image: NDArray[np.uint8], kps: NDArray[np.float32], crop: NDArray[np.uint8]) -> None: + cv2.warpAffine(image, umeyama(kps, ARCFACE_DST), (ALIGNED_SIZE, ALIGNED_SIZE), dst=crop) diff --git a/machine-learning/immich_ml/models/facial_recognition/detection.py b/machine-learning/immich_ml/models/facial_recognition/detection.py index a11375520..689b5c870 100644 --- a/machine-learning/immich_ml/models/facial_recognition/detection.py +++ b/machine-learning/immich_ml/models/facial_recognition/detection.py @@ -1,23 +1,30 @@ import numpy as np +from immich_model.constants import FACE_DETECTION_SIZE as DET_SIZE from numpy.typing import NDArray from immich_ml.models.base import InferenceModel -from immich_ml.models.transforms import decode_pil, letterbox, normalize -from immich_ml.schemas import FaceDetectionOutput, ModelTask, ModelType +from immich_ml.models.transforms import decode_pil, letterbox, normalize, widen +from immich_ml.schemas import FaceDetectionOutput, ModelTask, ModelType, Shape +from immich_ml.sessions.policy import ShapePolicy -from ._ops import DET_SIZE, decode_scrfd, nms +from ._ops import decode_scrfd, nms class FaceDetector(InferenceModel): depends = [] identity = (ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION) + shape_policy = ShapePolicy(dims=(Shape(batch=1, height=DET_SIZE, width=DET_SIZE),)) def _predict(self, inputs: NDArray[np.uint8] | bytes, minScore: float) -> FaceDetectionOutput: + session = self.session.for_shape(self.shape_policy.dims[0]) canvas, scale = letterbox(decode_pil(inputs), DET_SIZE) - blob = normalize(canvas.astype(np.float32), mean=127.5, std=128).transpose(2, 0, 1)[None] + blob: NDArray[np.float32] | NDArray[np.uint8] = ( + canvas[None] + if session.normalizes_input + else normalize(canvas.astype(np.float32), 127.5, 128).transpose(2, 0, 1)[None] + ) - input_name = self.session.get_inputs()[0].name - heads = self.session.run(None, {input_name: blob}) + heads = [widen(head) for head in session.run(None, {session.get_inputs()[0].name: blob})] scores, boxes, kps = decode_scrfd(heads, DET_SIZE) candidates = scores >= minScore diff --git a/machine-learning/immich_ml/models/facial_recognition/recognition.py b/machine-learning/immich_ml/models/facial_recognition/recognition.py index 09e48350a..9f8058dc2 100644 --- a/machine-learning/immich_ml/models/facial_recognition/recognition.py +++ b/machine-learning/immich_ml/models/facial_recognition/recognition.py @@ -1,26 +1,16 @@ -from pathlib import Path from typing import Any import numpy as np -import onnx -import onnxruntime as ort from numpy.typing import NDArray -from onnx.tools.update_model_dims import update_inputs_outputs_dims from PIL import Image -from immich_ml.config import log, settings +from immich_ml.config import settings from immich_ml.models.base import InferenceModel from immich_ml.models.transforms import decode_pil, normalize, serialize_np_array -from immich_ml.schemas import ( - FaceDetectionOutput, - FacialRecognitionOutput, - ModelFormat, - ModelSession, - ModelTask, - ModelType, -) +from immich_ml.schemas import FaceDetectionOutput, FacialRecognitionOutput, ModelTask, ModelType, Shape +from immich_ml.sessions.policy import ShapePolicy, batches, runs -from ._ops import align_face +from ._ops import ALIGNED_SIZE, align_face class FaceRecognizer(InferenceModel): @@ -29,15 +19,8 @@ class FaceRecognizer(InferenceModel): def __init__(self, model_name: str, **model_kwargs: Any) -> None: super().__init__(model_name, **model_kwargs) - max_batch_size = settings.max_batch_size and settings.max_batch_size.facial_recognition - self.batch_size = max_batch_size if max_batch_size else self._batch_size_default - - def _load(self) -> ModelSession: - session = self._make_session() - if (not self.batch_size or self.batch_size > 1) and str(session.get_inputs()[0].shape[0]) != "batch": - self._add_batch_axis(self.model_path) - session = self._make_session() - return session + sizes = batches(settings.max_batch_size.facial_recognition) + self.shape_policy = ShapePolicy(dims=tuple(Shape(batch) for batch in sizes)) def _predict( self, inputs: NDArray[np.uint8] | bytes | Image.Image, faces: FaceDetectionOutput @@ -45,20 +28,26 @@ class FaceRecognizer(InferenceModel): if faces["boxes"].shape[0] == 0: return [] image = np.asarray(decode_pil(inputs), dtype=np.uint8) - crops = np.stack([align_face(image, kps).transpose(2, 0, 1) for kps in faces["landmarks"]]) - embeddings = self._predict_batch(normalize(crops, mean=127.5, std=127.5)) + landmarks = faces["landmarks"] + sizes = runs(len(landmarks), self.session.batches) + crops: NDArray[np.float32] | NDArray[np.uint8] = np.empty( + (len(landmarks), ALIGNED_SIZE, ALIGNED_SIZE, 3), dtype=np.uint8 + ) + for crop, kps in zip(crops, landmarks): + align_face(image, kps, crop) + if not self.session.for_shape(Shape(sizes[0])).normalizes_input: + crops = normalize(crops.transpose(0, 3, 1, 2).astype(np.float32, order="C"), mean=127.5, std=127.5) + embeddings = self._predict_batch(crops, sizes) return self.postprocess(faces, embeddings) - def _predict_batch(self, crops: NDArray[np.float32]) -> NDArray[np.float32]: - input_name = self.session.get_inputs()[0].name - if not self.batch_size or crops.shape[0] <= self.batch_size: - return self.session.run(None, {input_name: crops})[0] - - batches = [ - self.session.run(None, {input_name: crops[i : i + self.batch_size]})[0] - for i in range(0, crops.shape[0], self.batch_size) - ] - return np.concatenate(batches, axis=0) + def _predict_batch(self, crops: NDArray[np.float32] | NDArray[np.uint8], sizes: list[int]) -> NDArray[np.float32]: + embeddings = [] + start = 0 + for size in sizes: + session = self.session.for_shape(Shape(size)) + embeddings.append(session.run(None, {session.get_inputs()[0].name: crops[start : start + size]})[0]) + start += size + return embeddings[0] if len(embeddings) == 1 else np.concatenate(embeddings, axis=0) def postprocess(self, faces: FaceDetectionOutput, embeddings: NDArray[np.float32]) -> FacialRecognitionOutput: return [ @@ -69,24 +58,3 @@ class FaceRecognizer(InferenceModel): } for (x1, y1, x2, y2), embedding, score in zip(faces["boxes"], embeddings, faces["scores"]) ] - - def _add_batch_axis(self, model_path: Path) -> None: - log.debug(f"Adding batch axis to model {model_path}") - proto = onnx.load(model_path) - static_input_dims = [shape.dim_value for shape in proto.graph.input[0].type.tensor_type.shape.dim[1:]] - static_output_dims = [shape.dim_value for shape in proto.graph.output[0].type.tensor_type.shape.dim[1:]] - input_dims = {proto.graph.input[0].name: ["batch"] + static_input_dims} - output_dims = {proto.graph.output[0].name: ["batch"] + static_output_dims} - updated_proto = update_inputs_outputs_dims(proto, input_dims, output_dims) - onnx.save(updated_proto, model_path) - - @property - def _batch_size_default(self) -> int | None: - providers = ort.get_available_providers() - if ( - self.model_format == ModelFormat.ONNX - and "MIGraphXExecutionProvider" not in providers - and "OpenVINOExecutionProvider" not in providers - ): - return None - return 1 diff --git a/machine-learning/immich_ml/models/ocr/ctc.py b/machine-learning/immich_ml/models/ocr/ctc.py index 90959b1b9..aa2d720fd 100644 --- a/machine-learning/immich_ml/models/ocr/ctc.py +++ b/machine-learning/immich_ml/models/ocr/ctc.py @@ -1,38 +1,83 @@ -from collections.abc import Sequence +from collections.abc import Callable, Sequence from pathlib import Path from typing import Any, Self import numpy as np from numpy.typing import NDArray +from immich_ml.models.transforms import widen -def greedy(outputs: Sequence[NDArray[Any]]) -> tuple[NDArray[np.int32], NDArray[np.float32]]: - """Models with an in-graph argmax head emit (indices, probs); the rest emit raw logits.""" - if len(outputs) == 2: - return outputs[0], outputs[1] - (probs,) = outputs - indices = probs.argmax(axis=2) - return indices.astype(np.int32), np.take_along_axis(probs, indices[:, :, None], axis=2)[:, :, 0] + +def picked(indices: NDArray[np.int32], confidence: NDArray[Any]) -> tuple[NDArray[np.int32], NDArray[np.float32]]: + return indices, widen(confidence) + + +def kept(indices: NDArray[np.integer[Any]]) -> NDArray[np.bool_]: + """The steps a greedy CTC decode reads a character from: not blank, and not a repeat of the step before.""" + keep = np.empty(indices.shape, dtype=bool) + keep[:, 0] = True + np.not_equal(indices[:, 1:], indices[:, :-1], out=keep[:, 1:]) + keep &= indices != 0 + return keep + + +def logits(raw: NDArray[Any]) -> tuple[NDArray[np.int32], NDArray[np.float32]]: + steps = raw.reshape(-1, raw.shape[-1]) # a binary keeps a unit axis between the steps and the classes + indices = _argmax(steps) + keep = kept(indices.reshape(raw.shape[:2])).ravel() + rows = widen(steps[keep]) # a confidence is only read where a character is kept + best = rows[np.arange(len(rows)), indices[keep]] + # the winner's softmax, summing only the few classes within exp(-16) of it rather than every class + near = rows > (best - 16)[:, None] + counts = np.count_nonzero(near, axis=1) + confidence = np.zeros(len(steps), np.float32) + if len(rows): + weights = np.exp(rows[near] - np.repeat(best, counts)) + confidence[keep] = 1 / np.add.reduceat(weights, np.cumsum(counts) - counts) + return indices.reshape(raw.shape[:2]).astype(np.int32), confidence.reshape(raw.shape[:2]) + + +def _argmax(steps: NDArray[Any]) -> NDArray[np.intp]: + if steps.dtype != np.float16: + return steps.argmax(axis=1) + bits = steps.view(np.int16) # a half's bits order like its value where it is not negative + indices = bits.argmax(axis=1) + negative = np.flatnonzero(bits[np.arange(len(bits)), indices] < 0) # a step with no logit above zero + if len(negative): + flipped = bits[negative] ^ ((bits[negative] >> 15) & np.int16(0x7FFF)) # negatives' order restored + indices[negative] = flipped.argmax(axis=1) + return indices + + +def probabilities(probs: NDArray[Any]) -> tuple[NDArray[np.int32], NDArray[np.float32]]: + # probabilities are non-negative, so half-precision bits order like their values and need no widening to compare + indices = (probs.view(np.uint16) if probs.dtype == np.float16 else probs).argmax(axis=2) + return indices.astype(np.int32), widen(np.take_along_axis(probs, indices[:, :, None], axis=2)[:, :, 0]) + + +Greedy = Callable[..., tuple[NDArray[np.int32], NDArray[np.float32]]] +GREEDY: dict[str, Greedy] = {"ctc_indices": picked, "ctc_logits": logits, "logits_softmax": probabilities} class CtcDecoder: - def __init__(self, charset: list[str]): + def __init__(self, charset: list[str], greedy: Greedy): self.charset = ["", *charset, " "] # PP-OCR: blank at 0, space last + self.greedy = greedy @classmethod - def from_file(cls, charset_path: Path) -> Self: + def from_file(cls, charset_path: Path, greedy: Greedy) -> Self: if not charset_path.is_file(): raise FileNotFoundError(f"Recognition charset not found: {charset_path}") - return cls(charset_path.read_text(encoding="utf-8").splitlines()) + return cls(charset_path.read_text(encoding="utf-8").splitlines(), greedy) def __len__(self) -> int: return len(self.charset) + def __call__(self, outputs: Sequence[NDArray[Any]]) -> tuple[list[str], NDArray[np.float32]]: + return self.decode(*self.greedy(*outputs)) + def decode(self, indices: NDArray[np.int32], probs: NDArray[np.float32]) -> tuple[list[str], NDArray[np.float32]]: - keep = np.empty(indices.shape, dtype=bool) - keep[:, 0] = True - np.not_equal(indices[:, 1:], indices[:, :-1], out=keep[:, 1:]) # repeats before blanks - keep &= indices != 0 + keep = kept(indices) scores: NDArray[np.float32] = np.where(keep, probs, 0).sum(1) scores /= np.maximum(keep.sum(1), 1) # an all-blank row sums to 0, so it stays 0 diff --git a/machine-learning/immich_ml/models/ocr/detection.py b/machine-learning/immich_ml/models/ocr/detection.py index 258058514..563f8a833 100644 --- a/machine-learning/immich_ml/models/ocr/detection.py +++ b/machine-learning/immich_ml/models/ocr/detection.py @@ -1,29 +1,31 @@ -from functools import cached_property +from collections.abc import Sequence from typing import Any import numpy as np +from immich_model.constants import ocr_canvases from numpy.typing import NDArray from PIL import Image -from rapidocr.inference_engine.base import FileInfo, InferSession -from rapidocr.utils.download_file import DownloadFile, DownloadFileInput -from rapidocr.utils.typings import EngineType, LangDet, OCRVersion, TaskType -from rapidocr.utils.typings import ModelType as RapidModelType -from immich_ml.config import log -from immich_ml.models.base import InferenceModel -from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType -from immich_ml.sessions.ort import OrtSession +from immich_ml.models.transforms import widen +from immich_ml.schemas import ModelGraph, ModelTask, ModelType, Shape +from immich_ml.sessions.policy import ShapePolicy +from .legacy import TextModel from .postprocess import DBPostProcess from .schemas import TextDetectionOutput -class TextDetector(InferenceModel): +class TextDetector(TextModel): depends = [] identity = (ModelType.DETECTION, ModelTask.OCR) + graph_options = ("maxResolution",) def __init__(self, model_name: str, **model_kwargs: Any) -> None: - super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX) + super().__init__(model_name, **model_kwargs) + short_side = model_kwargs.get("maxResolution", 736) + # RKNPU ships a binary per short side, which the label picks + canvases = tuple(Shape(batch=1, **canvas) for canvas in ocr_canvases(short_side)) + self.shape_policy = ShapePolicy(dims=canvases, label=f"res{short_side}") self.scale = np.float32(1.0 / 127.5) # (x/255 - 0.5) / 0.5 self._empty: TextDetectionOutput = { "boxes": np.empty(0, dtype=np.float32), @@ -31,28 +33,6 @@ class TextDetector(InferenceModel): } self.postprocess = DBPostProcess(thresh=0.3, max_candidates=1000, unclip_ratio=1.6, use_dilation=True) - def _download(self) -> None: - model_info = InferSession.get_model_url( - FileInfo( - engine_type=EngineType.ONNXRUNTIME, - ocr_version=OCRVersion.PPOCRV5, - task_type=TaskType.DET, - lang_type=LangDet.CH, - model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER, - ) - ) - download_params = DownloadFileInput( - file_url=model_info["model_dir"], - sha256=model_info["SHA256"], - save_path=self.model_path, - logger=log, - ) - DownloadFile.run(download_params) - - def _load(self) -> ModelSession: - # TODO: support other runtime sessions - return OrtSession(self.model_path) - def _predict( self, inputs: Image.Image, maxResolution: int = 736, minScore: float = 0.5, scoreMode: str = "fast" ) -> TextDetectionOutput: @@ -60,9 +40,8 @@ class TextDetector(InferenceModel): if width < 32 or height < 32: return self._empty - image, content = self._transform(inputs, maxResolution) - input_name = self.session.get_inputs()[0].name - probs = self.session.run(None, {input_name: image})[0][0] + session, image, content = self._transform(inputs, maxResolution) + probs = widen(session.run(None, {session.get_inputs()[0].name: image})[0][0]) if probs.ndim == 3: probs = probs[0] boxes, scores = self.postprocess(probs[: content[0], : content[1]], (height, width), minScore, scoreMode) @@ -71,27 +50,46 @@ class TextDetector(InferenceModel): order = self.reading_order(boxes) return {"boxes": boxes[order], "scores": scores[order]} - @cached_property - def raw_input(self) -> bool: - return self.session.get_inputs()[0].shape[-1] == 3 # NHWC models handle normalization and transpose internally - - def _transform(self, img: Image.Image, max_resolution: int) -> tuple[NDArray[Any], tuple[int, int]]: + def _transform(self, img: Image.Image, max_resolution: int) -> tuple[ModelGraph, NDArray[Any], tuple[int, int]]: ratio = min(max_resolution / min(img.height, img.width), 1.0) - resize_h = max(self._round32(img.height * ratio), 32) - resize_w = max(self._round32(img.width * ratio), 32) - resized = img.resize((resize_w, resize_h), resample=Image.Resampling.LANCZOS) + content_h = max(self._round32(img.height * ratio), 32) + content_w = max(self._round32(img.width * ratio), 32) + canvas = self._canvas(self.session.shapes, content_h, content_w) + session = self.session.for_shape(Shape(batch=1, height=canvas[0], width=canvas[1])) + # a compiled canvas may be smaller than the image was sized for, so it shrinks again + scale = min(canvas[0] / content_h, canvas[1] / content_w, 1.0) + content_h = min(max(self._round32(content_h * scale), 32), canvas[0]) + content_w = min(max(self._round32(content_w * scale), 32), canvas[1]) + resized = img.resize((content_w, content_h), resample=Image.Resampling.LANCZOS) if resized.mode != "RGB": resized = resized.convert("RGB") array = np.asarray(resized) - if self.raw_input: - return array[None], (resize_h, resize_w) + content = (content_h, content_w) + height, width = canvas + if session.normalizes_input: + if canvas == content: + return session, array[None], content + image = np.zeros((1, height, width, 3), dtype=np.uint8) # what is left of the canvas stays black + image[0, :content_h, :content_w] = array + return session, image, content # reverse plane order gets the BGR swap and the CHW transpose for free - out = np.empty((1, 3, resize_h, resize_w), dtype=np.float32) + out = np.empty((1, 3, height, width), dtype=np.float32) + if canvas != content: # -1.0 is what the graph would have made of a black pixel + out[0, :, content_h:] = -1.0 + out[0, :, :content_h, content_w:] = -1.0 for channel in range(3): - plane = out[0, 2 - channel] + plane = out[0, 2 - channel, :content_h, :content_w] np.multiply(array[:, :, channel], self.scale, out=plane) plane -= 1.0 - return out, (resize_h, resize_w) + return session, out, content + + @staticmethod + def _canvas(shapes: Sequence[Shape], height: int, width: int) -> tuple[int, int]: + return min( + ((shape.height, shape.width) for shape in shapes if shape.height and shape.width), + key=lambda canvas: (-min(canvas[0] / height, canvas[1] / width, 1.0), canvas[0] * canvas[1]), + default=(height, width), + ) @staticmethod def _round32(value: float) -> int: diff --git a/machine-learning/immich_ml/models/ocr/legacy.py b/machine-learning/immich_ml/models/ocr/legacy.py new file mode 100644 index 000000000..d17ec0339 --- /dev/null +++ b/machine-learning/immich_ml/models/ocr/legacy.py @@ -0,0 +1,42 @@ +from typing import Any + +from immich_ml.config import log, settings +from immich_ml.models.base import InferenceModel +from immich_ml.schemas import ModelFormat, ModelType + + +# TODO: Remove once everything uses the new model graphs +class TextModel(InferenceModel): + def __init__(self, model_name: str, **model_kwargs: Any) -> None: + if settings.legacy_models: + model_kwargs["model_format"] = ModelFormat.ONNX # the older exports come in no other format + super().__init__(model_name, **model_kwargs) + + def download(self) -> None: + if not settings.legacy_models or self.cached: + return super().download() + # here, as only an install still on the older exports has a use for it + from rapidocr.inference_engine.base import FileInfo, InferSession + from rapidocr.utils.download_file import DownloadFile, DownloadFileInput + from rapidocr.utils.typings import EngineType, LangDet, LangRec, OCRVersion, TaskType + from rapidocr.utils.typings import ModelType as RapidModelType + + detects = self.model_type == ModelType.DETECTION + language = self.model_name.split("__")[0] if "__" in self.model_name else "CH" + model_info = InferSession.get_model_url( + FileInfo( + engine_type=EngineType.ONNXRUNTIME, + ocr_version=OCRVersion.PPOCRV5, + task_type=TaskType.DET if detects else TaskType.REC, + lang_type=LangDet.CH if detects else LangRec[language], + model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER, + ) + ) + log.info( + f"Downloading {self.model_type} model '{self.model_name}' to {self.model_path}. This may take a while." + ) + DownloadFile.run( + DownloadFileInput( + file_url=model_info["model_dir"], sha256=model_info["SHA256"], save_path=self.model_path, logger=log + ) + ) diff --git a/machine-learning/immich_ml/models/ocr/recognition.py b/machine-learning/immich_ml/models/ocr/recognition.py index dc445e760..7cea06df2 100644 --- a/machine-learning/immich_ml/models/ocr/recognition.py +++ b/machine-learning/immich_ml/models/ocr/recognition.py @@ -1,77 +1,57 @@ import math -from functools import cached_property +from bisect import bisect_left from typing import Any import cv2 import numpy as np +from immich_model.constants import OCR_RECOGNITION_WIDTHS from numpy.typing import NDArray from PIL import Image -from rapidocr import LangRec -from rapidocr.inference_engine.base import FileInfo, InferSession -from rapidocr.utils.download_file import DownloadFile, DownloadFileInput -from rapidocr.utils.typings import EngineType, OCRVersion, TaskType -from rapidocr.utils.typings import ModelType as RapidModelType -from immich_ml.config import log, settings -from immich_ml.models.base import InferenceModel -from immich_ml.schemas import ModelFormat, ModelSession, ModelTask, ModelType -from immich_ml.sessions.ort import OrtSession +from immich_ml.config import settings +from immich_ml.schemas import ModelGraph, ModelSession, ModelTask, ModelType, Shape +from immich_ml.sessions.policy import ShapePolicy, batches, runs -from .ctc import CtcDecoder, greedy +from .ctc import GREEDY, CtcDecoder, probabilities +from .legacy import TextModel from .schemas import TextDetectionOutput, TextRecognitionOutput REC_HEIGHT = 48 -REC_BASE_RATIO = 320 / REC_HEIGHT # PP-OCR's rec_image_shape floor SCALE = np.float32(1.0 / 127.5) -class TextRecognizer(InferenceModel): +class TextRecognizer(TextModel): depends = [(ModelType.DETECTION, ModelTask.OCR)] identity = (ModelType.RECOGNITION, ModelTask.OCR) def __init__(self, model_name: str, **model_kwargs: Any) -> None: - self.language = LangRec[model_name.split("__")[0]] if "__" in model_name else LangRec.CH self._empty: TextRecognitionOutput = { "box": np.empty(0, dtype=np.float32), "boxScore": np.empty(0, dtype=np.float32), "text": [], "textScore": np.empty(0, dtype=np.float32), } - super().__init__(model_name, **model_kwargs, model_format=ModelFormat.ONNX) - max_batch_size = settings.max_batch_size and settings.max_batch_size.ocr - self.batch_size = max_batch_size if max_batch_size else 6 - - def _download(self) -> None: - model_info = InferSession.get_model_url( - FileInfo( - engine_type=EngineType.ONNXRUNTIME, - ocr_version=OCRVersion.PPOCRV5, - task_type=TaskType.REC, - lang_type=self.language, - model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER, - ) + super().__init__(model_name, **model_kwargs) + self.decoder: CtcDecoder | None = None + sizes = batches(settings.max_batch_size.ocr) + self.shape_policy = ShapePolicy( + dims=tuple(Shape(batch, width=width) for width in OCR_RECOGNITION_WIDTHS for batch in sizes) ) - download_params = DownloadFileInput( - file_url=model_info["model_dir"], - sha256=model_info["SHA256"], - save_path=self.model_path, - logger=log, - ) - DownloadFile.run(download_params) def _load(self) -> ModelSession: - # TODO: support other runtimes - session = OrtSession(self.model_path) - self.decoder = ( - CtcDecoder(character.splitlines()) - if (character := session.get_metadata().get("character")) is not None - else CtcDecoder.from_file(self.model_dir / "charset.txt") - ) + session = super()._load() + # the widths it was compiled for, narrowest first: none where it takes any + self.widths = tuple(sorted({shape.width for shape in session.shapes if shape.width is not None})) return session - @cached_property - def raw_input(self) -> bool: - return self.session.get_inputs()[0].shape[-1] == 3 # NHWC models handle normalization and transpose internally + def _decoder(self, session: ModelGraph) -> CtcDecoder: + character = session.get_metadata().get("character") + greedy = GREEDY.get(session.get_outputs()[0].name, probabilities) + return ( + CtcDecoder(character.splitlines(), greedy) + if character is not None + else CtcDecoder.from_file(self.model_dir / "charset.txt", greedy) + ) def _predict(self, img: Image.Image, texts: TextDetectionOutput, minScore: float = 0.9) -> TextRecognitionOutput: boxes, box_scores = texts["boxes"], texts["scores"] @@ -85,28 +65,35 @@ class TextRecognizer(InferenceModel): text_list: list[str] = [""] * len(order) score_list = np.zeros(len(order), dtype=np.float32) - input_name = self.session.get_inputs()[0].name - for start in range(0, len(order), self.batch_size): - chunk = order[start : start + self.batch_size] - width = int(REC_HEIGHT * max(REC_BASE_RATIO, ratios[chunk[-1]])) - images: NDArray[Any] = ( - np.full((len(chunk), REC_HEIGHT, width, 3), 127, dtype=np.uint8) - if self.raw_input - else np.zeros((len(chunk), 3, REC_HEIGHT, width), dtype=np.float32) - ) + sized = runs(len(order), self.session.batches) + bounds = [sum(sized[:index]) for index in range(len(sized))] + chunked = [order[start : start + size] for start, size in zip(bounds, sized)] + batch_widths = [self._width(int(REC_HEIGHT * ratios[chunk[-1]])) for chunk in chunked] + graphs = [self.session.for_shape(Shape(len(chunk), width=width)) for chunk, width in zip(chunked, batch_widths)] + raw = graphs[0].normalizes_input + # one buffer per request, sized for the largest batch and reshaped for each + sizes = [len(chunk) * REC_HEIGHT * width * 3 for chunk, width in zip(chunked, batch_widths)] + buffer = np.empty(max(sizes), dtype=np.uint8 if raw else np.float32) + for chunk, width, size, session in zip(chunked, batch_widths, sizes, graphs): + shape = (len(chunk), REC_HEIGHT, width, 3) if raw else (len(chunk), 3, REC_HEIGHT, width) + images: NDArray[Any] = buffer[:size].reshape(shape) for i, index in enumerate(chunk): crop = self._crop(img, coeffs[index], widths[index], heights[index], upright[index]) resized_w = max(1, min(width, math.ceil(REC_HEIGHT * ratios[index]))) - if self.raw_input: + if raw: cv2.resize(crop, (resized_w, REC_HEIGHT), dst=images[i, :, :resized_w]) + images[i, :, resized_w:] = 127 # the pad the crop is read against, over whatever was there continue resized = cv2.resize(crop, (resized_w, REC_HEIGHT)) view = images[i, :, :, :resized_w] np.multiply(resized.transpose(2, 0, 1)[::-1], SCALE, out=view) # [::-1] is the RGB -> BGR swap view -= 1.0 + images[i, :, :, resized_w:] = 0 - out_indices, out_probs = greedy(self.session.run(None, {input_name: images})) - chunk_texts, chunk_scores = self.decoder.decode(out_indices, out_probs) + outputs = session.run(None, {session.get_inputs()[0].name: images}) + if self.decoder is None: + self.decoder = self._decoder(session) + chunk_texts, chunk_scores = self.decoder(outputs) for index, text, score in zip(chunk, chunk_texts, chunk_scores): text_list[index] = text score_list[index] = score @@ -123,6 +110,11 @@ class TextRecognizer(InferenceModel): "textScore": score_list[valid], } + def _width(self, width: int) -> int: + if not self.widths: # room past the text in proportion to it, as the recognizer reads a line worse without + return min(OCR_RECOGNITION_WIDTHS[-1], max(OCR_RECOGNITION_WIDTHS[0], math.ceil(width * 1.25))) + return self.widths[min(bisect_left(self.widths, width), len(self.widths) - 1)] + def _crop_geometry( self, boxes: NDArray[np.float32] ) -> tuple[NDArray[np.int32], NDArray[np.int32], NDArray[np.float32]]: diff --git a/machine-learning/immich_ml/models/transforms.py b/machine-learning/immich_ml/models/transforms.py index ac220a381..7d473a90a 100644 --- a/machine-learning/immich_ml/models/transforms.py +++ b/machine-learning/immich_ml/models/transforms.py @@ -1,5 +1,6 @@ import string from io import BytesIO +from typing import Any import cv2 import numpy as np @@ -43,10 +44,6 @@ def get_pil_resampling(resample: str) -> Image.Resampling: return _PIL_RESAMPLING_METHODS[resample.lower()] -def pil_to_cv2(image: Image.Image) -> NDArray[np.uint8]: - return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # type: ignore - - def decode_pil(image_bytes: bytes | Image.Image | NDArray[np.uint8]) -> Image.Image: image: Image.Image match image_bytes: @@ -75,6 +72,13 @@ def serialize_np_array(arr: NDArray[np.float32]) -> str: return orjson.dumps(arr, option=orjson.OPT_SERIALIZE_NUMPY).decode() +def widen(array: NDArray[Any]) -> NDArray[Any]: + """Single precision from half through OpenCV, as numpy converts one value at a time without native half support.""" + if array.dtype != np.float16: + return array + return cv2.convertFp16(array.reshape(-1, array.shape[-1]).view(np.int16)).reshape(array.shape) # int16 carries half + + def letterbox(image: NDArray[np.uint8] | Image.Image, size: int) -> tuple[NDArray[np.uint8], float]: if isinstance(image, Image.Image): image = np.asarray(image) diff --git a/machine-learning/immich_ml/schemas.py b/machine-learning/immich_ml/schemas.py index a88faac9d..466e2553b 100644 --- a/machine-learning/immich_ml/schemas.py +++ b/machine-learning/immich_ml/schemas.py @@ -1,6 +1,7 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass from enum import Enum -from typing import Any, Literal, Protocol, TypeVar +from typing import Any, Literal, Protocol, TypeAlias, TypeVar import numpy as np import numpy.typing as npt @@ -44,11 +45,6 @@ class ModelSource(StrEnum): PADDLE = "paddle" -class ModelPrecision(StrEnum): - FP16 = "FP16" - FP32 = "FP32" - - class ModelOrganization(StrEnum): APP = "immich-app" TESTING = "immich-testing" @@ -62,16 +58,29 @@ class SessionNode(Protocol): def name(self) -> str: ... @property - def shape(self) -> tuple[int, ...]: ... + def shape(self) -> tuple[int | str, ...]: ... # ORT names a symbolic dim rather than sizing it -class ModelSession(Protocol): +ModelTensor: TypeAlias = npt.NDArray[np.float32] | npt.NDArray[np.int32] | npt.NDArray[np.uint8] +ModelInput = Mapping[str, ModelTensor] + + +@dataclass(frozen=True) +class Shape: + batch: int + height: int | None = None + width: int | None = None + + @property + def pins(self) -> dict[str, int]: + return {name: size for name, size in asdict(self).items() if size is not None} + + +class ModelGraph(Protocol): def run( self, output_names: list[str] | None, - input_feed: dict[str, npt.NDArray[np.float32]] - | dict[str, npt.NDArray[np.int32]] - | dict[str, npt.NDArray[np.uint8]], + input_feed: ModelInput, run_options: Any = None, ) -> list[npt.NDArray[np.float32]]: ... @@ -81,6 +90,21 @@ class ModelSession(Protocol): def get_metadata(self) -> dict[str, str]: ... + @property + def normalizes_input(self) -> bool: ... + + +class ModelSession(Protocol): + @property + def shapes(self) -> tuple[Shape, ...]: ... + + @property + def batches(self) -> tuple[int, ...]: ... + + def for_shape(self, shape: Shape) -> ModelGraph: ... + + def warm(self) -> None: ... + class FaceDetectionOutput(TypedDict): boxes: npt.NDArray[np.float32] diff --git a/machine-learning/immich_ml/sessions/ann/__init__.py b/machine-learning/immich_ml/sessions/ann/__init__.py index c516f9ba9..d921eaca0 100644 --- a/machine-learning/immich_ml/sessions/ann/__init__.py +++ b/machine-learning/immich_ml/sessions/ann/__init__.py @@ -7,7 +7,7 @@ import numpy as np from numpy.typing import NDArray from immich_ml.config import log, settings -from immich_ml.schemas import SessionNode +from immich_ml.schemas import ModelInput, SessionNode, Shape from .loader import Ann @@ -29,6 +29,9 @@ class AnnSession: fp16=settings.ann_fp16_turbo, ) log.info("Loaded ANN model with ID %d", self.model) + batch = self.ann.input_shapes[self.model][0][0] + self.shapes = (Shape(batch),) # the artifact is compiled for one shape + self.batches = (batch,) def __del__(self) -> None: self.ann.unload(self.model) @@ -46,10 +49,20 @@ class AnnSession: def get_metadata(self) -> dict[str, str]: return {} + @property + def normalizes_input(self) -> bool: + return False # no ARM-NN artifact is built from a graph that carries its own preprocessing + + def for_shape(self, shape: Shape) -> AnnSession: + return self + + def warm(self) -> None: + pass # the artifact is loaded whole when the session opens + def run( self, output_names: list[str] | None, - input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]], + input_feed: ModelInput, run_options: Any = None, ) -> list[NDArray[np.float32]]: inputs: list[NDArray[np.float32]] = [np.ascontiguousarray(v) for v in input_feed.values()] diff --git a/machine-learning/immich_ml/sessions/ort.py b/machine-learning/immich_ml/sessions/ort.py index 1f887256b..4ff1a6857 100644 --- a/machine-learning/immich_ml/sessions/ort.py +++ b/machine-learning/immich_ml/sessions/ort.py @@ -1,77 +1,56 @@ from __future__ import annotations import ctypes +import pickle import platform +import subprocess import sys -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import cache, cached_property from pathlib import Path from threading import Lock -from typing import Any +from typing import Any, NamedTuple import numpy as np import onnxruntime as ort +from immich_model.runtime import RewriteContext, RewritePlan, plan_rewrites from numpy.typing import NDArray +from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf +from pydantic import BaseModel -from immich_ml.models.constants import SUPPORTED_PROVIDERS -from immich_ml.schemas import ModelPrecision, SessionNode +from immich_ml.schemas import ModelInput, SessionNode, Shape from ..config import log, settings +from .policy import ShapePolicy -MigraphxInputSignature = tuple[tuple[str, str, tuple[int, ...]], ...] - -_migraphx_registry_lock = Lock() -_migraphx_model_locks: dict[str, Lock] = {} -_migraphx_compiled_inputs: set[tuple[str, MigraphxInputSignature]] = set() +# the one provider that handles a free dim well; the rest miscompile it, recompile inside the run, +# or partition around it, so they are handed one graph per shape instead +DYNAMIC_PROVIDERS = frozenset({"CPUExecutionProvider"}) -def _migraphx_get_model_lock(model_key: str) -> Lock: - with _migraphx_registry_lock: - lock = _migraphx_model_locks.get(model_key) - if lock is None: - lock = Lock() - _migraphx_model_locks[model_key] = lock - return lock +def _label(pins: Mapping[str, int]) -> str: + return "_".join(f"{name}{size}" for name, size in pins.items()) -def _migraphx_has_compiled_input(key: tuple[str, MigraphxInputSignature]) -> bool: - with _migraphx_registry_lock: - return key in _migraphx_compiled_inputs +UNREADABLE = 3 # how the preparing child says its source does not parse +DTYPES = { + "tensor(float)": np.float32, + "tensor(float16)": np.float16, + "tensor(uint8)": np.uint8, + "tensor(int32)": np.int32, + "tensor(int64)": np.int64, +} -def _migraphx_mark_compiled_input(key: tuple[str, MigraphxInputSignature]) -> None: - with _migraphx_registry_lock: - _migraphx_compiled_inputs.add(key) - - -def _migraphx_input_signature(input_feed: dict[str, NDArray[Any]]) -> MigraphxInputSignature: - return tuple((name, str(value.dtype), tuple(value.shape)) for name, value in sorted(input_feed.items())) - - -class OrtSession: - session: ort.InferenceSession - - def __init__( - self, - model_path: Path | str, - providers: list[str] | None = None, - provider_options: list[dict[str, Any]] | None = None, - sess_options: ort.SessionOptions | None = None, - threads: int = 2, - ): - self.model_path = Path(model_path) - self.threads = threads - self.providers = providers if providers is not None else self._providers_default - self.disabled_optimizers = _disabled_optimizers_default(self.providers) - log.debug(f"Setting disabled_optimizers to {self.disabled_optimizers}") - self.provider_options = provider_options if provider_options is not None else self._provider_options_default - self.sess_options = sess_options if sess_options is not None else self._sess_options_default - self.session = ort.InferenceSession( - self.model_path.as_posix(), - providers=self.providers, - provider_options=self.provider_options, - sess_options=self.sess_options, - disabled_optimizers=self.disabled_optimizers, - ) +class OrtGraph: + def __init__(self, spec: GraphSpec) -> None: + self.session = spec.session(prepared(spec)) + image = self.session.get_inputs()[0] + # whether the graph scales and shifts the image itself, so the host hands over raw pixels + self.normalizes_input = image.type == "tensor(uint8)" + # CoreML's rewrite moves the layout change out of the graph, so whoever applied it performs it + self.channels_first = self.normalizes_input and len(image.shape) == 4 and image.shape[1] == 3 def get_inputs(self) -> Sequence[SessionNode]: inputs: Sequence[SessionNode] = self.session.get_inputs() @@ -88,52 +67,91 @@ class OrtSession: def run( self, output_names: list[str] | None, - input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]], + input_feed: ModelInput, run_options: Any = None, ) -> list[NDArray[np.float32]]: - if "MIGraphXExecutionProvider" in self.providers: - model_key = self.model_path.resolve().as_posix() - input_key = (model_key, _migraphx_input_signature(input_feed)) - if not _migraphx_has_compiled_input(input_key): - model_lock = _migraphx_get_model_lock(model_key) - with model_lock: - if not _migraphx_has_compiled_input(input_key): - outputs: list[NDArray[np.float32]] = self.session.run(output_names, input_feed, run_options) - _migraphx_mark_compiled_input(input_key) - return outputs - - outputs = self.session.run(output_names, input_feed, run_options) - return outputs - - outputs = self.session.run(output_names, input_feed, run_options) + if self.channels_first: + input_feed = {name: frames.transpose(0, 3, 1, 2) for name, frames in input_feed.items()} + outputs: list[NDArray[Any]] = self.session.run(output_names, input_feed, run_options) return outputs - @property - def providers(self) -> list[str]: - return self._providers - @providers.setter - def providers(self, providers: list[str]) -> None: - log.info(f"Setting execution providers to {providers}, in descending order of preference") - self._providers = providers +class Device(NamedTuple): + kind: str # what the provider compiles for, which names the directory + version: str # what the facts go stale under + + +@dataclass(frozen=True) +class GraphSpec: + model_path: Path + pins: Mapping[str, int] + overrides: Sequence[tuple[str, int]] + providers: list[str] + disabled_optimizers: list[str] + threads: int = 2 # for the CPU, where they are not configured @property - def _providers_default(self) -> list[str]: - available_providers = set(ort.get_available_providers()) - log.debug(f"Available ORT providers: {available_providers}") - return [provider for provider in SUPPORTED_PROVIDERS if provider in available_providers] + def provider(self) -> str: + return self.providers[0] # the rest take only what it cannot run + + @cached_property + def directory(self) -> Path: + provider = self.provider.removesuffix("ExecutionProvider").lower() + kind = self.device.kind if self.device else "" + return self.model_path.parent / provider / kind / _label(self.pins) + + @cached_property + def device(self) -> Device | None: + match self.provider: + case "CPUExecutionProvider": + return Device(platform.machine(), _cpu()) # a prepack built for other instruction sets fails at runtime + case "OpenVINOExecutionProvider" if self.openvino_device.startswith("GPU"): + return _intel_gpu(self.openvino_device) + case "MIGraphXExecutionProvider": + return _amd_gpu(int(settings.device_id)) + return None + + @cached_property + def openvino_device(self) -> str: + device_ids: list[str] = ort.capi._pybind_state.get_available_openvino_device_ids() + return f"GPU.{settings.device_id}" if any(d.startswith("GPU") for d in device_ids) else "CPU" + + @cached_property + def manifest(self) -> Path: + return self.directory / "manifest.json" @property + def plan(self) -> RewritePlan: + return _plan(self.provider) + + @cached_property + def half(self) -> bool: + return not (settings.legacy_models or self.provider == "CPUExecutionProvider") + + @cached_property + def facts(self) -> Facts: + return Facts( + half=self.half, + onnxruntime=ort.__version__, + # workers on devices of one kind share what is prepared + options={key: value for key, value in self.provider_options[0].items() if not key.startswith("device")}, + overrides=list(self.overrides), + disabled_optimizers=self.disabled_optimizers, + rewrite=None if settings.legacy_models else self.plan.digest, # which were never tested with the rewriter + device=self.device.version if self.device else None, + ) + + def session(self, graph: Path, sess_options: ort.SessionOptions | None = None) -> ort.InferenceSession: + return ort.InferenceSession( + graph.as_posix(), + providers=self.providers, + provider_options=self.provider_options, + sess_options=sess_options or self.sess_options(), + disabled_optimizers=self.disabled_optimizers, + ) + + @cached_property def provider_options(self) -> list[dict[str, Any]]: - return self._provider_options - - @provider_options.setter - def provider_options(self, provider_options: list[dict[str, Any]]) -> None: - log.debug(f"Setting execution provider options to {provider_options}") - self._provider_options = provider_options - - @property - def _provider_options_default(self) -> list[dict[str, Any]]: provider_options = [] for provider in self.providers: match provider: @@ -142,80 +160,245 @@ class OrtSession: case "CUDAExecutionProvider": options = {"arena_extend_strategy": "kSameAsRequested", "device_id": settings.device_id} case "MIGraphXExecutionProvider": - migraphx_dir = self.model_path.parent / "migraphx" - # MIGraphX does not create the underlying folder and will crash if it does not exist - migraphx_dir.mkdir(parents=True, exist_ok=True) - options = { - "device_id": settings.device_id, - "migraphx_model_cache_dir": migraphx_dir.as_posix(), - "migraphx_fp16_enable": "1" if settings.rocm_precision == ModelPrecision.FP16 else "0", - } + options = {"device_id": settings.device_id, "migraphx_model_cache_dir": self.directory.as_posix()} case "OpenVINOExecutionProvider": - device_ids: list[str] = ort.capi._pybind_state.get_available_openvino_device_ids() - # Check for available devices, preferring GPU over CPU - gpu_devices = [d for d in device_ids if d.startswith("GPU")] - if gpu_devices: - device_type = f"GPU.{settings.device_id}" - log.debug(f"OpenVINO: Using GPU device {device_type}") - else: - device_type = "CPU" - log.debug("OpenVINO: No GPU found, using CPU") - options = { - "device_type": device_type, - "precision": settings.openvino_precision.value, - "cache_dir": (self.model_path.parent / "openvino").as_posix(), - } + options = {"device_type": self.openvino_device, "cache_dir": self.directory.as_posix()} + if not self.half: # the GPU would otherwise run an fp32 graph at fp16 + options["precision"] = "FP32" case "CoreMLExecutionProvider": options = { "ModelFormat": "MLProgram", "MLComputeUnits": "ALL", "SpecializationStrategy": "FastPrediction", "AllowLowPrecisionAccumulationOnGPU": "1", - "ModelCacheDirectory": (self.model_path.parent / "coreml").as_posix(), + "ModelCacheDirectory": self.directory.as_posix(), } case _: options = {} provider_options.append(options) + log.debug(f"Setting execution provider options to {provider_options}") return provider_options - @property def sess_options(self) -> ort.SessionOptions: - return self._sess_options - - @sess_options.setter - def sess_options(self, sess_options: ort.SessionOptions) -> None: - log.debug(f"Setting execution_mode to {sess_options.execution_mode.name}") - log.debug(f"Setting inter_op_num_threads to {sess_options.inter_op_num_threads}") - log.debug(f"Setting intra_op_num_threads to {sess_options.intra_op_num_threads}") - self._sess_options = sess_options - - @property - def _sess_options_default(self) -> ort.SessionOptions: sess_options = ort.SessionOptions() + sess_options.enable_cpu_mem_arena = settings.model_arena # some CPUs slow down many times over on subnormal operands, and ORT clears the flush in its threads otherwise sess_options.add_session_config_entry("session.set_denormal_as_zero", "1") - sess_options.enable_cpu_mem_arena = settings.model_arena # avoid thread contention between models # Set inter_op threads if settings.model_inter_op_threads > 0: sess_options.inter_op_num_threads = settings.model_inter_op_threads # these defaults work well for CPU, but bottleneck GPU - elif settings.model_inter_op_threads == 0 and self.providers == ["CPUExecutionProvider"]: + elif settings.model_inter_op_threads == 0 and self.provider == "CPUExecutionProvider": sess_options.inter_op_num_threads = 1 # Set intra_op threads if settings.model_intra_op_threads > 0: sess_options.intra_op_num_threads = settings.model_intra_op_threads - elif settings.model_intra_op_threads == 0 and self.providers == ["CPUExecutionProvider"]: + elif settings.model_intra_op_threads == 0 and self.provider == "CPUExecutionProvider": sess_options.intra_op_num_threads = self.threads if sess_options.inter_op_num_threads > 1: sess_options.execution_mode = ort.ExecutionMode.ORT_PARALLEL + for dimension, size in self.overrides: + sess_options.add_free_dimension_override_by_name(dimension, size) + + log.debug(f"Setting execution_mode to {sess_options.execution_mode.name}") + log.debug(f"Setting inter_op_num_threads to {sess_options.inter_op_num_threads}") + log.debug(f"Setting intra_op_num_threads to {sess_options.intra_op_num_threads}") return sess_options +class Facts(BaseModel): + onnxruntime: str + options: dict[str, Any] + overrides: list[tuple[str, int]] + disabled_optimizers: list[str] + rewrite: str | None = None + device: str | None = None + half: bool = False + + +class Manifest(BaseModel): + facts: Facts + graph: str # relative to the artifact's directory + + +def fresh(spec: GraphSpec) -> Path | None: + try: + manifest = Manifest.model_validate_json(spec.manifest.read_bytes()) + except (OSError, ValueError): + return None + return spec.model_path.parent / manifest.graph if manifest.facts == spec.facts else None + + +def prepared(spec: GraphSpec) -> Path: + if (graph := fresh(spec)) is None: + log.info(f"Preparing {spec.model_path} for {spec.provider} {dict(spec.pins)}") + # child process prepares the graph to avoid memory overhead in the main process + child = subprocess.run([sys.executable, "-m", "immich_ml.sessions.prepare"], input=pickle.dumps(spec)) + if (graph := fresh(spec)) is None: + if child.returncode == UNREADABLE: + raise InvalidProtobuf(f"{spec.model_path} could not be read") + raise RuntimeError(f"Preparing {spec.model_path} failed; the output above says why") + return graph + + +@cache +def _plan(provider: str) -> RewritePlan: + version = tuple(int(piece) for piece in ort.__version__.split(".")[:3]) + return plan_rewrites(RewriteContext(target=provider, ort_version=version)) + + +@cache +def _cpu() -> str: + try: + cpuinfo = Path("/proc/cpuinfo").read_text() + except OSError: + return platform.machine() + return next(line for line in cpuinfo.splitlines() if line.startswith(("flags", "Features"))) + + +class OrtSession: + def __init__( + self, model_path: Path | str, policy: ShapePolicy, providers: list[str] | None = None, threads: int = 2 + ) -> None: + self.model_path = Path(model_path) + self.policy = policy + self.providers = providers if providers is not None else _providers_default() + log.info(f"Setting execution providers to {self.providers}, in descending order of preference") + self.disabled_optimizers = _disabled_optimizers_default(self.providers) + log.debug(f"Setting disabled_optimizers to {self.disabled_optimizers}") + self.threads = threads + self.dynamic = self.providers[0] in DYNAMIC_PROVIDERS + # the shapes to snap onto, where a dim a shape leaves out takes any size + self.shapes = (policy.pinned,) if self.dynamic else policy.dims + self.batches = tuple(sorted({shape.batch for shape in self.shapes}, reverse=True)) + self.graphs: dict[Shape, OrtGraph] = {} + self.locks: dict[Shape, Lock] = {} + if self.dynamic or len(policy.dims) == 1: + self.for_shape(policy.dims[0]) # the one graph there is, so nothing waits on a shape + + def for_shape(self, shape: Shape) -> OrtGraph: + if self.dynamic: # a batch gains nothing here, and a graph that knows it has one row optimizes further + shape = self.shapes[0] + if (graph := self.graphs.get(shape)) is not None: + return graph + if not self.dynamic and shape not in self.policy.dims: # every shape admitted is another graph to compile + raise ValueError(f"Asked for {shape}, which is none of {list(self.policy.dims)}") + with self.locks.setdefault(shape, Lock()): # requests arrive together for a graph none of them has + if (graph := self.graphs.get(shape)) is None: + log.debug(f"Building a graph for {shape}") + spec = GraphSpec( + self.model_path, + shape.pins, + _overrides(self.policy, shape.pins), + self.providers, + self.disabled_optimizers, + self.threads, + ) + graph = self.graphs[shape] = OrtGraph(spec) + return graph + + def warm(self) -> None: + """Opens a graph for every shape and runs it once before opening the next: OpenVINO defers work to the first + run of a graph it imports from its cache, which importing another graph first can corrupt.""" + for shape in self.policy.dims[:1] if self.dynamic else self.policy.dims: + session = self.for_shape(shape).session + sizes = dict(_overrides(self.policy, shape.pins)) # for a dim the graph leaves free + feed = { + node.name: np.zeros( + [dim if isinstance(dim, int) else sizes.get(dim, 1) for dim in node.shape], DTYPES[node.type] + ) + for node in session.get_inputs() + } + session.run(None, feed) + + +def _overrides(policy: ShapePolicy, pins: Mapping[str, int]) -> list[tuple[str, int]]: + order = tuple(policy.dims[0].pins) + named = [] + for dimension, size in pins.items(): + position = order.index(dimension) # a graph that names several dims takes the first pin for all + named += [(dimension, size), (f"DynamicDimension.{position}", size)] + named.append(("None", size) if position == 0 else ("?", size)) + return named + + +@cache +def _intel_gpu(device: str) -> Device: + """The IP version and execution units OpenVINO keys its blobs by, and the driver it checks them against.""" + ov = ctypes.CDLL(str(Path(ort.__file__).parent / "capi" / "libopenvino_c.so")) + core, wanted = ctypes.c_void_p(), ctypes.c_char_p() + ov.ov_core_create(ctypes.byref(core)) + ov.ov_core_get_property(core, device.encode(), b"DEVICE_UUID", ctypes.byref(wanted)) + uuid = wanted.value + ov.ov_free(wanted) + ov.ov_core_free(core) + cl, gpu = ctypes.CDLL("libOpenCL.so.1"), ctypes.c_uint64(1 << 2) # CL_DEVICE_TYPE_GPU + platforms, devices, count = (ctypes.c_void_p * 8)(), (ctypes.c_void_p * 8)(), ctypes.c_uint32() + cl.clGetPlatformIDs(8, platforms, ctypes.byref(count)) + for vendor in platforms[: count.value]: + found = ctypes.c_uint32() + cl.clGetDeviceIDs(ctypes.c_void_p(vendor), gpu, 8, devices, ctypes.byref(found)) + for handle in devices[: found.value]: + own, driver = (ctypes.c_ubyte * 16)(), ctypes.create_string_buffer(64) + ip, units = ctypes.c_uint32(), ctypes.c_uint32() + # CL_DEVICE_UUID_KHR, CL_DRIVER_VERSION, CL_DEVICE_IP_VERSION_INTEL, CL_DEVICE_MAX_COMPUTE_UNITS + for key, value in ((0x106A, own), (0x102D, driver), (0x4250, ip), (0x1002, units)): + cl.clGetDeviceInfo(ctypes.c_void_p(handle), key, ctypes.sizeof(value), ctypes.byref(value), None) + if bytes(own).hex().encode() == uuid: + version = f"{ip.value >> 22}.{ip.value >> 14 & 0xFF}.{ip.value & 0x3FFF}" + return Device(f"{version}-{units.value}eu", driver.value.decode()) + raise LookupError(f"OpenCL has no GPU with the UUID of OpenVINO's {device}") + + +@cache +def _amd_gpu(index: int) -> Device: + """The arch MIGraphX keys its programs by, and the MIGraphX and HIP it loads, which its keys leave out.""" + # the provider finds ROCm through its runpath + ctypes.CDLL(str(Path(ort.__file__).parent / "capi" / "libonnxruntime_providers_migraphx.so")) + with open("/proc/self/maps") as maps: + loaded = {Path(path).name.split(".so")[0]: Path(path) for *_, path in map(str.split, maps) if ".so" in path} + hip, hsa = ctypes.CDLL(str(loaded["libamdhip64"])), ctypes.CDLL(str(loaded["libhsa-runtime64"])) + wanted, arches = ctypes.create_string_buffer(64), {} + hip.hipDeviceGetPCIBusId(wanted, 64, index) # honors HIP_VISIBLE_DEVICES, unlike HSA + + def visit(agent: int, _: int | None) -> int: + name, bdf, domain = ctypes.create_string_buffer(64), ctypes.c_uint32(), ctypes.c_uint32() + # HSA_AGENT_INFO_NAME, HSA_AMD_AGENT_INFO_BDFID, HSA_AMD_AGENT_INFO_DOMAIN + for key, value in ((0, name), (0xA006, bdf), (0xA00F, domain)): + hsa.hsa_agent_get_info(ctypes.c_uint64(agent), key, ctypes.byref(value)) + bus, device, function = bdf.value >> 8 & 0xFF, bdf.value >> 3 & 0x1F, bdf.value & 0x7 + arches[f"{domain.value:04x}:{bus:02x}:{device:02x}.{function:01x}"] = name.value.decode() + return 0 + + hsa.hsa_init() + hsa.hsa_iterate_agents(ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_uint64, ctypes.c_void_p)(visit), None) + return Device(arches[wanted.value.decode()], f"{loaded['libmigraphx'].name} {loaded['libamdhip64'].name}") + + +def flush_denormals() -> None: + """Reads subnormal floats as zero on the calling thread, which ORT does only on its own threads.""" + if sys.platform != "linux" or platform.machine() != "x86_64": + return + libm = ctypes.CDLL("libm.so.6") + env = (ctypes.c_uint32 * 8)() # glibc's fenv_t here: the x87 environment, then MXCSR + libm.fegetenv(env) + env[7] |= 0x8040 # denormals are zero, flush to zero + libm.fesetenv(env) + + +def _providers_default() -> list[str]: + from immich_ml.models.constants import SUPPORTED_PROVIDERS # here, as the preparing child has no use for the models + + available_providers = set(ort.get_available_providers()) + log.debug(f"Available ORT providers: {available_providers}") + return [provider for provider in SUPPORTED_PROVIDERS if provider in available_providers] + + def _disabled_optimizers_default(providers: list[str]) -> list[str]: disabled_optimizers: list[str] = [] if platform.machine() in ("arm64", "aarch64"): # as macOS and Linux name the same architecture @@ -226,14 +409,3 @@ def _disabled_optimizers_default(providers: list[str]) -> list[str]: disabled_optimizers.append("MatMulAddFusion") return disabled_optimizers - - -def flush_denormals() -> None: - """Reads subnormal floats as zero on the calling thread, which ORT does only on its own threads.""" - if sys.platform != "linux" or platform.machine() != "x86_64": - return - libm = ctypes.CDLL("libm.so.6") - env = (ctypes.c_uint32 * 8)() # glibc's fenv_t here: the x87 environment, then MXCSR - libm.fegetenv(env) - env[7] |= 0x8040 # denormals are zero, flush to zero - libm.fesetenv(env) diff --git a/machine-learning/immich_ml/sessions/policy.py b/machine-learning/immich_ml/sessions/policy.py new file mode 100644 index 000000000..8b6ea35c4 --- /dev/null +++ b/machine-learning/immich_ml/sessions/policy.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +from immich_ml.schemas import Shape + + +@dataclass(frozen=True) +class ShapePolicy: + dims: tuple[Shape, ...] = (Shape(batch=1),) + label: str = "" # the artifact subdirectory, for a submodel that ships more than one set + + @property + def pinned(self) -> Shape: + first = self.dims[0] + return Shape( + batch=1, + height=first.height if all(shape.height == first.height for shape in self.dims) else None, + width=first.width if all(shape.width == first.width for shape in self.dims) else None, + ) + + +def batches(maximum: int) -> tuple[int, ...]: + return (1, maximum) if maximum > 1 else (1,) + + +def runs(total: int, batches: Sequence[int]) -> list[int]: + runs = [] + while total > 0: + runs.append(max(size for size in batches if size <= total)) + total -= runs[-1] + return runs diff --git a/machine-learning/immich_ml/sessions/prepare.py b/machine-learning/immich_ml/sessions/prepare.py new file mode 100644 index 000000000..d2f974d98 --- /dev/null +++ b/machine-learning/immich_ml/sessions/prepare.py @@ -0,0 +1,107 @@ +import fcntl +import pickle +import sys +from pathlib import Path +from shutil import rmtree +from typing import Any + +import onnx +from google.protobuf.message import DecodeError +from immich_model.onnx.f16 import derive +from immich_model.runtime import apply_rewrites +from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf, NoSuchFile + +from ..config import settings +from .ort import UNREADABLE, GraphSpec, Manifest, fresh + + +def main() -> None: + spec: GraphSpec = pickle.loads(sys.stdin.buffer.read()) + with spec.model_path.with_suffix(".lock").open("w") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) # several workers may want the same graph + if fresh(spec) is None: + try: + prepare(spec) + except (InvalidProtobuf, NoSuchFile, DecodeError): + sys.exit(UNREADABLE) + + +def prepare(spec: GraphSpec) -> None: + rmtree(spec.directory, ignore_errors=True) # whatever is here was made under other facts, or not finished + spec.directory.mkdir(parents=True) + graph = rewritten(spec, narrowed(spec)) + if spec.provider == "CPUExecutionProvider": + graph = prepacked(spec, graph) # allows mmap'ing the file instead of prepacking to dirty heap memory + elif spec.directory.as_posix() in spec.provider_options[0].values(): + spec.session(graph) # the provider keeps what it compiles, in the directory it was given + manifest = Manifest(facts=spec.facts, graph=graph.relative_to(spec.model_path.parent).as_posix()) + spec.manifest.write_text(manifest.model_dump_json()) # last, so that it means complete + + +def narrowed(spec: GraphSpec) -> Path: + if not spec.half: + return spec.model_path + destination = spec.directory / "model_fp16.onnx" + derive(spec.model_path, destination) + return destination + + +def rewritten(spec: GraphSpec, graph: Path) -> Path: + written: Path = graph if settings.legacy_models else apply_rewrites(graph, spec.plan) + return written + + +def prepacked(spec: GraphSpec, graph: Path) -> Path: + destination = spec.directory / "model.onnx" + options = spec.sess_options() + options.optimized_model_filepath = destination.as_posix() + options.add_session_config_entry("session.optimized_model_external_initializers_file_name", "model.data") + options.add_session_config_entry("session.save_external_prepacked_constant_initializers", "1") + options.add_session_config_entry("session.optimized_model_external_initializers_min_size_in_bytes", "1024") + spec.session(graph, options) + align(destination) + return destination + + +def align(destination: Path) -> None: + """Re-lays the data file so every region starts aligned. ORT pads only large regions, so one odd length skews + all that follow, and MLAS reads a packed weight with aligned loads: it faults, and reads the rest slower.""" + model = onnx.load(destination.as_posix(), load_external_data=False) + length_at: dict[int, int] = {} + offsets: list[Any] = [] + packs: list[tuple[Any, str, list[list[str]]]] = [] + for tensor in model.graph.initializer: + fields = {entry.key: entry for entry in tensor.external_data} + for key, entry in fields.items(): + if key == "offset": + length_at[int(entry.value)] = int(fields["length"].value) + offsets.append(entry) + elif key.startswith("prepacked_"): + blob_key, *rest = entry.value.split("|") + triples = [item.split(";") for item in rest] + length_at.update((int(offset), int(length)) for offset, length, _ in triples) + packs.append((entry, blob_key, triples)) + + moved, cursor = {}, 0 + for offset in sorted(length_at): + cursor += -cursor % (4096 if length_at[offset] > 1 << 20 else 64) + moved[offset] = cursor + cursor += length_at[offset] + + data = destination.with_suffix(".data") + with data.open("rb") as source, data.with_suffix(".aligned").open("wb") as target: + for offset in sorted(length_at): + source.seek(offset) + target.seek(moved[offset]) # what it skips reads as zeros + target.write(source.read(length_at[offset])) + data.with_suffix(".aligned").replace(data) + + for entry in offsets: + entry.value = str(moved[int(entry.value)]) + for entry, blob_key, triples in packs: + entry.value = "|".join([blob_key, *(f"{moved[int(o)]};{n};{c}" for o, n, c in triples)]) + onnx.save(model, destination.as_posix()) + + +if __name__ == "__main__": + main() diff --git a/machine-learning/immich_ml/sessions/rknn/__init__.py b/machine-learning/immich_ml/sessions/rknn/__init__.py index 0d95e6cf9..ee0ef901a 100644 --- a/machine-learning/immich_ml/sessions/rknn/__init__.py +++ b/machine-learning/immich_ml/sessions/rknn/__init__.py @@ -1,78 +1,84 @@ from __future__ import annotations +from collections.abc import Sequence from pathlib import Path -from typing import Any, NamedTuple +from typing import Any import numpy as np +import orjson from numpy.typing import NDArray from immich_ml.config import log, settings -from immich_ml.schemas import SessionNode +from immich_ml.schemas import ModelInput, ModelTensor, SessionNode, Shape -from .rknnpool import RknnPoolExecutor, is_available, soc_name +from .rknnpool import RknnNode, RknnPoolExecutor, is_available, native_outputs, soc_name is_available = is_available and settings.rknn model_prefix = Path("rknpu") / soc_name if is_available and soc_name is not None else None -def run_inference(rknn_lite: Any, input: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]: - outputs: list[NDArray[np.float32]] = rknn_lite.inference(inputs=input, data_format="nchw") - return outputs +def model_path(model_dir: Path, variant: str = "") -> Path: + return (model_dir / model_prefix if model_prefix else model_dir) / variant / "model.rknn" -input_output_mapping: dict[str, dict[str, Any]] = { - "detection": { - "input": {"norm_tensor:0": (1, 3, 640, 640)}, - "output": { - "norm_tensor:1": (12800, 1), - "norm_tensor:2": (3200, 1), - "norm_tensor:3": (800, 1), - "norm_tensor:4": (12800, 4), - "norm_tensor:5": (3200, 4), - "norm_tensor:6": (800, 4), - "norm_tensor:7": (12800, 10), - "norm_tensor:8": (3200, 10), - "norm_tensor:9": (800, 10), - }, - }, - "recognition": {"input": {"norm_tensor:0": (1, 3, 112, 112)}, "output": {"norm_tensor:1": (1, 512)}}, -} +def run_inference(rknn_lite: Any, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[np.float32]]: + rknn_lite.rknn_runtime.set_inputs(inputs, None, data_format) + rknn_lite.rknn_runtime.run(False) + return native_outputs(rknn_lite) + + +def input_layout(compiled: tuple[int, ...], array: ModelTensor) -> str | None: + if array.ndim != 4: + return None # token ids carry no layout + if array.shape[3] == compiled[3]: + return "nhwc" + if array.shape[1] == compiled[3]: + return "nchw" + raise ValueError(f"binary takes {compiled[3]} channels, fed {list(array.shape)}") class RknnSession: def __init__(self, model_path: Path) -> None: - self.model_type = "detection" if "detection" in model_path.parts else "recognition" self.tpe = settings.rknn_threads log.info(f"Loading RKNN model from {model_path} with {self.tpe} threads.") self.rknnpool = RknnPoolExecutor(model_path=model_path.as_posix(), tpes=self.tpe, func=run_inference) log.info(f"Loaded RKNN model from {model_path} with {self.tpe} threads.") + # the shapes the binary was compiled for, among which the runtime routes itself + batch = self.rknnpool.inputs[0].shape[0] + declared = orjson.loads(self.rknnpool.custom_string) if self.rknnpool.custom_string else {} + self.shapes = tuple(Shape(batch, **dims) for dims in declared.get("dims", [{}])) + self.metadata = {key: value for key, value in declared.items() if key != "dims"} + self.batches = (batch,) - def get_inputs(self) -> list[SessionNode]: - return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["input"].items()] + def for_shape(self, shape: Shape) -> RknnSession: + return self - def get_outputs(self) -> list[SessionNode]: - return [RknnNode(name=k, shape=v) for k, v in input_output_mapping[self.model_type]["output"].items()] + def warm(self) -> None: + pass # the binary is loaded whole when the session opens + + def get_inputs(self) -> Sequence[SessionNode]: + return self.rknnpool.inputs + + def get_outputs(self) -> Sequence[SessionNode]: + return self.rknnpool.outputs def get_metadata(self) -> dict[str, str]: - return {} + return self.metadata + + @property + def normalizes_input(self) -> bool: + # the compiler transposes and casts every binary, so only what the exporter declared can say + return bool(self.rknnpool.custom_string) def run( self, output_names: list[str] | None, - input_feed: dict[str, NDArray[np.float32]] | dict[str, NDArray[np.int32]] | dict[str, NDArray[np.uint8]], + input_feed: ModelInput, run_options: Any = None, ) -> list[NDArray[np.float32]]: - input_data: list[NDArray[np.float32]] = [np.ascontiguousarray(v) for v in input_feed.values()] - res = self.rknnpool.run(input_data) - if res is None: - raise RuntimeError("RKNN inference failed!") - return res + inputs = [array if array.flags.c_contiguous else array.copy() for array in input_feed.values()] + return self.rknnpool.run(inputs, input_layout(self.rknnpool.inputs[0].shape, inputs[0])) -class RknnNode(NamedTuple): - name: str - shape: tuple[int, ...] - - -__all__ = ["RknnSession", "RknnNode", "is_available", "soc_name", "model_prefix"] +__all__ = ["RknnSession", "RknnNode", "is_available", "soc_name", "model_path"] diff --git a/machine-learning/immich_ml/sessions/rknn/rknnpool.py b/machine-learning/immich_ml/sessions/rknn/rknnpool.py index 1caec1abb..cf4844334 100644 --- a/machine-learning/immich_ml/sessions/rknn/rknnpool.py +++ b/machine-learning/immich_ml/sessions/rknn/rknnpool.py @@ -1,17 +1,19 @@ # This code is from leafqycc/rknn-multi-threaded # Following Apache License 2.0 +import ctypes import logging import threading from concurrent.futures import ThreadPoolExecutor from pathlib import Path -from typing import Callable +from typing import Any, Callable, NamedTuple import numpy as np from numpy.typing import NDArray from immich_ml.config import log from immich_ml.models.constants import RKNN_COREMASK_SUPPORTED_SOCS, RKNN_SUPPORTED_SOCS +from immich_ml.schemas import ModelTensor def get_soc(device_tree_path: Path | str) -> str | None: @@ -30,7 +32,7 @@ def get_soc(device_tree_path: Path | str) -> str | None: soc_name = None is_available = False try: - from rknnlite.api import RKNNLite + from rknnlite.api import RKNNLite, rknn_runtime soc_name = get_soc("/proc/device-tree/compatible") is_available = soc_name is not None @@ -38,6 +40,82 @@ except ImportError: log.debug("RKNN is not available") +class RknnNode(NamedTuple): + name: str + shape: tuple[int, ...] + + +_CUSTOM_STRING_QUERY = 7 # RKNN_QUERY_CUSTOM_STRING +# RKNN_QUERY_OUTPUT_ATTR, and RKNN_QUERY_CURRENT_OUTPUT_ATTR for the shape a dynamic binary last produced +_OUTPUT_QUERY, _CURRENT_OUTPUT_QUERY = 2, 15 +_DTYPES = {0: np.float32, 1: np.float16} # rknn_tensor_type + + +class _CustomString(ctypes.Structure): + _fields_ = [("string", ctypes.c_char * 1024)] + + +class _TensorAttr(ctypes.Structure): # rknn_tensor_attr + _fields_ = [ + ("index", ctypes.c_uint32), + ("n_dims", ctypes.c_uint32), + ("dims", ctypes.c_uint32 * 16), + ("name", ctypes.c_char * 256), + ("n_elems", ctypes.c_uint32), + ("size", ctypes.c_uint32), + ("fmt", ctypes.c_int), + ("type", ctypes.c_int), + ("qnt_type", ctypes.c_int), + ("fl", ctypes.c_int8), + ("zp", ctypes.c_int32), + ("scale", ctypes.c_float), + ("w_stride", ctypes.c_uint32), + ("size_with_stride", ctypes.c_uint32), + ("pass_through", ctypes.c_uint8), + ("h_stride", ctypes.c_uint32), + ] + + +def native_outputs(rknn_lite: "RKNNLite") -> list[NDArray[Any]]: + """The outputs in the precision the NPU computed them in, which RKNNLite would widen to fp32 on the CPU.""" + runtime: Any = rknn_lite.rknn_runtime + count = runtime.get_in_out_num()[1] + buffers = (rknn_runtime.RKNNOutput * count)() + for index, buffer in enumerate(buffers): + buffer.index = index + if runtime.lib.rknn_outputs_get(runtime.context, count, buffers, None) != 0: + raise RuntimeError("RKNN inference failed!") + query = _CURRENT_OUTPUT_QUERY if runtime.is_dynamic_shape else _OUTPUT_QUERY + outputs = [] + for buffer in buffers: + attr = _TensorAttr(index=buffer.index) + runtime.lib.rknn_query(runtime.context, query, ctypes.byref(attr), ctypes.sizeof(attr)) + data = ctypes.cast(buffer.buf, ctypes.POINTER(ctypes.c_char * buffer.size)).contents + outputs.append(np.frombuffer(data, _DTYPES[attr.type]).reshape(attr.dims[: attr.n_dims]).copy()) + runtime.lib.rknn_outputs_release(runtime.context, count, buffers) + return outputs + + +def tensor_nodes(rknn_lite: "RKNNLite") -> tuple[list[RknnNode], list[RknnNode]]: + runtime = rknn_lite.rknn_runtime + n_inputs, n_outputs = runtime.get_in_out_num() + + def node(index: int, is_output: bool) -> RknnNode: + attr = runtime.get_tensor_attr(index, is_output=is_output) + return RknnNode(attr.name.decode(), tuple(attr.dims[: attr.n_dims])) + + return [node(i, False) for i in range(n_inputs)], [node(i, True) for i in range(n_outputs)] + + +def custom_string(rknn_lite: "RKNNLite") -> str: + runtime: Any = rknn_lite.rknn_runtime + buffer = _CustomString() + # rknnlite declares argtypes on the shared handle, so pass plain values and let ctypes convert + runtime.lib.rknn_query(runtime.context, _CUSTOM_STRING_QUERY, ctypes.byref(buffer), ctypes.sizeof(buffer)) + declared: str = buffer.string.decode() + return declared + + def init_rknn(model_path: str) -> "RKNNLite": if not is_available: raise RuntimeError("rknn is not available!") @@ -63,21 +141,23 @@ class RknnPoolExecutor: self, model_path: str, tpes: int, - func: Callable[["RKNNLite", list[NDArray[np.float32]]], list[NDArray[np.float32]]], + func: Callable[["RKNNLite", list[ModelTensor], str | None], list[NDArray[np.float32]]], ) -> None: self.tpes = tpes self.rknn_pool = [init_rknn(model_path) for _ in range(tpes)] + self.inputs, self.outputs = tensor_nodes(self.rknn_pool[0]) + self.custom_string = custom_string(self.rknn_pool[0]) self.pool = ThreadPoolExecutor(max_workers=tpes) self.func = func self.num = 0 self.lock = threading.Lock() - def run(self, inputs: list[NDArray[np.float32]]) -> list[NDArray[np.float32]]: + def run(self, inputs: list[ModelTensor], data_format: str | None) -> list[NDArray[np.float32]]: with self.lock: idx = self.num % self.tpes self.num += 1 - fut = self.pool.submit(self.func, self.rknn_pool[idx], inputs) + fut = self.pool.submit(self.func, self.rknn_pool[idx], inputs, data_format) return fut.result() def release(self) -> None: diff --git a/machine-learning/pyproject.toml b/machine-learning/pyproject.toml index 84b1d2f4e..3271fbf96 100644 --- a/machine-learning/pyproject.toml +++ b/machine-learning/pyproject.toml @@ -9,6 +9,7 @@ dependencies = [ "fastapi>=0.95.2,<1.0", "gunicorn>=21.1.0", "huggingface-hub>=1.0,<2.0", + "immich-model", "numpy>=2.4.0,<3.0", "onnx>=1.22.0", "opencv-python-headless>=4.7.0.72,<5.0", @@ -58,6 +59,7 @@ compile-bytecode = true override-dependencies = ["opencv-python ; sys_platform == 'never'"] # only install the headless version [tool.uv.sources] +immich-model = { git = "https://github.com/immich-app/ml-models.git", rev = "88d2b64ff95ce7117760660681d95e0e1308819b" } onnxruntime-gpu = { index = "onnxruntime-cuda-12" } [[tool.uv.index]] diff --git a/machine-learning/responses.json b/machine-learning/responses.json index 8fe9cd61c..d46aa84ae 100644 --- a/machine-learning/responses.json +++ b/machine-learning/responses.json @@ -1,329 +1 @@ -{ - "clip": { - "image": [ - -0.013126737, -0.022995953, -0.0493738, -0.0063057775, 0.013601424, - -0.003761688, -0.03379882, 0.11106285, 0.024760082, 0.023903701, - 0.04418207, -0.013594999, 0.030850016, 0.0012876489, -0.012471005, - 0.009750715, 0.0095717255, 0.013320666, 0.0027921356, 0.03240264, - 0.033538498, 0.013624318, -0.0069946186, -0.0036184592, -0.009846507, - -0.017311024, -0.036686428, -0.0041808123, 0.030871637, -0.028624479, - -0.016515259, 0.014418001, -0.024542322, -0.0025438748, -0.049111884, - -0.023928944, 0.012270045, -0.016418075, 0.004895335, -0.15801854, - -0.026325515, 0.03166467, -0.017224329, 0.0411128, -0.022944424, - 0.015693054, -0.020919899, -0.010764121, -0.008499815, -0.020263294, - -0.009743323, -0.035395622, 0.03474742, 0.003049183, 0.009424564, - 0.010707678, 0.01664117, -0.0059027374, -0.013055344, 0.0018035833, - -0.003976456, -0.04325922, 0.014407317, 0.035728276, 0.002226939, - -0.006235411, -0.0073032235, -0.035129357, 0.001095443, -0.028552389, - -0.044300288, -0.012959393, 0.02259977, 0.017141517, -0.029432472, - -0.017583484, 0.010974336, 0.018141218, 0.0015389329, -0.008220305, - -0.0060965014, 0.004929384, 0.019477025, -0.033071984, 0.025183259, - 0.013607688, 0.01836233, 0.04586782, 0.0103442725, -0.036077496, - -0.029715508, 0.007203621, -0.7949153, 0.046866275, 0.026910711, - -0.0047834567, 0.033243995, 0.009379981, -0.03749048, -0.055274535, - -0.01955359, 0.012887587, 0.00922838, -0.0032776103, -0.011456734, - 0.0045412215, -0.11506394, 0.0348558, 0.029478835, -0.011811103, - -0.00483158, -0.010586126, -0.018853206, -0.01591496, -0.019360982, - -0.03211199, -0.013473663, -0.019630248, -0.017012835, 0.059128772, - -0.03396129, 0.0045991736, -0.015158291, 0.008241974, 0.004403056, - -0.007536049, -0.023821214, -0.0059521003, 0.015564905, 0.020600233, - 0.008175, 0.02100119, -0.0034459834, 0.1058016, 0.008383205, 0.03100292, - -0.023814196, -0.016157096, -0.008210107, -0.004146204, 0.016350364, - -0.056028433, 0.013261071, 0.034839876, -0.03236049, 0.026573967, - -0.018140865, 0.018515658, 0.013251766, 0.007693613, 0.0067239976, - -0.0013857568, 0.038114607, 0.0068016117, 0.036603037, 0.0040935865, - 0.010394745, -0.00082285365, -0.009811308, 0.020343611, -0.012164189, - -0.012208623, 0.0005465415, -0.015394064, 0.02499845, 0.021941017, - -0.016571017, -0.011810332, 0.017864, -0.010639794, -0.008609091, - -0.0007239709, 0.015229945, -0.0035874692, 0.018922612, -0.011209458, - -0.013052865, -0.009626533, -0.004419959, 0.007915186, 0.01094836, - 0.005509159, -0.0034862005, 0.01012292, -0.0059307595, -0.029599648, - 0.032845, -0.007011692, -0.014218981, 0.00790071, 0.017027052, - -0.022314077, -0.03041719, 0.015665755, 0.036747217, -0.018942915, - 0.008623111, 0.02179961, -0.022312569, 0.007024427, 0.016751591, - -0.0034192575, 0.024101255, -0.0046198783, 0.022274924, -0.015562676, - -0.0092551885, -0.0063787713, 0.045996074, 0.026235346, 0.009622556, - 0.05728027, 0.03168525, -0.017600676, 0.029278612, 0.01467962, - 0.032169178, 0.022459356, -0.012175933, -0.009438608, 0.027234165, - 0.013514767, -0.008831029, 0.010888894, 0.004518216, 0.009855367, - 0.012112431, -0.0073178695, 0.0072642234, 4.877679e-5, -0.01221576, - 0.023542404, -0.009026452, -0.055442516, 0.006579068, 0.033202186, - -0.007669379, 0.0010604112, -0.04271919, -0.029112164, 0.021844024, - 0.029739635, -0.026083348, 0.008940292, -0.039301652, -0.047215454, - 0.0018794702, -0.008740231, 0.029195482, 0.0037629968, -0.024923965, - -0.021407248, 0.009952853, -0.0055059143, 0.0044912454, 0.016966008, - -0.00081178773, -0.022250004, -0.014063889, -0.006170697, -0.0008208651, - -0.036218595, -0.0029040456, 0.03943083, -0.021814227, 0.017567957, - 0.035849728, -0.049075987, 0.0040634805, 0.009878297, 0.028557112, - 0.02336673, 0.010714448, 0.020129073, -0.030503238, 0.009350441, - 0.039086528, -0.0037483997, -0.0034365985, 0.019824414, 0.014027232, - 0.030565958, 0.0036307913, 0.0030920429, -0.009908996, 0.0027933475, - -7.140754e-5, -0.027733125, 0.0022445584, -0.032248124, 0.050226185, - 0.030529078, -0.040353864, 0.031086015, -0.0063569676, 0.031343475, - -0.020244656, -0.011442288, 0.018035123, -0.005479394, 0.01783419, - -0.036066547, -0.0106600635, 0.044636995, -0.030209303, -0.07192714, - 0.0128155155, 0.003505818, -0.0005725083, -0.01584388, -0.025725754, - 0.025868604, 0.10576061, -0.012738124, 0.0012224225, 0.0472961, - 0.021650923, 0.0061313445, 0.014010678, 0.016864019, 0.004049639, - 0.10989465, 0.011927816, 0.013589654, 0.011258818, 0.022496557, - -0.018828733, 0.021635532, 0.0116777215, 0.11320542, -0.0011280471, - 0.018990291, 0.001824643, -0.03793715, 0.0206918, -0.0050228164, - -0.013865701, 0.022277884, 0.019400347, 0.028610364, -0.023974052, - 0.0030309444, 0.027177742, 0.024541136, 0.023737634, 0.0012539584, - 0.0187086, -0.015451178, 0.015066189, -0.019812824, 0.050507285, - -0.0021846944, 0.041644007, -0.0070109894, -0.014599777, -0.05985813, - -0.036328156, -0.02293525, -0.0065515027, 0.016792618, -0.0059018973, - -0.008319917, 0.008072106, 0.0073447954, -0.052924518, -0.037344936, - 0.015524772, -0.0012835241, 0.014405327, 0.0057144985, 0.004945561, - -0.024654018, 0.011967616, 0.01832056, 0.019411784, 0.019788045, - -0.0006405928, -0.0015148119, -0.05064218, -0.031875107, -0.03803604, - -0.0096240705, 0.012371131, -0.019090319, 0.0075365147, -0.024229601, - 0.014469528, -0.004786435, 0.0011314518, 0.009256282, -0.04957284, - -0.0068631344, -0.010091242, -0.023295002, 0.03268865, 0.022269772, - 0.037733294, -0.015292435, -0.06330943, -0.00854154, 0.0027765913, - 0.0015374947, 0.0377278, 0.008772586, -0.01810512, -0.0025668603, - 0.014428339, 0.0027927365, 0.07493676, -0.022829408, -0.028912589, - 0.008928177, 0.011323267, 0.008405796, 0.016925976, 0.001739356, - -0.021090876, -0.0062678503, 0.010898773, -0.010470923, 0.015523946, - -0.027888289, 0.023872118, -0.048326068, 0.025968319, 0.0047795917, - -0.016123952, 0.00698612, -0.05154045, -0.003691712, -0.0101406425, - -0.0241034, 0.004006022, 0.0021649078, 0.0019942294, -0.009274028, - -0.006467623, -0.0010948133, -0.012350769, -0.0060371486, -0.0006392645, - 0.031422533, 0.015165475, -0.012650007, -0.005918423, 0.005781174, - -0.023262534, -0.0043034274, -0.010881872, -0.015937665, -0.0043740096, - -0.02981798, -0.0037422387, -0.029688178, 0.022320364, -0.0014900378, - -0.026122924, -0.04360404, 0.016354023, -0.02447563, 0.0205314, - 0.0042775236, -0.020184014, -0.0017819501, 0.036122557, 0.0036566693, - 0.07459051, -0.0035548757, 0.004874807, -0.028627345, -0.023153499, - 0.03710664, -0.000639956, -0.030509725, -0.005146651, -0.010251552, - 0.028408762, -0.008056198, -0.018420909, 0.02850364, -0.0075958185, - -0.008918139, 0.002778187, 0.06220242, -0.016280292, -0.026200369, - 0.05900717, -0.013802131, 0.005442568, -0.033114687, 0.010976371, - 0.008192846, 0.0031891295, 0.024811232, 0.009066575, -0.026441244, - 0.030676885, -0.014591597, -0.024314625, -0.037472498, -0.015021544, - -0.016501956, -0.0069196, 0.013831272, 0.056646723, 0.007946148, - -0.002477574, -0.030496774, -0.011770325, 0.06742689, -0.03180974, - -0.025615396 - ], - "text": [ - -0.0040579583, -0.00084722764, -0.008696951, -0.006850008, -0.010870523, - 0.014495447, -0.010678498, -0.09618138, 0.016697474, -0.014809047, - -0.0035991871, 0.020752821, 0.0020757387, 0.0018064519, 0.02969283, - -0.0040159826, 0.02335311, 0.015918557, -0.0015919582, 0.013545261, - 0.011341818, -0.006991808, 0.0020565446, -0.016662853, -0.0064206184, - 0.011536576, -0.01144098, 0.015054818, -0.013258694, 0.0046747606, - -0.00681864, -0.012852865, 0.012708946, 0.006093663, 0.0029938417, - 0.015458671, 0.0040865405, -0.0004354532, 0.0037405093, -0.015085074, - 0.0007998808, -0.021485215, -0.0066235093, 0.015721628, -0.002462181, - -0.0049815965, -0.011028703, 0.0041498104, -0.00070322485, 0.0031991813, - -0.0075132507, -0.008273014, 0.0125206, 0.019671565, 0.02124969, - 0.0076838327, 0.0015366874, 0.0004413452, -0.0027475145, 0.031049952, - 0.01782742, -0.01759819, 0.0040917504, -0.011803108, 0.0051114787, - -0.0075210207, -0.0062834355, 0.010283767, 0.02023746, -0.020258851, - 0.004795256, -0.011313993, -0.014636256, 0.004900588, -0.0026439666, - -0.0004062344, -0.040196564, -0.014539185, 0.014600707, 0.004162044, - -0.0155277, -0.016300475, 0.0039491425, 0.022403471, 0.0055926195, - 0.03543051, 0.01047029, -0.03262415, 0.02942466, -0.011440199, - -0.012713757, 0.0062652803, -0.15837249, 0.028388312, 7.452041e-5, - 0.007136255, -0.010753105, 0.016393846, 0.004432782, -0.010688704, - 0.015814407, 0.01654759, 0.0008900756, -0.007162077, -6.0264443e-5, - 6.894444e-5, -0.034889467, -0.026710762, 0.005202752, -0.012675916, - 0.010903986, 0.017916093, 0.005404525, -0.003624909, -0.012261727, - -0.021104869, -0.01593513, -0.009227664, -0.022844192, 0.008606035, - 0.0013098373, 0.00637147, 0.0027185818, -0.0032128745, 0.018255865, - -0.002964337, -0.006976183, -0.0063263937, 0.0075582284, 0.014295236, - 0.00485664, -0.029005948, -0.014672015, 0.61821824, -0.010903263, - 0.022125997, -0.018154604, -0.007414954, -0.012344926, 0.0029751004, - -0.010787629, -0.0056861844, -0.025484746, -0.0004887071, -0.036681578, - 0.010114145, -0.009012449, -0.00048479583, -0.011162415, -0.0057421126, - -0.0019520421, -0.0013580753, 0.0037870558, 0.0012404326, -0.00089634134, - -0.022112457, 0.0034537334, 0.014985147, -0.010455136, 0.018100852, - -0.010999219, -0.027524924, -0.009551776, -0.0047603208, -0.001092369, - 0.008849578, 0.021949856, -0.034437556, -0.0051499153, -0.0006772509, - -0.011200381, -0.009206776, 0.021897016, -0.0013931778, 0.0041396013, - -0.025534542, 0.0074160174, 0.00039215147, 0.025992293, 0.0069832364, - -0.0175616, 0.01272807, -0.020147255, 0.02455081, -0.01236127, - -0.011840565, 0.011820177, 0.018173985, 0.017230362, -0.016969377, - 0.0010091222, -0.04185319, -0.030467693, -0.012564729, 0.030740628, - -0.004086395, 0.0013223978, -0.0013041743, -0.01975558, -0.014959637, - 0.033446018, -0.014724724, -0.028613493, 0.010436393, -0.009841343, - -0.013723956, -0.0010025625, 0.016992576, 0.0056477, -0.026704773, - 0.018927934, 0.013758461, 0.016924908, -0.026889605, -0.01496036, - 0.02078507, -0.0149594685, -0.021289647, 0.027369255, -0.00205557, - 0.0129268635, -0.014446633, 0.0039108247, 0.014774828, 0.004396043, - 0.0038431762, 0.012223014, 0.0061016707, -0.006525442, -0.018426975, - 0.03081795, 0.024269402, -0.020132616, -0.008118887, 0.025062446, - -0.0033954307, -0.019662865, -0.0032548332, -0.008575233, -0.003158561, - 0.0012930515, 0.02213235, 0.017646195, -0.016638828, -0.0154889, - 0.031743307, 0.001081875, -0.0019133464, 0.034760594, -0.008569126, - -0.019119555, 0.020908207, -0.0047135833, 0.00984879, -0.016712308, - 0.028532412, 0.0038664932, -0.0071539935, 0.0013488994, 0.0060503725, - 0.0021401793, -0.032594826, 0.010918716, 0.0075080344, 0.00020341178, - -0.030393362, 0.014375046, 0.018798219, -0.0040685013, 0.020957684, - -0.012454064, 0.014856742, 0.0017268835, -0.008762698, 0.007062434, - 0.024501909, 0.0011791736, -0.023002177, -0.012701125, -0.0053904364, - 0.015551624, 0.018748082, 0.00704452, 0.0047835982, 0.0013530678, - 0.0033350172, 0.0056562345, 0.009079597, 0.0059383595, 0.011405316, - -0.004795079, 0.007274586, -0.0011659514, -0.0001364172, -0.0050517535, - -0.010681983, -0.023743946, -0.020241234, 0.0009631201, 0.014073974, - -0.00665422, 0.011845411, -0.0001289105, -0.024006248, -0.0009306585, - 0.0139923245, 0.020409467, -0.017118154, -0.0151973255, -0.016737074, - -0.002157259, -0.0060298163, 0.61768156, 0.01275426, -0.023746304, - -0.010786622, -0.00050265377, -0.010761652, -0.013012264, 0.013700237, - 0.019240098, 0.049388826, -0.007853694, -0.014966961, -0.026477624, - 0.02809707, 0.009291939, -0.028884491, 0.015102742, -0.27225503, - 0.01782282, -0.016989257, -0.0051341387, 0.0037056766, -0.004537146, - -0.026184445, 0.020256622, 0.0073146136, -0.0070027146, -0.009025792, - -0.015298917, -0.0052380697, -0.0005596046, -0.0041900063, 0.015934054, - 0.008158574, -0.0038807616, -0.0048019756, 0.0061978237, 0.022159556, - -0.02619826, 0.0013973896, -0.00012341494, 0.030957809, -0.009596324, - 0.008263321, 0.0017040323, -0.0010236687, 0.017982712, -0.012567677, - 0.007361281, 0.0028631007, 0.032613713, 0.035072606, -0.045417674, - 0.016303446, -0.009096281, 0.012163677, -0.008316459, 0.006423764, - -0.008586175, -0.0009862242, 0.009973197, 0.020825483, -0.005682246, - 0.0066081304, 0.0061441967, -0.00670868, -0.024878936, -0.024288971, - -0.009822955, -0.011659227, -0.0067634145, -0.0011930552, 0.017096667, - 0.01974797, -0.020388834, -0.008245143, -0.0071634515, 0.0012492571, - -0.010288493, -0.0025248309, -0.0039965925, 0.037344053, -0.019459987, - 0.022098366, -0.021084892, -0.014823354, -0.010007409, -0.005560381, - -0.012292843, 0.0132691385, 0.0066421456, 0.0045196814, 0.0044144704, - 0.0062646614, 0.0050272197, 0.020296281, 0.011412983, -0.0040745772, - 0.00542041, 0.0021500897, 0.005183101, 0.00985178, 0.014477596, - -0.0131016085, -0.0064126155, -0.004809687, -0.016441243, 0.010445765, - 0.0013761928, -0.0135576585, 0.0003352349, -0.010797083, -0.0058007324, - 0.021649584, 0.012650062, -0.009740497, -0.025809184, -0.026720846, - 0.029149767, 0.014593344, -0.0134959705, -0.004710099, -0.0062580137, - -0.0047687683, -0.029818097, -0.004622532, 0.02532894, 0.0051457905, - -0.0046252706, -0.02905562, -0.019097809, -0.035888474, -0.006897086, - -0.0035953831, -0.0013759647, 0.0027531807, -0.002395984, -0.040570017, - -0.02462688, 0.009387292, 0.0025142033, 0.02404064, -0.0014443685, - 1.0727288e-5, -0.024033979, 0.011659959, -0.016820917, 0.018782362, - -0.019061793, 0.0043488434, 0.00040266776, -0.0022744886, 0.0024185092, - -0.0041366024, 0.0028075825, -0.0085624885, -0.012087987, 0.013551666, - 0.0019014167, 0.007896904, 0.031102024, 0.0091334, -0.0030707342, - 0.0066130627, -0.002711352, -0.012097188, 0.017067473, 0.021030908, - 0.0014250687, -0.092848144, -0.0034704215, 0.013624546, 0.013779425, - -0.0025326884, -0.0018633928, 0.00014903376, 0.01547092, 0.008385425, - -0.0033495796, -0.015248458, -0.0356735, 0.005223496, -0.018293105, - -0.043073945, 0.016345823, -0.0050947615, 0.023554962, 0.034400985, - 0.0045644785, -0.00011241743, 0.0060564913, 0.0021182992, 0.01914424, - 0.019295372, -0.00551076, -0.00017086207, 0.0032044165, 0.010140755, - -0.022354674, 0.026089797 - ] - }, - "facial-recognition": [ - { - "boundingBox": { - "x1": 690.0, - "y1": -89.0, - "x2": 833.0, - "y2": 96.0 - }, - "score": 0.03575617074966431, - "embedding": [ - -0.43665668, -0.59305364, -0.12699714, 0.3985032, 0.1878969, - -0.25987914, 0.14818184, -0.542229, -0.06710237, -0.1319032, - 0.056408346, 0.046093762, -0.14984925, 0.043225512, 0.023826078, - -0.09063442, 0.07891726, -0.29357076, -0.6277133, -0.29042292, - 0.18038993, 0.21837695, 0.17909442, -0.040304773, -0.035560638, - -0.07568607, 0.1277122, -0.13466191, -0.2368693, 0.3642968, 0.29558533, - 0.20867407, 0.11252518, 0.47691494, -0.054775044, 0.030100197, - -0.049531147, 0.04045874, 0.23517768, 0.17130391, 0.17269331, - 0.08591308, 0.046999797, -0.17151847, -0.2443775, 0.3110528, - -0.23971468, -0.31744513, -0.026422635, -0.26203394, -0.18553479, - -0.31044272, 0.6385251, 0.27497086, 0.006674953, 0.053785797, - -0.20257844, -0.48399794, 0.21708605, -0.4781224, -0.12367296, - -0.099010885, 0.18633766, 0.31143454, -0.12165704, 0.13010044, - 0.12534627, 0.107288495, 0.37471777, -0.123026475, -0.1263274, - -0.15621608, 0.26027548, 0.15841314, 0.5164254, -0.31015784, 0.24754328, - 0.10240883, -0.1181829, -0.14073256, 0.027111322, 0.09927598, - -0.10066943, 0.4808423, -0.042361684, -0.08512197, 0.13695274, - 0.30378994, 0.11138052, -0.318214, -0.5708592, -0.14786953, 0.49985552, - -0.23231967, 0.13856675, -0.5383139, -0.059954256, 0.2796868, - -0.32447946, 0.16510965, 0.57146084, -0.15120608, 0.20110571, - -0.49805385, -0.2008879, -0.046678245, 0.24653266, 0.022508677, - -0.14091778, 0.38075653, 0.33811444, 0.05011098, -0.2371835, - -0.20052075, -0.14081016, -0.3422103, 0.11998144, 0.24423985, - 0.13769919, -0.25340003, -0.41080874, -0.28673622, -0.20673269, - 0.4604351, 0.4178845, 0.105202496, -0.1446912, 0.0807363, -0.37372503, - 0.13030809, -0.08456054, 0.21937889, -0.22700784, 0.3039499, - 0.009784861, -0.07245704, 0.50291365, -0.24968931, 0.3178813, - 0.12665558, -0.036484346, 0.21702805, -0.09277919, 0.17766781, - -0.12018812, 0.008044228, -0.26986086, 0.29888278, -0.28485933, - 0.30066437, -0.14316985, 0.53800535, 0.030840248, 0.023039162, - 0.73862207, 0.0034680888, 0.23797399, -0.11183337, 0.067846656, - -0.23546576, 0.39354736, 0.0053778216, 0.13494004, 0.1370637, - -0.029445097, 0.14705376, -0.48120612, 0.27262342, -0.05196667, - -0.3097266, 0.08714986, 0.10841283, -0.11757159, -0.5010461, - -0.32369986, -0.21964747, -0.19810468, 0.14780998, -0.04624281, - 0.24638015, -0.06710279, -0.31719172, 0.26955876, 0.37117082, - -0.3964724, 0.21541706, -0.12243534, -0.5392555, 0.04640211, 0.3657012, - -0.042127043, -0.030638859, 0.21909437, 0.16005577, -0.03320134, - -0.0949998, 0.33176076, 0.22538322, -0.016216129, -0.42417043, - 0.52940613, -0.011592716, -0.21875188, -0.06394625, 0.24449442, - -0.05658462, -0.09727913, -0.3978734, -0.11175068, 0.085142605, - -0.057618782, -0.0498557, 0.17287247, 0.41813853, -0.30433404, - 0.3087585, -0.6604493, -0.13869359, 0.072916515, -0.043251924, - 0.37401634, 0.17014223, -0.26469553, -0.34653437, 0.13010754, - 0.21517499, 0.74030113, 0.3460628, -0.5115478, 0.4696753, -0.009848075, - -0.1330159, -0.0061842054, 0.013667986, 0.16993025, -0.3161455, - 0.29015008, 0.65197945, 0.13776428, 0.5275149, 0.1472181, -0.114682674, - -0.05685012, 0.21696919, -0.34107065, 0.09352806, -0.03968816, - -0.13109599, 0.07406853, 0.15091223, 0.18835881, 0.19146737, -0.3898828, - 0.469747, -0.11145213, 0.039727956, 0.8268787, -0.09761663, - -0.043320894, 0.27001414, 0.12079324, 0.05877747, 0.028245524, - 0.20692128, 0.68440485, -0.34984088, -0.119763374, -0.39637753, - 0.23799005, 0.057573274, 0.07855352, 0.37982583, -0.0365879, - 0.068318695, 0.10845077, -0.18650186, 0.08927679, -0.27789003, - 0.31810492, 0.4251458, -0.03525705, -0.28072172, 0.07316002, 0.13499324, - -0.11333761, -0.0008841604, 0.10874095, 0.29681873, 0.008288942, - 0.24116173, 0.011309357, -0.3009541, -0.4752865, 0.19921738, - -0.16108191, 0.017838746, 0.51260126, -0.086799264, 0.34165853, - 0.32359147, 0.25770876, 0.21442738, -0.15971375, -0.26682994, - 0.22788364, -0.38956794, 0.084580205, -0.15929273, 0.24211408, - -0.24793725, -0.31528267, 0.15945697, -0.16866091, 0.19472758, 0.408394, - 0.24238603, -0.23643477, 0.29852632, 0.12915722, 0.327068, 0.501809, - -0.40538347, -0.023235738, -0.11315605, 0.007632144, -0.22626217, - 0.28817925, -0.5816528, 0.1551521, -0.016097836, -0.01634605, - 0.095855944, -0.010664792, 0.1402924, -0.22450349, -0.13961065, - -0.40732136, -0.24776831, 0.12040292, -0.06779129, 0.44510496, - 0.33206633, 0.19807269, -0.06460787, -0.2524265, -0.12726343, - 0.44656014, -0.09844789, -0.18762295, 0.16189753, 0.23589599, - -0.44798508, 0.2135099, -0.33205217, 0.28407755, -0.0951985, - 0.035582896, -0.51807857, -0.27382392, 0.03172898, 0.22928514, - 0.47157723, -0.48383215, 0.014225766, -0.08102345, 0.19384615, - -0.060681015, -0.037799604, -0.2875836, 0.024652202, 0.052712113, - -0.22610298, 0.46830428, 0.29616976, 0.14641494, -0.24234764, - 0.30126396, -0.011165038, 0.38622355, -0.12484505, 0.33650652, - 0.17399745, -0.2703057, 0.36919123, 0.26170117, -0.1537327, -0.43157104, - 0.35697, 0.043892622, -0.065475196, 0.5542902, 0.019970104, 0.43128124, - -0.014292087, -0.33983213, 0.3250854, 0.21585244, -0.34458104, - 0.23752448, 0.18115376, -0.2586738, -0.16033548, -0.16151018, - -0.23306333, 0.14865296, -0.31790328, 0.27215546, -0.059920013, - 0.16193654, 0.075943366, -0.16281635, 0.4489306, -0.43052202, - -0.038787995, -0.11722573, 0.07254093, -0.2997051, 0.16540596, - -0.15089649, 0.12507877, 0.43725327, 0.13540109, 0.13391787, - 0.013777234, 0.26951605, -0.2999856, -0.08645636, 0.12768297, - 0.23375636, -0.07325045, -0.04433371, 0.04709586, 0.09582621, - 0.23509142, 0.18061984, 0.35379466, 0.12938409, 0.33010754, 0.18966632, - 0.07585195, 0.0059688687, -0.13233723, 0.17105722, -0.020040989, - -0.2805646, -0.091034755, 0.1950869, -0.21115655, -0.16249251, - 0.07147664, -0.20138165, 0.15193966, 0.041464765, 0.01074836, - 0.029091328, -0.22078216, 0.06446775, -0.27403125, -0.51904315, - -0.20539844, 0.176225, -0.28688902, 0.030568387, 0.2964594, - -0.088931546, -0.4425866, 0.09070322, 0.08005672, 0.009866249, - -0.07386999, 0.06683251, -0.34370828, 0.23668535, -0.0847823, - -0.27400133, -0.31668398, -0.116622224, 0.20027944, 0.33772525, - -0.3041445, -0.61801887, 0.043022886, -0.24733649, -0.20657904, - -0.37058303, 0.00644885, 0.2548513, 0.029221226, -0.41749227, - 0.065117866, -0.3745206, 0.22699282, 0.22139677, 0.28097618, 0.10008535, - -0.039953396, -0.33505437, 0.28511694, 0.18131426, -0.879614, - -0.041319087, -0.62370497, 0.05170501, 0.23541749, -0.0033701807, - 0.15842043, 0.020002551, -0.22027364, -0.2730838, -0.23035137, - -0.077056274, 0.002099529 - ] - } - ], - "imageWidth": 600, - "imageHeight": 800 -} +{"clip": {"image": [-0.029240275, 0.02274738, -0.020163564, -0.016313331, 0.0009530146, 0.0010729627, 0.01782221, 0.04420342, 0.00032376708, 0.034060437, 0.040923152, -0.024181088, -0.0535712, -0.034794804, 0.010089799, -0.0028969506, -0.03506001, 0.030218706, 0.029766563, 0.023294315, 0.012413546, -0.004738038, -0.017942797, -0.036789585, 0.0052303826, -0.036181603, 0.009051782, 0.00932759, 0.0057503562, 0.015985414, 3.4025437e-05, 0.049436007, 0.015731763, -0.030918851, 0.019719727, -0.020576546, -0.005231761, -0.008333492, -0.029759789, 0.119099475, -0.0053320876, 0.004370226, -0.0022682582, 0.01789897, 0.02941454, -0.28252193, -0.06187134, 0.036058225, 0.0035019917, 0.004499765, 0.01413453, 0.03380403, -0.0032921762, -0.0024553582, -0.018557899, -0.001280334, 0.015386266, 0.0053381473, -0.002508861, -0.0050707674, -0.015148995, -0.042997293, 0.0124761835, -0.024046846, 0.0054675993, 0.00073952816, 0.043131523, -0.077897534, 0.026439417, -0.02380832, 0.007876205, -0.03306027, -0.015462417, -0.012476878, -0.024602083, -0.004384167, -0.019623984, 0.0063892286, -0.008998029, -0.05142428, 0.032637957, 0.032600347, 0.010663535, 0.026279477, 0.011462244, 0.02028843, -0.027531607, -0.012621107, 0.039038923, -0.025352659, -0.05156938, -0.011025616, -0.6725731, 0.040658288, -0.011361137, 0.061716076, 0.013698883, 0.004457988, -0.057969622, -0.10692411, -0.018529134, -0.004369318, -0.030310575, 0.0024083578, 0.019860104, 0.029141212, -0.07893188, 0.005939419, -0.011060325, -0.014225106, -0.024212057, -0.057036117, -0.03716643, -0.020800505, -0.01787459, -0.0024877065, -0.029530877, 0.009215086, 0.014535427, 0.043431558, 0.024118884, 0.011133975, -0.036812373, 0.019032136, 0.00082055596, -0.027445775, 0.02393861, -0.0048848973, 0.010327763, 0.012090686, 0.013239415, 0.017507533, -0.02365705, 0.078460515, 0.009419125, 0.0043909526, 0.018102704, -0.023108725, 0.008428085, -0.012307397, 0.0069907513, -0.0050128456, 0.003082945, 0.006579433, 0.0261621, 0.0075861793, 0.012343394, 0.026349016, -0.0063170977, -0.031295013, 0.061437856, -0.010111379, 0.092541546, -0.01952181, 0.035173517, 0.0036010533, 0.0045440826, -0.023346843, 0.00669373, -0.023700094, -0.0191347, 0.0072494824, -0.0028460417, -0.04747762, -0.010302901, 0.013440735, 0.0026138085, 0.028122144, 0.002231799, 0.017324915, 0.022839585, 0.011226531, -0.007104795, -0.010110564, -0.03943448, -0.027417231, 0.035023216, 0.007466257, -0.07993947, 0.006474745, 0.00438095, -0.0032703516, -0.059879813, -0.02456883, -0.011805905, -0.015260918, 0.0029254125, 0.008515833, 0.0062588216, -0.00930751, -0.008806945, 0.0023703414, -0.007838116, -0.0031276168, -0.06388368, 0.006610293, 0.031749472, 0.04533434, 0.011265556, -0.028518764, -0.003644105, -0.0042799907, -0.011587445, -0.01663729, 0.009182758, -0.012423823, -0.034826286, -0.032305446, -0.013168681, -0.0038836864, -0.02326945, -0.02496782, 0.0187495, -0.037348684, -0.032729626, -0.03352564, -0.0011865793, 0.022137962, 0.091489784, 0.012765105, 0.018697323, 0.045991085, -0.037717517, -0.053466607, 0.0392683, -0.0033215466, -0.014672436, 0.013624801, -0.004455726, 0.03386301, -0.006266096, -0.008634727, -0.0286226, 0.0002367907, 0.0054833116, 0.07611594, 0.020407192, 0.023071228, -0.03228062, -0.0010469863, -0.022632096, -0.0015129753, -0.0051687104, -0.031009467, -0.027093984, -0.03854705, 0.030401956, -0.014434913, -0.026718423, -0.005248297, -0.025074264, 0.02883739, 0.023266602, 0.018513815, -0.0016112248, -0.024298204, -0.027258217, 0.006657367, 0.055241838, 0.029772865, -0.014773467, 0.0003898727, -0.009236147, 0.11894643, -0.022672229, -0.039438512, -0.0056779194, 0.0298133, 0.020618236, 0.001813168, -0.017572122, -0.00958089, -0.020221429, 0.02207191, -0.017083341, 0.03473887, 0.024862533, -0.002068949, -0.05243768, 0.0015365004, 0.010737345, -0.011659201, 0.009595015, -0.02114795, 0.022264026, 0.020989275, -0.0950372, -0.010942789, -0.00563403, -0.024466451, 0.00057386636, 0.0010100049, 0.001916941, 0.006067649, -0.007447676, 0.03979671, -0.003139117, -0.020364363, 0.017236045, 0.011958324, 0.012651006, 0.032947008, -0.03181508, -0.03865767, 0.031411033, -0.052499652, 0.016763024, 0.022952998, 0.04835828, 0.022673093, -0.029616352, 0.02627719, 0.07818347, -0.02413108, 0.02037089, -0.021995421, 0.03421227, 0.023198899, -0.004189436, -0.018303256, 0.0057715955, 0.1230719, -0.022132011, -0.03965008, 0.0025436191, -0.017851725, 0.027440809, -0.031824205, 0.0039034756, 0.0003536298, -0.012361067, 0.02064873, 0.020198394, -0.0183519, -0.0026560547, -0.013674173, -0.0014906893, 0.030410081, -0.010374875, -0.026821297, -0.02242556, 0.041062664, -0.011202862, -0.02847564, 0.011953957, 0.018727072, 0.012663332, -0.011840388, -0.020886691, -0.0007219273, 0.0034005905, -0.003801304, 0.00090947666, 0.004812314, -0.0035383832, 0.0016857175, 0.025176654, -0.15410738, 0.006155143, -0.015186019, 0.048546694, -0.03260841, -0.008077512, 0.038287353, -0.05720707, -0.009176385, 0.008011662, -0.028211929, 0.025939306, -0.024916815, 0.0025883997, 0.012425659, 0.0047818753, -0.02380835, 0.03256469, 0.011183252, 0.113128014, -0.041271664, 0.005035802, 0.010346281, 0.01693021, -0.060056165, -0.01574106, -0.042067826, -0.010107939, -0.019539375, 0.012275092, -0.009981359, -0.014406553, 0.0377292, -0.05707011, 0.033400867, 0.009437524, -0.007360942, -0.030306397, -0.0042980304, -0.012032261, 0.017242596, -0.016564762, 0.007031159, 0.004422023, -0.023848636, 0.010254914, 0.0037647125, 0.036672875, -0.015325052, -0.014514223, 0.018100316, 0.015217877, -0.03895377, 0.01953362, 0.0025332512, -0.015211567, 0.022735089, -0.065954015, 0.01235176, 0.04623965, -0.029140005, 0.015668986, 0.0061620455, 0.009376337, -0.012881022, 0.014269646, -0.08320368, -0.048124664, -0.008810835, -0.024214962, -0.018557087, 0.029500315, 0.0064427103, 0.037454795, 0.025400363, 0.015685672, -0.025242975, -0.009204951, -0.011612312, -0.0049718856, -0.0076911235, -0.024881033, -0.02185726, 0.056367826, 0.04761963, -0.023618726, -0.00898248, -0.012248864, 0.008529039, -0.018900419, -0.008615848, 0.022764448, -0.048481625, -0.013486573, 0.021670641, -0.023852944, -0.037358228, 0.033274513, -0.026645055, 0.050488036, -0.019671183, 0.022968937, -0.0047809645, 0.0073866877, 0.013722248, -0.020764614, -0.0045974255, -0.0024552948, 0.0061884224, -0.026321784, -0.023174278, -0.01875276, -0.005241128, 0.04970341, 0.044514913, 0.0033109477, -0.02634, 0.029782368, 0.014291489, 0.00713931, -0.005261876, 0.0044720178, 0.017025232, -0.0015874332, -0.015341334, -0.0015963981, 0.046385005, 0.03236695, 0.0009339123, 0.017759783, 0.0082206065, -0.0008270097, -0.022638528, 0.00769505, 0.00788198, -0.017030075, -0.04163034, 0.0106496755, -0.034192007, -0.0015860915, -0.020517178, -0.032945078, -0.0053297435, 0.030102672, -0.029427446, -0.0044684974, 0.029017493, 0.14281422, 0.010522923, 0.027442401, -0.0006020587, -0.043298963, 0.011711217, -0.051137257, 0.01557905], "text": [-0.00377849, -0.0012915516, 0.013966385, 0.005374147, 0.003708035, -0.0146752475, -0.019110048, -0.08002025, -0.03458918, -0.0021400691, -0.0050363126, -0.049290482, -0.04186416, -0.034567855, 0.0193284, -0.013843339, 0.006643949, 0.015310201, -0.019513344, 0.03151202, 0.0482263, 0.011914014, 0.035720177, -0.043771014, -0.025255188, 0.032709885, 0.01917336, 0.03154133, -0.01403091, 0.016226673, 0.016670639, 0.02228914, -0.0022885373, 0.028669102, -0.03321755, 0.008693166, 0.047847517, 0.004506587, -0.04167899, 0.0019039687, -0.013117988, -0.014263629, 0.045430332, 0.014383425, 0.05601821, 0.00049987243, -0.009474628, 0.03574858, 0.033481594, -0.0015193712, -0.0013497402, 0.009645607, 0.025770081, -0.034270316, -0.009055372, -0.012197793, -0.000744126, -0.012675848, 0.027288472, 0.038235836, 0.010385463, -0.013494433, 0.0069217896, 0.0034014569, 0.018213907, -0.02141165, -0.005895786, -0.016048692, 0.03632707, 0.009623559, -0.009850707, -0.017509611, -0.03192185, 0.037554655, -0.024094664, -0.038104467, 0.009635034, 0.030345341, 0.017167449, -0.045698076, 0.0018091919, 0.05246557, 0.020726752, 0.0027810642, -0.026376907, -0.014414313, 0.020750526, -0.026784055, 0.03549227, -0.0069205053, -0.00020645747, -0.026004054, -0.11643398, 0.010628068, 0.0013605448, 0.0065109152, 0.013906813, -0.0089813685, -0.043357566, -0.035713397, -0.013247194, 0.008679616, 0.028527435, 0.002929705, -0.021830956, -0.023007216, -0.00068055827, 0.026913393, 0.0020245092, -0.00083614455, -0.053684544, -0.03807718, 0.01588328, 0.00246405, -0.024368426, -0.0013492755, -0.04259538, 0.031105986, -0.01810126, 0.007885868, 0.04096929, -0.055228896, -0.029127521, -0.040628076, 0.015561375, -0.016481364, -0.008955004, -0.028511193, 0.03761003, 0.04187235, 0.006159584, -0.015481074, -0.020519646, 0.5737937, 0.0019995875, 0.0055420515, 0.019577935, -0.04497258, 0.006789208, -0.029090641, -0.012311774, 0.009641061, -0.04204943, 0.006311922, -0.006451449, 0.00771893, 0.02750124, -0.0442429, 0.009942242, -0.035988983, 0.0034268626, -0.022107868, 0.025340268, 0.027656747, 0.030932598, 0.017699726, -0.005204463, -0.02277333, 0.0039381254, -0.004762562, -0.025038576, -0.010762719, -0.007172735, -0.035541486, -0.030383216, 0.0019866321, 0.020151345, -0.017473131, 0.009111604, 0.020751633, 0.005240196, -0.013846646, -0.025847463, -1.3745466e-05, -0.0039132237, 0.03454592, -0.025891943, 0.010132164, -0.022684896, 0.014084669, -0.001998648, 0.0015486531, -0.05517993, -0.0019543923, -0.0068913447, 0.015614538, 0.004655365, -0.009654506, -0.0056707286, -0.021899402, 0.020798812, 0.015453464, 0.043110102, -0.007159732, -0.03752856, 0.0019941884, -0.000685062, 0.014257045, 0.026522951, -0.024141379, 0.023501279, -0.049575273, -0.0007706183, 0.01265391, 0.011421983, -0.028066738, 0.0132201575, -0.042931046, -0.010257662, -0.0021566458, -0.002084107, 0.019634115, -0.010039014, 0.027098438, -0.018874772, -0.04714307, -0.010699856, -0.0058649103, 0.025120156, 0.007346442, -0.008592694, 0.008302162, -0.04190575, -0.04546809, 0.036745828, -0.020235257, 0.027421871, -0.006840629, -0.0027657642, 0.0016024954, 0.012953273, -0.01691378, 0.018790675, 0.0067636985, -0.029318608, 0.062098745, -0.007967272, 0.019493792, -0.036695562, -0.040483467, 0.0070916987, 0.035738114, -0.033661954, -0.00010214538, 0.020114105, -0.010346232, 0.0065442612, -0.027712272, -0.005088254, -0.017700551, -0.01533876, 0.015322339, 0.032468606, 0.014415187, -0.0071319668, 0.0029240989, -0.0032388463, -0.0055279015, 0.019732248, 0.06113575, -0.0064503155, 0.0066482383, 0.016748711, 0.0042133313, -0.021386685, 0.01769606, 0.0012047674, -0.017313577, -0.027136903, 0.008160362, -0.014473728, 0.032427505, -0.024649428, -0.0058570183, 0.00068549847, 0.040602386, 0.017727202, 0.018005662, 0.05793844, -0.00838984, -0.0006229641, 0.034855265, 0.010399062, 0.013727231, -0.03227891, -0.0150364125, 0.0012288038, 0.023029301, 0.024083996, -0.008781946, -0.04345279, 0.00011917178, -0.035171103, 0.012486819, -0.027204808, 0.01960962, 0.0013456346, 0.014400163, 0.012646647, -0.018884279, -0.016848212, -0.00072430464, -0.01609786, -0.011934695, -0.046755917, -0.0021253747, 0.012914511, -0.012384807, 0.033379596, -0.0045600235, 0.0018349845, -0.0041356944, 0.57307595, -0.010784544, 0.030332904, -0.031053256, 0.031344827, 0.046923596, 0.0333786, -0.014280093, 0.011716355, 0.04666055, 0.015004827, -0.04556153, -0.022805158, 0.018753275, 0.019795429, -0.03874178, -0.014518191, -0.13193808, -0.016464327, -0.013485259, -0.002359569, -0.03451385, -0.035092168, -0.028210662, 0.031448837, 0.04086557, -0.021704506, -0.015243849, -0.019880893, -0.0082949735, 0.024297263, -0.008026394, -0.027068129, 0.0037647572, 0.021149805, 0.013022537, -0.0014667712, 0.025677683, 0.010282272, -0.008790261, 0.042281196, 0.050509956, -0.047956854, -0.018188674, 0.023284355, -0.028096361, -0.04272614, 0.012470676, 0.04939491, -0.026531855, -0.012098173, 0.034240864, 0.014431801, -0.011577003, -0.011446267, 0.020599581, 0.022248501, -0.02279409, 0.017593509, 0.004292014, 0.007667772, 0.036062628, -0.010371202, 0.011502505, -0.014144773, -0.005132077, -0.006133113, 0.010585359, 0.0020189546, -0.009133387, 0.012453563, 0.009159339, 0.02919466, -0.0019923102, -0.010779508, 0.022078961, -0.022626825, 0.013201965, -0.06437432, 0.05752333, 0.0039885542, 0.016976768, -0.025057603, -0.0097831795, 0.007920791, 0.003893596, -0.01965289, 0.06852016, 0.013193639, -0.010074425, 0.0100556435, 0.02102327, 0.011133714, -0.007115316, -0.025725761, -0.02117057, 0.00640587, -0.011512984, 0.022911677, 0.014856382, 0.010215653, 0.0033845454, 0.02600312, 0.0087262355, 0.032080073, -0.027731141, 0.0015452767, 0.023681948, 0.022449773, -0.03942745, -0.0033734506, -0.090638615, 0.004891829, -0.006286816, -0.016344903, 0.006160707, 0.011177469, 0.011902139, 0.018553948, -0.005686419, -0.009042075, -0.0064733475, -0.017129358, -0.007700515, 0.013193817, 0.014671132, 0.017838879, 0.0010372871, -0.008472803, -0.0076532993, -0.0019450359, -0.026438195, -0.0039837547, 0.0160693, 0.00067757437, -0.020706588, 0.014210452, -0.022126926, -0.0238694, 0.022997815, -0.018322313, -0.004728166, 0.038256954, 0.0013424249, 0.028170351, -0.0061870427, 0.0090107545, -0.021891015, -0.002931492, 0.03504575, 0.031653132, 0.0062208693, 0.03467824, -0.0060050013, -0.002888752, -0.011318142, 0.0064615067, -0.0015551174, 0.022771595, -0.011298304, 0.021120522, -0.032193486, 0.0047918656, -0.0045222784, 0.0024909836, -0.010143601, -0.0012446165, 0.035437353, 0.003977082, -0.033944316, 0.012640697, -0.006762388, -0.00973952, 0.015522455, -0.005274172, 0.027863417, 0.027938526, -0.025374765, 0.062127434, 0.0020044039, -0.01779412, 0.10606827, 0.031732988, -0.014786097, 0.0100396145, -0.01334179, -0.030271435, -0.028912347, 0.04303614, -0.0035223681, -0.00508675, 0.03752229, 0.06613816, -0.020030629, -0.009390152, -0.0005004571, -0.020006225, -0.044624373, -0.039662942, 0.060866963]}, "facial-recognition": {"imageWidth": 840, "imageHeight": 700, "faces": [{"boundingBox": {"x1": 271.0, "y1": 159.0, "x2": 356.0, "y2": 273.0}, "embedding": [-0.02109016, 0.022240654, -0.026227806, -0.05155147, 0.049118947, 0.0029655523, 0.046891544, 0.020310389, -0.022646084, 0.01134918, -0.024934702, -0.030037647, -0.031309005, -0.0012668095, -0.026567042, -0.07635356, -0.02562877, -0.04984194, -0.0018215338, 0.05258435, -0.010422444, -0.05972049, -0.03207654, 0.016415888, -0.0018812593, 0.0820367, 0.0066207186, -0.055461194, -0.051618673, 0.041603304, 0.063333854, 0.064513415, -0.036522094, 0.046367053, 0.06820749, 0.06006578, 0.016158676, 0.017209368, 0.060429297, 0.08287553, -0.0022089495, -0.028104974, -0.0011498835, -0.01641029, -0.032944806, -0.014212797, -0.06011639, 0.03144668, 0.08421291, -0.006728484, -0.011458434, 0.02499428, 0.0402121, -0.015981779, 0.035817374, 0.08059203, 0.036166046, 0.051513035, 0.008670421, -0.0003494511, 0.009196479, -0.023113282, -0.014471686, -0.014376478, 0.02399979, -0.09826273, 0.0029255454, 0.06759321, 0.061925683, -0.042563133, 0.016879713, -0.00340339, -0.059502296, -0.014200149, -0.011815452, -0.028992267, 0.039617654, 0.03969662, -0.03325022, 0.05487882, 0.013048679, 0.028593926, -0.03801353, 0.018216677, 0.026949344, -0.056956366, -0.019216467, 0.029655717, 5.9537073e-05, 0.015320147, 0.06640795, 0.11813518, -0.049639888, -0.037013453, 0.036590148, -0.09359642, -0.013693963, -0.015224149, 0.06017036, -0.025735453, 0.08181334, -0.00954105, 0.026677586, -0.057009693, 0.04790535, -0.02605019, 0.041463066, 0.04516405, 0.04617457, 0.008583263, -0.011979799, -0.02104466, -0.0013846458, 0.025714133, 0.0111953085, 0.0003897781, 0.09347127, -0.045572978, -0.050724946, 0.01651668, 0.00842076, 0.07133321, -0.03769401, -0.12868567, -0.0055657336, -0.030446453, -0.038680833, 0.070453025, -0.046699736, -0.07327703, 0.07349414, -0.025136275, 0.041606314, 0.027192492, -0.03309477, -0.047578685, 0.020764457, -0.05969532, 0.036298834, -0.013373041, -0.039533332, 0.008016592, 0.0073562204, -0.0059271124, 0.030903146, 0.08904637, -0.06940869, -0.00879312, -0.13217753, 0.045521118, -0.0011331779, -0.01631236, 2.2618146e-05, 0.07898381, -0.02402306, -0.03648393, 0.04923067, 0.04355395, 0.011747625, -0.03877474, -0.029971004, 0.005220548, -0.013913924, -0.017296601, -0.059844308, -0.0059725675, -0.04702892, 0.0030567865, -0.079060726, -0.007029555, -0.050632905, -0.02699259, -0.0072723846, 0.051899362, -0.056258608, -0.020023236, 0.05895745, 0.049476776, -0.037481736, 0.054740164, -0.065221556, 0.004789071, -0.026941445, -0.05156219, -0.026733197, 0.0023032185, -0.04172011, -0.009579965, 0.043398514, 0.023691908, -0.031915914, 0.012704605, -0.07929288, -0.023400098, -0.017118972, -0.08946523, -0.06535939, 0.024354056, 0.007947206, -0.009458882, 0.037135456, -0.013232659, 0.018360876, 0.010736806, -0.073271245, -0.05172269, -0.0648574, -0.0262401, -0.0153322285, 0.05693816, 0.016138993, 0.05182635, -0.05635323, -0.016383078, -0.051853433, 0.021441665, 0.0059856814, 0.0345592, 0.020173322, 0.021153774, 0.06674612, 0.012637243, -0.02611364, 0.014571983, -0.035413478, 0.026567234, -0.019131972, -0.0046005896, -0.062435772, -0.026113842, -0.034159046, 0.009861883, 0.02600461, -0.038663063, 0.023845121, 0.038411032, 0.030545078, 0.0067734267, -0.0018502183, 0.07116862, 0.028058557, 0.051142003, 0.048419952, 0.036133725, 0.013653307, 0.052527156, 0.0028999227, 0.05951973, 0.023294574, -0.031254347, 0.020496838, 0.041185338, 0.012341979, -0.032141518, -0.058980536, -0.067163706, -0.0087312935, -0.049131174, -0.0014268243, 0.026549103, 0.0010296477, -0.018504227, -0.02070832, 0.10455088, 0.05974241, 0.02355615, -0.06970922, -0.018376207, 0.04915733, 0.007919119, -0.017177325, 0.041967094, -0.06020687, 0.004122747, -0.09327306, 0.036126312, 0.041461047, -0.058268014, 0.0057649505, -0.037257507, -0.10673296, -0.019929433, 0.031015912, -0.03131743, 0.004703155, 0.0021089648, 0.095528476, -0.037962947, -0.016162816, -0.04756664, 0.0060427883, -0.027776098, -0.022376593, -0.037984826, -0.03541779, 0.006211417, 0.0033174818, 0.05873058, 0.031586178, -0.03981285, 0.009193386, -0.087054215, -0.017509116, 0.023940355, -0.04521575, -0.035313338, -0.048166055, 0.059896406, -0.008972068, 0.046428304, -0.043441687, 0.070813425, 0.041591186, -0.03260059, 0.025136124, 0.06349667, -0.017939935, -0.019221002, 0.02087589, -0.031323336, 0.03951066, -0.015193063, 0.009318058, 0.05542475, -0.03079061, 0.038844343, 0.09238112, 0.05503655, 0.013255708, 0.07829127, 0.02934044, 0.053244863, 0.046482548, 0.035087477, -0.014583196, 0.05370969, 0.015286031, -0.015016701, -0.016063195, 0.0076204, -0.003374417, -0.0080213845, 0.024443319, -0.03873362, 0.08728878, 0.05896774, -0.05862249, -0.07154778, 0.025585681, -0.04580391, -0.103926145, 0.015800381, 0.056916792, -0.039125975, 0.05889627, -0.031576596, -0.0067394045, 0.021304216, -0.09174046, -0.049721844, -0.01496743, -0.010886794, 0.038958646, -0.0722713, -0.027069159, -0.03581728, -0.13875493, -0.09064049, -0.02001534, -0.05257496, 0.0069434205, 0.06362916, 0.017999075, 0.054439567, -0.042295467, 0.008325071, 0.074246295, 0.049776148, -0.044039264, 0.015481024, 0.11235533, 0.027103968, 0.024569018, 0.032471776, 0.014865656, 0.051862244, -0.029630529, -0.06417371, 0.003972846, -0.006920887, 0.019362615, 0.01369384, -0.059984356, 0.08630591, 0.074628614, 0.061739635, 0.008418423, -0.038174078, -0.08334343, 0.051286977, 0.013346276, 0.059975393, 0.043665156, -0.03155817, -0.00117017, 0.016538389, 0.014414994, -0.026954196, -0.051562443, -0.033782393, -0.01819859, 0.059424944, -0.018533932, -0.06567998, -0.014865563, -0.008695496, 0.04592779, -0.07062889, 0.0701535, -0.006305198, 0.017896414, -0.059151568, 0.0063205776, 0.010077412, 0.04935787, -0.04388182, -0.021789722, -0.050149724, 0.010672742, 0.027434545, -0.0029799917, -0.055056665, 0.008173985, -0.045736484, 0.029440103, 0.039946903, 0.069309875, -0.06679647, -0.009734461, -0.0010366557, -0.02580478, 0.037709497, 0.09016693, -0.061619915, 0.008214737, -0.0030659584, -0.01748481, -0.020389073, 0.017560003, -0.015404108, -0.0040513817, -0.037783537, 0.005536573, -0.027024597, -0.07526226, 0.028097268, -0.000898764, -0.039536834, -0.013770835, -0.08674664, 0.023270123, -0.013440714, -0.06855993, -0.040267516, 0.02090837, 0.0024361138, -0.014654124, -0.05955025, -0.008844592, 0.03937242, 0.012703652, 0.060103923, -0.07835982, -0.040523235, 0.039314635, 0.027467607, -0.033941492, -0.047454245, -0.038684025, 0.0976016, 0.031567827, 0.0036655331, -0.09418956, -0.03491417, 0.060269743, -0.027865143, 0.072315596, 0.029369606, 0.0591703, 0.03939378, 0.046697427, -0.010640258, 0.05346348, 0.040919587, -0.0018008988, -0.023286529, -0.03356717, -0.071884066, -0.008992476, -0.004773128, 0.06961985, -0.032190185, -0.07336817, -0.06919444, -0.03129296, -0.03060653, -0.055539172, -0.0029378114, -0.03436031, -0.028586183, -0.026332222, 0.021722393], "score": 0.8676442}]}, "ocr": {"box": [0.0076666665, 0.018, 0.068, 0.0155, 0.06866667, 0.04975, 0.008333334, 0.05225, 0.2045, 0.015, 0.338, 0.01675, 0.33783334, 0.05125, 0.20433334, 0.0495, 0.4795, 0.0165, 0.61, 0.0185, 0.60966665, 0.0525, 0.47916666, 0.0505, 0.7288333, 0.02325, 0.8948333, 0.02875, 0.89433336, 0.0615, 0.72833335, 0.056, 0.0, 0.06225, 0.049833335, 0.06225, 0.049833335, 0.08175, 0.0, 0.08175, 0.15916666, 0.06, 0.35283333, 0.06, 0.35283333, 0.091, 0.15916666, 0.091, 0.0053333333, 0.08, 0.0305, 0.08, 0.0305, 0.08975, 0.0053333333, 0.08975, 0.15383333, 0.09225, 0.25733334, 0.09725, 0.25616667, 0.15025, 0.15266667, 0.14525, 0.25716665, 0.0975, 0.298, 0.0975, 0.298, 0.13475, 0.25716665, 0.13475, 0.43283334, 0.0955, 0.5845, 0.10375, 0.5836667, 0.14175, 0.43183333, 0.13375, 0.0, 0.10375, 0.0475, 0.0965, 0.05, 0.1325, 0.0005, 0.13975, 0.15583333, 0.1475, 0.277, 0.1475, 0.277, 0.18125, 0.15583333, 0.18125, 0.15483333, 0.19075, 0.33233333, 0.19075, 0.33233333, 0.225, 0.15483333, 0.225, 0.7101667, 0.1945, 0.7711667, 0.19675, 0.7705, 0.233, 0.7095, 0.23075, 0.1825, 0.23825, 0.35316667, 0.23625, 0.35333332, 0.2695, 0.18266666, 0.2715, 0.5635, 0.246, 0.5936667, 0.246, 0.5936667, 0.26475, 0.5635, 0.26475, 0.196, 0.32125, 0.334, 0.3195, 0.33433333, 0.36825, 0.19616666, 0.37025, 0.0, 0.33125, 0.07783333, 0.336, 0.07683333, 0.3735, 0.0, 0.36875, 0.41983333, 0.5335, 0.5683333, 0.5335, 0.5683333, 0.61575, 0.41983333, 0.61575], "text": ["STAY", "MAIN STAY", "SEASONAL", "BARREL AGED", "lsch Style", "WEEKEND WARRIOR PaleAle", "LONDE", "BIGBRICHT", "IPA", "Mother-in-Lager", "adler", "HOPADILLO IPA", "RODEO CLOWN IPA", "B.B.H.", "KR&D SERIES", "NAL", "Session IPA", "achbrewing", "F.U.N."], "boxScore": [0.7672492, 0.736339, 0.77014536, 0.7674111, 0.7152776, 0.79979694, 0.5252589, 0.59605956, 0.84168875, 0.810685, 0.757814, 0.70302576, 0.79601747, 0.79040974, 0.8279445, 0.64908415, 0.8242554, 0.7369857, 0.85597765], "textScore": [0.9989463, 0.9239445, 0.9852257, 0.9771696, 0.91195405, 0.92228043, 0.9877033, 0.9603994, 0.9845086, 0.97863936, 0.9901719, 0.9738971, 0.9209216, 0.9732496, 0.9929817, 0.99414176, 0.9643116, 0.99681664, 0.92819786]}} \ No newline at end of file diff --git a/machine-learning/test_main.py b/machine-learning/test_main.py index 27191d67f..d4a77dd6a 100644 --- a/machine-learning/test_main.py +++ b/machine-learning/test_main.py @@ -1,15 +1,15 @@ import json import os +import pickle import platform import sys import threading import time from concurrent.futures import ThreadPoolExecutor -from io import BytesIO from pathlib import Path from random import randint from types import SimpleNamespace -from typing import Any, Callable +from typing import Any, Callable, TypeVar from unittest import mock import numpy as np @@ -18,13 +18,23 @@ import orjson import pytest from fastapi import HTTPException from fastapi.testclient import TestClient +from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf from PIL import Image from pytest import MonkeyPatch from pytest_mock import MockerFixture +import immich_ml.sessions.prepare as prepare_module from immich_ml import allocator from immich_ml.config import MaxBatchSize, PreloadModelData, Settings, settings -from immich_ml.main import app, lifespan, load, preload_models, update_state +from immich_ml.main import ( + MODEL_FILE_ERRORS, + app, + attempt, + lifespan, + load, + preload_models, + update_state, +) from immich_ml.main import run_inference as run_request from immich_ml.models.base import InferenceModel from immich_ml.models.cache import ModelCache @@ -32,24 +42,15 @@ from immich_ml.models.clip.textual import MClipTextualEncoder, OpenClipTextualEn from immich_ml.models.clip.visual import OpenClipVisualEncoder from immich_ml.models.facial_recognition.detection import FaceDetector from immich_ml.models.facial_recognition.recognition import FaceRecognizer +from immich_ml.models.ocr.ctc import logits, probabilities from immich_ml.models.ocr.detection import TextDetector from immich_ml.models.ocr.recognition import TextRecognizer -from immich_ml.schemas import ModelFormat, ModelPrecision, ModelTask, ModelType +from immich_ml.schemas import ModelFormat, ModelTask, ModelType, Shape from immich_ml.sessions.ann import AnnSession -from immich_ml.sessions.ort import OrtSession, flush_denormals +from immich_ml.sessions.ort import Device, GraphSpec, OrtSession, flush_denormals, fresh, prepared +from immich_ml.sessions.policy import ShapePolicy, batches, runs from immich_ml.sessions.rknn import RknnSession, run_inference - - -class FakeLock: - def __init__(self) -> None: - self.enter = mock.Mock() - self.exit = mock.Mock() - - def __enter__(self) -> None: - self.enter() - - def __exit__(self, *args: object) -> None: - self.exit(*args) +from immich_ml.sessions.rknn import model_path as rknn_model_path class TestBase: @@ -93,7 +94,6 @@ class TestBase: def test_sets_default_model_format_to_armnn_if_available(self, path: mock.Mock, mocker: MockerFixture) -> None: mocker.patch.object(settings, "ann", True) mocker.patch("immich_ml.sessions.ann.loader.is_available", True) - path.suffix = ".armnn" encoder = OpenClipTextualEncoder("ViT-B-32__openai", cache_dir=path) @@ -107,6 +107,26 @@ class TestBase: assert encoder.model_format == ModelFormat.ARMNN + def test_hands_the_session_the_shapes_the_model_feeds(self, path: mock.Mock, mocker: MockerFixture) -> None: + session = mocker.patch("immich_ml.models.base.OrtSession") + + FaceDetector("buffalo_l", cache_dir=path)._make_session() + + assert session.call_args.args[1].dims == (Shape(batch=1, height=640, width=640),) + + def test_opens_the_rknn_binary_compiled_for_the_requested_shapes(self, mocker: MockerFixture) -> None: + mocker.patch("immich_ml.sessions.rknn.model_prefix", Path("rknpu/rk3588")) + + assert rknn_model_path(Path("/cache/detection"), "res1088") == Path( + "/cache/detection/rknpu/rk3588/res1088/model.rknn" + ) + assert rknn_model_path(Path("/cache/visual")) == Path("/cache/visual/rknpu/rk3588/model.rknn") + mocker.patch.object(settings, "model_revision", "v2") # the older exports come as ONNX alone + detector = TextDetector( + "PP-OCRv5_mobile", cache_dir="/cache", model_format=ModelFormat.RKNN, maxResolution=1088 + ) + assert detector.model_path == Path("/cache/detection/rknpu/rk3588/res1088/model.rknn") + def test_sets_default_model_format_to_rknn_if_available(self, mocker: MockerFixture) -> None: mocker.patch.object(settings, "rknn", True) mocker.patch("immich_ml.sessions.rknn.is_available", True) @@ -163,6 +183,33 @@ class TestBase: path.return_value.mkdir.assert_called_once() warning.assert_called_once() + def test_unload_lets_go_of_the_session(self, stub_session: Callable[..., mock.Mock]) -> None: + encoder = OpenClipTextualEncoder("ViT-B-32__openai", session=stub_session((1, 77))) + + encoder.unload() + + assert not encoder.loaded + assert not hasattr(encoder, "session") + + def test_builds_by_warming_every_graph(self, stub_session: Callable[..., mock.Mock]) -> None: + session = stub_session((1, 112, 112, 3), shapes=(Shape(batch=1), Shape(batch=4))) + + FaceRecognizer("buffalo_l", session=session).build() + + session.warm.assert_called_once() + + @pytest.mark.parametrize( + ("symbols", "called"), [(["mi_collect", "malloc_trim"], "mi_collect"), (["malloc_trim"], "malloc_trim")] + ) + def test_returns_memory_through_the_allocator_the_process_runs_on( + self, mocker: MockerFixture, symbols: list[str], called: str + ) -> None: + process = mocker.patch("immich_ml.allocator._process", mock.Mock(spec=symbols)) + + allocator.release() + + getattr(process, called).assert_called_once() + def test_download(self, snapshot_download: mock.Mock) -> None: encoder = OpenClipTextualEncoder("ViT-B-32__openai", cache_dir="/path/to/cache") encoder.download() @@ -201,21 +248,32 @@ class TestBase: snapshot_download.assert_called_once() ort_session.assert_not_called() + +class TestShapePolicy: + def test_pins_only_what_cannot_vary(self) -> None: + policy = ShapePolicy(dims=(Shape(batch=1, height=64, width=64), Shape(batch=1, height=64, width=128))) + + assert policy.pinned == Shape(batch=1, height=64) + + def test_batches_follow_the_configured_maximum(self) -> None: + assert batches(1) == (1,) + assert batches(6) == (1, 6) # a single row as well, so a remainder is never padded + @pytest.mark.parametrize( - ("symbols", "called"), [(["mi_collect", "malloc_trim"], "mi_collect"), (["malloc_trim"], "malloc_trim")] + ("total", "sizes", "expected"), + [ + (10, (1, 6), [6, 1, 1, 1, 1]), # the remainder runs a row at a time + (12, (1, 6), [6, 6]), + ], ) - def test_returns_memory_through_the_allocator_the_process_runs_on( - self, mocker: MockerFixture, symbols: list[str], called: str + def test_runs_split_the_rows_the_session_will_take( + self, total: int, sizes: tuple[int, ...], expected: list[int] ) -> None: - process = mocker.patch("immich_ml.allocator._process", mock.Mock(spec=symbols)) - - allocator.release() - - getattr(process, called).assert_called_once() + assert runs(total, sizes) == expected @pytest.mark.usefixtures("ort_session") -class TestOrtSession: +class TestOrtSessions: CPU_EP = ["CPUExecutionProvider"] CUDA_EP = ["CUDAExecutionProvider", "CPUExecutionProvider"] OV_EP = ["OpenVINOExecutionProvider", "CPUExecutionProvider"] @@ -225,161 +283,367 @@ class TestOrtSession: COREML_EP = ["CoreMLExecutionProvider", "CPUExecutionProvider"] @pytest.mark.providers(CPU_EP) - def test_sets_cpu_provider(self, providers: list[str]) -> None: - session = OrtSession("ViT-B-32__openai") + def test_sets_cpu_provider(self, ort_session: mock.Mock, providers: list[str]) -> None: + ort_sessions("ViT-B-32__openai") - assert session.providers == self.CPU_EP + assert given_providers(ort_session) == self.CPU_EP @pytest.mark.providers(CUDA_EP) - def test_sets_cuda_provider_if_available(self, providers: list[str]) -> None: - session = OrtSession("ViT-B-32__openai") + def test_sets_cuda_provider_if_available(self, ort_session: mock.Mock, providers: list[str]) -> None: + ort_sessions("ViT-B-32__openai") - assert session.providers == self.CUDA_EP + assert given_providers(ort_session) == self.CUDA_EP @pytest.mark.ov_device_ids(["GPU.0", "CPU"]) @pytest.mark.providers(OV_EP) - def test_sets_openvino_provider_if_available(self, providers: list[str], ov_device_ids: list[str]) -> None: - session = OrtSession("ViT-B-32__openai") + def test_sets_openvino_provider_if_available( + self, ort_session: mock.Mock, providers: list[str], ov_device_ids: list[str] + ) -> None: + ort_sessions("ViT-B-32__openai") - assert session.providers == self.OV_EP + assert given_providers(ort_session) == self.OV_EP @pytest.mark.providers(CUDA_EP_OUT_OF_ORDER) - def test_sets_providers_in_correct_order(self, providers: list[str]) -> None: - session = OrtSession("ViT-B-32__openai") + def test_sets_providers_in_correct_order(self, ort_session: mock.Mock, providers: list[str]) -> None: + ort_sessions("ViT-B-32__openai") - assert session.providers == self.CUDA_EP + assert given_providers(ort_session) == self.CUDA_EP @pytest.mark.providers(TRT_EP) - def test_ignores_unsupported_providers(self, providers: list[str]) -> None: - session = OrtSession("ViT-B-32__openai") + def test_ignores_unsupported_providers(self, ort_session: mock.Mock, providers: list[str]) -> None: + ort_sessions("ViT-B-32__openai") - assert session.providers == self.CUDA_EP + assert given_providers(ort_session) == self.CUDA_EP @pytest.mark.providers(ROCM_EP) - def test_uses_rocm(self, providers: list[str]) -> None: - session = OrtSession("ViT-B-32__openai") + def test_uses_rocm(self, ort_session: mock.Mock, providers: list[str]) -> None: + ort_sessions("ViT-B-32__openai") - assert session.providers == self.ROCM_EP + assert given_providers(ort_session) == self.ROCM_EP @pytest.mark.providers(COREML_EP) - def test_uses_coreml(self, providers: list[str]) -> None: - session = OrtSession("ViT-B-32__openai") + def test_uses_coreml(self, ort_session: mock.Mock, providers: list[str]) -> None: + ort_sessions("ViT-B-32__openai") - assert session.providers == self.COREML_EP + assert given_providers(ort_session) == self.COREML_EP - def test_sets_provider_kwarg(self) -> None: - providers = ["CUDAExecutionProvider"] - session = OrtSession("ViT-B-32__openai", providers=providers) + def test_leaves_a_dimension_free_when_the_model_feeds_several_sizes( + self, ort_session: mock.Mock, mocker: MockerFixture + ) -> None: + sess_options = mocker.patch("immich_ml.sessions.ort.ort.SessionOptions").return_value + policy = ShapePolicy(dims=(Shape(batch=1, height=64, width=64), Shape(batch=1, height=64, width=128))) - assert session.providers == providers + ort_sessions("PP-OCRv5_mobile", providers=["CPUExecutionProvider"], shape_policy=policy) + + assert sess_options.add_free_dimension_override_by_name.call_args_list == [ + mock.call("batch", 1), + mock.call("DynamicDimension.0", 1), + mock.call("None", 1), + mock.call("height", 64), + mock.call("DynamicDimension.1", 64), + mock.call("?", 64), + ] + + def test_builds_one_graph_on_cpu_that_runs_a_row_at_a_time_at_any_size_of_what_varies( + self, ort_session: mock.Mock + ) -> None: + widths = (Shape(batch=1, height=64, width=64), Shape(batch=4, height=64, width=128)) + session = ort_sessions("PP-OCRv5_mobile", providers=["CPUExecutionProvider"], shape_policy=ShapePolicy(widths)) + + assert session.shapes == (Shape(batch=1, height=64),) # the width is left out, so it takes any + session.for_shape(Shape(batch=4, height=64, width=128)) + session.for_shape(Shape(batch=1, height=64, width=96)) + + ort_session.assert_called_once() + + def test_builds_a_graph_once_for_the_requests_that_arrive_wanting_it(self, ort_session: mock.Mock) -> None: + together = threading.Barrier(2) + policy = ShapePolicy(dims=(Shape(batch=1, height=64, width=64), Shape(batch=1, height=64, width=128))) + session = ort_sessions("PP-OCRv5_mobile", providers=["CoreMLExecutionProvider"], shape_policy=policy) + + def slow(*args: Any, **kwargs: Any) -> Any: # long enough for the other to arrive + time.sleep(0.05) + return mock.DEFAULT + + ort_session.side_effect = slow + + def request(_: int) -> Any: + together.wait() + return session.for_shape(Shape(batch=1, height=64, width=64)) + + with ThreadPoolExecutor(2) as pool: + first, second = pool.map(request, range(2)) + + assert first is second + ort_session.assert_called_once() + + @pytest.mark.parametrize( + ("provider", "expected"), + [ + ("CoreMLExecutionProvider", ["open", "run at 64", "open", "run at 128"]), + ("CPUExecutionProvider", ["run at 64"]), # the one graph, opened when the session is + ], + ) + def test_warms_each_graph_before_opening_the_next( + self, ort_session: mock.Mock, provider: str, expected: list[str] + ) -> None: + policy = ShapePolicy(dims=(Shape(batch=1, height=64, width=64), Shape(batch=1, height=64, width=128))) + events: list[str] = [] + + def open_graph(*args: Any, **kwargs: Any) -> Any: + events.append("open") + return mock.DEFAULT + + def run(output_names: Any, feed: dict[str, np.ndarray]) -> list[np.ndarray]: + events.append(f"run at {feed['image'].shape[2]}") # a dim the graph leaves free takes the shape's size + return [np.zeros(1)] + + ort_session.side_effect = open_graph + ort_session.return_value.get_inputs.return_value = [ + SimpleNamespace(name="image", shape=[1, 64, "width", 3], type="tensor(uint8)") + ] + ort_session.return_value.run.side_effect = run + session = ort_sessions("PP-OCRv5_mobile", providers=[provider], shape_policy=policy) + events.clear() + + session.warm() + + assert events == expected + + def test_opens_a_shape_while_another_is_still_being_built(self, ort_session: mock.Mock) -> None: + policy = ShapePolicy(dims=(Shape(batch=1, height=64, width=64), Shape(batch=1, height=64, width=128))) + session = ort_sessions("PP-OCRv5_mobile", providers=["CoreMLExecutionProvider"], shape_policy=policy) + building, built = threading.Event(), threading.Event() + + def build(*args: Any, **kwargs: Any) -> Any: + building.set() + built.wait(5) + return mock.DEFAULT + + ort_session.side_effect = build + + with ThreadPoolExecutor(1) as pool: + slow = pool.submit(session.for_shape, Shape(batch=1, height=64, width=64)) + building.wait(5) + ort_session.side_effect = None + session.for_shape(Shape(batch=1, height=64, width=128)) + assert not slow.done() + built.set() + + def test_builds_one_session_per_shape_it_is_asked_for(self, ort_session: mock.Mock) -> None: + policy = ShapePolicy(dims=(Shape(batch=1, height=64, width=64), Shape(batch=1, height=64, width=128))) + sessions = ort_sessions("PP-OCRv5_mobile", providers=["CoreMLExecutionProvider"], shape_policy=policy) + assert sessions.shapes == policy.dims # the policy speaks for graphs that do not exist yet + ort_session.assert_not_called() + + for width in (128, 64, 128): + sessions.for_shape(Shape(batch=1, height=64, width=width)) + + assert ort_session.call_count == 2 + caches = [call.kwargs["provider_options"][0]["ModelCacheDirectory"] for call in ort_session.call_args_list] + assert len(set(caches)) == 2 # CoreML keys its compiled model on the path, so variants cannot share one @pytest.mark.ov_device_ids(["GPU.0", "CPU"]) - def test_sets_default_provider_options(self, ov_device_ids: list[str]) -> None: + def test_pins_clip_on_openvino( + self, ort_session: mock.Mock, ov_device_ids: list[str], mocker: MockerFixture + ) -> None: + sess_options = mocker.patch("immich_ml.sessions.ort.ort.SessionOptions").return_value + sess_options.inter_op_num_threads = 0 + + # older Intel GPUs compute garbage for a free dim, and CLIP is the only model that would leave one + ort_sessions("/cache/ViT-B-32__openai/textual/model.onnx", providers=["OpenVINOExecutionProvider"]) + + assert sess_options.add_free_dimension_override_by_name.called + + def test_sets_provider_kwarg(self, ort_session: mock.Mock) -> None: + providers = ["CUDAExecutionProvider"] + ort_sessions("ViT-B-32__openai", providers=providers) + + assert given_providers(ort_session) == providers + + @pytest.mark.ov_device_ids(["GPU.0", "CPU"]) + def test_sets_default_provider_options(self, ort_session: mock.Mock, ov_device_ids: list[str]) -> None: model_path = "/cache/ViT-B-32__openai/textual/model.onnx" - session = OrtSession(model_path, providers=["OpenVINOExecutionProvider", "CPUExecutionProvider"]) + ort_sessions(model_path, providers=["OpenVINOExecutionProvider", "CPUExecutionProvider"]) - assert session.provider_options == [ + assert given_options(ort_session) == [ { "device_type": "GPU.0", + "cache_dir": "/cache/ViT-B-32__openai/textual/openvino/12.71.4-128eu/batch1", "precision": "FP32", - "cache_dir": "/cache/ViT-B-32__openai/textual/openvino", }, {"arena_extend_strategy": "kSameAsRequested"}, ] @pytest.mark.ov_device_ids(["GPU.0", "GPU.1", "CPU"]) - def test_sets_provider_options_for_openvino(self, ov_device_ids: list[str]) -> None: + @pytest.mark.parametrize(("revision", "precision"), [("main", {"precision": "FP32"}), ("v2", {})]) + def test_sets_provider_options_for_openvino( + self, + ort_session: mock.Mock, + ov_device_ids: list[str], + mocker: MockerFixture, + revision: str, + precision: dict[str, str], + ) -> None: + mocker.patch.object(settings, "model_revision", revision) # the older exports are fp32, the others fp16 model_path = "/cache/ViT-B-32__openai/textual/model.onnx" os.environ["MACHINE_LEARNING_DEVICE_ID"] = "1" - session = OrtSession(model_path, providers=["OpenVINOExecutionProvider"]) + ort_sessions(model_path, providers=["OpenVINOExecutionProvider"]) - assert session.provider_options == [ + assert given_options(ort_session) == [ { "device_type": "GPU.1", - "precision": "FP32", - "cache_dir": "/cache/ViT-B-32__openai/textual/openvino", - } - ] - - @pytest.mark.ov_device_ids(["GPU.0", "GPU.1", "CPU"]) - def test_sets_openvino_to_fp16_if_enabled(self, ov_device_ids: list[str], mocker: MockerFixture) -> None: - model_path = "/cache/ViT-B-32__openai/textual/model.onnx" - os.environ["MACHINE_LEARNING_DEVICE_ID"] = "1" - mocker.patch.object(settings, "openvino_precision", ModelPrecision.FP16) - - session = OrtSession(model_path, providers=["OpenVINOExecutionProvider"]) - - assert session.provider_options == [ - { - "device_type": "GPU.1", - "precision": "FP16", - "cache_dir": "/cache/ViT-B-32__openai/textual/openvino", + "cache_dir": "/cache/ViT-B-32__openai/textual/openvino/12.71.4-128eu/batch1", + **precision, } ] @pytest.mark.ov_device_ids(["CPU"]) - def test_sets_provider_options_for_openvino_cpu(self, ov_device_ids: list[str]) -> None: + def test_sets_provider_options_for_openvino_cpu(self, ort_session: mock.Mock, ov_device_ids: list[str]) -> None: model_path = "/cache/ViT-B-32__openai/model.onnx" - session = OrtSession(model_path, providers=["OpenVINOExecutionProvider"]) + ort_sessions(model_path, providers=["OpenVINOExecutionProvider"]) - assert session.provider_options == [ - { - "device_type": "CPU", - "precision": "FP32", - "cache_dir": "/cache/ViT-B-32__openai/openvino", - } + assert given_options(ort_session) == [ + {"device_type": "CPU", "cache_dir": "/cache/ViT-B-32__openai/openvino/batch1", "precision": "FP32"} ] - def test_sets_provider_options_for_cuda(self) -> None: + def test_sets_provider_options_for_cuda(self, ort_session: mock.Mock) -> None: os.environ["MACHINE_LEARNING_DEVICE_ID"] = "1" - session = OrtSession("ViT-B-32__openai", providers=["CUDAExecutionProvider"]) + ort_sessions("ViT-B-32__openai", providers=["CUDAExecutionProvider"]) - assert session.provider_options == [{"arena_extend_strategy": "kSameAsRequested", "device_id": "1"}] + assert given_options(ort_session) == [{"arena_extend_strategy": "kSameAsRequested", "device_id": "1"}] - def test_sets_provider_options_for_rocm(self, mocker: MockerFixture) -> None: + def test_sets_provider_options_for_rocm(self, ort_session: mock.Mock, mocker: MockerFixture) -> None: model_path = "/cache/ViT-B-32__openai/textual/model.onnx" os.environ["MACHINE_LEARNING_DEVICE_ID"] = "1" - mkdir = mocker.patch("immich_ml.sessions.ort.Path.mkdir") - session = OrtSession(model_path, providers=["MIGraphXExecutionProvider"]) + ort_sessions(model_path, providers=["MIGraphXExecutionProvider"]) - assert session.provider_options == [ - { - "device_id": "1", - "migraphx_model_cache_dir": "/cache/ViT-B-32__openai/textual/migraphx", - "migraphx_fp16_enable": "0", - } + assert given_options(ort_session) == [ + {"device_id": "1", "migraphx_model_cache_dir": "/cache/ViT-B-32__openai/textual/migraphx/gfx1100/batch1"} ] - mkdir.assert_called_once_with(parents=True, exist_ok=True) - def test_sets_rocm_to_fp16_if_enabled(self, path: mock.Mock, mocker: MockerFixture) -> None: - model_path = "/cache/ViT-B-32__openai/textual/model.onnx" - os.environ["MACHINE_LEARNING_DEVICE_ID"] = "1" - mocker.patch.object(settings, "rocm_precision", ModelPrecision.FP16) - mkdir = mocker.patch("immich_ml.sessions.ort.Path.mkdir") + @pytest.mark.skipif(sys.platform != "linux" or platform.machine() != "x86_64", reason="an x86 register") + def test_flushes_denormals_on_the_calling_thread(self) -> None: + results: list[float] = [] - session = OrtSession(model_path, providers=["MIGraphXExecutionProvider"]) + def probe() -> None: # a fresh thread, as a request thread is + flush_denormals() + results.append(float((np.array([1e-39], np.float32) * np.float32(1))[0])) - assert session.provider_options == [ - { - "device_id": "1", - "migraphx_model_cache_dir": "/cache/ViT-B-32__openai/textual/migraphx", - "migraphx_fp16_enable": "1", - } + thread = threading.Thread(target=probe) + thread.start() + thread.join() + + assert results == [0.0] + + def test_sets_default_sess_options_if_cpu(self, ort_session: mock.Mock) -> None: + ort_sessions("ViT-B-32__openai", providers=["CPUExecutionProvider"]) + + assert given_sess_options(ort_session).execution_mode == ort.ExecutionMode.ORT_SEQUENTIAL + assert given_sess_options(ort_session).inter_op_num_threads == 1 + assert given_sess_options(ort_session).intra_op_num_threads == 2 + assert given_sess_options(ort_session).get_session_config_entry("session.set_denormal_as_zero") == "1" + + def test_gives_a_model_the_cpu_threads_it_asks_for(self, ort_session: mock.Mock) -> None: + ort_sessions("ViT-B-32__openai", providers=["CPUExecutionProvider"], threads=4) + + assert given_sess_options(ort_session).intra_op_num_threads == 4 + assert OpenClipTextualEncoder.threads == 4 and OpenClipVisualEncoder.threads == 2 + + @pytest.mark.ov_device_ids(["CPU"]) + def test_sets_default_sess_options_if_openvino_cpu(self, ort_session: mock.Mock, ov_device_ids: list[str]) -> None: + model_path = "/cache/ViT-B-32__openai/model.onnx" + ort_sessions(model_path, providers=["OpenVINOExecutionProvider"]) + + assert given_sess_options(ort_session).execution_mode == ort.ExecutionMode.ORT_SEQUENTIAL + assert given_sess_options(ort_session).inter_op_num_threads == 0 + assert given_sess_options(ort_session).intra_op_num_threads == 0 + + @pytest.mark.ov_device_ids(["GPU.0", "CPU"]) + def test_sets_default_sess_options_if_openvino_gpu(self, ort_session: mock.Mock, ov_device_ids: list[str]) -> None: + model_path = "/cache/ViT-B-32__openai/model.onnx" + ort_sessions(model_path, providers=["OpenVINOExecutionProvider"]) + + assert given_sess_options(ort_session).inter_op_num_threads == 0 + assert given_sess_options(ort_session).intra_op_num_threads == 0 + + def test_sets_default_sess_options_does_not_set_threads_if_non_cpu_and_default_threads( + self, ort_session: mock.Mock + ) -> None: + ort_sessions("ViT-B-32__openai", providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) + + assert given_sess_options(ort_session).inter_op_num_threads == 0 + assert given_sess_options(ort_session).intra_op_num_threads == 0 + + def test_sets_default_sess_options_sets_threads_if_non_cpu_and_set_threads( + self, ort_session: mock.Mock, mocker: MockerFixture + ) -> None: + mock_settings = mocker.patch("immich_ml.sessions.ort.settings", autospec=True) + mock_settings.model_inter_op_threads = 2 + mock_settings.model_intra_op_threads = 4 + + ort_sessions("ViT-B-32__openai", providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) + + assert given_sess_options(ort_session).inter_op_num_threads == 2 + assert given_sess_options(ort_session).intra_op_num_threads == 4 + + def test_uses_arena_if_enabled(self, ort_session: mock.Mock, mocker: MockerFixture) -> None: + mock_settings = mocker.patch("immich_ml.sessions.ort.settings", autospec=True) + mock_settings.model_inter_op_threads = 0 + mock_settings.model_intra_op_threads = 0 + mock_settings.model_arena = True + + ort_sessions("ViT-B-32__openai", providers=["CPUExecutionProvider"]) + + assert given_sess_options(ort_session).enable_cpu_mem_arena + + def test_does_not_use_arena_if_disabled(self, ort_session: mock.Mock, mocker: MockerFixture) -> None: + mock_settings = mocker.patch("immich_ml.sessions.ort.settings", autospec=True) + mock_settings.model_inter_op_threads = 0 + mock_settings.model_intra_op_threads = 0 + mock_settings.model_arena = False + + ort_sessions("ViT-B-32__openai", providers=["CPUExecutionProvider"]) + + assert not given_sess_options(ort_session).enable_cpu_mem_arena + + def test_feeds_channels_first_where_the_rewriter_moved_the_layout_out_of_the_graph( + self, ort_session: mock.Mock + ) -> None: + # CoreML's plan retypes the image input, and the models go on handing over frames as they decoded them + ort_session.return_value.get_inputs.return_value = [ + SimpleNamespace(name="image", type="tensor(uint8)", shape=["batch", 3, 224, 224]) ] - mkdir.assert_called_once_with(parents=True, exist_ok=True) + session = ort_sessions("ViT-B-32__openai", providers=["CoreMLExecutionProvider"]).for_shape(Shape(batch=1)) + frame = np.zeros((1, 224, 224, 3), dtype=np.uint8) + frame[0, 5, 7] = (1, 2, 3) - def test_sets_provider_options_kwarg(self) -> None: - session = OrtSession( - "ViT-B-32__openai", - providers=["OpenVINOExecutionProvider", "CPUExecutionProvider"], - provider_options=[], - ) + session.run(None, {"image": frame}) - assert session.provider_options == [] + fed = ort_session.return_value.run.call_args.args[1]["image"] + assert fed.shape == (1, 3, 224, 224) + assert fed[0, :, 5, 7].tolist() == [1, 2, 3] + assert np.shares_memory(fed, frame) # a view: ORT copies it once, as it would have copied the frame + + @pytest.mark.parametrize( + "image", + [ + SimpleNamespace(name="image", type="tensor(uint8)", shape=["batch", 224, 224, 3]), + SimpleNamespace(name="input.1", type="tensor(float)", shape=[1, 3, 224, 224]), + ], + ) + def test_feeds_every_other_graph_what_it_was_handed(self, ort_session: mock.Mock, image: SimpleNamespace) -> None: + ort_session.return_value.get_inputs.return_value = [image] + session = ort_sessions("ViT-B-32__openai", providers=["CPUExecutionProvider"]).for_shape(Shape(batch=1)) + feed = {image.name: np.zeros((1, 3, 224, 224), dtype=np.float32)} + + session.run(None, feed) + + assert ort_session.return_value.run.call_args.args[1] is feed @pytest.mark.parametrize( ("machine", "disabled"), @@ -390,7 +654,7 @@ class TestOrtSession: ) -> None: mocker.patch("immich_ml.sessions.ort.platform.machine", return_value=machine) - OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"]) + ort_sessions("ViT-B-32__openai", providers=["CPUExecutionProvider"]) assert ort_session.call_args.kwargs["disabled_optimizers"] == disabled @@ -407,147 +671,206 @@ class TestOrtSession: ) -> None: mocker.patch("immich_ml.sessions.ort.platform.machine", return_value="x86_64") - OrtSession("ViT-B-32__openai", providers=providers) + ort_sessions("ViT-B-32__openai", providers=providers) assert ort_session.call_args.kwargs["disabled_optimizers"] == disabled - def test_sets_default_sess_options_if_cpu(self) -> None: - session = OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"]) + def test_opens_the_graph_that_was_prepared_for_it(self, ort_session: mock.Mock, mocker: MockerFixture) -> None: + mocker.patch("immich_ml.sessions.ort.prepared", return_value=Path("/cache/visual/cpu/free/model.onnx")) - assert session.sess_options.execution_mode == ort.ExecutionMode.ORT_SEQUENTIAL - assert session.sess_options.inter_op_num_threads == 1 - assert session.sess_options.intra_op_num_threads == 2 - assert session.sess_options.get_session_config_entry("session.set_denormal_as_zero") == "1" + ort_sessions("/cache/visual/model.onnx", providers=["CPUExecutionProvider"]) - @pytest.mark.skipif(sys.platform != "linux" or platform.machine() != "x86_64", reason="an x86 register") - def test_flushes_denormals_on_the_calling_thread(self) -> None: - results: list[float] = [] + assert ort_session.call_args.args == ("/cache/visual/cpu/free/model.onnx",) - def probe() -> None: # a fresh thread, as a request thread is - flush_denormals() - results.append(float((np.array([1e-39], np.float32) * np.float32(1))[0])) - thread = threading.Thread(target=probe) - thread.start() - thread.join() +def linear_graph(directory: Path) -> Path: + import onnx + from onnx import TensorProto, helper, numpy_helper - assert results == [0.0] + weight = numpy_helper.from_array(np.random.default_rng(0).random((64, 64), dtype=np.float32), name="weight") + graph = helper.make_graph( + [helper.make_node("MatMul", ["input", "weight"], ["output"])], + "linear", + [helper.make_tensor_value_info("input", TensorProto.FLOAT, ["batch", 64])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT, ["batch", 64])], + [weight], + ) + directory.mkdir(parents=True, exist_ok=True) + onnx.save( + helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)], ir_version=9), directory / "model.onnx" + ) + return directory / "model.onnx" - def test_gives_a_model_the_cpu_threads_it_asks_for(self) -> None: - session = OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"], threads=4) - assert session.sess_options.intra_op_num_threads == 4 - assert OpenClipTextualEncoder.threads == 4 and OpenClipVisualEncoder.threads == 2 +def graph_spec(model_path: Path, providers: list[str] | None = None, **fields: Any) -> GraphSpec: + fields = {"pins": {}, "overrides": [], "disabled_optimizers": [], **fields} + return GraphSpec(model_path, providers=providers or ["CPUExecutionProvider"], **fields) - @pytest.mark.ov_device_ids(["CPU"]) - def test_sets_default_sess_options_if_openvino_cpu(self, ov_device_ids: list[str]) -> None: - model_path = "/cache/ViT-B-32__openai/model.onnx" - session = OrtSession(model_path, providers=["OpenVINOExecutionProvider"]) - assert session.sess_options.execution_mode == ort.ExecutionMode.ORT_SEQUENTIAL - assert session.sess_options.inter_op_num_threads == 0 - assert session.sess_options.intra_op_num_threads == 0 +class TestPreparedGraphs: + def test_opens_what_a_child_prepared_until_the_facts_change(self, tmp_path: Path, mocker: MockerFixture) -> None: + graph = linear_graph(tmp_path) + + def spec() -> GraphSpec: # a graph is opened once, so what holds while it is cannot change under it + return graph_spec(graph, overrides=[("batch", 2)]) + + def stand_in(*args: Any, **kwargs: Any) -> SimpleNamespace: + prepare_module.prepare(pickle.loads(kwargs["input"])) + return SimpleNamespace(returncode=0) + + child = mocker.patch("immich_ml.sessions.ort.subprocess.run", side_effect=stand_in) + + assert prepared(spec()) == prepared(spec()) == tmp_path / "cpu" / platform.machine() / "model.onnx" + child.assert_called_once() + assert child.call_args.args[0][1:] == ["-m", "immich_ml.sessions.prepare"] + + mocker.patch("immich_ml.sessions.ort.ort.__version__", "9.9.9") # what it packed may not be read back the same + prepared(spec()) + assert child.call_count == 2 + + def test_shares_what_it_prepared_between_workers_on_different_devices(self, monkeypatch: MonkeyPatch) -> None: + facts = [] + for device in ("0", "1"): + monkeypatch.setenv("MACHINE_LEARNING_DEVICE_ID", device) + facts.append(graph_spec(Path("/cache/model.onnx"), ["CUDAExecutionProvider"]).facts) + + assert facts[0] == facts[1] @pytest.mark.ov_device_ids(["GPU.0", "CPU"]) - def test_sets_default_sess_options_if_openvino_gpu(self, ov_device_ids: list[str]) -> None: - model_path = "/cache/ViT-B-32__openai/model.onnx" - session = OrtSession(model_path, providers=["OpenVINOExecutionProvider"]) + @pytest.mark.parametrize( + ("provider", "reader"), + [("OpenVINOExecutionProvider", "_intel_gpu"), ("MIGraphXExecutionProvider", "_amd_gpu")], + ) + def test_prepares_once_per_kind_of_device_and_again_for_another_version( + self, provider: str, reader: str, ov_device_ids: mock.Mock, mocker: MockerFixture + ) -> None: + mocker.patch( + f"immich_ml.sessions.ort.{reader}", side_effect=[Device("a", "1"), Device("a", "2"), Device("b", "1")] + ) + first, updated, other = [graph_spec(Path("/cache/model.onnx"), [provider]) for _ in range(3)] - assert session.sess_options.inter_op_num_threads == 0 - assert session.sess_options.intra_op_num_threads == 0 + assert first.directory == updated.directory and first.facts != updated.facts + assert first.directory != other.directory - def test_sets_default_sess_options_does_not_set_threads_if_non_cpu_and_default_threads(self) -> None: - session = OrtSession("ViT-B-32__openai", providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) + @pytest.mark.parametrize(("returncode", "error"), [(3, InvalidProtobuf), (1, RuntimeError)]) + def test_tells_an_unreadable_source_from_a_failed_preparation( + self, tmp_path: Path, mocker: MockerFixture, returncode: int, error: type[Exception] + ) -> None: + mocker.patch("immich_ml.sessions.ort.subprocess.run", return_value=SimpleNamespace(returncode=returncode)) - assert session.sess_options.inter_op_num_threads == 0 - assert session.sess_options.intra_op_num_threads == 0 + with pytest.raises(error): # the first has the download cleared and fetched again + prepared(graph_spec(linear_graph(tmp_path))) - def test_sets_default_sess_options_sets_threads_if_non_cpu_and_set_threads(self, mocker: MockerFixture) -> None: - mock_settings = mocker.patch("immich_ml.sessions.ort.settings", autospec=True) - mock_settings.model_inter_op_threads = 2 - mock_settings.model_intra_op_threads = 4 - - session = OrtSession("ViT-B-32__openai", providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) - - assert session.sess_options.inter_op_num_threads == 2 - assert session.sess_options.intra_op_num_threads == 4 - - def test_uses_arena_if_enabled(self, mocker: MockerFixture) -> None: - mock_settings = mocker.patch("immich_ml.sessions.ort.settings", autospec=True) - mock_settings.model_inter_op_threads = 0 - mock_settings.model_intra_op_threads = 0 - mock_settings.model_arena = True - - session = OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"]) - - assert session.sess_options.enable_cpu_mem_arena - - def test_does_not_use_arena_if_disabled(self, mocker: MockerFixture) -> None: - mock_settings = mocker.patch("immich_ml.sessions.ort.settings", autospec=True) - mock_settings.model_inter_op_threads = 0 - mock_settings.model_intra_op_threads = 0 - mock_settings.model_arena = False - - session = OrtSession("ViT-B-32__openai", providers=["CPUExecutionProvider"]) - - assert not session.sess_options.enable_cpu_mem_arena - - def test_sets_sess_options_kwarg(self) -> None: - sess_options = ort.SessionOptions() - session = OrtSession( - "ViT-B-32__openai", - providers=["OpenVINOExecutionProvider", "CPUExecutionProvider"], - provider_options=[], - sess_options=sess_options, + def test_the_child_reports_an_unreadable_source(self, tmp_path: Path, mocker: MockerFixture) -> None: + tmp_path.joinpath("model.onnx").write_bytes(b"not a graph") + request = pickle.dumps(graph_spec(tmp_path / "model.onnx")) + mocker.patch( + "immich_ml.sessions.prepare.sys.stdin", SimpleNamespace(buffer=SimpleNamespace(read=lambda: request)) ) - assert sess_options is session.sess_options + with pytest.raises(SystemExit) as left: + prepare_module.main() - def test_serializes_rocm_first_run_for_new_input_signature(self, mocker: MockerFixture) -> None: - lock = FakeLock() - get_model_lock = mocker.patch("immich_ml.sessions.ort._migraphx_get_model_lock", return_value=lock) - mocker.patch("immich_ml.sessions.ort._migraphx_compiled_inputs", set()) - mocker.patch("immich_ml.sessions.ort.Path.mkdir") - session = OrtSession("/cache/ViT-B-32__openai/model.onnx", providers=["MIGraphXExecutionProvider"]) - input_feed = {"input": np.random.rand(1, 3, 224, 224).astype(np.float32)} + assert left.value.code == 3 - session.run(None, input_feed) - session.run(None, input_feed) + def test_packs_the_weights_into_an_artifact_that_gives_the_same_answer(self, tmp_path: Path) -> None: + spec = graph_spec(linear_graph(tmp_path)) + feed = {"input": np.random.default_rng(1).random((2, 64), dtype=np.float32)} + expected = ort.InferenceSession(spec.model_path.as_posix(), providers=spec.providers).run(None, feed) + spec.directory.mkdir(parents=True) + spec.directory.joinpath("left-by-a-child-that-did-not-finish").touch() - lock.enter.assert_called_once() - lock.exit.assert_called_once() - get_model_lock.assert_called_once() - session.session.run.assert_has_calls([mock.call(None, input_feed, None), mock.call(None, input_feed, None)]) + prepare_module.prepare(spec) - def test_serializes_rocm_run_for_each_new_input_signature(self, mocker: MockerFixture) -> None: - lock = FakeLock() - mocker.patch("immich_ml.sessions.ort._migraphx_get_model_lock", return_value=lock) - mocker.patch("immich_ml.sessions.ort._migraphx_compiled_inputs", set()) - mocker.patch("immich_ml.sessions.ort.Path.mkdir") - session = OrtSession("/cache/ViT-B-32__openai/model.onnx", providers=["MIGraphXExecutionProvider"]) - input_feed = {"input": np.random.rand(1, 3, 224, 224).astype(np.float32)} - new_shape_input_feed = {"input": np.random.rand(2, 3, 224, 224).astype(np.float32)} - - session.run(None, input_feed) - session.run(None, new_shape_input_feed) - - assert lock.enter.call_count == 2 - assert lock.exit.call_count == 2 - session.session.run.assert_has_calls( - [mock.call(None, input_feed, None), mock.call(None, new_shape_input_feed, None)] + assert sorted(file.name for file in spec.directory.iterdir()) == ["manifest.json", "model.data", "model.onnx"] + np.testing.assert_array_equal( + spec.session(tmp_path / "cpu" / platform.machine() / "model.onnx").run(None, feed)[0], expected[0] ) - def test_does_not_serialize_non_rocm_run(self, mocker: MockerFixture) -> None: - lock = FakeLock() - get_model_lock = mocker.patch("immich_ml.sessions.ort._migraphx_get_model_lock", return_value=lock) - session = OrtSession("/cache/ViT-B-32__openai/model.onnx", providers=["CPUExecutionProvider"]) - input_feed = {"input": np.random.rand(1, 3, 224, 224).astype(np.float32)} + def test_compiles_where_the_provider_keeps_the_result(self, tmp_path: Path, mocker: MockerFixture) -> None: + session = mocker.patch("immich_ml.sessions.ort.ort.InferenceSession") + spec = graph_spec(linear_graph(tmp_path), ["CoreMLExecutionProvider", "CPUExecutionProvider"]) - session.run(None, input_feed) + prepare_module.prepare(spec) - get_model_lock.assert_not_called() - lock.enter.assert_not_called() - session.session.run.assert_called_once_with(None, input_feed, None) + assert session.call_args.args == (spec.model_path.as_posix(),) + assert ( + session.call_args.kwargs["provider_options"][0]["ModelCacheDirectory"] == (tmp_path / "coreml").as_posix() + ) + assert fresh(spec) == spec.model_path + + def test_rewrites_only_the_exports_the_rewriter_was_written_for( + self, tmp_path: Path, mocker: MockerFixture + ) -> None: + apply = mocker.patch("immich_ml.sessions.prepare.apply_rewrites") + apply.side_effect = lambda source, plan: source.with_name("model.rw-abc.onnx") + spec = graph_spec(linear_graph(tmp_path), ["CUDAExecutionProvider"]) + + assert prepare_module.rewritten(spec, spec.model_path) == spec.model_path + mocker.patch.object(settings, "model_revision", "v2") + assert prepare_module.rewritten(spec, spec.model_path) == tmp_path / "model.rw-abc.onnx" + assert apply.call_args.args[1] is spec.plan + + @pytest.mark.parametrize( + ("providers", "revision", "half"), + [ + (["CUDAExecutionProvider", "CPUExecutionProvider"], "v2", True), + (["CPUExecutionProvider"], "v2", False), + (["CUDAExecutionProvider", "CPUExecutionProvider"], "main", False), + ], + ) + def test_narrows_to_half_precision_for_an_accelerator( + self, tmp_path: Path, mocker: MockerFixture, providers: list[str], revision: str, half: bool + ) -> None: + mocker.patch.object(settings, "model_revision", revision) + derive = mocker.patch("immich_ml.sessions.prepare.derive") + spec = graph_spec(tmp_path / "model.onnx", providers) + + narrowed = prepare_module.narrowed(spec) + + assert spec.facts.half is half + assert narrowed == (spec.directory / "model_fp16.onnx" if half else spec.model_path) + assert derive.called is half + + def test_moves_skewed_regions_onto_the_boundary_without_losing_a_byte(self, tmp_path: Path) -> None: + import onnx + from onnx import TensorProto, helper, numpy_helper + from onnx.external_data_helper import set_external_data + + # an odd-length region first, which is what skews every region after it + arrays: dict[str, np.ndarray[Any, Any]] = { + "odd": np.arange(7, dtype=np.uint8), + "weight": np.arange(256, dtype=np.float32), + } + packed = bytes(range(200)) # the form MLAS packed `weight` into, which ORT files under the weight + data, tensors = b"", [] + for name, array in arrays.items(): + tensor = numpy_helper.from_array(array, name=name) + set_external_data(tensor, location="model.data", offset=len(data), length=array.nbytes) + tensor.ClearField("raw_data") + tensors.append(tensor) + data += array.tobytes() + tensors[1].external_data.add(key="prepacked_0", value=f"MatMul+hash|{len(data)};{len(packed)};checksum") + (tmp_path / "model.data").write_bytes(data + packed) + graph = helper.make_graph( + [helper.make_node("Identity", ["weight"], ["output"])], + "skewed", + [], + [helper.make_tensor_value_info("output", TensorProto.FLOAT, [256])], + tensors, + ) + onnx.save(helper.make_model(graph), (tmp_path / "model.onnx").as_posix()) + + prepare_module.align(tmp_path / "model.onnx") + + moved = onnx.load((tmp_path / "model.onnx").as_posix(), load_external_data=False) + entries = [{entry.key: entry.value for entry in tensor.external_data} for tensor in moved.graph.initializer] + assert [entry["offset"] for entry in entries] == ["0", "64"] # 0 and 7 before + assert entries[1]["prepacked_0"] == "MatMul+hash|1088;200;checksum" # at 1031 before + written = (tmp_path / "model.data").read_bytes() + assert written[0:7] == arrays["odd"].tobytes() + assert written[64 : 64 + 1024] == arrays["weight"].tobytes() + assert written[1088:] == packed and not any(written[7:64]) class TestAnnSession: @@ -618,7 +941,6 @@ class TestRknnSession: def test_run_rknn(self, rknn_session: mock.Mock, mocker: MockerFixture) -> None: rknn_session.return_value.load.return_value = 123 - np_spy = mocker.spy(np, "ascontiguousarray") mocker.patch("immich_ml.sessions.rknn.soc_name", "rk3566") session = RknnSession(Path("ViT-B-32__openai")) [input1, input2] = [np.random.rand(1, 3, 224, 224).astype(np.float32) for _ in range(2)] @@ -626,9 +948,42 @@ class TestRknnSession: session.run(None, input_feed) - rknn_session.return_value.run.assert_called_once_with([input1, input2]) - assert np_spy.call_count == 2 - np_spy.assert_has_calls([mock.call(input1), mock.call(input2)]) + rknn_session.return_value.run.assert_called_once_with([input1, input2], "nchw") + assert all(fed.flags.c_contiguous for fed in rknn_session.return_value.run.call_args.args[0]) + + def test_run_rknn_rejects_a_shape_the_binary_was_not_compiled_for( + self, rknn_session: mock.Mock, mocker: MockerFixture + ) -> None: + mocker.patch("immich_ml.sessions.rknn.soc_name", "rk3566") + session = RknnSession(Path("ViT-B-32__openai")) + + # a wrong data_format label is not something librknnrt reports; it reinterprets the buffer + with pytest.raises(ValueError, match="takes 3 channels"): + session.run(None, {"input.1": np.zeros((1, 5, 112, 112), dtype=np.float32)}) + + def test_shapes_come_from_the_binary(self, rknn_session: mock.Mock) -> None: + rknn_session.return_value.custom_string = '{"dims":[{"height":736,"width":1472},{"height":736,"width":736}]}' + + session = RknnSession(Path("PP-OCRv5_mobile")) + + assert session.shapes == ( + Shape(batch=1, height=736, width=1472), + Shape(batch=1, height=736, width=736), + ) + + def test_names_the_token_table_the_binary_left_to_the_host(self, rknn_session: mock.Mock) -> None: + rknn_session.return_value.custom_string = '{"dims":[{}],"embedding":"token_embedding.weight_fp16"}' + + session = RknnSession(Path("ViT-B-32__openai")) + + assert session.get_metadata() == {"embedding": "token_embedding.weight_fp16"} + assert session.shapes == (Shape(batch=1),) + + def test_offers_no_choice_of_shape_when_the_binary_has_one(self, rknn_session: mock.Mock) -> None: + # a graph the exporter left no dim free in is stamped with one empty set + rknn_session.return_value.custom_string = '{"dims":[{}]}' + + assert RknnSession(Path("buffalo_l")).shapes == (Shape(batch=1),) class TestCLIP: @@ -647,6 +1002,7 @@ class TestCLIP: mocker.patch.object(OpenClipVisualEncoder, "preprocess_cfg", clip_preprocess_cfg) mocked = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + mocked.for_shape.return_value = mocked mocked.run.return_value = [[self.embedding]] clip_encoder = OpenClipVisualEncoder("ViT-B-32__openai", cache_dir="test_cache") @@ -657,21 +1013,52 @@ class TestCLIP: assert len(embedding) == clip_model_cfg["embed_dim"] mocked.run.assert_called_once() + @pytest.mark.parametrize( + ("normalizes_input", "expected_dtype", "expected_shape"), + [(False, np.float32, (1, 3, 224, 224)), (True, np.uint8, (1, 224, 224, 3))], + ) + def test_visual_feeds_the_contract_the_session_declares( + self, + normalizes_input: bool, + expected_dtype: type[np.generic], + expected_shape: tuple[int, ...], + pil_image: Image.Image, + mocker: MockerFixture, + stub_session: Callable[..., mock.Mock], + clip_model_cfg: dict[str, Any], + clip_preprocess_cfg: dict[str, Any], + ) -> None: + mocker.patch.object(OpenClipVisualEncoder, "download") + mocker.patch.object(OpenClipVisualEncoder, "model_cfg", clip_model_cfg) + mocker.patch.object(OpenClipVisualEncoder, "preprocess_cfg", clip_preprocess_cfg) + session = stub_session(expected_shape, outputs=[[self.embedding]], normalizes_input=normalizes_input) + mocker.patch.object(InferenceModel, "_make_session", return_value=session) + + OpenClipVisualEncoder("ViT-B-32__openai", cache_dir="test_cache").predict(pil_image) + + fed = session.run.call_args.args[1]["image"] + assert fed.dtype == expected_dtype + assert fed.shape == expected_shape + def test_visual_squashes_when_the_tower_was_calibrated_that_way( - self, mocker: MockerFixture, clip_model_cfg: dict[str, Any], clip_preprocess_cfg: dict[str, Any] + self, + mocker: MockerFixture, + stub_session: Callable[..., mock.Mock], + clip_model_cfg: dict[str, Any], + clip_preprocess_cfg: dict[str, Any], ) -> None: mocker.patch.object(OpenClipVisualEncoder, "download") mocker.patch.object(OpenClipVisualEncoder, "model_cfg", clip_model_cfg) mocker.patch.object(OpenClipVisualEncoder, "preprocess_cfg", clip_preprocess_cfg | {"resize_mode": "squash"}) - session = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value - session.run.return_value = [[self.embedding]] + session = stub_session((1, 224, 224, 3), outputs=[[self.embedding]], normalizes_input=True) + mocker.patch.object(InferenceModel, "_make_session", return_value=session) # a stripe far enough left that a shortest-side crop of this frame would discard it image = Image.new("RGB", (600, 200), "black") image.paste(Image.new("RGB", (20, 200), "white"), (0, 0)) OpenClipVisualEncoder("ViT-B-32__openai", cache_dir="test_cache").predict(image) - assert (session.run.call_args.args[1]["image"][0, :, 0, 0] > 0).all() # white, where black normalizes below 0 + assert session.run.call_args.args[1]["image"][0, 0, 0].tolist() == [255, 255, 255] def test_basic_text( self, @@ -684,6 +1071,9 @@ class TestCLIP: mocker.patch.object(OpenClipTextualEncoder, "tokenizer_cfg", clip_tokenizer_cfg) mocked = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + mocked.for_shape.return_value = mocked + mocked.get_metadata.return_value = {} + mocked.get_inputs.return_value = [SimpleNamespace(name="text")] mocked.run.return_value = [[self.embedding]] mocker.patch("immich_ml.models.clip.textual.Tokenizer.from_file", autospec=True) @@ -695,6 +1085,45 @@ class TestCLIP: assert len(embedding) == clip_model_cfg["embed_dim"] mocked.run.assert_called_once() + def test_feeds_the_token_rows_a_graph_left_to_the_host( + self, + mocker: MockerFixture, + tmp_path: Path, + clip_model_cfg: dict[str, Any], + clip_tokenizer_cfg: Callable[[Path], dict[str, Any]], + ) -> None: + import onnx + from onnx.external_data_helper import set_external_data + + table = np.arange(40, dtype=np.float16).reshape(10, 4) + (tmp_path / "textual").mkdir() + (tmp_path / "textual/model.safetensors").write_bytes(b"\0" * 16 + table.tobytes()) # the table past a header + weights = onnx.numpy_helper.from_array(table, "tok") + set_external_data(weights, "model.safetensors", offset=16, length=table.nbytes) + weights.ClearField("raw_data") # only the file holds it, as in an export + onnx.save( + onnx.helper.make_model(onnx.helper.make_graph([], "textual", [], [], [weights])), + tmp_path / "textual/model.onnx", + ) + mocker.patch.object(OpenClipTextualEncoder, "download") + mocker.patch.object(OpenClipTextualEncoder, "model_cfg", clip_model_cfg) + mocker.patch.object(OpenClipTextualEncoder, "tokenizer_cfg", clip_tokenizer_cfg) + mocker.patch("immich_ml.models.clip.textual.Tokenizer.from_file", autospec=True) + session = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session.for_shape.return_value = session + session.get_metadata.return_value = {"embedding": "tok"} + session.get_inputs.return_value = [SimpleNamespace(name="text"), SimpleNamespace(name="token_embeds")] + session.run.return_value = [[self.embedding]] + encoder = OpenClipTextualEncoder("ViT-B-32__openai", cache_dir=tmp_path) + mocker.patch.object(encoder, "tokenize", return_value={"text": np.array([[1, 3, 2]], np.int32)}) + + encoder.predict("test search query") + + feed = session.run.call_args.args[1] + assert list(feed) == ["text", "token_embeds"] # a binary reads its inputs in order + assert feed["token_embeds"].dtype == np.float16 + assert feed["token_embeds"].tolist() == [[[4, 5, 6, 7], [12, 13, 14, 15], [8, 9, 10, 11]]] + def test_reads_model_configs_as_utf8(self, mocker: MockerFixture, tmp_path: Path) -> None: original_open = Path.open @@ -736,7 +1165,8 @@ class TestCLIP: mocker.patch.object(OpenClipTextualEncoder, "download") mocker.patch.object(OpenClipTextualEncoder, "model_cfg", clip_model_cfg) mocker.patch.object(OpenClipTextualEncoder, "tokenizer_cfg", clip_tokenizer_cfg) - mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session.for_shape.return_value.get_metadata.return_value = {} # no table left to the host mock_tokenizer = mocker.patch("immich_ml.models.clip.textual.Tokenizer.from_file", autospec=True).return_value mock_ids = [randint(0, 50000) for _ in range(77)] mock_tokenizer.encode.return_value = SimpleNamespace(ids=mock_ids) @@ -762,7 +1192,8 @@ class TestCLIP: mocker.patch.object(OpenClipTextualEncoder, "download") mocker.patch.object(OpenClipTextualEncoder, "model_cfg", clip_model_cfg) mocker.patch.object(OpenClipTextualEncoder, "tokenizer_cfg", clip_tokenizer_cfg) - mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session.for_shape.return_value.get_metadata.return_value = {} # no table left to the host mock_tokenizer = mocker.patch("immich_ml.models.clip.textual.Tokenizer.from_file", autospec=True).return_value mock_ids = [randint(0, 50000) for _ in range(77)] mock_tokenizer.encode.return_value = SimpleNamespace(ids=mock_ids) @@ -787,7 +1218,8 @@ class TestCLIP: mocker.patch.object(OpenClipTextualEncoder, "download") mocker.patch.object(OpenClipTextualEncoder, "model_cfg", clip_model_cfg) mocker.patch.object(OpenClipTextualEncoder, "tokenizer_cfg", clip_tokenizer_cfg) - mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session.for_shape.return_value.get_metadata.return_value = {} # no table left to the host mock_tokenizer = mocker.patch("immich_ml.models.clip.textual.Tokenizer.from_file", autospec=True).return_value mock_ids = [randint(0, 50000) for _ in range(77)] mock_tokenizer.encode.return_value = SimpleNamespace(ids=mock_ids) @@ -807,7 +1239,8 @@ class TestCLIP: mocker.patch.object(OpenClipTextualEncoder, "download") mocker.patch.object(OpenClipTextualEncoder, "model_cfg", clip_model_cfg) mocker.patch.object(OpenClipTextualEncoder, "tokenizer_cfg", clip_tokenizer_cfg) - mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session.for_shape.return_value.get_metadata.return_value = {} # no table left to the host mock_tokenizer = mocker.patch("immich_ml.models.clip.textual.Tokenizer.from_file", autospec=True).return_value mock_ids = [randint(0, 50000) for _ in range(77)] mock_tokenizer.encode.return_value = SimpleNamespace(ids=mock_ids) @@ -828,7 +1261,8 @@ class TestCLIP: mocker.patch.object(OpenClipTextualEncoder, "download") mocker.patch.object(OpenClipTextualEncoder, "model_cfg", clip_model_cfg) mocker.patch.object(OpenClipTextualEncoder, "tokenizer_cfg", clip_tokenizer_cfg) - mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session.for_shape.return_value.get_metadata.return_value = {} # no table left to the host mock_tokenizer = mocker.patch("immich_ml.models.clip.textual.Tokenizer.from_file", autospec=True).return_value mock_ids = [randint(0, 50000) for _ in range(77)] mock_tokenizer.encode.return_value = SimpleNamespace(ids=mock_ids) @@ -849,7 +1283,8 @@ class TestCLIP: mocker.patch.object(OpenClipTextualEncoder, "download") mocker.patch.object(OpenClipTextualEncoder, "model_cfg", clip_model_cfg) mocker.patch.object(OpenClipTextualEncoder, "tokenizer_cfg", clip_tokenizer_cfg) - mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session.for_shape.return_value.get_metadata.return_value = {} # no table left to the host mock_tokenizer = mocker.patch("immich_ml.models.clip.textual.Tokenizer.from_file", autospec=True).return_value mock_ids = [randint(0, 50000) for _ in range(77)] mock_tokenizer.encode.return_value = SimpleNamespace(ids=mock_ids) @@ -869,7 +1304,8 @@ class TestCLIP: mocker.patch.object(MClipTextualEncoder, "download") mocker.patch.object(MClipTextualEncoder, "model_cfg", clip_model_cfg) mocker.patch.object(MClipTextualEncoder, "tokenizer_cfg", clip_tokenizer_cfg) - mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session = mocker.patch.object(InferenceModel, "_make_session", autospec=True).return_value + session.for_shape.return_value.get_metadata.return_value = {} # no table left to the host mock_tokenizer = mocker.patch("immich_ml.models.clip.textual.Tokenizer.from_file", autospec=True).return_value mock_ids = [randint(0, 50000) for _ in range(77)] mock_attention_mask = [randint(0, 1) for _ in range(77)] @@ -916,30 +1352,52 @@ def expected_box(cell_x: int, cell_y: int) -> list[float]: return [cx - 8, cy - 16, cx + 24, cy + 32] +M = TypeVar("M", bound=InferenceModel) + + +def given_options(ort_session: mock.Mock) -> Any: + return ort_session.call_args.kwargs["provider_options"] + + +def given_sess_options(ort_session: mock.Mock) -> Any: + return ort_session.call_args.kwargs["sess_options"] + + +def given_providers(ort_session: mock.Mock) -> Any: + return ort_session.call_args.kwargs["providers"] + + +def loaded(model: M, session: mock.Mock, mocker: MockerFixture) -> M: + mocker.patch.object(model, "download") + mocker.patch.object(model, "_make_session", return_value=session) + model.load() + return model + + +def ort_sessions( + model_path: str, + providers: list[str] | None = None, + shape_policy: ShapePolicy = ShapePolicy(), + threads: int = 2, +) -> OrtSession: + return OrtSession(model_path, shape_policy, providers=providers, threads=threads) + + def expected_landmarks(cell_x: int, cell_y: int) -> np.ndarray: cx, cy = cell_x * 8, cell_y * 8 return (np.tile([cx, cy], 5) + np.arange(10) * 8).reshape(5, 2).astype(np.float32) -@pytest.fixture -def batch_axis(mocker: MockerFixture) -> SimpleNamespace: - """Patches for the `_add_batch_axis` path: the onnx module it rewrites the model - through, and the download it would otherwise trigger on load.""" - mocker.patch("immich_ml.models.base.InferenceModel.download") - return SimpleNamespace( - onnx=mocker.patch("immich_ml.models.facial_recognition.recognition.onnx", autospec=True), - update_dims=mocker.patch( - "immich_ml.models.facial_recognition.recognition.update_inputs_outputs_dims", autospec=True - ), - ) - - class TestFaceRecognition: - def test_detection(self, stub_session: Callable[..., mock.Mock], mocker: MockerFixture) -> None: + @pytest.mark.parametrize("dtype", [np.float32, np.float16]) # a graph narrowed to half answers in it + def test_detection( + self, stub_session: Callable[..., mock.Mock], mocker: MockerFixture, dtype: type[np.generic] + ) -> None: mocker.patch.object(FaceDetector, "load") face_detector = FaceDetector("buffalo_s", cache_dir="test_cache") - session = stub_session((1, 3, 640, 640), outputs=make_scrfd_heads([(10, 10, 0.9), (50, 50, 0.8)])) + heads = [head.astype(dtype) for head in make_scrfd_heads([(10, 10, 0.9), (50, 50, 0.8)])] + session = stub_session((1, 3, 640, 640), outputs=heads) face_detector.session = session faces = face_detector.predict(Image.new("RGB", (640, 640)), minScore=0.7) @@ -948,7 +1406,7 @@ class TestFaceRecognition: assert set(faces) == {"boxes", "scores", "landmarks"} # NMS returns highest score first assert faces["boxes"].tolist() == [expected_box(10, 10), expected_box(50, 50)] - assert np.allclose(faces["scores"], [0.9, 0.8]) + assert np.allclose(faces["scores"], [0.9, 0.8], atol=1e-3) assert faces["landmarks"].shape == (2, 5, 2) assert np.allclose(faces["landmarks"][0], expected_landmarks(10, 10)) assert np.allclose(faces["landmarks"][1], expected_landmarks(50, 50)) @@ -983,10 +1441,6 @@ class TestFaceRecognition: def test_recognition(self, stub_session: Callable[..., mock.Mock], mocker: MockerFixture) -> None: mocker.patch.object(FaceRecognizer, "load") - mocker.patch( - "immich_ml.models.facial_recognition.recognition.ort.get_available_providers", - return_value=["CPUExecutionProvider"], - ) face_recognizer = FaceRecognizer("buffalo_s", cache_dir="test_cache") # a uniform grey image whose crops land wholly inside it, so every sampled @@ -1002,7 +1456,7 @@ class TestFaceRecognition: scores = np.array([0.67] * num_faces).astype(np.float32) embeddings = np.random.rand(num_faces, 512).astype(np.float32) - session = stub_session(("batch", 3, 112, 112), outputs=[embeddings]) + session = stub_session(("batch", 3, 112, 112), outputs=[embeddings], shapes=(Shape(batch=num_faces),)) face_recognizer.session = session faces = face_recognizer.predict(image, {"boxes": bbox, "landmarks": kpss, "scores": scores}) @@ -1025,12 +1479,67 @@ class TestFaceRecognition: # in normalize's scale-then-subtract, but still rejects a wrong mean or std assert np.allclose(crops, (128 - 127.5) / 127.5, atol=1e-6) + @pytest.mark.parametrize( + ("normalizes_input", "expected_dtype", "expected_shape"), + [(False, np.float32, (1, 3, 640, 640)), (True, np.uint8, (1, 640, 640, 3))], + ) + def test_detection_feeds_the_contract_the_session_declares( + self, + normalizes_input: bool, + expected_dtype: type[np.generic], + expected_shape: tuple[int, ...], + stub_session: Callable[..., mock.Mock], + mocker: MockerFixture, + ) -> None: + mocker.patch.object(FaceDetector, "load") + face_detector = FaceDetector("buffalo_s", cache_dir="test_cache") + face_detector.session = stub_session( + expected_shape, outputs=make_scrfd_heads([(10, 10, 0.9)]), normalizes_input=normalizes_input + ) + + faces = face_detector.predict(Image.new("RGB", (640, 640)), minScore=0.7) + + fed = face_detector.session.run.call_args.args[1]["input.1"] + assert fed.dtype == expected_dtype + assert fed.shape == expected_shape + # whichever contract it fed, the decode is the same and lands on the same box + assert faces["boxes"].tolist() == [expected_box(10, 10)] + + @pytest.mark.parametrize( + ("normalizes_input", "expected_dtype", "expected_shape"), + [(False, np.float32, (2, 3, 112, 112)), (True, np.uint8, (2, 112, 112, 3))], + ) + def test_recognition_feeds_the_contract_the_session_declares( + self, + normalizes_input: bool, + expected_dtype: type[np.generic], + expected_shape: tuple[int, ...], + stub_session: Callable[..., mock.Mock], + mocker: MockerFixture, + ) -> None: + mocker.patch.object(FaceRecognizer, "load") + face_recognizer = FaceRecognizer("buffalo_s", cache_dir="test_cache") + face_recognizer.session = stub_session( + expected_shape, + outputs=[np.zeros((2, 512), dtype=np.float32)], + normalizes_input=normalizes_input, + shapes=(Shape(batch=2),), + ) + faces: Any = { + "boxes": np.array([[0, 0, 60, 60], [60, 60, 120, 120]], dtype=np.float32), + "scores": np.array([0.9, 0.8], dtype=np.float32), + "landmarks": np.stack([expected_landmarks(1, 1), expected_landmarks(8, 8)]), + } + + face_recognizer.predict(Image.new("RGB", (200, 200)), faces) + + fed = face_recognizer.session.run.call_args.args[1]["input.1"] + assert fed.dtype == expected_dtype + assert fed.shape == expected_shape + assert fed.flags.c_contiguous + def test_recognition_returns_early_without_faces(self, pil_image: Image.Image, mocker: MockerFixture) -> None: mocker.patch.object(FaceRecognizer, "load") - mocker.patch( - "immich_ml.models.facial_recognition.recognition.ort.get_available_providers", - return_value=["CPUExecutionProvider"], - ) face_recognizer = FaceRecognizer("buffalo_s", cache_dir="test_cache") session = mock.Mock() face_recognizer.session = session @@ -1048,15 +1557,10 @@ class TestFaceRecognition: self, pil_image: Image.Image, stub_session: Callable[..., mock.Mock], mocker: MockerFixture ) -> None: mocker.patch.object(FaceRecognizer, "load") - mocker.patch( - "immich_ml.models.facial_recognition.recognition.ort.get_available_providers", - return_value=["CPUExecutionProvider"], - ) face_recognizer = FaceRecognizer("buffalo_s", cache_dir="test_cache") - face_recognizer.batch_size = 2 num_faces = 5 - session = stub_session((1, 3, 112, 112)) + session = stub_session((1, 3, 112, 112), shapes=(Shape(batch=1), Shape(batch=2))) session.run.side_effect = lambda _, feed: [np.zeros((feed["input.1"].shape[0], 512), dtype=np.float32)] face_recognizer.session = session @@ -1069,133 +1573,69 @@ class TestFaceRecognition: assert session.run.call_count == 3 # 2 + 2 + 1 assert [c.args[1]["input.1"].shape[0] for c in session.run.call_args_list] == [2, 2, 1] - def test_recognition_adds_batch_axis_for_ort( - self, batch_axis: SimpleNamespace, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture - ) -> None: - mocker.patch( - "immich_ml.models.facial_recognition.recognition.ort.get_available_providers", - return_value=["CPUExecutionProvider"], - ) - ort_session.return_value.get_inputs.return_value = [SimpleNamespace(name="input.1", shape=(1, 3, 224, 224))] - ort_session.return_value.get_outputs.return_value = [SimpleNamespace(name="output.1", shape=(1, 800))] - path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx" - - proto = mock.Mock() - - input_dims = mock.Mock() - input_dims.name = "input.1" - input_dims.type.tensor_type.shape.dim = [SimpleNamespace(dim_value=size) for size in [1, 3, 224, 224]] - proto.graph.input = [input_dims] - - output_dims = mock.Mock() - output_dims.name = "output.1" - output_dims.type.tensor_type.shape.dim = [SimpleNamespace(dim_value=size) for size in [1, 800]] - proto.graph.output = [output_dims] - - batch_axis.onnx.load.return_value = proto - - face_recognizer = FaceRecognizer("buffalo_s", cache_dir=path) - face_recognizer.load() - - assert face_recognizer.batch_size is None - batch_axis.update_dims.assert_called_once_with( - proto, {"input.1": ["batch", 3, 224, 224]}, {"output.1": ["batch", 800]} - ) - batch_axis.onnx.save.assert_called_once_with(batch_axis.update_dims.return_value, face_recognizer.model_path) - - def test_recognition_does_not_add_batch_axis_if_exists( - self, batch_axis: SimpleNamespace, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture - ) -> None: - mocker.patch( - "immich_ml.models.facial_recognition.recognition.ort.get_available_providers", - return_value=["CPUExecutionProvider"], - ) - path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx" - - inputs = [SimpleNamespace(name="input.1", shape=("batch", 3, 224, 224))] - outputs = [SimpleNamespace(name="output.1", shape=("batch", 800))] - ort_session.return_value.get_inputs.return_value = inputs - ort_session.return_value.get_outputs.return_value = outputs - - face_recognizer = FaceRecognizer("buffalo_s", cache_dir=path) - face_recognizer.load() - - # batching is available here, so the axis is only skipped because it already exists - assert face_recognizer.batch_size is None - batch_axis.update_dims.assert_not_called() - batch_axis.onnx.load.assert_not_called() - batch_axis.onnx.save.assert_not_called() - @pytest.mark.parametrize( - ("session_fixture", "suffix", "model_kwargs", "available_providers"), + ("provider", "shapes"), [ - pytest.param("ann_session", ".armnn", {"model_format": ModelFormat.ARMNN}, None, id="armnn"), - pytest.param( - "ort_session", ".onnx", {}, ["OpenVINOExecutionProvider", "CPUExecutionProvider"], id="openvino" - ), - pytest.param( - "ort_session", ".onnx", {}, ["MIGraphXExecutionProvider", "CPUExecutionProvider"], id="migraphx" - ), + ("CPUExecutionProvider", (Shape(batch=1),)), # a row at a time + ("CUDAExecutionProvider", (Shape(batch=1), Shape(batch=4))), # the configured size, and one for the rest ], ) - def test_recognition_does_not_add_batch_axis_when_batching_is_unsupported( - self, - request: pytest.FixtureRequest, - batch_axis: SimpleNamespace, - path: mock.Mock, - ort_pybind: mock.Mock, - mocker: MockerFixture, - session_fixture: str, - suffix: str, - model_kwargs: dict[str, Any], - available_providers: list[str] | None, + def test_recognition_is_built_the_batches_it_asks_for( + self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture, provider: str, shapes: tuple[Shape, ...] ) -> None: - session = request.getfixturevalue(session_fixture) - ort_pybind.get_available_openvino_device_ids.return_value = ["CPU"] - mocker.patch( - "immich_ml.models.facial_recognition.recognition.ort.get_available_providers", - return_value=available_providers, - ) - path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = suffix - session.return_value.get_inputs.return_value = [SimpleNamespace(name="input.1", shape=(1, 3, 224, 224))] - session.return_value.get_outputs.return_value = [SimpleNamespace(name="output.1", shape=(1, 800))] + mocker.patch("immich_ml.models.base.InferenceModel.download") + mocker.patch("immich_ml.sessions.ort.ort.get_available_providers", return_value=[provider]) - face_recognizer = FaceRecognizer("buffalo_s", cache_dir=path, **model_kwargs) + face_recognizer = FaceRecognizer("buffalo_s", cache_dir=path) face_recognizer.load() - assert face_recognizer.batch_size == 1 - batch_axis.update_dims.assert_not_called() - batch_axis.onnx.load.assert_not_called() - batch_axis.onnx.save.assert_not_called() + assert face_recognizer.session.shapes == shapes + + @pytest.mark.parametrize("model_format", [ModelFormat.RKNN, ModelFormat.ARMNN]) + def test_recognition_takes_the_batch_the_artifact_was_compiled_for( + self, + rknn_session: mock.Mock, + ann_session: mock.Mock, + path: mock.Mock, + mocker: MockerFixture, + model_format: ModelFormat, + ) -> None: + mocker.patch("immich_ml.models.base.InferenceModel.download") + mocker.patch.object(settings, "max_batch_size", MaxBatchSize(facial_recognition=8)) # which cannot raise it + rknn_session.return_value.inputs = [SimpleNamespace(name="image", shape=(1, 112, 112, 3))] + + face_recognizer = FaceRecognizer("buffalo_s", cache_dir=path, model_format=model_format) + face_recognizer.load() + + assert face_recognizer.session.shapes == (Shape(batch=1),) def test_set_custom_max_batch_size(self, mocker: MockerFixture) -> None: mocker.patch.object(settings, "max_batch_size", MaxBatchSize(facial_recognition=2)) recognizer = FaceRecognizer("buffalo_l", cache_dir="test_cache") - assert recognizer.batch_size == 2 + assert recognizer.shape_policy.dims == (Shape(batch=1), Shape(batch=2)) def test_ignore_other_custom_max_batch_size(self, mocker: MockerFixture) -> None: mocker.patch.object(settings, "max_batch_size", MaxBatchSize(ocr=2)) - mocker.patch( - "immich_ml.models.facial_recognition.recognition.ort.get_available_providers", - return_value=["CPUExecutionProvider"], - ) recognizer = FaceRecognizer("buffalo_l", cache_dir="test_cache") - assert recognizer.batch_size is None + assert recognizer.shape_policy.dims == ( + Shape(batch=1), + Shape(batch=4), + ) # another task's setting is not this one's class TestOcr: - def test_det_min_score_is_per_request(self, path: mock.Mock) -> None: - path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx" + @pytest.mark.parametrize("dtype", [np.float32, np.float16]) # a graph narrowed to half answers in it + def test_det_min_score_is_per_request( + self, path: mock.Mock, stub_session: Callable[..., mock.Mock], dtype: type[np.generic] + ) -> None: text_detector = TextDetector("PP-OCRv5_mobile", cache_dir="test_cache") - probs = np.zeros((1, 1, 64, 64), dtype=np.float32) + probs = np.zeros((1, 1, 64, 64), dtype=dtype) probs[..., 16:32, 8:56] = 0.6 - text_detector.session = mock.Mock() - text_detector.session.get_inputs.return_value = [SimpleNamespace(name="input.1", shape=(1, 3, 64, 64))] - text_detector.session.run.return_value = [probs] + text_detector.session = stub_session((1, 3, 64, 64), outputs=[probs]) image = Image.new("RGB", (64, 64)) assert len(text_detector._predict(image, minScore=0.5)["boxes"]) == 1 @@ -1203,14 +1643,92 @@ class TestOcr: # the default must be unaffected by the request that just ran assert len(text_detector._predict(image)["boxes"]) == 1 - def test_rec_min_score_is_per_request(self, path: mock.Mock, mocker: MockerFixture) -> None: - path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx" - text_recognizer = TextRecognizer("PP-OCRv5_mobile", cache_dir="test_cache") - text_recognizer.session = mock.Mock() - text_recognizer.session.get_inputs.return_value = [SimpleNamespace(name="input.1", shape=(1, 3, 48, 96))] - text_recognizer.session.run.return_value = [np.zeros((1, 4, 8), dtype=np.float32)] + def test_fetches_the_older_exports_from_where_rapidocr_hosts_them( + self, tmp_path: Path, snapshot_download: mock.Mock, mocker: MockerFixture + ) -> None: + fetch = mocker.patch("rapidocr.utils.download_file.DownloadFile.run") + + TextRecognizer("EN__PP-OCRv5_mobile", cache_dir=tmp_path).download() + + fetched = fetch.call_args.args[0] + assert fetched.file_url.endswith("/onnx/PP-OCRv5/rec/en_PP-OCRv5_rec_mobile.onnx") + assert fetched.save_path == tmp_path / "recognition/model.onnx" + snapshot_download.assert_not_called() + + mocker.patch.object(settings, "model_revision", "v2") + TextDetector("PP-OCRv5_mobile", cache_dir=tmp_path).download() + snapshot_download.assert_called_once() + + @pytest.mark.parametrize(("crop", "fed_width"), [(96, 224), (384, 480)]) # the floor, then a quarter past the text + def test_rec_feeds_raw_rgb_padded_to_the_batch_width( + self, path: mock.Mock, mocker: MockerFixture, stub_session: Callable[..., mock.Mock], crop: int, fed_width: int + ) -> None: + session = stub_session((1, 48, 224, 3), outputs=[np.zeros((1, 4, 8), np.float32)], normalizes_input=True) + text_recognizer = loaded(TextRecognizer("PP-OCRv5_mobile", cache_dir=path), session, mocker) text_recognizer.decoder = mock.Mock() - text_recognizer.decoder.decode.return_value = (["hello"], np.array([0.8], dtype=np.float32)) + text_recognizer.decoder.return_value = (["hi"], np.array([0.95], dtype=np.float32)) + image = Image.new("RGB", (500, 100), (7, 8, 9)) + box = np.array([[[0, 0], [crop, 0], [crop, 48], [0, 48]]], dtype=np.float32) + texts: Any = {"boxes": box, "scores": np.array([0.9], dtype=np.float32)} + + text_recognizer._predict(image, texts) + + fed = session.run.call_args.args[1]["input.1"] + assert fed.dtype == np.uint8 and fed.shape == (1, 48, fed_width, 3) + assert fed[0, 0, 0].tolist() == [7, 8, 9] # the crop, unnormalized + assert fed[0, 0, -1].tolist() == [127, 127, 127] # and the pad the batch was filled with + + def test_det_letterboxes_onto_the_canvas_the_session_takes( + self, path: mock.Mock, stub_session: Callable[..., mock.Mock] + ) -> None: + text_detector = TextDetector("PP-OCRv5_mobile", cache_dir=path) + text_detector.session = stub_session( + (1, 64, 128, 3), + outputs=[np.zeros((1, 64, 128), dtype=np.float32)], + name="image", + normalizes_input=True, + shapes=(Shape(batch=1, height=64, width=128), Shape(batch=1, height=128, width=64)), + ) + + text_detector._predict(Image.new("RGB", (100, 100), (10, 20, 30)), maxResolution=64) + + fed = text_detector.session.run.call_args.args[1]["image"] + assert fed.dtype == np.uint8 and fed.shape == (1, 64, 128, 3) # the canvas costing it the least downscale + assert fed[0, :, :64].all() and not fed[0, :, 64:].any() # in the corner, with the rest left black + # the legacy path swaps to BGR on the way in; the fused graph does its own + assert fed[0, 0, 0].tolist() == [10, 20, 30] + + def test_rec_runs_the_batch_at_a_compiled_width( + self, path: mock.Mock, mocker: MockerFixture, stub_session: Callable[..., mock.Mock] + ) -> None: + session = stub_session( + (1, 48, 400, 3), + outputs=[np.zeros((1, 4, 8), np.float32)], + normalizes_input=True, + shapes=(Shape(batch=1, width=512), Shape(batch=1, width=400)), # as a binary declares them, widest first + ) + text_recognizer = loaded(TextRecognizer("PP-OCRv5_mobile", cache_dir=path), session, mocker) + text_recognizer.decoder = mock.Mock() + text_recognizer.decoder.return_value = (["hi"], np.array([0.95], dtype=np.float32)) + image = Image.new("RGB", (500, 100), (7, 8, 9)) + box = np.array([[[0, 0], [384, 0], [384, 48], [0, 48]]], dtype=np.float32) + texts: Any = {"boxes": box, "scores": np.array([0.9], dtype=np.float32)} + + text_recognizer._predict(image, texts) + + fed = session.run.call_args.args[1]["input.1"] + assert fed.shape == (1, 48, 400, 3) # 384 wide in its own right, run at the compiled width above it + + def test_rec_min_score_is_per_request( + self, path: mock.Mock, mocker: MockerFixture, stub_session: Callable[..., mock.Mock] + ) -> None: + text_recognizer = loaded( + TextRecognizer("PP-OCRv5_mobile", cache_dir="test_cache"), + stub_session((1, 3, 48, 96), outputs=[np.zeros((1, 4, 8), dtype=np.float32)]), + mocker, + ) + text_recognizer.decoder = mock.Mock() + text_recognizer.decoder.return_value = (["hello"], np.array([0.8], dtype=np.float32)) mocker.patch.object(text_recognizer, "_crop", return_value=np.zeros((48, 96, 3), dtype=np.uint8)) image = Image.new("RGB", (100, 50)) box = np.array([[[0, 0], [96, 0], [96, 48], [0, 48]]], dtype=np.float32) @@ -1222,35 +1740,56 @@ class TestOcr: # the default (0.9) rejects a 0.8 score, and must be unaffected by the 0.7 request assert text_recognizer._predict(image, texts())["text"] == [] + def test_rec_decodes_the_half_precision_probabilities_a_host_decode_graph_emits(self) -> None: + probs = np.zeros((1, 3, 4), dtype=np.float16) + probs[0, 0, 2] = 0.75 + probs[0, 1, 3] = 2**-20 # subnormal in half precision, and still the likeliest class + probs[0, 2, 1] = 0.5 + + indices, picked = probabilities(probs) + + assert indices.tolist() == [[2, 3, 1]] + assert picked.dtype == np.float32 and picked.tolist() == [[0.75, 2**-20, 0.5]] + + @pytest.mark.parametrize("dtype", [np.float32, np.float16]) # a binary hands back what the NPU computed in + def test_rec_decodes_the_raw_logits_a_binary_emits(self, dtype: type[np.generic]) -> None: + raw = np.full((1, 4, 1, 3), -5.0, dtype=dtype) # steps, then a unit axis, then the classes + raw[0, 0, 0, 2] = 5.0 + raw[0, 1, 0] = [-3.0, -1.0, -2.0] # no logit above zero, which the half's bits order backwards + raw[0, 2, 0] = [-3.0, -1.0, -2.0] # a repeat, which the decode does not read + raw[0, 3, 0] = [1.0, 1.0, 0.0] # a tie goes to the first class, here the blank + + indices, confidence = logits(raw) + + assert indices.tolist() == [[2, 1, 1, 0]] + assert np.allclose(confidence, [[1 / (1 + 2 * np.exp(-10)), 1 / (1 + np.exp(-1) + np.exp(-2)), 0, 0]]) + def test_set_rec_set_default_max_batch_size( self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture ) -> None: - path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx" mocker.patch("immich_ml.models.base.InferenceModel.download") text_recognizer = TextRecognizer("PP-OCRv5_mobile", cache_dir="test_cache") - assert text_recognizer.batch_size == 6 + assert {shape.batch for shape in text_recognizer.shape_policy.dims} == {1, 6} def test_set_custom_max_batch_size(self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture) -> None: - path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx" mocker.patch("immich_ml.models.base.InferenceModel.download") mocker.patch.object(settings, "max_batch_size", MaxBatchSize(ocr=4)) text_recognizer = TextRecognizer("PP-OCRv5_mobile", cache_dir="test_cache") - assert text_recognizer.batch_size == 4 + assert {shape.batch for shape in text_recognizer.shape_policy.dims} == {1, 4} def test_ignore_other_custom_max_batch_size( self, ort_session: mock.Mock, path: mock.Mock, mocker: MockerFixture ) -> None: - path.return_value.__truediv__.return_value.__truediv__.return_value.suffix = ".onnx" mocker.patch("immich_ml.models.base.InferenceModel.download") mocker.patch.object(settings, "max_batch_size", MaxBatchSize(facial_recognition=3)) text_recognizer = TextRecognizer("PP-OCRv5_mobile", cache_dir="test_cache") - assert text_recognizer.batch_size == 6 + assert {shape.batch for shape in text_recognizer.shape_policy.dims} == {1, 6} @pytest.mark.asyncio @@ -1262,6 +1801,15 @@ class TestCache: assert len(model_cache._models) == 1 mock_get_model.return_value.assert_called_once() + async def test_separate_instances_per_graph_option(self) -> None: + # the real TextDetector, so that dropping its graph_options fails this + model_cache = ModelCache() + + for max_resolution in (736, 1088, 1088): + model_cache.get("PP-OCRv5_mobile", ModelType.DETECTION, ModelTask.OCR, maxResolution=max_resolution) + # maxResolution picks the graphs a detector builds, so instances cannot share them + assert len(model_cache._models) == 2 + async def test_kwargs_used(self, mock_get_model: mock.Mock) -> None: model_cache = ModelCache() model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, cache_dir="test_cache") @@ -1283,6 +1831,8 @@ class TestCache: events: list[str] = [] class Model: + graph_options = () + def __init__(self, *args: Any, **kwargs: Any) -> None: pass @@ -1408,6 +1958,7 @@ class TestCache: ], any_order=True, ) + mock_get_model.return_value.return_value.build.assert_called() # so that no request waits on a graph async def test_preloads_all_models(self, monkeypatch: MonkeyPatch, mock_get_model: mock.Mock) -> None: os.environ["MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL"] = "ViT-B-32__openai" @@ -1477,6 +2028,7 @@ class TestLoad: res = await load(mock_model) assert res is mock_model + mock_model.unload.assert_called_once() # so that the retry holds one copy while it loads mock_model.clear_cache.assert_called_once() assert mock_model.load.call_count == 2 @@ -1508,124 +2060,22 @@ class TestLoad: await load(mock_model) mock_model.clear_cache.assert_not_called() + mock_model.unload.assert_called_once() assert mock_model.load.call_count == 2 warning.assert_called_once_with( "ARMNN is available, but model 'test_model_name' does not support it.", exc_info=error ) - mock_model.model_format = ModelFormat.ONNX + assert mock_model.model_format == ModelFormat.ONNX + async def test_leaves_the_cache_alone_when_a_request_fails_for_another_reason(self) -> None: + mock_model = mock.Mock(spec=InferenceModel) + mock_model.loaded = True + mock_model.predict.side_effect = OSError("image file is truncated") -@pytest.mark.parametrize("size", [(0, 100), (100, 0), (0, 0)]) -def test_predict_rejects_empty_image(size: tuple[int, int], deployed_app: TestClient) -> None: - with mock.patch("immich_ml.main.decode_pil", return_value=Image.new("RGB", size)): - response = deployed_app.post( - "http://localhost:3003/predict", - data={"entries": json.dumps({"clip": {"visual": {"modelName": "ViT-B-32__openai"}}})}, - files={"image": b"fake image bytes"}, - ) + with pytest.raises(OSError): + await attempt(mock_model, mock_model.predict, MODEL_FILE_ERRORS) - assert response.status_code == 400 - assert "zero" in response.json()["detail"].lower() - - -def test_root_endpoint(deployed_app: TestClient) -> None: - response = deployed_app.get("http://localhost:3003") - - body = response.json() - assert response.status_code == 200 - assert body == {"message": "Immich ML"} - - -def test_ping_endpoint(deployed_app: TestClient) -> None: - response = deployed_app.get("http://localhost:3003/ping") - - assert response.status_code == 200 - assert response.text == "pong" - - -@pytest.mark.skipif( - not settings.test_full, - reason="More time-consuming since it deploys the app and loads models.", -) -class TestPredictionEndpoints: - def test_clip_image_endpoint( - self, pil_image: Image.Image, responses: dict[str, Any], deployed_app: TestClient - ) -> None: - byte_image = BytesIO() - pil_image.save(byte_image, format="jpeg") - expected = responses["clip"]["image"] - - response = deployed_app.post( - "http://localhost:3003/predict", - data={"entries": json.dumps({"clip": {"visual": {"modelName": "ViT-B-32__openai"}}})}, - files={"image": byte_image.getvalue()}, - ) - - actual = response.json() - assert response.status_code == 200 - assert isinstance(actual, dict) - embedding = actual.get("clip", None) - assert isinstance(embedding, str) - parsed_embedding = orjson.loads(embedding) - assert np.allclose(expected, parsed_embedding) - - def test_clip_text_endpoint(self, responses: dict[str, Any], deployed_app: TestClient) -> None: - expected = responses["clip"]["text"] - - response = deployed_app.post( - "http://localhost:3003/predict", - data={ - "entries": json.dumps( - { - "clip": {"textual": {"modelName": "ViT-B-32__openai"}}, - }, - ), - "text": "test search query", - }, - ) - - actual = response.json() - assert response.status_code == 200 - assert isinstance(actual, dict) - embedding = actual.get("clip", None) - assert isinstance(embedding, str) - parsed_embedding = orjson.loads(embedding) - assert np.allclose(expected, parsed_embedding) - - def test_face_endpoint(self, pil_image: Image.Image, responses: dict[str, Any], deployed_app: TestClient) -> None: - byte_image = BytesIO() - pil_image.save(byte_image, format="jpeg") - - response = deployed_app.post( - "http://localhost:3003/predict", - data={ - "entries": json.dumps( - { - "facial-recognition": { - "detection": {"modelName": "buffalo_l", "options": {"minScore": 0.034}}, - "recognition": {"modelName": "buffalo_l"}, - } - } - ) - }, - files={"image": byte_image.getvalue()}, - ) - - actual = response.json() - assert response.status_code == 200 - assert isinstance(actual, dict) - assert actual.get("imageHeight", None) == responses["imageHeight"] - assert actual.get("imageWidth", None) == responses["imageWidth"] - assert "facial-recognition" in actual and isinstance(actual["facial-recognition"], list) - assert len(actual["facial-recognition"]) == len(responses["facial-recognition"]) - - for expected_face, actual_face in zip(responses["facial-recognition"], actual["facial-recognition"]): - assert expected_face["boundingBox"] == actual_face["boundingBox"] - embedding = actual_face.get("embedding", None) - assert isinstance(embedding, str) - parsed_embedding = orjson.loads(embedding) - assert np.allclose(expected_face["embedding"], parsed_embedding) - assert np.allclose(expected_face["score"], actual_face["score"]) + mock_model.clear_cache.assert_not_called() @pytest.mark.asyncio @@ -1687,3 +2137,122 @@ async def test_returns_memory_once_requests_stop_unless_another_arrives(mocker: await anext(third) loop.call_later.return_value.cancel.assert_called_once() await third.aclose() + + +@pytest.mark.parametrize("size", [(0, 100), (100, 0), (0, 0)]) +def test_predict_rejects_empty_image(size: tuple[int, int], deployed_app: TestClient) -> None: + with mock.patch("immich_ml.main.decode_pil", return_value=Image.new("RGB", size)): + response = deployed_app.post( + "http://localhost:3003/predict", + data={"entries": json.dumps({"clip": {"visual": {"modelName": "ViT-B-32__openai"}}})}, + files={"image": b"fake image bytes"}, + ) + + assert response.status_code == 400 + assert "zero" in response.json()["detail"].lower() + + +def test_root_endpoint(deployed_app: TestClient) -> None: + response = deployed_app.get("http://localhost:3003") + + body = response.json() + assert response.status_code == 200 + assert body == {"message": "Immich ML"} + + +def test_ping_endpoint(deployed_app: TestClient) -> None: + response = deployed_app.get("http://localhost:3003/ping") + + assert response.status_code == 200 + assert response.text == "pong" + + +@pytest.mark.skipif( + not settings.test_full, + reason="More time-consuming since it deploys the app and loads models.", +) +class TestPredictionEndpoints: + def test_clip_image_endpoint( + self, asset: Callable[[str], bytes], responses: dict[str, Any], deployed_app: TestClient + ) -> None: + response = deployed_app.post( + "http://localhost:3003/predict", + data={"entries": json.dumps({"clip": {"visual": {"modelName": "ViT-B-32__openai"}}})}, + files={"image": asset("albums/nature/silver_fir.jpg")}, + ) + + assert response.status_code == 200 + assert np.allclose(orjson.loads(response.json()["clip"]), responses["clip"]["image"], atol=1e-3) + + def test_clip_text_endpoint(self, responses: dict[str, Any], deployed_app: TestClient) -> None: + response = deployed_app.post( + "http://localhost:3003/predict", + data={ + "entries": json.dumps({"clip": {"textual": {"modelName": "ViT-B-32__openai"}}}), + "text": "a photo of a forest", + }, + ) + + assert response.status_code == 200 + assert np.allclose(orjson.loads(response.json()["clip"]), responses["clip"]["text"], atol=1e-3) + + def test_face_endpoint( + self, asset: Callable[[str], bytes], responses: dict[str, Any], deployed_app: TestClient + ) -> None: + expected = responses["facial-recognition"] + + response = deployed_app.post( + "http://localhost:3003/predict", + data={ + "entries": json.dumps( + { + "facial-recognition": { + "detection": {"modelName": "buffalo_l", "options": {"minScore": 0.7}}, + "recognition": {"modelName": "buffalo_l"}, + } + } + ) + }, + files={"image": asset("metadata/faces/portrait.jpg")}, + ) + + actual = response.json() + assert response.status_code == 200 + assert actual["imageWidth"] == expected["imageWidth"] + assert actual["imageHeight"] == expected["imageHeight"] + assert len(actual["facial-recognition"]) == len(expected["faces"]) + + for expected_face, actual_face in zip(expected["faces"], actual["facial-recognition"]): + assert actual_face["boundingBox"] == expected_face["boundingBox"] + assert actual_face["score"] == pytest.approx(expected_face["score"], abs=1e-3) + assert np.allclose(orjson.loads(actual_face["embedding"]), expected_face["embedding"], atol=1e-3) + + def test_ocr_endpoint( + self, asset: Callable[[str], bytes], responses: dict[str, Any], deployed_app: TestClient + ) -> None: + expected = responses["ocr"] + + response = deployed_app.post( + "http://localhost:3003/predict", + data={ + "entries": json.dumps( + { + "ocr": { + "detection": { + "modelName": "PP-OCRv5_mobile", + "options": {"maxResolution": 736, "minScore": 0.5}, + }, + "recognition": {"modelName": "PP-OCRv5_mobile", "options": {"minScore": 0.9}}, + } + } + ) + }, + files={"image": asset("albums/text/craft-beer.jpg")}, + ) + + actual = response.json()["ocr"] + assert response.status_code == 200 + assert actual["text"] == expected["text"] + assert np.allclose(actual["box"], expected["box"], atol=1e-3) + assert np.allclose(actual["boxScore"], expected["boxScore"], atol=1e-3) + assert np.allclose(actual["textScore"], expected["textScore"], atol=1e-3) diff --git a/machine-learning/uv.lock b/machine-learning/uv.lock index 5c84de4a4..00105a84e 100644 --- a/machine-learning/uv.lock +++ b/machine-learning/uv.lock @@ -682,6 +682,7 @@ dependencies = [ { name = "fastapi" }, { name = "gunicorn" }, { name = "huggingface-hub" }, + { name = "immich-model" }, { name = "numpy" }, { name = "onnx" }, { name = "opencv-python-headless" }, @@ -762,6 +763,7 @@ requires-dist = [ { name = "fastapi", specifier = ">=0.95.2,<1.0" }, { name = "gunicorn", specifier = ">=21.1.0" }, { name = "huggingface-hub", specifier = ">=1.0,<2.0" }, + { name = "immich-model", git = "https://github.com/immich-app/ml-models.git?rev=88d2b64ff95ce7117760660681d95e0e1308819b" }, { name = "numpy", specifier = ">=2.4.0,<3.0" }, { name = "onnx", specifier = ">=1.22.0" }, { name = "onnxruntime", marker = "extra == 'armnn'", specifier = ">=1.23.2,<2" }, @@ -824,6 +826,18 @@ types = [ { name = "types-ujson", specifier = ">=5.10.0.20240515" }, ] +[[package]] +name = "immich-model" +version = "0.2.0" +source = { git = "https://github.com/immich-app/ml-models.git?rev=88d2b64ff95ce7117760660681d95e0e1308819b#88d2b64ff95ce7117760660681d95e0e1308819b" } +dependencies = [ + { name = "numpy" }, + { name = "onnx" }, + { name = "onnx-ir" }, + { name = "onnxscript" }, + { name = "safetensors" }, +] + [[package]] name = "iniconfig" version = "2.0.0" @@ -1284,6 +1298,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/00/50/257a880384a1dd502d543b0067945074d63cd17d0840e958355bc8197da8/onnx-1.22.0-cp314-cp314t-win_arm64.whl", hash = "sha256:8e268cdc0547e3949799ffd4a44451dc2b9080b57d0824a2db680b6ec65506f0", size = 17231391, upload-time = "2026-06-15T12:50:03.047Z" }, ] +[[package]] +name = "onnx-ir" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "sympy" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d6/c2/61194cec0dbc5622273c0ebd592d37cc1dca0d7f1a744f02edd45ac905a3/onnx_ir-1.0.0.tar.gz", hash = "sha256:9e261f25fde8da9612ae5cb43b3b374d5ff469c04af0363cad588b2bb000b812", size = 163121, upload-time = "2026-08-11T14:49:46.895Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/cd/6d1637172eb59c7b18ac90ed089d1f599a11fe0e63b4db2d017f3bb38a32/onnx_ir-1.0.0-py3-none-any.whl", hash = "sha256:e578f0d608d3062866b48223616eb2d10a6d6d01f8b8faac596129034f483cc7", size = 185849, upload-time = "2026-08-11T14:49:45.524Z" }, +] + [[package]] name = "onnxruntime" version = "1.30.0" @@ -1371,6 +1401,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3e/92/46ae2cd565961a89189900f385bb2f13a9fa731ea4674001d23720fbb1e0/onnxruntime_openvino-1.24.1-cp313-cp313-win_amd64.whl", hash = "sha256:434bf49aa71393c577a456c9d76c98e6d6958a833fa0876793e3d5437b5a511a", size = 13658485, upload-time = "2026-02-26T13:44:53.889Z" }, ] +[[package]] +name = "onnxscript" +version = "0.7.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes" }, + { name = "numpy" }, + { name = "onnx" }, + { name = "onnx-ir" }, + { name = "packaging" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0a/01/3e3fab8d643ca097ea4aa9e51246643699dfaaa0650589744fe44bc46651/onnxscript-0.7.2.tar.gz", hash = "sha256:2c664f6383d10f332a4d47b2876dcab16dba84909fe703656b19abc281fda165", size = 646719, upload-time = "2026-09-09T17:06:44.567Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b9/3b/06260997cdc41138e58718588a6c87d0eb342bbe0dda8a6aae91d163c384/onnxscript-0.7.2-py3-none-any.whl", hash = "sha256:d0e7121c6a1eefd608058928e111cbdb76709f70d269ff0d07aee493bd1d13c9", size = 754215, upload-time = "2026-09-09T17:06:46.442Z" }, +] + [[package]] name = "opencv-python" version = "4.11.0.86" @@ -2034,6 +2081,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/19/016553f86f207450aebebc2b2b5088d086b901cc8186c02ac4284db3bd88/ruff-0.15.16-py3-none-win_arm64.whl", hash = "sha256:8cd61783afb39638a7133ef0d2dfb1e91277593962f81b5a8423eb0b888a6121", size = 11134555, upload-time = "2026-06-04T16:33:00.136Z" }, ] +[[package]] +name = "safetensors" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/29/9c/6e74567782559a63bd040a236edca26fd71bc7ba88de2ef35d75df3bca5e/safetensors-0.7.0.tar.gz", hash = "sha256:07663963b67e8bd9f0b8ad15bb9163606cd27cc5a1b96235a50d8369803b96b0", size = 200878, upload-time = "2025-11-19T15:18:43.199Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/47/aef6c06649039accf914afef490268e1067ed82be62bcfa5b7e886ad15e8/safetensors-0.7.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c82f4d474cf725255d9e6acf17252991c3c8aac038d6ef363a4bf8be2f6db517", size = 467781, upload-time = "2025-11-19T15:18:35.84Z" }, + { url = "https://files.pythonhosted.org/packages/e8/00/374c0c068e30cd31f1e1b46b4b5738168ec79e7689ca82ee93ddfea05109/safetensors-0.7.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:94fd4858284736bb67a897a41608b5b0c2496c9bdb3bf2af1fa3409127f20d57", size = 447058, upload-time = "2025-11-19T15:18:34.416Z" }, + { url = "https://files.pythonhosted.org/packages/f1/06/578ffed52c2296f93d7fd2d844cabfa92be51a587c38c8afbb8ae449ca89/safetensors-0.7.0-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e07d91d0c92a31200f25351f4acb2bc6aff7f48094e13ebb1d0fb995b54b6542", size = 491748, upload-time = "2025-11-19T15:18:09.79Z" }, + { url = "https://files.pythonhosted.org/packages/ae/33/1debbbb70e4791dde185edb9413d1fe01619255abb64b300157d7f15dddd/safetensors-0.7.0-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8469155f4cb518bafb4acf4865e8bb9d6804110d2d9bdcaa78564b9fd841e104", size = 503881, upload-time = "2025-11-19T15:18:16.145Z" }, + { url = "https://files.pythonhosted.org/packages/8e/1c/40c2ca924d60792c3be509833df711b553c60effbd91da6f5284a83f7122/safetensors-0.7.0-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:54bef08bf00a2bff599982f6b08e8770e09cc012d7bba00783fc7ea38f1fb37d", size = 623463, upload-time = "2025-11-19T15:18:21.11Z" }, + { url = "https://files.pythonhosted.org/packages/9b/3a/13784a9364bd43b0d61eef4bea2845039bc2030458b16594a1bd787ae26e/safetensors-0.7.0-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:42cb091236206bb2016d245c377ed383aa7f78691748f3bb6ee1bfa51ae2ce6a", size = 532855, upload-time = "2025-11-19T15:18:25.719Z" }, + { url = "https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dac7252938f0696ddea46f5e855dd3138444e82236e3be475f54929f0c510d48", size = 507152, upload-time = "2025-11-19T15:18:33.023Z" }, + { url = "https://files.pythonhosted.org/packages/3c/a8/4b45e4e059270d17af60359713ffd83f97900d45a6afa73aaa0d737d48b6/safetensors-0.7.0-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1d060c70284127fa805085d8f10fbd0962792aed71879d00864acda69dbab981", size = 541856, upload-time = "2025-11-19T15:18:31.075Z" }, + { url = "https://files.pythonhosted.org/packages/06/87/d26d8407c44175d8ae164a95b5a62707fcc445f3c0c56108e37d98070a3d/safetensors-0.7.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:cdab83a366799fa730f90a4ebb563e494f28e9e92c4819e556152ad55e43591b", size = 674060, upload-time = "2025-11-19T15:18:37.211Z" }, + { url = "https://files.pythonhosted.org/packages/11/f5/57644a2ff08dc6325816ba7217e5095f17269dada2554b658442c66aed51/safetensors-0.7.0-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:672132907fcad9f2aedcb705b2d7b3b93354a2aec1b2f706c4db852abe338f85", size = 771715, upload-time = "2025-11-19T15:18:38.689Z" }, + { url = "https://files.pythonhosted.org/packages/86/31/17883e13a814bd278ae6e266b13282a01049b0c81341da7fd0e3e71a80a3/safetensors-0.7.0-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:5d72abdb8a4d56d4020713724ba81dac065fedb7f3667151c4a637f1d3fb26c0", size = 714377, upload-time = "2025-11-19T15:18:40.162Z" }, + { url = "https://files.pythonhosted.org/packages/4a/d8/0c8a7dc9b41dcac53c4cbf9df2b9c83e0e0097203de8b37a712b345c0be5/safetensors-0.7.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b0f6d66c1c538d5a94a73aa9ddca8ccc4227e6c9ff555322ea40bdd142391dd4", size = 677368, upload-time = "2025-11-19T15:18:41.627Z" }, + { url = "https://files.pythonhosted.org/packages/05/e5/cb4b713c8a93469e3c5be7c3f8d77d307e65fe89673e731f5c2bfd0a9237/safetensors-0.7.0-cp38-abi3-win32.whl", hash = "sha256:c74af94bf3ac15ac4d0f2a7c7b4663a15f8c2ab15ed0fc7531ca61d0835eccba", size = 326423, upload-time = "2025-11-19T15:18:45.74Z" }, + { url = "https://files.pythonhosted.org/packages/5d/e6/ec8471c8072382cb91233ba7267fd931219753bb43814cbc71757bfd4dab/safetensors-0.7.0-cp38-abi3-win_amd64.whl", hash = "sha256:d1239932053f56f3456f32eb9625590cc7582e905021f94636202a864d470755", size = 341380, upload-time = "2025-11-19T15:18:44.427Z" }, +] + [[package]] name = "setuptools" version = "80.9.0" @@ -2123,14 +2192,14 @@ wheels = [ [[package]] name = "sympy" -version = "1.12" +version = "1.14.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "mpmath" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e5/57/3485a1a3dff51bfd691962768b14310dae452431754bfc091250be50dd29/sympy-1.12.tar.gz", hash = "sha256:ebf595c8dac3e0fdc4152c51878b498396ec7f30e7a914d6071e674d49420fb8", size = 6722203, upload-time = "2023-05-10T18:23:00.378Z" } +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d2/05/e6600db80270777c4a64238a98d442f0fd07cc8915be2a1c16da7f2b9e74/sympy-1.12-py3-none-any.whl", hash = "sha256:c3588cd4295d0c0f603d0f2ae780587e64e2efeedb3521e46b9bb1d08d184fa5", size = 5742435, upload-time = "2023-05-10T18:22:14.76Z" }, + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] [[package]]