From ab159b81182f26a8943f5dd66b719e8cdaad4626 Mon Sep 17 00:00:00 2001 From: Mert <101130780+mertalev@users.noreply.github.com> Date: Mon, 28 Sep 2026 14:14:01 -0400 Subject: [PATCH] fix(ml): stricter typing and request validation (#31856) --- machine-learning/conftest.py | 7 - machine-learning/immich_ml/main.py | 166 ++++---- machine-learning/immich_ml/models/__init__.py | 42 -- machine-learning/immich_ml/models/base.py | 40 +- machine-learning/immich_ml/models/cache.py | 22 +- .../immich_ml/models/clip/textual.py | 12 +- .../immich_ml/models/clip/visual.py | 8 +- .../models/facial_recognition/detection.py | 9 +- .../models/facial_recognition/recognition.py | 18 +- .../immich_ml/models/ocr/detection.py | 34 +- .../immich_ml/models/ocr/legacy.py | 16 +- .../immich_ml/models/ocr/recognition.py | 10 +- machine-learning/immich_ml/pipeline.py | 85 ++++ machine-learning/immich_ml/schemas.py | 41 +- machine-learning/mise.toml | 6 +- machine-learning/test_main.py | 372 ++++++++++-------- 16 files changed, 511 insertions(+), 377 deletions(-) create mode 100644 machine-learning/immich_ml/pipeline.py diff --git a/machine-learning/conftest.py b/machine-learning/conftest.py index 34de2cbbb..5668e6a77 100644 --- a/machine-learning/conftest.py +++ b/machine-learning/conftest.py @@ -27,13 +27,6 @@ def pil_image() -> Image.Image: return Image.new("RGB", (600, 800)) -@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 - - @pytest.fixture(scope="session") def deployed_app() -> Iterator[TestClient]: with TestClient(app) as client: diff --git a/machine-learning/immich_ml/main.py b/machine-learning/immich_ml/main.py index 8f27cb0d9..dc946be98 100644 --- a/machine-learning/immich_ml/main.py +++ b/machine-learning/immich_ml/main.py @@ -14,27 +14,29 @@ from fastapi import Depends, FastAPI, File, Form, HTTPException from fastapi.responses import JSONResponse, PlainTextResponse from onnxruntime.capi.onnxruntime_pybind11_state import InvalidProtobuf, NoSuchFile from PIL.Image import Image -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError from starlette.formparsers import MultiPartParser from immich_ml import allocator -from immich_ml.models import get_model_deps from immich_ml.models.base import InferenceModel from immich_ml.models.transforms import decode_pil from immich_ml.sessions.ort import flush_denormals from .config import PreloadModelData, log, settings +from .models.base import InferenceEntry from .models.cache import ModelCache +from .pipeline import Clip, FacialRecognition, Ocr, PipelineRequest, Slot from .schemas import ( - InferenceEntries, - InferenceEntry, + FaceDetectionOptions, + FaceRecognitionOptions, InferenceResponse, ModelFormat, ModelIdentity, - ModelTask, - ModelType, - PipelineRequest, - T, + Options, + TextDetectionOptions, + TextRecognitionOptions, + TextualOptions, + VisualOptions, ) @@ -43,6 +45,7 @@ class ORJSONResponse(JSONResponse): return orjson.dumps(content, option=orjson.OPT_SERIALIZE_NUMPY) +PIPELINE_REQUEST = TypeAdapter(PipelineRequest) MultiPartParser.spool_max_size = 2**26 # spools to disk if payload is 64 MiB or larger model_cache = ModelCache() @@ -84,49 +87,40 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]: async def preload_models(preload: PreloadModelData) -> None: - log.info(f"Preloading models: clip:{preload.clip} facial_recognition:{preload.facial_recognition}") + requests = [ + *(PipelineRequest(clip=Clip(textual=Slot(name, TextualOptions()))) for name in _names(preload.clip.textual)), + *(PipelineRequest(clip=Clip(visual=Slot(name, VisualOptions()))) for name in _names(preload.clip.visual)), + *( + PipelineRequest(facial_recognition=FacialRecognition(detection=Slot(name, FaceDetectionOptions()))) + for name in _names(preload.facial_recognition.detection) + ), + *( + PipelineRequest(facial_recognition=FacialRecognition(recognition=Slot(name, FaceRecognitionOptions()))) + for name in _names(preload.facial_recognition.recognition) + ), + *( + PipelineRequest( + ocr=Ocr(detection=Slot(name, TextDetectionOptions(max_resolution=preload.ocr.max_resolution))) + ) + for name in _names(preload.ocr.detection) + ), + *( + PipelineRequest(ocr=Ocr(recognition=Slot(name, TextRecognitionOptions()))) + for name in _names(preload.ocr.recognition) + ), + ] + for entry in (entry for request in requests for entry in request.entries()): + await _preload(entry) - async def load_models(model_string: str, model_type: ModelType, model_task: ModelTask, **options: Any) -> None: - for model_name in model_string.split(","): - 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) +async def _preload[O: Options](entry: InferenceEntry[O]) -> None: + log.info(f"Preloading: {entry}") + model = await load(model_cache.get(entry)) + await attempt(model, model.build, MODEL_FILE_ERRORS) - if preload.clip.visual is not None: - await load_models(preload.clip.visual, ModelType.VISUAL, ModelTask.SEARCH) - if preload.facial_recognition.detection is not None: - await load_models( - preload.facial_recognition.detection, - ModelType.DETECTION, - ModelTask.FACIAL_RECOGNITION, - ) - - if preload.facial_recognition.recognition is not None: - await load_models( - preload.facial_recognition.recognition, - ModelType.RECOGNITION, - ModelTask.FACIAL_RECOGNITION, - ) - - if preload.ocr.detection is not None: - await load_models( - preload.ocr.detection, - ModelType.DETECTION, - ModelTask.OCR, - maxResolution=preload.ocr.max_resolution, - ) - - if preload.ocr.recognition is not None: - await load_models( - preload.ocr.recognition, - ModelType.RECOGNITION, - ModelTask.OCR, - ) +def _names(setting: str | None) -> list[str]: + return [] if setting is None else [name.strip() for name in setting.split(",")] async def update_state() -> AsyncGenerator[None, None]: @@ -144,25 +138,15 @@ async def update_state() -> AsyncGenerator[None, None]: release = asyncio.get_running_loop().call_later(5, allocator.release) -def get_entries(entries: str = Form()) -> InferenceEntries: +def get_entries(entries: str = Form()) -> list[InferenceEntry[Any]]: try: - request: PipelineRequest = orjson.loads(entries) - without_deps: list[InferenceEntry] = [] - with_deps: list[InferenceEntry] = [] - for task, types in request.items(): - for type, entry in types.items(): - parsed: InferenceEntry = { - "name": entry["modelName"], - "task": task, - "type": type, - "options": entry.get("options", {}), - } - dep = get_model_deps(parsed["name"], type, task) - (with_deps if dep else without_deps).append(parsed) - return without_deps, with_deps - except (orjson.JSONDecodeError, ValidationError, KeyError, AttributeError) as e: + found = list(PIPELINE_REQUEST.validate_json(entries).entries()) + except (ValidationError, ValueError) as e: log.error(f"Invalid request format: {e}") raise HTTPException(422, "Invalid request format.") + if not found: + raise HTTPException(422, "No model task requested.") + return found app = FastAPI(lifespan=lifespan) @@ -180,7 +164,7 @@ def ping() -> PlainTextResponse: @app.post("/predict", dependencies=[Depends(update_state)]) async def predict( - entries: InferenceEntries = Depends(get_entries), + entries: list[InferenceEntry[Any]] = Depends(get_entries), image: bytes | None = File(default=None), text: str | None = Form(default=None), ) -> Any: @@ -197,50 +181,50 @@ async def predict( return ORJSONResponse(response) -async def run_inference(payload: Image | str, entries: InferenceEntries) -> InferenceResponse: - outputs: dict[ModelIdentity, Any] = {} - response: InferenceResponse = {} +async def run_inference(payload: Image | str, entries: list[InferenceEntry[Any]]) -> InferenceResponse: + read: set[ModelIdentity] = set() - async def _run_inference(entry: InferenceEntry) -> None: - model = model_cache.get(entry["name"], entry["type"], entry["task"], ttl=settings.model_ttl, **entry["options"]) - inputs = [payload] - for dep in model.depends: - try: - inputs.append(outputs[dep]) - except KeyError: - message = f"Task {entry['task']} of type {entry['type']} depends on output of {dep}" - raise HTTPException(400, message) + async def _run_inference[O: Options](entry: InferenceEntry[O]) -> Any: + model = model_cache.get(entry, ttl=settings.model_ttl) + read.update(model.depends) + if missing := [dep for dep in model.depends if dep not in runs]: + raise HTTPException(400, f"{entry.model.__name__} depends on output of {missing[0]}") + inputs = [payload, *[await runs[dep] for dep in model.depends]] model = await load(model) - output = await attempt(model, partial(model.predict, *inputs, **entry["options"]), MODEL_FILE_ERRORS) - outputs[model.identity] = output - response[entry["task"]] = output + return await attempt(model, partial(model.predict, *inputs, options=entry.options), MODEL_FILE_ERRORS) - for stage in entries: # those that depend on another's output run after it - results = await asyncio.gather(*[_run_inference(entry) for entry in stage], return_exceptions=True) - for result in results: - if isinstance(result, BaseException): - raise result + runs = {entry.model.identity: asyncio.create_task(_run_inference(entry)) for entry in entries} + outputs = await asyncio.gather(*runs.values(), return_exceptions=True) + for output in outputs: + if isinstance(output, BaseException): + raise output + response: InferenceResponse = { # a task answers with the output no other model of it reads + entry.model.identity[1].value: output + for entry, output in zip(entries, outputs) + if entry.model.identity not in read + } if isinstance(payload, Image): response["imageHeight"], response["imageWidth"] = payload.height, payload.width return response -async def run(func: Callable[..., T], *args: Any, **kwargs: Any) -> T: +async def run[R](func: Callable[[], R]) -> R: if thread_pool is None: - return func(*args, **kwargs) - partial_func = partial(func, *args, **kwargs) - return await asyncio.get_running_loop().run_in_executor(thread_pool, partial_func) + return func() + return await asyncio.get_running_loop().run_in_executor(thread_pool, func) -async def load(model: InferenceModel) -> InferenceModel: +async def load[O: Options](model: InferenceModel[O]) -> InferenceModel[O]: if not model.loaded: await attempt(model, model.load, (OSError, BadZipFile, *MODEL_FILE_ERRORS)) return model -async def attempt(model: InferenceModel, func: Callable[[], T], corrupt: tuple[type[Exception], ...]) -> T: - def _attempt() -> T: +async def attempt[O: Options, R]( + model: InferenceModel[O], func: Callable[[], R], corrupt: tuple[type[Exception], ...] +) -> R: + def _attempt() -> R: if not model.loaded and model.load_attempts > 1: raise HTTPException(500, f"Failed to load model '{model.model_name}'") try: diff --git a/machine-learning/immich_ml/models/__init__.py b/machine-learning/immich_ml/models/__init__.py index f93939130..e69de29bb 100644 --- a/machine-learning/immich_ml/models/__init__.py +++ b/machine-learning/immich_ml/models/__init__.py @@ -1,42 +0,0 @@ -from immich_ml.models.base import InferenceModel -from immich_ml.models.clip.textual import MClipTextualEncoder, OpenClipTextualEncoder -from immich_ml.models.clip.visual import OpenClipVisualEncoder -from immich_ml.models.ocr.detection import TextDetector -from immich_ml.models.ocr.recognition import TextRecognizer -from immich_ml.schemas import ModelSource, ModelTask, ModelType - -from .constants import get_model_source -from .facial_recognition.detection import FaceDetector -from .facial_recognition.recognition import FaceRecognizer - - -def get_model_class(model_name: str, model_type: ModelType, model_task: ModelTask) -> type[InferenceModel]: - source = get_model_source(model_name) - match source, model_type, model_task: - case ModelSource.OPENCLIP | ModelSource.MCLIP, ModelType.VISUAL, ModelTask.SEARCH: - return OpenClipVisualEncoder - - case ModelSource.OPENCLIP, ModelType.TEXTUAL, ModelTask.SEARCH: - return OpenClipTextualEncoder - - case ModelSource.MCLIP, ModelType.TEXTUAL, ModelTask.SEARCH: - return MClipTextualEncoder - - case ModelSource.INSIGHTFACE, ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION: - return FaceDetector - - case ModelSource.INSIGHTFACE, ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION: - return FaceRecognizer - - case ModelSource.PADDLE, ModelType.DETECTION, ModelTask.OCR: - return TextDetector - - case ModelSource.PADDLE, ModelType.RECOGNITION, ModelTask.OCR: - return TextRecognizer - - case _: - raise ValueError(f"Unknown model combination: {source}, {model_type}, {model_task}") - - -def get_model_deps(model_name: str, model_type: ModelType, model_task: ModelTask) -> list[tuple[ModelType, ModelTask]]: - return get_model_class(model_name, model_type, model_task).depends diff --git a/machine-learning/immich_ml/models/base.py b/machine-learning/immich_ml/models/base.py index 2e999c859..66dc1026c 100644 --- a/machine-learning/immich_ml/models/base.py +++ b/machine-learning/immich_ml/models/base.py @@ -1,10 +1,12 @@ from __future__ import annotations from abc import ABC, abstractmethod +from collections.abc import Hashable +from dataclasses import dataclass from pathlib import Path from shutil import rmtree from threading import Lock -from typing import Any, ClassVar +from typing import Any, ClassVar, Self from huggingface_hub import snapshot_download @@ -13,9 +15,10 @@ import immich_ml.sessions.rknn as rknn from immich_ml.sessions.ort import OrtSession from ..config import clean_name, log, settings -from ..schemas import ModelFormat, ModelIdentity, ModelSession, ModelTask, ModelType +from ..schemas import ModelFormat, ModelIdentity, ModelSession, ModelSource, ModelTask, ModelType, Options from ..sessions.ann import AnnSession from ..sessions.policy import ShapePolicy +from .constants import get_model_source _IGNORED_PATTERNS: dict[ModelFormat, list[str]] = { ModelFormat.ONNX: ["*.armnn", "*.rknn"], @@ -24,11 +27,10 @@ _IGNORED_PATTERNS: dict[ModelFormat, list[str]] = { } -class InferenceModel(ABC): +class InferenceModel[O: Options](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, ...]] = () + sources: ClassVar[tuple[ModelSource, ...]] # 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 @@ -39,7 +41,6 @@ class InferenceModel(ABC): cache_dir: Path | str | None = None, model_format: ModelFormat | None = None, session: ModelSession | None = None, - **model_kwargs: Any, ) -> None: self.loaded = session is not None self.load_attempts = 0 @@ -50,6 +51,15 @@ class InferenceModel(ABC): if session is not None: self.session = session + @classmethod + def graph(cls, options: O) -> Hashable: + """What of its options the graph is built for: an instance serves only requests that agree on it.""" + return None + + @classmethod + def create(cls, model_name: str, options: O) -> Self: + return cls(model_name) + def download(self) -> None: if self.cached: return @@ -87,9 +97,9 @@ class InferenceModel(ABC): self.load() self.session.warm() - def predict(self, *inputs: Any, **model_kwargs: Any) -> Any: + def predict(self, *inputs: Any, options: O) -> Any: self.load() - return self._predict(*inputs, **model_kwargs) + return self._predict(*inputs, options=options) @abstractmethod def _predict(self, *inputs: Any, **model_kwargs: Any) -> Any: ... @@ -180,3 +190,17 @@ class InferenceModel(ABC): return ModelFormat.ARMNN else: return ModelFormat.ONNX + + +@dataclass(frozen=True) +class InferenceEntry[O: Options]: + model: type[InferenceModel[O]] + name: str + options: O + + def __post_init__(self) -> None: + if get_model_source(self.name) not in self.model.sources: + raise ValueError(f"'{self.name}' is not a model {self.model.__name__} runs") + + def __str__(self) -> str: + return f"{self.model.__name__}(name={self.name!r}, options={self.options})" diff --git a/machine-learning/immich_ml/models/cache.py b/machine-learning/immich_ml/models/cache.py index a433411b9..db2ae46e6 100644 --- a/machine-learning/immich_ml/models/cache.py +++ b/machine-learning/immich_ml/models/cache.py @@ -1,25 +1,23 @@ import asyncio +from collections.abc import Hashable from typing import Any from immich_ml import allocator -from immich_ml.models import get_model_class -from immich_ml.models.base import InferenceModel +from immich_ml.models.base import InferenceEntry, InferenceModel +from immich_ml.schemas import Options -from ..schemas import ModelTask, ModelType +Key = tuple[type[InferenceModel[Any]], str, Hashable] class ModelCache: def __init__(self) -> None: - self._models: dict[tuple[Any, ...], InferenceModel] = {} - self._expiries: dict[tuple[Any, ...], asyncio.TimerHandle] = {} + self._models: dict[Key, InferenceModel[Any]] = {} + self._expiries: dict[Key, asyncio.TimerHandle] = {} - def get( - 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, *(options.get(option) for option in model_cls.graph_options)) + def get[O: Options](self, entry: InferenceEntry[O], ttl: int | None = None) -> InferenceModel[O]: + key = (entry.model, entry.name, entry.model.graph(entry.options)) if key not in self._models: - self._models[key] = model_cls(model_name, **options) + self._models[key] = entry.model.create(entry.name, entry.options) elif key not in self._expiries: ttl = None # a model that was preloaded stays if ttl: @@ -36,7 +34,7 @@ class ModelCache: self._expiries.pop(key).cancel() allocator.release() - def _evict(self, key: tuple[Any, ...]) -> None: + def _evict(self, key: Key) -> None: if key not in self._models: return del self._models[key], self._expiries[key] diff --git a/machine-learning/immich_ml/models/clip/textual.py b/machine-learning/immich_ml/models/clip/textual.py index d5ad3400f..789d13394 100644 --- a/machine-learning/immich_ml/models/clip/textual.py +++ b/machine-learning/immich_ml/models/clip/textual.py @@ -12,7 +12,7 @@ from immich_ml.config import log from immich_ml.models.base import InferenceModel from immich_ml.models.constants import WEBLATE_TO_FLORES200 from immich_ml.models.transforms import clean_text, serialize_np_array -from immich_ml.schemas import ModelSession, ModelTask, ModelType +from immich_ml.schemas import ModelSession, ModelSource, ModelTask, ModelType, TextualOptions def _mapped(model_path: Path, name: str) -> NDArray[Any]: @@ -25,13 +25,13 @@ def _mapped(model_path: Path, name: str) -> NDArray[Any]: return np.memmap(model_path.parent / data["location"], dtype=dtype, mode="r", offset=offset, shape=shape) -class BaseCLIPTextualEncoder(InferenceModel): +class BaseCLIPTextualEncoder(InferenceModel[TextualOptions]): depends = [] identity = (ModelType.TEXTUAL, ModelTask.SEARCH) threads = 4 # lower search latency - def _predict(self, inputs: str, language: str | None = None) -> str: - tokens: dict[str, NDArray[Any]] = self.tokenize(inputs, language=language) + def _predict(self, inputs: str, options: TextualOptions) -> str: + tokens: dict[str, NDArray[Any]] = self.tokenize(inputs, language=options.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()))] @@ -99,6 +99,8 @@ class BaseCLIPTextualEncoder(InferenceModel): class OpenClipTextualEncoder(BaseCLIPTextualEncoder): + sources = (ModelSource.OPENCLIP,) + def _load_tokenizer(self) -> Tokenizer: context_length: int = self.text_cfg.get("context_length", 77) pad_token: str = self.tokenizer_cfg["pad_token"] @@ -129,6 +131,8 @@ class OpenClipTextualEncoder(BaseCLIPTextualEncoder): class MClipTextualEncoder(OpenClipTextualEncoder): + sources = (ModelSource.MCLIP,) + def tokenize(self, text: str, language: str | None = None) -> dict[str, NDArray[np.int32]]: text = clean_text(text, canonicalize=self.canonicalize) tokens: Encoding = self.tokenizer.encode(text) diff --git a/machine-learning/immich_ml/models/clip/visual.py b/machine-learning/immich_ml/models/clip/visual.py index fe627b738..c4774b779 100644 --- a/machine-learning/immich_ml/models/clip/visual.py +++ b/machine-learning/immich_ml/models/clip/visual.py @@ -19,14 +19,14 @@ from immich_ml.models.transforms import ( serialize_np_array, to_numpy, ) -from immich_ml.schemas import ModelGraph, ModelInput, ModelSession, ModelTask, ModelType +from immich_ml.schemas import ModelGraph, ModelInput, ModelSession, ModelSource, ModelTask, ModelType, VisualOptions -class BaseCLIPVisualEncoder(InferenceModel): +class BaseCLIPVisualEncoder(InferenceModel[VisualOptions]): depends = [] identity = (ModelType.VISUAL, ModelTask.SEARCH) - def _predict(self, inputs: Image.Image | bytes) -> str: + def _predict(self, inputs: Image.Image | bytes, options: VisualOptions) -> str: image = decode_pil(inputs) session = self.session.for_shape(self.shape_policy.dims[0]) res: NDArray[np.float32] = session.run(None, self.transform(session, image))[0][0] @@ -60,6 +60,8 @@ class BaseCLIPVisualEncoder(InferenceModel): class OpenClipVisualEncoder(BaseCLIPVisualEncoder): + sources = (ModelSource.OPENCLIP, ModelSource.MCLIP) + def _load(self) -> ModelSession: size: list[int] | int = self.preprocess_cfg["size"] self.size = size[0] if isinstance(size, list) else size diff --git a/machine-learning/immich_ml/models/facial_recognition/detection.py b/machine-learning/immich_ml/models/facial_recognition/detection.py index 689b5c870..2ef5e50c6 100644 --- a/machine-learning/immich_ml/models/facial_recognition/detection.py +++ b/machine-learning/immich_ml/models/facial_recognition/detection.py @@ -4,18 +4,19 @@ from numpy.typing import NDArray from immich_ml.models.base import InferenceModel from immich_ml.models.transforms import decode_pil, letterbox, normalize, widen -from immich_ml.schemas import FaceDetectionOutput, ModelTask, ModelType, Shape +from immich_ml.schemas import FaceDetectionOptions, FaceDetectionOutput, ModelSource, ModelTask, ModelType, Shape from immich_ml.sessions.policy import ShapePolicy from ._ops import decode_scrfd, nms -class FaceDetector(InferenceModel): +class FaceDetector(InferenceModel[FaceDetectionOptions]): depends = [] identity = (ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION) + sources = (ModelSource.INSIGHTFACE,) shape_policy = ShapePolicy(dims=(Shape(batch=1, height=DET_SIZE, width=DET_SIZE),)) - def _predict(self, inputs: NDArray[np.uint8] | bytes, minScore: float) -> FaceDetectionOutput: + def _predict(self, inputs: NDArray[np.uint8] | bytes, options: FaceDetectionOptions) -> FaceDetectionOutput: session = self.session.for_shape(self.shape_policy.dims[0]) canvas, scale = letterbox(decode_pil(inputs), DET_SIZE) blob: NDArray[np.float32] | NDArray[np.uint8] = ( @@ -27,7 +28,7 @@ class FaceDetector(InferenceModel): 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 + candidates = scores >= options.min_score scores, boxes, kps = scores[candidates], boxes[candidates] / scale, kps[candidates] / scale keep = nms(boxes, scores) diff --git a/machine-learning/immich_ml/models/facial_recognition/recognition.py b/machine-learning/immich_ml/models/facial_recognition/recognition.py index 9f8058dc2..b957bbd5e 100644 --- a/machine-learning/immich_ml/models/facial_recognition/recognition.py +++ b/machine-learning/immich_ml/models/facial_recognition/recognition.py @@ -7,15 +7,24 @@ from PIL import Image 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, ModelTask, ModelType, Shape +from immich_ml.schemas import ( + FaceDetectionOutput, + FaceRecognitionOptions, + FacialRecognitionOutput, + ModelSource, + ModelTask, + ModelType, + Shape, +) from immich_ml.sessions.policy import ShapePolicy, batches, runs from ._ops import ALIGNED_SIZE, align_face -class FaceRecognizer(InferenceModel): +class FaceRecognizer(InferenceModel[FaceRecognitionOptions]): depends = [(ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION)] identity = (ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION) + sources = (ModelSource.INSIGHTFACE,) def __init__(self, model_name: str, **model_kwargs: Any) -> None: super().__init__(model_name, **model_kwargs) @@ -23,7 +32,10 @@ class FaceRecognizer(InferenceModel): self.shape_policy = ShapePolicy(dims=tuple(Shape(batch) for batch in sizes)) def _predict( - self, inputs: NDArray[np.uint8] | bytes | Image.Image, faces: FaceDetectionOutput + self, + inputs: NDArray[np.uint8] | bytes | Image.Image, + faces: FaceDetectionOutput, + options: FaceRecognitionOptions, ) -> FacialRecognitionOutput: if faces["boxes"].shape[0] == 0: return [] diff --git a/machine-learning/immich_ml/models/ocr/detection.py b/machine-learning/immich_ml/models/ocr/detection.py index 563f8a833..de4de264f 100644 --- a/machine-learning/immich_ml/models/ocr/detection.py +++ b/machine-learning/immich_ml/models/ocr/detection.py @@ -1,5 +1,5 @@ -from collections.abc import Sequence -from typing import Any +from collections.abc import Hashable, Sequence +from typing import Any, Self import numpy as np from immich_model.constants import ocr_canvases @@ -7,7 +7,7 @@ from numpy.typing import NDArray from PIL import Image from immich_ml.models.transforms import widen -from immich_ml.schemas import ModelGraph, ModelTask, ModelType, Shape +from immich_ml.schemas import ModelGraph, ModelTask, ModelType, Shape, TextDetectionOptions from immich_ml.sessions.policy import ShapePolicy from .legacy import TextModel @@ -15,17 +15,17 @@ from .postprocess import DBPostProcess from .schemas import TextDetectionOutput -class TextDetector(TextModel): +class TextDetector(TextModel[TextDetectionOptions]): depends = [] identity = (ModelType.DETECTION, ModelTask.OCR) - graph_options = ("maxResolution",) - def __init__(self, model_name: str, **model_kwargs: Any) -> None: + def __init__( + self, model_name: str, max_resolution: int = TextDetectionOptions.max_resolution, **model_kwargs: Any + ) -> None: 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}") + canvases = tuple(Shape(batch=1, **canvas) for canvas in ocr_canvases(max_resolution)) + self.shape_policy = ShapePolicy(dims=canvases, label=f"res{max_resolution}") self.scale = np.float32(1.0 / 127.5) # (x/255 - 0.5) / 0.5 self._empty: TextDetectionOutput = { "boxes": np.empty(0, dtype=np.float32), @@ -33,18 +33,24 @@ class TextDetector(TextModel): } self.postprocess = DBPostProcess(thresh=0.3, max_candidates=1000, unclip_ratio=1.6, use_dilation=True) - def _predict( - self, inputs: Image.Image, maxResolution: int = 736, minScore: float = 0.5, scoreMode: str = "fast" - ) -> TextDetectionOutput: + @classmethod + def graph(cls, options: TextDetectionOptions) -> Hashable: + return options.max_resolution + + @classmethod + def create(cls, model_name: str, options: TextDetectionOptions) -> Self: + return cls(model_name, max_resolution=options.max_resolution) + + def _predict(self, inputs: Image.Image, options: TextDetectionOptions) -> TextDetectionOutput: width, height = inputs.size if width < 32 or height < 32: return self._empty - session, image, content = self._transform(inputs, maxResolution) + session, image, content = self._transform(inputs, options.max_resolution) 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) + boxes, scores = self.postprocess(probs[: content[0], : content[1]], (height, width), options.min_score) if len(boxes) == 0: return self._empty order = self.reading_order(boxes) diff --git a/machine-learning/immich_ml/models/ocr/legacy.py b/machine-learning/immich_ml/models/ocr/legacy.py index d17ec0339..db76feb1b 100644 --- a/machine-learning/immich_ml/models/ocr/legacy.py +++ b/machine-learning/immich_ml/models/ocr/legacy.py @@ -1,16 +1,16 @@ -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 +from immich_ml.schemas import ModelFormat, ModelSource, ModelType, Options # 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) +class TextModel[O: Options](InferenceModel[O]): + sources = (ModelSource.PADDLE,) + + @property + def _model_format_default(self) -> ModelFormat: + # the older exports come in no other format + return ModelFormat.ONNX if settings.legacy_models else super()._model_format_default def download(self) -> None: if not settings.legacy_models or self.cached: diff --git a/machine-learning/immich_ml/models/ocr/recognition.py b/machine-learning/immich_ml/models/ocr/recognition.py index 7cea06df2..6ce3ccd5c 100644 --- a/machine-learning/immich_ml/models/ocr/recognition.py +++ b/machine-learning/immich_ml/models/ocr/recognition.py @@ -9,7 +9,7 @@ from numpy.typing import NDArray from PIL import Image from immich_ml.config import settings -from immich_ml.schemas import ModelGraph, ModelSession, ModelTask, ModelType, Shape +from immich_ml.schemas import ModelGraph, ModelSession, ModelTask, ModelType, Shape, TextRecognitionOptions from immich_ml.sessions.policy import ShapePolicy, batches, runs from .ctc import GREEDY, CtcDecoder, probabilities @@ -20,7 +20,7 @@ REC_HEIGHT = 48 SCALE = np.float32(1.0 / 127.5) -class TextRecognizer(TextModel): +class TextRecognizer(TextModel[TextRecognitionOptions]): depends = [(ModelType.DETECTION, ModelTask.OCR)] identity = (ModelType.RECOGNITION, ModelTask.OCR) @@ -53,7 +53,9 @@ class TextRecognizer(TextModel): else CtcDecoder.from_file(self.model_dir / "charset.txt", greedy) ) - def _predict(self, img: Image.Image, texts: TextDetectionOutput, minScore: float = 0.9) -> TextRecognitionOutput: + def _predict( + self, img: Image.Image, texts: TextDetectionOutput, options: TextRecognitionOptions + ) -> TextRecognitionOutput: boxes, box_scores = texts["boxes"], texts["scores"] if boxes.shape[0] == 0: return self._empty @@ -101,7 +103,7 @@ class TextRecognizer(TextModel): boxes[:, :, 0] /= img.width boxes[:, :, 1] /= img.height - valid = score_list > minScore + valid = score_list > options.min_score valid_list = valid.tolist() return { "box": boxes.reshape(-1, 8)[valid].reshape(-1), diff --git a/machine-learning/immich_ml/pipeline.py b/machine-learning/immich_ml/pipeline.py new file mode 100644 index 000000000..a970723bf --- /dev/null +++ b/machine-learning/immich_ml/pipeline.py @@ -0,0 +1,85 @@ +from collections.abc import Iterator +from dataclasses import dataclass, field +from typing import Annotated, Any + +from pydantic import ConfigDict, Field, with_config + +from .models.base import InferenceEntry, InferenceModel +from .models.clip.textual import BaseCLIPTextualEncoder, MClipTextualEncoder, OpenClipTextualEncoder +from .models.clip.visual import OpenClipVisualEncoder +from .models.constants import get_model_source +from .models.facial_recognition.detection import FaceDetector +from .models.facial_recognition.recognition import FaceRecognizer +from .models.ocr.detection import TextDetector +from .models.ocr.recognition import TextRecognizer +from .schemas import ( + FaceDetectionOptions, + FaceRecognitionOptions, + ModelSource, + Options, + TextDetectionOptions, + TextRecognitionOptions, + TextualOptions, + VisualOptions, +) + +STRICT = ConfigDict(extra="forbid") + + +def textual(model_name: str) -> type[BaseCLIPTextualEncoder]: + return MClipTextualEncoder if get_model_source(model_name) == ModelSource.MCLIP else OpenClipTextualEncoder + + +@with_config(STRICT) +@dataclass(frozen=True) +class Slot[O: Options]: + model_name: Annotated[str, Field(alias="modelName")] + options: Annotated[O, Field(default_factory=dict, validate_default=True)] # left out, it takes their defaults + + def entry(self, model: type[InferenceModel[O]]) -> InferenceEntry[O]: + return InferenceEntry(model, self.model_name, self.options) + + +@with_config(STRICT) +@dataclass(frozen=True) +class Clip: + visual: Slot[VisualOptions] | None = None + textual: Slot[TextualOptions] | None = None + + +@with_config(STRICT) +@dataclass(frozen=True) +class FacialRecognition: + detection: Slot[FaceDetectionOptions] | None = None + recognition: Slot[FaceRecognitionOptions] | None = None + + +@with_config(STRICT) +@dataclass(frozen=True) +class Ocr: + detection: Slot[TextDetectionOptions] | None = None + recognition: Slot[TextRecognitionOptions] | None = None + + +@with_config(STRICT) +@dataclass(frozen=True) +class PipelineRequest: + clip: Clip = field(default_factory=Clip) + facial_recognition: Annotated[FacialRecognition, Field(alias="facial-recognition")] = field( + default_factory=FacialRecognition + ) + ocr: Ocr = field(default_factory=Ocr) + + def entries(self) -> Iterator[InferenceEntry[Any]]: + if (visual := self.clip.visual) is not None: + yield visual.entry(OpenClipVisualEncoder) + if (text := self.clip.textual) is not None: + yield text.entry(textual(text.model_name)) + if (detection := self.facial_recognition.detection) is not None: + yield detection.entry(FaceDetector) + if (recognition := self.facial_recognition.recognition) is not None: + yield recognition.entry(FaceRecognizer) + if (boxes := self.ocr.detection) is not None: + yield boxes.entry(TextDetector) + if (reading := self.ocr.recognition) is not None: + yield reading.entry(TextRecognizer) diff --git a/machine-learning/immich_ml/schemas.py b/machine-learning/immich_ml/schemas.py index 466e2553b..2364d7ccc 100644 --- a/machine-learning/immich_ml/schemas.py +++ b/machine-learning/immich_ml/schemas.py @@ -1,10 +1,11 @@ from collections.abc import Mapping, Sequence from dataclasses import asdict, dataclass from enum import Enum -from typing import Any, Literal, Protocol, TypeAlias, TypeVar +from typing import Annotated, Any, Protocol, TypeAlias import numpy as np import numpy.typing as npt +from pydantic import ConfigDict, Field from typing_extensions import TypedDict @@ -121,25 +122,39 @@ class DetectedFace(TypedDict): FacialRecognitionOutput = list[DetectedFace] -class PipelineEntry(TypedDict): - modelName: str - options: dict[str, Any] +class Options: + __pydantic_config__ = ConfigDict(extra="ignore") # so that an older service ignores an option a newer server adds -PipelineRequest = dict[ModelTask, dict[ModelType, PipelineEntry]] +@dataclass(frozen=True) +class VisualOptions(Options): + pass -class InferenceEntry(TypedDict): - name: str - task: ModelTask - type: ModelType - options: dict[str, Any] +@dataclass(frozen=True) +class TextualOptions(Options): + language: str | None = None -InferenceEntries = tuple[list[InferenceEntry], list[InferenceEntry]] +@dataclass(frozen=True) +class FaceDetectionOptions(Options): + min_score: Annotated[float, Field(alias="minScore")] = 0.7 -InferenceResponse = dict[ModelTask | Literal["imageHeight"] | Literal["imageWidth"], Any] +@dataclass(frozen=True) +class FaceRecognitionOptions(Options): + pass -T = TypeVar("T") +@dataclass(frozen=True) +class TextDetectionOptions(Options): + max_resolution: Annotated[int, Field(alias="maxResolution")] = 736 + min_score: Annotated[float, Field(alias="minScore")] = 0.5 + + +@dataclass(frozen=True) +class TextRecognitionOptions(Options): + min_score: Annotated[float, Field(alias="minScore")] = 0.9 + + +InferenceResponse = dict[str, Any] diff --git a/machine-learning/mise.toml b/machine-learning/mise.toml index 6470a71b7..d57e9fe7c 100644 --- a/machine-learning/mise.toml +++ b/machine-learning/mise.toml @@ -6,16 +6,16 @@ uv = "0.8.15" run = "uv sync --locked" [tasks.lint] -run = "uv run ruff check immich_ml" +run = "uv run ruff check immich_ml scripts test_main.py conftest.py" [tasks.test] run = "uv run pytest --cov=immich_ml --cov-report term-missing" [tasks.format] -run = "uv run ruff format immich_ml" +run = "uv run ruff format immich_ml scripts test_main.py conftest.py" [tasks.check] -run = "uv run mypy --strict immich_ml/" +run = "uv run mypy --strict immich_ml/ scripts/ test_main.py conftest.py" [tasks.ci-unit] run = [ diff --git a/machine-learning/test_main.py b/machine-learning/test_main.py index d4a77dd6a..d4e2e9360 100644 --- a/machine-learning/test_main.py +++ b/machine-learning/test_main.py @@ -28,15 +28,17 @@ from immich_ml import allocator from immich_ml.config import MaxBatchSize, PreloadModelData, Settings, settings from immich_ml.main import ( MODEL_FILE_ERRORS, + PIPELINE_REQUEST, app, attempt, + get_entries, 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.base import InferenceEntry, InferenceModel from immich_ml.models.cache import ModelCache from immich_ml.models.clip.textual import MClipTextualEncoder, OpenClipTextualEncoder from immich_ml.models.clip.visual import OpenClipVisualEncoder @@ -45,7 +47,19 @@ 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, ModelTask, ModelType, Shape +from immich_ml.schemas import ( + FaceDetectionOptions, + FaceRecognitionOptions, + ModelFormat, + ModelSource, + ModelTask, + ModelType, + Shape, + TextDetectionOptions, + TextRecognitionOptions, + TextualOptions, + VisualOptions, +) from immich_ml.sessions.ann import AnnSession from immich_ml.sessions.ort import Device, GraphSpec, OrtSession, flush_denormals, fresh, prepared from immich_ml.sessions.policy import ShapePolicy, batches, runs @@ -123,7 +137,7 @@ class TestBase: 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 + "PP-OCRv5_mobile", cache_dir="/cache", model_format=ModelFormat.RKNN, max_resolution=1088 ) assert detector.model_path == Path("/cache/detection/rknpu/rk3588/res1088/model.rknn") @@ -1006,7 +1020,7 @@ class TestCLIP: mocked.run.return_value = [[self.embedding]] clip_encoder = OpenClipVisualEncoder("ViT-B-32__openai", cache_dir="test_cache") - embedding_str = clip_encoder.predict(pil_image) + embedding_str = clip_encoder.predict(pil_image, options=VisualOptions()) assert isinstance(embedding_str, str) embedding = orjson.loads(embedding_str) assert isinstance(embedding, list) @@ -1034,7 +1048,7 @@ class TestCLIP: 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) + OpenClipVisualEncoder("ViT-B-32__openai", cache_dir="test_cache").predict(pil_image, options=VisualOptions()) fed = session.run.call_args.args[1]["image"] assert fed.dtype == expected_dtype @@ -1056,7 +1070,7 @@ class TestCLIP: # 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) + OpenClipVisualEncoder("ViT-B-32__openai", cache_dir="test_cache").predict(image, options=VisualOptions()) assert session.run.call_args.args[1]["image"][0, 0, 0].tolist() == [255, 255, 255] @@ -1078,7 +1092,7 @@ class TestCLIP: mocker.patch("immich_ml.models.clip.textual.Tokenizer.from_file", autospec=True) clip_encoder = OpenClipTextualEncoder("ViT-B-32__openai", cache_dir="test_cache") - embedding_str = clip_encoder.predict("test search query") + embedding_str = clip_encoder.predict("test search query", options=TextualOptions()) assert isinstance(embedding_str, str) embedding = orjson.loads(embedding_str) assert isinstance(embedding, list) @@ -1117,7 +1131,7 @@ class TestCLIP: 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") + encoder.predict("test search query", options=TextualOptions()) feed = session.run.call_args.args[1] assert list(feed) == ["text", "token_embeds"] # a binary reads its inputs in order @@ -1352,7 +1366,7 @@ 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) +M = TypeVar("M", bound=InferenceModel[Any]) def given_options(ort_session: mock.Mock) -> Any: @@ -1400,7 +1414,7 @@ class TestFaceRecognition: 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) + faces = face_detector.predict(Image.new("RGB", (640, 640)), options=FaceDetectionOptions(min_score=0.7)) assert isinstance(faces, dict) assert set(faces) == {"boxes", "scores", "landmarks"} @@ -1421,8 +1435,18 @@ class TestFaceRecognition: face_detector.session = session # the threshold is a request parameter, so the same loaded model must honour both - assert face_detector.predict(Image.new("RGB", (640, 640)), minScore=0.7)["boxes"].shape[0] == 1 - assert face_detector.predict(Image.new("RGB", (640, 640)), minScore=0.4)["boxes"].shape[0] == 2 + assert ( + face_detector.predict(Image.new("RGB", (640, 640)), options=FaceDetectionOptions(min_score=0.7))[ + "boxes" + ].shape[0] + == 1 + ) + assert ( + face_detector.predict(Image.new("RGB", (640, 640)), options=FaceDetectionOptions(min_score=0.4))[ + "boxes" + ].shape[0] + == 2 + ) def test_detection_scales_boxes_back_to_the_original_image( self, stub_session: Callable[..., mock.Mock], mocker: MockerFixture @@ -1434,7 +1458,7 @@ class TestFaceRecognition: face_detector.session = session # a 320x320 image is letterboxed up to 640, so coordinates come back halved - faces = face_detector.predict(Image.new("RGB", (320, 320)), minScore=0.7) + faces = face_detector.predict(Image.new("RGB", (320, 320)), options=FaceDetectionOptions(min_score=0.7)) assert faces["boxes"].tolist() == [[v / 2 for v in expected_box(10, 10)]] assert np.allclose(faces["landmarks"][0], expected_landmarks(10, 10) / 2) @@ -1459,7 +1483,9 @@ class TestFaceRecognition: 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}) + faces = face_recognizer.predict( + image, {"boxes": bbox, "landmarks": kpss, "scores": scores}, options=FaceRecognitionOptions() + ) assert isinstance(faces, list) assert len(faces) == num_faces @@ -1497,7 +1523,7 @@ class TestFaceRecognition: 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) + faces = face_detector.predict(Image.new("RGB", (640, 640)), options=FaceDetectionOptions(min_score=0.7)) fed = face_detector.session.run.call_args.args[1]["input.1"] assert fed.dtype == expected_dtype @@ -1531,7 +1557,7 @@ class TestFaceRecognition: "landmarks": np.stack([expected_landmarks(1, 1), expected_landmarks(8, 8)]), } - face_recognizer.predict(Image.new("RGB", (200, 200)), faces) + face_recognizer.predict(Image.new("RGB", (200, 200)), faces, options=FaceRecognitionOptions()) fed = face_recognizer.session.run.call_args.args[1]["input.1"] assert fed.dtype == expected_dtype @@ -1550,7 +1576,7 @@ class TestFaceRecognition: "scores": np.empty(0, dtype=np.float32), } - assert face_recognizer.predict(pil_image, empty) == [] + assert face_recognizer.predict(pil_image, empty, options=FaceRecognitionOptions()) == [] session.run.assert_not_called() def test_recognition_batches_when_batch_size_is_set( @@ -1569,7 +1595,7 @@ class TestFaceRecognition: "landmarks": (np.random.rand(num_faces, 5, 2) * 100).astype(np.float32), "scores": np.array([0.67] * num_faces, dtype=np.float32), } - assert len(face_recognizer.predict(pil_image, faces)) == num_faces + assert len(face_recognizer.predict(pil_image, faces, options=FaceRecognitionOptions())) == num_faces 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] @@ -1638,10 +1664,10 @@ class TestOcr: 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 - assert len(text_detector._predict(image, minScore=0.9)["boxes"]) == 0 + assert len(text_detector._predict(image, TextDetectionOptions(min_score=0.5))["boxes"]) == 1 + assert len(text_detector._predict(image, TextDetectionOptions(min_score=0.9))["boxes"]) == 0 # the default must be unaffected by the request that just ran - assert len(text_detector._predict(image)["boxes"]) == 1 + assert len(text_detector._predict(image, TextDetectionOptions())["boxes"]) == 1 def test_fetches_the_older_exports_from_where_rapidocr_hosts_them( self, tmp_path: Path, snapshot_download: mock.Mock, mocker: MockerFixture @@ -1671,7 +1697,7 @@ class TestOcr: 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) + text_recognizer._predict(image, texts, TextRecognitionOptions()) fed = session.run.call_args.args[1]["input.1"] assert fed.dtype == np.uint8 and fed.shape == (1, 48, fed_width, 3) @@ -1681,7 +1707,7 @@ class TestOcr: 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 = TextDetector("PP-OCRv5_mobile", max_resolution=64, cache_dir=path) text_detector.session = stub_session( (1, 64, 128, 3), outputs=[np.zeros((1, 64, 128), dtype=np.float32)], @@ -1690,7 +1716,7 @@ class TestOcr: 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) + text_detector._predict(Image.new("RGB", (100, 100), (10, 20, 30)), TextDetectionOptions(max_resolution=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 @@ -1714,7 +1740,7 @@ class TestOcr: 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) + text_recognizer._predict(image, texts, TextRecognitionOptions()) 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 @@ -1736,9 +1762,9 @@ class TestOcr: def texts() -> Any: # _predict normalizes the boxes in place, so each call needs its own return {"boxes": box.copy(), "scores": np.array([0.9], dtype=np.float32)} - assert text_recognizer._predict(image, texts(), minScore=0.7)["text"] == ["hello"] + assert text_recognizer._predict(image, texts(), TextRecognitionOptions(min_score=0.7))["text"] == ["hello"] # 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"] == [] + assert text_recognizer._predict(image, texts(), TextRecognitionOptions())["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) @@ -1792,59 +1818,62 @@ class TestOcr: assert {shape.batch for shape in text_recognizer.shape_policy.dims} == {1, 6} +def stub_model() -> mock.MagicMock: + """A model class that makes mocks, where only what the cache does with its entries matters.""" + model = mock.MagicMock(sources=(ModelSource.INSIGHTFACE,)) + model.graph.return_value = None + return model + + @pytest.mark.asyncio class TestCache: - async def test_caches(self, mock_get_model: mock.Mock) -> None: - model_cache = ModelCache() - model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION) - model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION) + async def test_caches(self) -> None: + model_cache, model = ModelCache(), stub_model() + model_cache.get(InferenceEntry(model, "buffalo_l", FaceRecognitionOptions())) + model_cache.get(InferenceEntry(model, "buffalo_l", FaceRecognitionOptions())) assert len(model_cache._models) == 1 - mock_get_model.return_value.assert_called_once() + model.create.assert_called_once() async def test_separate_instances_per_graph_option(self) -> None: - # the real TextDetector, so that dropping its graph_options fails this + # the real TextDetector, so that keying it on anything but its resolution 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 + for options in ( + TextDetectionOptions(), + TextDetectionOptions(max_resolution=736, min_score=0.3), + TextDetectionOptions(max_resolution=1088), + ): + model_cache.get(InferenceEntry(TextDetector, "PP-OCRv5_mobile", options)) + # the resolution picks the graphs a detector builds, so instances cannot share them; a score does not 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") - mock_get_model.return_value.assert_called_once_with("test_model_name", cache_dir="test_cache") + async def test_creates_the_model_from_the_request_options(self) -> None: + model = stub_model() + ModelCache().get(InferenceEntry(model, "buffalo_l", FaceDetectionOptions(min_score=0.3))) + model.create.assert_called_once_with("buffalo_l", FaceDetectionOptions(min_score=0.3)) - async def test_different_clip(self, mock_get_model: mock.Mock) -> None: + async def test_separate_instances_per_model(self) -> None: model_cache = ModelCache() - model_cache.get("test_model_name", ModelType.VISUAL, ModelTask.SEARCH) - model_cache.get("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH) - assert mock_get_model.call_args_list == [ - mock.call("test_model_name", ModelType.VISUAL, ModelTask.SEARCH), - mock.call("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH), - ] + visual = model_cache.get(InferenceEntry(OpenClipVisualEncoder, "ViT-B-32__openai", VisualOptions())) + text = model_cache.get(InferenceEntry(OpenClipTextualEncoder, "ViT-B-32__openai", TextualOptions())) + assert isinstance(visual, OpenClipVisualEncoder) and isinstance(text, OpenClipTextualEncoder) assert len(model_cache._models) == 2 - async def test_lets_go_of_a_model_unused_for_its_ttl_and_returns_its_memory( - self, mock_get_model: mock.Mock, mocker: MockerFixture - ) -> None: + async def test_lets_go_of_a_model_unused_for_its_ttl_and_returns_its_memory(self, mocker: MockerFixture) -> None: events: list[str] = [] - class Model: - graph_options = () - - def __init__(self, *args: Any, **kwargs: Any) -> None: - pass - + class Loaded: def __del__(self) -> None: events.append("destroyed") - mock_get_model.return_value = Model + model = stub_model() + model.create.side_effect = lambda name, options: Loaded() + mocker.patch("immich_ml.models.cache.allocator.release", side_effect=lambda: events.append("released")) loop = mocker.patch("immich_ml.models.cache.asyncio.get_running_loop").return_value model_cache = ModelCache() - model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100) + model_cache.get(InferenceEntry(model, "buffalo_l", FaceRecognitionOptions()), ttl=100) delay, evict, key = loop.call_later.call_args.args evict(key) @@ -1852,49 +1881,45 @@ class TestCache: assert events == ["destroyed", "released"] assert not model_cache._models - async def test_a_use_starts_the_ttl_over(self, mock_get_model: mock.Mock, mocker: MockerFixture) -> None: + async def test_a_use_starts_the_ttl_over(self, mocker: MockerFixture) -> None: loop = mocker.patch("immich_ml.models.cache.asyncio.get_running_loop").return_value - model_cache = ModelCache() + model_cache, entry = ModelCache(), InferenceEntry(stub_model(), "buffalo_l", FaceRecognitionOptions()) - model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100) - model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100) + model_cache.get(entry, ttl=100) + model_cache.get(entry, ttl=100) loop.call_later.return_value.cancel.assert_called_once() assert loop.call_later.call_count == 2 - async def test_keeps_a_preloaded_model_whatever_ttl_a_request_names( - self, mock_get_model: mock.Mock, mocker: MockerFixture - ) -> None: + async def test_keeps_a_preloaded_model_whatever_ttl_a_request_names(self, mocker: MockerFixture) -> None: loop = mocker.patch("immich_ml.models.cache.asyncio.get_running_loop").return_value - model_cache = ModelCache() + model_cache, entry = ModelCache(), InferenceEntry(stub_model(), "buffalo_l", FaceRecognitionOptions()) - model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION) - model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100) + model_cache.get(entry) + model_cache.get(entry, ttl=100) loop.call_later.assert_not_called() async def test_loads_mclip(self) -> None: - model_cache = ModelCache() + request = PIPELINE_REQUEST.validate_python( + {"clip": {"textual": {"modelName": "XLM-Roberta-Large-Vit-B-32", "options": {}}}} + ) - model = model_cache.get("XLM-Roberta-Large-Vit-B-32", ModelType.TEXTUAL, ModelTask.SEARCH) + [entry] = request.entries() + model = ModelCache().get(entry) assert isinstance(model, MClipTextualEncoder) assert model.model_name == "XLM-Roberta-Large-Vit-B-32" - async def test_raises_exception_if_invalid_model_type(self) -> None: - invalid: Any = SimpleNamespace(value="invalid") - model_cache = ModelCache() - + async def test_refuses_a_model_its_slot_does_not_run(self) -> None: with pytest.raises(ValueError): - model_cache.get("XLM-Roberta-Large-Vit-B-32", ModelType.TEXTUAL, invalid) - - async def test_raises_exception_if_unknown_model_name(self) -> None: - model_cache = ModelCache() + InferenceEntry(OpenClipVisualEncoder, "buffalo_l", VisualOptions()) + async def test_refuses_an_unknown_model_name(self) -> None: with pytest.raises(ValueError): - model_cache.get("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH) + InferenceEntry(OpenClipTextualEncoder, "test_model_name", TextualOptions()) - async def test_preloads_clip_models(self, monkeypatch: MonkeyPatch, mock_get_model: mock.Mock) -> None: + async def test_preloads_clip_models(self, mocker: MockerFixture) -> None: os.environ["MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL"] = "ViT-B-32__openai" os.environ["MACHINE_LEARNING_PRELOAD__CLIP__VISUAL"] = "ViT-B-32__openai" @@ -1902,22 +1927,16 @@ class TestCache: assert settings.preload is not None assert settings.preload.clip.textual == "ViT-B-32__openai" assert settings.preload.clip.visual == "ViT-B-32__openai" - - model_cache = ModelCache() - monkeypatch.setattr("immich_ml.main.model_cache", model_cache) + get = mocker.patch("immich_ml.main.model_cache.get") await preload_models(settings.preload) - mock_get_model.assert_has_calls( - [ - mock.call("ViT-B-32__openai", ModelType.TEXTUAL, ModelTask.SEARCH), - mock.call("ViT-B-32__openai", ModelType.VISUAL, ModelTask.SEARCH), - ], - any_order=True, - ) - async def test_preloads_facial_recognition_models( - self, monkeypatch: MonkeyPatch, mock_get_model: mock.Mock - ) -> None: + assert {call.args[0] for call in get.call_args_list} >= { + InferenceEntry(OpenClipTextualEncoder, "ViT-B-32__openai", TextualOptions()), + InferenceEntry(OpenClipVisualEncoder, "ViT-B-32__openai", VisualOptions()), + } + + async def test_preloads_facial_recognition_models(self, mocker: MockerFixture) -> None: os.environ["MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__DETECTION"] = "buffalo_s" os.environ["MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__RECOGNITION"] = "buffalo_s" @@ -1925,20 +1944,16 @@ class TestCache: assert settings.preload is not None assert settings.preload.facial_recognition.detection == "buffalo_s" assert settings.preload.facial_recognition.recognition == "buffalo_s" - - model_cache = ModelCache() - monkeypatch.setattr("immich_ml.main.model_cache", model_cache) + get = mocker.patch("immich_ml.main.model_cache.get") await preload_models(settings.preload) - mock_get_model.assert_has_calls( - [ - mock.call("buffalo_s", ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION), - mock.call("buffalo_s", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION), - ], - any_order=True, - ) - async def test_preloads_ocr_models(self, monkeypatch: MonkeyPatch, mock_get_model: mock.Mock) -> None: + assert {call.args[0] for call in get.call_args_list} >= { + InferenceEntry(FaceDetector, "buffalo_s", FaceDetectionOptions()), + InferenceEntry(FaceRecognizer, "buffalo_s", FaceRecognitionOptions()), + } + + async def test_preloads_ocr_models(self, mocker: MockerFixture) -> None: os.environ["MACHINE_LEARNING_PRELOAD__OCR__DETECTION"] = "PP-OCRv5_mobile" os.environ["MACHINE_LEARNING_PRELOAD__OCR__RECOGNITION"] = "PP-OCRv5_mobile" @@ -1946,52 +1961,15 @@ class TestCache: assert settings.preload is not None assert settings.preload.ocr.detection == "PP-OCRv5_mobile" assert settings.preload.ocr.recognition == "PP-OCRv5_mobile" - - model_cache = ModelCache() - monkeypatch.setattr("immich_ml.main.model_cache", model_cache) + get = mocker.patch("immich_ml.main.model_cache.get") await preload_models(settings.preload) - mock_get_model.assert_has_calls( - [ - mock.call("PP-OCRv5_mobile", ModelType.DETECTION, ModelTask.OCR), - mock.call("PP-OCRv5_mobile", ModelType.RECOGNITION, ModelTask.OCR), - ], - 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" - os.environ["MACHINE_LEARNING_PRELOAD__CLIP__VISUAL"] = "ViT-B-32__openai" - os.environ["MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__RECOGNITION"] = "buffalo_s" - os.environ["MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__DETECTION"] = "buffalo_s" - os.environ["MACHINE_LEARNING_PRELOAD__OCR__DETECTION"] = "PP-OCRv5_mobile" - os.environ["MACHINE_LEARNING_PRELOAD__OCR__RECOGNITION"] = "PP-OCRv5_mobile" - - settings = Settings() - assert settings.preload is not None - assert settings.preload.clip.visual == "ViT-B-32__openai" - assert settings.preload.clip.textual == "ViT-B-32__openai" - assert settings.preload.facial_recognition.recognition == "buffalo_s" - assert settings.preload.facial_recognition.detection == "buffalo_s" - assert settings.preload.ocr.detection == "PP-OCRv5_mobile" - assert settings.preload.ocr.recognition == "PP-OCRv5_mobile" - - model_cache = ModelCache() - monkeypatch.setattr("immich_ml.main.model_cache", model_cache) - - await preload_models(settings.preload) - mock_get_model.assert_has_calls( - [ - mock.call("ViT-B-32__openai", ModelType.TEXTUAL, ModelTask.SEARCH), - mock.call("ViT-B-32__openai", ModelType.VISUAL, ModelTask.SEARCH), - mock.call("buffalo_s", ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION), - mock.call("buffalo_s", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION), - mock.call("PP-OCRv5_mobile", ModelType.DETECTION, ModelTask.OCR), - mock.call("PP-OCRv5_mobile", ModelType.RECOGNITION, ModelTask.OCR), - ], - any_order=True, - ) + assert {call.args[0] for call in get.call_args_list} >= { + InferenceEntry(TextDetector, "PP-OCRv5_mobile", TextDetectionOptions(settings.preload.ocr.max_resolution)), + InferenceEntry(TextRecognizer, "PP-OCRv5_mobile", TextRecognitionOptions()), + } + get.return_value.build.assert_called() # so that no request waits on a graph @pytest.mark.asyncio @@ -2001,7 +1979,7 @@ class TestLoad: mock_model.loaded = False mock_model.load_attempts = 0 - res = await load(mock_model) + res: InferenceModel[Any] = await load(mock_model) assert res is mock_model mock_model.load.assert_called_once() @@ -2011,7 +1989,7 @@ class TestLoad: mock_model = mock.Mock(spec=InferenceModel) mock_model.loaded = True - res = await load(mock_model) + res: InferenceModel[Any] = await load(mock_model) assert res is mock_model mock_model.load.assert_not_called() @@ -2025,7 +2003,7 @@ class TestLoad: mock_model.loaded = False mock_model.load_attempts = 0 - res = await load(mock_model) + res: InferenceModel[Any] = 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 @@ -2078,14 +2056,86 @@ class TestLoad: mock_model.clear_cache.assert_not_called() +def test_reads_each_entrys_options_as_its_model_takes_them(tmp_path: Path) -> None: + request = { + "ocr": { + "detection": {"modelName": "PP-OCRv5_mobile", "options": {"minScore": 0.3, "cache_dir": str(tmp_path)}}, + "recognition": {"modelName": "PP-OCRv5_mobile"}, + }, + "clip": {"textual": {"modelName": "ViT-B-32__openai", "options": {"language": "de"}}}, + "facial-recognition": {"recognition": {"modelName": "buffalo_l"}}, + } + + entries = get_entries(json.dumps(request)) + + # what a request leaves out keeps the default, which a preload of the same model shares; what no model takes, + # such as where a model is cached, goes nowhere + assert {entry.options for entry in entries} == { + TextDetectionOptions(max_resolution=736, min_score=0.3), + TextRecognitionOptions(min_score=0.9), + TextualOptions(language="de"), + FaceRecognitionOptions(), + } + + +@pytest.mark.parametrize( + "request_", + [ + {"ocr": {"detection": {"modelName": "PP-OCRv5_mobile", "options": {"minScore": "high"}}}}, + {"clip": {"visual": {"modelName": "buffalo_l"}}}, # a model its slot does not run + {"facial_recognition": {"recognition": {"modelName": "buffalo_l"}}}, + {"clip": {"visaul": {"modelName": "ViT-B-32__openai"}}}, + {"clip": {"visual": {"modelName": "ViT-B-32__openai", "option": {}}}}, + {}, + ], +) +def test_refuses_what_a_model_cannot_take(request_: dict[str, Any]) -> None: + with pytest.raises(HTTPException) as refused: + get_entries(json.dumps(request_)) + + assert refused.value.status_code == 422 + + +@pytest.mark.asyncio +async def test_feeds_a_model_the_output_of_the_one_it_depends_on(mocker: MockerFixture) -> None: + detection = (ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION) + models = { + FaceDetector: mock.Mock(spec=InferenceModel, depends=[], loaded=True), + FaceRecognizer: mock.Mock(spec=InferenceModel, depends=[detection], loaded=True), + } + models[FaceDetector].predict.return_value = "faces" + models[FaceRecognizer].predict.return_value = "embeddings" + mocker.patch("immich_ml.main.model_cache.get", side_effect=lambda entry, ttl: models[entry.model]) + entries: list[InferenceEntry[Any]] = [ # the recognition first, so that only waiting on the detection feeds it + InferenceEntry(FaceRecognizer, "buffalo_l", FaceRecognitionOptions()), + InferenceEntry(FaceDetector, "buffalo_l", FaceDetectionOptions()), + ] + + response = await run_request("image", entries) + + models[FaceRecognizer].predict.assert_called_once_with("image", "faces", options=FaceRecognitionOptions()) + assert response == {"facial-recognition": "embeddings"} + + +@pytest.mark.asyncio +async def test_refuses_a_model_whose_input_the_request_does_not_ask_for(mocker: MockerFixture) -> None: + recognition = mock.Mock(spec=InferenceModel, depends=[FaceDetector.identity], loaded=True) + mocker.patch("immich_ml.main.model_cache.get", return_value=recognition) + entries: list[InferenceEntry[Any]] = [InferenceEntry(FaceRecognizer, "buffalo_l", FaceRecognitionOptions())] + + with pytest.raises(HTTPException) as refused: + await run_request("image", entries) + + assert refused.value.status_code == 400 + recognition.predict.assert_not_called() + + @pytest.mark.asyncio async def test_waits_for_every_entry_before_failing_the_request(mocker: MockerFixture) -> None: finished = threading.Event() - def model(task: ModelTask, load: Callable[[], None]) -> mock.Mock: - stub = mock.Mock( - spec=InferenceModel, depends=[], loaded=False, load_attempts=0, identity=(ModelType.VISUAL, task) - ) + def model(load: Callable[[], None]) -> mock.Mock: + stub = mock.Mock(spec=InferenceModel, depends=[], loaded=False, load_attempts=0) stub.load.side_effect = load return stub @@ -2094,18 +2144,18 @@ async def test_waits_for_every_entry_before_failing_the_request(mocker: MockerFi finished.set() models = { - ModelTask.SEARCH: model(ModelTask.SEARCH, slow), - ModelTask.OCR: model(ModelTask.OCR, mock.Mock(side_effect=RuntimeError("fails at once"))), + OpenClipVisualEncoder: model(slow), + TextDetector: model(mock.Mock(side_effect=RuntimeError("fails at once"))), } - mocker.patch("immich_ml.main.model_cache.get", side_effect=lambda name, type, task, **options: models[task]) - entries: Any = [ - {"name": "a", "task": ModelTask.SEARCH, "type": ModelType.VISUAL, "options": {}}, - {"name": "b", "task": ModelTask.OCR, "type": ModelType.DETECTION, "options": {}}, + mocker.patch("immich_ml.main.model_cache.get", side_effect=lambda entry, ttl: models[entry.model]) + entries: list[InferenceEntry[Any]] = [ + InferenceEntry(OpenClipVisualEncoder, "ViT-B-32__openai", VisualOptions()), + InferenceEntry(TextDetector, "PP-OCRv5_mobile", TextDetectionOptions()), ] with ThreadPoolExecutor(2) as pool, pytest.raises(RuntimeError, match="fails at once"): mocker.patch("immich_ml.main.thread_pool", pool) - await run_request("text", (entries, [])) + await run_request("text", entries) assert finished.is_set() # or it would still be loading with no request counted as active