fix(ml): stricter typing and request validation (#31856)

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:

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,
+    FaceDetectionOptions,
-    InferenceEntry,
+    FaceRecognitionOptions,
     InferenceResponse,
     ModelFormat,
     ModelIdentity,
-    ModelTask,
+    Options,
-    ModelType,
+    TextDetectionOptions,
-    PipelineRequest,
+    TextRecognitionOptions,
-    T,
+    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:
+async def _preload[O: Options](entry: InferenceEntry[O]) -> None:
-        await load_models(preload.clip.textual, ModelType.TEXTUAL, ModelTask.SEARCH)
+    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:
+def _names(setting: str | None) -> list[str]:
-        await load_models(
+    return [] if setting is None else [name.strip() for name in setting.split(",")]
-            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,
-        )
 
 
 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)
+        found = list(PIPELINE_REQUEST.validate_json(entries).entries())
-        without_deps: list[InferenceEntry] = []
+    except (ValidationError, ValueError) as e:
-        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:
         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:
+async def run_inference(payload: Image | str, entries: list[InferenceEntry[Any]]) -> InferenceResponse:
-    outputs: dict[ModelIdentity, Any] = {}
+    read: set[ModelIdentity] = set()
-    response: InferenceResponse = {}
 
-    async def _run_inference(entry: InferenceEntry) -> None:
+    async def _run_inference[O: Options](entry: InferenceEntry[O]) -> Any:
-        model = model_cache.get(entry["name"], entry["type"], entry["task"], ttl=settings.model_ttl, **entry["options"])
+        model = model_cache.get(entry, ttl=settings.model_ttl)
-        inputs = [payload]
+        read.update(model.depends)
-        for dep in model.depends:
+        if missing := [dep for dep in model.depends if dep not in runs]:
-            try:
+            raise HTTPException(400, f"{entry.model.__name__} depends on output of {missing[0]}")
-                inputs.append(outputs[dep])
+        inputs = [payload, *[await runs[dep] for dep in model.depends]]
-            except KeyError:
-                message = f"Task {entry['task']} of type {entry['type']} depends on output of {dep}"
-                raise HTTPException(400, message)
         model = await load(model)
-        output = await attempt(model, partial(model.predict, *inputs, **entry["options"]), MODEL_FILE_ERRORS)
+        return await attempt(model, partial(model.predict, *inputs, options=entry.options), MODEL_FILE_ERRORS)
-        outputs[model.identity] = output
-        response[entry["task"]] = output
 
-    for stage in entries:  # those that depend on another's output run after it
+    runs = {entry.model.identity: asyncio.create_task(_run_inference(entry)) for entry in entries}
-        results = await asyncio.gather(*[_run_inference(entry) for entry in stage], return_exceptions=True)
+    outputs = await asyncio.gather(*runs.values(), return_exceptions=True)
-        for result in results:
+    for output in outputs:
-            if isinstance(result, BaseException):
+        if isinstance(output, BaseException):
-                raise result
+            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)
+        return func()
-    partial_func = partial(func, *args, **kwargs)
+    return await asyncio.get_running_loop().run_in_executor(thread_pool, func)
-    return await asyncio.get_running_loop().run_in_executor(thread_pool, partial_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:
+async def attempt[O: Options, R](
-    def _attempt() -> T:
+    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:

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

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
+    sources: ClassVar[tuple[ModelSource, ...]]
-    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
@@ -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})"

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 InferenceEntry, InferenceModel
-from immich_ml.models.base import 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._models: dict[Key, InferenceModel[Any]] = {}
-        self._expiries: dict[tuple[Any, ...], asyncio.TimerHandle] = {}
+        self._expiries: dict[Key, asyncio.TimerHandle] = {}
 
-    def get(
+    def get[O: Options](self, entry: InferenceEntry[O], ttl: int | None = None) -> InferenceModel[O]:
-        self, model_name: str, model_type: ModelType, model_task: ModelTask, ttl: int | None = None, **options: Any
+        key = (entry.model, entry.name, entry.model.graph(entry.options))
-    ) -> 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))
         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]

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:
+    def _predict(self, inputs: str, options: TextualOptions) -> str:
-        tokens: dict[str, NDArray[Any]] = self.tokenize(inputs, language=language)
+        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)

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

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)
 

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 []

machine-learning/immich_ml/models/ocr/detection.py

@@ -1,5 +1,5 @@
-from collections.abc import Sequence
+from collections.abc import Hashable, Sequence
-from typing import Any
+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))
+        canvases = tuple(Shape(batch=1, **canvas) for canvas in ocr_canvases(max_resolution))
-        self.shape_policy = ShapePolicy(dims=canvases, label=f"res{short_side}")
+        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(
+    @classmethod
-        self, inputs: Image.Image, maxResolution: int = 736, minScore: float = 0.5, scoreMode: str = "fast"
+    def graph(cls, options: TextDetectionOptions) -> Hashable:
-    ) -> TextDetectionOutput:
+        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)

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):
+class TextModel[O: Options](InferenceModel[O]):
-    def __init__(self, model_name: str, **model_kwargs: Any) -> None:
+    sources = (ModelSource.PADDLE,)
-        if settings.legacy_models:
+
-            model_kwargs["model_format"] = ModelFormat.ONNX  # the older exports come in no other format
+    @property
-        super().__init__(model_name, **model_kwargs)
+    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:

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),

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)

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):
+class Options:
-    modelName: str
+    __pydantic_config__ = ConfigDict(extra="ignore")  # so that an older service ignores an option a newer server adds
-    options: dict[str, Any]
 
 
-PipelineRequest = dict[ModelTask, dict[ModelType, PipelineEntry]]
+@dataclass(frozen=True)
+class VisualOptions(Options):
+    pass
 
 
-class InferenceEntry(TypedDict):
+@dataclass(frozen=True)
-    name: str
+class TextualOptions(Options):
-    task: ModelTask
+    language: str | None = None
-    type: ModelType
-    options: dict[str, Any]
 
 
-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]

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 = [

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 (
-        assert face_detector.predict(Image.new("RGB", (640, 640)), minScore=0.4)["boxes"].shape[0] == 2
+            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, TextDetectionOptions(min_score=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.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:
+    async def test_caches(self) -> None:
-        model_cache = ModelCache()
+        model_cache, model = ModelCache(), stub_model()
-        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION)
+        model_cache.get(InferenceEntry(model, "buffalo_l", FaceRecognitionOptions()))
-        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION)
+        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):
+        for options in (
-            model_cache.get("PP-OCRv5_mobile", ModelType.DETECTION, ModelTask.OCR, maxResolution=max_resolution)
+            TextDetectionOptions(),
-        # maxResolution picks the graphs a detector builds, so instances cannot share them
+            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:
+    async def test_creates_the_model_from_the_request_options(self) -> None:
-        model_cache = ModelCache()
+        model = stub_model()
-        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, cache_dir="test_cache")
+        ModelCache().get(InferenceEntry(model, "buffalo_l", FaceDetectionOptions(min_score=0.3)))
-        mock_get_model.return_value.assert_called_once_with("test_model_name", cache_dir="test_cache")
+        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)
+        visual = model_cache.get(InferenceEntry(OpenClipVisualEncoder, "ViT-B-32__openai", VisualOptions()))
-        model_cache.get("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH)
+        text = model_cache.get(InferenceEntry(OpenClipTextualEncoder, "ViT-B-32__openai", TextualOptions()))
-        assert mock_get_model.call_args_list == [
+        assert isinstance(visual, OpenClipVisualEncoder) and isinstance(text, OpenClipTextualEncoder)
-            mock.call("test_model_name", ModelType.VISUAL, ModelTask.SEARCH),
-            mock.call("test_model_name", ModelType.TEXTUAL, ModelTask.SEARCH),
-        ]
         assert len(model_cache._models) == 2
 
-    async def test_lets_go_of_a_model_unused_for_its_ttl_and_returns_its_memory(
+    async def test_lets_go_of_a_model_unused_for_its_ttl_and_returns_its_memory(self, mocker: MockerFixture) -> None:
-        self, mock_get_model: mock.Mock, mocker: MockerFixture
-    ) -> None:
         events: list[str] = []
 
-        class Model:
+        class Loaded:
-            graph_options = ()
-
-            def __init__(self, *args: Any, **kwargs: Any) -> None:
-                pass
-
             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(entry, ttl=100)
-        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, 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(
+    async def test_keeps_a_preloaded_model_whatever_ttl_a_request_names(self, mocker: MockerFixture) -> None:
-        self, mock_get_model: mock.Mock, 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(entry)
-        model_cache.get("test_model_name", ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION, ttl=100)
+        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:
+    async def test_refuses_a_model_its_slot_does_not_run(self) -> None:
-        invalid: Any = SimpleNamespace(value="invalid")
-        model_cache = ModelCache()
-
         with pytest.raises(ValueError):
-            model_cache.get("XLM-Roberta-Large-Vit-B-32", ModelType.TEXTUAL, invalid)
+            InferenceEntry(OpenClipVisualEncoder, "buffalo_l", VisualOptions())
-
-    async def test_raises_exception_if_unknown_model_name(self) -> None:
-        model_cache = ModelCache()
 
+    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"
-
+        get = mocker.patch("immich_ml.main.model_cache.get")
-        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),
-            ],
-            any_order=True,
-        )
 
-    async def test_preloads_facial_recognition_models(
+        assert {call.args[0] for call in get.call_args_list} >= {
-        self, monkeypatch: MonkeyPatch, mock_get_model: mock.Mock
+            InferenceEntry(OpenClipTextualEncoder, "ViT-B-32__openai", TextualOptions()),
-    ) -> None:
+            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"
-
+        get = mocker.patch("immich_ml.main.model_cache.get")
-        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("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"
-
+        get = mocker.patch("immich_ml.main.model_cache.get")
-        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("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:
+        assert {call.args[0] for call in get.call_args_list} >= {
-        os.environ["MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL"] = "ViT-B-32__openai"
+            InferenceEntry(TextDetector, "PP-OCRv5_mobile", TextDetectionOptions(settings.preload.ocr.max_resolution)),
-        os.environ["MACHINE_LEARNING_PRELOAD__CLIP__VISUAL"] = "ViT-B-32__openai"
+            InferenceEntry(TextRecognizer, "PP-OCRv5_mobile", TextRecognitionOptions()),
-        os.environ["MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__RECOGNITION"] = "buffalo_s"
+        }
-        os.environ["MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__DETECTION"] = "buffalo_s"
+        get.return_value.build.assert_called()  # so that no request waits on a graph
-        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,
-        )
 
 
 @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:
+    def model(load: Callable[[], None]) -> mock.Mock:
-        stub = mock.Mock(
+        stub = mock.Mock(spec=InferenceModel, depends=[], loaded=False, load_attempts=0)
-            spec=InferenceModel, depends=[], loaded=False, load_attempts=0, identity=(ModelType.VISUAL, task)
-        )
         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),
+        OpenClipVisualEncoder: model(slow),
-        ModelTask.OCR: model(ModelTask.OCR, mock.Mock(side_effect=RuntimeError("fails at once"))),
+        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])
+    mocker.patch("immich_ml.main.model_cache.get", side_effect=lambda entry, ttl: models[entry.model])
-    entries: Any = [
+    entries: list[InferenceEntry[Any]] = [
-        {"name": "a", "task": ModelTask.SEARCH, "type": ModelType.VISUAL, "options": {}},
+        InferenceEntry(OpenClipVisualEncoder, "ViT-B-32__openai", VisualOptions()),
-        {"name": "b", "task": ModelTask.OCR, "type": ModelType.DETECTION, "options": {}},
+        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